diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h index cbd80ae4d47..bbc6375660c 100644 --- a/enzyme/Enzyme/AdjointGenerator.h +++ b/enzyme/Enzyme/AdjointGenerator.h @@ -6309,6 +6309,17 @@ 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 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; 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..0f2aea42363 --- /dev/null +++ b/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: } 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..a2d9bfedc7e --- /dev/null +++ b/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: } 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: }