compile: four more Espresso workarounds from fuzz round 2; matmul saturation characterized

aneforge/_compile.py

@@ -181,15 +181,32 @@ def weight(self, name: str, W: np.ndarray, allow_int8: bool,
 # param-free unary ops: graph op name == MIL op name, signature op(x = ...)
 @op("relu", "silu", "sigmoid", "tanh", "exp", "sqrt", "abs",
   "sin", "cos", "erf", "relu6", "softsign", "atan", "exp2",
-  "floor", "ceil", "round", "sign")
+  "floor", "ceil", "sign")
 def _e_unary(em, t, n, s):
   em.line(f'{em.ty(t.shape)} {n} = {t.op}(x = {s[0]})[name = string("{n}")];')
 
 
+@op("round")
+def _e_round(em, t, n, s):
+  """The ANE round kernel corrupts values with |x| >= 1024 (measured on M5/H17s:
+  round(1024)=1025, round(1025)=1026, round(-2047)=-2048 - it adds 0.5 in fp16, where the tie
+  is unrepresentable). Every fp16 value with |x| >= 1024 is already an integer, so route them
+  through unchanged: select(|x| < 1024, round(x), x). Exact on the full fp16 range, and safe on
+  every family regardless of whether its kernel shares the bug (the select form is semantically
+  identity there); cross-compiles for H13-H16s (see test_espresso_workarounds)."""
+  thr = float(np.float16(1024.0)).hex()
+  em.line(f'fp16 {n}_c = const()[name = string("{n}_c"), val = fp16({thr})];')
+  em.line(f'{em.ty(t.shape)} {n}_a = abs(x = {s[0]})[name = string("{n}_a")];')
+  em.line(f'{em.ty_dt(t.shape, "bool")} {n}_lt = less(x = {n}_a, y = {n}_c)[name = string("{n}_lt")];')
+  em.line(f'{em.ty(t.shape)} {n}_r = round(x = {s[0]})[name = string("{n}_r")];')
+  em.line(f'{em.ty(t.shape)} {n} = select(cond = {n}_lt, a = {n}_r, b = {s[0]})[name = string("{n}")];')
+
+
 @op("softplus")
 def _e_softplus(em, t, n, s):
   """Espresso's ANE softplus computes exp(x) in the fp16 datapath and returns 0 for every
-  x >= ln(65504) ~= 11.09 (measured on M5 / macOS 26.5: sp(10)=10, sp(11)=0 - silent, not an
+  x >= ln(65504) ~= 11.09 (measured on M5/H17s, macOS 26.5; other families may differ but the
+  stable form is exact everywhere: sp(10)=10, sp(11)=0 - silent, not an
   error). Lower the stable split instead: softplus(x) = softplus(min(x, 10)) + relu(x - 10),
   whose error is <= log1p(e^-10) ~= 4.5e-5 and whose exp argument never overflows."""
   ten = float(np.float16(10.0)).hex(); nten = float(np.float16(-10.0)).hex()
