submission 759702

  ⋯ 3 unchanged lines
   
   
  @triton.jit
- def _conv2d_implicit_gemm(
+ def _conv2d_kernel(
      input_ptr, weight_ptr, output_ptr,
-     batch, in_channels, out_channels,
-     in_h, in_w, out_h, out_w,
-     kernel_h, kernel_w,
-     K,  # in_channels * kernel_h * kernel_w
-     M,  # 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_oc, stride_out_h, stride_out_w,
+     total_spatial,  # batch * out_h * out_w
+     K_total,        # in_channels * kernel_h * kernel_w
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      BLOCK_K: tl.constexpr,
  ):
-     pid_b = tl.program_id(2)
      pid_m = tl.program_id(0)
      pid_n = tl.program_id(1)
   
      offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
   
-     m_mask = offs_m < M
+     m_mask = offs_m < total_spatial
      n_mask = offs_n < out_channels
   
-     oh = offs_m // out_w
-     ow = offs_m % out_w
+     # Decompose spatial index: offs_m -> (b, oh, ow)
+     b_idx = offs_m // (out_h * out_w)
+     rem = offs_m % (out_h * out_w)
+     oh_idx = rem // out_w
+     ow_idx = rem % out_w
   
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     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, K_total, BLOCK_K):
+         offs_k = k_start + tl.arange(0, BLOCK_K)
+         k_mask = offs_k < K_total
   
+         # Decompose k -> (ic, kh, kw)
          ic = offs_k // (kernel_h * kernel_w)
-         rem = offs_k % (kernel_h * kernel_w)
-         kh = rem // kernel_w
-         kw = rem % kernel_w
+         rem_k = offs_k % (kernel_h * kernel_w)
+         kh = rem_k // kernel_w
+         kw = rem_k % kernel_w
   
-         ih = oh[:, None] + kh[None, :]
-         iw = ow[:, None] + kw[None, :]
+         # Load input patch: [BLOCK_M, BLOCK_K]
+         ih = oh_idx[:, None] + kh[None, :]
+         iw = ow_idx[:, None] + kw[None, :]
+         in_ptrs = (input_ptr
+                    + b_idx[:, None] * stride_in_b
+                    + ic[None, :] * stride_in_c
+                    + ih * stride_in_h
+                    + iw * stride_in_w)
+         a = tl.load(in_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0)
   
-         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)
+         # Load weight: [BLOCK_K, BLOCK_N]
+         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_val = tl.load(w_ptrs, mask=k_mask[:, None] & n_mask[None, :], 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)
+         acc += tl.dot(a, b_val, allow_tf32=False)
   
-         acc += tl.dot(a, b, allow_tf32=False)
+     out_ptrs = (output_ptr
+                 + b_idx[:, None] * stride_out_b
+                 + offs_n[None, :] * stride_out_oc
+                 + oh_idx[:, None] * stride_out_h
+                 + ow_idx[:, None] * stride_out_w)
+     tl.store(out_ptrs, acc, mask=m_mask[:, None] & n_mask[None, :])
   
-     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)
   
-  
  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
+         )
   
-     inp = input_tensor.contiguous()
-     wgt = kernel_weights.contiguous()
+     total_spatial = batch * out_h * out_w
+     K_total = in_channels * kernel_h * kernel_w
   
      BLOCK_M = 64
      BLOCK_N = 64
      BLOCK_K = 32
   
-     grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(out_channels, BLOCK_N), batch)
+     grid = (triton.cdiv(total_spatial, BLOCK_M), triton.cdiv(out_channels, BLOCK_N))
   
-     _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,
-         K, M,
+     _conv2d_kernel[grid](
+         input_tensor, kernel_weights, output_tensor,
+         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_spatial, K_total,
          BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
      )
+  
      return output_tensor
   
  import inspect