# mypy: allow-untyped-defs
from __future__ import annotations

import logging
from typing import TYPE_CHECKING, TypedDict

import torch
from torch._inductor.codegen.rocm.ck_conv_template import CKGroupedConvFwdTemplate

from .. import config, ir
from ..lowering import (
    add_layout_constraint,
    constrain_to_fx_strides,
    fallback_handler,
    lowerings as L,
    register_lowering,
)
from ..select_algorithm import (
    autotune_select_algorithm,
    ExternKernelChoice,
    SymbolicGridFn,
    TritonTemplate,
)
from ..utils import (
    is_ones,
    is_zeros,
    pad_listlike,
    sympy_product,
    use_ck_conv_template,
    use_triton_template,
)
from ..virtualized import V
from .mm_common import load_kernel_template


if TYPE_CHECKING:
    from collections.abc import Sequence

    from ..ir import TensorBox

log = logging.getLogger(__name__)


aten = torch.ops.aten


@SymbolicGridFn
def conv2d_grid(n, c, h, w, meta, *, cdiv):
    return (
        cdiv(n * h * w, meta["BLOCK_M"]),
        cdiv(c, meta["BLOCK_N"]),
        meta["GROUPS"],
    )


@SymbolicGridFn
def conv3d_grid(n, c, d, h, w, meta, *, cdiv):
    return (
        cdiv(n * d * h * w, meta["BLOCK_M"]),
        cdiv(c, meta["BLOCK_N"]),
        meta["GROUPS"],
    )


# =============================================================================
# Depthwise conv1d (groups == in_channels == out_channels)
# Uses direct element-wise multiply-accumulate instead of implicit GEMM.
# Channels-last (NLC) layout with 3D tiling: BLOCK_N x BLOCK_L x BLOCK_C.
# =============================================================================


@SymbolicGridFn
def depthwise_conv1d_grid(n, c, l, meta, *, cdiv):
    return (
        cdiv(n, meta["BLOCK_N"]),
        cdiv(l, meta["BLOCK_L"]),
        cdiv(c, meta["BLOCK_C"]),
    )


depthwise_conv1d_template = TritonTemplate(
    name="depthwise_conv1d",
    grid=depthwise_conv1d_grid,
    source=load_kernel_template("triton_depthwise_conv"),
    cache_codegen_enabled_for_template=True,
)

LOOP_BODY_2D = """
        idx_x_h = i - PADDING_H + idx_y_h * STRIDE_H
        idx_x_w = j - PADDING_W + idx_y_w * STRIDE_W
        idx_x_c = tl.arange(0, BLOCK_K) + k

        x_ptrs = x_base + (
            (idx_x_h * stride_xh)[:, None]
            + (idx_x_w * stride_xw)[:, None]
            + (idx_x_c * stride_xc)[None, :]
        )
        mask_x = (
            (idx_n < BATCH)[:, None]
            & (idx_x_h >= 0)[:, None]
            & (idx_x_h < IN_H)[:, None]
            & (idx_x_w >= 0)[:, None]
            & (idx_x_w < IN_W)[:, None]
            & (idx_x_c < GROUP_IN_C)[None, :]
        )
        matrix_x = tl.load(x_ptrs, mask=mask_x, other=0.0)

        w_ptrs = w_base + (
            (idx_x_c * stride_wc_in)[:, None] + (i * stride_wh) + (j * stride_ww)
        )
        mask_w = (idx_x_c[:, None] < GROUP_IN_C) & (idx_y_c[None, :] < GROUP_OUT_C)
        matrix_w = tl.load(w_ptrs, mask=mask_w, other=0.0)
        acc += tl.dot(matrix_x, matrix_w, allow_tf32=ALLOW_TF32)
"""

