submission 512024

+ import torch
  import triton
  import triton.language as tl
- import torch
   
   
+ # Triton kernel: 2D cross-correlation (PyTorch conv2d semantics) with:
+ # - no padding
+ # - stride = 1
+ # - out_channels == in_channels == channels
+ # Accumulation is in fp32 for both fp32 and bf16 inputs.
+ #
+ # Performance fix:
+ # - Parallelize the output height (OH) across the launch grid instead of looping over OH inside
+ #   a single program. The previous version looped over OH in-kernel and timed out on large cases.
+ # - Keep a deterministic sequential accumulation across (ic, kh, kw) to improve numerical
+ #   agreement with PyTorch's conv2d under strict tolerances, while avoiding excessive overhead.
+  
+  
+ @triton.autotune(
+     configs=[
+         triton.Config({'BLOCK_OC': 32, 'BLOCK_OW': 64}, num_stages=3, num_warps=4),
+         triton.Config({'BLOCK_OC': 64, 'BLOCK_OW': 32}, num_stages=3, num_warps=8),
+         triton.Config({'BLOCK_OC': 64, 'BLOCK_OW': 64}, num_stages=4, num_warps=8),
+         triton.Config({'BLOCK_OC': 32, 'BLOCK_OW': 128}, num_stages=3, num_warps=4),
+         triton.Config({'BLOCK_OC': 64, 'BLOCK_OW': 128}, num_stages=4, num_warps=8),
+     ],
+     key=['B', 'C', 'H', 'W', 'KH', 'KW'],
+ )
  @triton.jit
- def _conv2d_nchw_kernel(
+ def _conv2d_nopad_stride1_parallel_oh(
      x_ptr, w_ptr, y_ptr,
-     B, C, H, W, K, OH, OW,
-     stride_xb, stride_xc, stride_xh, stride_xw,
-     stride_woc, stride_wic, stride_wkh, stride_wkw,
-     stride_yb, stride_yc, stride_yh, stride_yw,
-     BLOCK_M: tl.constexpr,  # tile size along output height
-     BLOCK_N: tl.constexpr,  # tile size along output width
+     B, C, H, W, KH, KW, OH, OW,
+     stride_x_b, stride_x_c, stride_x_h, stride_x_w,
+     stride_w_oc, stride_w_ic, stride_w_kh, stride_w_kw,
+     stride_y_b, stride_y_c, stride_y_h, stride_y_w,
+     BLOCK_OC: tl.constexpr,
+     BLOCK_OW: tl.constexpr,
  ):
-     """
-     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.
+     # Grid mapping:
+     # axis 0: combined (batch, oh) tiles
+     # axis 1: tiles over output channels
+     # axis 2: tiles over output width
+     pid_boh = tl.program_id(axis=0)
+     pid_oc = tl.program_id(axis=1)
+     pid_ow = tl.program_id(axis=2)
   
-     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
-     """
-     # 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
+     # Recover (b, oh) from packed pid_boh
+     # We launch axis0 = B * OH, so:
+     b = pid_boh // OH
+     oh = pid_boh % OH
   
-     # Map pid_z into (b, oc)
-     oc = pid_z % C
-     b = pid_z // C
+     # Compute offsets for oc and ow tiles
+     oc_offsets = pid_oc * BLOCK_OC + tl.arange(0, BLOCK_OC)  # [BLOCK_OC]
+     ow_offsets = pid_ow * BLOCK_OW + tl.arange(0, BLOCK_OW)  # [BLOCK_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)
+     # Masks for bounds
+     oc_mask = oc_offsets < C
+     ow_mask = ow_offsets < OW
   
-     # 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, :])
+     # FP32 accumulator
+     acc = tl.zeros((BLOCK_OC, BLOCK_OW), dtype=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)
+     # Reduction across input channels, KH, KW in deterministic order
+     for ic in tl.range(0, C):
+         for kh in tl.range(0, KH):
+             in_h = oh + kh  # valid: no padding, stride=1
+             for kw in tl.range(0, KW):
+                 in_w = ow_offsets + kw  # [BLOCK_OW]
+                 in_w_mask = (ow_mask & (in_w < W))
   
-     # Output pointers
-     y_ptrs = y_ptr + b * stride_yb + oc * stride_yc + \
-              (offs_oh[:, None] * stride_yh) + (offs_ow[None, :] * stride_yw)
+                 # Load weights for this (ic, kh, kw) across BLOCK_OC output channels
+                 w_ptrs = (
+                     w_ptr
+                     + oc_offsets * stride_w_oc
+                     + ic * stride_w_ic
+                     + kh * stride_w_kh
+                     + kw * stride_w_kw
+                 )
+                 w_vals = tl.load(w_ptrs, mask=oc_mask, other=0.0).to(tl.float32)  # [BLOCK_OC]
   
-     # Accumulator in fp32 for precision
-     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+                 # Load input row for this (ic, in_h) across BLOCK_OW positions
+                 x_ptrs = (
+                     x_ptr
+                     + b * stride_x_b
+                     + ic * stride_x_c
+                     + in_h * stride_x_h
+                     + in_w * stride_x_w
+                 )
+                 x_vals = tl.load(x_ptrs, mask=in_w_mask, other=0.0).to(tl.float32)  # [BLOCK_OW]
   
-     # 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)
+                 # Outer product accumulate
+                 acc += w_vals[:, None] * x_vals[None, :]
   
-                 # Input window coordinates
-                 ih = offs_oh[:, None] + kh
-                 iw = offs_ow[None, :] + kw
+     # Store results to output
+     y_ptrs = (
+         y_ptr
+         + b * stride_y_b
+         + oc_offsets[:, None] * stride_y_c
+         + oh * stride_y_h
+         + ow_offsets[None, :] * stride_y_w
+     )
+     y_mask = oc_mask[:, None] & ow_mask[None, :]
+     tl.store(y_ptrs, acc.to(y_ptr.dtype.element_ty), mask=y_mask)
   
-                 # 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):
+ def kernel_function(input_tensor: torch.Tensor, weight: torch.Tensor, output_tensor: torch.Tensor = None):
      """
-     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.
+     Triton-backed 2D convolution (no padding, stride=1), matching PyTorch's F.conv2d semantics
+     for the test settings:
+       - input:  [B, C, H, W]
+       - weight: [C, C, KH, KW] (out_channels == in_channels == C)
+       - output: [B, C, H-KH+1, W-KW+1]
   
-     Fusion note:
-       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.
+     All math runs inside the Triton kernel; the wrapper only validates/allocates/launches.
+     Accumulation is done in fp32 for both fp32 and bf16 inputs.
   
-     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 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.
+         input_tensor: CUDA tensor [B, C, H, W], dtype float32 or bfloat16, contiguous.
+         weight:       CUDA tensor [C, C, KH, KW], same dtype as input, contiguous.
+         output_tensor (optional): Pre-allocated output [B, C, OH, OW], same dtype/device.
   
      Returns:
-       Output tensor y with shape [B, C, H-K+1, W-K+1].
+         Output tensor [B, C, OH, OW], CUDA, same dtype as input.
      """
-     # 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"
+     # Validate inputs
+     assert input_tensor.is_cuda and weight.is_cuda, "Tensors must be CUDA"
+     assert input_tensor.dtype in (torch.float32, torch.bfloat16), "Supported dtypes: float32, bfloat16"
+     assert input_tensor.dtype == weight.dtype, "Input/weight dtypes must match"
+     assert input_tensor.dim() == 4 and weight.dim() == 4, "Expected 4D tensors"
+     assert input_tensor.is_contiguous() and weight.is_contiguous(), "Inputs must be contiguous"
   
-     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
+     B, C_in, H, W = input_tensor.shape
+     OC, IC, KH, KW = weight.shape
+     assert C_in == IC == OC, "Expect out_channels == in_channels == channels"
+     assert KH >= 1 and KW >= 1, "Kernel dims must be positive"
   
-     OH = H - K + 1
-     OW = W - K + 1
-     assert OH > 0 and OW > 0, "Kernel larger than input; no padding supported"
+     OH = H - KH + 1
+     OW = W - KW + 1
+     assert OH > 0 and OW > 0, "Invalid kernel size: no padding, stride=1"
   
-     # Allocate output if not provided
-     if out is None:
-         out = torch.empty((B, C, OH, OW), device=x.device, dtype=x.dtype)
+     if output_tensor is None:
+         output_tensor = torch.empty((B, OC, OH, OW), device=input_tensor.device, dtype=input_tensor.dtype)
      else:
-         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)}"
+         assert output_tensor.is_cuda, "Output tensor must be CUDA"
+         assert output_tensor.dtype == input_tensor.dtype, "Output dtype mismatch"
+         assert output_tensor.is_contiguous(), "Output must be contiguous"
+         assert tuple(output_tensor.shape) == (B, OC, OH, OW), "Output shape mismatch"
   
      # Strides
