submission 553113

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),
}


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

        out_flat = torch.empty(BD, S, dtype=x_pad_flat.dtype, device=x_pad_flat.device)

        for tile_bd in hl.tile(BD, block_size=1):
            i_bd = tile_bd.id
            b_val = b[i_bd % D].to(torch.float32)

            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)