submission 99213

import torch
import triton
import triton.language as tl

@triton.autotune(
    configs=[
        # === H100 (Hopper) Aggressive Configs ===
        # H100 любит большие тайлы и высокий prefetch для скрытия HBM3 латентности
        
        # 1. Massive Width Vectorization: Отлично для memory bandwidth
        triton.Config({'BLOCK_H': 8, 'BLOCK_W': 256}, num_warps=8, num_stages=6),
        
        # 2. Large Vertical Tiling: Максимальный weight reuse (веса в L1/Regs)
        triton.Config({'BLOCK_H': 16, 'BLOCK_W': 128}, num_warps=8, num_stages=5),
        
        # 3. Balanced High-Throughput: Золотая середина для Hopper
        triton.Config({'BLOCK_H': 16, 'BLOCK_W': 64},  num_warps=8, num_stages=6),
        
        # 4. Extreme Prefetching: Если кернел ограничен latency
        triton.Config({'BLOCK_H': 8, 'BLOCK_W': 128}, num_warps=8, num_stages=7),
        
        # === Fallback / Safe Configs ===
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 128}, num_warps=4, num_stages=4),
    ],
    key=['W_OUT', 'H_OUT', 'C_IN', 'K'],
)
@triton.jit
def conv2d_kernel_h100(
    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
):
    """
    Hopper-Optimized Direct Convolution Kernel.
    Использует Pointer Induction и Aggressive Tiling.
    """
    
    # --- 1. Grid & Indices ---
    pid_w = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_z = tl.program_id(2)

    # Разложение Z координаты
    batch_idx = pid_z // C_OUT
    out_ch = pid_z % C_OUT

    # --- 2. Coordinate Generation (Broadcasting Setup) ---
    # Генерируем векторы координат для тайла
    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
    offs_w = pid_w * BLOCK_W + tl.arange(0, BLOCK_W)

    # --- 3. Mask Hoisting ---
    # Вычисляем маски ОДИН раз.
    # Валидность output координат гарантирует валидность input координат
    # при stride=1, padding=0 и корректных размерах тензоров.
    mask_h = offs_h < H_OUT
    mask_w = offs_w < W_OUT
    # Комбинированная маска [BLOCK_H, BLOCK_W]
    mask_block = mask_h[:, None] & mask_w[None, :]

    # --- 4. Pointer Setup (Base Calculation) ---
    
    # Output Pointer: [BLOCK_H, BLOCK_W]
    # Используем broadcasting для создания 2D сетки указателей назначения
    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]
    # Начальная позиция окна свертки.
    # ptr_in[row, col] соответствует input[row, col]
    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: Scalar (будет обновляться в цикле)
    ptr_wei_base = weight_ptr + out_ch * stride_w_out

    # --- 5. Accumulation Loop ---
    acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)

    # Текущие указатели каналов (Pointer Chasing State)
    curr_in_ch = ptr_in_base
    curr_wei_ch = ptr_wei_base

    # Основной цикл по входным каналам (Reduction Dimension)
    for cin in range(C_IN):
        
        # Локальные копии указателей для Spatial Loop
        # Мы не модифицируем curr_in_ch внутри внутреннего цикла, чтобы сохранить базу канала
        curr_in_row = curr_in_ch
        curr_wei_row = curr_wei_ch
        
        # Spatial Loop: Kernel Height
        for kh in range(K):
            
            # Spatial Loop: Kernel Width
            # Этот цикл обычно полностью разворачивается (unrolled) компилятором для малых K (1,3,5,7)
            for kw in range(K):
                # A. Load Weight (Scalar Broadcast)
                # Загружаем 1 float, рассылаем на весь Grid [BLOCK_H, BLOCK_W]
                # Смещение веса: kw * stride_w_w
                w_val = tl.load(curr_wei_row + kw * stride_w_w)
                
                # B. Load Input (Vectorized 2D Block)
                # Загружаем блок данных.
                # Смещение входа: kw * stride_in_w
                # Contiguous load по оси W критичен для H100 HBM
                in_val = tl.load(curr_in_row + kw * stride_in_w, mask=mask_block, other=0.0)
                
                # C. FMA
                acc = acc + in_val * w_val
            
            # Pointer Update: Move down vertically
            # Сдвиг указателей на следующую строку
            curr_in_row += stride_in_h
            curr_wei_row += stride_w_h
            
        # Pointer Update: Move to next channel
        # Замена тяжелого умножения (cin * stride) на быстрое сложение
        curr_in_ch += stride_in_c
        curr_wei_ch += stride_w_in

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


def custom_kernel(data):
    """
    H100-Ready Wrapper.
    """
    input_tensor, kernel, output_tensor = data
    
    # Memory Coalescing is non-negotiable on H100
    if not input_tensor.is_contiguous():
        input_tensor = input_tensor.contiguous()
    if not kernel.is_contiguous():
        kernel = kernel.contiguous()
    
    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
    
    # Grid: Tiles (W, H), Batch*OutCh
    grid = lambda META: (
        triton.cdiv(w_out, META['BLOCK_W']),
        triton.cdiv(h_out, META['BLOCK_H']),
        batch * c_out
    )
    
    conv2d_kernel_h100[grid](
        input_tensor, kernel, output_tensor,
        *input_tensor.stride(),
        *kernel.stride(),
        *output_tensor.stride(),
        h_in, w_in, h_out, w_out,
        c_in, c_out, k_h,
    )
    
    return output_tensor