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

enzyme/Enzyme/AdjointGenerator.h

@@ -6309,6 +6309,17 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
         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 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::OUT_DIFF) &&
+              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;

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<unsigned, 2> 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<StructType>(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<Type *, 2> 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<ReverseCacheKey, Function *>(

enzyme/test/Enzyme/ReverseMode/custom-const-float.ll

@@ -0,0 +1,58 @@
+; 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:        ret { double }
+; 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: }

enzyme/test/Enzyme/ReverseMode/custom-split-const-float.ll

@@ -0,0 +1,77 @@
+; 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:        ret { double }
+; 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: }

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