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]>
84 lines
4.1 KiB
Shell
Executable file
84 lines
4.1 KiB
Shell
Executable file
#!/usr/bin/env bash
|
|
#
|
|
# check.sh: run one build of the gate's tensor program on one device, both of
|
|
# its functions, and hold the outcome to the one it must have.
|
|
#
|
|
# check.sh iree <device> <module.vmfb> <expect>
|
|
# check.sh xla <platform> <module.mlir> <expect>
|
|
#
|
|
# Run in the directory the stablehlo_gate example wrote. <expect> is two
|
|
# characters, `=` exact or `!` differs. For a gate module, the rows' outcome
|
|
# then the probe's: the conforming program must give `==`; mutants m1 and m2
|
|
# (the rounding constant, the shift) `!!`; m3 (row 0's last sign) `!=`, since
|
|
# its probe is the program's. For the limb module (a name starting `limbs`),
|
|
# the coefficients of eight-bit limbs, then of nine-bit limbs, which bf16
|
|
# cannot hold: a conforming build gives `=!`.
|
|
#
|
|
# iree: the module was built by that compiler; its runtime (`iree-run-module`,
|
|
# on PATH) runs it on <device> and compares the outputs, and an exact outcome
|
|
# must say so in its log. xla: xla_run.py, beside this script, builds and runs
|
|
# the module on <platform> under ${XLA_PYTHON:-python3}.
|
|
#
|
|
# Exit status: 0 the outcome is the expected one; 1 it is not; 2 a usage error.
|
|
|
|
set -euo pipefail
|
|
|
|
[ $# = 4 ] || { sed -n '3,20p' "$0" >&2; exit 2; }
|
|
how=$1 dev=$2 module=$3 expect=$4
|
|
case "$expect" in [=!][=!]) ;; *) echo "check.sh: <expect> is two of = and !" >&2; exit 2 ;; esac
|
|
|
|
size() { stat -c %s "$1" 2>/dev/null || stat -f %z "$1"; }
|
|
n=$(( $(size r0.bin) / 8 ))
|
|
case "$(basename "$module")" in limbs*) kind=limbs ;; *) kind=gate ;; esac
|
|
rows_in=(e0.npy e1.npy e2.npy e3.npy e4.npy e5.npy e6.npy e7.npy a0r.npy a0i.npy a1r.npy a1i.npy)
|
|
rows_out=(r0.bin r1.bin r2.bin r3.bin)
|
|
probe_in=(pa.npy pb.npy)
|
|
probe_out=(prod.bin fx.bin)
|
|
|
|
# run <function> <log> <inputs...> -- <expected...>: 0 exact, 1 differs. A
|
|
# run that ends any other way (a device that will not open, a module that will
|
|
# not load, an exception) is neither, and fails the check: otherwise it would
|
|
# read as a mutant caught.
|
|
run() {
|
|
local fn=$1 log=$2; shift 2
|
|
local status=0
|
|
if [ "$how" = iree ]; then
|
|
local args=() seen=0 a
|
|
for a in "$@"; do
|
|
if [ "$a" = -- ]; then seen=1
|
|
elif [ $seen = 0 ]; then case "$a" in *=@*) args+=("--input=$a") ;; *) args+=("--input=@$a") ;; esac
|
|
else case "$a" in *=@*) args+=("--expected_output=$a") ;; *) args+=("--expected_output=${n}xi64=@$a") ;; esac
|
|
fi
|
|
done
|
|
# Floats are compared exactly too.
|
|
iree-run-module --device="$dev" --module="$module" --function="$fn" --expected_f32_threshold=0 "${args[@]}" > "$log" 2>&1 || status=$?
|
|
if [ $status = 0 ] && grep -q '^\[SUCCESS\] all function outputs matched' "$log"; then return 0; fi
|
|
if [ $status != 0 ] && grep -q '^\[FAILED\] result\[[0-9]*\]: .* does not match the expected' "$log"; then return 1; fi
|
|
else
|
|
"${XLA_PYTHON:-python3}" -I "$(dirname "$0")/xla_run.py" "$dev" "$module" "$fn" "$@" > "$log" 2>&1 || status=$?
|
|
if [ $status = 0 ] && grep -q "^$fn on .*: exact\$" "$log"; then return 0; fi
|
|
if [ $status = 1 ] && grep -q "^$fn on .*: differs\$" "$log"; then return 1; fi
|
|
fi
|
|
echo "::error::$(basename "$module") $fn on $dev ($how): the run did not complete (exit $status)"
|
|
head -30 "$log"
|
|
exit 1
|
|
}
|
|
|
|
tag="$(basename "$module")"
|
|
word() { [ "$1" = '=' ] && echo exact || echo differ; }
|
|
if [ "$kind" = limbs ]; then
|
|
ln=$(( $(size c.bin) / 4 / 112 ))
|
|
r='='; run coefficients "$tag.limbs8.log" "[email protected]" "16x${ln}[email protected]" -- "112x${ln}[email protected]" || r='!'
|
|
p='='; run coefficients "$tag.limbs9.log" "[email protected]" "16x${ln}[email protected]" -- "112x${ln}[email protected]" || p='!'
|
|
first="eight-bit limbs" second="nine-bit limbs" logs=("$tag.limbs8.log" "$tag.limbs9.log")
|
|
else
|
|
r='='; run rows "$tag.rows.log" "${rows_in[@]}" -- "${rows_out[@]}" || r='!'
|
|
p='='; run probe "$tag.probe.log" "${probe_in[@]}" -- "${probe_out[@]}" || p='!'
|
|
first=rows second=probe logs=("$tag.rows.log" "$tag.probe.log")
|
|
fi
|
|
echo "$tag on $dev ($how): $first $(word $r), $second $(word $p)"
|
|
if [ "$r$p" != "$expect" ]; then
|
|
echo "::error::$tag on $dev ($how): $first $(word $r), $second $(word $p); expected $first $(word "${expect:0:1}"), $second $(word "${expect:1:1}")"
|
|
head -20 "${logs[@]}"
|
|
exit 1
|
|
fi
|