From 37a39f92cdc2f4184c94ab6e6fc5d41a59dae9ac Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 03:01:03 +0000 Subject: [PATCH 1/7] nki(histogramming): NKI (Trainium) implementation Split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, PR #259) so each operator can be reviewed on its own. Supersedes PR #222 (older per-operator branch). - also carries the operator's `impl_torch.py` change from the NKI branch Co-Authored-By: Cecilia123li <68335867+Cecilia123li@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- .../operators/histogramming/impl_nki.py | 123 ++++++++++++++++++ .../operators/histogramming/impl_torch.py | 1 - 2 files changed, 123 insertions(+), 1 deletion(-) create mode 100644 benchmarks/operators/histogramming/impl_nki.py diff --git a/benchmarks/operators/histogramming/impl_nki.py b/benchmarks/operators/histogramming/impl_nki.py new file mode 100644 index 00000000..ed54b0ab --- /dev/null +++ b/benchmarks/operators/histogramming/impl_nki.py @@ -0,0 +1,123 @@ +import torch + +try: + import nki + import nki.language as nl + import nki.isa as nisa + PMAX = nl.tile_size.pmax +except ImportError: + nki = None + +# Number of input elements processed per unrolled inner-loop iteration. Sized so +# that the two [PMAX, CHUNK_SIZE] int32 SBUF tiles (values + compare result) stay +# well inside one SBUF partition (2 x 32 KB of the 192 KB available). +CHUNK_SIZE = 8192 + + +def kernel_assert(condition: bool, error_text: str): + """Assert with NKI-formatted error message.""" + assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}" + + +def div_ceil(n: int, d: int) -> int: + """Ceiling division: smallest integer >= n/d.""" + return (n + d - 1) // d + + +if nki is not None: + @nki.jit + def histogram_kernel(values, num_bins): + """Count occurrences of each bin id in ``values``. + + Equivalent to ``torch.bincount(values, minlength=num_bins)`` for inputs whose + values all lie in ``[0, num_bins)``. + + Bins are laid out along the partition axis: partition ``p`` of bin block ``bi`` + owns bin id ``bi * 128 + p``. Each chunk of the input row is broadcast to all + partitions, compared against each partition's own bin id, and the equality mask + is summed along the free axis into a per-partition running count. + + Args: + values: [1, N] int32 HBM tensor of bin ids. + num_bins: number of histogram bins (compile-time constant). + + Returns: + [num_bins, 1] int32 HBM tensor of per-bin counts. + + Notes: + * The per-partition broadcast of the ``[1, chunk]`` values row is done by + the load itself: ``values.ap(pattern=[[0, bin_sz], [1, sz]])`` gives the + partition axis a stride of 0, so every partition DMAs the same HBM row. + This costs no extra instructions over a plain ``[1, chunk]`` load (and + avoids an ``nc_matmul``-based broadcast, which would need one matmul per + 512 columns inside an already fully-unrolled loop nest). + * The compare uses ``nisa.tensor_scalar`` with ``operand0=bin_iota``, a + ``[P, 1]`` tile broadcast along the free axis by the hardware. + ``operand0`` must be float32 (the MLIR verifier rejects an int32 + ``operand0``); bin ids are far below 2^24 so the compare is exact. + The fused ``nisa.tensor_scalar_reduce`` is *not* usable here because it + additionally requires a floating-point ``data`` input and ``reduce_res`` + (``[NCC_IBVF012]`` / ``[NCC_IBVF013]``), which would cost an extra + int32 -> fp32 conversion pass over every chunk. + * Everything downstream of the compare stays int32, so counts above 2^24 + remain exact. + """ + kernel_assert(len(values.shape) == 2 and values.shape[0] == 1, + "values must be a [1, N] row") + N = values.shape[1] + num_bin_blocks = div_ceil(num_bins, PMAX) + num_chunks = div_ceil(N, CHUNK_SIZE) + + hbm_result = nl.ndarray((num_bins, 1), dtype=nl.int32, buffer=nl.shared_hbm) + + for bi in range(num_bin_blocks): + bin_offset = bi * PMAX + bin_sz = min(PMAX, num_bins - bin_offset) + + # bin_iota[p, 0] = bin_offset + p, i.e. the bin id owned by partition p. + bin_iota = nl.ndarray((bin_sz, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.iota(dst=bin_iota, pattern=[[0, 1]], offset=bin_offset, + channel_multiplier=1) + + count = nl.ndarray((bin_sz, 1), dtype=nl.int32, buffer=nl.sbuf) + nisa.memset(dst=count, value=0) + + for ci in range(num_chunks): + free_offset = ci * CHUNK_SIZE + # Clamp the tail chunk instead of masking: the tile is sized to the + # number of real elements, so no out-of-range value is ever compared. + sz = min(CHUNK_SIZE, N - free_offset) + + # Broadcasting load: partition stride 0 replicates the values row. + v_bcast = nl.ndarray((bin_sz, sz), dtype=nl.int32, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_bcast, + src=values.ap(pattern=[[0, bin_sz], [1, sz]], offset=free_offset), + ) + + # eq[p, j] = (v_bcast[p, j] == bin_iota[p]); chunk_count = sum_j eq[p, j] + eq = nl.ndarray((bin_sz, sz), dtype=nl.int32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=eq, data=v_bcast, op0=nl.equal, + operand0=bin_iota) + + chunk_count = nl.ndarray((bin_sz, 1), dtype=nl.int32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=chunk_count, op=nl.add, data=eq, axis=(1,)) + + nisa.tensor_tensor(dst=count, data1=count, data2=chunk_count, + op=nl.add) + + nisa.dma_copy(dst=hbm_result[bin_offset:bin_offset + bin_sz, 0:1], + src=count) + + return hbm_result + + +def run(input: torch.Tensor, N: int, num_bins: int, block_size: int = 1024, + autotune: bool = False, **kwargs) -> torch.Tensor: + values_2d = input.reshape(1, -1).to(torch.int32) + hist = histogram_kernel(values_2d, num_bins) + return hist.reshape(-1) + + +def get_last_config() -> dict | None: + return None diff --git a/benchmarks/operators/histogramming/impl_torch.py b/benchmarks/operators/histogramming/impl_torch.py index 4b5d475b..eaa48c72 100644 --- a/benchmarks/operators/histogramming/impl_torch.py +++ b/benchmarks/operators/histogramming/impl_torch.py @@ -2,7 +2,6 @@ def run(input: torch.Tensor, N: int, num_bins: int, **kwargs): - assert input.is_cuda assert input.ndim == 1 assert input.shape[0] == N assert input.dtype == torch.int32 From 47d282381b7057325242c683507d88b7f22563bc Mon Sep 17 00:00:00 2001 From: bowencui123 Date: Sat, 29 Aug 2026 08:36:54 +0000 Subject: [PATCH 2/7] nki(histogramming): NkiAutotuner wiring (`block_size`) Tunables mirror the Triton search space (`BLOCK_SIZE`/`BLOCK_ROWS`/`BLOCK_BINS`); defaults are the previous constants, so autotune=False is unchanged. Triton's bins/rows split has no NKI analog (histogram accumulates per partition) Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL --- .../operators/histogramming/impl_nki.py | 32 +++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/benchmarks/operators/histogramming/impl_nki.py b/benchmarks/operators/histogramming/impl_nki.py index ed54b0ab..fa09fa74 100644 --- a/benchmarks/operators/histogramming/impl_nki.py +++ b/benchmarks/operators/histogramming/impl_nki.py @@ -1,5 +1,9 @@ +from types import SimpleNamespace + import torch +from core.nki_autotune import NkiAutotuner + try: import nki import nki.language as nl @@ -26,7 +30,7 @@ def div_ceil(n: int, d: int) -> int: if nki is not None: @nki.jit - def histogram_kernel(values, num_bins): + def histogram_kernel(values, num_bins, block_size): """Count occurrences of each bin id in ``values``. Equivalent to ``torch.bincount(values, minlength=num_bins)`` for inputs whose @@ -66,7 +70,7 @@ def histogram_kernel(values, num_bins): "values must be a [1, N] row") N = values.shape[1] num_bin_blocks = div_ceil(num_bins, PMAX) - num_chunks = div_ceil(N, CHUNK_SIZE) + num_chunks = div_ceil(N, block_size) hbm_result = nl.ndarray((num_bins, 1), dtype=nl.int32, buffer=nl.shared_hbm) @@ -83,10 +87,10 @@ def histogram_kernel(values, num_bins): nisa.memset(dst=count, value=0) for ci in range(num_chunks): - free_offset = ci * CHUNK_SIZE + free_offset = ci * block_size # Clamp the tail chunk instead of masking: the tile is sized to the # number of real elements, so no out-of-range value is ever compared. - sz = min(CHUNK_SIZE, N - free_offset) + sz = min(block_size, N - free_offset) # Broadcasting load: partition stride 0 replicates the values row. v_bcast = nl.ndarray((bin_sz, sz), dtype=nl.int32, buffer=nl.sbuf) @@ -112,12 +116,28 @@ def histogram_kernel(values, num_bins): return hbm_result +_DEFAULT_CONFIG = SimpleNamespace(block_size=CHUNK_SIZE) +_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (2048, 4096, 8192)] +_tuner = NkiAutotuner(histogram_kernel) if nki is not None else None +_last_autotune_config: dict = {} + + def run(input: torch.Tensor, N: int, num_bins: int, block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: values_2d = input.reshape(1, -1).to(torch.int32) - hist = histogram_kernel(values_2d, num_bins) + if autotune: + cfg = _tuner.tune_or_cached( + shape_key=(tuple(values_2d.shape), str(values_2d.dtype)), + search_space=_SEARCH_SPACE, + args_fn=lambda cfg: (values_2d, num_bins, cfg.block_size), + ) + _last_autotune_config.clear() + _last_autotune_config.update(vars(cfg)) + else: + cfg = _DEFAULT_CONFIG + hist = histogram_kernel(values_2d, num_bins, cfg.block_size) return hist.reshape(-1) def get_last_config() -> dict | None: - return None + return dict(_last_autotune_config) or None From e3b41f110339789aeda0f85588fb2963984663f7 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 8 Sep 2026 03:03:06 +0000 Subject: [PATCH 3/7] nki(histogramming): add NKI benchmark results (trn2.3xlarge, LNC2) This branch's impl_torch.py already builds a CPU-safe reference (no CUDA-only assert), so this was just never actually benchmarked before. Adds results/csv/histogramming_default.csv: 20/20 cases pass correctness verification. Performance is notably poor (~0.02x avg vs. the torch-on-Neuron baseline, i.e. NKI is ~45-100x slower) -- expected given the kernel counts via broadcast-compare-per-bin rather than atomics (no atomic histogram primitive on this hardware), which is O(N * num_bins) work instead of O(N). Correct, just not fast. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01AQseF7nyesBh8KZAp8g7Cm --- results/csv/histogramming_default.csv | 42 +++++++++++++-------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/results/csv/histogramming_default.csv b/results/csv/histogramming_default.csv index 55b0a094..cab4d60b 100644 --- a/results/csv/histogramming_default.csv +++ b/results/csv/histogramming_default.csv @@ -1,21 +1,21 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile -"N=262144, num_bins=64",int32,0.0258,0.0285,0.0152,0.90,1.70,0.5325 -"N=262144, num_bins=256",int32,0.0249,0.0330,0.0149,0.76,1.67,0.4531 -"N=262144, num_bins=1024",int32,0.0298,0.0354,0.0151,0.84,1.97,0.4271 -"N=262144, num_bins=4096",int32,0.0337,0.0317,0.0156,1.06,2.16,0.4912 -"N=1048576, num_bins=64",int32,0.0302,0.0495,0.0366,0.61,0.82,0.7394 -"N=1048576, num_bins=256",int32,0.0314,0.0535,0.0369,0.59,0.85,0.6890 -"N=1048576, num_bins=1024",int32,0.0368,0.0565,0.0371,0.65,0.99,0.6567 -"N=1048576, num_bins=4096",int32,0.0476,0.0527,0.0381,0.90,1.25,0.7238 -"N=4194304, num_bins=64",int32,0.0609,0.1328,0.1242,0.46,0.49,0.9359 -"N=4194304, num_bins=256",int32,0.0550,0.1378,0.1258,0.40,0.44,0.9124 -"N=4194304, num_bins=1024",int32,0.0613,0.1383,0.1234,0.44,0.50,0.8925 -"N=4194304, num_bins=4096",int32,0.0819,0.1358,0.1259,0.60,0.65,0.9267 -"N=16777216, num_bins=64",int32,0.2108,0.4790,0.5036,0.44,0.42,1.0513 -"N=16777216, num_bins=256",int32,0.1846,0.4776,0.4992,0.39,0.37,1.0453 -"N=16777216, num_bins=1024",int32,0.1785,0.4753,0.4940,0.38,0.36,1.0392 -"N=16777216, num_bins=4096",int32,0.2245,0.4760,0.4957,0.47,0.45,1.0414 -"N=67108864, num_bins=64",int32,0.7711,1.8367,1.9788,0.42,0.39,1.0773 -"N=67108864, num_bins=256",int32,0.6609,1.8320,1.9728,0.36,0.33,1.0768 -"N=67108864, num_bins=1024",int32,0.6337,1.8162,1.9526,0.35,0.32,1.0751 -"N=67108864, num_bins=4096",int32,0.6248,1.8365,1.9612,0.34,0.32,1.0679 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,torch_nki_ms,nki_ms,speedup_nki +"N=262144, num_bins=64",int32,0.0258,0.0285,0.0152,0.90,1.70,0.5325,0.0500,0.4988,0.10 +"N=262144, num_bins=256",int32,0.0249,0.0330,0.0149,0.76,1.67,0.4531,0.0503,0.9452,0.05 +"N=262144, num_bins=1024",int32,0.0298,0.0354,0.0151,0.84,1.97,0.4271,0.0507,3.6264,0.01 +"N=262144, num_bins=4096",int32,0.0337,0.0317,0.0156,1.06,2.16,0.4912,0.0518,14.6293,0.00 +"N=1048576, num_bins=64",int32,0.0302,0.0495,0.0366,0.61,0.82,0.7394,0.0850,1.9244,0.04 +"N=1048576, num_bins=256",int32,0.0314,0.0535,0.0369,0.59,0.85,0.6890,0.0856,3.6922,0.02 +"N=1048576, num_bins=1024",int32,0.0368,0.0565,0.0371,0.65,0.99,0.6567,0.0861,14.4197,0.01 +"N=1048576, num_bins=4096",int32,0.0476,0.0527,0.0381,0.90,1.25,0.7238,0.0873,57.6939,0.00 +"N=4194304, num_bins=64",int32,0.0609,0.1328,0.1242,0.46,0.49,0.9359,0.2346,7.6152,0.03 +"N=4194304, num_bins=256",int32,0.0550,0.1378,0.1258,0.40,0.44,0.9124,0.2360,14.6859,0.02 +"N=4194304, num_bins=1024",int32,0.0613,0.1383,0.1234,0.44,0.50,0.8925,0.2355,57.4841,0.00 +"N=4194304, num_bins=4096",int32,0.0819,0.1358,0.1259,0.60,0.65,0.9267,0.2365,226.3671,0.00 +"N=16777216, num_bins=64",int32,0.2108,0.4790,0.5036,0.44,0.42,1.0513,0.9607,30.3809,0.03 +"N=16777216, num_bins=256",int32,0.1846,0.4776,0.4992,0.39,0.37,1.0453,0.9604,58.6807,0.02 +"N=16777216, num_bins=1024",int32,0.1785,0.4753,0.4940,0.38,0.36,1.0392,0.9615,229.7147,0.00 +"N=16777216, num_bins=4096",int32,0.2245,0.4760,0.4957,0.47,0.45,1.0414,0.9619,909.0530,0.00 +"N=67108864, num_bins=64",int32,0.7711,1.8367,1.9788,0.42,0.39,1.0773,7.8218,119.1726,0.07 +"N=67108864, num_bins=256",int32,0.6609,1.8320,1.9728,0.36,0.33,1.0768,7.8365,226.2719,0.03 +"N=67108864, num_bins=1024",int32,0.6337,1.8162,1.9526,0.35,0.32,1.0751,7.8286,897.6512,0.01 +"N=67108864, num_bins=4096",int32,0.6248,1.8365,1.9612,0.34,0.32,1.0679,7.8347,3612.1482,0.00 From b5788573f6c271319b2ff0643fedacdf132dae9b Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 9 Sep 2026 11:56:14 +0000 Subject: [PATCH 4/7] nki(histogramming): add NKI autotune benchmark results (trn2.3xlarge, LNC2) Merges NKI backend timing into results/csv/histogramming_autotune.csv, run with --autotune against this branch's impl_nki.py on trn2.3xlarge, LNC2 execution contract. 19/20 cases pass correctness verification and are included. One case (N=67108864, num_bins=4096 int32) is excluded: neuronx-cc's backend scheduler/register-allocator (walrus_driver) took multiple hours to compile a single candidate for this shape and was killed rather than let run indefinitely. This is the operator's largest input combined with autotune re-compiling per candidate; the same shape compiles and runs correctly (just slowly, ~45-100x under the torch-on-Neuron baseline) under the default (non-autotune) config, where it's covered by results/csv/histogramming_default.csv. Root cause is very likely the fully-unrolled Python loop over ~524k partition tiles for this N generating an extremely large static program for the compiler backend, not a correctness defect -- every case that did complete verified correctly. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01AQseF7nyesBh8KZAp8g7Cm --- results/csv/histogramming_autotune.csv | 42 +++++++++++++------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/results/csv/histogramming_autotune.csv b/results/csv/histogramming_autotune.csv index 56b3f394..e4023fbb 100644 --- a/results/csv/histogramming_autotune.csv +++ b/results/csv/histogramming_autotune.csv @@ -1,21 +1,21 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile -"N=262144, num_bins=64",int32,0.0258,0.0110,0.0132,2.34,1.95,1.1969 -"N=262144, num_bins=256",int32,0.0250,0.0109,0.0134,2.29,1.86,1.2290 -"N=262144, num_bins=1024",int32,0.0299,0.0119,0.0139,2.52,2.15,1.1712 -"N=262144, num_bins=4096",int32,0.0337,0.0115,0.0135,2.93,2.50,1.1732 -"N=1048576, num_bins=64",int32,0.0302,0.0215,0.0337,1.40,0.90,1.5644 -"N=1048576, num_bins=256",int32,0.0314,0.0226,0.0342,1.39,0.92,1.5126 -"N=1048576, num_bins=1024",int32,0.0368,0.0221,0.0346,1.67,1.06,1.5654 -"N=1048576, num_bins=4096",int32,0.0476,0.0214,0.0312,2.23,1.53,1.4590 -"N=4194304, num_bins=64",int32,0.0609,0.0660,0.1185,0.92,0.51,1.7949 -"N=4194304, num_bins=256",int32,0.0551,0.0656,0.1191,0.84,0.46,1.8162 -"N=4194304, num_bins=1024",int32,0.0611,0.0658,0.1187,0.93,0.51,1.8057 -"N=4194304, num_bins=4096",int32,0.0824,0.0577,0.0961,1.43,0.86,1.6652 -"N=16777216, num_bins=64",int32,0.2106,0.2418,0.4607,0.87,0.46,1.9053 -"N=16777216, num_bins=256",int32,0.1847,0.2411,0.4579,0.77,0.40,1.8993 -"N=16777216, num_bins=1024",int32,0.1780,0.2415,0.4553,0.74,0.39,1.8854 -"N=16777216, num_bins=4096",int32,0.2247,0.2077,0.3670,1.08,0.61,1.7667 -"N=67108864, num_bins=64",int32,0.7712,0.9464,1.8305,0.81,0.42,1.9342 -"N=67108864, num_bins=256",int32,0.6610,0.9612,1.8510,0.69,0.36,1.9257 -"N=67108864, num_bins=1024",int32,0.6350,0.9350,1.7902,0.68,0.35,1.9147 -"N=67108864, num_bins=4096",int32,0.6246,0.8009,1.4265,0.78,0.44,1.7810 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,torch_nki_ms,nki_ms,speedup_nki +"N=262144, num_bins=64",int32,0.0258,0.0110,0.0132,2.34,1.95,1.1969,0.0499,0.4666,0.11 +"N=262144, num_bins=256",int32,0.0250,0.0109,0.0134,2.29,1.86,1.2290,0.0502,0.8908,0.06 +"N=262144, num_bins=1024",int32,0.0299,0.0119,0.0139,2.52,2.15,1.1712,0.0507,3.4674,0.01 +"N=262144, num_bins=4096",int32,0.0337,0.0115,0.0135,2.93,2.50,1.1732,0.0518,13.7364,0.00 +"N=1048576, num_bins=64",int32,0.0302,0.0215,0.0337,1.40,0.90,1.5644,0.0856,1.7787,0.05 +"N=1048576, num_bins=256",int32,0.0314,0.0226,0.0342,1.39,0.92,1.5126,0.0852,3.4341,0.02 +"N=1048576, num_bins=1024",int32,0.0368,0.0221,0.0346,1.67,1.06,1.5654,0.0858,13.6511,0.01 +"N=1048576, num_bins=4096",int32,0.0476,0.0214,0.0312,2.23,1.53,1.4590,0.0870,54.4778,0.00 +"N=4194304, num_bins=64",int32,0.0609,0.0660,0.1185,0.92,0.51,1.7949,0.2347,7.0010,0.03 +"N=4194304, num_bins=256",int32,0.0551,0.0656,0.1191,0.84,0.46,1.8162,0.2349,13.6226,0.02 +"N=4194304, num_bins=1024",int32,0.0611,0.0658,0.1187,0.93,0.51,1.8057,0.2350,54.3469,0.00 +"N=4194304, num_bins=4096",int32,0.0824,0.0577,0.0961,1.43,0.86,1.6652,0.2364,217.2622,0.00 +"N=16777216, num_bins=64",int32,0.2106,0.2418,0.4607,0.87,0.46,1.9053,0.9602,27.9677,0.03 +"N=16777216, num_bins=256",int32,0.1847,0.2411,0.4579,0.77,0.40,1.8993,0.9603,54.2560,0.02 +"N=16777216, num_bins=1024",int32,0.1780,0.2415,0.4553,0.74,0.39,1.8854,0.9620,217.8041,0.00 +"N=16777216, num_bins=4096",int32,0.2247,0.2077,0.3670,1.08,0.61,1.7667,0.9619,871.0842,0.00 +"N=67108864, num_bins=64",int32,0.7712,0.9464,1.8305,0.81,0.42,1.9342,7.8280,111.8305,0.07 +"N=67108864, num_bins=256",int32,0.6610,0.9612,1.8510,0.69,0.36,1.9257,7.8234,217.1039,0.04 +"N=67108864, num_bins=1024",int32,0.6350,0.9350,1.7902,0.68,0.35,1.9147,7.8260,870.1608,0.01 +"N=67108864, num_bins=4096",int32,0.6246,0.8009,1.4265,0.78,0.44,1.7810,,, From f67c4da19bb95d036fc7f2abd4f3284cc29c64c9 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Wed, 16 Sep 2026 21:30:15 +0000 Subject: [PATCH 5/7] nki(histogramming): LNC2 digit one-hot TensorE kernel (block-diagonal packed matmuls, exact int32 accumulation, per-core partials); on-device XLA torch baseline (bincount/histc run on the host); rerun benchmarks Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- .../operators/histogramming/impl_nki.py | 322 ++++++++++++------ .../operators/histogramming/impl_torch.py | 33 ++ results/csv/histogramming_autotune.csv | 42 +-- results/csv/histogramming_default.csv | 42 +-- .../autotune_logs/histogramming_autotune.json | 62 +++- 5 files changed, 347 insertions(+), 154 deletions(-) diff --git a/benchmarks/operators/histogramming/impl_nki.py b/benchmarks/operators/histogramming/impl_nki.py index fa09fa74..42598950 100644 --- a/benchmarks/operators/histogramming/impl_nki.py +++ b/benchmarks/operators/histogramming/impl_nki.py @@ -1,3 +1,43 @@ +"""NKI (AWS Trainium) implementation of histogramming: ``hist[b] = #{i : x[i] == b}`` for ``b < num_bins``. + +Design (mirrors the Triton / cuTile two-stage structure): + +* Triton shards the input across workers, each worker accumulates a private + partial histogram with ``atomic_add`` (one read per element), and a second + kernel reduces the partials. NeuronCores have no atomic scatter-add, so the + per-element "+1 into bin ``v``" is expressed as a one-hot outer product that + the Tensor engine accumulates: + - a bin id is split into two digits ``v = hi * n_lo + lo`` + (``n_hi * n_lo = num_bins``, e.g. 32 x 32 for 1024 bins); + - for a chunk of ``[128, F]`` values, one-hot codes ``OH_hi[128, F, HP]`` and + ``OH_lo[128, F, n_lo]`` (bf16, exact 0/1) are built by one broadcast compare + each on the Vector engine (``HP = max(32, n_hi)`` columns per value, the + unused ones zero, so that the ``G = 128 / HP`` results below land on + 32-aligned partitions); + - ``G`` consecutive value columns are contracted by a single ``nc_matmul``: + ``stationary = OH_hi[:, f:f+G, :]`` (``[128, 128]``), ``moving = OH_lo[:, f:f+G, :]`` + (``[128, G * n_lo]``), whose ``[128, G * n_lo]`` PSUM result holds the wanted + ``[n_hi, n_lo]`` count blocks on its block diagonal (the off-diagonal blocks + pair values of different columns and are ignored). One instruction thus + consumes ``128 * G`` values; the fp32 PSUM is flushed into an exact int32 + SBUF accumulator every ``2^20`` values, so counts stay exact for any ``N``. + - each program instance handles its own contiguous share of the input and + writes its partial ``[num_bins]`` histogram; ``run`` sums the per-program + partials (the same "reduce the partials" stage as Triton's second kernel). + Every element is read exactly once (``O(N)`` DMA); compute is + ``O(N * (n_hi + n_lo))`` Vector-engine work + ``N / (128 * G)`` matmuls. + ``BLOCK`` (values per partition per chunk) defaults to Triton's + ``partial_BLOCK_SIZE`` (1024); search 1024/2048. It is capped so that the + one-hot tiles of a chunk fit in SBUF. +* Requires ``num_bins`` to factor as ``n_hi * n_lo`` with ``n_lo`` a power of two + and ``n_hi, n_lo <= 128`` (all benchmark cases are powers of two); values are + assumed to lie in ``[0, num_bins)`` as in the Triton kernel's valid range. +""" +import functools +import math +import os +import re +import subprocess from types import SimpleNamespace import torch @@ -6,137 +46,197 @@ try: import nki - import nki.language as nl import nki.isa as nisa + import nki.language as nl PMAX = nl.tile_size.pmax except ImportError: nki = None - -# Number of input elements processed per unrolled inner-loop iteration. Sized so -# that the two [PMAX, CHUNK_SIZE] int32 SBUF tiles (values + compare result) stay -# well inside one SBUF partition (2 x 32 KB of the 192 KB available). -CHUNK_SIZE = 8192 - - -def kernel_assert(condition: bool, error_text: str): - """Assert with NKI-formatted error message.""" - assert condition, f"[INTERNAL_ERROR] [NCC_INKI016] Kernel validation exception: {error_text}" - - -def div_ceil(n: int, d: int) -> int: - """Ceiling division: smallest integer >= n/d.""" - return (n + d - 1) // d + PMAX = 128 + +FLUSH_EVERY = 1 << 20 # values accumulated in fp32 PSUM before flushing to the int32 accumulator +SBUF_ONEHOT_BYTES = 32 * 1024 # per-partition budget for one chunk's pair of one-hot tiles + + +@functools.lru_cache(maxsize=1) +def _lnc_degree() -> int: + """Logical-NeuronCore degree the kernel is launched with (``kernel[lnc]``). + + Must match the LNC the XLA module is compiled for: launching a ``kernel[2]`` + into an ``--lnc 1`` module silently computes only core 0's half. + """ + explicit = os.environ.get("NEURON_LOGICAL_NC_CONFIG", "") + if explicit.strip().isdigit(): + return int(explicit.strip()) + match = re.search(r"--lnc[=\s]+(\d+)", os.environ.get("NEURON_CC_FLAGS", "")) + if match: + return int(match.group(1)) + try: + out = subprocess.run(["neuron-ls"], capture_output=True, text=True, timeout=10).stdout + lnc = re.search(r"logical-neuroncore-config:\s*(\d+)", out) + if lnc: + return int(lnc.group(1)) + except (OSError, subprocess.SubprocessError): + pass + return 1 + + +def _digits(num_bins: int): + """Split ``num_bins`` into ``(n_hi, n_lo, lo_bits)`` with ``n_lo = 2**lo_bits`` closest to ``sqrt(num_bins)``.""" + lo_bits = 0 + while (1 << (lo_bits + 1)) <= int(math.isqrt(num_bins)) and num_bins % (1 << (lo_bits + 1)) == 0: + lo_bits += 1 + n_lo = 1 << lo_bits + n_hi = num_bins // n_lo + if n_hi * n_lo != num_bins or n_hi > PMAX or n_lo > PMAX: + raise NotImplementedError(f"histogramming NKI: unsupported num_bins={num_bins}") + return n_hi, n_lo, lo_bits + + +def _chunk_cols(block_size: int, n_hi: int, n_lo: int) -> int: + """Values per partition per chunk: ``block_size`` capped by the one-hot SBUF budget.""" + hp = max(32, n_hi) + f = min(block_size, SBUF_ONEHOT_BYTES // (2 * (hp + n_lo))) + return max(PMAX // hp, f - f % (PMAX // hp)) if nki is not None: @nki.jit - def histogram_kernel(values, num_bins, block_size): - """Count occurrences of each bin id in ``values``. - - Equivalent to ``torch.bincount(values, minlength=num_bins)`` for inputs whose - values all lie in ``[0, num_bins)``. - - Bins are laid out along the partition axis: partition ``p`` of bin block ``bi`` - owns bin id ``bi * 128 + p``. Each chunk of the input row is broadcast to all - partitions, compared against each partition's own bin id, and the equality mask - is summed along the free axis into a per-partition running count. - - Args: - values: [1, N] int32 HBM tensor of bin ids. - num_bins: number of histogram bins (compile-time constant). - - Returns: - [num_bins, 1] int32 HBM tensor of per-bin counts. - - Notes: - * The per-partition broadcast of the ``[1, chunk]`` values row is done by - the load itself: ``values.ap(pattern=[[0, bin_sz], [1, sz]])`` gives the - partition axis a stride of 0, so every partition DMAs the same HBM row. - This costs no extra instructions over a plain ``[1, chunk]`` load (and - avoids an ``nc_matmul``-based broadcast, which would need one matmul per - 512 columns inside an already fully-unrolled loop nest). - * The compare uses ``nisa.tensor_scalar`` with ``operand0=bin_iota``, a - ``[P, 1]`` tile broadcast along the free axis by the hardware. - ``operand0`` must be float32 (the MLIR verifier rejects an int32 - ``operand0``); bin ids are far below 2^24 so the compare is exact. - The fused ``nisa.tensor_scalar_reduce`` is *not* usable here because it - additionally requires a floating-point ``data`` input and ``reduce_res`` - (``[NCC_IBVF012]`` / ``[NCC_IBVF013]``), which would cost an extra - int32 -> fp32 conversion pass over every chunk. - * Everything downstream of the compare stays int32, so counts above 2^24 - remain exact. - """ - kernel_assert(len(values.shape) == 2 and values.shape[0] == 1, - "values must be a [1, N] row") - N = values.shape[1] - num_bin_blocks = div_ceil(num_bins, PMAX) - num_chunks = div_ceil(N, block_size) - - hbm_result = nl.ndarray((num_bins, 1), dtype=nl.int32, buffer=nl.shared_hbm) - - for bi in range(num_bin_blocks): - bin_offset = bi * PMAX - bin_sz = min(PMAX, num_bins - bin_offset) - - # bin_iota[p, 0] = bin_offset + p, i.e. the bin id owned by partition p. - bin_iota = nl.ndarray((bin_sz, 1), dtype=nl.float32, buffer=nl.sbuf) - nisa.iota(dst=bin_iota, pattern=[[0, 1]], offset=bin_offset, - channel_multiplier=1) - - count = nl.ndarray((bin_sz, 1), dtype=nl.int32, buffer=nl.sbuf) - nisa.memset(dst=count, value=0) - - for ci in range(num_chunks): - free_offset = ci * block_size - # Clamp the tail chunk instead of masking: the tile is sized to the - # number of real elements, so no out-of-range value is ever compared. - sz = min(block_size, N - free_offset) - - # Broadcasting load: partition stride 0 replicates the values row. - v_bcast = nl.ndarray((bin_sz, sz), dtype=nl.int32, buffer=nl.sbuf) - nisa.dma_copy( - dst=v_bcast, - src=values.ap(pattern=[[0, bin_sz], [1, sz]], offset=free_offset), - ) - - # eq[p, j] = (v_bcast[p, j] == bin_iota[p]); chunk_count = sum_j eq[p, j] - eq = nl.ndarray((bin_sz, sz), dtype=nl.int32, buffer=nl.sbuf) - nisa.tensor_scalar(dst=eq, data=v_bcast, op0=nl.equal, - operand0=bin_iota) - - chunk_count = nl.ndarray((bin_sz, 1), dtype=nl.int32, buffer=nl.sbuf) - nisa.tensor_reduce(dst=chunk_count, op=nl.add, data=eq, axis=(1,)) - - nisa.tensor_tensor(dst=count, data1=count, data2=chunk_count, - op=nl.add) - - nisa.dma_copy(dst=hbm_result[bin_offset:bin_offset + bin_sz, 0:1], - src=count) - - return hbm_result - - -_DEFAULT_CONFIG = SimpleNamespace(block_size=CHUNK_SIZE) -_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (2048, 4096, 8192)] -_tuner = NkiAutotuner(histogram_kernel) if nki is not None else None + def histogram_kernel(values, num_bins, n_hi, n_lo, lo_bits, chunk_cols): + """Per-program partial histograms of the flat ``(N,)`` int32 ``values`` -> ``[lnc, num_bins]`` int32.""" + N = values.shape[0] + num_programs = nl.num_programs() + pid = nl.program_id(0) + out = nl.ndarray((num_programs, num_bins), dtype=nl.int32, buffer=nl.shared_hbm) + HP = max(32, n_hi) # one-hot width of the hi digit (32-aligned PSUM blocks) + G = PMAX // HP # value columns contracted per matmul + + # digit iotas [128, n_d] (bf16: exact small integers, enables the 2x Vector-engine mode) + iota_hi = nl.ndarray((PMAX, n_hi), dtype=nl.bfloat16, buffer=nl.sbuf) + _iota_bf16(iota_hi, n_hi) + iota_lo = nl.ndarray((PMAX, n_lo), dtype=nl.bfloat16, buffer=nl.sbuf) + _iota_bf16(iota_lo, n_lo) + # ping-pong one-hot buffers, zeroed once (the compare only writes the first n_hi columns) + oh_hi = [] + oh_lo = [] + for _ in range(2): + t = nl.ndarray((PMAX, chunk_cols, HP), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.memset(dst=t, value=0.0) + oh_hi.append(t) + oh_lo.append(nl.ndarray((PMAX, chunk_cols, n_lo), dtype=nl.bfloat16, buffer=nl.sbuf)) + acc = nl.ndarray((n_hi, n_lo), dtype=nl.int32, buffer=nl.sbuf) + nisa.memset(dst=acc, value=0) + psum = nl.ndarray((PMAX, G * n_lo), dtype=nl.float32, buffer=nl.psum) + + # This program's contiguous share of the values. + per_core = (N + num_programs - 1) // num_programs + lo = pid * per_core + hi = min(N, lo + per_core) + chunk = PMAX * chunk_cols + since_flush = 0 + group_start = True # next matmul overwrites (rather than accumulates into) PSUM + for ci in range(max(0, (hi - lo + chunk - 1) // chunk)): + start = lo + ci * chunk + count = min(chunk, hi - start) + q = count // PMAX # [128, q] values + r = count - q * PMAX # tail values (< 128): processed as a [1, r] chunk + buf = ci % 2 + if q > 0: + _chunk(values, PMAX, q, start, n_hi, n_lo, lo_bits, HP, G, + iota_hi, iota_lo, oh_hi[buf], oh_lo[buf], psum, group_start) + group_start = False + if r > 0: + _chunk(values, 1, r, start + q * PMAX, n_hi, n_lo, lo_bits, HP, G, + iota_hi, iota_lo, oh_hi[buf], oh_lo[buf], psum, group_start) + group_start = False + since_flush += count + if since_flush >= FLUSH_EVERY: + # fp32 PSUM stays exact below 2^24: flush into the int32 accumulator. + _flush(psum, acc, n_hi, n_lo, HP, G) + since_flush = 0 + group_start = True + if not group_start: + _flush(psum, acc, n_hi, n_lo, HP, G) + # counts[hi * n_lo + lo] <- acc[hi, lo], written to this program's row of ``out``. + nisa.dma_copy(dst=out.ap(pattern=[[n_lo, n_hi], [1, n_lo]], offset=pid * num_bins), src=acc) + return out + + def _iota_bf16(dst, n): + """``dst[p, d] = d`` (bf16) on every partition.""" + tmp = nl.ndarray((PMAX, n), dtype=nl.int32, buffer=nl.sbuf) + nisa.iota(dst=tmp, pattern=[[1, n]], offset=0, channel_multiplier=0) + nisa.tensor_copy(dst=dst, src=tmp) + + def _chunk(values, p, f, start, n_hi, n_lo, lo_bits, HP, G, iota_hi, iota_lo, oh_hi, oh_lo, psum, group_start): + """Accumulate the ``[p, f]`` chunk of values at flat offset ``start`` into ``psum``.""" + v = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) + nisa.dma_copy(dst=v, src=values.ap(pattern=[[f, p], [1, f]], offset=start)) + # digits hi = v >> lo_bits, lo = v & (n_lo - 1), converted to bf16 (exact: < 128) + hi_i = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) + lo_i = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) + nisa.tensor_scalar(dst=hi_i, data=v, op0=nl.right_shift, operand0=lo_bits) + nisa.tensor_scalar(dst=lo_i, data=v, op0=nl.bitwise_and, operand0=n_lo - 1) + hi_d = nl.ndarray((p, f), dtype=nl.bfloat16, buffer=nl.sbuf) + lo_d = nl.ndarray((p, f), dtype=nl.bfloat16, buffer=nl.sbuf) + nisa.tensor_copy(dst=hi_d, src=hi_i) + nisa.tensor_copy(dst=lo_d, src=lo_i) + # one-hots oh[p, f, d] = (digit[p, f] == d): one broadcast compare per digit + nisa.tensor_tensor(dst=oh_hi[0:p, 0:f, 0:n_hi], + data1=hi_d.ap(pattern=[[f, p], [1, f], [0, n_hi]]), + data2=iota_hi.ap(pattern=[[n_hi, p], [0, f], [1, n_hi]]), + op=nl.equal) + nisa.tensor_tensor(dst=oh_lo[0:p, 0:f, 0:n_lo], + data1=lo_d.ap(pattern=[[f, p], [1, f], [0, n_lo]]), + data2=iota_lo.ap(pattern=[[n_lo, p], [0, f], [1, n_lo]]), + op=nl.equal) + # G value columns per matmul: psum[g*HP + hi, g*n_lo + lo] += #{values in column f+g with digits (hi, lo)} + for f0 in range(0, f, G): + g = min(G, f - f0) + nisa.nc_matmul(dst=psum[0:g * HP, 0:g * n_lo], + stationary=oh_hi.ap(pattern=[[chunk_stride(oh_hi), p], [1, g * HP]], offset=f0 * HP), + moving=oh_lo.ap(pattern=[[chunk_stride(oh_lo), p], [1, g * n_lo]], offset=f0 * n_lo), + accumulate=(f0 > 0 or not group_start)) + + def chunk_stride(t): + """Elements per partition of a ``[128, F, D]`` SBUF tile (partition stride of its flat view).""" + return t.shape[1] * t.shape[2] + + def _flush(psum, acc, n_hi, n_lo, HP, G): + """``acc[hi, lo] += sum_g psum[g*HP + hi, g*n_lo + lo]`` (diagonal blocks), then PSUM is reused.""" + for g in range(G): + tmp = nl.ndarray((n_hi, n_lo), dtype=nl.int32, buffer=nl.sbuf) + nisa.tensor_copy(dst=tmp, src=psum[g * HP:g * HP + n_hi, g * n_lo:(g + 1) * n_lo]) + nisa.tensor_tensor(dst=acc, data1=acc, data2=tmp, op=nl.add) + + +# Values per partition per chunk: Triton's partial_BLOCK_SIZE (default 1024; search 1024/2048). +_DEFAULT_CONFIG = SimpleNamespace(block_size=1024) +_SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (1024, 2048)] +_kernel = histogram_kernel[_lnc_degree()] if nki is not None else None +_tuner = NkiAutotuner(_kernel) if nki is not None else None _last_autotune_config: dict = {} def run(input: torch.Tensor, N: int, num_bins: int, block_size: int = 1024, autotune: bool = False, **kwargs) -> torch.Tensor: - values_2d = input.reshape(1, -1).to(torch.int32) + assert input.ndim == 1 and input.shape[0] == N and input.dtype == torch.int32 + n_hi, n_lo, lo_bits = _digits(num_bins) if autotune: + # The SBUF cap can map several block sizes onto the same chunk width: keep one config each. + seen = set() + space = [c for c in _SEARCH_SPACE + if not (_chunk_cols(c.block_size, n_hi, n_lo) in seen or seen.add(_chunk_cols(c.block_size, n_hi, n_lo)))] cfg = _tuner.tune_or_cached( - shape_key=(tuple(values_2d.shape), str(values_2d.dtype)), - search_space=_SEARCH_SPACE, - args_fn=lambda cfg: (values_2d, num_bins, cfg.block_size), + shape_key=(N, num_bins), + search_space=space, + args_fn=lambda cfg: (input, num_bins, n_hi, n_lo, lo_bits, _chunk_cols(cfg.block_size, n_hi, n_lo)), ) _last_autotune_config.clear() _last_autotune_config.update(vars(cfg)) else: cfg = _DEFAULT_CONFIG - hist = histogram_kernel(values_2d, num_bins, cfg.block_size) - return hist.reshape(-1) + partials = _kernel(input, num_bins, n_hi, n_lo, lo_bits, _chunk_cols(cfg.block_size, n_hi, n_lo)) + # Stage 2 (Triton's reduce kernel): sum the per-program partial histograms. + return partials.sum(dim=0, dtype=torch.int32) def get_last_config() -> dict | None: diff --git a/benchmarks/operators/histogramming/impl_torch.py b/benchmarks/operators/histogramming/impl_torch.py index eaa48c72..38595a30 100644 --- a/benchmarks/operators/histogramming/impl_torch.py +++ b/benchmarks/operators/histogramming/impl_torch.py @@ -1,8 +1,41 @@ +import math + import torch +def _histogram_xla(x: torch.Tensor, num_bins: int, chunk: int = 1 << 21) -> torch.Tensor: + """On-device histogram for Neuron/XLA. + + ``torch.bincount`` / ``torch.histc`` have no XLA lowering (they silently run on + the host), ``scatter_add_`` / ``index_add_`` lose updates on duplicate indices, + and ``sort`` is unsupported on trn2 -- so the bin id is split into two digits + ``v = hi * n_lo + lo`` and the counts are the matmul of the two one-hot codes + (``counts[hi, lo] = sum_i OH_hi[i, hi] * OH_lo[i, lo]``, exact in fp32 + accumulation), computed in chunks so the one-hot tensors stay small. + """ + n_lo = 1 + while n_lo * 2 <= math.isqrt(num_bins) and num_bins % (n_lo * 2) == 0: + n_lo *= 2 + n_hi = num_bins // n_lo + lo_bits = n_lo.bit_length() - 1 + ar_hi = torch.arange(n_hi, device=x.device, dtype=torch.int32).view(1, n_hi) + ar_lo = torch.arange(n_lo, device=x.device, dtype=torch.int32).view(1, n_lo) + acc = torch.zeros(n_hi, n_lo, dtype=torch.float32, device=x.device) + for i in range(0, x.numel(), chunk): + xc = x[i:i + chunk] + # values outside [0, num_bins) match no one-hot column and are dropped + hi = (xc >> lo_bits).view(-1, 1) + lo = (xc & (n_lo - 1)).view(-1, 1) + oh_hi = (hi == ar_hi).to(torch.bfloat16) + oh_lo = (lo == ar_lo).to(torch.bfloat16) + acc = acc + (oh_hi.t() @ oh_lo).to(torch.float32) + return acc.reshape(-1).to(torch.int32) + + def run(input: torch.Tensor, N: int, num_bins: int, **kwargs): assert input.ndim == 1 assert input.shape[0] == N assert input.dtype == torch.int32 + if input.device.type == "xla": + return _histogram_xla(input, num_bins) return torch.bincount(input.to(torch.int64), minlength=num_bins).to(torch.int32) diff --git a/results/csv/histogramming_autotune.csv b/results/csv/histogramming_autotune.csv index 346448b9..47fe78b4 100644 --- a/results/csv/histogramming_autotune.csv +++ b/results/csv/histogramming_autotune.csv @@ -1,21 +1,21 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -"N=262144, num_bins=64",int32,0.0258,0.0110,0.0132,2.34,1.95,1.1969,0.0077,3.35 -"N=262144, num_bins=256",int32,0.0250,0.0109,0.0134,2.29,1.86,1.2290,0.0080,3.11 -"N=262144, num_bins=1024",int32,0.0299,0.0119,0.0139,2.52,2.15,1.1712,0.0091,3.28 -"N=262144, num_bins=4096",int32,0.0337,0.0115,0.0135,2.93,2.50,1.1732,0.0088,3.82 -"N=1048576, num_bins=64",int32,0.0302,0.0215,0.0337,1.40,0.90,1.5644,0.0091,3.31 -"N=1048576, num_bins=256",int32,0.0314,0.0226,0.0342,1.39,0.92,1.5126,0.0092,3.43 -"N=1048576, num_bins=1024",int32,0.0368,0.0221,0.0346,1.67,1.06,1.5654,0.0091,4.03 -"N=1048576, num_bins=4096",int32,0.0476,0.0214,0.0312,2.23,1.53,1.4590,0.0100,4.78 -"N=4194304, num_bins=64",int32,0.0609,0.0660,0.1185,0.92,0.51,1.7949,0.0134,4.53 -"N=4194304, num_bins=256",int32,0.0551,0.0656,0.1191,0.84,0.46,1.8162,0.0133,4.13 -"N=4194304, num_bins=1024",int32,0.0611,0.0658,0.1187,0.93,0.51,1.8057,0.0138,4.41 -"N=4194304, num_bins=4096",int32,0.0824,0.0577,0.0961,1.43,0.86,1.6652,0.0143,5.77 -"N=16777216, num_bins=64",int32,0.2106,0.2418,0.4607,0.87,0.46,1.9053,0.0316,6.66 -"N=16777216, num_bins=256",int32,0.1847,0.2411,0.4579,0.77,0.40,1.8993,0.0322,5.74 -"N=16777216, num_bins=1024",int32,0.1780,0.2415,0.4553,0.74,0.39,1.8854,0.0327,5.45 -"N=16777216, num_bins=4096",int32,0.2247,0.2077,0.3670,1.08,0.61,1.7667,0.0334,6.72 -"N=67108864, num_bins=64",int32,0.7712,0.9464,1.8305,0.81,0.42,1.9342,0.0954,8.09 -"N=67108864, num_bins=256",int32,0.6610,0.9612,1.8510,0.69,0.36,1.9257,0.0955,6.92 -"N=67108864, num_bins=1024",int32,0.6350,0.9350,1.7902,0.68,0.35,1.9147,0.0955,6.65 -"N=67108864, num_bins=4096",int32,0.6246,0.8009,1.4265,0.78,0.44,1.7810,0.0961,6.50 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang,torch_nki_ms,nki_ms,speedup_nki +"N=262144, num_bins=64",int32,0.0258,0.0110,0.0132,2.34,1.95,1.1969,0.0077,3.35,0.2542,0.0564,4.51 +"N=262144, num_bins=256",int32,0.0250,0.0109,0.0134,2.29,1.86,1.2290,0.0080,3.11,0.3846,0.0839,4.58 +"N=262144, num_bins=1024",int32,0.0299,0.0119,0.0139,2.52,2.15,1.1712,0.0091,3.28,0.6149,0.1569,3.92 +"N=262144, num_bins=4096",int32,0.0337,0.0115,0.0135,2.93,2.50,1.1732,0.0088,3.82,1.2115,0.1722,7.04 +"N=1048576, num_bins=64",int32,0.0302,0.0215,0.0337,1.40,0.90,1.5644,0.0091,3.31,0.8258,0.1184,6.98 +"N=1048576, num_bins=256",int32,0.0314,0.0226,0.0342,1.39,0.92,1.5126,0.0092,3.43,1.4596,0.1982,7.36 +"N=1048576, num_bins=1024",int32,0.0368,0.0221,0.0346,1.67,1.06,1.5654,0.0091,4.03,2.5813,0.3739,6.90 +"N=1048576, num_bins=4096",int32,0.0476,0.0214,0.0312,2.23,1.53,1.4590,0.0100,4.78,4.9714,0.5989,8.30 +"N=4194304, num_bins=64",int32,0.0609,0.0660,0.1185,0.92,0.51,1.7949,0.0134,4.53,3.7932,0.3645,10.41 +"N=4194304, num_bins=256",int32,0.0551,0.0656,0.1191,0.84,0.46,1.8162,0.0133,4.13,5.8816,0.6516,9.03 +"N=4194304, num_bins=1024",int32,0.0611,0.0658,0.1187,0.93,0.51,1.8057,0.0138,4.41,10.3676,1.2422,8.35 +"N=4194304, num_bins=4096",int32,0.0824,0.0577,0.0961,1.43,0.86,1.6652,0.0143,5.77,19.8474,2.3084,8.60 +"N=16777216, num_bins=64",int32,0.2106,0.2418,0.4607,0.87,0.46,1.9053,0.0316,6.66,15.9280,1.3650,11.67 +"N=16777216, num_bins=256",int32,0.1847,0.2411,0.4579,0.77,0.40,1.8993,0.0322,5.74,24.5002,2.4854,9.86 +"N=16777216, num_bins=1024",int32,0.1780,0.2415,0.4553,0.74,0.39,1.8854,0.0327,5.45,43.3601,4.7177,9.19 +"N=16777216, num_bins=4096",int32,0.2247,0.2077,0.3670,1.08,0.61,1.7667,0.0334,6.72,82.2719,9.1437,9.00 +"N=67108864, num_bins=64",int32,0.7712,0.9464,1.8305,0.81,0.42,1.9342,0.0954,8.09,62.3042,5.3560,11.63 +"N=67108864, num_bins=256",int32,0.6610,0.9612,1.8510,0.69,0.36,1.9257,0.0955,6.92,100.5773,9.8051,10.26 +"N=67108864, num_bins=1024",int32,0.6350,0.9350,1.7902,0.68,0.35,1.9147,0.0955,6.65,178.3460,18.6167,9.58 +"N=67108864, num_bins=4096",int32,0.6246,0.8009,1.4265,0.78,0.44,1.7810,0.0961,6.50,322.3469,36.4823,8.84 diff --git a/results/csv/histogramming_default.csv b/results/csv/histogramming_default.csv index fb2343ba..a31814c0 100644 --- a/results/csv/histogramming_default.csv +++ b/results/csv/histogramming_default.csv @@ -1,21 +1,21 @@ -params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang -"N=262144, num_bins=64",int32,0.0258,0.0285,0.0152,0.90,1.70,0.5325,0.0087,2.96 -"N=262144, num_bins=256",int32,0.0249,0.0330,0.0149,0.76,1.67,0.4531,0.0087,2.87 -"N=262144, num_bins=1024",int32,0.0298,0.0354,0.0151,0.84,1.97,0.4271,0.0097,3.06 -"N=262144, num_bins=4096",int32,0.0337,0.0317,0.0156,1.06,2.16,0.4912,0.0095,3.56 -"N=1048576, num_bins=64",int32,0.0302,0.0495,0.0366,0.61,0.82,0.7394,0.0099,3.06 -"N=1048576, num_bins=256",int32,0.0314,0.0535,0.0369,0.59,0.85,0.6890,0.0103,3.06 -"N=1048576, num_bins=1024",int32,0.0368,0.0565,0.0371,0.65,0.99,0.6567,0.0100,3.69 -"N=1048576, num_bins=4096",int32,0.0476,0.0527,0.0381,0.90,1.25,0.7238,0.0110,4.34 -"N=4194304, num_bins=64",int32,0.0609,0.1328,0.1242,0.46,0.49,0.9359,0.0165,3.69 -"N=4194304, num_bins=256",int32,0.0550,0.1378,0.1258,0.40,0.44,0.9124,0.0175,3.15 -"N=4194304, num_bins=1024",int32,0.0613,0.1383,0.1234,0.44,0.50,0.8925,0.0175,3.51 -"N=4194304, num_bins=4096",int32,0.0819,0.1358,0.1259,0.60,0.65,0.9267,0.0188,4.35 -"N=16777216, num_bins=64",int32,0.2108,0.4790,0.5036,0.44,0.42,1.0513,0.0484,4.35 -"N=16777216, num_bins=256",int32,0.1846,0.4776,0.4992,0.39,0.37,1.0453,0.0489,3.78 -"N=16777216, num_bins=1024",int32,0.1785,0.4753,0.4940,0.38,0.36,1.0392,0.0485,3.68 -"N=16777216, num_bins=4096",int32,0.2245,0.4760,0.4957,0.47,0.45,1.0414,0.0497,4.51 -"N=67108864, num_bins=64",int32,0.7711,1.8367,1.9788,0.42,0.39,1.0773,0.1581,4.88 -"N=67108864, num_bins=256",int32,0.6609,1.8320,1.9728,0.36,0.33,1.0768,0.1583,4.17 -"N=67108864, num_bins=1024",int32,0.6337,1.8162,1.9526,0.35,0.32,1.0751,0.1588,3.99 -"N=67108864, num_bins=4096",int32,0.6248,1.8365,1.9612,0.34,0.32,1.0679,0.1596,3.91 +params,dtype,torch_ms,triton_ms,cutile_ms,speedup_triton,speedup_cutile,triton_vs_cutile,tilelang_ms,speedup_tilelang,torch_nki_ms,nki_ms,speedup_nki +"N=262144, num_bins=64",int32,0.0258,0.0285,0.0152,0.90,1.70,0.5325,0.0087,2.96,0.2539,0.0564,4.50 +"N=262144, num_bins=256",int32,0.0249,0.0330,0.0149,0.76,1.67,0.4531,0.0087,2.87,0.3849,0.0840,4.58 +"N=262144, num_bins=1024",int32,0.0298,0.0354,0.0151,0.84,1.97,0.4271,0.0097,3.06,0.6151,0.1570,3.92 +"N=262144, num_bins=4096",int32,0.0337,0.0317,0.0156,1.06,2.16,0.4912,0.0095,3.56,1.2115,0.1723,7.03 +"N=1048576, num_bins=64",int32,0.0302,0.0495,0.0366,0.61,0.82,0.7394,0.0099,3.06,0.8256,0.1184,6.98 +"N=1048576, num_bins=256",int32,0.0314,0.0535,0.0369,0.59,0.85,0.6890,0.0103,3.06,1.4592,0.1982,7.36 +"N=1048576, num_bins=1024",int32,0.0368,0.0565,0.0371,0.65,0.99,0.6567,0.0100,3.69,2.5810,0.3739,6.90 +"N=1048576, num_bins=4096",int32,0.0476,0.0527,0.0381,0.90,1.25,0.7238,0.0110,4.34,4.9717,0.5992,8.30 +"N=4194304, num_bins=64",int32,0.0609,0.1328,0.1242,0.46,0.49,0.9359,0.0165,3.69,3.7934,0.3645,10.41 +"N=4194304, num_bins=256",int32,0.0550,0.1378,0.1258,0.40,0.44,0.9124,0.0175,3.15,5.8810,0.6515,9.03 +"N=4194304, num_bins=1024",int32,0.0613,0.1383,0.1234,0.44,0.50,0.8925,0.0175,3.51,10.3664,1.2422,8.34 +"N=4194304, num_bins=4096",int32,0.0819,0.1358,0.1259,0.60,0.65,0.9267,0.0188,4.35,19.8482,2.3085,8.60 +"N=16777216, num_bins=64",int32,0.2108,0.4790,0.5036,0.44,0.42,1.0513,0.0484,4.35,15.9309,1.3651,11.67 +"N=16777216, num_bins=256",int32,0.1846,0.4776,0.4992,0.39,0.37,1.0453,0.0489,3.78,24.5020,2.4854,9.86 +"N=16777216, num_bins=1024",int32,0.1785,0.4753,0.4940,0.38,0.36,1.0392,0.0485,3.68,43.3609,4.7177,9.19 +"N=16777216, num_bins=4096",int32,0.2245,0.4760,0.4957,0.47,0.45,1.0414,0.0497,4.51,82.2721,9.1443,9.00 +"N=67108864, num_bins=64",int32,0.7711,1.8367,1.9788,0.42,0.39,1.0773,0.1581,4.88,62.3085,5.3561,11.63 +"N=67108864, num_bins=256",int32,0.6609,1.8320,1.9728,0.36,0.33,1.0768,0.1583,4.17,100.5742,9.8051,10.26 +"N=67108864, num_bins=1024",int32,0.6337,1.8162,1.9526,0.35,0.32,1.0751,0.1588,3.99,178.3415,18.6166,9.58 +"N=67108864, num_bins=4096",int32,0.6248,1.8365,1.9612,0.34,0.32,1.0679,0.1596,3.91,322.3446,36.4824,8.84 diff --git a/results/logs/autotune_logs/histogramming_autotune.json b/results/logs/autotune_logs/histogramming_autotune.json index 0751e29a..b5df44c4 100644 --- a/results/logs/autotune_logs/histogramming_autotune.json +++ b/results/logs/autotune_logs/histogramming_autotune.json @@ -20,6 +20,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -43,6 +46,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -66,6 +72,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -89,6 +98,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -112,6 +124,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -135,6 +150,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -158,6 +176,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -181,6 +202,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -204,6 +228,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -227,6 +254,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -250,6 +280,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -273,6 +306,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -296,6 +332,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -319,6 +358,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -342,6 +384,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -365,6 +410,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -388,6 +436,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 4 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -411,6 +462,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -434,6 +488,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } }, { @@ -457,6 +514,9 @@ "reduce_block_rows": 128, "reduce_block_bins": 64, "reduce_occupancy": 8 + }, + "nki_autotune_cfg": { + "block_size": 1024 } } -] \ No newline at end of file +] From 28c6df82c394d63dd0ba6be158b591afa99f0194 Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Mon, 21 Sep 2026 20:38:17 +0000 Subject: [PATCH 6/7] nki(histogramming): migrate to the tilebench package layout Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- tilebench/benchmarks/operators/histogramming/impl_nki.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tilebench/benchmarks/operators/histogramming/impl_nki.py b/tilebench/benchmarks/operators/histogramming/impl_nki.py index 42598950..07fd1050 100644 --- a/tilebench/benchmarks/operators/histogramming/impl_nki.py +++ b/tilebench/benchmarks/operators/histogramming/impl_nki.py @@ -42,7 +42,7 @@ import torch -from core.nki_autotune import NkiAutotuner +from tilebench.core.nki_autotune import NkiAutotuner try: import nki From 96924c23c8fe381836ea9e561bf2ef047cc16dee Mon Sep 17 00:00:00 2001 From: Bowen Cui Date: Tue, 22 Sep 2026 04:57:45 +0000 Subject: [PATCH 7/7] nki(histogramming): drop docstrings and inline comments Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01ScXYNjrrKGgDUVNHxv7HJt --- .../operators/histogramming/impl_nki.py | 72 ++----------------- .../operators/histogramming/impl_torch.py | 10 --- 2 files changed, 7 insertions(+), 75 deletions(-) diff --git a/tilebench/benchmarks/operators/histogramming/impl_nki.py b/tilebench/benchmarks/operators/histogramming/impl_nki.py index 07fd1050..a35746af 100644 --- a/tilebench/benchmarks/operators/histogramming/impl_nki.py +++ b/tilebench/benchmarks/operators/histogramming/impl_nki.py @@ -1,38 +1,3 @@ -"""NKI (AWS Trainium) implementation of histogramming: ``hist[b] = #{i : x[i] == b}`` for ``b < num_bins``. - -Design (mirrors the Triton / cuTile two-stage structure): - -* Triton shards the input across workers, each worker accumulates a private - partial histogram with ``atomic_add`` (one read per element), and a second - kernel reduces the partials. NeuronCores have no atomic scatter-add, so the - per-element "+1 into bin ``v``" is expressed as a one-hot outer product that - the Tensor engine accumulates: - - a bin id is split into two digits ``v = hi * n_lo + lo`` - (``n_hi * n_lo = num_bins``, e.g. 32 x 32 for 1024 bins); - - for a chunk of ``[128, F]`` values, one-hot codes ``OH_hi[128, F, HP]`` and - ``OH_lo[128, F, n_lo]`` (bf16, exact 0/1) are built by one broadcast compare - each on the Vector engine (``HP = max(32, n_hi)`` columns per value, the - unused ones zero, so that the ``G = 128 / HP`` results below land on - 32-aligned partitions); - - ``G`` consecutive value columns are contracted by a single ``nc_matmul``: - ``stationary = OH_hi[:, f:f+G, :]`` (``[128, 128]``), ``moving = OH_lo[:, f:f+G, :]`` - (``[128, G * n_lo]``), whose ``[128, G * n_lo]`` PSUM result holds the wanted - ``[n_hi, n_lo]`` count blocks on its block diagonal (the off-diagonal blocks - pair values of different columns and are ignored). One instruction thus - consumes ``128 * G`` values; the fp32 PSUM is flushed into an exact int32 - SBUF accumulator every ``2^20`` values, so counts stay exact for any ``N``. - - each program instance handles its own contiguous share of the input and - writes its partial ``[num_bins]`` histogram; ``run`` sums the per-program - partials (the same "reduce the partials" stage as Triton's second kernel). - Every element is read exactly once (``O(N)`` DMA); compute is - ``O(N * (n_hi + n_lo))`` Vector-engine work + ``N / (128 * G)`` matmuls. - ``BLOCK`` (values per partition per chunk) defaults to Triton's - ``partial_BLOCK_SIZE`` (1024); search 1024/2048. It is capped so that the - one-hot tiles of a chunk fit in SBUF. -* Requires ``num_bins`` to factor as ``n_hi * n_lo`` with ``n_lo`` a power of two - and ``n_hi, n_lo <= 128`` (all benchmark cases are powers of two); values are - assumed to lie in ``[0, num_bins)`` as in the Triton kernel's valid range. -""" import functools import math import os @@ -53,17 +18,12 @@ nki = None PMAX = 128 -FLUSH_EVERY = 1 << 20 # values accumulated in fp32 PSUM before flushing to the int32 accumulator -SBUF_ONEHOT_BYTES = 32 * 1024 # per-partition budget for one chunk's pair of one-hot tiles +FLUSH_EVERY = 1 << 20 +SBUF_ONEHOT_BYTES = 32 * 1024 @functools.lru_cache(maxsize=1) def _lnc_degree() -> int: - """Logical-NeuronCore degree the kernel is launched with (``kernel[lnc]``). - - Must match the LNC the XLA module is compiled for: launching a ``kernel[2]`` - into an ``--lnc 1`` module silently computes only core 0's half. - """ explicit = os.environ.get("NEURON_LOGICAL_NC_CONFIG", "") if explicit.strip().isdigit(): return int(explicit.strip()) @@ -81,7 +41,6 @@ def _lnc_degree() -> int: def _digits(num_bins: int): - """Split ``num_bins`` into ``(n_hi, n_lo, lo_bits)`` with ``n_lo = 2**lo_bits`` closest to ``sqrt(num_bins)``.""" lo_bits = 0 while (1 << (lo_bits + 1)) <= int(math.isqrt(num_bins)) and num_bins % (1 << (lo_bits + 1)) == 0: lo_bits += 1 @@ -93,7 +52,6 @@ def _digits(num_bins: int): def _chunk_cols(block_size: int, n_hi: int, n_lo: int) -> int: - """Values per partition per chunk: ``block_size`` capped by the one-hot SBUF budget.""" hp = max(32, n_hi) f = min(block_size, SBUF_ONEHOT_BYTES // (2 * (hp + n_lo))) return max(PMAX // hp, f - f % (PMAX // hp)) @@ -102,20 +60,17 @@ def _chunk_cols(block_size: int, n_hi: int, n_lo: int) -> int: if nki is not None: @nki.jit def histogram_kernel(values, num_bins, n_hi, n_lo, lo_bits, chunk_cols): - """Per-program partial histograms of the flat ``(N,)`` int32 ``values`` -> ``[lnc, num_bins]`` int32.""" N = values.shape[0] num_programs = nl.num_programs() pid = nl.program_id(0) out = nl.ndarray((num_programs, num_bins), dtype=nl.int32, buffer=nl.shared_hbm) - HP = max(32, n_hi) # one-hot width of the hi digit (32-aligned PSUM blocks) - G = PMAX // HP # value columns contracted per matmul + HP = max(32, n_hi) + G = PMAX // HP - # digit iotas [128, n_d] (bf16: exact small integers, enables the 2x Vector-engine mode) iota_hi = nl.ndarray((PMAX, n_hi), dtype=nl.bfloat16, buffer=nl.sbuf) _iota_bf16(iota_hi, n_hi) iota_lo = nl.ndarray((PMAX, n_lo), dtype=nl.bfloat16, buffer=nl.sbuf) _iota_bf16(iota_lo, n_lo) - # ping-pong one-hot buffers, zeroed once (the compare only writes the first n_hi columns) oh_hi = [] oh_lo = [] for _ in range(2): @@ -127,18 +82,17 @@ def histogram_kernel(values, num_bins, n_hi, n_lo, lo_bits, chunk_cols): nisa.memset(dst=acc, value=0) psum = nl.ndarray((PMAX, G * n_lo), dtype=nl.float32, buffer=nl.psum) - # This program's contiguous share of the values. per_core = (N + num_programs - 1) // num_programs lo = pid * per_core hi = min(N, lo + per_core) chunk = PMAX * chunk_cols since_flush = 0 - group_start = True # next matmul overwrites (rather than accumulates into) PSUM + group_start = True for ci in range(max(0, (hi - lo + chunk - 1) // chunk)): start = lo + ci * chunk count = min(chunk, hi - start) - q = count // PMAX # [128, q] values - r = count - q * PMAX # tail values (< 128): processed as a [1, r] chunk + q = count // PMAX + r = count - q * PMAX buf = ci % 2 if q > 0: _chunk(values, PMAX, q, start, n_hi, n_lo, lo_bits, HP, G, @@ -150,27 +104,22 @@ def histogram_kernel(values, num_bins, n_hi, n_lo, lo_bits, chunk_cols): group_start = False since_flush += count if since_flush >= FLUSH_EVERY: - # fp32 PSUM stays exact below 2^24: flush into the int32 accumulator. _flush(psum, acc, n_hi, n_lo, HP, G) since_flush = 0 group_start = True if not group_start: _flush(psum, acc, n_hi, n_lo, HP, G) - # counts[hi * n_lo + lo] <- acc[hi, lo], written to this program's row of ``out``. nisa.dma_copy(dst=out.ap(pattern=[[n_lo, n_hi], [1, n_lo]], offset=pid * num_bins), src=acc) return out def _iota_bf16(dst, n): - """``dst[p, d] = d`` (bf16) on every partition.""" tmp = nl.ndarray((PMAX, n), dtype=nl.int32, buffer=nl.sbuf) nisa.iota(dst=tmp, pattern=[[1, n]], offset=0, channel_multiplier=0) nisa.tensor_copy(dst=dst, src=tmp) def _chunk(values, p, f, start, n_hi, n_lo, lo_bits, HP, G, iota_hi, iota_lo, oh_hi, oh_lo, psum, group_start): - """Accumulate the ``[p, f]`` chunk of values at flat offset ``start`` into ``psum``.""" v = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) nisa.dma_copy(dst=v, src=values.ap(pattern=[[f, p], [1, f]], offset=start)) - # digits hi = v >> lo_bits, lo = v & (n_lo - 1), converted to bf16 (exact: < 128) hi_i = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) lo_i = nl.ndarray((p, f), dtype=nl.int32, buffer=nl.sbuf) nisa.tensor_scalar(dst=hi_i, data=v, op0=nl.right_shift, operand0=lo_bits) @@ -179,7 +128,6 @@ def _chunk(values, p, f, start, n_hi, n_lo, lo_bits, HP, G, iota_hi, iota_lo, oh lo_d = nl.ndarray((p, f), dtype=nl.bfloat16, buffer=nl.sbuf) nisa.tensor_copy(dst=hi_d, src=hi_i) nisa.tensor_copy(dst=lo_d, src=lo_i) - # one-hots oh[p, f, d] = (digit[p, f] == d): one broadcast compare per digit nisa.tensor_tensor(dst=oh_hi[0:p, 0:f, 0:n_hi], data1=hi_d.ap(pattern=[[f, p], [1, f], [0, n_hi]]), data2=iota_hi.ap(pattern=[[n_hi, p], [0, f], [1, n_hi]]), @@ -188,7 +136,6 @@ def _chunk(values, p, f, start, n_hi, n_lo, lo_bits, HP, G, iota_hi, iota_lo, oh data1=lo_d.ap(pattern=[[f, p], [1, f], [0, n_lo]]), data2=iota_lo.ap(pattern=[[n_lo, p], [0, f], [1, n_lo]]), op=nl.equal) - # G value columns per matmul: psum[g*HP + hi, g*n_lo + lo] += #{values in column f+g with digits (hi, lo)} for f0 in range(0, f, G): g = min(G, f - f0) nisa.nc_matmul(dst=psum[0:g * HP, 0:g * n_lo], @@ -197,18 +144,15 @@ def _chunk(values, p, f, start, n_hi, n_lo, lo_bits, HP, G, iota_hi, iota_lo, oh accumulate=(f0 > 0 or not group_start)) def chunk_stride(t): - """Elements per partition of a ``[128, F, D]`` SBUF tile (partition stride of its flat view).""" return t.shape[1] * t.shape[2] def _flush(psum, acc, n_hi, n_lo, HP, G): - """``acc[hi, lo] += sum_g psum[g*HP + hi, g*n_lo + lo]`` (diagonal blocks), then PSUM is reused.""" for g in range(G): tmp = nl.ndarray((n_hi, n_lo), dtype=nl.int32, buffer=nl.sbuf) nisa.tensor_copy(dst=tmp, src=psum[g * HP:g * HP + n_hi, g * n_lo:(g + 1) * n_lo]) nisa.tensor_tensor(dst=acc, data1=acc, data2=tmp, op=nl.add) -# Values per partition per chunk: Triton's partial_BLOCK_SIZE (default 1024; search 1024/2048). _DEFAULT_CONFIG = SimpleNamespace(block_size=1024) _SEARCH_SPACE = [SimpleNamespace(block_size=b) for b in (1024, 2048)] _kernel = histogram_kernel[_lnc_degree()] if nki is not None else None @@ -221,7 +165,6 @@ def run(input: torch.Tensor, N: int, num_bins: int, block_size: int = 1024, assert input.ndim == 1 and input.shape[0] == N and input.dtype == torch.int32 n_hi, n_lo, lo_bits = _digits(num_bins) if autotune: - # The SBUF cap can map several block sizes onto the same chunk width: keep one config each. seen = set() space = [c for c in _SEARCH_SPACE if not (_chunk_cols(c.block_size, n_hi, n_lo) in seen or seen.add(_chunk_cols(c.block_size, n_hi, n_lo)))] @@ -235,7 +178,6 @@ def run(input: torch.Tensor, N: int, num_bins: int, block_size: int = 1024, else: cfg = _DEFAULT_CONFIG partials = _kernel(input, num_bins, n_hi, n_lo, lo_bits, _chunk_cols(cfg.block_size, n_hi, n_lo)) - # Stage 2 (Triton's reduce kernel): sum the per-program partial histograms. return partials.sum(dim=0, dtype=torch.int32) diff --git a/tilebench/benchmarks/operators/histogramming/impl_torch.py b/tilebench/benchmarks/operators/histogramming/impl_torch.py index 38595a30..29c4cb0d 100644 --- a/tilebench/benchmarks/operators/histogramming/impl_torch.py +++ b/tilebench/benchmarks/operators/histogramming/impl_torch.py @@ -4,15 +4,6 @@ def _histogram_xla(x: torch.Tensor, num_bins: int, chunk: int = 1 << 21) -> torch.Tensor: - """On-device histogram for Neuron/XLA. - - ``torch.bincount`` / ``torch.histc`` have no XLA lowering (they silently run on - the host), ``scatter_add_`` / ``index_add_`` lose updates on duplicate indices, - and ``sort`` is unsupported on trn2 -- so the bin id is split into two digits - ``v = hi * n_lo + lo`` and the counts are the matmul of the two one-hot codes - (``counts[hi, lo] = sum_i OH_hi[i, hi] * OH_lo[i, lo]``, exact in fp32 - accumulation), computed in chunks so the one-hot tensors stay small. - """ n_lo = 1 while n_lo * 2 <= math.isqrt(num_bins) and num_bins % (n_lo * 2) == 0: n_lo *= 2 @@ -23,7 +14,6 @@ def _histogram_xla(x: torch.Tensor, num_bins: int, chunk: int = 1 << 21) -> torc acc = torch.zeros(n_hi, n_lo, dtype=torch.float32, device=x.device) for i in range(0, x.numel(), chunk): xc = x[i:i + chunk] - # values outside [0, num_bins) match no one-hot column and are dropped hi = (xc >> lo_bits).view(-1, 1) lo = (xc & (n_lo - 1)).view(-1, 1) oh_hi = (hi == ar_hi).to(torch.bfloat16)