Skip to content

nki(histogramming): NKI (Trainium) implementation #280

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bowencui123 wants to merge 9 commits into
base: main
Choose a base branch
from
Open

nki(histogramming): NKI (Trainium) implementation #280

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
37a39f9
nki(histogramming): NKI (Trainium) implementation
bowencui123 Aug 29, 2026
47d2823
nki(histogramming): NkiAutotuner wiring (`block_size`)
bowencui123 Aug 29, 2026
e3b41f1
nki(histogramming): add NKI benchmark results (trn2.3xlarge, LNC2)
Sep 8, 2026
b578857
nki(histogramming): add NKI autotune benchmark results (trn2.3xlarge,…
Sep 9, 2026
f24eded
Merge remote-tracking branch 'origin/main' into bowen/nki/histogramming
bowencui123 Sep 15, 2026
f67c4da
nki(histogramming): LNC2 digit one-hot TensorE kernel (block-diagonal…
bowencui123 Sep 16, 2026
56abdf8
Merge remote-tracking branch 'origin/main' into bowen/nki/histogramming
bowencui123 Sep 21, 2026
28c6df8
nki(histogramming): migrate to the tilebench package layout
bowencui123 Sep 21, 2026
96924c2
nki(histogramming): drop docstrings and inline comments
bowencui123 Sep 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 21 additions & 21 deletions results/B200/csv/histogramming_autotune.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
42 changes: 21 additions & 21 deletions results/B200/csv/histogramming_default.csv
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
185 changes: 185 additions & 0 deletions tilebench/benchmarks/operators/histogramming/impl_nki.py
View file
Open in desktop
Comment thread
bowencui123 marked this conversation as resolved.
Original file line number Diff line number Diff line change
@@ -0,0 +1,185 @@
import functools
import math
import os
import re
import subprocess
from types import SimpleNamespace

import torch

from tilebench.core.nki_autotune import NkiAutotuner

try:
import nki
import nki.isa as nisa
import nki.language as nl
PMAX = nl.tile_size.pmax
except ImportError:
nki = None
PMAX = 128

FLUSH_EVERY = 1 << 20
SBUF_ONEHOT_BYTES = 32 * 1024


@functools.lru_cache(maxsize=1)
def _lnc_degree() -> int:
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):
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:
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, n_hi, n_lo, lo_bits, chunk_cols):
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)
G = PMAX // HP

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)
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)

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
for ci in range(max(0, (hi - lo + chunk - 1) // chunk)):
start = lo + ci * chunk
count = min(chunk, hi - start)
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,
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:
_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)
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):
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):
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))
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)
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)
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):
return t.shape[1] * t.shape[2]

def _flush(psum, acc, n_hi, n_lo, HP, G):
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)


_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:
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:
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=(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
partials = _kernel(input, num_bins, n_hi, n_lo, lo_bits, _chunk_cols(cfg.block_size, n_hi, n_lo))
return partials.sum(dim=0, dtype=torch.int32)


def get_last_config() -> dict | None:
return dict(_last_autotune_config) or None
24 changes: 23 additions & 1 deletion tilebench/benchmarks/operators/histogramming/impl_torch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,31 @@
import math

import torch


def _histogram_xla(x: torch.Tensor, num_bins: int, chunk: int = 1 << 21) -> torch.Tensor:
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]
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.is_cuda
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)