@@ -310,6 +327,11 @@ def scalar_const(v): return v.op == "const_array" and np.asarray(v.attrs["value"
 
 @op("muls")
 def _e_muls(em, t, n, s):
+  if float(t.attrs["k"]) == 0.0:
+    # mul-by-zero after a reduce crashes Espresso (ANECCompile FAILED, measured on M5); x - x is the
+    # same zeros for every finite x and lowers everywhere
+    em.line(f'{em.ty(t.shape)} {n} = sub(x = {s[0]}, y = {s[0]})[name = string("{n}")];')
+    return
   if t.srcs[0].op in _ROUNDING_OPS:                # the same Espresso scale-slot mis-fusion
     _e_mul_full_const(em, t, n, s[0], float(np.float16(t.attrs["k"])))
     return
@@ -900,6 +922,7 @@ def cross_compile_check(out: Tensor, target, int8: bool = False) -> bool:
   """Does the graph compile for another ANE family, checked from this host (compile-level only)?"""
   from . import _runtime
   from . import _targets as TG
+  if out.op == "input": out = out * 1.0            # same empty-program guard as compile()
   if isinstance(target, str):
     arch = target.strip().lower()
     # e5rt silently falls back to the host target on an unknown arch string (false pass); gate first.
@@ -1179,6 +1202,11 @@ def compile(out: Tensor, int8: bool = False, build_dir=None, opt: "str | int | N
       block_size: int = 32, validate: bool = False, target=None,
       _check_precision: bool = True):
   """Lower `out` into ONE fused ANE program (or a segmented plan if it has `af.sdpa` nodes)."""
+  if out.op == "input":
+    # A graph whose output IS an input lowers to an empty MIL body, which crashes Espresso's
+    # ANE compiler ("unordered_map::at: key not found"). Wrap in an exact scalar mul(1.0) so
+    # the program always has at least one op (and downstream port naming stays consistent).
+    out = out * 1.0
   if _check_precision:                 # once per user compile (internal re-entries pass False)
     _precision_signal(out, strict=validate)
     _dispatch_floor_signal(out)

aneforge/linalg.py

@@ -86,7 +86,7 @@ def conjugate_gradient(A, b, iters: int = 20, x0=None, refine: int = 0):
 
   bT = af.input((1, n))
   x0T = af.input((1, n)) if x0 is not None else None
-  dot = lambda u, v: (u * v) @ ones                       # ANE matmul dot (wide accum)
+  dot = lambda u, v: (u * v) @ ones                       # ANE matmul dot (saturates above ~32752 = fp16_max/2; scaling keeps sums in range)
 
   def cg_block(x, r, K):
     """K unrolled CG steps from residual r (p=r), accumulating into x."""

scripts/fuzz.py

@@ -135,6 +135,29 @@ def _rand_shape(rng):
     s = tuple(int(rng.choice(DIM_POOL)) for _ in range(rank))
     if np.prod(s) <= MAX_ELEMS: return s
 
+def _acc_bound(mode): return INT_MAX if mode == "int" else FLOAT_MAX
+
+def _acc_violation(spec, feed):
+  """True if any matmul/reduce-sum node's worst-case |partial sum| exceeds the mode bound.
+  Measured on M5/H17s - per-family datapath formats differ across ANE generations (see the
+  guide's datapath chapter), so these cliffs may sit elsewhere on other chips; a community run
+  that diverges from this model on another family is DATA, not noise. On H17s: matmul results
+  SATURATE to inf above ~32752 = fp16_max/2 for every
+  K (the matmul sibling of the documented slice-x16 saturation at 4094), and integer reduce sums
+  crossing 2048 lose bit-exactness. The oracle only judges graphs whose accumulations provably
+  stay in range under ANY summation order; the bounds here sit below both cliffs. (Two known
+  transpose-fed matmul cases return inf even BELOW these bounds - an open finding.)"""
+  vs = [np.asarray(feed, np.float64)]
+  for nd in spec["nodes"]:
+    x = vs[nd["src"][0]]
+    if nd["op"] == "matmul":
+      W = np.asarray(_weight(x.shape[1], nd["n"], nd["wseed"], spec["mode"]), np.float64)
+      if (np.abs(x) @ np.abs(W)).max() > _acc_bound(spec["mode"]): return True
+    if nd["op"] in ("rsum", "rmean"):
+      if np.abs(x).sum(axis=nd["axis"]).max() > _acc_bound(spec["mode"]): return True
+    vs.append(_mirror_node(nd, vs, spec))
+  return False
+
 def _ok(y, mode):
   if not np.isfinite(y).all(): return False
   if mode == "int":
@@ -206,6 +229,10 @@ def gen_spec(seed):
       if np.prod(x.shape) * np.prod(reps) > MAX_ELEMS: continue
       y = np.tile(x, reps); node = {"op": "tile", "src": [i], "reps": reps}
     if node is None or y is None or not _ok(y, mode): continue
+    if node["op"] == "matmul":
+      W = np.asarray(_weight(x.shape[1], node["n"], node["wseed"], mode), np.float64)
+      if (np.abs(x) @ np.abs(W)).max() > _acc_bound(mode): continue
+    if node["op"] in ("rsum", "rmean") and np.abs(x).sum(axis=node["axis"]).max() > _acc_bound(mode): continue
     vals.append(y); spec["nodes"].append(node)
   return spec
 
@@ -221,7 +248,12 @@ def _mirror(spec, feed, dtype):
   """Evaluate the spec's numpy mirror at `dtype`; returns (all_values, output)."""
   vs = [np.asarray(feed, dtype)]
   for nd in spec["nodes"]:
-    op = nd["op"]; x = vs[nd["src"][0]]
+    vs.append(_mirror_node(nd, vs, spec, dtype))
+  return vs, vs[-1]
+
+def _mirror_node(nd, vs, spec, dtype=np.float64):
+  op = nd["op"]; x = vs[nd["src"][0]]
+  if True:
     if op in UNARY: y = UNARY[op][1](x)
     elif op in BINARY: y = BINARY[op][1](x, vs[nd["src"][1]])
     elif op in REDUCE: y = REDUCE[op][1](x, nd["axis"])
@@ -236,8 +268,7 @@ def _mirror(spec, feed, dtype):
     elif op == "slice": y = x[tuple(slice(b, b + z) for b, z in zip(nd["begin"], nd["size"]))]
     elif op == "tile": y = np.tile(x, nd["reps"])
     else: raise ValueError(f"unknown op {op!r}")
-    vs.append(np.asarray(y, dtype))
-  return vs, vs[-1]
+  return np.asarray(y, dtype)
 
 def build_graph(spec, emi=False):
   """Build the aneforge graph for a spec; `emi` appends an identity chain (muls(1), adds(0),
@@ -293,6 +324,8 @@ def run_case(spec, opts=(0, 1), emi=True):
   budget (float mode); opt=1 answers to the autotuner's accuracy gate (GATE_TOL), because lossy
   gated variants are part of its contract. Returns failure dicts (empty = ok)."""
   feed = _feed_for(spec)
+  if _acc_violation(spec, feed):
+    return []                                      # fp16 accumulation out of range: engine-undefined, not judged
   vs64, ref = _mirror(spec, feed, np.float64)
   scale = max(float(np.abs(v).max()) for v in vs64) + 1e-3
   if spec["mode"] == "float":                      # conditioning screen: fp32 and fp64 mirrors agree?
@@ -302,6 +335,8 @@ def run_case(spec, opts=(0, 1), emi=True):
   def judge(got, opt, prefix=""):
     if got.shape != ref.shape:
       return {"opt": opt, "kind": prefix + "shape", "error": f"{got.shape} != {ref.shape}"}
+    if not np.isfinite(got).all() and np.isfinite(ref).all():
+      return {"opt": opt, "kind": prefix + "non-finite", "error": "output has NaN/inf, reference does not"}
     if opt == 0 and spec["mode"] == "int":
       if not np.array_equal(got, ref):
         n_bad = int((got != ref).sum())

tests/test_espresso_workarounds.py

@@ -63,3 +63,48 @@ def test_softplus_normal_range_unchanged():
   got = _run(x.softplus(), xv)
   ref = np.logaddexp(0.0, xv.astype(np.float64))
   assert np.abs(got - ref).max() < 1e-2
+
+
+# -- round 2: findings from the second fuzz batch -- #
+
+def test_round_large_magnitudes_exact():
+  # native ANE round corrupts |x| >= 1024 (round(1024)=1025, round(-2047)=-2048); the select
+  # routing keeps every fp16 value exact - values >= 1024 are already integers
+  xv = np.array([[1023.0, 1024.0, 1025.0, 2047.0, -1024.0, -2047.0, 3.4, -2.6]], np.float16)
+  x = af.input((1, 8))
+  got = _run(x.round(), xv)
+  want = np.array([[1023.0, 1024.0, 1025.0, 2047.0, -1024.0, -2047.0, 3.0, -3.0]])
+  assert np.array_equal(got, want), got.ravel()
+
+def test_reduce_round_scalar_mul_chain():
+  # rmax -> round -> muls previously returned the bare rmax (both epilogues dropped)
+  iv = np.array([[1, 2, 0, -1, 2, 1, 0, 1], [2, -3, 1, 0, 2, 2, 1, 0], [3, 1, 1, 2, 0, 1, 3, 2]], np.float16)
+  x = af.input((3, 8))
+  got = _run(x.amax((1,)).round() * -3.0, iv)
+  assert np.array_equal(got.ravel(), np.array([-6.0, -6.0, -9.0])), got.ravel()
+
+def test_muls_zero_after_reduce_compiles_and_is_zero():
+  # mul-by-zero after a reduce crashed ANECCompile; the sub(x,x) form compiles and is exact
+  xv = np.random.default_rng(0).integers(-3, 4, size=(4, 8, 5)).astype(np.float16)
+  x = af.input((4, 8, 5))
+  got = _run(x.sum((1,)) * 0.0, xv)
+  assert got.shape == (4, 1, 5) and np.array_equal(got, np.zeros((4, 1, 5)))
+
+def test_empty_program_gets_identity():
+  # a graph whose output is its input lowered to an empty MIL body, which crashes Espresso
+  # ("unordered_map::at: key not found"); the emitter now inserts an explicit identity
+  xv = np.arange(12, dtype=np.float16).reshape(3, 4)
+  x = af.input((3, 4))
+  got = _run(x, xv)
+  assert np.array_equal(got, xv.astype(np.float64))
+
+
+def test_workarounds_cross_compile_for_all_families():
+  # the workaround emissions must not break any target family's compile (measured behavior is
+  # per-family; the workaround FORMS are semantically exact everywhere, so compiling is the bar)
+  from aneforge._compile import cross_compile_check
+  for fam in ("h13", "h14", "h15", "h16", "h16s"):
+    assert cross_compile_check(af.input((1, 8)).round(), fam), f"round select-form: {fam}"
+    assert cross_compile_check(af.input((1, 8)).softplus(), fam), f"stable softplus: {fam}"
+    assert cross_compile_check(af.input((4, 8)).sum((1,)) * 0.0, fam), f"muls-zero sub-form: {fam}"
+    assert cross_compile_check(af.input((3, 4)), fam), f"empty-program guard: {fam}"