From adc86f01751220fea94171fed43c16411f58d568 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Sat, 17 Oct 2026 11:17:34 -0400 Subject: [PATCH] wai-quantum: the tensor-engine member's kernel compiles for a device The helper defined its kernel inside main() and imported the kernel toolchain there. The simulator runs the kernel as Python and found the names; compiling for a device traces the kernel's source and resolves its names in the kernel's module, so it found none ("failed to resolve name 'nl.ndarray'") and every device test failed at the first request. The kernel now sits at module level; main() binds the toolchain as module globals after the target is chosen. Compiling for a device also runs the toolchain's compiler from PATH, and the interpreter the member is given need not be on it, so the helper puts that interpreter's directory first on PATH for a device target. Checked off the device with the toolchain's compiler: the kernel compiles for the device target; the simulator path is unchanged. Co-Authored-By: Claude Opus 5.5 --- wai/CHANGELOG.md | 19 ++++++++++ .../quantum_fabric/tensor_engine_helper.py | 37 ++++++++++++------- 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/wai/CHANGELOG.md b/wai/CHANGELOG.md index 953adf182..c3a9eafcf 100644 --- a/wai/CHANGELOG.md +++ b/wai/CHANGELOG.md @@ -112,6 +112,25 @@ wire format may change and such changes are recorded here. and its wire identifiers (`WEV1`, `WEP1`, `urn:wai:event:1`, `science.transaction.wai.event`, `X-WAI-EVENT`). +### wai-quantum (next release): the tensor-engine member's kernel compiles for a device + +Fix. Every pinned digest is unchanged. + +- **The kernel compiled for the simulator only.** The helper + (`quantum_fabric/tensor_engine_helper.py`) defined its kernel inside `main()` and + imported the kernel toolchain there. The simulator runs the kernel as Python, so it found + the names. Compiling for a device traces the kernel's source and looks its names up in + the kernel's module, so it found none and refused the first request (`failed to resolve + name`). On a device every test of the member failed at that request. The kernel now sits + at module level. `main()` still imports the toolchain after the target is chosen, but + binds it as module globals. +- **The compiler is found beside the interpreter.** Compiling for a device runs the + toolchain's compiler from `PATH`, and the interpreter the member is given need not be on + it. The helper now puts that interpreter's directory first on `PATH` for a device target. +- **Checked** off the device with the toolchain's own compiler: the kernel now compiles + for the device target. Execution there needs the device's runtime. The simulator path is + unchanged. + ### wai-quantum 0.6.2: the HIP member's matrix-instruction kernel, and the Metal member's tensor-operation kernel Additive. Every pinned digest is unchanged. diff --git a/wai/wai-quantum/src/quantum_fabric/tensor_engine_helper.py b/wai/wai-quantum/src/quantum_fabric/tensor_engine_helper.py index 6a3da3308..009e259d5 100644 --- a/wai/wai-quantum/src/quantum_fabric/tensor_engine_helper.py +++ b/wai/wai-quantum/src/quantum_fabric/tensor_engine_helper.py @@ -31,7 +31,25 @@ M = 112 MAGIC = 0x45545157 +def coefficients(w, a): + """The kernel. It sits at module level, and the toolchain's names are + module globals (bound in main), because compiling for a device resolves + the kernel's names in its module, not in an enclosing function.""" + tiles = a.shape[1] // TILE + out = nl.ndarray((M, a.shape[1]), dtype=nl.float32, buffer=nl.shared_hbm) + ws = nl.load(w) + for t in nl.affine_range(tiles): + xs = nl.load(a[:, t * TILE:(t + 1) * TILE]) + acc = nl.ndarray((M, TILE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(acc, ws, xs) + nl.store(out[:, t * TILE:(t + 1) * TILE], nl.copy(acc)) + return out + + def main(): + # Imported here, after the target is chosen (the toolchain reads it at + # import), into the module's namespace, where the kernel looks them up. + global np, ml_dtypes, nki, nisa, nl import os target = "simulator" @@ -45,6 +63,10 @@ def main(): return 2 if target != "simulator": os.environ.setdefault("NEURON_PLATFORM_TARGET_OVERRIDE", target) + # Compiling for a device runs the toolchain's compiler, found on PATH; + # it is installed beside this interpreter, which need not be on PATH. + here = os.path.dirname(sys.executable) + os.environ["PATH"] = here + os.pathsep + os.environ.get("PATH", "") import numpy as np import ml_dtypes @@ -52,21 +74,10 @@ def main(): import nki.isa as nisa import nki.language as nl - @nki.jit - def coefficients(w, a): - tiles = a.shape[1] // TILE - out = nl.ndarray((M, a.shape[1]), dtype=nl.float32, buffer=nl.shared_hbm) - ws = nl.load(w) - for t in nl.affine_range(tiles): - xs = nl.load(a[:, t * TILE:(t + 1) * TILE]) - acc = nl.ndarray((M, TILE), dtype=nl.float32, buffer=nl.psum) - nisa.nc_matmul(acc, ws, xs) - nl.store(out[:, t * TILE:(t + 1) * TILE], nl.copy(acc)) - return out - inp, out = sys.stdin.buffer, sys.stdout.buffer version = getattr(nki, "__version__", "unknown") - run = nki.simulate(coefficients) if target == "simulator" else coefficients + kernel = nki.jit(coefficients) + run = nki.simulate(kernel) if target == "simulator" else kernel out.write(f"wq-tensor-engine 1 {target} nki {version} numpy {np.__version__} ml_dtypes {ml_dtypes.__version__}\n".encode()) out.flush() -- 2.54.0 (Apple Git-157)