submission 99105

  ⋯ 3 unchanged lines
   
  @triton.autotune(
      configs=[
-         # === A100/H100/B200 - максимум ===
-         triton.Config({'BLOCK_H': 8, 'BLOCK_W': 128}, num_warps=8, num_stages=5),
+         # === 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': 16, 'BLOCK_W': 64}, num_warps=8, num_stages=4),
-         triton.Config({'BLOCK_H': 8, 'BLOCK_W': 256}, num_warps=8, num_stages=4),
+         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),
          
-         # === L4 / средние ===
-         triton.Config({'BLOCK_H': 4, 'BLOCK_W': 128}, num_warps=8, num_stages=4),
-         triton.Config({'BLOCK_H': 8, 'BLOCK_W': 64}, num_warps=4, num_stages=4),
-         triton.Config({'BLOCK_H': 2, 'BLOCK_W': 256}, num_warps=4, num_stages=4),
+         # === 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),
          
-         # === Маленькие ===
-         triton.Config({'BLOCK_H': 2, 'BLOCK_W': 64}, num_warps=4, num_stages=3),
-         triton.Config({'BLOCK_H': 1, 'BLOCK_W': 128}, num_warps=2, num_stages=3),
+         # === 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_size'],
+     key=['W_OUT', 'H_OUT', 'C_IN', 'K'],
  )
  @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-БЛОЧНОЕ ЯДРО CONV2D
-     
-     Ключевые оптимизации:
-     ✅ 2D-блокировка для максимальной L1 cache локальности
-     ✅ Порядок циклов: cin → kh → kw (оптимально для NCHW)
-     ✅ Pre-computed offsets вне горячих циклов
-     ✅ 2D broadcasting для эффективных операций
-     ✅ num_stages=5 для максимального prefetching
-     ✅ Правильные маски для граничных условий
-     ✅ += вместо acc = acc + для компиляторной оптимизации
+     Ultimate Optimized Conv2D Kernel (2D Tiling + Pre-calc Masks + Pointer Arithmetic).
      """
      
-     # === Grid Decoding ===
+     # 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
      
-     # === 2D Output Offsets ===
+     # 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
      
-     # === Base Pointers ===
-     out_base = output_ptr + batch_idx * stride_out_n + out_ch * stride_out_c
-     in_base = input_ptr + batch_idx * stride_in_n
-     wei_base = weight_ptr + out_ch * stride_w_out
+     # Combined Mask [BLOCK_H, BLOCK_W]
+     # Вычисляем один раз и используем везде.
+     # При stride=1 и padding=0 валидность выхода гарантирует валидность входа.
+     mask_block = mask_h[:, None] & mask_w[None, :]
      
-     # === 2D Accumulator ===
-     acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
+     # 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
      
-     # === Combined 2D Mask (early compute) ===
-     mask_2d = (mask_h[:, None] & mask_w[None, :])
+     # 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
      
-     # === Main Convolution Loop ===
+     # 5. Accumulator
+     acc = tl.zeros([BLOCK_H, BLOCK_W], dtype=tl.float32)
+     
+     # 6. Main Loop
      for cin in range(C_IN):
-         in_ch = in_base + cin * stride_in_c
-         wei_ch = wei_base + cin * stride_w_in
+         # Сдвигаем указатели каналов
+         src_ch = src_ptr_base + cin * stride_in_c
+         wei_ch = wei_ptr_base + cin * stride_w_in
          
          for kh in range(K):
-             # Pre-compute height offsets for this kernel position
-             in_h_off = (offs_h[:, None] + kh) * stride_in_h
-             in_row = in_ch + in_h_off
+             # Смещение по вертикали ядра
+             # Для входа: добавляем 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):
-                 # Load weight (scalar, broadcasts to [BLOCK_H, BLOCK_W])
+                 # --- A. Load Weight (Scalar) ---
+                 # Загружаем [1] скаляр и "размножаем" его неявно при умножении
                  wei_val = tl.load(wei_row + kw * stride_w_w)
                  
-                 # Load input block (2D) with proper indexing
-                 in_w_off = (offs_w[None, :] + kw) * stride_in_w
-                 in_ptrs = in_row + in_w_off
-                 in_val = tl.load(in_ptrs, mask=mask_2d, other=0.0)
+                 # --- 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
                  
-                 # FMA with += for better compiler optimization
-                 acc += in_val * wei_val
-     
-     # === Store Result ===
-     out_h_off = offs_h[:, None] * stride_out_h
-     out_w_off = offs_w[None, :] * stride_out_w
-     out_ptrs = out_base + out_h_off + out_w_off
-     
-     tl.store(out_ptrs, acc, mask=mask_2d)
+                 # Используем пре-калькулированную маску!
+                 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):
-     """Optimized wrapper function."""
-     
      input_tensor, kernel, output_tensor = data
      
-     # Ensure contiguous memory layout
-     input_tensor = input_tensor.contiguous()
-     kernel = kernel.contiguous()
+     # Contiguous check - критично для Triton
+     if not input_tensor.is_contiguous():
+         input_tensor = input_tensor.contiguous()
+     if not kernel.is_contiguous():
+         kernel = kernel.contiguous()
      
-     # Extract dimensions
+     # Dimensions
      batch, c_in, h_in, w_in = input_tensor.shape
      c_out, _, k_h, k_w = kernel.shape
      
-     # Calculate output dimensions (stride=1, padding=0)
      h_out = h_in - k_h + 1
      w_out = w_in - k_w + 1
      
-     # Grid configuration: (W_blocks, H_blocks, Batch*Out_Channels)
+     # 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
      )
      
-     # Launch kernel
-     conv2d_kernel_tiled[grid](
+     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
-