From e6be4acffce7dd5f363d856e9b0be831cdc889b2 Mon Sep 17 00:00:00 2001 From: Tunan Date: Sun, 13 Sep 2026 12:15:03 -0400 Subject: [PATCH 1/2] Do not trust write-only/no-capture on Const args of resolved callees A Const argument that a custom EnzymeRules derivative rule reads (or writes) but the primal never touches was being dropped from a differentiated function's signature and replaced with jl_nothing, causing a segfault, whenever the call crossed a real (non-inlined) function boundary. The write-only/no-capture attributes trusted by shouldDisableNoWrite are inferred from the callee's primal semantics alone and cannot see that a custom rule reachable from within an otherwise-resolved callee may need the argument's real value. shouldDisableNoWrite (Utils.h) now also disables the no-write shortcut for any resolved callee with a body Enzyme will still analyze (not just truly opaque external declarations), covering DifferentialUseAnalysis.cpp, EnzymeLogic.cpp, and AdjointGenerator.h, which all share this predicate. Since the argument is now known to potentially be read/written despite carrying stale readnone/readonly/writeonly attributes copied from the primal, also drop those specific attributes (not enzyme_inactive or nocapture) wherever they get copied onto a Const argument: in FunctionUtils.cpp's CloneFunctionWithReturns, and in the two call-site attribute-copy loops in AdjointGenerator.h. Reported in EnzymeAD/Enzyme.jl#3570. --- enzyme/Enzyme/AdjointGenerator.h | 39 ++++++++++++++++++++++++++++++++ enzyme/Enzyme/FunctionUtils.cpp | 15 ++++++++++++ enzyme/Enzyme/Utils.h | 17 ++++++++++++++ 3 files changed, 71 insertions(+) diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index cb1ddef8b82..057980a659c 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -5230,6 +5230,27 @@ class AdjointGenerator : public llvm::InstVisitor { auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction); + if (argTy == DIFFE_TYPE::CONSTANT) { + // The ReadNone/ReadOnly/WriteOnly entries just copied into + // structAttrs above (from PrimalParamAttrsToPreserve) describe the + // *original* call's primal behavior only. Differentiating this + // call may route through a custom derivative rule that reads (or + // writes) this constant argument despite the primal never doing + // so, so those copied attributes are not known to still hold on + // the newly generated call -- drop them here rather than carry an + // unverified claim onto it. + auto &attrs = structAttrs[args.size()]; + attrs.erase(std::remove_if( + attrs.begin(), attrs.end(), + [](const Attribute &A) { + return A.isEnumAttribute() && + (A.getKindAsEnum() == Attribute::ReadNone || + A.getKindAsEnum() == Attribute::ReadOnly || + A.getKindAsEnum() == Attribute::WriteOnly); + }), + attrs.end()); + } + bool replace = (argTy == DIFFE_TYPE::DUP_NONEED && (writeOnlyNoCapture || @@ -5504,6 +5525,24 @@ class AdjointGenerator : public llvm::InstVisitor { auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction); + if (argTy == DIFFE_TYPE::CONSTANT) { + // See the matching comment in the forward-mode call handling above: + // these ReadNone/ReadOnly/WriteOnly entries were just copied from + // the original call's primal attributes and are not known to still + // hold once this call is routed through differentiation (e.g. a + // custom rule reading/writing this constant argument). + auto &attrs = structAttrs[pre_args.size()]; + attrs.erase( + std::remove_if(attrs.begin(), attrs.end(), + [](const Attribute &A) { + return A.isEnumAttribute() && + (A.getKindAsEnum() == Attribute::ReadNone || + A.getKindAsEnum() == Attribute::ReadOnly || + A.getKindAsEnum() == Attribute::WriteOnly); + }), + attrs.end()); + } + bool writeOnlyNoCapture = true; bool readNoneNoCapture = false; if (!isNoCapture(&call, i)) { diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp index 0746cd8711d..1a18a23bb82 100644 --- a/enzyme/Enzyme/FunctionUtils.cpp +++ b/enzyme/Enzyme/FunctionUtils.cpp @@ -3155,6 +3155,21 @@ Function *PreProcessCache::CloneFunctionWithReturns( if (EnzymePrintActivity) llvm::errs() << "in new function " << NewF->getName() << " constant arg " << *j << "\n"; + + // Read/write-effect attributes on this parameter (e.g. `readnone`) + // were inferred from the *original* function's primal body alone. + // Differentiating this function can introduce genuine reads/writes + // of a `CONSTANT` argument that the primal never had (e.g. a custom + // derivative rule reachable from within this body reading it), which + // those attributes know nothing about. Since we cannot yet tell + // whether that will happen for this particular argument, drop the + // now-unverified claim rather than carry a possibly-false promise + // onto the differentiated function; `enzyme_inactive`/`nocapture` + // are untouched, only what could allow the optimizer to elide or + // reorder an actual memory access is. + NewF->removeParamAttr(jj, Attribute::ReadNone); + NewF->removeParamAttr(jj, Attribute::ReadOnly); + NewF->removeParamAttr(jj, Attribute::WriteOnly); } else { nonconstant.insert(i); if (EnzymePrintActivity) diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index 5bfb4df4ab1..dd846c1b3b8 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -1618,6 +1618,23 @@ static inline bool shouldDisableNoWrite(const llvm::CallInst *CI) { if (funcName == "MPI_Wait" || funcName == "MPI_Waitall") { return true; } + + // At this point `F` is resolved (the `!F` case above already returned + // true) and carries none of the existing custom-derivative markers. If + // it still has a body Enzyme can (and will) look into -- an ordinary + // resolved function, including one with `available_externally` linkage + // -- do not trust write-only/no-capture attribute inference on this call + // to prune an argument's primal value. Such attributes are typically + // computed from the callee's own primal semantics alone and cannot see + // that a call reachable from within this body may carry a custom + // derivative rule (e.g. an EnzymeRules-style rule) that reads the + // argument at differentiation time despite the primal never doing so. + // Only calls to truly opaque functions (declarations with no body to + // look into at all) are safe to trust here. + if (!F->isDeclaration()) { + return true; + } + return false; } From 95a7a168f83c761e539a295dbdb96ec8bb04bc34 Mon Sep 17 00:00:00 2001 From: Tunan Date: Mon, 14 Sep 2026 23:32:57 -0500 Subject: [PATCH 2/2] remove_read_write_under_conditions --- enzyme/Enzyme/AdjointGenerator.h | 19 +-- enzyme/Enzyme/FunctionUtils.cpp | 36 +++-- enzyme/Enzyme/Utils.h | 13 +- .../Enzyme/ForwardMode/callee_const_arg.ll | 39 +++++ .../Enzyme/ReverseMode/custom_const_state.ll | 138 ++++++++++++++++++ 5 files changed, 204 insertions(+), 41 deletions(-) create mode 100644 enzyme/test/Enzyme/ForwardMode/callee_const_arg.ll create mode 100644 enzyme/test/Enzyme/ReverseMode/custom_const_state.ll diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index 057980a659c..6b0f8edfaa8 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -5230,15 +5230,8 @@ class AdjointGenerator : public llvm::InstVisitor { auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction); - if (argTy == DIFFE_TYPE::CONSTANT) { - // The ReadNone/ReadOnly/WriteOnly entries just copied into - // structAttrs above (from PrimalParamAttrsToPreserve) describe the - // *original* call's primal behavior only. Differentiating this - // call may route through a custom derivative rule that reads (or - // writes) this constant argument despite the primal never doing - // so, so those copied attributes are not known to still hold on - // the newly generated call -- drop them here rather than carry an - // unverified claim onto it. + if (argTy == DIFFE_TYPE::CONSTANT && shouldDisableNoWrite(&call)) { + // Delete the read/write labels for `CONSTANT`. auto &attrs = structAttrs[args.size()]; attrs.erase(std::remove_if( attrs.begin(), attrs.end(), @@ -5525,12 +5518,8 @@ class AdjointGenerator : public llvm::InstVisitor { auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction); - if (argTy == DIFFE_TYPE::CONSTANT) { - // See the matching comment in the forward-mode call handling above: - // these ReadNone/ReadOnly/WriteOnly entries were just copied from - // the original call's primal attributes and are not known to still - // hold once this call is routed through differentiation (e.g. a - // custom rule reading/writing this constant argument). + if (argTy == DIFFE_TYPE::CONSTANT && shouldDisableNoWrite(&call)) { + // Delete the read/write labels for `CONSTANT` auto &attrs = structAttrs[pre_args.size()]; attrs.erase( std::remove_if(attrs.begin(), attrs.end(), diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp index 1a18a23bb82..1d69ae904b8 100644 --- a/enzyme/Enzyme/FunctionUtils.cpp +++ b/enzyme/Enzyme/FunctionUtils.cpp @@ -3156,20 +3156,28 @@ Function *PreProcessCache::CloneFunctionWithReturns( llvm::errs() << "in new function " << NewF->getName() << " constant arg " << *j << "\n"; - // Read/write-effect attributes on this parameter (e.g. `readnone`) - // were inferred from the *original* function's primal body alone. - // Differentiating this function can introduce genuine reads/writes - // of a `CONSTANT` argument that the primal never had (e.g. a custom - // derivative rule reachable from within this body reading it), which - // those attributes know nothing about. Since we cannot yet tell - // whether that will happen for this particular argument, drop the - // now-unverified claim rather than carry a possibly-false promise - // onto the differentiated function; `enzyme_inactive`/`nocapture` - // are untouched, only what could allow the optimizer to elide or - // reorder an actual memory access is. - NewF->removeParamAttr(jj, Attribute::ReadNone); - NewF->removeParamAttr(jj, Attribute::ReadOnly); - NewF->removeParamAttr(jj, Attribute::WriteOnly); + // Delete read/write labels for the `Const` + bool usedByUnverifiedCall = false; + for (auto *U : i->users()) { + auto *CI = dyn_cast(U); + if (!CI) + continue; + bool isArgOperand = false; + for (auto &Use : CI->args()) + if (Use.get() == &*i) { + isArgOperand = true; + break; + } + if (isArgOperand && shouldDisableNoWrite(CI)) { + usedByUnverifiedCall = true; + break; + } + } + if (usedByUnverifiedCall) { + NewF->removeParamAttr(jj, Attribute::ReadNone); + NewF->removeParamAttr(jj, Attribute::ReadOnly); + NewF->removeParamAttr(jj, Attribute::WriteOnly); + } } else { nonconstant.insert(i); if (EnzymePrintActivity) diff --git a/enzyme/Enzyme/Utils.h b/enzyme/Enzyme/Utils.h index dd846c1b3b8..4eb1b950cc6 100644 --- a/enzyme/Enzyme/Utils.h +++ b/enzyme/Enzyme/Utils.h @@ -1619,18 +1619,7 @@ static inline bool shouldDisableNoWrite(const llvm::CallInst *CI) { return true; } - // At this point `F` is resolved (the `!F` case above already returned - // true) and carries none of the existing custom-derivative markers. If - // it still has a body Enzyme can (and will) look into -- an ordinary - // resolved function, including one with `available_externally` linkage - // -- do not trust write-only/no-capture attribute inference on this call - // to prune an argument's primal value. Such attributes are typically - // computed from the callee's own primal semantics alone and cannot see - // that a call reachable from within this body may carry a custom - // derivative rule (e.g. an EnzymeRules-style rule) that reads the - // argument at differentiation time despite the primal never doing so. - // Only calls to truly opaque functions (declarations with no body to - // look into at all) are safe to trust here. + // Get rid of readnone/writeonly for `Const` if (!F->isDeclaration()) { return true; } diff --git a/enzyme/test/Enzyme/ForwardMode/callee_const_arg.ll b/enzyme/test/Enzyme/ForwardMode/callee_const_arg.ll new file mode 100644 index 00000000000..1af2ab27366 --- /dev/null +++ b/enzyme/test/Enzyme/ForwardMode/callee_const_arg.ll @@ -0,0 +1,39 @@ +; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s; fi +; RUN: %opt < %s %newLoadEnzyme -passes="enzyme" -enzyme-preopt=false -enzyme-detect-readthrow=0 -S | FileCheck %s + +; @inner has no custom rule and never touches %a0, so @outer's readnone on +; the call is accurate -- this does not reproduce the original bug by itself +; (that needs a custom rule whose behavior diverges from the primal; see +; ReverseMode/custom_const_state.ll). It tests shouldDisableNoWrite's +; strategy: once a callee is resolved and has a body Enzyme will analyze, it +; cannot tell an accurate readnone from one hiding a custom rule, so it never +; lets either prune a Const pointer via write-only/no-capture attributes. +; Before the fix, this accurate readnone alone let Enzyme drop %a1. See +; EnzymeAD/Enzyme.jl#3570. + +declare double @__enzyme_fwddiff(...) + +define double @dsquare(double %x) { +entry: + %0 = tail call double (...) @__enzyme_fwddiff(ptr @outer, metadata !"enzyme_dup", double %x, double 1.0, metadata !"enzyme_const", ptr null) + ret double %0 +} + +define dso_local double @inner(ptr nocapture %a0, double %a1) { +top: + %a3 = fmul double %a1, %a1 + ret double %a3 +} + +; Function Attrs: nosync readnone +define double @outer(double %a0, ptr %a1) { +entry: + %a7 = call double @inner(ptr %a1, double %a0) readnone + ret double %a7 +} + +; CHECK: define internal double @fwddiffeouter(double %a0, double %"a0'", ptr %a1) +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call fast double @fwddiffeinner(ptr %a1, double %a0, double %"a0'") +; CHECK-NEXT: ret double %0 +; CHECK-NEXT: } diff --git a/enzyme/test/Enzyme/ReverseMode/custom_const_state.ll b/enzyme/test/Enzyme/ReverseMode/custom_const_state.ll new file mode 100644 index 00000000000..c2085254f6f --- /dev/null +++ b/enzyme/test/Enzyme/ReverseMode/custom_const_state.ll @@ -0,0 +1,138 @@ +; 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 + +; Verifies the combined reverse-mode custom-rule path +; (!enzyme_augment/!enzyme_gradient): a Const `ptr %state` argument the +; custom rule reads and writes, but the primal never touches, reaches the +; generated augment/gradient call as a real pointer, and the generated +; derivative definition/call site don't keep a readnone/readonly/writeonly +; claim that no longer holds once the custom rule's memory access counts. +; +; %state carries readnone/readonly/writeonly as a *parameter* attribute, on +; both @outer_*'s definition and the argument at its call to @atom -- not as +; a whole-call function attribute. The fix reads parameter attributes off +; the function being cloned (FunctionUtils.cpp) and off the specific call +; (AdjointGenerator.h); asserting them directly exercises both copy sites +; without depending on some earlier pass inferring one from the other. +; +; @atom is a pure `declare`, so asserting these on a call to it doesn't +; contradict any body -- unlike a defined callee (see +; ForwardMode/callee_const_arg.ll, which covers that case). The three +; variants (_rn/_ro/_wo) each stimulate one of the three attributes the fix +; removes. `nocapture` is asserted alongside each and must survive on the +; generated derivative, since the fix never touches it. +; +; Does not exercise tape encoding/data flow -- the fix only touches which +; ReadNone/ReadOnly/WriteOnly attributes survive on a Const argument, not +; how augment's tape reaches gradient, so an empty `{}` tape is sufficient. +; Does not cover the split augment/reverse API. See EnzymeAD/Enzyme.jl#3570. + +define internal { {}, double } @augment_atom(double %x, ptr nocapture %state, ptr nocapture %dstate) { +entry: + %old = load double, ptr %state + %new = fadd fast double %old, 1.000000e+00 + store double %new, ptr %state + %struct1 = insertvalue { {}, double } undef, double %new, 1 + ret { {}, double } %struct1 +} + +define internal { double } @gradient_atom(double %x, ptr nocapture %state, ptr nocapture %dstate, double %differet, {} %tapeArg) { +entry: + %struct1 = insertvalue { double } undef, double %differet, 0 + ret { double } %struct1 +} + +declare !enzyme_augment !{ptr @augment_atom} !enzyme_gradient !{ptr @gradient_atom} double @atom(double %x, ptr %state) + +define dso_local double @outer_rn(double %x, ptr nocapture readnone %state) { +entry: + %call = tail call fast double @atom(double %x, ptr nocapture readnone %state) + ret double %call +} + +define dso_local double @outer_ro(double %x, ptr nocapture readonly %state) { +entry: + %call = tail call fast double @atom(double %x, ptr nocapture readonly %state) + ret double %call +} + +define dso_local double @outer_wo(double %x, ptr nocapture writeonly %state) { +entry: + %call = tail call fast double @atom(double %x, ptr nocapture writeonly %state) + ret double %call +} + +define dso_local double @douter_rn(double %x, ptr %state) { +entry: + %0 = tail call double (ptr, ...) @__enzyme_autodiff(ptr nonnull @outer_rn, double %x, metadata !"enzyme_const", ptr %state) + ret double %0 +} + +define dso_local double @douter_ro(double %x, ptr %state) { +entry: + %0 = tail call double (ptr, ...) @__enzyme_autodiff(ptr nonnull @outer_ro, double %x, metadata !"enzyme_const", ptr %state) + ret double %0 +} + +define dso_local double @douter_wo(double %x, ptr %state) { +entry: + %0 = tail call double (ptr, ...) @__enzyme_autodiff(ptr nonnull @outer_wo, double %x, metadata !"enzyme_const", ptr %state) + ret double %0 +} + +declare double @__enzyme_autodiff(ptr, ...) + +; CHECK directives follow the order these definitions appear in the +; generated module (FileCheck scans forward only). Generated call lines end +; in `{{$}}` so a reintroduced readnone/readonly/writeonly on %state can't +; slip in at the end either (LLVM always prints it inline before the +; parameter, never after). +; +; augment_atom/gradient_atom's own `nocapture` on %state/%dstate is real +; (neither stores the pointer) and unaffected by the fix -- confirm it +; survives, unlike ReadNone/ReadOnly/WriteOnly. +; CHECK: define internal { {}, double } @augment_atom(double %x, ptr nocapture %state, ptr nocapture{{.*}} %dstate) +; CHECK: define internal { double } @gradient_atom(double %x, ptr nocapture{{.*}} %state, ptr nocapture{{.*}} %dstate, double %differet, {} %tapeArg) + +; CHECK: define internal { double } @diffeouter_rn(double %x, ptr nocapture %state, double %differeturn) #{{[0-9]+}} { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { double } @fixgradient_atom(double %x, ptr nocapture %state, double %differeturn){{$}} +; CHECK-NEXT: %1 = extractvalue { double } %0, 0 +; CHECK-NEXT: %2 = insertvalue { double } undef, double %1, 0 +; CHECK-NEXT: ret { double } %2 +; CHECK-NEXT: } + +; CHECK: define internal { double } @fixgradient_atom(double %arg0, ptr %arg1, double %arg2) { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { {}, double } @fixaugment_atom(double %arg0, ptr %arg1) +; CHECK-NEXT: %1 = call { double } @fixgradient_atom.1(double %arg0, ptr %arg1, double %arg2, {} undef) +; CHECK-NEXT: ret { double } %1 +; CHECK-NEXT: } + +; CHECK: define internal { {}, double } @fixaugment_atom(double %arg0, ptr %arg1) { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { {}, double } @augment_atom(double %arg0, ptr %arg1, ptr %arg1) +; CHECK-NEXT: ret { {}, double } %0 +; CHECK-NEXT: } + +; CHECK: define internal { double } @fixgradient_atom.1(double %arg0, ptr %arg1, double %postarg0, {} %postarg1) { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { double } @gradient_atom(double %arg0, ptr %arg1, ptr %arg1, double %postarg0, {} %postarg1) +; CHECK-NEXT: ret { double } %0 +; CHECK-NEXT: } + +; CHECK: define internal { double } @diffeouter_ro(double %x, ptr nocapture %state, double %differeturn) #{{[0-9]+}} { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { double } @fixgradient_atom(double %x, ptr nocapture %state, double %differeturn){{$}} +; CHECK-NEXT: %1 = extractvalue { double } %0, 0 +; CHECK-NEXT: %2 = insertvalue { double } undef, double %1, 0 +; CHECK-NEXT: ret { double } %2 +; CHECK-NEXT: } + +; CHECK: define internal { double } @diffeouter_wo(double %x, ptr nocapture %state, double %differeturn) #{{[0-9]+}} { +; CHECK-NEXT: entry: +; CHECK-NEXT: %0 = call { double } @fixgradient_atom(double %x, ptr nocapture %state, double %differeturn){{$}} +; CHECK-NEXT: %1 = extractvalue { double } %0, 0 +; CHECK-NEXT: %2 = insertvalue { double } undef, double %1, 0 +; CHECK-NEXT: ret { double } %2 +; CHECK-NEXT: }