#!/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" "16x112xbf16=@w.bin" "16x${ln}xbf16=@a.bin" -- "112x${ln}xf32=@c.bin" || r='!'
  p='='; run coefficients "$tag.limbs9.log" "16x112xbf16=@w9.bin" "16x${ln}xbf16=@a9.bin" -- "112x${ln}xf32=@c9.bin" || 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