From 469e13e7b71a4434292e2df499401fecb56a9324 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 8 Jun 2026 20:03:24 +0000 Subject: [PATCH 01/12] Initial commit for nki mul2 --- benchmarks/operators/mul2/impl_nki.py | 85 +++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 benchmarks/operators/mul2/impl_nki.py diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py new file mode 100644 index 00000000..8b9eab7a --- /dev/null +++ b/benchmarks/operators/mul2/impl_nki.py @@ -0,0 +1,85 @@ +import os +os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2" + +import torch +from torch_xla.core import xla_model as xm + +import neuronxcc.nki as nki +import neuronxcc.nki.language as nl +import neuronxcc.nki.isa as nisa + +#defines maximum partitions supported +PMAX = nl.tile_size.pmax + +#vector_add +@nki.jit +def _mul2_kernel(a_input, b_input): + #assert a_input & b_input are the same shape + assert a_input.shape == b_input.shape + assert a_input.dtype == b_input.dtype + + num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX + + #allocate a result tensor in HBM (move result back to HBM) + hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) + + + #Allocate sbuf space for partitions of the dataset at a time + for i in nl.affine_range(num_blocks): + offset = i*PMAX + + mask = nl.arange(PMAX)[:, None] < (a_input.shape[0] - offset) + + a_tile = nl.load(a_input[offset : offset + PMAX, :], mask=mask) + b_tile = nl.load(b_input[offset : offset + PMAX, :], mask=mask) + + #compute addition, returns a tensor + sbuf_result_tile = nisa.tensor_tensor(data1=a_tile, data2=b_tile, op=nl.multiply) + + #copy the calculated addition result back into hbm in 128 partition increments + nl.store(hbm_result_tile[offset:offset+PMAX, :], value=sbuf_result_tile, mask=mask) + + return hbm_result_tile + +def run(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + assert a.shape == b.shape, f"shape mismatch: {a.shape} vs {b.shape}" + assert a.dtype == b.dtype, f"dtype mismatch: {a.dtype} vs {b.dtype}" + + device = xm.xla_device() + a = a.to(device) + b = b.to(device) + + return _mul2_kernel(a, b) + +#local main for testing +if __name__ == "__main__": + a = torch.ones((256, 10), dtype=torch.float32) + b = torch.ones((256, 10), dtype=torch.float32) + result = run(a, b) + print("compiled OK") + + """ + device = xm.xla_device() + #test 1: exact multiple of PMAX + a = torch.ones((256, 10), dtype=torch.float32).to(device) + b = torch.ones((256, 10), dtype=torch.float32).to(device) + result = _mul2_kernel(a, b) + assert torch.allclose(result, a * b), "test 1 failed" # change 2: a+b -> a*b + print("test 1 passed: exact multiple of PMAX") + + #test 2: non-multiple of PMAX + a = torch.ones((244, 10), dtype=torch.float32).to(device) + b = torch.ones((244, 10), dtype=torch.float32).to(device) + result = _mul2_kernel(a, b) + assert torch.allclose(result, a * b), "test 2 failed" + print("test 2 passed: non-multiple of PMAX") + + #test 3: known values - change 3: 3*5=15 instead of 3+5=8 + a = torch.full((128, 10), 3.0, dtype=torch.float32).to(device) + b = torch.full((128, 10), 5.0, dtype=torch.float32).to(device) + result = _mul2_kernel(a, b) + assert torch.all(result == 15.0), "test 3 failed" + print("test 3 passed: known values 3*5=15") + + print("all tests passed") + """ \ No newline at end of file From ed3f5a3d0373060bd338478d8d3a6c7ed0f4dcab Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 10 Jun 2026 18:17:28 +0000 Subject: [PATCH 02/12] Revised nki mul2 to fit the kernel guidelines. Fixed logical flaws in the mul2 kernel --- benchmarks/operators/mul2/impl_nki.py | 101 +++++++------------------- 1 file changed, 27 insertions(+), 74 deletions(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index 8b9eab7a..4a797dd1 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -1,85 +1,38 @@ -import os -os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2" - import torch from torch_xla.core import xla_model as xm -import neuronxcc.nki as nki -import neuronxcc.nki.language as nl -import neuronxcc.nki.isa as nisa - -#defines maximum partitions supported -PMAX = nl.tile_size.pmax - -#vector_add -@nki.jit -def _mul2_kernel(a_input, b_input): - #assert a_input & b_input are the same shape - assert a_input.shape == b_input.shape - assert a_input.dtype == b_input.dtype - - num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX - - #allocate a result tensor in HBM (move result back to HBM) - hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) - - - #Allocate sbuf space for partitions of the dataset at a time - for i in nl.affine_range(num_blocks): - offset = i*PMAX - - mask = nl.arange(PMAX)[:, None] < (a_input.shape[0] - offset) - - a_tile = nl.load(a_input[offset : offset + PMAX, :], mask=mask) - b_tile = nl.load(b_input[offset : offset + PMAX, :], mask=mask) - - #compute addition, returns a tensor - sbuf_result_tile = nisa.tensor_tensor(data1=a_tile, data2=b_tile, op=nl.multiply) - - #copy the calculated addition result back into hbm in 128 partition increments - nl.store(hbm_result_tile[offset:offset+PMAX, :], value=sbuf_result_tile, mask=mask) +try: + import neuronxcc.nki as nki + import neuronxcc.nki.language as nl + import neuronxcc.nki.isa as nisa + PMAX = nl.tile_size.pmax +except ImportError: + nki = None - return hbm_result_tile +if nki is not None: + @nki.jit + def mul2_kernel(a_input): + num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX -def run(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - assert a.shape == b.shape, f"shape mismatch: {a.shape} vs {b.shape}" - assert a.dtype == b.dtype, f"dtype mismatch: {a.dtype} vs {b.dtype}" + hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) - device = xm.xla_device() - a = a.to(device) - b = b.to(device) + for i in range(num_blocks): + offset = i*PMAX - return _mul2_kernel(a, b) + partition_index = nl.arange(PMAX)[:, None] + free_dim_index = nl.arange(a_input.shape[1])[None, :] + + mask = partition_index < (a_input.shape[0] - offset) -#local main for testing -if __name__ == "__main__": - a = torch.ones((256, 10), dtype=torch.float32) - b = torch.ones((256, 10), dtype=torch.float32) - result = run(a, b) - print("compiled OK") - - """ - device = xm.xla_device() - #test 1: exact multiple of PMAX - a = torch.ones((256, 10), dtype=torch.float32).to(device) - b = torch.ones((256, 10), dtype=torch.float32).to(device) - result = _mul2_kernel(a, b) - assert torch.allclose(result, a * b), "test 1 failed" # change 2: a+b -> a*b - print("test 1 passed: exact multiple of PMAX") + a_tile = nl.load(a_input[offset + partition_index, free_dim_index], mask=mask) - #test 2: non-multiple of PMAX - a = torch.ones((244, 10), dtype=torch.float32).to(device) - b = torch.ones((244, 10), dtype=torch.float32).to(device) - result = _mul2_kernel(a, b) - assert torch.allclose(result, a * b), "test 2 failed" - print("test 2 passed: non-multiple of PMAX") + result_tile = nl.multiply(a_tile, 2, mask=mask) + + nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) - #test 3: known values - change 3: 3*5=15 instead of 3+5=8 - a = torch.full((128, 10), 3.0, dtype=torch.float32).to(device) - b = torch.full((128, 10), 5.0, dtype=torch.float32).to(device) - result = _mul2_kernel(a, b) - assert torch.all(result == 15.0), "test 3 failed" - print("test 3 passed: known values 3*5=15") + return hbm_result_tile - print("all tests passed") - """ \ No newline at end of file +def run(x: torch.Tensor, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor: + x_2d = x.reshape(-1, 1) + result = mul2_kernel(x_2d) + return result.reshape(-1) From 1db381e4f613c1dd14c4eb963e7acb10a4da484f Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Jun 2026 20:37:18 +0000 Subject: [PATCH 03/12] removed unnecessary import and added get_last_config() --- benchmarks/operators/mul2/impl_nki.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index 4a797dd1..f202f48b 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -1,5 +1,4 @@ import torch -from torch_xla.core import xla_model as xm try: import neuronxcc.nki as nki @@ -36,3 +35,6 @@ def run(x: torch.Tensor, block_size: int = 1024, autotune=False, **kwargs) -> to x_2d = x.reshape(-1, 1) result = mul2_kernel(x_2d) return result.reshape(-1) + +def get_last_config() -> dict | None: + return None From 55acbc582c2f421f39ffec23c3cc0198c250bed8 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 16 Jun 2026 21:18:15 +0000 Subject: [PATCH 04/12] Changed to free dimension tiling --- benchmarks/operators/mul2/impl_nki.py | 35 ++++++++++++++++++--------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index f202f48b..1bef74d8 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -5,36 +5,49 @@ import neuronxcc.nki.language as nl import neuronxcc.nki.isa as nisa PMAX = nl.tile_size.pmax + FMAX_SBUF = 64000 except ImportError: nki = None if nki is not None: @nki.jit def mul2_kernel(a_input): - num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX + free_dim = min(a_input.shape[1], FMAX_SBUF) + + num_free_blocks = (a_input.shape[1] + (free_dim - 1)) // free_dim hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) - for i in range(num_blocks): - offset = i*PMAX + partition_index = nl.arange(PMAX)[:, None] + + for j in range(num_free_blocks): + free_offset = j * free_dim - partition_index = nl.arange(PMAX)[:, None] - free_dim_index = nl.arange(a_input.shape[1])[None, :] + free_dim_index = nl.arange(free_dim)[None, :] - mask = partition_index < (a_input.shape[0] - offset) + free_mask = free_dim_index < (a_input.shape[1] - free_offset) - a_tile = nl.load(a_input[offset + partition_index, free_dim_index], mask=mask) + a_tile = nl.load(a_input[partition_index, free_offset + free_dim_index], mask=free_mask) - result_tile = nl.multiply(a_tile, 2, mask=mask) + result_tile = nl.multiply(a_tile, 2, mask=free_mask) - nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) + nl.store(hbm_result_tile[partition_index, free_offset + free_dim_index], value=result_tile, mask=free_mask) return hbm_result_tile def run(x: torch.Tensor, block_size: int = 1024, autotune=False, **kwargs) -> torch.Tensor: - x_2d = x.reshape(-1, 1) + n = x.numel() + + free_dim = (n + (PMAX -1)) // PMAX + padded_size = PMAX * free_dim + + if padded_size > n: + x = torch.nn.functional.pad(x, (0, padded_size - n)) + + x_2d = x.reshape(PMAX, free_dim) + result = mul2_kernel(x_2d) - return result.reshape(-1) + return result.reshape(-1)[:n] def get_last_config() -> dict | None: return None From 5dace0c5ab425db5814584723f51133ecc9a6a3a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Thu, 18 Jun 2026 16:21:32 +0000 Subject: [PATCH 05/12] fixed tiling logic --- benchmarks/operators/mul2/impl_nki.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index 1bef74d8..5a3b4c85 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -5,33 +5,29 @@ import neuronxcc.nki.language as nl import neuronxcc.nki.isa as nisa PMAX = nl.tile_size.pmax - FMAX_SBUF = 64000 except ImportError: nki = None if nki is not None: @nki.jit def mul2_kernel(a_input): - free_dim = min(a_input.shape[1], FMAX_SBUF) - - num_free_blocks = (a_input.shape[1] + (free_dim - 1)) // free_dim + num_blocks = (a_input.shape[1] + (PMAX - 1)) // PMAX hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) - partition_index = nl.arange(PMAX)[:, None] - - for j in range(num_free_blocks): - free_offset = j * free_dim + for i in range(num_blocks): + offset = i * PMAX - free_dim_index = nl.arange(free_dim)[None, :] + partition_index = nl.arange(PMAX)[:, None] + free_dim_index = nl.arange(a_input.shape[1])[None, :] - free_mask = free_dim_index < (a_input.shape[1] - free_offset) + mask = free_dim_index < (a_input.shape[0] - offset) - a_tile = nl.load(a_input[partition_index, free_offset + free_dim_index], mask=free_mask) + a_tile = nl.load(a_input[offset + partition_index, free_dim_index], mask=mask) - result_tile = nl.multiply(a_tile, 2, mask=free_mask) + result_tile = nl.multiply(a_tile, 2, mask=mask) - nl.store(hbm_result_tile[partition_index, free_offset + free_dim_index], value=result_tile, mask=free_mask) + nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) return hbm_result_tile From 13421f72c4297f64dff1cb408fcf7b8c49fe436c Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 1 Jul 2026 20:35:40 +0000 Subject: [PATCH 06/12] modified json parsing in _total_time_ms to fit trainium conventions --- core/nki_timer.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index ce3711a4..463ac737 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -130,9 +130,17 @@ def _total_time_ms(summary_json_text: str) -> float: total_time across rows and convert seconds -> ms. """ data = json.loads(summary_json_text) - rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - if isinstance(rows, dict): - rows = [rows] + + #rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + #if isinstance(rows, dict): + #rows = [rows] + if isinstance(data, dict) and not any(k in data for k in ["total_time", "summary", "rows"]): + rows = [v for v in data.values() if isinstance(v, dict)] + else: + rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + if isinstance(rows, dict): + rows = [rows] + times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: raise RuntimeError(f"no 'total_time' in neuron-profile summary-json: {summary_json_text[:300]}") From 44b94e4b5ac6e58f1b7e83840cc45e5760e983a9 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 1 Jul 2026 20:39:38 +0000 Subject: [PATCH 07/12] set NEURON_RT_NUM_CORES=1 to resolve neuron-profile core allocation error on trn2.3xlarge --- scripts/run_bench.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/scripts/run_bench.py b/scripts/run_bench.py index 7e06f091..d6e98da4 100644 --- a/scripts/run_bench.py +++ b/scripts/run_bench.py @@ -1,6 +1,7 @@ import argparse import csv import json +import os from pathlib import Path from core.engine import run_benchmark_suite @@ -23,6 +24,8 @@ def _split(results: list[dict], active: list[str]) -> tuple[list[dict], list[dic def main(): + os.environ["NEURON_RT_NUM_CORES"] = "1" + parser = argparse.ArgumentParser(description="Run TileBench benchmarks") parser.add_argument("--operator", type=str, default="vector_add", help="Operator to benchmark") From 4b07dccd1ba5ff53a0776fb45566af188415fc0a Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Mon, 6 Jul 2026 03:00:00 +0000 Subject: [PATCH 08/12] Added free dim tiling --- benchmarks/operators/mul2/impl_nki.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index 5a3b4c85..4c87f75e 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -13,21 +13,33 @@ def mul2_kernel(a_input): num_blocks = (a_input.shape[1] + (PMAX - 1)) // PMAX + free_tile_size = 16384 + + num_free_blocks = (a_input.shape[1] + free_tile_size - 1) // free_tile_size + hbm_result_tile = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.hbm) for i in range(num_blocks): offset = i * PMAX partition_index = nl.arange(PMAX)[:, None] - free_dim_index = nl.arange(a_input.shape[1])[None, :] + + mask_p = partition_index < (a_input.shape[0] - offset) + + for j in range(num_free_blocks): + free_offset = j * free_tile_size + + free_dim_index = nl.arange(free_tile_size)[None, :] - mask = free_dim_index < (a_input.shape[0] - offset) + mask_f = free_dim_index < (a_input.shape[1] - free_offset) + + mask = mask_p & mask_f - a_tile = nl.load(a_input[offset + partition_index, free_dim_index], mask=mask) + a_tile = nl.load(a_input[offset + partition_index, free_offset + free_dim_index], mask=mask) - result_tile = nl.multiply(a_tile, 2, mask=mask) + result_tile = nl.multiply(a_tile, 2, mask=mask) - nl.store(hbm_result_tile[offset + partition_index, free_dim_index], value=result_tile, mask=mask) + nl.store(hbm_result_tile[offset + partition_index, free_offset + free_dim_index], value=result_tile, mask=mask) return hbm_result_tile From a598f558a2b303bed09b62dcb5b7e9b61cbeeae6 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Wed, 8 Jul 2026 08:43:42 +0000 Subject: [PATCH 09/12] Fixed tiling bug --- benchmarks/operators/mul2/impl_nki.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmarks/operators/mul2/impl_nki.py b/benchmarks/operators/mul2/impl_nki.py index 4c87f75e..770f58be 100644 --- a/benchmarks/operators/mul2/impl_nki.py +++ b/benchmarks/operators/mul2/impl_nki.py @@ -11,7 +11,7 @@ if nki is not None: @nki.jit def mul2_kernel(a_input): - num_blocks = (a_input.shape[1] + (PMAX - 1)) // PMAX + num_blocks = (a_input.shape[0] + (PMAX - 1)) // PMAX free_tile_size = 16384 From b3e4dbe433258d9f9a83fe650fca2547276b92af Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 4 Aug 2026 18:42:32 +0000 Subject: [PATCH 10/12] removed code in run_bench and nki_timer --- core/nki_timer.py | 13 +++---------- scripts/run_bench.py | 3 --- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index 463ac737..d00c9ff8 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -130,16 +130,9 @@ def _total_time_ms(summary_json_text: str) -> float: total_time across rows and convert seconds -> ms. """ data = json.loads(summary_json_text) - - #rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - #if isinstance(rows, dict): - #rows = [rows] - if isinstance(data, dict) and not any(k in data for k in ["total_time", "summary", "rows"]): - rows = [v for v in data.values() if isinstance(v, dict)] - else: - rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) - if isinstance(rows, dict): - rows = [rows] + rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) + if isinstance(rows, dict): + rows = [rows] times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: diff --git a/scripts/run_bench.py b/scripts/run_bench.py index d6e98da4..7e06f091 100644 --- a/scripts/run_bench.py +++ b/scripts/run_bench.py @@ -1,7 +1,6 @@ import argparse import csv import json -import os from pathlib import Path from core.engine import run_benchmark_suite @@ -24,8 +23,6 @@ def _split(results: list[dict], active: list[str]) -> tuple[list[dict], list[dic def main(): - os.environ["NEURON_RT_NUM_CORES"] = "1" - parser = argparse.ArgumentParser(description="Run TileBench benchmarks") parser.add_argument("--operator", type=str, default="vector_add", help="Operator to benchmark") From 1ae27c815a35ca23ce61f2a3685c08407043cbb7 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 4 Aug 2026 18:43:34 +0000 Subject: [PATCH 11/12] removed space --- core/nki_timer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index d00c9ff8..2fdda4d3 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -132,8 +132,7 @@ def _total_time_ms(summary_json_text: str) -> float: data = json.loads(summary_json_text) rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) if isinstance(rows, dict): - rows = [rows] - + rows = [rows] times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: raise RuntimeError(f"no 'total_time' in neuron-profile summary-json: {summary_json_text[:300]}") From ce3a933a27b24c682727c367206757ad14983474 Mon Sep 17 00:00:00 2001 From: Ubuntu Date: Tue, 4 Aug 2026 18:43:57 +0000 Subject: [PATCH 12/12] removed space --- core/nki_timer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/nki_timer.py b/core/nki_timer.py index 2fdda4d3..ce3711a4 100644 --- a/core/nki_timer.py +++ b/core/nki_timer.py @@ -132,7 +132,7 @@ def _total_time_ms(summary_json_text: str) -> float: data = json.loads(summary_json_text) rows = data if isinstance(data, list) else data.get("summary", data.get("rows", [data])) if isinstance(rows, dict): - rows = [rows] + rows = [rows] times = [float(r["total_time"]) for r in rows if isinstance(r, dict) and "total_time" in r] if not times: raise RuntimeError(f"no 'total_time' in neuron-profile summary-json: {summary_json_text[:300]}")