"""
This is a relatively simple conv implementation that can likely be
improved.  Many alternate conv versions can be found here:
https://github.com/pytorch/torchdynamo/pull/971
"""
conv2d_template = TritonTemplate(
    name="convolution2d",
    grid=conv2d_grid,
    source=r"""
{{def_kernel("X", "W")}}
    # Tensor dimensions
    BATCH = {{size("X", 0)}}
    IN_C = {{size("X", 1)}}
    IN_H = {{size("X", 2)}}
    IN_W = {{size("X", 3)}}
    OUT_C = {{size(None, 1)}}
    OUT_H = {{size(None, 2)}}
    OUT_W = {{size(None, 3)}}

    # Strides:
    stride_xn = {{stride("X", 0)}}
    stride_xc = {{stride("X", 1)}}
    stride_xh = {{stride("X", 2)}}
    stride_xw = {{stride("X", 3)}}
    stride_wc_out = {{stride("W", 0)}}
    stride_wc_in = {{stride("W", 1)}}
    stride_wh = {{stride("W", 2)}}
    stride_ww = {{stride("W", 3)}}

    nhw = tl.program_id(0).to(INDEX_DTYPE) * BLOCK_M + tl.arange(0, BLOCK_M)
    idx_y_w = nhw % OUT_W
    nh = nhw // OUT_W
    idx_y_h = nh % OUT_H
    idx_n = nh // OUT_H
    idx_y_c = tl.program_id(1).to(INDEX_DTYPE) * BLOCK_N + tl.arange(0, BLOCK_N)

{% if GROUPS == 1 %}
    group = 0
    GROUP_IN_C = IN_C
    GROUP_OUT_C = OUT_C
{% else %}
    group = tl.program_id(2).to(INDEX_DTYPE)
    GROUP_IN_C = IN_C // GROUPS
    GROUP_OUT_C = OUT_C // GROUPS
{% endif %}

    x_base = X + (group * stride_xc * GROUP_IN_C + idx_n * stride_xn)[:, None]
    w_base = (
        W + (group * stride_wc_out * GROUP_OUT_C + idx_y_c * stride_wc_out)[None, :]
    )

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

{% if UNROLL %}
{% for i in range(KERNEL_H) %}
{% for j in range(KERNEL_W) %}
    i = {{i}}
    j = {{j}}
    for k in range(0, GROUP_IN_C, BLOCK_K):
        """
    + LOOP_BODY_2D
    + """
{% endfor %}
{% endfor %}
{% else %}
    # Could be simplified, but slightly slower:
    # for i in range(KERNEL_H):
    #     for j in range(KERNEL_W):
    #         for k in range(0, GROUP_IN_C, BLOCK_K):
    BLOCK_K_COUNT = (GROUP_IN_C + BLOCK_K - 1) // BLOCK_K
    for ijk in range(KERNEL_H * KERNEL_W * BLOCK_K_COUNT):
        k = (ijk % BLOCK_K_COUNT) * BLOCK_K
        ij = ijk // BLOCK_K_COUNT
        i = ij // KERNEL_W
        j = ij % KERNEL_W
        """
    + LOOP_BODY_2D
    + """
{% endif %}

    mask = (
        (idx_n < BATCH)[:, None]
        & (idx_y_h < OUT_H)[:, None]
        & (idx_y_w < OUT_W)[:, None]
        & (idx_y_c < GROUP_OUT_C)[None, :]
    )
    idx_n = idx_n[:, None]
    idx_c = idx_y_c[None, :] + group * GROUP_OUT_C
    idx_h = idx_y_h[:, None]
    idx_w = idx_y_w[:, None]

    # inductor generates a suffix
    {{store_output(("idx_n", "idx_c", "idx_h", "idx_w"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"))}}
""",
)

LOOP_BODY_3D = """
        idx_x_d = d - PADDING_D + idx_y_d * STRIDE_D
        idx_x_h = i - PADDING_H + idx_y_h * STRIDE_H
        idx_x_w = j - PADDING_W + idx_y_w * STRIDE_W
        idx_x_c = tl.arange(0, BLOCK_K) + k

        x_ptrs = x_base + (
            (idx_x_d * stride_xd)[:, None]
            + (idx_x_h * stride_xh)[:, None]
            + (idx_x_w * stride_xw)[:, None]
            + (idx_x_c * stride_xc)[None, :]
        )
        mask_x = (
            (idx_n < BATCH)[:, None]
            & (idx_x_d >= 0)[:, None]
            & (idx_x_d < IN_D)[:, None]
            & (idx_x_h >= 0)[:, None]
            & (idx_x_h < IN_H)[:, None]
            & (idx_x_w >= 0)[:, None]
            & (idx_x_w < IN_W)[:, None]
            & (idx_x_c < GROUP_IN_C)[None, :]
        )
        matrix_x = tl.load(x_ptrs, mask=mask_x, other=0.0)

        w_ptrs = w_base + (
            (idx_x_c * stride_wc_in)[:, None] +
            (d * stride_wd) + (i * stride_wh) + (j * stride_ww)
        )
        mask_w = (idx_x_c[:, None] < GROUP_IN_C) & (idx_y_c[None, :] < GROUP_OUT_C)
        matrix_w = tl.load(w_ptrs, mask=mask_w, other=0.0)
        acc += tl.dot(matrix_x, matrix_w, allow_tf32=ALLOW_TF32)
"""

