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