# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Bit-exact data-movement kernels for the Wan causal VAE decoder.

Adapted from SGLang PR #34125:
https://github.com/sgl-project/sglang/pull/34125

The kernels in this module only move values, fill zeros, or perform the same
two-operand addition as the PyTorch reference path. They preserve the input
layout instead of selecting a new Conv3d layout, so enabling them does not
change the decoder's cuDNN accumulation order.
"""

from __future__ import annotations

import torch
import triton
import triton.language as tl

_MAX_INT32 = 2**31 - 1
_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)


@triton.jit
def _cat_pad_5d_kernel(
    x_ptr,
    cache_ptr,
    out_ptr,
    keep_ptr,
    total,
    channels,
    frames,
    height,
    width,
    cache_frames,
    out_frames,
    out_height,
    out_width,
    zero_front_frames,
    pad_height,
    pad_width,
    stride_x_batch,
    stride_x_channel,
    stride_x_frame,
    stride_x_height,
    stride_x_width,
    stride_cache_batch,
    stride_cache_channel,
    stride_cache_frame,
    stride_cache_height,
    stride_cache_width,
    has_cache: tl.constexpr,
    keep_frames_const: tl.constexpr,
    channels_inner: tl.constexpr,
    index_64bit: tl.constexpr,
    block_size_const: tl.constexpr,
):
    if index_64bit:
        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
    else:
        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
    mask = offsets < total

    if channels_inner:
        out_channel = offsets % channels
        remaining = offsets // channels
        out_width_idx = remaining % out_width
        remaining = remaining // out_width
        out_height_idx = remaining % out_height
        remaining = remaining // out_height
        out_frame = remaining % out_frames
        out_batch = remaining // out_frames
    else:
        out_width_idx = offsets % out_width
        remaining = offsets // out_width
        out_height_idx = remaining % out_height
        remaining = remaining // out_height
        out_frame = remaining % out_frames
        remaining = remaining // out_frames
        out_channel = remaining % channels
        out_batch = remaining // channels

    input_width = out_width_idx - pad_width
    input_height = out_height_idx - pad_height
    input_frame = out_frame - zero_front_frames

    spatial_valid = (input_width >= 0) & (input_width < width) & (input_height >= 0) & (input_height < height)
    from_cache = spatial_valid & (input_frame >= 0) & (input_frame < cache_frames)
    from_x = spatial_valid & (input_frame >= cache_frames) & (input_frame < cache_frames + frames)

    x_frame = input_frame - cache_frames
    x_offset = (
        out_batch * stride_x_batch
        + out_channel * stride_x_channel
        + x_frame * stride_x_frame
        + input_height * stride_x_height
        + input_width * stride_x_width
    )
    values = tl.load(x_ptr + x_offset, mask=mask & from_x, other=0.0)
    if has_cache:
        cache_offset = (
            out_batch * stride_cache_batch
            + out_channel * stride_cache_channel
            + input_frame * stride_cache_frame
            + input_height * stride_cache_height
            + input_width * stride_cache_width
        )
        cache_values = tl.load(cache_ptr + cache_offset, mask=mask & from_cache, other=0.0)
        values = tl.where(from_cache, cache_values, values)
    tl.store(out_ptr + offsets, values, mask=mask)

    if keep_frames_const > 0:
        keep_frame = out_frame - (out_frames - keep_frames_const)
        keep_mask = mask & spatial_valid & (keep_frame >= 0)
        if channels_inner:
            keep_offset = (
                ((out_batch * keep_frames_const + keep_frame) * height + input_height) * width + input_width
            ) * channels + out_channel
        else:
            keep_offset = (
                ((out_batch * channels + out_channel) * keep_frames_const + keep_frame) * height + input_height
            ) * width + input_width
        tl.store(keep_ptr + keep_offset, values, mask=keep_mask)


def _uses_channels_last_3d(tensor: torch.Tensor) -> bool:
    return tensor.shape[1] > 1 and tensor.stride(1) == 1


def cat_pad_5d(
    x: torch.Tensor,
    cache_x: torch.Tensor | None,
    padding: list[int] | tuple[int, ...],
    keep_cache_frames: int = 0,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None:
    """Fuse temporal concat, constant padding, and optional cache refresh.

    ``padding`` follows the Wan causal Conv3d convention
    ``(w_left, w_right, h_top, h_bottom, t_front, t_back)``. The output uses
    the same dense layout that the reference ``cat``/``pad`` path selects.
    Unsupported inputs return ``None`` so callers can execute the reference
    implementation.
    """
    if len(padding) != 6:
        return None
    pad_width_left, pad_width_right, pad_height_top, pad_height_bottom, pad_front, pad_back = padding
    if any(value < 0 for value in padding):
        return None
    if pad_width_left != pad_width_right or pad_height_top != pad_height_bottom or pad_back != 0:
        return None
    if x.dim() != 5 or not x.is_cuda or x.dtype not in _SUPPORTED_DTYPES:
        return None
    if keep_cache_frames < 0:
        return None

    channels_inner = _uses_channels_last_3d(x)
    cache_frames = 0
    if cache_x is not None:
        if (
            cache_x.dim() != 5
            or cache_x.dtype != x.dtype
            or cache_x.device != x.device
            or cache_x.shape[0] != x.shape[0]
            or cache_x.shape[1] != x.shape[1]
            or cache_x.shape[3:] != x.shape[3:]
            or _uses_channels_last_3d(cache_x) != channels_inner
        ):
            return None
        cache_frames = cache_x.shape[2]

    zero_front_frames = pad_front - cache_frames
    if zero_front_frames < 0:
        return None

    batch, channels, frames, height, width = x.shape
    out_frames = pad_front + frames
    out_height = height + 2 * pad_height_top
    out_width = width + 2 * pad_width_left
    keep_frames = min(keep_cache_frames, out_frames)
    memory_format = torch.channels_last_3d if channels_inner else torch.contiguous_format
    out = torch.empty(
        (batch, channels, out_frames, out_height, out_width),
        device=x.device,
        dtype=x.dtype,
        memory_format=memory_format,
    )
    total = out.numel()
    if total == 0 or total > _MAX_INT32 * 4:
        return None

    if keep_frames:
        keep_arg = torch.empty(
            (batch, channels, keep_frames, height, width),
            device=x.device,
            dtype=x.dtype,
            memory_format=memory_format,
        )
    else:
        keep_arg = out

    if cache_x is None:
        cache_arg = x
        cache_strides = (0, 0, 0, 0, 0)
    else:
        cache_arg = cache_x
        cache_strides = cache_x.stride()

    block_size = 512
    grid = (triton.cdiv(total, block_size),)
    with torch.get_device_module().device(x.device):
        _cat_pad_5d_kernel[grid](
            x,
            cache_arg,
            out,
            keep_arg,
            total,
            channels,
            frames,
            height,
            width,
            cache_frames,
            out_frames,
            out_height,
            out_width,
            zero_front_frames,
            pad_height_top,
            pad_width_left,
            *x.stride(),
            *cache_strides,
            has_cache=cache_x is not None,
            keep_frames_const=keep_frames,
            channels_inner=channels_inner,
            index_64bit=total >= _MAX_INT32,
            block_size_const=block_size,
        )

    if keep_cache_frames:
        return out, keep_arg
    return out


@triton.jit
def _dup_up3d_add_kernel(
    main_ptr,
    source_ptr,
    out_ptr,
    total,
    out_channels,
    out_frames,
    out_height,
    out_width,
    frame_offset,
    stride_main_batch,
    stride_main_channel,
    stride_main_frame,
    stride_main_height,
    stride_main_width,
    stride_source_batch,
    stride_source_channel,
    stride_source_frame,
    stride_source_height,
    stride_source_width,
    stride_out_batch,
    stride_out_channel,
    stride_out_frame,
    stride_out_height,
    stride_out_width,
    factor_temporal_const: tl.constexpr,
    factor_spatial_const: tl.constexpr,
    repeats_const: tl.constexpr,
    channels_inner: tl.constexpr,
    index_64bit: tl.constexpr,
    block_size_const: tl.constexpr,
):
    if index_64bit:
        offsets = tl.program_id(0).to(tl.int64) * block_size_const + tl.arange(0, block_size_const).to(tl.int64)
    else:
        offsets = tl.program_id(0) * block_size_const + tl.arange(0, block_size_const)
    mask = offsets < total

    if channels_inner:
        out_channel = offsets % out_channels
        remaining = offsets // out_channels
        out_width_idx = remaining % out_width
        remaining = remaining // out_width
        out_height_idx = remaining % out_height
        remaining = remaining // out_height
        out_frame = remaining % out_frames
        out_batch = remaining // out_frames
    else:
        out_width_idx = offsets % out_width
        remaining = offsets // out_width
        out_height_idx = remaining % out_height
        remaining = remaining // out_height
        out_frame = remaining % out_frames
        remaining = remaining // out_frames
        out_channel = remaining % out_channels
        out_batch = remaining // out_channels

    unsliced_frame = out_frame + frame_offset
    source_frame = unsliced_frame // factor_temporal_const
    temporal_remainder = unsliced_frame % factor_temporal_const
    source_height = out_height_idx // factor_spatial_const
    height_remainder = out_height_idx % factor_spatial_const
    source_width = out_width_idx // factor_spatial_const
    width_remainder = out_width_idx % factor_spatial_const
    repeated_channel = (
        (out_channel * factor_temporal_const + temporal_remainder) * factor_spatial_const + height_remainder
    ) * factor_spatial_const + width_remainder
    source_channel = repeated_channel // repeats_const

    main_offset = (
        out_batch * stride_main_batch
        + out_channel * stride_main_channel
        + out_frame * stride_main_frame
        + out_height_idx * stride_main_height
        + out_width_idx * stride_main_width
    )
    source_offset = (
        out_batch * stride_source_batch
        + source_channel * stride_source_channel
        + source_frame * stride_source_frame
        + source_height * stride_source_height
        + source_width * stride_source_width
    )
    out_offset = (
        out_batch * stride_out_batch
        + out_channel * stride_out_channel
        + out_frame * stride_out_frame
        + out_height_idx * stride_out_height
        + out_width_idx * stride_out_width
    )
    main_value = tl.load(main_ptr + main_offset, mask=mask, other=0.0)
    source_value = tl.load(source_ptr + source_offset, mask=mask, other=0.0)
    value = main_value.to(tl.float32) + source_value.to(tl.float32)
    tl.store(out_ptr + out_offset, value, mask=mask)


def dup_up3d_add(
    main: torch.Tensor,
    source: torch.Tensor,
    factor_temporal: int,
    factor_spatial: int,
    repeats: int,
    drop_first_frames: bool,
) -> torch.Tensor | None:
    """Evaluate ``main + DupUp3D(source)`` without materializing the shortcut."""
    if main.dim() != 5 or source.dim() != 5:
        return None
    if factor_temporal <= 0 or factor_spatial <= 0 or repeats <= 0:
        return None
    if factor_temporal & (factor_temporal - 1) or factor_spatial & (factor_spatial - 1):
        return None
    if repeats & (repeats - 1):
        return None
    if not main.is_cuda or not source.is_cuda:
        return None
    if main.dtype != source.dtype or main.dtype not in _SUPPORTED_DTYPES or main.device != source.device:
        return None

    batch, source_channels, frames, height, width = source.shape
    frame_offset = factor_temporal - 1 if drop_first_frames else 0
    numerator = source_channels * repeats
    denominator = factor_temporal * factor_spatial * factor_spatial
    if numerator % denominator:
        return None
    expected_shape = (
        batch,
        numerator // denominator,
        frames * factor_temporal - frame_offset,
        height * factor_spatial,
        width * factor_spatial,
    )
    if tuple(main.shape) != expected_shape:
        return None

    out = torch.empty_like(main)
    total = out.numel()
    if total == 0 or total > _MAX_INT32 * 4:
        return None

    block_size = 512
    grid = (triton.cdiv(total, block_size),)
    with torch.get_device_module().device(main.device):
        _dup_up3d_add_kernel[grid](
            main,
            source,
            out,
            total,
            expected_shape[1],
            expected_shape[2],
            expected_shape[3],
            expected_shape[4],
            frame_offset,
            *main.stride(),
            *source.stride(),
            *out.stride(),
            factor_temporal_const=factor_temporal,
            factor_spatial_const=factor_spatial,
            repeats_const=repeats,
            channels_inner=out.stride(1) == 1 and expected_shape[1] > 1,
            index_64bit=total >= _MAX_INT32,
            block_size_const=block_size,
        )
    return out