conv3d_template = TritonTemplate(
    name="convolution3d",
    grid=conv3d_grid,
    source=r"""
{{def_kernel("X", "W")}}
    # Tensor dimensions
    BATCH = {{size("X", 0)}}
    IN_C = {{size("X", 1)}}
    IN_D = {{size("X", 2)}}
    IN_H = {{size("X", 3)}}
    IN_W = {{size("X", 4)}}
    OUT_C = {{size(None, 1)}}
    OUT_D = {{size(None, 2)}}
    OUT_H = {{size(None, 3)}}
    OUT_W = {{size(None, 4)}}

    # Strides:
    stride_xn = {{stride("X", 0)}}
    stride_xc = {{stride("X", 1)}}
    stride_xd = {{stride("X", 2)}}
    stride_xh = {{stride("X", 3)}}
    stride_xw = {{stride("X", 4)}}
    stride_wc_out = {{stride("W", 0)}}
    stride_wc_in = {{stride("W", 1)}}
    stride_wd = {{stride("W", 2)}}
    stride_wh = {{stride("W", 3)}}
    stride_ww = {{stride("W", 4)}}

    ndhw = tl.program_id(0).to(INDEX_DTYPE) * BLOCK_M + tl.arange(0, BLOCK_M)
    idx_y_w = ndhw % OUT_W
    ndh = ndhw // OUT_W
    idx_y_h = ndh % OUT_H
    nd = ndh // OUT_H
    idx_y_d = nd % OUT_D
    idx_n = nd // OUT_D
    idx_y_c = tl.program_id(1).to(INDEX_DTYPE) * BLOCK_N + tl.arange(0, BLOCK_N)

{% if GROUPS == 1 %}
    group = 0
    GROUP_IN_C = IN_C
    GROUP_OUT_C = OUT_C
{% else %}
    group = tl.program_id(2).to(INDEX_DTYPE)
    GROUP_IN_C = IN_C // GROUPS
    GROUP_OUT_C = OUT_C // GROUPS
{% endif %}

    x_base = X + (group * stride_xc * GROUP_IN_C + idx_n * stride_xn)[:, None]
    w_base = (
        W + (group * stride_wc_out * GROUP_OUT_C + idx_y_c * stride_wc_out)[None, :]
    )

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)

{% if UNROLL %}
{% for d in range(KERNEL_D) %}
{% for i in range(KERNEL_H) %}
{% for j in range(KERNEL_W) %}
    d = {{d}}
    i = {{i}}
    j = {{j}}
    for k in range(0, GROUP_IN_C, BLOCK_K):
        """
    + LOOP_BODY_3D
    + """
{% endfor %}
{% endfor %}
{% endfor %}
{% else %}
    # Could be simplified, but slightly slower:
    # for d in range(KERNEL_D):
    #   for i in range(KERNEL_H):
    #     for j in range(KERNEL_W):
    #         for k in range(0, GROUP_IN_C, BLOCK_K):
    BLOCK_K_COUNT = (GROUP_IN_C + BLOCK_K - 1) // BLOCK_K
    for dijk in range(KERNEL_D * KERNEL_H * KERNEL_W * BLOCK_K_COUNT):
        k = (dijk % BLOCK_K_COUNT) * BLOCK_K
        dij = dijk // BLOCK_K_COUNT
        j = dij % KERNEL_W
        di = dij // KERNEL_W
        i = di % KERNEL_H
        d = di // KERNEL_H
        """
    + LOOP_BODY_3D
    + """
{% endif %}

    mask = (
        (idx_n < BATCH)[:, None]
        & (idx_y_d < OUT_D)[:, None]
        & (idx_y_h < OUT_H)[:, None]
        & (idx_y_w < OUT_W)[:, None]
        & (idx_y_c < GROUP_OUT_C)[None, :]
    )
    idx_n = idx_n[:, None]
    idx_c = idx_y_c[None, :] + group * GROUP_OUT_C
    idx_d = idx_y_d[:, None]
    idx_h = idx_y_h[:, None]
    idx_w = idx_y_w[:, None]

    # inductor generates a suffix
    {{store_output(("idx_n", "idx_c", "idx_d", "idx_h", "idx_w"), "acc", "mask", val_shape=("BLOCK_M", "BLOCK_N"))}}
""",
)

aten_convolution = ExternKernelChoice(
    torch.convolution,
    "at::convolution",
    has_out_variant=False,
    op_overload=aten.convolution.default,
)


def conv1x1_via_mm(x, w, *, out):
    w = torch.squeeze(torch.squeeze(w, -1), -1)
    return torch.matmul(
        x.permute(0, 2, 3, 1), w.permute(1, 0), out=out.permute(0, 2, 3, 1)
    )


aten_conv1x1_via_mm = ExternKernelChoice(conv1x1_via_mm, None)


class ConvLayoutParams(TypedDict):
    stride: tuple[int, ...]
    padding: tuple[int, ...]
    dilation: tuple[int, ...]
    transposed: bool
    output_padding: tuple[int, ...]
    groups: int


def conv_layout(
    x: TensorBox,
    weight: TensorBox,
    bias: TensorBox | None,
    stride: Sequence[int],
    padding: tuple[int, ...],
    dilation: tuple[int, ...],
    transposed: bool,
    output_padding: tuple[int, ...],
    groups: int,
) -> ir.Layout:
    """Determine output layout for a convolution"""
    # We use guard_int_seq rather than size_hints because the output shape
    # depends on these values — if they ever contained symbols, size_hints
    # would silently substitute a hint that could be wrong, producing an
    # incorrect layout. guard_int_seq will install a proper guard instead.
    # Note: stride and padding are already guarded via guard_int_seq in
    # convolution() above, but we guard all four here so conv_layout is
    # self-contained and doesn't rely on callers.
    guard = V.graph.sizevars.guard_int_seq
    with V.graph.fake_mode:
        output = torch.ops.aten.convolution(
            ir.ir_node_to_tensor(x),
            ir.ir_node_to_tensor(weight),
            ir.ir_node_to_tensor(bias),
            guard(stride),
            guard(padding),
            guard(dilation),
            transposed,
            guard(output_padding),
            groups,
        )
        sizes = ir.convert_shape_to_inductor(output.size())
        stride = ir.convert_shape_to_inductor(output.stride())  # type: ignore[assignment]

    return ir.FixedLayout(
        x.get_device_or_error(),
        x.get_dtype(),
        sizes,
        stride,
    )


def channels_last_order(rank):
    order = list(reversed(range(rank)))
    order.insert(1, order.pop(-1))
    return order


