submission 639581

  ⋯ 3 unchanged lines
   
   
  @triton.jit
- def _conv2d_implicit_gemm_kernel(
+ def _conv2d_kernel(
      input_ptr, weight_ptr, output_ptr,
-     in_channels, out_channels,
-     in_h, in_w, out_h, out_w,
+     in_channels, in_h, in_w,
+     out_channels,
      kernel_h, kernel_w,
+     out_h, out_w,
      stride_in_b, stride_in_c, stride_in_h, stride_in_w,
      stride_w_oc, stride_w_ic, stride_w_kh, stride_w_kw,
      stride_out_b, stride_out_c, stride_out_h, stride_out_w,
-     total_out_pixels,
-     reduction_size,
+     total_k,
+     total_m,
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
  ):
-     """
-     Fused 2D convolution via implicit GEMM.
-     M = out_h*out_w (output spatial), N = out_channels, K = in_channels*kernel_h*kernel_w
-     A[M,K] = im2col(input), B[K,N] = weight reshaped, C[M,N] = output
-     """
+     # Grid: (num_m_tiles, num_n_tiles, batch)
+     # M = out_h * out_w (spatial output positions)
+     # N = out_channels
+     # K = in_channels * kernel_h * kernel_w (reduction)
      pid_m = tl.program_id(0)
      pid_n = tl.program_id(1)
-     pid_b = tl.program_id(2)
+     b = tl.program_id(2)
   
-     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
-     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)  # spatial positions
+     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)  # output channels
   
-     m_mask = offs_m < total_out_pixels
-     n_mask = offs_n < out_channels
+     mask_m = offs_m < total_m
+     mask_n = offs_n < out_channels
   
+     # Precompute oh, ow for each spatial position
      oh = offs_m // out_w
      ow = offs_m % out_w
   
-     acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
      kh_kw = kernel_h * kernel_w
-     base_in = input_ptr + pid_b * stride_in_b
   
-     for k_start in range(0, reduction_size, BLOCK_K):
+     for k_start in range(0, total_k, BLOCK_K):
          offs_k = k_start + tl.arange(0, BLOCK_K)
-         k_mask = offs_k < reduction_size
+         mask_k = offs_k < total_k
   
          ic = offs_k // kh_kw
          rem = offs_k % kh_kw
          kh = rem // kernel_w
          kw = rem % kernel_w
   
-         # A tile [BLOCK_M, BLOCK_K]: implicit im2col
+         # Load input tile [BLOCK_M, BLOCK_K]
+         # input[b, ic, oh + kh, ow + kw]
          ih = oh[:, None] + kh[None, :]
          iw = ow[:, None] + kw[None, :]
-         in_ptrs = base_in + ic[None, :] * stride_in_c + ih * stride_in_h + iw * stride_in_w
-         a_mask = m_mask[:, None] & k_mask[None, :]
-         a = tl.load(in_ptrs, mask=a_mask, other=0.0)
+         ic_bc = ic[None, :]
   
-         # B tile [BLOCK_K, BLOCK_N]: weight[oc, ic, kh, kw]
+         in_ptrs = (input_ptr
+                    + b * stride_in_b
+                    + ic_bc * stride_in_c
+                    + ih * stride_in_h
+                    + iw * stride_in_w)
+  
+         in_mask = mask_m[:, None] & mask_k[None, :]
+         a = tl.load(in_ptrs, mask=in_mask, other=0.0)
+  
+         # Load weight tile [BLOCK_K, BLOCK_N]
+         # weight[oc, ic, kh, kw] -> we need [K, N] layout
          w_ptrs = (weight_ptr
                    + offs_n[None, :] * stride_w_oc
                    + ic[:, None] * stride_w_ic
                    + kh[:, None] * stride_w_kh
                    + kw[:, None] * stride_w_kw)
-         b_mask = k_mask[:, None] & n_mask[None, :]
-         b = tl.load(w_ptrs, mask=b_mask, other=0.0)
   
-         acc = tl.dot(a, b, acc, allow_tf32=False)
+         w_mask = mask_k[:, None] & mask_n[None, :]
+         b_tile = tl.load(w_ptrs, mask=w_mask, other=0.0)
   
-     # Store output
-     out_ptrs = (output_ptr + pid_b * stride_out_b
+         acc += tl.dot(a, b_tile, allow_tf32=False)
+  
+     # Store output [BLOCK_M, BLOCK_N]
+     out_ptrs = (output_ptr
+                 + b * stride_out_b
                  + offs_n[None, :] * stride_out_c
                  + oh[:, None] * stride_out_h
                  + ow[:, None] * stride_out_w)
-     out_mask = m_mask[:, None] & n_mask[None, :]
+     out_mask = mask_m[:, None] & mask_n[None, :]
      tl.store(out_ptrs, acc, mask=out_mask)
   
   
  ⋯ 6 unchanged lines
      if output_tensor is None:
          output_tensor = torch.empty(
              (batch, out_channels, out_h, out_w),
-             device=input_tensor.device, dtype=input_tensor.dtype,
-         )
+             device=input_tensor.device, dtype=input_tensor.dtype)
   
-     total_out_pixels = out_h * out_w
-     reduction_size = in_channels * kernel_h * kernel_w
+     total_k = in_channels * kernel_h * kernel_w
+     total_m = out_h * out_w
   
      BLOCK_M = 64
-     BLOCK_N = min(64, triton.next_power_of_2(out_channels))
-     BLOCK_K = min(32, triton.next_power_of_2(reduction_size))
-     if BLOCK_K < 16:
-         BLOCK_K = 16
+     BLOCK_N = 64
+     BLOCK_K = 32
   
-     grid = (triton.cdiv(total_out_pixels, BLOCK_M),
-             triton.cdiv(out_channels, BLOCK_N),
-             batch)
+     grid = (triton.cdiv(total_m, BLOCK_M), triton.cdiv(out_channels, BLOCK_N), batch)
   
-     _conv2d_implicit_gemm_kernel[grid](
+     _conv2d_kernel[grid](
          input_tensor, kernel_weights, output_tensor,
-         in_channels, out_channels,
-         in_h, in_w, out_h, out_w,
+         in_channels, in_h, in_w,
+         out_channels,
          kernel_h, kernel_w,
+         out_h, out_w,
          input_tensor.stride(0), input_tensor.stride(1),
          input_tensor.stride(2), input_tensor.stride(3),
          kernel_weights.stride(0), kernel_weights.stride(1),
          kernel_weights.stride(2), kernel_weights.stride(3),
          output_tensor.stride(0), output_tensor.stride(1),
          output_tensor.stride(2), output_tensor.stride(3),
-         total_out_pixels,
-         reduction_size,
+         total_k,
+         total_m,
          BLOCK_M=BLOCK_M,
          BLOCK_N=BLOCK_N,
          BLOCK_K=BLOCK_K,