submission 758655

  ⋯ 10 unchanged lines
   
  @triton.jit
  def conv2d_kernel(
-     input_ptr, weight_ptr, output_ptr,
-     batch,
-     H, W,
-     out_H, out_W,
-     stride_h, stride_w,
+     input_ptr,
+     weight_ptr,
+     output_ptr,
+     B: tl.constexpr,
+     C: tl.constexpr,
+     H: tl.constexpr,
+     W: tl.constexpr,
+     kH: tl.constexpr,
+     kW: tl.constexpr,
+     out_H: tl.constexpr,
+     out_W: tl.constexpr,
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
-     BLOCK_K: tl.constexpr,
-     C: tl.constexpr,
-     KH: tl.constexpr,
-     KW: tl.constexpr,
+     BLOCK_O: tl.constexpr,
  ):
-     """Direct 2‑D convolution (no padding, stride = 1).
+     """
+     Triton kernel for a 2‑D convolution (no padding, stride = 1).
   
-     Shapes (contiguous):
-         input  : [batch, C, H, W]
-         weight : [C, C, KH, KW]   (output channels == input channels == C)
-         output : [batch, C, out_H, out_W] where out_H = H‑KH+1, out_W = W‑KW+1
+     Each program instance computes a tile of shape
+         BLOCK_O (output channels) × BLOCK_M (output height) × BLOCK_N (output width)
+     for one batch element.
      """
-     pid_b = tl.program_id(0)   # batch index
-     pid_ocb = tl.program_id(1) # output‑channel block index
-     pid_sp = tl.program_id(2)  # flattened spatial block index
+     pid0 = tl.program_id(0)  # batch * output‑channel‑tiles
+     pid1 = tl.program_id(1)  # tile over output height
+     pid2 = tl.program_id(2)  # tile over output width
   
-     # output‑channel base for this block
-     oc_start = pid_ocb * BLOCK_K
+     # -------------------------------------------------------------------------
+     # Output‑channel tiling
+     # -------------------------------------------------------------------------
+     num_oc_tiles = (C + BLOCK_O - 1) // BLOCK_O
+     batch = pid0 // num_oc_tiles
+     oc_tile = pid0 % num_oc_tiles
+     oc_start = oc_tile * BLOCK_O
   
-     # number of column blocks (computed at runtime)
-     num_col_blocks = tl.cdiv(out_W, BLOCK_N)
+     oc_range = oc_start + tl.arange(0, BLOCK_O)               # (BLOCK_O,)
+     oc_mask = oc_range < C                                    # (BLOCK_O,)
   
-     # derive row / column block indices from flattened spatial id
-     col_block = pid_sp % num_col_blocks
-     row_block = pid_sp // num_col_blocks
+     # -------------------------------------------------------------------------
+     # Spatial tile coordinates
+     # -------------------------------------------------------------------------
+     y = tl.arange(0, BLOCK_M)                                 # (BLOCK_M,)
+     x = tl.arange(0, BLOCK_N)                                 # (BLOCK_N,)
+     out_y = pid1 * BLOCK_M + y                                # (BLOCK_M,)
+     out_x = pid2 * BLOCK_N + x                                # (BLOCK_N,)
   
-     # top‑left corner of the tile in output coordinates
-     row_start = row_block * BLOCK_M
-     col_start = col_block * BLOCK_N
+     mask_y = out_y < out_H
+     mask_x = out_x < out_W
+     mask_spatial = mask_y[:, None] & mask_x[None, :]          # (BLOCK_M, BLOCK_N)
   
-     # per‑thread offsets within the tile
-     row_offsets = row_start + tl.arange(0, BLOCK_M)
-     col_offsets = col_start + tl.arange(0, BLOCK_N)
+     # -------------------------------------------------------------------------
+     # Accumulator for the output tile
+     # -------------------------------------------------------------------------
+     acc = tl.zeros((BLOCK_O, BLOCK_M, BLOCK_N), dtype=tl.float32)
   
-     # masks for valid output positions
-     mask_row = row_offsets < out_H
-     mask_col = col_offsets < out_W
-     mask_rc = mask_row[:, None] & mask_col[None, :]   # (BLOCK_M, BLOCK_N)
+     # Pre‑compute constants for address calculations
+     C_kH_kW = C * kH * kW
+     kH_kW = kH * kW
   
-     # output‑channel offsets handled by this block
-     oc_offsets = oc_start + tl.arange(0, BLOCK_K)
-     mask_oc = oc_offsets < C
+     # -------------------------------------------------------------------------
+     # Main convolution loops (static – unrolled)
+     # -------------------------------------------------------------------------
+     for ic in range(C):
+         ic_offset = ic * kH_kW
+         for kh in range(kH):
+             for kw in range(kW):
+                 # Input coordinates for the current kernel position
+                 in_y = out_y + kh                                 # (BLOCK_M,)
+                 in_x = out_x + kw                                 # (BLOCK_N,)
   
-     # accumulator for the tile (BLOCK_K, BLOCK_M, BLOCK_N)
-     acc = tl.zeros((BLOCK_K, BLOCK_M, BLOCK_N), dtype=tl.float32)
+                 # Linear offset into the input tensor:
+                 # ((batch*C + ic) * H + in_y) * W + in_x
+                 row = ((batch * C + ic) * H + in_y) * W
+                 offset = row[:, None] + in_x[None, :]            # (BLOCK_M, BLOCK_N)
   
-     # strides for the input tensor (NCHW layout)
-     stride_input_batch = C * H * W
-     stride_input_c = H * W
-     stride_input_h = W
-     stride_input_w = 1
+                 a = tl.load(
+                     input_ptr + offset,
+                     mask=mask_spatial,
+                     other=0.0,
+                 )                                                # (BLOCK_M, BLOCK_N)
   
-     # strides for the weight tensor (OC, IC, KH, KW)
-     stride_weight_oc = C * KH * KW
-     stride_weight_ic = KH * KW
-     stride_weight_kh = KW
-     stride_weight_kw = 1
+                 # Linear offset into the weight tensor for each output channel:
+                 # oc * C * kH * kW + ic * kH * kW + kh * kW + kw
+                 weight_offset = oc_range * C_kH_kW + ic_offset + kh * kW + kw
+                 w = tl.load(
+                     weight_ptr + weight_offset,
+                     mask=oc_mask,
+                     other=0.0,
+                 )                                                # (BLOCK_O,)
   
-     # reduction over input channel and kernel spatial dimensions
-     for ic in range(C):
-         for kh in range(KH):
-             for kw in range(KW):
-                 # input coordinates for this kernel element
-                 in_row = row_offsets + kh
-                 in_col = col_offsets + kw
+                 # Accumulate: broadcast w over the spatial tile
+                 acc += w[:, None, None] * a[None, :, :]
   
-                 # flat offsets for the input tile
-                 offset_input = (
-                     pid_b * stride_input_batch
-                     + ic * stride_input_c
-                     + in_row[:, None] * stride_input_h
-                     + in_col[None, :] * stride_input_w
-                 )
-                 # load input tile (masked)
-                 inp = tl.load(
-                     input_ptr + offset_input,
-                     mask=mask_rc,
-                     other=0.0
-                 )  # (BLOCK_M, BLOCK_N)
+     # -------------------------------------------------------------------------
+     # Write the result back to the output tensor
+     # -------------------------------------------------------------------------
+     oc_b = oc_range[:, None, None]               # (BLOCK_O, 1, 1)
+     out_y_b = out_y[None, :, None]               # (1, BLOCK_M, 1)
+     out_x_b = out_x[None, None, :]               # (1, 1, BLOCK_N)
   
-                 # flat offsets for the weight slice (all output channels in this block)
-                 offset_weight = (
-                     oc_offsets * stride_weight_oc
-                     + ic * stride_weight_ic
-                     + kh * stride_weight_kh
-                     + kw * stride_weight_kw
-                 )
-                 w = tl.load(
-                     weight_ptr + offset_weight,
-                     mask=mask_oc,
-                     other=0.0
-                 )   # (BLOCK_K,)
+     out_offset = ((batch * C + oc_b) * out_H + out_y_b) * out_W + out_x_b  # (BLOCK_O, BLOCK_M, BLOCK_N)
   
-                 # accumulate
-                 acc += w[:, None, None] * inp[None, :, :]
+     # Combine masks for channels and spatial positions
+     mask = mask_spatial[None, :, :] & oc_mask[:, None, None]               # (BLOCK_O, BLOCK_M, BLOCK_N)
   
-     # write the result to the output tensor
-     stride_out_batch = C * out_H * out_W
-     stride_out_oc = out_H * out_W
-     stride_out_h = out_W
-     stride_out_w = 1
+     tl.store(output_ptr + out_offset, acc, mask=mask)
   
-     offset_out = (
-         pid_b * stride_out_batch
-         + oc_offsets[:, None, None] * stride_out_oc
-         + row_offsets[None, :, None] * stride_out_h
-         + col_offsets[None, None, :] * stride_out_w
-     )
-     mask_store = mask_oc[:, None, None] & mask_row[None, :, None] & mask_col[None, None, :]
-     tl.store(output_ptr + offset_out, acc, mask=mask_store)
   
-  
  def custom_kernel(data: input_t) -> output_t:
-     """Triton implementation of `torch.nn.functional.conv2d` (stride = 1, no padding).
+     """
+     Compute a 2‑D convolution (no padding, stride = 1) using a Triton kernel.
   
      Args:
-         data: tuple of (input_tensor, kernel, output) where
-             * input_tensor shape = [batch, C, H, W]
-             * kernel shape      = [C, C, kH, kW] (output channels == input channels)
-             * output shape      = [batch, C, H‑kH+1, W‑kW+1]
+         data: Tuple (input_tensor, kernel, output) where
+             input_tensor: [B, C, H, W]  float32 CUDA tensor
+             kernel:       [C, C, kH, kW] float32 CUDA tensor
+             output:       pre‑allocated [B, C, H‑kH+1, W‑kW+1] float32 CUDA tensor
   
      Returns:
-         The `output` tensor filled with the convolution result.
+         The output tensor filled with the convolution result.
      """
      input_tensor, kernel, output = data
   