def convert_1x1_conv_to_mm(x, weight, bias):
    # special case for 1x1 convolution, which is actually just a matmul
    rank = len(weight.get_size())
    for _ in range(rank - 2):
        weight = L[aten.squeeze](weight, dim=-1)
    weight = L[aten.permute](weight, [1, 0])

    x = ir.ExternKernel.require_stride_order(x, channels_last_order(rank))
    x_permute = list(range(rank))
    x_permute.append(x_permute.pop(1))
    x = L[aten.permute](x, x_permute)
    *sizes, in_chan = x.get_size()
    x = L[aten.reshape](x, [sympy_product(sizes), in_chan])
    if bias is None:
        result = L[aten.mm](x, weight)
    else:
        result = L[aten.addmm](bias, x, weight)
    result = L[aten.reshape](result, [*sizes, -1])
    result_permute = list(range(rank))
    result_permute.insert(1, result_permute.pop(-1))
    return L[aten.permute](result, result_permute)


@register_lowering(aten.convolution)
def convolution(
    x: TensorBox,
    weight: TensorBox,
    bias: TensorBox | None,
    stride: Sequence[int],
    padding: Sequence[int],
    dilation: Sequence[int],
    transposed: bool,
    output_padding: Sequence[int],
    groups: int,
):
    """Lower aten.convolution using Inductor convolution kernels or fallbacks."""
    stride = tuple(stride)
    padding = tuple(padding)
    dilation = tuple(dilation)
    output_padding = tuple(output_padding)
    if not isinstance(groups, int):
        groups = V.graph.sizevars.guard_int(groups)
    assert isinstance(groups, int)

    # Need use hint for triton template since the template does not
    # work with a dynamic shape.
    #
    # No need to guard_int for dilation and output_padding
    # since the template is only used when dilation is 1 and output_padding
    # is 0.
    stride = tuple(V.graph.sizevars.guard_int_seq(stride))
    padding = tuple(V.graph.sizevars.guard_int_seq(padding))

    kwargs: ConvLayoutParams = {
        "stride": stride,
        "padding": padding,
        "dilation": dilation,
        "transposed": transposed,
        "output_padding": output_padding,
        "groups": groups,
    }

    device_type = ir.get_device_type(x)

    if len(x.get_size()) == len(weight.get_size()) - 1:
        # add batch dimension to simplify rest of function
        return L[aten.squeeze](
            convolution(L[aten.expand](x, [1, *x.get_size()]), weight, bias, **kwargs),
            dim=0,
        )

    out_chan, in_chan, *kernel_shape = V.graph.sizevars.guard_int_seq(weight.get_size())

    # Always convert conv1D to 2D for Intel GPU.
    # Only conv2D can be converted to channel last layout,
    # which have much better performance.
    if len(x.get_size()) == 3 and len(kernel_shape) == 1 and device_type == "xpu":
        kwargs.update(
            {
                "stride": (1,) + stride,
                "padding": (0,) + padding,
                "dilation": (1,) + dilation,
                "output_padding": (0,) + output_padding,
            }
        )
        # (N, C, L) -> (N, C, 1, L)
        x = L[aten.unsqueeze](x, dim=2)
        weight = L[aten.unsqueeze](weight, dim=2)

        return L[aten.squeeze](
            convolution(x, weight, bias, **kwargs),
            dim=2,
        )

    ndim = len(kernel_shape)
    stride = pad_listlike(stride, ndim)
    padding = pad_listlike(padding, ndim)
    dilation = pad_listlike(dilation, ndim)
    output_padding = pad_listlike(output_padding, ndim)

    def channels_last_conv():
        if V.graph.layout_opt and ndim == 2:
            return True

        layout = conv_layout(x, weight, None, **kwargs)
        # TODO: This does not guard on the stride order decision,
        # shall we use optimization_hint to handle unbacked?
        req_stride_order = ir.get_stride_order(
            V.graph.sizevars.guarding_hints_or_throw(layout.stride)
        )
        return req_stride_order == ir.NHWC_STRIDE_ORDER

    autotuning_gemm = config.max_autotune or config.max_autotune_gemm

    if (
        (config.conv_1x1_as_mm or (autotuning_gemm and channels_last_conv()))
        and is_ones(kernel_shape)
        and is_ones(stride)
        and is_zeros(padding)
        and is_ones(dilation)
        and not transposed
        and is_zeros(output_padding)
        and groups == 1
        and V.graph.sizevars.statically_known_gt(sympy_product(x.get_size()), 0)
    ):
        return convert_1x1_conv_to_mm(x, weight, bias)

    if bias is not None and device_type != "cpu":
        # peel off the bias, cudnn is slower with it
        result = convolution(x, weight, None, **kwargs)
        if V.graph.sizevars.statically_known_equals(result.get_size()[1], 0):
            # we should not add bias when the output channel is 0
            return result
        return L[aten.add](
            result, L[aten.view](bias, [result.get_size()[1]] + ndim * [1])
        )

    x.realize()
    weight.realize()

    # ndim can be 1 for convolution in models such as demucs
    # TODO: check if it's beneficial to convert Conv1d to Conv2d and then
    # apply channels last.
    if V.graph.layout_opt and ndim == 2:
        V.graph.num_channels_last_conv += 1
        x = ir.ExternKernel.require_channels_last(x)  # type: ignore[assignment]
        # TODO maybe we can convert weights to channels last just once before
        # running the model.
        weight = ir.ExternKernel.require_channels_last(weight)  # type: ignore[assignment]
        layout = conv_layout(x, weight, None, **kwargs)
    else:
        layout = conv_layout(x, weight, None, **kwargs)
        # TODO: This does not guard on the stride order decision,
        # shall we use optimization_hint to handle unbacked?
        req_stride_order = ir.get_stride_order(
            V.graph.sizevars.guarding_hints_or_throw(layout.stride)
        )
        x = ir.ExternKernel.require_stride_order(x, req_stride_order)  # type: ignore[assignment]
        weight = ir.ExternKernel.require_stride_order(weight, req_stride_order)  # type: ignore[assignment]

    ordered_kwargs_for_cpp_kernel = [
        "stride",
        "padding",
        "dilation",
        "transposed",
        "output_padding",
        "groups",
    ]
    if bias is None:
        args = [x, weight]
        kwargs["bias"] = None  # type: ignore[typeddict-unknown-key]
        ordered_kwargs_for_cpp_kernel.insert(0, "bias")
    else:
        bias = ir.ExternKernel.realize_input(bias)  # type: ignore[assignment]
        assert bias is not None
        args = [x, weight, bias]
        bias.freeze_layout()
        V.graph.sizevars.guard_int_seq(bias.get_size())

    choices = []
    if torch._inductor.utils._use_conv_autotune_backend("ATEN"):
        choices = [
            aten_convolution.bind(
                args,
                layout,
                ordered_kwargs_for_cpp_kernel,
                **kwargs,
            )
        ]

    if (
        torch._inductor.utils._use_conv_autotune_backend("TRITON")
        and use_triton_template(layout)
        # templates only support these:
        and is_ones(dilation)
        and not transposed
        and is_zeros(output_padding)
        # there are some odd models where this check fails (e.g. shufflenet_v2_x1_0)
        and V.graph.sizevars.statically_known_equals(in_chan * groups, x.get_size()[1])  # type: ignore[arg-type]
    ):
        if (
            is_ones(kernel_shape)
            and is_ones(stride)
            and is_zeros(padding)
            and groups == 1
        ):
            choices.append(aten_conv1x1_via_mm.bind(args, layout))

        is_depthwise = groups > 1 and in_chan == 1 and out_chan == groups
        if is_depthwise and ndim == 1:
            depthwise_configs = V.choices.get_depthwise_conv_configs(device_type)
            for cfg in depthwise_configs:
                depthwise_conv1d_template.maybe_append_choice(
                    choices,
                    input_nodes=(x, weight),
                    layout=layout,
                    KERNEL_SIZE=kernel_shape[0],
                    CONV_STRIDE=stride[0],
                    PADDING=padding[0],
                    num_stages=cfg.num_stages,
                    num_warps=cfg.num_warps,
                    **cfg.kwargs,
                )

        conv_configs = V.choices.get_conv_configs(device_type)

        dtype_size = x.get_dtype().itemsize
        for cfg in conv_configs(
            sympy_product([x.get_size()[0], *x.get_size()[2:]]),
            out_chan,
            in_chan,
            dtype_size=dtype_size,
        ):
            unroll = is_ones(kernel_shape)
            # The non-unrolled loop in these templates triggers triton#1254
            # with 8 warps, producing incorrect results for non-1x1 kernels.
            num_warps = cfg.num_warps if unroll else min(cfg.num_warps, 4)

            if ndim == 2:
                conv2d_template.maybe_append_choice(
                    choices,
                    input_nodes=(x, weight),
                    layout=layout,
                    KERNEL_H=kernel_shape[0],
                    KERNEL_W=kernel_shape[1],
                    STRIDE_H=stride[0],
                    STRIDE_W=stride[1],
                    PADDING_H=padding[0],
                    PADDING_W=padding[1],
                    GROUPS=groups,
                    # TODO(jansel): try unroll for bigger kernels once fixed:
                    #               https://github.com/triton-lang/triton/issues/1254
                    UNROLL=unroll,
                    ALLOW_TF32=torch.backends.cudnn.fp32_precision == "tf32",
                    num_stages=cfg.num_stages,
                    num_warps=num_warps,
                    **cfg.kwargs,
                )
            elif ndim == 3:
                conv3d_template.maybe_append_choice(
                    choices,
                    input_nodes=(x, weight),
                    layout=layout,
                    KERNEL_D=kernel_shape[0],
                    KERNEL_H=kernel_shape[1],
                    KERNEL_W=kernel_shape[2],
                    STRIDE_D=stride[0],
                    STRIDE_H=stride[1],
                    STRIDE_W=stride[2],
                    PADDING_D=padding[0],
                    PADDING_H=padding[1],
                    PADDING_W=padding[2],
                    GROUPS=groups,
                    # TODO(jansel): try unroll for bigger kernels once fixed:
                    #               https://github.com/triton-lang/triton/issues/1254
                    UNROLL=unroll,
                    ALLOW_TF32=torch.backends.cudnn.fp32_precision == "tf32",
                    num_stages=cfg.num_stages,
                    num_warps=num_warps,
                    **cfg.kwargs,
                )
    if use_ck_conv_template(layout):
        CKGroupedConvFwdTemplate.add_ck_conv_choices(
            choices,
            layout,
            input_nodes=(x, weight) + ((bias,) if bias is not None else tuple()),
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            n_spatial_dimensions=ndim,
        )
    node, _ = autotune_select_algorithm("convolution", choices, args, layout)
    return node


