Function kernel_fn

kernels/reduce.py:51–140

Entry point called by bench.py. Must match reference.reduce_sum_ref signature. Args: x: Input tensor of any shape dim: Dimension to reduce over (default: -1, last dim) Returns: Tensor with the specified dimension reduced (summed).

(x: torch.Tensor, dim: int = -1)

def kernel_fn(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
    """
    Entry point called by bench.py. Must match reference.reduce_sum_ref signature.

    Args:
        x: Input tensor of any shape
        dim: Dimension to reduce over (default: -1, last dim)

    Returns:
        Tensor with the specified dimension reduced (summed).
    """
    assert x.is_cuda

    # Normalize dim
    if dim < 0:
        dim = x.ndim + dim
    assert 0 <= dim < x.ndim, f"dim {dim} out of range for tensor with {x.ndim} dims"

    # Compute shapes
    # We want to reshape to [outer, reduce_size, inner] then reduce the middle dim
    outer_size = 1
    for i in range(dim):
        outer_size *= x.size(i)

    reduce_size = x.size(dim)

    inner_size = 1
    for i in range(dim + 1, x.ndim):
        inner_size *= x.size(i)

    # Make contiguous and reshape
    x_contig = x.contiguous()

    # Output shape: same as input but with dim removed
    out_shape = list(x.shape)
    out_shape.pop(dim)
    if len(out_shape) == 0:
        out_shape = [1]

    # Total number of output elements
    n_output = outer_size * inner_size

    # For the simple case where we reduce over the last dimension
    # and inner_size == 1, we can use a straightforward approach
    if inner_size == 1:
        x_2d = x_contig.view(outer_size, reduce_size)
        out_flat = torch.empty(n_output, device=x.device, dtype=torch.float32)

        BLOCK_SIZE = triton.next_power_of_2(min(reduce_size, 8192))

        grid = (n_output,)
        reduce_sum_kernel[grid](
            x_2d,
            out_flat,
            reduce_size,
            x_2d.stride(0),
            x_2d.stride(1),
            BLOCK_SIZE=BLOCK_SIZE,