[AutoDiff] Recompute tanh/exp on the operand in the reverse pass

[duburcqa]

Summary

Previously the reverse-mode formulas for `tanh(x)` and `exp(x)` reused the forward stmt value (`sqr(stmt)`, `stmt`). `BackupSSA` spills that forward value to a single plain alloca, so inside a dynamic loop the reversed iteration reads the last-iteration value for every backward step - giving silently wrong gradients for all earlier iterations.

Recompute `tanh(operand)` / `exp(operand)` instead. The operand rides the adstack through `LocalLoad`, so a fresh call on it is per-iteration correct. Both ops are already in `NonLinearOps::unary_collections` so the operand alloca is promoted.

Test plan

• `test_adstack_tanh_exp_loop_carried[tanh]` and `[exp]` (new, added here): 3-iteration dynamic loop where `a` is redefined from scratch each iteration (read-only alloca pattern). Verifies gradient against `torch.autograd`; the pre-fix formula would produce silently wrong values.

[claude]

🔴 The PR fixes the stale-read bug for tanh and exp by requiring their operand allocas to be promoted to AdStacks (because both are in NonLinearOps::unary_collections), but rsqrt has the identical bug and is absent from unary_collections. In any dynamic loop, rsqrt's backward pass computes rsqrt(stmt->operand) on a plain alloca that BackupSSA overwrites every forward iteration, so the reversed loop reads the last-iteration operand for every backward step and produces silently wrong gradients. Fix: add UnaryOpType::rsqrt to unary_collections alongside sqrt.

Extended reasoning...

What the bug is and how it manifests

NonLinearOps::unary_collections (lines 44-46) lists: abs, sin, cos, tan, tanh, asin, acos, exp, log, sqrt — rsqrt is absent. AdStackAllocaJudger::visit(UnaryOpStmt) (lines 438-444) only sets is_stack_needed_=true when the consuming op is in unary_collections. As a result, when an alloca is exclusively consumed by a rsqrt op inside a loop, it is NOT promoted to an AdStackAllocaStmt. It stays as a plain AllocaStmt.

The specific code path that triggers it

PromoteSSA2LocalVar creates an alloca for the operand value (e.g., a GlobalLoadStmt result). ReplaceLocalVarWithStacks calls AdStackAllocaJudger::run on every AllocaStmt. Because rsqrt is absent from unary_collections, the judger returns false, and the alloca stays plain. BackupSSA then spills the operand with a plain LocalStoreStmt that is overwritten each forward iteration. In the backward pass, MakeAdjoint::visit(UnaryOpStmt) for rsqrt (lines 1298-1299) emits:

accumulate(stmt->operand, mul(adjoint(stmt), mul(constant(-0.5f, stmt->ret_type),
                                                 pow(rsqrt(stmt->operand), constant(3, stmt->ret_type)))));

stmt->operand is a LocalLoadStmt from the plain alloca. BackupSSA's generic_visit skips AllocaStmt operands (condition !op->is<AllocaStmt>()), so the reversed loop reads the plain alloca directly — which holds the last-forward-iteration value after all N forward iterations have completed.

Why existing code does not prevent it

The PR's own description states the fix works because "both ops are already in NonLinearOps::unary_collections so the operand alloca is promoted." This is exactly the precondition rsqrt violates. sqrt (present in unary_collections) works correctly; rsqrt (absent) does not, despite both being in the same visit(UnaryOpStmt) handler and the backward formula for rsqrt also recomputing on the operand.

Impact

For any loop computing y[i] = rsqrt(x[i]), the reverse-mode gradient of x is computed as -0.5 * rsqrt(x[i])^3 * dy[i]. With this bug, rsqrt(x[i]) becomes rsqrt(x[N-1]) for all i != N-1, giving wrong gradients for all but the final iteration. The error is silent — no assertion fires.

Step-by-step proof

1. Forward loop, i=0..N-1: each iteration stores x[i] into the plain alloca. After the loop completes, alloca holds x[N-1].
2. Backward loop, i=N-1..0: each backward step calls rsqrt(stmt->operand). stmt->operand loads the plain alloca, always returning x[N-1].
3. For all i < N-1, the computed gradient is -0.5 * rsqrt(x[N-1])^3 * dy[i] instead of the correct -0.5 * rsqrt(x[i])^3 * dy[i].

How to fix it

Add UnaryOpType::rsqrt to NonLinearOps::unary_collections. This causes AdStackAllocaJudger to promote the operand alloca to an AdStackAllocaStmt, which BackupSSA then handles with per-iteration push/pop semantics — exactly the mechanism the PR already applies to tanh and exp.