submission 99099

  ⋯ 1 unchanged lines
  import triton
  import triton.language as tl
   
- # Добавляем BLOCK_H в конфигурации автотюнинга
  @triton.autotune(
      configs=[
-         # Конфигурации для больших тензоров (A100/H100)
+         # === A100/H100/B200 (максимальная производительность) ===
+         triton.Config({'BLOCK_H': 8, 'BLOCK_W': 128}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_H': 16, 'BLOCK_W': 64}, num_warps=8, num_stages=5),
+         triton.Config({'BLOCK_H': 4, 'BLOCK_W': 256}, num_warps=8, num_stages=5),
          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),
+         # === L4 / средние размеры ===
+         triton.Config({'BLOCK_H': 4, 'BLOCK_W': 64}, num_warps=8, num_stages=4),
          triton.Config({'BLOCK_H': 2, 'BLOCK_W': 128}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_H': 8, 'BLOCK_W': 32}, num_warps=4, num_stages=4),
          
-         # Fallback для маленьких размеров
-         triton.Config({'BLOCK_H': 1, 'BLOCK_W': 32},  num_warps=2, num_stages=2),
+         # === Маленькие тензоры / fallback ===
+         triton.Config({'BLOCK_H': 1, 'BLOCK_W': 64}, num_warps=2, num_stages=3),
+         triton.Config({'BLOCK_H': 2, 'BLOCK_W': 32}, num_warps=2, num_stages=3),
      ],
-     key=['W_OUT', 'H_OUT', 'C_IN', 'K'],
+     key=['w_out', 'h_out', 'c_in', 'k_size'],
  )
  @triton.jit
- def conv2d_kernel_tiled(
+ 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,
  ⋯ 2 unchanged lines
      BLOCK_H: tl.constexpr, BLOCK_W: tl.constexpr
  ):
      """
-     Супер-оптимизированное ядро с 2D-блокировкой (H x W).
-     Обрабатывает прямоугольный тайл выхода за один раз.
+     ✅ МАКСИМАЛЬНО ОПТИМИЗИРОВАННОЕ ЯДРО CONV2D
+     
+     Оптимизации:
+     1. 2D-блокировка (BLOCK_H × BLOCK_W) для максимальной локальности
+     2. Правильный порядок циклов: cin → kh → kw
+     3. Pre-computed pointers для минимума arithmetic
+     4. 2D broadcasting для эффективных загрузок/записей
+     5. num_stages=5 для максимального prefetching
+     6. Правильные маски для граничных условий
      """
      
-     # Grid IDs
+     # === Grid IDs ===
      pid_w = tl.program_id(0)
      pid_h = tl.program_id(1)
      pid_z = tl.program_id(2)
      
-     # Декодируем Batch и Output Channel
+     # Декодирование Batch и Output Channel
      batch_idx = pid_z // C_OUT
      out_ch = pid_z % C_OUT
      
-     # --- 1. Координаты и маски (2D) ---
-     
-     # Output Height offsets: [BLOCK_H, 1]
+     # === 2D Offsets ===
+     # Height: [BLOCK_H] -> reshape в [BLOCK_H, 1] для broadcasting
      offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
      mask_h = offs_h < H_OUT
      
-     # Output Width offsets: [1, BLOCK_W]
+     # Width: [BLOCK_W] -> reshape в [1, BLOCK_W] для broadcasting
      offs_w = pid_w * BLOCK_W + tl.arange(0, BLOCK_W)
      mask_w = offs_w < W_OUT
      
-     # --- 2. Базовые указатели ---
+     # === Базовые Pointers ===
+     # Output: [batch, out_ch, :, :]
+     out_base = output_ptr + batch_idx * stride_out_n + out_ch * stride_out_c
      
-     # Указатель на начало 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, :, :, :]
+     in_base = input_ptr + batch_idx * stride_in_n
      
-     # Указатель на начало Input для текущего Batch
-     # Input: [batch, cin, h, w]
-     src_base = input_ptr + batch_idx * stride_in_n
-     
-     # Указатель на начало Weights для текущего Out Channel
-     # Weight: [cout, cin, kh, kw]
+     # Weight: [out_ch, :, :, :]
      wei_base = weight_ptr + out_ch * stride_w_out
      
-     # --- 3. Аккумулятор (Register File) ---
-     # Теперь это 2D массив [BLOCK_H, BLOCK_W]
+     # === 2D Accumulator в Register File ===
+     # [BLOCK_H, BLOCK_W] - максимум параллелизма
      acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
      
-     # --- 4. Основной цикл ---
+     # === ОСНОВНОЙ ЦИКЛ СВЕРТКИ ===
+     # Порядок: cin → kh → kw (оптимально для памяти NCHW)
      for cin in range(C_IN):
-         # Сдвигаем указатели на текущий входной канал
-         src_ch = src_base + cin * stride_in_c
+         # Pre-compute channel-shifted pointers
+         in_ch = in_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
+             # Pre-compute Height offset (Broadcasting: [BLOCK_H, 1])
+             # offs_h[:, None] превращает [BLOCK_H] в [BLOCK_H, 1]
+             in_h = (offs_h[:, None] + kh) * stride_in_h
+             in_row = in_ch + in_h
              
-             # Смещение веса по высоте
-             wei_row_ptr = wei_ch + kh * stride_w_h
+             # Pre-compute Weight row offset
+             wei_row = 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)
+                 # === Загрузка Веса (Scalar broadcast) ===
+                 # Один скаляр используется для всего блока [BLOCK_H, BLOCK_W]
+                 wei_val = tl.load(wei_row + 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
+                 # === Загрузка Входа (2D Block) ===
+                 # offs_w[None, :] превращает [BLOCK_W] в [1, BLOCK_W]
+                 # Broadcasting: [BLOCK_H, 1] + [1, BLOCK_W] = [BLOCK_H, BLOCK_W]
+                 in_w = (offs_w[None, :] + kw) * stride_in_w
+                 in_ptrs = in_row + 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)
+                 # Загрузка с правильной маской
+                 in_val = tl.load(
+                     in_ptrs,
+                     mask=(mask_h[:, None] & mask_w[None, :]),
+                     other=0.0
+                 )
                  
-                 # --- C. FMA ---
-                 acc = acc + val_in * wei_val
+                 # === FMA (Fused Multiply-Add) ===
+                 acc += in_val * wei_val
+     
+     # === ЗАПИСЬ РЕЗУЛЬТАТА ===
+     # 2D адреса: base + h_offset * stride_h + w_offset * stride_w
+     out_h = offs_h[:, None] * stride_out_h
+     out_w = offs_w[None, :] * stride_out_w
+     out_ptrs = out_base + out_h + out_w
+     
+     # Запись с маской для граничных условий
+     tl.store(
+         out_ptrs,
+         acc,
+         mask=(mask_h[:, None] & mask_w[None, :])
+     )
   
-     # --- 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):
+     """✅ Оптимальная wrapper для Conv2D."""
+     
      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()
+     # Гарантируем контигуозность памяти
+     input_tensor = input_tensor.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: (W_blocks, H_blocks, Batch*C_out)
      grid = lambda META: (
          triton.cdiv(w_out, META['BLOCK_W']),
          triton.cdiv(h_out, META['BLOCK_H']),
          batch * c_out
      )
      
-     conv2d_kernel_tiled[grid](
+     # Запуск ядра
+     conv2d_kernel_optimized[grid](
          input_tensor, kernel, output_tensor,
          *input_tensor.stride(),
          *kernel.stride(),
  ⋯ 3 unchanged lines
      )
      
      return output_tensor
+