submission 664306

  ⋯ 3 unchanged lines
   
   
  @triton.jit
- def _conv2d_kernel(
+ def _conv2d_implicit_gemm(
      input_ptr, weight_ptr, output_ptr,
-     in_channels, in_h, in_w,
-     out_channels,
+     batch, in_channels, out_channels,
+     in_h, in_w, out_h, out_w,
      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_k,
-     total_m,
+     K,  # in_channels * kernel_h * kernel_w
+     M,  # out_h * out_w
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
  ):
-     # 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_b = tl.program_id(2)
      pid_m = tl.program_id(0)
      pid_n = tl.program_id(1)
-     b = tl.program_id(2)
   
-     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
+     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
   
-     mask_m = offs_m < total_m
-     mask_n = offs_n < out_channels
+     m_mask = offs_m < M
+     n_mask = 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)
   
-     kh_kw = kernel_h * kernel_w
+     num_k_tiles = tl.cdiv(K, BLOCK_K)
+     for k_tile in range(num_k_tiles):
+         offs_k = k_tile * BLOCK_K + tl.arange(0, BLOCK_K)
+         k_mask = offs_k < K
   
-     for k_start in range(0, total_k, BLOCK_K):
-         offs_k = k_start + tl.arange(0, BLOCK_K)
-         mask_k = offs_k < total_k
-  
-         ic = offs_k // kh_kw
-         rem = offs_k % kh_kw
+         ic = offs_k // (kernel_h * kernel_w)
+         rem = offs_k % (kernel_h * kernel_w)
          kh = rem // kernel_w
          kw = rem % kernel_w
   
-         # Load input tile [BLOCK_M, BLOCK_K]
-         # input[b, ic, oh + kh, ow + kw]
          ih = oh[:, None] + kh[None, :]
          iw = ow[:, None] + kw[None, :]
-         ic_bc = ic[None, :]
   
-         in_ptrs = (input_ptr
-                    + b * stride_in_b
-                    + ic_bc * stride_in_c
-                    + ih * stride_in_h
-                    + iw * stride_in_w)
+         inp_idx = pid_b * (in_channels * in_h * in_w) + ic[None, :] * (in_h * in_w) + ih * in_w + iw
+         inp_mask = m_mask[:, None] & k_mask[None, :]
+         a = tl.load(input_ptr + inp_idx, mask=inp_mask, other=0.0)
   
-         in_mask = mask_m[:, None] & mask_k[None, :]
-         a = tl.load(in_ptrs, mask=in_mask, other=0.0)
+         w_idx = offs_n[None, :] * K + offs_k[:, None]
+         w_mask = n_mask[None, :] & k_mask[:, None]
+         b = tl.load(weight_ptr + w_idx, mask=w_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)
+         acc += tl.dot(a, b, allow_tf32=False)
   
-         w_mask = mask_k[:, None] & mask_n[None, :]
-         b_tile = tl.load(w_ptrs, mask=w_mask, other=0.0)
+     out_idx = pid_b * (out_channels * M) + offs_n[None, :] * M + offs_m[:, None]
+     out_mask = m_mask[:, None] & n_mask[None, :]
+     tl.store(output_ptr + out_idx, acc, mask=out_mask)
   
-         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 = mask_m[:, None] & mask_n[None, :]
-     tl.store(out_ptrs, acc, mask=out_mask)
-  
-  
  def kernel_function(input_tensor, kernel_weights, output_tensor=None):
      batch, in_channels, in_h, in_w = input_tensor.shape
      out_channels, _, kernel_h, kernel_w = kernel_weights.shape
      out_h = in_h - kernel_h + 1
      out_w = in_w - kernel_w + 1
+     M = out_h * out_w
+     K = in_channels * kernel_h * kernel_w
   
      if output_tensor is None:
-         output_tensor = torch.empty(
-             (batch, out_channels, out_h, out_w),
-             device=input_tensor.device, dtype=input_tensor.dtype)
+         output_tensor = torch.empty(batch, out_channels, out_h, out_w,
+                                     device=input_tensor.device, dtype=input_tensor.dtype)
   
-     total_k = in_channels * kernel_h * kernel_w
-     total_m = out_h * out_w
+     inp = input_tensor.contiguous()
+     wgt = kernel_weights.contiguous()
   
      BLOCK_M = 64
      BLOCK_N = 64
      BLOCK_K = 32
   
-     grid = (triton.cdiv(total_m, BLOCK_M), triton.cdiv(out_channels, BLOCK_N), batch)
+     grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(out_channels, BLOCK_N), batch)
   
-     _conv2d_kernel[grid](
-         input_tensor, kernel_weights, output_tensor,
-         in_channels, in_h, in_w,
-         out_channels,
+     _conv2d_implicit_gemm[grid](
+         inp, wgt, output_tensor,
+         batch, in_channels, out_channels,
+         in_h, in_w, out_h, out_w,
          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_k,
-         total_m,
-         BLOCK_M=BLOCK_M,
-         BLOCK_N=BLOCK_N,
-         BLOCK_K=BLOCK_K,
+         K, M,
+         BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
      )
      return output_tensor