From c2702843cb1aa37eebfa7ba5de5f3f6f0480cf9e Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 1 Oct 2026 17:30:10 +0200 Subject: [PATCH 1/3] Keep massaged constant args of a custom rule constant in split reverse A custom augmented forward pass (enzyme_augment) that is called with a constant pointer argument is massaged in CreateAugmentedPrimal: the fixaugment_ wrapper takes the constant argument and passes it as its own shadow, and the AugmentedReturn records the activity of the custom rule (DUP_ARG), not the activity of the call (CONSTANT). When the caller is differentiated in split mode, its reverse pass compares the activities of the call with the activities recorded in the augmented forward pass. It only expects DUP_NONEED -> DUP_ARG, so the massaged argument hit Assertion `argsInverted[i] == DIFFE_TYPE::DUP_NONEED' failed. CreatePrimalAndGradient massages the custom reverse pass (enzyme_gradient) from the CONSTANT activity of the call in the same way, and the fixgradient_ wrapper has the signature of the call. Keep the argument constant in this case. Assisted-by: Claude Code (Opus 5.5) --- enzyme/Enzyme/AdjointGenerator.h | 9 ++ .../Enzyme/ReverseMode/custom-split-const.ll | 85 +++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 enzyme/test/Enzyme/ReverseMode/custom-split-const.ll diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index cbd80ae4d47..963b10cde8b 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -6309,6 +6309,15 @@ class AdjointGenerator : public llvm::InstVisitor { for (size_t i = 0; i < argsInverted.size(); i++) { if (subdata->constant_args[i] == argsInverted[i]) continue; + // A custom augmented forward pass with a constant pointer argument + // is massaged to take that argument as its own shadow, so the + // augmented call records it as DUP_ARG. The custom reverse pass is + // massaged the same way from the constant argument, so keep it + // constant here. + if (argsInverted[i] == DIFFE_TYPE::CONSTANT && + subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG && + hasMetadata(called, "enzyme_gradient")) + continue; assert(subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG); assert(argsInverted[i] == DIFFE_TYPE::DUP_NONEED); argsInverted[i] = DIFFE_TYPE::DUP_ARG; diff --git a/enzyme/test/Enzyme/ReverseMode/custom-split-const.ll b/enzyme/test/Enzyme/ReverseMode/custom-split-const.ll new file mode 100644 index 00000000000..4204a729751 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/custom-split-const.ll @@ -0,0 +1,85 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme-preopt=false -enzyme -mem2reg -instsimplify -simplifycfg -S | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s + +; @scale has a custom augmented forward pass and a custom reverse pass that +; take a shadow for %p. @caller passes a constant %p, so Enzyme massages both +; custom passes to take a constant %p. @top uses the result of @caller, so +; Enzyme differentiates @caller in split mode (augmented forward pass + +; reverse pass). The reverse pass of @caller must call the massaged custom +; reverse pass with a constant %p. + +define internal { {}, double } @augment_scale(double %x, double* %p, double* %dp) { +entry: + %s = load double, double* %p + %m = fmul double %x, %s + %r = insertvalue { {}, double } undef, double %m, 1 + ret { {}, double } %r +} + +define internal { double } @gradient_scale(double %x, double* %p, double* %dp, double %differet, {} %tape) { +entry: + %s = load double, double* %p + %dx = fmul double %differet, %s + %r = insertvalue { double } undef, double %dx, 0 + ret { double } %r +} + +declare !enzyme_augment !{{ {}, double } (double, double*, double*)* @augment_scale} !enzyme_gradient !{{ double } (double, double*, double*, double, {})* @gradient_scale} double @scale(double %x, double* %p) + +define double @caller(double %x, double* %p) { +entry: + %r = call double @scale(double %x, double* %p) + ret double %r +} + +define void @top(double %x, double* %p, double* %out) { +entry: + %r = call double @caller(double %x, double* %p) + %sq = fmul double %r, %r + store double %sq, double* %out + ret void +} + +define double @dtop(double %x, double* %p, double* %out, double* %dout) { +entry: + %r = call double (...) @__enzyme_autodiff(void (double, double*, double*)* @top, double %x, metadata !"enzyme_const", double* %p, double* %out, double* %dout) + ret double %r +} + +declare double @__enzyme_autodiff(...) + +; CHECK: define internal { double } @diffetop(double %x, double* %p, double* nocapture writeonly %out, double* nocapture %"out'") +; CHECK-NEXT: entry: +; CHECK-NEXT: %r_augmented = call { {}, double } @augmented_caller(double %x, double* %p) +; CHECK-NEXT: %r = extractvalue { {}, double } %r_augmented, 1 +; CHECK-NEXT: %sq = fmul double %r, %r +; CHECK-NEXT: store double %sq, double* %out +; CHECK-NEXT: %[[i0:.+]] = load double, double* %"out'" +; CHECK-NEXT: store double 0.000000e+00, double* %"out'" +; CHECK-NEXT: %[[i1:.+]] = fmul fast double %[[i0]], %r +; CHECK-NEXT: %[[i2:.+]] = fmul fast double %[[i0]], %r +; CHECK-NEXT: %[[i3:.+]] = fadd fast double %[[i1]], %[[i2]] +; CHECK-NEXT: %[[i4:.+]] = call { double } @diffecaller(double %x, double* %p, double %[[i3]], {} undef) +; CHECK-NEXT: ret { double } %[[i4]] +; CHECK-NEXT: } + +; CHECK: define internal { {}, double } @fixaugment_scale(double %arg0, double* %arg1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { {}, double } @augment_scale(double %arg0, double* %arg1, double* %arg1) +; CHECK-NEXT: ret { {}, double } %[[i0]] +; CHECK-NEXT: } + +; CHECK: define internal { {}, double } @augmented_caller(double %x, double* %p) +; CHECK: %r_augmented = call { {}, double } @fixaugment_scale(double %x, double* %p) + +; CHECK: define internal { double } @diffecaller(double %x, double* %p, double %differeturn, {} %tapeArg) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { double } @fixgradient_scale(double %x, double* %p, double %differeturn, {} undef) +; CHECK-NEXT: ret { double } %[[i0]] +; CHECK-NEXT: } + +; CHECK: define internal { double } @fixgradient_scale(double %arg0, double* %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { double } @gradient_scale(double %arg0, double* %arg1, double* %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: ret { double } %[[i0]] +; CHECK-NEXT: } From 768df328d2df96e8f976f45b65809e0219d483d8 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 1 Oct 2026 18:11:52 +0000 Subject: [PATCH 2/3] Drop the adjoint of a massaged constant float arg of a custom rule A constant floating-point argument of a function with 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, but the caller expects no adjoint for a constant argument: - in combined mode, the struct had one element too many (`getNumElements() == structidx` assertion, or an invalid return in the fixgradient_ wrapper without assertions); - in split mode, the augmented call records the argument as OUT_DIFF, which hit `subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG`. Return only the adjoints of the arguments that are OUT_DIFF in the call, and keep a massaged OUT_DIFF argument constant in the split reverse pass, as for a massaged DUP_ARG pointer. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Ld2d8Gu8C8e7rT5maHr2UZ --- enzyme/Enzyme/AdjointGenerator.h | 10 ++- enzyme/Enzyme/EnzymeLogic.cpp | 54 ++++++++++++- .../Enzyme/ReverseMode/custom-const-float.ll | 60 ++++++++++++++ .../ReverseMode/custom-split-const-float.ll | 79 +++++++++++++++++++ 4 files changed, 195 insertions(+), 8 deletions(-) create mode 100644 enzyme/test/Enzyme/ReverseMode/custom-const-float.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index 963b10cde8b..bbc6375660c 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -6309,13 +6309,15 @@ class AdjointGenerator : public llvm::InstVisitor { for (size_t i = 0; i < argsInverted.size(); i++) { if (subdata->constant_args[i] == argsInverted[i]) continue; - // A custom augmented forward pass with a constant pointer argument - // is massaged to take that argument as its own shadow, so the - // augmented call records it as DUP_ARG. The custom reverse pass is + // A custom augmented forward pass with a constant argument is + // massaged to take a constant pointer as its own shadow (DUP_ARG) + // and a constant float as active (OUT_DIFF), and the augmented call + // records the massaged activity. The custom reverse pass is // massaged the same way from the constant argument, so keep it // constant here. if (argsInverted[i] == DIFFE_TYPE::CONSTANT && - subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG && + (subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG || + subdata->constant_args[i] == DIFFE_TYPE::OUT_DIFF) && hasMetadata(called, "enzyme_gradient")) continue; assert(subdata->constant_args[i] == DIFFE_TYPE::DUP_ARG); diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp index 55af7fb75d6..98a4e60dbe6 100644 --- a/enzyme/Enzyme/EnzymeLogic.cpp +++ b/enzyme/Enzyme/EnzymeLogic.cpp @@ -4050,9 +4050,48 @@ Function *EnzymeLogic::CreatePrimalAndGradient( } } - FunctionType *FTy = - FunctionType::get(revfn->getReturnType(), dupargs, - revfn->getFunctionType()->isVarArg()); + // A constant floating-point argument is massaged to OUT_DIFF, so the + // custom reverse pass also returns its adjoint. The caller does not + // expect an adjoint for a constant argument, so only return the + // adjoints of the arguments that are OUT_DIFF in the call. + Type *RetTy = revfn->getReturnType(); + SmallVector keptAdjoints; + bool dropAdjoints = false; + { + unsigned idx = 0; + for (size_t act_idx = 0; act_idx < key.constant_args.size(); + act_idx++) { + if (next_constant_args[act_idx] != DIFFE_TYPE::OUT_DIFF) + continue; + if (key.constant_args[act_idx] == DIFFE_TYPE::OUT_DIFF) + keptAdjoints.push_back(idx); + else + dropAdjoints = true; + idx++; + } + if (dropAdjoints) { + auto ST = dyn_cast(RetTy); + if (!ST || ST->getNumElements() != idx) { + std::string s; + llvm::raw_string_ostream ss(s); + ss << "Unexpected return type of custom reverse pass for function " + << key.todiff->getName() << ": " << *RetTy << ", expected " + << idx << " adjoints\n"; + if (context.req) + ss << " at context: " << *context.req; + EmitNoDerivativeError(ss.str(), key.todiff, context); + dropAdjoints = false; + } else { + SmallVector keptTys; + for (auto kept : keptAdjoints) + keptTys.push_back(ST->getElementType(kept)); + RetTy = StructType::get(ST->getContext(), keptTys); + } + } + } + + FunctionType *FTy = FunctionType::get( + RetTy, dupargs, revfn->getFunctionType()->isVarArg()); Function *NewF = Function::Create( FTy, Function::LinkageTypes::InternalLinkage, "fixgradient_" + key.todiff->getName(), key.todiff->getParent()); @@ -4102,7 +4141,14 @@ Function *EnzymeLogic::CreatePrimalAndGradient( bb.CreateRet(UndefValue::get(NewF->getReturnType())); else if (NewF->getReturnType()->isVoidTy()) bb.CreateRetVoid(); - else + else if (dropAdjoints) { + Value *res = UndefValue::get(NewF->getReturnType()); + for (auto en : llvm::enumerate(keptAdjoints)) + res = bb.CreateInsertValue(res, + bb.CreateExtractValue(cal, {en.value()}), + {(unsigned)en.index()}); + bb.CreateRet(res); + } else bb.CreateRet(cal); return insert_or_assign2( diff --git a/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll b/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll new file mode 100644 index 00000000000..973368d712e --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll @@ -0,0 +1,60 @@ +; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %OPnewLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s; fi + +; @scale has a custom augmented forward pass and a custom reverse pass that +; treat %c as active. @caller passes a constant %c, so Enzyme massages the +; custom reverse pass to treat %c as active (OUT_DIFF). The massaged custom +; reverse pass must not return the adjoint of %c. + +define internal { {}, double } @augment_scale(double %x, double %c) { +entry: + %m = fmul double %x, %c + %r = insertvalue { {}, double } undef, double %m, 1 + ret { {}, double } %r +} + +define internal { double, double } @gradient_scale(double %x, double %c, double %differet, {} %tape) { +entry: + %dx = fmul double %differet, %c + %dc = fmul double %differet, %x + %r0 = insertvalue { double, double } undef, double %dx, 0 + %r1 = insertvalue { double, double } %r0, double %dc, 1 + ret { double, double } %r1 +} + +declare !enzyme_augment !{ptr @augment_scale} !enzyme_gradient !{ptr @gradient_scale} double @scale(double %x, double %c) + +define double @caller(double %x, double %c) { +entry: + %r = call double @scale(double %x, double %c) + %sq = fmul double %r, %r + ret double %sq +} + +define double @dcaller(double %x, double %c) { +entry: + %r = call double (...) @__enzyme_autodiff(ptr @caller, double %x, metadata !"enzyme_const", double %c) + ret double %r +} + +declare double @__enzyme_autodiff(...) + +; CHECK: define internal { double } @diffecaller(double %x, double %c, double %differeturn) +; CHECK-NEXT: entry: +; CHECK-NEXT: %r_augmented = call { {}, double } @fixaugment_scale(double %x, double %c) +; CHECK-NEXT: %r = extractvalue { {}, double } %r_augmented, 1 +; CHECK-NEXT: %[[i0:.+]] = fmul fast double %differeturn, %r +; CHECK-NEXT: %[[i1:.+]] = fmul fast double %differeturn, %r +; CHECK-NEXT: %[[i2:.+]] = fadd fast double %[[i0]], %[[i1]] +; CHECK-NEXT: %[[i3:.+]] = call { double } @fixgradient_scale(double %x, double %c, double %[[i2]], {} undef) +; CHECK-NEXT: %[[i4:.+]] = extractvalue { double } %[[i3]], 0 +; CHECK-NEXT: %[[i5:.+]] = insertvalue { double } undef, double %[[i4]], 0 +; CHECK-NEXT: ret { double } %[[i5]] +; CHECK-NEXT: } + +; CHECK: define internal { double } @fixgradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { double, double } @gradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: %[[i1:.+]] = extractvalue { double, double } %[[i0]], 0 +; CHECK-NEXT: %[[i2:.+]] = insertvalue { double } undef, double %[[i1]], 0 +; CHECK-NEXT: ret { double } %[[i2]] +; CHECK-NEXT: } diff --git a/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll b/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll new file mode 100644 index 00000000000..d552c82e0d4 --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll @@ -0,0 +1,79 @@ +; RUN: if [ %llvmver -ge 15 ]; then %opt < %s %OPnewLoadEnzyme -enzyme-preopt=false -passes="enzyme,function(mem2reg,instsimplify,%simplifycfg)" -S | FileCheck %s; fi + +; @scale has a custom augmented forward pass and a custom reverse pass that +; treat %c as active. @caller passes a constant %c, so Enzyme massages both +; custom passes to treat %c as active (OUT_DIFF). @top uses the result of +; @caller, so Enzyme differentiates @caller in split mode. The reverse pass of +; @caller must keep %c constant, and the massaged custom reverse pass must not +; return the adjoint of %c. + +define internal { {}, double } @augment_scale(double %x, double %c) { +entry: + %m = fmul double %x, %c + %r = insertvalue { {}, double } undef, double %m, 1 + ret { {}, double } %r +} + +define internal { double, double } @gradient_scale(double %x, double %c, double %differet, {} %tape) { +entry: + %dx = fmul double %differet, %c + %dc = fmul double %differet, %x + %r0 = insertvalue { double, double } undef, double %dx, 0 + %r1 = insertvalue { double, double } %r0, double %dc, 1 + ret { double, double } %r1 +} + +declare !enzyme_augment !{ptr @augment_scale} !enzyme_gradient !{ptr @gradient_scale} double @scale(double %x, double %c) + + +define double @caller(double %x, double %c) { +entry: + %r = call double @scale(double %x, double %c) + ret double %r +} + +define void @top(double %x, double %c, ptr %out) { +entry: + %r = call double @caller(double %x, double %c) + %sq = fmul double %r, %r + store double %sq, ptr %out + ret void +} + +define double @dtop(double %x, double %c, ptr %out, ptr %dout) { +entry: + %r = call double (...) @__enzyme_autodiff(ptr @top, double %x, metadata !"enzyme_const", double %c, ptr %out, ptr %dout) + ret double %r +} + +declare double @__enzyme_autodiff(...) + +; CHECK: define internal { double } @diffetop(double %x, double %c, ptr nocapture writeonly %out, ptr nocapture %"out'") +; CHECK-NEXT: entry: +; CHECK-NEXT: %r_augmented = call { {}, double } @augmented_caller(double %x, double %c) +; CHECK: %{{.+}} = call { double } @diffecaller(double %x, double %c, double %{{.+}}, {} undef) + +; CHECK: define internal { {}, double } @fixaugment_scale(double %arg0, double %arg1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { {}, double } @augment_scale(double %arg0, double %arg1) +; CHECK-NEXT: ret { {}, double } %[[i0]] +; CHECK-NEXT: } + +; CHECK: define internal { {}, double } @augmented_caller(double %x, double %c) +; CHECK: %r_augmented = call { {}, double } @fixaugment_scale(double %x, double %c) + +; CHECK: define internal { double } @diffecaller(double %x, double %c, double %differeturn, {} %tapeArg) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { double } @fixgradient_scale(double %x, double %c, double %differeturn, {} undef) +; CHECK-NEXT: %[[i1:.+]] = extractvalue { double } %[[i0]], 0 +; CHECK-NEXT: %[[i2:.+]] = insertvalue { double } undef, double %[[i1]], 0 +; CHECK-NEXT: ret { double } %[[i2]] +; CHECK-NEXT: } + +; CHECK: define internal { double } @fixgradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %[[i0:.+]] = call { double, double } @gradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: %[[i1:.+]] = extractvalue { double, double } %[[i0]], 0 +; CHECK-NEXT: %[[i2:.+]] = insertvalue { double } undef, double %[[i1]], 0 +; CHECK-NEXT: ret { double } %[[i2]] +; CHECK-NEXT: } From b6addb5bd6ffffb3ae549ec6fcac35170836d43d Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 1 Oct 2026 18:25:09 +0000 Subject: [PATCH 3/3] Accept either return form of diffecaller in the const-float tests LLVM 16 folds the extractvalue/insertvalue pair after the fixgradient_ call, newer LLVM keeps it. Only check the call and the return. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Ld2d8Gu8C8e7rT5maHr2UZ --- enzyme/test/Enzyme/ReverseMode/custom-const-float.ll | 4 +--- enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll b/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll index 973368d712e..0f2aea42363 100644 --- a/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll +++ b/enzyme/test/Enzyme/ReverseMode/custom-const-float.ll @@ -46,9 +46,7 @@ declare double @__enzyme_autodiff(...) ; CHECK-NEXT: %[[i1:.+]] = fmul fast double %differeturn, %r ; CHECK-NEXT: %[[i2:.+]] = fadd fast double %[[i0]], %[[i1]] ; CHECK-NEXT: %[[i3:.+]] = call { double } @fixgradient_scale(double %x, double %c, double %[[i2]], {} undef) -; CHECK-NEXT: %[[i4:.+]] = extractvalue { double } %[[i3]], 0 -; CHECK-NEXT: %[[i5:.+]] = insertvalue { double } undef, double %[[i4]], 0 -; CHECK-NEXT: ret { double } %[[i5]] +; CHECK: ret { double } ; CHECK-NEXT: } ; CHECK: define internal { double } @fixgradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1) diff --git a/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll b/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll index d552c82e0d4..a2d9bfedc7e 100644 --- a/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll +++ b/enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll @@ -65,9 +65,7 @@ declare double @__enzyme_autodiff(...) ; CHECK: define internal { double } @diffecaller(double %x, double %c, double %differeturn, {} %tapeArg) ; CHECK-NEXT: entry: ; CHECK-NEXT: %[[i0:.+]] = call { double } @fixgradient_scale(double %x, double %c, double %differeturn, {} undef) -; CHECK-NEXT: %[[i1:.+]] = extractvalue { double } %[[i0]], 0 -; CHECK-NEXT: %[[i2:.+]] = insertvalue { double } undef, double %[[i1]], 0 -; CHECK-NEXT: ret { double } %[[i2]] +; CHECK: ret { double } ; CHECK-NEXT: } ; CHECK: define internal { double } @fixgradient_scale(double %arg0, double %arg1, double %postarg0, {} %postarg1)