Function reduce_sum_kernel

kernels/reduce.py:21–48

Parallel sum reduction. One program per output element. Reduces over `reduce_size` elements with stride `stride_x_col`.

(
    X_ptr,
    OUT_ptr,
    reduce_size,
    stride_x_row,
    stride_x_col,
    BLOCK_SIZE: tl.constexpr,
)

@triton.jit
def reduce_sum_kernel(
    X_ptr,
    OUT_ptr,
    reduce_size,
    stride_x_row,
    stride_x_col,
    BLOCK_SIZE: tl.constexpr,
):
    """
    Parallel sum reduction. One program per output element.
    Reduces over `reduce_size` elements with stride `stride_x_col`.
    """
    row_idx = tl.program_id(0)

    # Base pointer for this row
    row_start = X_ptr + row_idx * stride_x_row

    # Accumulate in float32 for stability
    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)

    for offset in range(0, reduce_size, BLOCK_SIZE):
        col_offsets = offset + tl.arange(0, BLOCK_SIZE)
        mask = col_offsets < reduce_size
        x = tl.load(row_start + col_offsets * stride_x_col, mask=mask, other=0.0).to(tl.float32)
        acc += x

    result = tl.sum(acc, axis=0)
    tl.store(OUT_ptr + row_idx, result)


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