submission 732334

  ⋯ 1 unchanged lines
  #!POPCORN gpu H100
   
  import torch
- import torch.nn.functional as F
+ import triton
+ import triton.language as tl
   
  from task import input_t, output_t
   
- # Disable TF32: must match reference cuDNN computation exactly.
  torch.backends.cudnn.allow_tf32 = False
  torch.backends.cuda.matmul.allow_tf32 = False
- # Enable cuDNN benchmarking to select the fastest algorithm.
- torch.backends.cudnn.benchmark = True
   
   
+ def _ceildiv(a: int, b: int) -> int:
+     return -(a // -b)
+  
+  
+ @triton.jit
+ def _conv2d_general(
+     X, W, OUT,
+     BATCH, IN_C, IN_H, IN_W, OUT_C, OUT_H, OUT_W, KERNEL_H, KERNEL_W,
+     sXn, sXc, sXh, sXw,
+     sWci, sWkh, sWkw, sWco,
+     sOn, sOc, sOh, sOw,
+     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+     NUM_STAGES: tl.constexpr,
+ ):
+     """General conv2d with full vector K-decomposition (handles any KW vs BK)."""
+     pid_m = tl.program_id(0)
+     pid_n = tl.program_id(1)
+     nhw = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     ow = nhw % OUT_W
+     nh = nhw // OUT_W
+     oh = nh % OUT_H
+     n = nh // OUT_H
+     co = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+     KHW = KERNEL_H * KERNEL_W
+     K_TOTAL = KHW * IN_C
+     K_BLOCKS = (K_TOTAL + BLOCK_K - 1) // BLOCK_K
+     valid_m = n < BATCH
+     valid_n = co < OUT_C
+  
+     # Precompute per-M base (constant across K loop)
+     x_base = n * sXn + oh * sXh + ow * sXw
+     w_co_base = co * sWco
+  
+     for k_blk in tl.range(0, K_BLOCKS, num_stages=NUM_STAGES):
+         offs_k = k_blk * BLOCK_K + tl.arange(0, BLOCK_K)
+         valid_k = offs_k < K_TOTAL
+         # Full vector decomposition (correct for any BK vs KW)
+         ci = offs_k // KHW
+         tmp = offs_k % KHW
+         ki = tmp // KERNEL_W
+         kj = tmp % KERNEL_W
+  
+         x_ptrs = X + x_base[:, None] + ci[None, :] * sXc + ki[None, :] * sXh + kj[None, :] * sXw
+         x_tile = tl.load(x_ptrs, mask=valid_m[:, None] & valid_k[None, :], other=0.0)
+  
+         w_ptrs = W + ci[:, None] * sWci + ki[:, None] * sWkh + kj[:, None] * sWkw + w_co_base[None, :]
+         w_tile = tl.load(w_ptrs, mask=valid_k[:, None] & valid_n[None, :], other=0.0)
+  
+         acc = tl.dot(x_tile, w_tile, acc, allow_tf32=False)
+  
+     out_ptrs = (
+         OUT
+         + n[:, None] * sOn
+         + co[None, :] * sOc
+         + oh[:, None] * sOh
+         + ow[:, None] * sOw
+     )
+     tl.store(out_ptrs, acc, mask=valid_m[:, None] & valid_n[None, :])
+  
+  
+ @triton.jit
+ def _conv2d_optimized(
+     X, W, OUT,
+     BATCH, IN_C, IN_H, IN_W, OUT_C, OUT_H, OUT_W, KERNEL_H, KERNEL_W,
+     sXn, sXc, sXh, sXw,
+     sWci, sWkh, sWkw, sWco,
+     sOn, sOc, sOh, sOw,
+     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+     NUM_STAGES: tl.constexpr,
+ ):
+     """Optimized for scoring shape: BK<=KW so ki is scalar within a K-block."""
+     pid_m = tl.program_id(0)
+     pid_n = tl.program_id(1)
+     nhw = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     ow = nhw % OUT_W
+     nh = nhw // OUT_W
+     oh = nh % OUT_H
+     n = nh // OUT_H
+     co = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+     KHW = KERNEL_H * KERNEL_W
+     K_TOTAL = KHW * IN_C
+     K_BLOCKS = K_TOTAL // BLOCK_K
+     valid_m = n < BATCH
+  
+     x_base = n * sXn + oh * sXh + ow * sXw
+     w_co_base = co * sWco
+  
+     for k_blk in tl.range(0, K_BLOCKS, num_stages=NUM_STAGES):
+         flat_base = k_blk * BLOCK_K
+         ci_s = flat_base // KHW
+         rem = flat_base - ci_s * KHW
+         ki_s = rem // KERNEL_W
+         kj_start = rem - ki_s * KERNEL_W
+         kj = kj_start + tl.arange(0, BLOCK_K)
+  
+         x_ptrs = X + x_base[:, None] + ci_s * sXc + ki_s * sXh + kj[None, :] * sXw
+         x_tile = tl.load(x_ptrs, mask=valid_m[:, None], other=0.0)
+  
+         w_ptrs = W + ci_s * sWci + ki_s * sWkh + kj[:, None] * sWkw + w_co_base[None, :]
+         w_tile = tl.load(w_ptrs)
+  
+         acc = tl.dot(x_tile, w_tile, acc, allow_tf32=False)
+  
+     tl.store(
+         OUT + n[:, None] * sOn + co[None, :] * sOc + oh[:, None] * sOh + ow[:, None] * sOw,
+         acc, mask=valid_m[:, None],
+     )
+  
+  
  def custom_kernel(data: input_t) -> output_t:
      input_tensor, kernel, output = data
-     output[...] = F.conv2d(input_tensor, kernel, stride=1, padding=0)
+     batch, in_c, in_h, in_w = input_tensor.shape
+     out_c, _, kh, kw = kernel.shape
+     out_h = in_h - kh + 1
+     out_w = in_w - kw + 1
+  
+     wt = kernel.permute(1, 2, 3, 0).contiguous()
+  
+     K_TOTAL = in_c * kh * kw
+     # Use optimized path when BK divides KW and K_TOTAL, and BN divides OUT_C
+     if kw >= 32 and K_TOTAL % 32 == 0 and out_c % 128 == 0:
+         BM, BN, BK, NS, NW = 32, 128, 32, 3, 4
+         kern = _conv2d_optimized
+     else:
+         BM, BN, BK, NS, NW = 64, 128, 16, 2, 8
+         kern = _conv2d_general
+  
+     grid = (
+         _ceildiv(batch * out_h * out_w, BM),
+         _ceildiv(out_c, BN),
+         1,
+     )
+     kern[grid](
+         input_tensor, wt, output,
+         batch, in_c, in_h, in_w, out_c, out_h, out_w, kh, kw,
+         input_tensor.stride(0), input_tensor.stride(1),
+         input_tensor.stride(2), input_tensor.stride(3),
+         wt.stride(0), wt.stride(1), wt.stride(2), wt.stride(3),
+         output.stride(0), output.stride(1), output.stride(2), output.stride(3),
+         BLOCK_M=BM, BLOCK_N=BN, BLOCK_K=BK, NUM_STAGES=NS,
+         num_warps=NW,
+     )
      return output