Keep massaged constant args of a custom rule constant in split reverse

[vchuravy]

Requested by Valentin · project thread

What

A function with a custom rule (enzyme_augment + enzyme_gradient) is called with a constant pointer argument, from a caller that Enzyme differentiates in split mode (augmented forward pass + reverse pass). This hits an assertion in recursivelyHandleSubfunction:

Assertion `argsInverted[i] == DIFFE_TYPE::DUP_NONEED' failed.

Why

CreateAugmentedPrimal massages the custom augmented forward pass to take the constant argument. The fixaugment_ wrapper has the signature of the call (no shadow for the constant argument) and passes the primal as the shadow. But the AugmentedReturn records the activity that it gives to the custom rule (DUP_ARG), not the activity of the call (CONSTANT).

In the reverse pass of the caller (ReverseModeGradient), the activities of the call are compared with subdata->constant_args. The loop only expects the difference DUP_NONEED (reverse) vs DUP_ARG (augmented), so CONSTANT vs DUP_ARG fails the assertion.

CreatePrimalAndGradient massages the custom reverse pass from the CONSTANT activity of the call in the same way. Its fixgradient_ wrapper has the signature of the call, and it uses the augmented data with the massaged activities internally. So the correct activity for the reverse call is CONSTANT. The fix keeps the argument constant when the call activity is CONSTANT, the augmented activity is DUP_ARG, and the callee has enzyme_gradient. The fix does not only remove the assertion: changing the activity to DUP_ARG (as the code does for DUP_NONEED) would request a non-massaged reverse pass with a shadow argument that the caller does not have.

In combined mode there is no comparison, so the bug only shows when the caller is in split mode, for example when its result is used later in a differentiated function.

Found while differentiating MITgcm global_ocean.cs32x15 with LLVM flang + full LTO.

Constant floating-point arguments (second commit)

A constant floating-point argument of a custom rule is massaged to OUT_DIFF, so the custom reverse pass also returns the adjoint of that argument. The fixgradient_ wrapper returned that struct as is, so the caller got one adjoint too many:

• combined mode: getNumElements() == structidx fails (without assertions, the verifier reports Function return type does not match operand type of return inst);
• split mode: the augmented call records OUT_DIFF, so subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG fails.

The fixgradient_ wrapper now returns only the adjoints of the arguments that are OUT_DIFF in the call, and the split reverse pass keeps a massaged OUT_DIFF argument constant, like a massaged DUP_ARG pointer.

Test

• ReverseMode/custom-split-const.ll: @scale has a custom augmented forward pass and reverse pass with a shadow for a pointer argument. @caller passes a constant pointer to it. @top uses the result of @caller, so @caller is differentiated in split mode. The test checks that @diffecaller calls @fixgradient_scale with only the primal pointer. LLVM 16 build with assertions: fails with the assertion above without the fix, passes with it.
• ReverseMode/custom-const-float.ll (combined) and ReverseMode/custom-split-const-float.ll (split): a custom rule with a constant double argument; the fixgradient_ wrapper must return only the adjoint of the active argument. Opaque pointers, LLVM >= 15. LLVM 18 with Enzyme assertions: both assert without the second commit and pass with it. ReverseMode, ForwardMode, TypeAnalysis and ActivityAnalysis lit results are otherwise unchanged against the base.

Related

• Fix: avoid pruning Const arguments needed by a custom derivative rule #3224 (avoid pruning Const arguments needed by a custom derivative rule) is a different problem: it removes readnone/readonly/writeonly and the no-write shortcut for constant arguments of resolved callees. It does not touch the activity comparison in the split reverse pass, and does not fix this assertion.

Not fixed here

• When an augmented forward pass returns its tape directly (tapeIdx == -1) and the tape type is empty ({}), the tape->getType()->isEmptyTy() branch in recursivelyHandleSubfunction erases the augmented call itself, so the primal side effects of the callee are lost. The tests avoid that path. This is handled separately.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Ld2d8Gu8C8e7rT5maHr2UZ

[vchuravy]

Fortran Debug ubuntu-22.04 ifx 2023.2.4 llvm 17 failed during setup, before Enzyme was built: apt-get update fails because http://apt.llvm.org/jammy llvm-toolchain-jammy-17 "no longer has a Release file", so /opt/intel/oneapi/setvars.sh is never installed. This is an infrastructure problem in the job's apt sources, not something this PR changes, and a re-run would hit the same missing repository. I'm leaving it as is.

---

Generated by Claude Code