submission 597516

  ⋯ 3 unchanged lines
   
   
  @triton.jit
- def _conv2d_nchw_str1_nopad_kernel(
-     x_ptr,  # *const T_x
-     w_ptr,  # *const T_w
-     y_ptr,  # *T_y
-     # sizes
-     N, C_in, H, W,
-     C_out, K, H_out, W_out,
-     # strides for input x (NCHW)
-     stride_xn, stride_xc, stride_xh, stride_xw,
-     # strides for weight w (O, I, KH, KW)
-     stride_wo, stride_wi, stride_wkh, stride_wkw,
-     # strides for output y (N, O, H_out, W_out)
-     stride_yn, stride_yc, stride_yh, stride_yw,
-     # compile-time tile/block params
-     BLOCK_H: tl.constexpr,
-     BLOCK_W: tl.constexpr,
-     BLOCK_OC: tl.constexpr,
-     M: tl.constexpr,  # = BLOCK_H * BLOCK_W (compile-time)
+ def _conv2d_kernel(
+     input_ptr, weight_ptr, output_ptr,
+     in_channels, out_channels,
+     in_h, in_w, out_h, out_w,
+     kernel_h, kernel_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,
+     K,  # in_channels * kernel_h * kernel_w
+     M,  # out_h * out_w
+     BLOCK_M: tl.constexpr,
+     BLOCK_N: tl.constexpr,
+     BLOCK_K: tl.constexpr,
  ):
-     # Program IDs
-     pid_n = tl.program_id(0)
-     pid_sp = tl.program_id(1)
-     pid_ocg = tl.program_id(2)
+     pid_m = tl.program_id(0)
+     pid_n = tl.program_id(1)
+     b = tl.program_id(2)
   
-     # Tile decomposition for spatial output grid
-     num_tiles_w = tl.cdiv(W_out, BLOCK_W)
-     tile_h_idx = pid_sp // num_tiles_w
-     tile_w_idx = pid_sp % num_tiles_w
-     oh0 = tile_h_idx * BLOCK_H
-     ow0 = tile_w_idx * BLOCK_W
+     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
   
-     # Output channel group origin
-     oc0 = pid_ocg * BLOCK_OC
+     mask_m = offs_m < M
+     mask_n = offs_n < out_channels
   
-     # Output channel range and mask
-     oc_range = oc0 + tl.arange(0, BLOCK_OC)
-     oc_mask = oc_range < C_out
+     oh = offs_m // out_w
+     ow = offs_m % out_w
   
-     # Flatten spatial tile indexing: indices 0..M-1 map to (mh, mw)
-     m_idx = tl.arange(0, M)
-     mh = m_idx // BLOCK_W  # [0..BLOCK_H-1]
-     mw = m_idx % BLOCK_W   # [0..BLOCK_W-1]
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     # Per-lane output coords and mask for in-bounds (avoid indexing tensors with tensors)
-     oh_m = oh0 + mh
-     ow_m = ow0 + mw
-     mask_m = (oh_m < H_out) & (ow_m < W_out)
+     kh_kw = kernel_h * kernel_w
   
-     # Base pointers for this program instance
-     x_n_base = x_ptr + pid_n * stride_xn
-     y_n_base = y_ptr + pid_n * stride_yn
+     for k_start in range(0, K, BLOCK_K):
+         offs_k = k_start + tl.arange(0, BLOCK_K)
+         mask_k = offs_k < K
   
-     # Accumulator (M x BLOCK_OC) in fp32 for numerical stability
-     acc = tl.zeros((M, BLOCK_OC), dtype=tl.float32)
+         ic = offs_k // kh_kw
+         rem = offs_k % kh_kw
+         kh = rem // kernel_w
+         kw = rem % kernel_w
   
-     # Convolution:
-     # y[n, oc, oh, ow] = sum_{ic, kh, kw} x[n, ic, oh+kh, ow+kw] * w[oc, ic, kh, kw]
-     for ic in tl.range(0, C_in):
-         x_ic_base = x_n_base + ic * stride_xc
-         w_ic_base = w_ptr + ic * stride_wi
-         for kh in tl.range(0, K):
-             ih_m = oh_m + kh  # (M,)
-             x_ptrs_row = x_ic_base + ih_m * stride_xh  # (M,)
-             for kw in tl.range(0, K):
-                 iw_m = ow_m + kw  # (M,)
-                 x_ptrs_m = x_ptrs_row + iw_m * stride_xw  # (M,)
-                 x_vec = tl.load(x_ptrs_m, mask=mask_m, other=0.0)
+         ih = oh[:, None] + kh[None, :]
+         iw = ow[:, None] + kw[None, :]
   
-                 # Weight vector for this (ic, kh, kw) and output channel group
-                 w_ptrs_vec = (
-                     w_ic_base
-                     + oc0 * stride_wo
-                     + kh * stride_wkh
-                     + kw * stride_wkw
-                     + tl.arange(0, BLOCK_OC) * stride_wo
-                 )
-                 w_vec = tl.load(w_ptrs_vec, mask=oc_mask, other=0.0)
+         in_offset = (b * stride_in_b +
+                      ic[None, :] * stride_in_c +
+                      ih * stride_in_h +
+                      iw * stride_in_w)
+         a_tile = tl.load(input_ptr + in_offset,
+                          mask=mask_m[:, None] & mask_k[None, :], other=0.0)
   
-                 # Outer product accumulation into (M, BLOCK_OC)
-                 acc += x_vec.to(tl.float32)[:, None] * w_vec.to(tl.float32)[None, :]
+         w_offset = (offs_n[:, None] * stride_w_oc +
+                     ic[None, :] * stride_w_ic +
+                     kh[None, :] * stride_w_kh +
+                     kw[None, :] * stride_w_kw)
+         w_tile = tl.load(weight_ptr + w_offset,
+                          mask=mask_n[:, None] & mask_k[None, :], other=0.0)
   
-     # Store results
-     y_ptrs = (
-         y_n_base
-         + oc_range[None, :] * stride_yc
-         + oh_m[:, None] * stride_yh
-         + ow_m[:, None] * stride_yw
-     )
-     out_dtype = y_ptr.dtype.element_ty
-     y_vals = acc.to(out_dtype)
-     y_mask = mask_m[:, None] & oc_mask[None, :]
-     tl.store(y_ptrs, y_vals, mask=y_mask)
+         acc += tl.dot(a_tile, tl.trans(w_tile), allow_tf32=False)
   
+     out_offset = (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(output_ptr + out_offset, acc, mask=out_mask)
   
- def kernel_function(input_tensor: torch.Tensor, kernel: torch.Tensor, out: torch.Tensor = None):
-     """
-     Triton 2D convolution (NCHW) with stride=1 and padding=0.
-     - Single fused kernel: performs full accumulation over input channels and kernel spatial dims.
-     - No math or reductions in the wrapper: only validation, allocation, and launch.
   
-     Args:
-         input_tensor: (N, C_in, H, W), CUDA tensor (float32 / float16 / bfloat16)
-         kernel:       (C_out, C_in, K, K), contiguous CUDA tensor, same dtype/device as input
-         out:          Optional preallocated tensor (N, C_out, H-K+1, W-K+1)
+ 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
   
-     Returns:
-         (N, C_out, H-K+1, W-K+1) tensor on CUDA with same dtype as input.
-     """
-     # Validation (no compute)
-     assert isinstance(input_tensor, torch.Tensor) and isinstance(kernel, torch.Tensor)
-     assert input_tensor.is_cuda and kernel.is_cuda
-     assert input_tensor.dim() == 4 and kernel.dim() == 4
-     assert input_tensor.device == kernel.device
-     assert input_tensor.dtype == kernel.dtype
-     assert input_tensor.is_contiguous() and kernel.is_contiguous()
+     if output_tensor is None:
+         output_tensor = torch.empty(
+             (batch, out_channels, out_h, out_w),
+             device=input_tensor.device, dtype=input_tensor.dtype)
   
-     N, C_in, H, W = input_tensor.shape
-     C_out, C_in_w, K, K_w = kernel.shape
-     assert C_in == C_in_w and K == K_w, "Incompatible kernel shape"
-     assert H >= K and W >= K, "Kernel larger than input"
-     H_out = H - K + 1
-     W_out = W - K + 1
+     M = out_h * out_w
+     K = in_channels * kernel_h * kernel_w
   
-     # Supported dtypes
-     if input_tensor.dtype not in (torch.float32, torch.float16, torch.bfloat16):
-         raise TypeError(f"Unsupported dtype: {input_tensor.dtype}")
-     out_dtype = input_tensor.dtype
+     BLOCK_M = 64
+     BLOCK_N = 64
+     BLOCK_K = 32
   
-     # Allocate output if needed
-     if out is None:
-         out = torch.empty((N, C_out, H_out, W_out), device=input_tensor.device, dtype=out_dtype)
-     else:
-         assert isinstance(out, torch.Tensor) and out.is_cuda
-         assert out.device == input_tensor.device
-         assert out.dtype == out_dtype
-         assert out.shape == (N, C_out, H_out, W_out)
-         assert out.is_contiguous()
+     grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(out_channels, BLOCK_N), batch)
   
-     # Element-wise strides
-     stride_xn, stride_xc, stride_xh, stride_xw = input_tensor.stride()
-     stride_wo, stride_wi, stride_wkh, stride_wkw = kernel.stride()
-     stride_yn, stride_yc, stride_yh, stride_yw = out.stride()
-  
-     # Tile sizes (powers of two, masks handle tails)
-     BLOCK_H = 16
-     BLOCK_W = 16
-     BLOCK_OC = 32
-     M = BLOCK_H * BLOCK_W  # constexpr
-  
-     def grid(meta):
-         num_tiles_h = triton.cdiv(H_out, meta["BLOCK_H"])
-         num_tiles_w = triton.cdiv(W_out, meta["BLOCK_W"])
-         return (
-             N,                                       # batch
-             num_tiles_h * num_tiles_w,               # spatial tiles
-             triton.cdiv(C_out, meta["BLOCK_OC"]),    # output channel tiles
-         )
-  
-     # Launch kernel
-     _conv2d_nchw_str1_nopad_kernel[grid](
-         input_tensor, kernel, out,
-         N, C_in, H, W,
-         C_out, K, H_out, W_out,
-         stride_xn, stride_xc, stride_xh, stride_xw,
-         stride_wo, stride_wi, stride_wkh, stride_wkw,
-         stride_yn, stride_yc, stride_yh, stride_yw,
-         BLOCK_H=BLOCK_H, BLOCK_W=BLOCK_W, BLOCK_OC=BLOCK_OC, M=M,
-         num_warps=4, num_stages=2,
+     _conv2d_kernel[grid](
+         input_tensor, kernel_weights, output_tensor,
+         in_channels, out_channels,
+         in_h, in_w, out_h, out_w,
+         kernel_h, kernel_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),
+         K, M,
+         BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
      )
-     return out
+     return output_tensor
   
  import inspect
  def custom_kernel(input):