Fix: avoid pruning Const arguments needed by a custom derivative rule

enzyme/Enzyme/AdjointGenerator.h

@@ -5230,6 +5230,20 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
         auto argTy =
             gutils->getDiffeType(call.getArgOperand(i), foreignFunction);
 
+        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(),
+                          [](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 +5518,20 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
 
       auto argTy = gutils->getDiffeType(call.getArgOperand(i), foreignFunction);
 
+      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(),
+                           [](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)) {

enzyme/Enzyme/FunctionUtils.cpp

@@ -3155,6 +3155,29 @@ Function *PreProcessCache::CloneFunctionWithReturns(
       if (EnzymePrintActivity)
         llvm::errs() << "in new function " << NewF->getName()
                      << " constant arg " << *j << "\n";
+
+      // Delete read/write labels for the `Const`
+      bool usedByUnverifiedCall = false;
+      for (auto *U : i->users()) {
+        auto *CI = dyn_cast<CallInst>(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)

enzyme/Enzyme/Utils.h

@@ -1618,6 +1618,12 @@ static inline bool shouldDisableNoWrite(const llvm::CallInst *CI) {
   if (funcName == "MPI_Wait" || funcName == "MPI_Waitall") {
     return true;
   }
+
+  // Get rid of readnone/writeonly for `Const`
+  if (!F->isDeclaration()) {
+    return true;
+  }
+
   return false;
 }
 

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: }

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: }