Custom Enzyme inactive marker

Author: wsmoses

Diff for: enzyme/Enzyme/ActivityAnalysis.cpp

@@ -143,6 +143,11 @@ const char *KnownInactiveFunctions[] = {"__assert_fail",
 /// This tool can only be used when in DOWN mode
 bool ActivityAnalyzer::isFunctionArgumentConstant(CallInst *CI, Value *val) {
   assert(directions & DOWN);
+
+  if (CI->hasFnAttr("enzyme_inactive")) {
+    return true;
+  }
+
   Function *F = CI->getCalledFunction();
 
   // Indirect function calls may actively use the argument
@@ -1133,11 +1138,15 @@ bool ActivityAnalyzer::isInstructionInactiveFromOrigin(TypeResults &TR,
 
   // Calls to print/assert/cxa guard are definitionally inactive
   if (auto op = dyn_cast<CallInst>(inst)) {
+    if (op->hasFnAttr("enzyme_inactive")) {
+      return true;
+    }
     if (auto called = op->getCalledFunction()) {
       if (called->getName() == "free" || called->getName() == "_ZdlPv" ||
           called->getName() == "_ZdlPvm" || called->getName() == "munmap") {
         return true;
       }
+
       for (auto FuncName : KnownInactiveFunctionsStartingWith) {
         if (called->getName().startswith(FuncName)) {
           return true;

Diff for: enzyme/test/Enzyme/enzyme_inactive.ll

@@ -0,0 +1,28 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
+
+; Function Attrs: noinline nounwind readnone uwtable
+define double @tester(double %x) {
+entry:
+  tail call void @myprint(double %x)
+  ret double %x
+}
+
+define double @test_derivative(double %x) {
+entry:
+  %0 = tail call double (double (double)*, ...) @__enzyme_autodiff(double (double)* nonnull @tester, double %x)
+  ret double %0
+}
+
+declare void @myprint(double %x) #0
+
+; Function Attrs: nounwind
+declare double @__enzyme_autodiff(double (double)*, ...)
+
+attributes #0 = { "enzyme_inactive" }
+
+; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(double %x, double %[[differet:.+]])
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   tail call void @myprint(double %x)
+; CHECK-NEXT:   %0 = insertvalue { double } undef, double %differeturn, 0
+; CHECK-NEXT:   ret { double } %0
+; CHECK-NEXT: }

Diff for: enzyme/test/Enzyme/enzyme_inactive2.ll

@@ -0,0 +1,28 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
+
+; Function Attrs: noinline nounwind readnone uwtable
+define double @tester(double %x) {
+entry:
+  tail call void @myprint(double %x) #0
+  ret double %x
+}
+
+define double @test_derivative(double %x) {
+entry:
+  %0 = tail call double (double (double)*, ...) @__enzyme_autodiff(double (double)* nonnull @tester, double %x)
+  ret double %0
+}
+
+declare void @myprint(double %x)
+
+; Function Attrs: nounwind
+declare double @__enzyme_autodiff(double (double)*, ...)
+
+attributes #0 = { "enzyme_inactive" }
+
+; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(double %x, double %[[differet:.+]])
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   tail call void @myprint(double %x)
+; CHECK-NEXT:   %0 = insertvalue { double } undef, double %differeturn, 0
+; CHECK-NEXT:   ret { double } %0
+; CHECK-NEXT: }