submission 99081

import torch
import triton
import triton.language as tl

# Добавляем BLOCK_H в конфигурации автотюнинга
@triton.autotune(
    configs=[
        # Конфигурации для больших тензоров (A100/H100)
        triton.Config({'BLOCK_H': 4, 'BLOCK_W': 128}, num_warps=8, num_stages=4),
        triton.Config({'BLOCK_H': 8, 'BLOCK_W': 64},  num_warps=8, num_stages=4),
        triton.Config({'BLOCK_H': 2, 'BLOCK_W': 256}, num_warps=8, num_stages=3),
        
        # Конфигурации для средних размеров / L4
        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),
        
        # Fallback для маленьких размеров
        triton.Config({'BLOCK_H': 1, 'BLOCK_W': 32},  num_warps=2, num_stages=2),
    ],
    key=['W_OUT', 'H_OUT', 'C_IN', 'K'],
)
@triton.jit
def conv2d_kernel_tiled(
    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
):
    """
    Супер-оптимизированное ядро с 2D-блокировкой (H x W).
    Обрабатывает прямоугольный тайл выхода за один раз.
    """
    
    # Grid IDs
    pid_w = tl.program_id(0)
    pid_h = tl.program_id(1)
    pid_z = tl.program_id(2)
    
    # Декодируем Batch и Output Channel
    batch_idx = pid_z // C_OUT
    out_ch = pid_z % C_OUT
    
    # --- 1. Координаты и маски (2D) ---
    
    # Output Height offsets: [BLOCK_H, 1]
    offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
    mask_h = offs_h < H_OUT
    
    # Output Width offsets: [1, BLOCK_W]
    offs_w = pid_w * BLOCK_W + tl.arange(0, BLOCK_W)
    mask_w = offs_w < W_OUT
    
    # --- 2. Базовые указатели ---
    
    # Указатель на начало Output для текущего Batch/Channel
    # output[b, c, h, w] -> используем 2D broadcasting для записи
    dst_base = output_ptr + batch_idx * stride_out_n + out_ch * stride_out_c
    
    # Указатель на начало Input для текущего Batch
    # Input: [batch, cin, h, w]
    src_base = input_ptr + batch_idx * stride_in_n
    
    # Указатель на начало Weights для текущего Out Channel
    # Weight: [cout, cin, kh, kw]
    wei_base = weight_ptr + out_ch * stride_w_out
    
    # --- 3. Аккумулятор (Register File) ---
    # Теперь это 2D массив [BLOCK_H, BLOCK_W]
    acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
    
    # --- 4. Основной цикл ---
    for cin in range(C_IN):
        # Сдвигаем указатели на текущий входной канал
        src_ch = src_base + cin * stride_in_c
        wei_ch = wei_base + cin * stride_w_in
        
        for kh in range(K):
            # Pre-calc смещения по вертикали для входа
            # Input H index = (offs_h + kh)
            # Broadcasting: (BLOCK_H, 1)
            curr_h_in = offs_h[:, None] + kh
            src_row_ptr = src_ch + curr_h_in * stride_in_h
            
            # Смещение веса по высоте
            wei_row_ptr = wei_ch + kh * stride_w_h
            
            for kw in range(K):
                # --- A. Загрузка Веса (Scalar broadcast) ---
                # Загружаем 1 скаляр, используем для всего блока [BLOCK_H, BLOCK_W]
                wei_val = tl.load(wei_row_ptr + kw * stride_w_w)
                
                # --- B. Загрузка Входа (2D Block Load) ---
                # Input W index = (offs_w + kw)
                # Ptr = src_row_ptr (зависит от H) + (W смещение)
                # Размерность: [BLOCK_H, 1] + [1, BLOCK_W] -> [BLOCK_H, BLOCK_W]
                src_ptrs = src_row_ptr + (offs_w[None, :] + kw) * stride_in_w
                
                # Маска: проверяем валидность H (выхода) и W (выхода)
                # Примечание: padding=0, stride=1 гарантируют, что если output внутри границ,
                # то input (out + k) тоже внутри границ (при валидных размерах тензоров).
                # Используем комбинированную маску mask_h & mask_w
                val_in = tl.load(src_ptrs, mask=(mask_h[:, None] & mask_w[None, :]), other=0.0)
                
                # --- C. FMA ---
                acc = acc + val_in * wei_val

    # --- 5. Запись результата ---
    # Output Ptr: dst_base + h * stride_h + w * stride_w
    offs_out = dst_base + offs_h[:, None] * stride_out_h + offs_w[None, :] * stride_out_w
    tl.store(offs_out, acc, mask=(mask_h[:, None] & mask_w[None, :]))


def custom_kernel(data):
    input_tensor, kernel, output_tensor = data
    
    # Обеспечиваем memory layout
    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 Calculation
    # X: Тайлы по ширине
    # Y: Тайлы по высоте (теперь делим на BLOCK_H!)
    # Z: Batch * OutChannels
    grid = lambda META: (
        triton.cdiv(w_out, META['BLOCK_W']),
        triton.cdiv(h_out, META['BLOCK_H']),
        batch * c_out
    )
    
    conv2d_kernel_tiled[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