submission 490602

+ 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_kernel(
-     input_ptr,      # Pointer to input tensor [B, C, H, W]
-     kernel_ptr,     # Pointer to kernel tensor [C_out, C_in, KH, KW]
-     output_ptr,     # Pointer to output tensor [B, C_out, OH, OW]
-     # Dimensions
-     batch: tl.constexpr,
-     in_channels: tl.constexpr,
-     out_channels: tl.constexpr,
-     in_height: tl.constexpr,
-     in_width: tl.constexpr,
-     kernel_size: tl.constexpr,
-     out_height: tl.constexpr,
-     out_width: tl.constexpr,
-     # Block sizes
-     BLOCK_OH: tl.constexpr,
-     BLOCK_OW: tl.constexpr,
+ def conv2d_nopad_kernel(
+     x_ptr, w_ptr, y_ptr,
+     B, IC, OC, H, W, OH, OW,
+     stride_xb, stride_xc, stride_xh, stride_xw,
+     stride_wo, stride_wc, stride_wkh, stride_wkw,
+     stride_yb, stride_yc, stride_yh, stride_yw,
+     BLOCK_W: tl.constexpr,
+     K: tl.constexpr,
  ):
      """
-     Fused 2D convolution kernel.
-     Each program computes a block of output pixels for one (batch, out_channel) pair.
-     The convolution sum over in_channels and kernel spatial dimensions is computed inline.
+     Direct 2D convolution without padding and stride=1.
+  
+     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)
      """
-     # Program IDs
-     pid_b = tl.program_id(0)  # batch index
-     pid_oc = tl.program_id(1)  # output channel index
-     pid_spatial = tl.program_id(2)  # spatial block index
-     
-     # Calculate spatial block position
-     num_blocks_ow = tl.cdiv(out_width, BLOCK_OW)
-     pid_oh = pid_spatial // num_blocks_ow
-     pid_ow = pid_spatial % num_blocks_ow
-     
-     # Output pixel offsets within this block
-     offs_oh = pid_oh * BLOCK_OH + tl.arange(0, BLOCK_OH)
-     offs_ow = pid_ow * BLOCK_OW + tl.arange(0, BLOCK_OW)
-     
-     # Masks for valid output positions
-     mask_oh = offs_oh < out_height
-     mask_ow = offs_ow < out_width
-     
-     # Initialize accumulator for this block of output pixels
-     acc = tl.zeros((BLOCK_OH, BLOCK_OW), dtype=tl.float32)
-     
-     # Loop over input channels
-     for ic in range(in_channels):
-         # Loop over kernel height
-         for kh in range(kernel_size):
-             # Loop over kernel width
-             for kw in range(kernel_size):
-                 # Load kernel weight for this (oc, ic, kh, kw)
-                 # Kernel layout: [out_channels, in_channels, kernel_size, kernel_size]
-                 kernel_offset = (pid_oc * in_channels * kernel_size * kernel_size +
-                                 ic * kernel_size * kernel_size +
-                                 kh * kernel_size + kw)
-                 w = tl.load(kernel_ptr + kernel_offset)
-                 
-                 # Input positions: ih = oh + kh, iw = ow + kw
-                 # Input layout: [batch, channels, height, width]
-                 # offs_ih = offs_oh + kh (shape: BLOCK_OH)
-                 # offs_iw = offs_ow + kw (shape: BLOCK_OW)
-                 
-                 # Load input values for this block
-                 # We need to load input[pid_b, ic, offs_oh + kh, offs_ow + kw]
-                 input_base = (pid_b * in_channels * in_height * in_width +
-                              ic * in_height * in_width)
-                 
-                 # Calculate input offsets for each output position
-                 # input_offset[i, j] = input_base + (offs_oh[i] + kh) * in_width + (offs_ow[j] + kw)
-                 offs_ih = offs_oh + kh  # [BLOCK_OH]
-                 offs_iw = offs_ow + kw  # [BLOCK_OW]
-                 
-                 # Create 2D offset grid
-                 input_offsets = input_base + offs_ih[:, None] * in_width + offs_iw[None, :]
-                 
-                 # Create mask (input positions are always valid since we only compute valid output positions)
-                 mask = mask_oh[:, None] & mask_ow[None, :]
-                 
-                 # Load input values
-                 x = tl.load(input_ptr + input_offsets, mask=mask, other=0.0)
-                 
-                 # Accumulate: acc += x * w
-                 acc += x * w
-     
-     # Store output
-     # Output layout: [batch, out_channels, out_height, out_width]
-     output_base = (pid_b * out_channels * out_height * out_width +
-                    pid_oc * out_height * out_width)
-     output_offsets = output_base + offs_oh[:, None] * out_width + offs_ow[None, :]
-     output_mask = mask_oh[:, None] & mask_ow[None, :]
-     
-     tl.store(output_ptr + output_offsets, acc, mask=output_mask)
+     pid_m = tl.program_id(axis=0)  # flatten over (B, OC, OH)
+     pid_n = tl.program_id(axis=1)  # tile along OW
   
+     # Decompose pid_m into (b, oc, oh)
+     tmp = pid_m
+     oh = tmp % OH
+     tmp = tmp // OH
+     oc = tmp % OC
+     b = tmp // OC
   
- def kernel_function(input_tensor: torch.Tensor, kernel: torch.Tensor, output_tensor: torch.Tensor) -> torch.Tensor:
+     # Compute tile offsets along width
+     col_start = pid_n * BLOCK_W
+     offs_ow = col_start + tl.arange(0, BLOCK_W)
+     mask = offs_ow < OW
+  
+     # Initialize accumulator (vector across BLOCK_W)
+     acc = tl.zeros((BLOCK_W,), dtype=tl.float32)
+  
+     # 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)
+  
+                 # 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)
+  
+                 # FMA accumulate
+                 acc += x_vec_f32 * w_val_f32
+  
+     # 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)
+  
+  
+ def kernel_function(x: torch.Tensor, w: torch.Tensor, out: torch.Tensor = None):
      """
-     Wrapper for 2D convolution using Triton.
-     
-     This is a fused implementation that computes the entire convolution in a single kernel:
-     - For each output position, accumulates over all input channels and kernel spatial positions
-     - No separate im2col or matrix multiplication steps
-     
+     Triton-backed 2D convolution without padding and stride=1.
+  
+     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.
+  
      Args:
-         input_tensor: Input tensor of shape [batch, in_channels, height, width]
-         kernel: Convolution kernel of shape [out_channels, in_channels, kH, kW]
-         output_tensor: Pre-allocated output tensor of shape [batch, out_channels, oH, oW]
-         
+         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].
+  
      Returns:
-         output_tensor filled with convolution result
+         Output tensor y of shape [B, C_out, H-K+1, W-K+1], dtype = x.dtype, device = x.device.
      """
-     # Extract dimensions
-     batch, in_channels, in_height, in_width = input_tensor.shape
-     out_channels, _, kernel_h, kernel_w = kernel.shape
-     
-     # For this problem, kernel is square and in_channels == out_channels
-     assert kernel_h == kernel_w, "Only square kernels supported"
-     kernel_size = kernel_h
-     
-     # Output dimensions (stride=1, padding=0)
-     out_height = in_height - kernel_size + 1
-     out_width = in_width - kernel_size + 1
-     
-     # Verify output shape
-     assert output_tensor.shape == (batch, out_channels, out_height, out_width)
-     
-     # Block sizes for output spatial dimensions
-     BLOCK_OH = 8
-     BLOCK_OW = 8
-     
-     # Grid dimensions
-     num_blocks_oh = triton.cdiv(out_height, BLOCK_OH)
-     num_blocks_ow = triton.cdiv(out_width, BLOCK_OW)
-     num_spatial_blocks = num_blocks_oh * num_blocks_ow
-     
-     grid = (batch, out_channels, num_spatial_blocks)
-     
-     # Launch kernel
-     conv2d_kernel[grid](
-         input_tensor,
-         kernel,
-         output_tensor,
-         batch,
-         in_channels,
-         out_channels,
-         in_height,
-         in_width,
-         kernel_size,
-         out_height,
-         out_width,
-         BLOCK_OH,
-         BLOCK_OW,
+     # 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"
+  
+     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
+     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"
+  
+     # Prepare output
+     if out is None:
+         out = torch.empty((B, OC, 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"
+  
+     # Extract strides
+     stride_xb, stride_xc, stride_xh, stride_xw = x.stride()
+     stride_wo, stride_wc, 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"]))
+  
+     # Launch the kernel
+     conv2d_nopad_kernel[grid](
+         x, w, out,
+         B, IC, OC, H, W, OH, OW,
+         stride_xb, stride_xc, stride_xh, stride_xw,
+         stride_wo, stride_wc, stride_wkh, stride_wkw,
+         stride_yb, stride_yc, stride_yh, stride_yw,
+         K=K,
      )
-     
-     return output_tensor
   
+     return out
+  
  import inspect
   
  def custom_kernel(input):