submission 513373

  import torch
- import torch.nn.functional as F
+ import triton
+ import triton.language as tl
  from task import input_t, output_t
   
   
- def custom_kernel(data: input_t) -> output_t:
-     """
-     Fast 2D convolution using cuDNN with benchmark mode enabled.
-     Disable TF32 to match reference precision (rtol=1e-3, atol=1e-3).
-     Enable benchmark mode so cuDNN can select the fastest algorithm.
-     """
-     input_tensor, kernel, output = data
-     
-     # Save state
-     old_tf32_cudnn = torch.backends.cudnn.allow_tf32
-     old_tf32_matmul = torch.backends.cuda.matmul.allow_tf32
-     old_benchmark = torch.backends.cudnn.benchmark
-     
-     # Disable TF32 for precision, enable benchmark for speed
-     torch.backends.cudnn.allow_tf32 = False
-     torch.backends.cuda.matmul.allow_tf32 = False
-     torch.backends.cudnn.benchmark = True
-     
-     try:
-         result = F.conv2d(
-             input_tensor,
-             kernel,
-             stride=1,
-             padding=0,
+ @triton.autotune(
+     configs=[
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=4, num_stages=4),
+         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_warps=8, num_stages=3),
+         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_warps=16, num_stages=2),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 16}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 16}, num_warps=16, num_stages=3),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 16}, num_warps=8, num_stages=4),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=16, num_stages=2),
+         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 256, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_warps=16, num_stages=2),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_warps=8, num_stages=2),
+         triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 128}, num_warps=4, num_stages=3),
+         triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 128}, num_warps=8, num_stages=2),
+     ],
+     key=['M', 'N', 'K'],
+ )
+ @triton.jit
+ def conv2d_fused_kernel(
+     input_ptr, weight_ptr, output_ptr,
+     M, N, K,
+     IC, H, W, OH, OW, KH, KW,
+     stride_ib, stride_ic, stride_ih, stride_iw,
+     stride_ob, stride_ooc, stride_ooh, stride_oow,
+     BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+ ):
+     pid = tl.program_id(0)
+     num_pid_m = tl.cdiv(M, BLOCK_M)
+     num_pid_n = tl.cdiv(N, BLOCK_N)
+     GROUP_SIZE: tl.constexpr = 8
+     num_pid_in_group = GROUP_SIZE * num_pid_n
+     group_id = pid // num_pid_in_group
+     first_pid_m = group_id * GROUP_SIZE
+     group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE)
+     pid_m = first_pid_m + (pid % group_size_m)
+     pid_n = (pid % num_pid_in_group) // group_size_m
+  
+     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+  
+     ohow = OH * OW
+     b = offs_m // ohow
+     r = offs_m % ohow
+     oh = r // OW
+     ow = r % OW
+  
+     mm = offs_m < M
+     nm = offs_n < N
+  
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+  
+     inp_b = input_ptr + b * stride_ib + oh * stride_ih + ow * stride_iw
+  
+     offs_k = tl.arange(0, BLOCK_K)
+     kh_kw = KH * KW
+  
+     for k0 in range(0, K, BLOCK_K):
+         k = k0 + offs_k
+         km = k < K
+         ic = k // kh_kw
+         rk = k % kh_kw
+         kh = rk // KW
+         kw = rk % KW
+  
+         a = tl.load(
+             inp_b[:, None] + ic[None, :] * stride_ic + kh[None, :] * stride_ih + kw[None, :] * stride_iw,
+             mask=mm[:, None] & km[None, :], other=0.0
          )
-         output.copy_(result)
-     finally:
-         torch.backends.cudnn.allow_tf32 = old_tf32_cudnn
-         torch.backends.cuda.matmul.allow_tf32 = old_tf32_matmul
-         torch.backends.cudnn.benchmark = old_benchmark
-     
-     return output
   No newline at end of file
+         w = tl.load(
+             weight_ptr + k[:, None] * N + offs_n[None, :],
+             mask=km[:, None] & nm[None, :], other=0.0
+         )
+         acc += tl.dot(a, w, allow_tf32=False)
+  
+     tl.store(
+         output_ptr + b[:, None] * stride_ob + offs_n[None, :] * stride_ooc + oh[:, None] * stride_ooh + ow[:, None] * stride_oow,
+         acc, mask=mm[:, None] & nm[None, :]
+     )
+  
+  
+ def custom_kernel(data: input_t) -> output_t:
+     x, w, out = data
+     B, IC, H, W_ = x.shape
+     OC, _, KH, KW = w.shape
+     OH, OW = H - KH + 1, W_ - KW + 1
+     K = IC * KH * KW
+     M = B * OH * OW
+     N = OC
+  
+     wt_flat = w.reshape(OC, K).t().contiguous()
+  
+     grid = lambda META: (triton.cdiv(M, META['BLOCK_M']) * triton.cdiv(N, META['BLOCK_N']),)
+     conv2d_fused_kernel[grid](
+         x, wt_flat, out,
+         M, N, K,
+         IC, H, W_, OH, OW, KH, KW,
+         x.stride(0), x.stride(1), x.stride(2), x.stride(3),
+         out.stride(0), out.stride(1), out.stride(2), out.stride(3),
+     )
+     return out
   No newline at end of file