-     # Ensure contiguous layout (required for pointer arithmetic)
+     # -------------------------------------------------------------------------
+     # Ensure tensors are on CUDA, contiguous and have matching device
+     # -------------------------------------------------------------------------
+     if not input_tensor.is_cuda:
+         input_tensor = input_tensor.cuda()
+     if not kernel.is_cuda:
+         kernel = kernel.cuda()
+     if not output.is_cuda:
+         output = output.cuda()
      input_tensor = input_tensor.contiguous()
      kernel = kernel.contiguous()
      output = output.contiguous()
   
-     # Extract shapes
-     batch, C, H, W = input_tensor.shape
+     B, C, H, W = input_tensor.shape
      _, _, kH, kW = kernel.shape
      out_H = H - kH + 1
      out_W = W - kW + 1
   
-     # Tiling configuration (tuned for Hopper SMs)
-     BLOCK_M = 32   # rows of the output tile
-     BLOCK_N = 32   # columns of the output tile
-     BLOCK_K = 16   # number of output channels processed per block
+     assert output.shape == (B, C, out_H, out_W), "Output tensor has incorrect shape"
   
+     # -------------------------------------------------------------------------
+     # Tiling configuration (product BLOCK_O*BLOCK_M*BLOCK_N = 2048 to stay within
+     # register budget on Hopper). The new configuration widens the width tile
+     # (BLOCK_N) to improve memory coalescing while keeping the total tile size
+     # unchanged.
+     # -------------------------------------------------------------------------
+     BLOCK_O = 16   # output channels per program
+     BLOCK_M = 8    # output height per program
+     BLOCK_N = 16   # output width per program  (8*16*16 = 2048)
+  
      # Grid dimensions:
