submission 757227

# EVOLVE-BLOCK-START

import torch
import triton
import triton.language as tl
from typing import TypeVar, Tuple

input_t = TypeVar("input_t", bound=Tuple[torch.Tensor, torch.Tensor, torch.Tensor])
output_t = TypeVar("output_t", bound=torch.Tensor)


@triton.jit
def conv2d_kernel(
    input_ptr, weight_ptr, output_ptr,
    batch,
    H, W,
    out_H, out_W,
    stride_h, stride_w,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    C: tl.constexpr,
    KH: tl.constexpr,
    KW: tl.constexpr,
):
    """Direct 2‑D convolution (no padding, stride = 1).

    Shapes (contiguous):
        input  : [batch, C, H, W]
        weight : [C, C, KH, KW]   (output channels == input channels == C)
        output : [batch, C, out_H, out_W] where out_H = H‑KH+1, out_W = W‑KW+1
    """
    pid_b = tl.program_id(0)   # batch index
    pid_ocb = tl.program_id(1) # output‑channel block index
    pid_sp = tl.program_id(2)  # flattened spatial block index

    # output‑channel base for this block
    oc_start = pid_ocb * BLOCK_K

    # number of column blocks (computed at runtime)
    num_col_blocks = tl.cdiv(out_W, BLOCK_N)

    # derive row / column block indices from flattened spatial id
    col_block = pid_sp % num_col_blocks
    row_block = pid_sp // num_col_blocks

    # top‑left corner of the tile in output coordinates
    row_start = row_block * BLOCK_M
    col_start = col_block * BLOCK_N

    # per‑thread offsets within the tile
    row_offsets = row_start + tl.arange(0, BLOCK_M)
    col_offsets = col_start + tl.arange(0, BLOCK_N)

    # masks for valid output positions
    mask_row = row_offsets < out_H
    mask_col = col_offsets < out_W
    mask_rc = mask_row[:, None] & mask_col[None, :]   # (BLOCK_M, BLOCK_N)

    # output‑channel offsets handled by this block
    oc_offsets = oc_start + tl.arange(0, BLOCK_K)
    mask_oc = oc_offsets < C

    # accumulator for the tile (BLOCK_K, BLOCK_M, BLOCK_N)
    acc = tl.zeros((BLOCK_K, BLOCK_M, BLOCK_N), dtype=tl.float32)

    # strides for the input tensor (NCHW layout)
    stride_input_batch = C * H * W
    stride_input_c = H * W
    stride_input_h = W
    stride_input_w = 1

    # strides for the weight tensor (OC, IC, KH, KW)
    stride_weight_oc = C * KH * KW
    stride_weight_ic = KH * KW
    stride_weight_kh = KW
    stride_weight_kw = 1

    # reduction over input channel and kernel spatial dimensions
    for ic in range(C):
        for kh in range(KH):
            for kw in range(KW):
                # input coordinates for this kernel element
                in_row = row_offsets + kh
                in_col = col_offsets + kw

                # flat offsets for the input tile
                offset_input = (
                    pid_b * stride_input_batch
                    + ic * stride_input_c
                    + in_row[:, None] * stride_input_h
                    + in_col[None, :] * stride_input_w
                )
                # load input tile (masked)
                inp = tl.load(
                    input_ptr + offset_input,
                    mask=mask_rc,
                    other=0.0
                )  # (BLOCK_M, BLOCK_N)

                # flat offsets for the weight slice (all output channels in this block)
                offset_weight = (
                    oc_offsets * stride_weight_oc
                    + ic * stride_weight_ic
                    + kh * stride_weight_kh
                    + kw * stride_weight_kw
                )
                w = tl.load(
                    weight_ptr + offset_weight,
                    mask=mask_oc,
                    other=0.0
                )   # (BLOCK_K,)

                # accumulate
                acc += w[:, None, None] * inp[None, :, :]

    # write the result to the output tensor
    stride_out_batch = C * out_H * out_W
    stride_out_oc = out_H * out_W
    stride_out_h = out_W
    stride_out_w = 1

    offset_out = (
        pid_b * stride_out_batch
        + oc_offsets[:, None, None] * stride_out_oc
        + row_offsets[None, :, None] * stride_out_h
        + col_offsets[None, None, :] * stride_out_w
    )
    mask_store = mask_oc[:, None, None] & mask_row[None, :, None] & mask_col[None, None, :]
    tl.store(output_ptr + offset_out, acc, mask=mask_store)


def custom_kernel(data: input_t) -> output_t:
    """Triton implementation of `torch.nn.functional.conv2d` (stride = 1, no padding).

    Args:
        data: tuple of (input_tensor, kernel, output) where
            * input_tensor shape = [batch, C, H, W]
            * kernel shape      = [C, C, kH, kW] (output channels == input channels)
            * output shape      = [batch, C, H‑kH+1, W‑kW+1]

    Returns:
        The `output` tensor filled with the convolution result.
    """
    input_tensor, kernel, output = data

    # Ensure contiguous layout (required for pointer arithmetic)
    input_tensor = input_tensor.contiguous()
    kernel = kernel.contiguous()
    output = output.contiguous()

    # Extract shapes
    batch, C, H, W = input_tensor.shape
    _, _, kH, kW = kernel.shape
    out_H = H - kH + 1
    out_W = W - kW + 1

    # Tiling configuration (tuned for Hopper SMs)
    BLOCK_M = 32   # rows of the output tile
    BLOCK_N = 32   # columns of the output tile
    BLOCK_K = 16   # number of output channels processed per block

    # Grid dimensions:
    #   dim0 -> batch
    #   dim1 -> output‑channel blocks
    #   dim2 -> spatial blocks (row blocks * col blocks)
    num_oc_blocks = (C + BLOCK_K - 1) // BLOCK_K
    num_row_blocks = (out_H + BLOCK_M - 1) // BLOCK_M
    num_col_blocks = (out_W + BLOCK_N - 1) // BLOCK_N
    num_spatial_blocks = num_row_blocks * num_col_blocks

    grid = (batch, num_oc_blocks, num_spatial_blocks)

    # Launch the Triton kernel
    conv2d_kernel[grid](
        input_tensor,
        kernel,
        output,
        batch,
        H,
        W,
        out_H,
        out_W,
        1,  # stride_h (unused, kept for compatibility)
        1,  # stride_w (unused)
        BLOCK_M=BLOCK_M,
        BLOCK_N=BLOCK_N,
        BLOCK_K=BLOCK_K,
        C=C,
        KH=kH,
        KW=kW,
        num_warps=4,
        num_stages=3,
    )
    return output
# EVOLVE-BLOCK-END