submission 554374

+ """Variant D: No pad + pre-expanded weight/bias — no F.pad, no i_bd % D."""
  from task import input_t, output_t
   
  import os
  os.environ["HELION_AUTOTUNE_EFFORT"] = "none"
   
  import torch
- import torch.nn.functional as F
  import helion
  import helion.language as hl
   
   
- # Per-shape configs: map (B, D, S, W) to optimized helion.Config objects.
- # Autotune locally for each shape, then paste the best config here.
- SHAPE_CONFIGS: dict[tuple, helion.Config] = {
-     (1, 64, 64, 4):    helion.Config(block_sizes=[64],   num_warps=16, num_stages=2),
-     (2, 128, 128, 4):  helion.Config(block_sizes=[128],  num_warps=4,  num_stages=1),
-     (1, 256, 256, 3):  helion.Config(block_sizes=[256],  num_warps=1,  num_stages=1),
-     (1, 128, 64, 8):   helion.Config(block_sizes=[256],  num_warps=16, num_stages=1),
-     (4, 64, 128, 4):   helion.Config(block_sizes=[128],  num_warps=1,  num_stages=2),
-     (1, 1536, 2048, 4): helion.Config(block_sizes=[1024], num_warps=16, num_stages=1),
-     (1, 2560, 2048, 4): helion.Config(block_sizes=[1024], num_warps=16, num_stages=1),
-     (1, 2560, 4096, 4): helion.Config(block_sizes=[1024], num_warps=16, num_stages=1),
- }
+ @helion.kernel(
+     config=helion.Config(block_sizes=[512], num_warps=4, num_stages=1),
+     static_shapes=False,
+ )
+ def _kernel(
+     x_flat: torch.Tensor,  # (B*D, S) original input, flattened — no padding
+     w_exp: torch.Tensor,   # (B*D, W) expanded weight
+     b_exp: torch.Tensor,   # (B*D,) expanded bias
+ ) -> torch.Tensor:
+     BD, S = x_flat.shape
+     W = hl.specialize(w_exp.size(1))
   
+     out_flat = torch.empty(BD, S, dtype=x_flat.dtype, device=x_flat.device)
   
- def _make_kernel(config: helion.Config):
-     @helion.kernel(static_shapes=True, config=config)
-     def kernel(
-         x_pad_flat: torch.Tensor,  # (B*D, L) flattened padded input
-         w: torch.Tensor,           # (D, W) filter coefficients
-         b: torch.Tensor,           # (D,) additive offset
-     ) -> torch.Tensor:
-         BD, L = x_pad_flat.shape
-         D = hl.specialize(w.size(0))
-         W = hl.specialize(w.size(1))
-         S = L - W + 1
+     for tile_bd in hl.tile(BD, block_size=1):
+         i_bd = tile_bd.id
+         b_val = b_exp[i_bd]
   
-         out_flat = torch.empty(BD, S, dtype=x_pad_flat.dtype, device=x_pad_flat.device)
+         for tile_s in hl.tile(S):
+             x_cur = hl.load(x_flat, [i_bd, tile_s.index])
+             # Last tap (j = W-1): always in-bounds
+             acc = x_cur * w_exp[i_bd, W - 1]
+             # Earlier taps: masked (causal zeros)
+             for j in range(W - 1):
+                 loaded = hl.inline_triton(
+                     """
+                     idx = {s} + {offset}
+                     tl.load({x} + {ibd} * {stride} + idx, mask=(idx >= 0), other=0.0)
+                     """,
+                     args={
+                         "x": x_flat,
+                         "ibd": i_bd,
+                         "stride": x_flat.stride(0),
+                         "s": tile_s.index,
+                         "offset": j - (W - 1),  # static compile-time int
+                     },
+                     output_like=x_cur,
+                 )
+                 acc = acc + loaded * w_exp[i_bd, j]
+             out_flat[i_bd, tile_s] = acc + b_val
   
-         for tile_bd in hl.tile(BD, block_size=1):
-             i_bd = tile_bd.id
-             b_val = b[i_bd % D].to(torch.float32)
+     return out_flat
   
-             for tile_s in hl.tile(S):
-                 acc = hl.zeros([tile_s], dtype=torch.float32)
-                 for j in range(W):
-                     coeff = w[i_bd % D, j].to(torch.float32)
-                     xv = hl.load(x_pad_flat, [i_bd, tile_s.index + j]).to(torch.float32)
-                     acc = acc + xv * coeff
-                 out_flat[i_bd, tile_s] = (acc + b_val).to(out_flat.dtype)
   
-         return out_flat
-  
-     return kernel
-  
-  
- _KERNELS = {shape: _make_kernel(cfg) for shape, cfg in SHAPE_CONFIGS.items()}
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      x, weight, bias = data
      B, D, S = x.shape
      W = weight.shape[1]
-     kernel = _KERNELS[(B, D, S, W)]
-     x_pad = F.pad(x, (W - 1, 0))
-     x_pad_flat = x_pad.reshape(B * D, S + W - 1)
-     return kernel(x_pad_flat, weight, bias).reshape(B, D, S)
+  
+     x_flat = x.reshape(B * D, S)
+     w_exp = weight.unsqueeze(0).expand(B, -1, -1).reshape(B * D, W).contiguous()
+     b_exp = bias.unsqueeze(0).expand(B, -1).reshape(B * D).contiguous()
+  
+     return _kernel(x_flat, w_exp, b_exp).reshape(B, D, S)