submission 99116

import torch
import triton
import triton.language as tl

@triton.autotune(
    configs=[
        # === High-end (A100, H100, B200) ===
        # Большой тайл по ширине (W) и средний по высоте (H) + высокий prefetch
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 256}, num_warps=8, num_stages=5),
        triton.Config({'BLOCK_H': 8, 'BLOCK_W': 128}, num_warps=8, num_stages=5),
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 128}, num_warps=8, num_stages=6),
        
        # === Balanced (L4, A10) ===
        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 / Latency optimized ===
        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_optimized(
    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 (2D Tiling + Pre-calc Masks + Pointer Arithmetic).
    """
    
    # 1. Grid IDs
    pid_w = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_z = tl.program_id(2)
    
    # 2. Decode Dimensions
    batch_idx = pid_z // C_OUT
    out_ch = pid_z % C_OUT
    
    # 3. Calculate Offsets & Masks (Pre-calculated!)
    # Output Y coords [BLOCK_H]
    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
    mask_h = offs_h < H_OUT
    
    # Output X coords [BLOCK_W]
    offs_w = pid_w * BLOCK_W + tl.arange(0, BLOCK_W)
    mask_w = offs_w < W_OUT
    
    # Combined Mask [BLOCK_H, BLOCK_W]
    # Вычисляем один раз и используем везде.
    # При stride=1 и padding=0 валидность выхода гарантирует валидность входа.
    mask_block = mask_h[:, None] & mask_w[None, :]
    
    # 4. Base Pointers Setup
    # Output Ptr: Base + Batch offset + Channel offset
    dst_ptr_base = output_ptr + batch_idx * stride_out_n + out_ch * stride_out_c
    
    # Input Ptr: Base + Batch offset + (Initial H offset) + (Initial W offset)
    # Входной H начинается там же, где выходной H (offs_h), так как stride=1
    # Входной W начинается там же, где выходной W (offs_w)
    # Мы используем broadcasting для создания 2D сетки указателей
    # Input Ptrs [BLOCK_H, BLOCK_W]
    src_ptr_base = input_ptr + batch_idx * stride_in_n + \
                   (offs_h[:, None] * stride_in_h) + \
                   (offs_w[None, :] * stride_in_w)
                   
    # Weight Ptr Base: Channel Offset
    wei_ptr_base = weight_ptr + out_ch * stride_w_out
    
    # 5. Accumulator
    acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
    
    # 6. Main Loop
    for cin in range(C_IN):
        # Сдвигаем указатели каналов
        src_ch = src_ptr_base + cin * stride_in_c
        wei_ch = wei_ptr_base + cin * stride_w_in
        
        for kh in range(K):
            # Смещение по вертикали ядра
            # Для входа: добавляем stride_in_h * kh
            # Для веса: добавляем stride_w_h * kh
            src_row = src_ch + kh * stride_in_h
            wei_row = wei_ch + kh * stride_w_h
            
            for kw in range(K):
                # --- A. Load Weight (Scalar) ---
                # Загружаем [1] скаляр и "размножаем" его неявно при умножении
                wei_val = tl.load(wei_row + kw * stride_w_w)
                
                # --- B. Load Input (2D Block) ---
                # Указатель уже содержит offs_h и offs_w.
                # Нам нужно только добавить смещение текущего kw
                # src_row [BLOCK_H, BLOCK_W] + scalar offset
                src_ptrs = src_row + kw * stride_in_w
                
                # Используем пре-калькулированную маску!
                val_in = tl.load(src_ptrs, mask=mask_block, other=0.0)
                
                # --- C. FMA ---
                acc = acc + val_in * wei_val

    # 7. Store Result
    # Вычисляем указатели назначения
    dst_ptrs = dst_ptr_base + \
               (offs_h[:, None] * stride_out_h) + \
               (offs_w[None, :] * stride_out_w)
               
    tl.store(dst_ptrs, acc, mask=mask_block)


def custom_kernel(data):
    input_tensor, kernel, output_tensor = data
    
    # Contiguous check - критично для Triton
    if not input_tensor.is_contiguous():
        input_tensor = input_tensor.contiguous()
    if not kernel.is_contiguous():
        kernel = kernel.contiguous()
    
    # Dimensions
    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: (W_tiles, H_tiles, Batch*OutCh)
    grid = lambda META: (
        triton.cdiv(w_out, META['BLOCK_W']),
        triton.cdiv(h_out, META['BLOCK_H']),
        batch * c_out
    )
    
    conv2d_kernel_optimized[grid](
        input_tensor, kernel, output_tensor,
        # Strides
        *input_tensor.stride(),
        *kernel.stride(),
        *output_tensor.stride(),
        # Dimensions
        h_in, w_in, h_out, w_out,
        c_in, c_out, k_h,
    )
    
    return output_tensor