submission 511426

- import torch
  import triton
  import triton.language as tl
+ import torch
   
   
- @triton.autotune(
-     configs=[
-         triton.Config({"BLOCK_W": 32}, num_warps=4, num_stages=2),
-         triton.Config({"BLOCK_W": 64}, num_warps=4, num_stages=3),
-         triton.Config({"BLOCK_W": 128}, num_warps=8, num_stages=3),
-     ],
-     key=["OW", "K"],
- )
  @triton.jit
- def conv2d_nopad_kernel(
+ def _conv2d_nchw_kernel(
      x_ptr, w_ptr, y_ptr,
-     B, IC, OC, H, W, OH, OW,
+     B, C, H, W, K, OH, OW,
      stride_xb, stride_xc, stride_xh, stride_xw,
-     stride_wo, stride_wc, stride_wkh, stride_wkw,
+     stride_woc, stride_wic, stride_wkh, stride_wkw,
      stride_yb, stride_yc, stride_yh, stride_yw,
-     BLOCK_W: tl.constexpr,
-     K: tl.constexpr,
+     BLOCK_M: tl.constexpr,  # tile size along output height
+     BLOCK_N: tl.constexpr,  # tile size along output width
  ):
      """
-     Direct 2D convolution without padding and stride=1.
+     Direct NCHW 2D convolution kernel (stride=1, no padding).
+     Computes one output channel (oc) and one batch index (b) per program block,
+     over a BLOCK_M x BLOCK_N tile of the output spatial domain.
   
-     Each program instance computes a vector tile of OW of size BLOCK_W for a single (b, oc, oh).
-     Accumulation is done in fp32 for numerical stability and then cast to the output dtype.
-  
-     Parameters:
-         x_ptr: [B, IC, H, W]
-         w_ptr: [OC, IC, K, K]
-         y_ptr: [B, OC, OH, OW]
-         K: kernel size (square, tl.constexpr)
-         BLOCK_W: width tile size (tl.constexpr)
+     Shapes:
+       x: [B, C, H, W]
+       w: [C, C, K, K]  (out_channels == in_channels == C)
+       y: [B, C, OH, OW], where OH = H - K + 1, OW = W - K + 1
      """
-     pid_m = tl.program_id(axis=0)  # flatten over (B, OC, OH)
-     pid_n = tl.program_id(axis=1)  # tile along OW
+     # Program ids for tiling
+     pid_m = tl.program_id(axis=0)  # output height tiles
+     pid_n = tl.program_id(axis=1)  # output width tiles
+     pid_z = tl.program_id(axis=2)  # packed batch/output-channel index
   
-     # Decompose pid_m into (b, oc, oh)
-     tmp = pid_m
-     oh = tmp % OH
-     tmp = tmp // OH
-     oc = tmp % OC
-     b = tmp // OC
+     # Map pid_z into (b, oc)
+     oc = pid_z % C
+     b = pid_z // C
   
-     # Compute tile offsets along width
-     col_start = pid_n * BLOCK_W
-     offs_ow = col_start + tl.arange(0, BLOCK_W)
-     mask = offs_ow < OW
+     # Compute output tile offsets
+     start_oh = pid_m * BLOCK_M
+     start_ow = pid_n * BLOCK_N
+     offs_oh = start_oh + tl.arange(0, BLOCK_M)
+     offs_ow = start_ow + tl.arange(0, BLOCK_N)
   
-     # Initialize accumulator (vector across BLOCK_W)
-     acc = tl.zeros((BLOCK_W,), dtype=tl.float32)
+     # Masks for out-of-bounds in output tile
+     mask_oh = offs_oh < OH
+     mask_ow = offs_ow < OW
+     store_mask = (mask_oh[:, None]) & (mask_ow[None, :])
   
-     # Loop over input channels and kernel spatial dims
-     # IC can be dynamic; kernel dims K are constexpr for better unrolling.
-     for ic in tl.range(0, IC):
-         # Base pointers that don't depend on kx/ky or ow
-         x_base = x_ptr + b * stride_xb + ic * stride_xc
-         w_ic_base = w_ptr + oc * stride_wo + ic * stride_wc
-         for ky in range(0, K):
-             x_row_base = x_base + (oh + ky) * stride_xh
-             w_row_base = w_ic_base + ky * stride_wkh
-             for kx in range(0, K):
-                 # Load scalar weight w[oc, ic, ky, kx]
-                 w_val = tl.load(w_row_base + kx * stride_wkw)
-                 w_val_f32 = w_val.to(tl.float32)
+     # Improve contiguity/coalescing
+     offs_oh = tl.max_contiguous(tl.multiple_of(tl.where(mask_oh, offs_oh, 0), BLOCK_M), BLOCK_M)
+     offs_ow = tl.max_contiguous(tl.multiple_of(tl.where(mask_ow, offs_ow, 0), BLOCK_N), BLOCK_N)
   
-                 # Load input vector x[b, ic, oh+ky, offs_ow + kx]
-                 x_ptrs = x_row_base + (offs_ow + kx) * stride_xw
-                 x_vec = tl.load(x_ptrs, mask=mask, other=0.0)
-                 x_vec_f32 = x_vec.to(tl.float32)
+     # Output pointers
+     y_ptrs = y_ptr + b * stride_yb + oc * stride_yc + \
+              (offs_oh[:, None] * stride_yh) + (offs_ow[None, :] * stride_yw)
   
-                 # FMA accumulate
-                 acc += x_vec_f32 * w_val_f32
+     # Accumulator in fp32 for precision
+     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     # Store results
-     y_ptrs = y_ptr + b * stride_yb + oc * stride_yc + oh * stride_yh + offs_ow * stride_yw
-     tl.store(y_ptrs, acc.to(y_ptr.dtype.element_ty), mask=mask)
+     # Convolution: reduction over input channels and kernel spatial dims
+     # For off-tile elements (masked), loads are also masked to avoid OOB reads.
+     for ic in range(0, C):
+         w_ic_base = w_ptr + oc * stride_woc + ic * stride_wic
+         # Iterate over kernel elements
+         for kh in range(0, K):
+             for kw in range(0, K):
+                 # Load weight scalar
+                 w_val = tl.load(w_ic_base + kh * stride_wkh + kw * stride_wkw)
+                 w_f32 = w_val.to(tl.float32)
   
+                 # Input window coordinates
+                 ih = offs_oh[:, None] + kh
+                 iw = offs_ow[None, :] + kw
   
+                 # Input pointers for the entire tile patch
+                 x_ptrs = x_ptr + b * stride_xb + ic * stride_xc + \
+                          (ih * stride_xh) + (iw * stride_xw)
+  
+                 # Load and mask invalid output positions to avoid OOB input
+                 x_vals = tl.load(x_ptrs, mask=store_mask, other=0.0)
+                 x_f32 = x_vals.to(tl.float32)
+  
+                 # Accumulate
+                 acc += x_f32 * w_f32
+  
+     # Cast and store
+     if y_ptr.dtype.element_ty == tl.float32:
+         out_vals = acc
+     elif y_ptr.dtype.element_ty == tl.bfloat16:
+         out_vals = acc.to(tl.bfloat16)
+     elif y_ptr.dtype.element_ty == tl.float16:
+         out_vals = acc.to(tl.float16)
+     else:
+         # Fallback to float32 if unexpected dtype
+         out_vals = acc
+  
+     tl.store(y_ptrs, out_vals, mask=store_mask)
+  
+  
  def kernel_function(x: torch.Tensor, w: torch.Tensor, out: torch.Tensor = None):
      """
-     Triton-backed 2D convolution without padding and stride=1.
+     Triton wrapper that performs a stride-1, no-padding, NCHW 2D convolution:
+       y[b, oc, oh, ow] = sum_{ic, kh, kw} x[b, ic, oh+kh, ow+kw] * w[oc, ic, kh, kw]
+     where out_channels == in_channels == C.
   
-     This wrapper:
-     - Validates shapes/dtypes/devices.
-     - Allocates the output if not provided.
-     - Configures the launch grid and meta-parameters.
-     - Launches a single fused Triton kernel that performs the full convolution.
-       No intermediate PyTorch compute ops are used; all math happens in the Triton kernel.
-  
      Fusion note:
-     - There are no extra operations (bias, activation) in the test pipeline.
-       As such, the implementation consists of a single pass conv2d kernel.
-       If bias/activation existed, they would be natural to fuse in the epilogue.
+       The operation performed is a plain convolution. There is no bias, activation,
+       or postprocessing pipeline in the test requirements, so there is nothing
+       meaningful to fuse. We implement the complete convolution in a single Triton
+       kernel pass. If future tests add bias/activation, those can be fused into the
+       epilogue (e.g., BN or ReLU) to reduce memory traffic.
   
+     Runtime behavior:
+       - Validates shapes, dtypes, and device.
+       - Allocates output if not provided.
+       - Launches a single Triton kernel over tiled output height/width, batching and
+         output channels packed along the third grid axis.
+       - All math is performed inside the Triton kernel; no PyTorch compute ops are
+         invoked in the execution path.
+  
      Args:
-         x: Input tensor of shape [B, C_in, H, W], contiguous.
-         w: Weight tensor of shape [C_out, C_in, K, K], contiguous.
-         out: Optional preallocated output tensor [B, C_out, H-K+1, W-K+1].
+       x: Input tensor with shape [B, C, H, W], dtype float32 or bfloat16, device CUDA.
+       w: Kernel tensor with shape [C, C, K, K], same dtype/device as x.
+       out: Optional preallocated output buffer with shape [B, C, H-K+1, W-K+1], same dtype/device.
   
      Returns:
-         Output tensor y of shape [B, C_out, H-K+1, W-K+1], dtype = x.dtype, device = x.device.
+       Output tensor y with shape [B, C, H-K+1, W-K+1].
      """
-     # Basic checks
-     assert isinstance(x, torch.Tensor) and isinstance(w, torch.Tensor)
-     assert x.is_cuda and w.is_cuda, "Input and weights must be on CUDA"
-     assert x.is_contiguous() and w.is_contiguous(), "Input and weights must be contiguous"
+     # Basic validations (wrapper is allowed to perform checks and allocation)
+     assert isinstance(x, torch.Tensor) and isinstance(w, torch.Tensor), "Inputs must be torch tensors"
+     assert x.device.type == "cuda" and w.device.type == "cuda", "CUDA device required"
+     assert x.is_contiguous() and w.is_contiguous(), "Inputs must be contiguous"
+     assert x.dtype == w.dtype, "x and w must have the same dtype"
+     assert x.dtype in (torch.float32, torch.bfloat16), "Supported dtypes: float32, bfloat16"
   
-     B, IC, H, W = x.shape
-     OC, IC_w, KH, KW = w.shape
-     assert IC_w == IC, "Weight in_channels must match input channels"
-     assert KH == KW, "Only square kernels are supported"
-     K = KH
+     B, Cx, H, W = x.shape
+     Co, Ci, Kh, Kw = w.shape
+     assert Cx == Ci == Co, "This kernel requires out_channels == in_channels == C"
+     assert Kh == Kw, "Kernel must be square (KxK)"
+     K = Kh
+     C = Cx
+  
      OH = H - K + 1
      OW = W - K + 1
-     assert OH > 0 and OW > 0, "Kernel larger than input (no padding) produces non-positive output spatial size"
+     assert OH > 0 and OW > 0, "Kernel larger than input; no padding supported"
   
-     # Prepare output
+     # Allocate output if not provided
      if out is None:
-         out = torch.empty((B, OC, OH, OW), device=x.device, dtype=x.dtype)
+         out = torch.empty((B, C, OH, OW), device=x.device, dtype=x.dtype)
      else:
-         assert out.is_cuda, "Output must be on CUDA"
-         assert out.dtype == x.dtype, "Output dtype must match input dtype"
-         assert out.is_contiguous(), "Output must be contiguous"
-         assert tuple(out.shape) == (B, OC, OH, OW), "Output shape mismatch"
+         assert isinstance(out, torch.Tensor), "out must be a tensor if provided"
+         assert out.device == x.device, "out must be on the same device as x"
+         assert out.dtype == x.dtype, "out must have the same dtype as x"
+         assert out.shape == (B, C, OH, OW), f"out must have shape {(B, C, OH, OW)}"
   
-     # Extract strides
+     # Strides
      stride_xb, stride_xc, stride_xh, stride_xw = x.stride()
-     stride_wo, stride_wc, stride_wkh, stride_wkw = w.stride()
+     stride_woc, stride_wic, stride_wkh, stride_wkw = w.stride()
      stride_yb, stride_yc, stride_yh, stride_yw = out.stride()
   
-     # Launch configuration: 2D grid over (B * OC * OH) x ceil(OW / BLOCK_W)
-     def grid(META):
-         return (B * OC * OH, triton.cdiv(OW, META["BLOCK_W"]))
+     # Kernel launch configuration: tile output height/width; pack batch*channels along axis-2
+     # Choose power-of-two tile sizes for better performance; autotuning could be added if needed.
+     # Keep tiles moderate to balance occupancy and register pressure.
+     BLOCK_M = 32
+     BLOCK_N = 32
   
-     # Launch the kernel
-     conv2d_nopad_kernel[grid](
+     grid = (triton.cdiv(OH, BLOCK_M), triton.cdiv(OW, BLOCK_N), B * C)
+  
+     _conv2d_nchw_kernel[grid](
          x, w, out,
-         B, IC, OC, H, W, OH, OW,
+         B, C, H, W, K, OH, OW,
          stride_xb, stride_xc, stride_xh, stride_xw,
-         stride_wo, stride_wc, stride_wkh, stride_wkw,
+         stride_woc, stride_wic, stride_wkh, stride_wkw,
          stride_yb, stride_yc, stride_yh, stride_yw,
-         K=K,
+         BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
+         num_warps=4,  # conservative default; can be tuned
+         num_stages=3,  # pipeline stages
      )
   
      return out