@register_lowering(aten._convolution)
def _convolution(
    x,
    weight,
    bias,
    stride,
    padding,
    dilation,
    transposed,
    output_padding,
    groups,
    benchmark,
    deterministic,
    cudnn_enabled,
    allow_tf32,
):
    return convolution(
        x, weight, bias, stride, padding, dilation, transposed, output_padding, groups
    )


def constrain_conv_to_fx_strides(fx_node, *args, **kwargs):
    assert fx_node.target is torch.ops.aten.convolution.default
    if V.graph.layout_opt:
        return args, kwargs
    else:
        return constrain_to_fx_strides(fx_node, *args, **kwargs)


add_layout_constraint(aten.convolution, constrain_conv_to_fx_strides)


def conv_bwd_input_layout(
    grad_out: TensorBox,
    input: TensorBox,
    weight: TensorBox,
    stride: Sequence[int],
    padding: tuple[int, ...],
    dilation: tuple[int, ...],
    transposed: bool,
    output_padding: tuple[int, ...],
    groups: int,
) -> ir.Layout:
    guard = V.graph.sizevars.guard_int_seq
    with V.graph.fake_mode:
        go = ir.ir_node_to_tensor(grad_out)
        x = ir.ir_node_to_tensor(input)
        w = ir.ir_node_to_tensor(weight)

        dx, _, _ = torch.ops.aten.convolution_backward(
            go,
            x,
            w,
            None,  # bias_sizes
            guard(stride),
            guard(padding),
            guard(dilation),
            transposed,
            guard(output_padding),
            groups,
            (True, False, False),
        )
        sizes = ir.convert_shape_to_inductor(dx.size())
        stride_ = ir.convert_shape_to_inductor(dx.stride())

    return ir.FixedLayout(
        input.get_device_or_error(),
        input.get_dtype(),
        sizes,
        stride_,
    )


