submission 554356

#!POPCORN leaderboard causal_conv1d
#!POPCORN gpu B200_Nebius

import os

os.environ["ENABLE_TILE"] = "1"
os.environ["HELION_BACKEND"] = "tileir"

from task import input_t, output_t

import torch
import helion
import helion.language as hl


# VG2: TileIR — hardcoded best config to avoid autotuning timeout on KernelBot.
_BEST_CONFIG = helion.Config(block_sizes=[8, 128], num_ctas=1, occupancy=4, num_stages=4, indexing="tensor_descriptor")


@helion.kernel(static_shapes=True, config=_BEST_CONFIG)
def _causal_conv1d(
    x_pad: torch.Tensor,  # [B, D, S+W-1] fp32
    w: torch.Tensor,       # [D, W] fp32
    b: torch.Tensor,       # [D] fp32
) -> torch.Tensor:
    B = x_pad.size(0)
    D = x_pad.size(1)
    L = x_pad.size(2)
    W = hl.specialize(w.size(1))
    S = L - W + 1

    y = torch.empty(B, D, S, dtype=x_pad.dtype, device=x_pad.device)

    for rb, rd, rs in hl.tile([B, D, S], block_size=[1, None, None]):
        bi = rb.begin
        acc = hl.zeros([rd, rs], dtype=torch.float32)

        for j in range(W):
            coeff = w[rd, j].to(torch.float32)
            x_val = hl.load(x_pad, [bi, rd, rs.index + j]).to(torch.float32)
            acc = acc + x_val * coeff[:, None]

        acc = acc + b[rd].to(torch.float32)[:, None]
        y[rb, rd, rs] = acc[None, :, :].to(y.dtype)

    return y


def custom_kernel(data: input_t) -> output_t:
    x, weight, bias = data
    B, D, S = x.shape
    W = weight.shape[1]

    pad_zeros = torch.zeros(B, D, W - 1, dtype=x.dtype, device=x.device)
    padded = torch.cat([pad_zeros, x], dim=2)

    return _causal_conv1d(padded, weight, bias)