submission 554527

from task import input_t, output_t

import torch
import torch.nn.functional as F
import helion
import helion.language as hl


SHAPE_CONFIGS: dict[tuple, helion.Config] = {
    # Test shapes
    (1, 64, 64, 4): helion.Config(block_sizes=[8, 2], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[2], load_eviction_policies=['last', 'first', ''], loop_orders=[[1, 0]], num_stages=1, num_warps=8, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[True]),
    (2, 128, 128, 4): helion.Config(block_sizes=[16, 32], indexing=['tensor_descriptor', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[0, 1]], num_stages=8, num_warps=16, pid_type='xyz', range_flattens=[None, None], range_multi_buffers=[None, False], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, False], static_ranges=[False]),
    (1, 256, 256, 3): helion.Config(block_sizes=[32, 32], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[False]),
    (1, 128, 64, 8): helion.Config(block_sizes=[16, 8], indexing=['tensor_descriptor', 'pointer', 'pointer', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=2, num_warps=4, pid_type='xyz', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[True]),
    (4, 64, 128, 4): helion.Config(block_sizes=[32, 32], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[False]),
    # Benchmark shapes
    (1, 1536, 2048, 4): helion.Config(block_sizes=[64, 32], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=16, pid_type='flat', range_flattens=[None, False], range_multi_buffers=[None, True], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[False]),
    (1, 2560, 2048, 4): helion.Config(block_sizes=[8, 64], indexing=['tensor_descriptor', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'first'], loop_orders=[[1, 0]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 2], range_warp_specializes=[None, None], static_ranges=[False]),
    (1, 2560, 4096, 4): helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', ''], loop_orders=[[1, 0]], num_stages=6, num_warps=1, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[True]),
}


def _make_kernel(config: helion.Config):
    @helion.kernel(static_shapes=True, config=config)
    def kernel(
        x_pad: torch.Tensor,  # (BD, L) zero-padded input, flattened B*D
        w: torch.Tensor,      # (BD, W) filter coefficients, expanded for B
        b: torch.Tensor,      # (BD,) bias, expanded for B
        y: torch.Tensor,      # (BD, S) output (pre-allocated)
    ) -> None:
        BD = x_pad.size(0)
        L = x_pad.size(1)
        W = hl.specialize(w.size(1))
        S = L - W + 1

        for rd, rs in hl.tile([BD, S]):
            acc = hl.zeros([rd, rs], dtype=torch.float32)
            for k in range(W):
                wk = w[rd, k].to(torch.float32)[:, None]
                xk = hl.load(x_pad, [rd, rs.index + k]).to(torch.float32)
                acc = acc + xk * wk
            acc = acc + b[rd].to(torch.float32)[:, None]
            y[rd, rs] = acc.to(y.dtype)

    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]
    key = (B, D, S, W)
    if key not in _KERNELS:
        _KERNELS[key] = _make_kernel(helion.Config(block_sizes=[8, 128], num_warps=4, num_stages=1))
    kernel = _KERNELS[key]
    x_pad = F.pad(x, (W - 1, 0)).reshape(B * D, S + W - 1)
    # Expand weight and bias for B batches
    if B > 1:
        w_expanded = weight.unsqueeze(0).expand(B, -1, -1).reshape(B * D, W)
        b_expanded = bias.unsqueeze(0).expand(B, -1).reshape(B * D)
    else:
        w_expanded = weight
        b_expanded = bias
    y = torch.empty(B * D, S, dtype=x.dtype, device=x.device)
    kernel(x_pad, w_expanded, b_expanded, y)
    return y.reshape(B, D, S)