def conv_bwd_weight_layout(
    grad_out: TensorBox,
    input: TensorBox,
    weight: TensorBox,
    stride: Sequence[int],
    padding: tuple[int, ...],
    dilation: tuple[int, ...],
    transposed: bool,
    output_padding: tuple[int, ...],
    groups: int,
) -> ir.Layout:
    guard = V.graph.sizevars.guard_int_seq
    with V.graph.fake_mode:
        go = ir.ir_node_to_tensor(grad_out)
        x = ir.ir_node_to_tensor(input)
        w = ir.ir_node_to_tensor(weight)

        _, dw, _ = torch.ops.aten.convolution_backward(
            go,
            x,
            w,
            None,  # bias_sizes
            guard(stride),
            guard(padding),
            guard(dilation),
            transposed,
            guard(output_padding),
            groups,
            (False, True, False),
        )
        sizes = ir.convert_shape_to_inductor(dw.size())
        stride_ = ir.convert_shape_to_inductor(dw.stride())

    return ir.FixedLayout(
        weight.get_device_or_error(),
        weight.get_dtype(),
        sizes,
        stride_,
    )


def call_aten_dw(
    x_t,
    go_t,
    *,
    w_shape,
    stride,
    padding,
    dilation,
    transposed,
    output_padding,
    groups,
    out,
):
    if x_t.is_contiguous(memory_format=torch.channels_last):
        memory_fmt = torch.channels_last
    else:
        memory_fmt = torch.contiguous_format

    dummy_weight = torch.empty(
        w_shape, dtype=out.dtype, device=x_t.device, memory_format=memory_fmt
    )

    torch.ops.aten.convolution_backward.out(
        out1=None,
        out2=out,
        out3=None,
        grad_output=go_t,
        input=x_t,
        weight=dummy_weight,
        bias_sizes=None,
        stride=stride,
        padding=padding,
        dilation=dilation,
        transposed=transposed,
        output_padding=output_padding,
        groups=groups,
        output_mask=(False, True, False),
    )
    return out


ext_kn_aten_dw = ExternKernelChoice(call_aten_dw, None)


def call_aten_dx(
    go_t,
    w_t,
    *,
    x_shape,
    stride,
    padding,
    dilation,
    transposed,
    output_padding,
    groups,
    out,
):
    if go_t.is_contiguous(memory_format=torch.channels_last):
        memory_fmt = torch.channels_last
    else:
        memory_fmt = torch.contiguous_format

    dummy_input = torch.empty(
        x_shape, dtype=out.dtype, device=go_t.device, memory_format=memory_fmt
    )

    torch.ops.aten.convolution_backward.out(
        out1=out,
        out2=None,
        out3=None,
        grad_output=go_t,
        input=dummy_input,
        weight=w_t,
        bias_sizes=None,
        stride=stride,
        padding=padding,
        dilation=dilation,
        transposed=transposed,
        output_padding=output_padding,
        groups=groups,
        output_mask=(True, False, False),
    )
    return out


ext_kn_aten_dx = ExternKernelChoice(call_aten_dx, None)


