diff --git a/results/B200/csv/histogramming_autotune.csv b/results/B200/csv/histogramming_autotune.csv index 346448b9..47fe78b4 100644 --- a/results/B200/csv/histogramming_autotune.csv +++ b/results/B200/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/B200/csv/histogramming_default.csv b/results/B200/csv/histogramming_default.csv index fb2343ba..a31814c0 100644 --- a/results/B200/csv/histogramming_default.csv +++ b/results/B200/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/tilebench/benchmarks/operators/histogramming/impl_nki.py b/tilebench/benchmarks/operators/histogramming/impl_nki.py new file mode 100644 index 00000000..a35746af --- /dev/null +++ b/tilebench/benchmarks/operators/histogramming/impl_nki.py @@ -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 diff --git a/tilebench/benchmarks/operators/histogramming/impl_torch.py b/tilebench/benchmarks/operators/histogramming/impl_torch.py index 4b5d475b..29c4cb0d 100644 --- a/tilebench/benchmarks/operators/histogramming/impl_torch.py +++ b/tilebench/benchmarks/operators/histogramming/impl_torch.py @@ -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)