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.

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