@SymbolicGridFn
def conv2d_bwd_weight_grid(cout, cin, kh, kw, meta, *, cdiv):
    g = meta["GROUPS"]
    return (
        cdiv(cin * kh * kw, meta["BLOCK_M"]),
        cdiv(cout // g, meta["BLOCK_N"]),
        g,
    )


conv2d_bwd_weight_template = TritonTemplate(
    name="convolution2d_bwd_weight",
    grid=conv2d_bwd_weight_grid,
    source=load_kernel_template("triton_conv2d_bwd_weight"),
)


@SymbolicGridFn
def conv2d_bwd_input_grid(n, cin, h, w, meta, *, cdiv):
    g = meta["GROUPS"]
    return (
        cdiv(n * h * w, meta["BLOCK_M"]),
        cdiv(cin // g, meta["BLOCK_N"]),
        g,
    )


conv2d_bwd_input_template = TritonTemplate(
    name="convolution2d_bwd_input",
    grid=conv2d_bwd_input_grid,
    source=load_kernel_template("triton_conv2d_bwd_input"),
)


aten_convolution_backward_fallback = fallback_handler(aten.convolution_backward.default)


@register_lowering(aten.convolution_backward.default)
def convolution_backward_lowering(
    grad_out: TensorBox,
    input: TensorBox,
    weight: TensorBox,
    bias_sizes: Sequence[int] | None,
    stride: Sequence[int],
    padding: Sequence[int],
    dilation: Sequence[int],
    transposed: bool,
    output_padding: Sequence[int],
    groups: int,
    output_mask: Sequence[bool],
):
    """
    Lowering function for backward convolution operator.

    TRITON kernels are only registered for supported configurations (currently 2D convolutions).
    For unsupported dimensions or configurations, the choices list remains empty,
    triggering an automatic fallback to the ATen reference implementation.
    This ensures correctness for all cases while enabling TRITON optimizations only where implemented.
    """
    stride = tuple(stride)
    padding = tuple(padding)
    dilation = tuple(dilation)
    output_padding = tuple(output_padding)
    if not isinstance(groups, int):
        groups = V.graph.sizevars.guard_int(groups)

    out_chan, in_chan, *kernel_shape = V.graph.sizevars.guard_int_seq(weight.get_size())

    # The Triton bwd templates substitute DILATION_H/W into the generated
    # kernel source, so they must be concrete Python ints. The fwd template
    # gates on is_ones(dilation) and hard-codes dilation=1, so it can skip
    # this guard.
    stride = tuple(V.graph.sizevars.guard_int_seq(stride))
    padding = tuple(V.graph.sizevars.guard_int_seq(padding))
    dilation = tuple(V.graph.sizevars.guard_int_seq(dilation))

    input.realize()
    weight.realize()
    grad_out.realize()

    kwargs: ConvLayoutParams = {
        "stride": stride,
        "padding": padding,
        "dilation": dilation,
        "transposed": transposed,
        "output_padding": output_padding,
        "groups": groups,
    }

    ndim = len(kernel_shape)
    stride = pad_listlike(stride, ndim)
    padding = pad_listlike(padding, ndim)
    dilation = pad_listlike(dilation, ndim)
    output_padding = pad_listlike(output_padding, ndim)

    device_type = ir.get_device_type(input)

    conv_configs = V.choices.get_conv_configs(device_type)
    dtype_size = input.get_dtype().itemsize

    has_triton_dw_choices = False
    dw = None
    choices_dw = []
    args_w = []
    layout_dw = conv_bwd_weight_layout(grad_out, input, weight, **kwargs)
    if output_mask[1]:
        if V.graph.layout_opt and ndim == 2:
            V.graph.num_channels_last_conv += 1
            input = ir.ExternKernel.require_channels_last(input)  # type: ignore[assignment]
            grad_out = ir.ExternKernel.require_channels_last(grad_out)  # type: ignore[assignment]
            layout_dw = conv_bwd_weight_layout(grad_out, input, weight, **kwargs)
        else:
            guard = V.graph.sizevars.guard_int_seq
            stride_order = ir.get_stride_order(guard(layout_dw.stride))
            input = ir.ExternKernel.require_stride_order(input, stride_order)  # type: ignore[assignment]
            grad_out = ir.ExternKernel.require_stride_order(grad_out, stride_order)  # type: ignore[assignment]

        args_w = [input, grad_out]

        if (
            torch._inductor.utils._use_conv_bwd_weight_autotune_backend("TRITON")
            and use_triton_template(layout_dw)
            and not transposed
            and is_zeros(output_padding)
        ):
            for cfg in conv_configs(
                sympy_product([input.get_size()[0], *input.get_size()[2:]]),
                out_chan,
                in_chan,
                dtype_size=dtype_size,
            ):
                if ndim == 2:
                    has_triton_dw_choices = True
                    conv2d_bwd_weight_template.maybe_append_choice(
                        choices_dw,
                        input_nodes=(input, grad_out),
                        layout=layout_dw,
                        KERNEL_H=kernel_shape[0],
                        KERNEL_W=kernel_shape[1],
                        PADDING_H=padding[0],
                        PADDING_W=padding[1],
                        STRIDE_H=stride[0],
                        STRIDE_W=stride[1],
                        DILATION_H=dilation[0],
                        DILATION_W=dilation[1],
                        GROUPS=groups,
                        ALLOW_TF32=torch.backends.cudnn.allow_tf32,
                        num_stages=cfg.num_stages,
                        num_warps=cfg.num_warps,
                        **cfg.kwargs,
                    )

                # TODO: backward weight 3D

    has_triton_dx_choices = False
    dx = None
    choices_dx = []
    args_x = []
    layout_dx = conv_bwd_input_layout(grad_out, input, weight, **kwargs)
    if output_mask[0]:
        if V.graph.layout_opt and ndim == 2:
            V.graph.num_channels_last_conv += 1
            grad_out = ir.ExternKernel.require_channels_last(grad_out)  # type: ignore[assignment]
            weight = ir.ExternKernel.require_channels_last(weight)  # type: ignore[assignment]
            layout_dx = conv_bwd_input_layout(grad_out, input, weight, **kwargs)
        else:
            guard = V.graph.sizevars.guard_int_seq
            stride_order = ir.get_stride_order(guard(layout_dx.stride))
            grad_out = ir.ExternKernel.require_stride_order(grad_out, stride_order)  # type: ignore[assignment]
            weight = ir.ExternKernel.require_stride_order(weight, stride_order)  # type: ignore[assignment]

        args_x = [grad_out, weight]

        if (
            torch._inductor.utils._use_conv_bwd_input_autotune_backend("TRITON")
            and use_triton_template(layout_dx)
            and not transposed
            and is_zeros(output_padding)
        ):
            # TODO: Use the autotune configuration specific to backward convolution.
            for cfg in conv_configs(
                sympy_product([input.get_size()[0], *input.get_size()[2:]]),
                out_chan,
                in_chan,
                dtype_size=dtype_size,
            ):
                if ndim == 2:
                    has_triton_dx_choices = True
                    conv2d_bwd_input_template.maybe_append_choice(
                        choices_dx,
                        input_nodes=(grad_out, weight),
                        layout=layout_dx,
                        KERNEL_H=kernel_shape[0],
                        KERNEL_W=kernel_shape[1],
                        PADDING_H=padding[0],
                        PADDING_W=padding[1],
                        STRIDE_H=stride[0],
                        STRIDE_W=stride[1],
                        DILATION_H=dilation[0],
                        DILATION_W=dilation[1],
                        GROUPS=groups,
                        ALLOW_TF32=torch.backends.cudnn.allow_tf32,
                        num_stages=cfg.num_stages,
                        num_warps=cfg.num_warps,
                        **cfg.kwargs,
                    )

                # TODO: backward input 3D

    # Fallback when no TRITON choices available, i.e., ndim != 2, backend config = ATEN,...
    if not has_triton_dx_choices and not has_triton_dw_choices:
        return aten_convolution_backward_fallback(
            grad_out,
            input,
            weight,
            bias_sizes,
            stride,
            padding,
            dilation,
            transposed,
            output_padding,
            groups,
            output_mask,
        )

    if output_mask[1]:
        if (
            torch._inductor.utils._use_conv_bwd_weight_autotune_backend("ATEN")
            or not has_triton_dw_choices
        ):
            choices_dw.append(
                ext_kn_aten_dw.bind(
                    input_nodes=args_w,
                    layout=layout_dw,
                    ordered_kwargs_for_cpp_kernel=[
                        "w_shape",
                        "stride",
                        "padding",
                        "dilation",
                        "transposed",
                        "output_padding",
                        "groups",
                    ],
                    w_shape=weight.get_size(),
                    stride=stride,
                    padding=padding,
                    dilation=dilation,
                    transposed=transposed,
                    output_padding=output_padding,
                    groups=groups,
                )
            )

        # TODO: use_ck_conv_template for bwd conv

        dw, _ = autotune_select_algorithm(
            "convolution_bwd_weight", choices_dw, args_w, layout_dw
        )

    if output_mask[0]:
        if (
            torch._inductor.utils._use_conv_bwd_input_autotune_backend("ATEN")
            or not has_triton_dx_choices
        ):
            choices_dx.append(
                ext_kn_aten_dx.bind(
                    input_nodes=args_x,
                    layout=layout_dx,
                    ordered_kwargs_for_cpp_kernel=[
                        "x_shape",
                        "stride",
                        "padding",
                        "dilation",
                        "transposed",
                        "output_padding",
                        "groups",
                    ],
                    x_shape=input.get_size(),
                    stride=stride,
                    padding=padding,
                    dilation=dilation,
                    transposed=transposed,
                    output_padding=output_padding,
                    groups=groups,
                )
            )

        # TODO: use_ck_conv_template for bwd conv

        dx, _ = autotune_select_algorithm(
            "convolution_bwd_input", choices_dx, args_x, layout_dx
        )

    db = None
    if output_mask[2] and bias_sizes is not None:
        db = L[aten.sum](grad_out, axis=[0] + list(range(2, ndim + 2)))

    return (dx, dw, db)


def constrain_conv_bwd_to_fx_strides(fx_node, *args, **kwargs):
    assert fx_node.target == torch.ops.aten.convolution_backward.default
    if V.graph.layout_opt:
        return args, kwargs
    else:
        return constrain_to_fx_strides(fx_node, *args, **kwargs)


add_layout_constraint(aten.convolution_backward, constrain_conv_bwd_to_fx_strides)
