|
| 1 | +""" |
| 2 | +Jagged Mean Example |
| 3 | +=============== |
| 4 | +
|
| 5 | +This example demonstrates how to compute the mean of each row in a jagged tensor |
| 6 | +with variable features per row using Helion. |
| 7 | +""" |
| 8 | + |
| 9 | +# %% |
| 10 | +# Imports |
| 11 | +# ------- |
| 12 | +from __future__ import annotations |
| 13 | + |
| 14 | +from typing import Callable |
| 15 | + |
| 16 | +import torch |
| 17 | + |
| 18 | +import helion |
| 19 | +from helion._testing import run_example |
| 20 | +import helion.language as hl |
| 21 | + |
| 22 | + |
| 23 | +# %% |
| 24 | +# Jagged Mean Kernel |
| 25 | +# --------------- |
| 26 | +@helion.kernel() |
| 27 | +def jagged_sum_kernel( |
| 28 | + x_data: torch.Tensor, |
| 29 | + x_offsets: torch.Tensor, |
| 30 | +) -> torch.Tensor: |
| 31 | + """ |
| 32 | + Compute the mean of each row in a jagged tensor with variable features per row. |
| 33 | +
|
| 34 | + Args: |
| 35 | + x_data: 2-D tensor of shape (total_elements, M) holding all elements |
| 36 | + x_offsets: (num_rows + 1) tensor. Row i is the slice |
| 37 | + x_data[x_offsets[i] : x_offsets[i+1], :] |
| 38 | +
|
| 39 | + Returns: |
| 40 | + 2-D tensor of shape (num_rows, M) containing the sum of jagged dimension. |
| 41 | + """ |
| 42 | + M = x_data.shape[1] |
| 43 | + num_rows = x_offsets.size(0) - 1 |
| 44 | + |
| 45 | + out = torch.zeros([num_rows, M], dtype=x_data.dtype, device=x_data.device) |
| 46 | + |
| 47 | + # Flatten x_data for easier indexing |
| 48 | + x_flat = x_data.view(-1) |
| 49 | + |
| 50 | + # Process rows in tiles |
| 51 | + for tile_b in hl.tile(num_rows): |
| 52 | + starts = x_offsets[tile_b] |
| 53 | + ends = x_offsets[tile_b.index + 1] |
| 54 | + nnz = ends - starts |
| 55 | + max_nnz = nnz.amax() |
| 56 | + |
| 57 | + # Process features in tiles |
| 58 | + for tile_m in hl.tile(M): |
| 59 | + # Initialize accumulator |
| 60 | + row_sums = hl.zeros([tile_b, tile_m], dtype=x_data.dtype) |
| 61 | + |
| 62 | + # Process elements within each row |
| 63 | + for tile_k in hl.tile(0, max_nnz): |
| 64 | + # Compute flattened indices |
| 65 | + base_indices = starts[:, None] + tile_k.index[None, :] |
| 66 | + flat_indices = ( |
| 67 | + base_indices[:, :, None] * M + tile_m.index[None, None, :] |
| 68 | + ) |
| 69 | + |
| 70 | + # Combined mask: valid row element AND valid feature |
| 71 | + row_mask = tile_k.index[None, :] < nnz[:, None] |
| 72 | + combined_mask = row_mask[:, :, None] |
| 73 | + |
| 74 | + x_slice = hl.load( |
| 75 | + x_flat, |
| 76 | + [flat_indices], |
| 77 | + extra_mask=combined_mask, |
| 78 | + ) |
| 79 | + # Accumulate - sum across the k dimension (dim=1) |
| 80 | + row_sums = row_sums + x_slice.sum(dim=1) |
| 81 | + |
| 82 | + # Apply feature mask to output |
| 83 | + out[tile_b, tile_m] = row_sums |
| 84 | + |
| 85 | + return out |
| 86 | + |
| 87 | + |
| 88 | +# %% |
| 89 | +# Reference Implementation |
| 90 | +# -------------------- |
| 91 | +def reference_jagged_sum_kernel_pytorch( |
| 92 | + x_data: torch.Tensor, |
| 93 | + x_offsets: torch.Tensor, |
| 94 | +) -> torch.Tensor: |
| 95 | + """ |
| 96 | + PyTorch reference implementation for jagged mean with variable features. |
| 97 | +
|
| 98 | + Args: |
| 99 | + x_data: 2-D tensor holding all elements |
| 100 | + x_offsets: Offsets tensor for row indexing |
| 101 | +
|
| 102 | + Returns: |
| 103 | + Tensor containing the mean of each row |
| 104 | + """ |
| 105 | + num_rows = x_offsets.numel() - 1 |
| 106 | + M = x_data.size(1) |
| 107 | + out = torch.zeros((num_rows, M), dtype=x_data.dtype, device=x_data.device) |
| 108 | + for i in range(num_rows): |
| 109 | + start = int(x_offsets[i]) |
| 110 | + end = int(x_offsets[i + 1]) |
| 111 | + if end > start: |
| 112 | + out[i, :] = x_data[start:end, :].sum(dim=0) |
| 113 | + return out |
| 114 | + |
| 115 | + |
| 116 | +# %% |
| 117 | +# Benchmark Wrapper |
| 118 | +# -------------- |
| 119 | +def jagged_sum_tritonbench( |
| 120 | + tb_op: object, x: torch.Tensor, B: int, M: int, seqlen: int, sparsity: float |
| 121 | +) -> Callable[[], torch.Tensor]: |
| 122 | + """ |
| 123 | + Wrapper for tritonbench that matches the expected interface. |
| 124 | +
|
| 125 | + Args: |
| 126 | + tb_op: TritonBench operator instance |
| 127 | + x: Nested tensor in jagged format with shape (B, *, M) |
| 128 | + B: Batch size |
| 129 | + M: Number of features |
| 130 | + seqlen: Maximum sequence length |
| 131 | + sparsity: Sparsity factor (not used) |
| 132 | +
|
| 133 | + Returns: |
| 134 | + Callable that returns tensor of shape (B, M) with mean values per row and feature |
| 135 | + """ |
| 136 | + x_values = x._values |
| 137 | + x_offsets = x._offsets # pyright: ignore[reportAttributeAccessIssue] |
| 138 | + |
| 139 | + return lambda: jagged_sum_kernel(x_values, x_offsets) |
| 140 | + |
| 141 | + |
| 142 | +# %% |
| 143 | +# Helper function to create test data |
| 144 | +# --------------------------------- |
| 145 | +def create_test_jagged_tensor( |
| 146 | + B: int, |
| 147 | + M: int, |
| 148 | + max_seqlen: int, |
| 149 | + device: str = "cuda", |
| 150 | + dtype: torch.dtype = torch.float32, |
| 151 | +) -> tuple[torch.Tensor, torch.Tensor]: |
| 152 | + """Create test jagged tensor data.""" |
| 153 | + |
| 154 | + # Generate random sequence lengths |
| 155 | + seq_lengths = torch.randint(1, max_seqlen + 1, (B,), device=device) |
| 156 | + |
| 157 | + # Create offsets |
| 158 | + x_offsets = torch.cat( |
| 159 | + [ |
| 160 | + torch.zeros(1, dtype=torch.long, device=device), |
| 161 | + torch.cumsum(seq_lengths, dim=0), |
| 162 | + ] |
| 163 | + ) |
| 164 | + |
| 165 | + # Create values |
| 166 | + nnz = int(x_offsets[-1]) |
| 167 | + x_data = torch.randn(nnz, M, dtype=dtype, device=device) |
| 168 | + |
| 169 | + return x_data, x_offsets |
| 170 | + |
| 171 | + |
| 172 | +# %% |
| 173 | +# Main Function |
| 174 | +# ----------- |
| 175 | +def main() -> None: |
| 176 | + """ |
| 177 | + Main entry point that runs the jagged mean kernel verification. |
| 178 | +
|
| 179 | + Creates test data with random jagged tensors and feature counts, then compares |
| 180 | + the kernel implementation against the PyTorch reference implementation. |
| 181 | + """ |
| 182 | + B, M, max_seqlen = 8, 128, 64 |
| 183 | + device = "cuda" |
| 184 | + |
| 185 | + x_data, x_offsets = create_test_jagged_tensor( |
| 186 | + B, M, max_seqlen, device, dtype=torch.float32 |
| 187 | + ) |
| 188 | + |
| 189 | + run_example( |
| 190 | + lambda x, o: jagged_sum_kernel(x, o), |
| 191 | + lambda x, o: reference_jagged_sum_kernel_pytorch(x, o), |
| 192 | + (x_data, x_offsets), |
| 193 | + ) |
| 194 | + |
| 195 | + |
| 196 | +if __name__ == "__main__": |
| 197 | + main() |
0 commit comments