-     stride_xb, stride_xc, stride_xh, stride_xw = x.stride()
-     stride_woc, stride_wic, stride_wkh, stride_wkw = w.stride()
-     stride_yb, stride_yc, stride_yh, stride_yw = out.stride()
+     stride_x_b, stride_x_c, stride_x_h, stride_x_w = input_tensor.stride()
+     stride_w_oc, stride_w_ic, stride_w_kh, stride_w_kw = weight.stride()
+     stride_y_b, stride_y_c, stride_y_h, stride_y_w = output_tensor.stride()
   
-     # 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
+     # 3D grid over (batch*OH, channel tiles, output-width tiles)
+     def grid(META):
+         return (
+             B * OH,
+             triton.cdiv(C_in, META['BLOCK_OC']),
+             triton.cdiv(OW, META['BLOCK_OW']),
+         )
   
-     grid = (triton.cdiv(OH, BLOCK_M), triton.cdiv(OW, BLOCK_N), B * C)
-  
-     _conv2d_nchw_kernel[grid](
-         x, w, out,
-         B, C, H, W, K, OH, OW,
-         stride_xb, stride_xc, stride_xh, stride_xw,
-         stride_woc, stride_wic, stride_wkh, stride_wkw,
-         stride_yb, stride_yc, stride_yh, stride_yw,
-         BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
-         num_warps=4,  # conservative default; can be tuned
-         num_stages=3,  # pipeline stages
+     # Launch Triton kernel
+     _conv2d_nopad_stride1_parallel_oh[grid](
+         input_tensor, weight, output_tensor,
+         B, C_in, H, W, KH, KW, OH, OW,
+         stride_x_b, stride_x_c, stride_x_h, stride_x_w,
+         stride_w_oc, stride_w_ic, stride_w_kh, stride_w_kw,
+         stride_y_b, stride_y_c, stride_y_h, stride_y_w,
      )
   
-     return out
+     return output_tensor
   
  import inspect