-     #   dim0 -> batch
-     #   dim1 -> output‑channel blocks
-     #   dim2 -> spatial blocks (row blocks * col blocks)
-     num_oc_blocks = (C + BLOCK_K - 1) // BLOCK_K
-     num_row_blocks = (out_H + BLOCK_M - 1) // BLOCK_M
-     num_col_blocks = (out_W + BLOCK_N - 1) // BLOCK_N
-     num_spatial_blocks = num_row_blocks * num_col_blocks
+     #   dim0: batch * output‑channel‑tiles
+     #   dim1: tiles over output height
+     #   dim2: tiles over output width
+     num_oc_tiles = (C + BLOCK_O - 1) // BLOCK_O
+     grid = (
+         B * num_oc_tiles,
+         (out_H + BLOCK_M - 1) // BLOCK_M,
+         (out_W + BLOCK_N - 1) // BLOCK_N,
+     )
   
-     grid = (batch, num_oc_blocks, num_spatial_blocks)
-  
-     # Launch the Triton kernel
+     # Launch the kernel
      conv2d_kernel[grid](
          input_tensor,
          kernel,
          output,
-         batch,
-         H,
-         W,
-         out_H,
-         out_W,
-         1,  # stride_h (unused, kept for compatibility)
-         1,  # stride_w (unused)
-         BLOCK_M=BLOCK_M,
-         BLOCK_N=BLOCK_N,
-         BLOCK_K=BLOCK_K,
-         C=C,
-         KH=kH,
-         KW=kW,
-         num_warps=4,
-         num_stages=3,
+         B, C, H, W, kH, kW, out_H, out_W,
+         BLOCK_M, BLOCK_N, BLOCK_O,
+         num_warps=1,      # 32 threads per block (1 warp)
+         num_stages=5,     # deeper pipeline to hide memory latency
      )
      return output
  # EVOLVE-BLOCK-END