2a30dcd360
stablehlo::limb_module: the tensor-engine member's limb layout as one StableHLO dot_general of bf16 operands, accumulated and returned in f32, the program a compiler maps to a matrix unit. stablehlo_gate writes it with the batch's limb weights and limbs and the exact coefficients, and a nine-bit-limb data set whose exact coefficients bf16 cannot reproduce (a run that matches them is not comparing). check.sh holds the limb module to "=!" (exact on eight-bit limbs, different on nine); both drivers read typed raw files, and floats are compared exactly (IREE's --expected_f32_threshold=0). Measured: exact under both compilers on the x86-64 and arm64 CPUs and on the discrete GPU through CUDA; the nine-bit limbs differ everywhere. Not exact: the IREE 3.12 HIP build for gfx1151, whose bf16 multiply returns non-integers (-14.0002 for -14), so the workflow shows that step and does not require it. Not built: the IREE Vulkan and Metal builds do not lower bf16 in this version. Co-Authored-By: Claude Opus 5.5 <[email protected]>
81 lines
3 KiB
Python
81 lines
3 KiB
Python
"""Build one function of the gate's tensor program with a second, independent
|
|
tensor compiler (XLA, through its Python package), run it on one device, and
|
|
compare every output with the exact one.
|
|
|
|
python -I xla_run.py <platform> <module.mlir> <function> <input>... -- <expected>...
|
|
|
|
<platform> is the package's name for the device kind (`cpu`, `cuda`, `tpu`).
|
|
This compiler builds a module's `main`, so the chosen function is renamed
|
|
`main` and every other public one made private (and unused). An input is a
|
|
`.npy` file, or `<shape>x<type>=@<raw file>` (`[email protected]`); an
|
|
expected output is a raw file of little-endian 64-bit integers, or
|
|
`<shape>x<type>=@<raw file>`. Every output must equal its expected values
|
|
exactly, floats included.
|
|
|
|
Exit status: 0 every output exact; 1 an output differs; 2 a usage error.
|
|
"""
|
|
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
|
|
def main(argv):
|
|
if len(argv) < 4 or "--" not in argv:
|
|
print(__doc__, file=sys.stderr)
|
|
return 2
|
|
platform, path, fn = argv[:3]
|
|
rest = argv[3:]
|
|
k = rest.index("--")
|
|
inputs, expected = rest[:k], rest[k + 1:]
|
|
|
|
import jax
|
|
|
|
jax.config.update("jax_enable_x64", True)
|
|
import re
|
|
|
|
text = open(path).read()
|
|
if text.count(f"func.func @{fn}(") != 1:
|
|
print(f"{path}: expected one public function @{fn}", file=sys.stderr)
|
|
return 2
|
|
text = re.sub(r"func\.func @(\w+)\(", lambda g: "func.func @main(" if g.group(1) == fn else f"func.func private @{g.group(1)}(", text)
|
|
|
|
import ml_dtypes
|
|
|
|
types = {"bf16": ml_dtypes.bfloat16, "f32": "<f4", "i32": "<i4", "i64": "<i8"}
|
|
|
|
def typed(spec):
|
|
"""`<shape>x<type>=@<file>` as an array, or None for a plain path."""
|
|
if "=@" not in spec:
|
|
return None
|
|
head, file = spec.split("=@", 1)
|
|
*dims, ty = head.split("x")
|
|
return np.fromfile(file, dtype=types[ty]).reshape([int(d) for d in dims])
|
|
|
|
device = jax.devices(platform)[0]
|
|
exe = device.client.compile_and_load(text, [device])
|
|
args = [typed(p) if typed(p) is not None else np.load(p) for p in inputs]
|
|
results = exe.execute([jax.device_put(x, device) for x in args])
|
|
if len(results) != len(expected):
|
|
print(f"{fn}: {len(results)} outputs, {len(expected)} expected", file=sys.stderr)
|
|
return 1
|
|
bad = 0
|
|
for got, name in zip(results, expected):
|
|
want = typed(name)
|
|
if want is None:
|
|
want = np.fromfile(name, dtype="<i8")
|
|
got = np.asarray(got)
|
|
if got.dtype != want.dtype or got.size != want.size:
|
|
print(f"{name}: {got.dtype}{list(got.shape)}, not {want.dtype}{list(want.shape)}")
|
|
bad += 1
|
|
continue
|
|
got = got.reshape(want.shape)
|
|
if not np.array_equal(got, want):
|
|
print(f"{name}: {int(np.sum(got != want))} of {want.size} values differ")
|
|
bad += 1
|
|
print(f"{fn} on {device.platform} ({device.device_kind}): {'exact' if bad == 0 else 'differs'}")
|
|
return 1 if bad else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main(sys.argv[1:]))
|