submission 99156

import torch
import triton
import triton.language as tl

@triton.autotune(
    configs=[
        # === High-End (A100, H100, B200) ===
        # Максимальный prefetch (stages=6) скрывает латентность HBM
        triton.Config({'BLOCK_H': 8, 'BLOCK_W': 128}, num_warps=8, num_stages=6),
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 256}, num_warps=8, num_stages=6),
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 128}, num_warps=8, num_stages=5),
        
        # === Mid-Range / General ===
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 64},  num_warps=4, num_stages=4),
        triton.Config({'BLOCK_H': 2, 'BLOCK_W': 128}, num_warps=4, num_stages=4),
        
        # === Small sizes / Latency ===
        triton.Config({'BLOCK_H': 2, 'BLOCK_W': 64},  num_warps=4, num_stages=3),
    ],
    key=['W_OUT', 'H_OUT', 'C_IN', 'K'],
)
@triton.jit
def conv2d_kernel_ultimate(
    input_ptr, weight_ptr, output_ptr,
    stride_in_n, stride_in_c, stride_in_h, stride_in_w,
    stride_w_out, stride_w_in, stride_w_h, stride_w_w,
    stride_out_n, stride_out_c, stride_out_h, stride_out_w,
    H_IN, W_IN, H_OUT, W_OUT, C_IN, C_OUT, K,
    BLOCK_H: tl.constexpr, BLOCK_W: tl.constexpr
):
    """
    Ultimate Optimized Conv2D Kernel.
    Особенности:
    1. 2D Tiling (H, W) для переиспользования данных в L1 кэше.
    2. Pointer Induction: замена умножения на сложение в циклах.
    3. Pre-calculated Masks: вынос логики масок из горячих циклов.
    """
    
    # --- 1. Setup Grid & Indices ---
    pid_w = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_z = tl.program_id(2)
    
    batch_idx = pid_z // C_OUT
    out_ch = pid_z % C_OUT
    
    # --- 2. Offsets & Masks (Computed ONCE) ---
    
    # Output Y offsets [BLOCK_H]
    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
    # Output X offsets [BLOCK_W]
    offs_w = pid_w * BLOCK_W + tl.arange(0, BLOCK_W)
    
    # Pre-calculate mask [BLOCK_H, BLOCK_W]
    # Примечание: при валидных размерах тензоров и stride=1, 
    # проверка выхода гарантирует валидность входа для padding=0.
    mask_h = offs_h < H_OUT
    mask_w = offs_w < W_OUT
    mask_block = mask_h[:, None] & mask_w[None, :]
    
    # --- 3. Initial Pointers Setup ---
    
    # Output Pointer [BLOCK_H, BLOCK_W] (Broadcasting)
    # Base + Batch + Channel + H_offset + W_offset
    ptr_out = output_ptr + \
              batch_idx * stride_out_n + \
              out_ch * stride_out_c + \
              (offs_h[:, None] * stride_out_h) + \
              (offs_w[None, :] * stride_out_w)

    # Input Pointer Base [BLOCK_H, BLOCK_W]
    # Мы начинаем с позиции, соответствующей верхнему левому углу окна для первого пикселя блока.
    # Так как stride=1, input_h == output_h
    ptr_in_base = input_ptr + \
                  batch_idx * stride_in_n + \
                  (offs_h[:, None] * stride_in_h) + \
                  (offs_w[None, :] * stride_in_w)

    # Weight Pointer Base
    ptr_wei_base = weight_ptr + out_ch * stride_w_out
    
    # Accumulator
    acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
    
    # --- 4. Main Loop (Pointer Chasing) ---
    
    # Текущие указатели для начала канала
    curr_in_ch = ptr_in_base
    curr_wei_ch = ptr_wei_base
    
    for cin in range(C_IN):
        # Временные указатели для Spatial Loop
        curr_in_row = curr_in_ch
        curr_wei_row = curr_wei_ch
        
        for kh in range(K):
            # Смещение внутри строки (kw)
            # Мы не меняем указатель строки, а вычисляем смещения от него,
            # так как K обычно мал, и компилятор хорошо разворачивает это.
            # Однако для строки мы делаем инкремент.
            
            for kw in range(K):
                # Load Weight: Scalar -> Broadcast
                # ptr + kw * stride_w
                wei_val = tl.load(curr_wei_row + kw * stride_w_w)
                
                # Load Input: 2D Block
                # ptr + kw * stride_in (т.к. contiguous по W, это просто смещение на kw)
                # Input Stride W обычно равен 1, но используем переменную для универсальности.
                val_in = tl.load(curr_in_row + kw * stride_in_w, mask=mask_block, other=0.0)
                
                # FMA
                acc = acc + val_in * wei_val
            
            # Инкремент указателей строк (сдвиг вниз по H)
            curr_in_row += stride_in_h
            curr_wei_row += stride_w_h

        # Инкремент указателей каналов (Pointer Induction)
        # Это заменяет умножение `cin * stride` на сложение
        curr_in_ch += stride_in_c
        curr_wei_ch += stride_w_in

    # --- 5. Store Result ---
    tl.store(ptr_out, acc, mask=mask_block)


def custom_kernel(data):
    """
    Optimized Conv2D entry point.
    """
    input_tensor, kernel, output_tensor = data
    
    # 1. Ensure contiguous layout (Critical for vectorization)
    if not input_tensor.is_contiguous():
        input_tensor = input_tensor.contiguous()
    if not kernel.is_contiguous():
        kernel = kernel.contiguous()
    
    # 2. Extract shapes
    batch, c_in, h_in, w_in = input_tensor.shape
    c_out, _, k_h, k_w = kernel.shape
    
    h_out = h_in - k_h + 1
    w_out = w_in - k_w + 1
    
    # 3. Grid definition
    grid = lambda META: (
        triton.cdiv(w_out, META['BLOCK_W']),
        triton.cdiv(h_out, META['BLOCK_H']),
        batch * c_out
    )
    
    # 4. Launch
    conv2d_kernel_ultimate[grid](
        input_tensor, kernel, output_tensor,
        # Strides
        *input_tensor.stride(),
        *kernel.stride(),
        *output_tensor.stride(),
        # Dims
        h_in, w_in, h_out, w_out,
        c_in, c_out, k_h,
    )
    
    return output_tensor