Forward mode load store (#193)

* added getForwardBuilder

* storeInst forward mode

* loadInst forward mod

* moved forward and reverse builder helpers into GradientUtils

Author: tgymnich

enzyme/Enzyme/AdjointGenerator.h

@@ -269,10 +269,15 @@ class AdjointGenerator
         IRBuilder<> BuilderZ(newi);
         Value *newip = nullptr;
 
-        bool needShadow = is_value_needed_in_reverse<ValueType::ShadowPtr>(
-            TR, gutils, &I,
-            /*toplevel*/ Mode == DerivativeMode::ReverseModeCombined,
-            oldUnreachable);
+        // TODO: In the case of fwd mode this should be true if the loaded value
+        // itself is used as a pointer.
+        bool needShadow =
+            Mode == DerivativeMode::ForwardMode
+                ? false
+                : is_value_needed_in_reverse<ValueType::ShadowPtr>(
+                      TR, gutils, &I,
+                      /*toplevel*/ Mode == DerivativeMode::ReverseModeCombined,
+                      oldUnreachable);
 
         switch (Mode) {
 
@@ -291,8 +296,8 @@ class AdjointGenerator
           gutils->invertedPointers[&I] = newip;
           break;
         }
-
-        case DerivativeMode::ReverseModeGradient: {
+        case DerivativeMode::ReverseModeGradient:
+        case DerivativeMode::ForwardMode: {
           // only make shadow where caching needed
           if (can_modref && needShadow) {
             newip = gutils->cacheForReverse(BuilderZ, placeholder,
@@ -322,13 +327,18 @@ class AdjointGenerator
 
     Value *inst = newi;
 
+    // TODO: In the case of fwd mode this should be true if the loaded value
+    // itself is used as a pointer.
+    bool primalNeededInReverse =
+        Mode == DerivativeMode::ForwardMode
+            ? false
+            : is_value_needed_in_reverse<ValueType::Primal>(
+                  TR, gutils, &I,
+                  /*toplevel*/ Mode == DerivativeMode::ReverseModeCombined,
+                  oldUnreachable);
     //! Store loads that need to be cached for use in reverse pass
     if (cache_reads_always ||
-        (!cache_reads_never && can_modref &&
-         is_value_needed_in_reverse<ValueType::Primal>(
-             TR, gutils, &I,
-             /*toplevel*/ Mode == DerivativeMode::ReverseModeCombined,
-             oldUnreachable))) {
+        (!cache_reads_never && can_modref && primalNeededInReverse)) {
       if (!gutils->unnecessaryIntermediates.count(&I)) {
         IRBuilder<> BuilderZ(gutils->getNewFromOriginal(&I)->getNextNode());
         // auto tbaa = inst->getMetadata(LLVMContext::MD_tbaa);
@@ -379,15 +389,36 @@ class AdjointGenerator
     }
 
     if (isfloat) {
-      IRBuilder<> Builder2(parent);
-      getReverseBuilder(Builder2);
-      auto prediff = diffe(&I, Builder2);
-      setDiffe(&I, Constant::getNullValue(type), Builder2);
 
-      if (!gutils->isConstantValue(I.getOperand(0))) {
-        ((DiffeGradientUtils *)gutils)
-            ->addToInvertedPtrDiffe(I.getOperand(0), prediff, Builder2,
-                                    alignment, OrigOffset);
+      switch (Mode) {
+      case DerivativeMode::ForwardMode: {
+        IRBuilder<> Builder2(&I);
+        getForwardBuilder(Builder2);
+
+        if (!gutils->isConstantValue(&I)) {
+          auto diff = Builder2.CreateLoad(
+              gutils->invertPointerM(I.getOperand(0), Builder2));
+          setDiffe(&I, diff, Builder2);
+        }
+        break;
+      }
+      case DerivativeMode::ReverseModeGradient:
+      case DerivativeMode::ReverseModeCombined: {
+        IRBuilder<> Builder2(parent);
+        getReverseBuilder(Builder2);
+
+        auto prediff = diffe(&I, Builder2);
+        setDiffe(&I, Constant::getNullValue(type), Builder2);
+
+        if (!gutils->isConstantValue(I.getOperand(0))) {
+          ((DiffeGradientUtils *)gutils)
+              ->addToInvertedPtrDiffe(I.getOperand(0), prediff, Builder2,
+                                      alignment, OrigOffset);
+        }
+        break;
+      }
+      case DerivativeMode::ReverseModePrimal:
+        break;
       }
     }
   }
@@ -494,8 +525,9 @@ class AdjointGenerator
 
     if (FT) {
       //! Only need to update the reverse function
-      if (Mode == DerivativeMode::ReverseModeGradient ||
-          Mode == DerivativeMode::ReverseModeCombined) {
+      switch (Mode) {
+      case DerivativeMode::ReverseModeGradient:
+      case DerivativeMode::ReverseModeCombined: {
         IRBuilder<> Builder2(SI.getParent());
         getReverseBuilder(Builder2);
 
@@ -512,13 +544,29 @@ class AdjointGenerator
           ts = setPtrDiffe(orig_ptr, Constant::getNullValue(valType), Builder2);
           addToDiffe(orig_val, dif1, Builder2, FT);
         }
+        break;
+      }
+      case DerivativeMode::ForwardMode: {
+        IRBuilder<> Builder2(&SI);
+        getForwardBuilder(Builder2);
+
+        if (constantval) {
+          ts = setPtrDiffe(orig_ptr, Constant::getNullValue(valType), Builder2);
+        } else {
+          auto diff = diffe(orig_val, Builder2);
+
+          ts = setPtrDiffe(orig_ptr, diff, Builder2);
+        }
+        break;
+      }
       }
 
       //! Storing an integer or pointer
     } else {
       //! Only need to update the forward function
       if (Mode == DerivativeMode::ReverseModePrimal ||
-          Mode == DerivativeMode::ReverseModeCombined) {
+          Mode == DerivativeMode::ReverseModeCombined ||
+          Mode == DerivativeMode::ForwardMode) {
         IRBuilder<> storeBuilder(gutils->getNewFromOriginal(&SI));
 
         Value *valueop = nullptr;
@@ -935,25 +983,12 @@ class AdjointGenerator
     setDiffe(&IVI, Constant::getNullValue(IVI.getType()), Builder2);
   }
 
-  inline void getReverseBuilder(IRBuilder<> &Builder2, bool original = true) {
-    BasicBlock *BB = Builder2.GetInsertBlock();
-    if (original)
-      BB = gutils->getNewFromOriginal(BB);
-    BasicBlock *BB2 = gutils->reverseBlocks[BB].back();
-    if (!BB2) {
-      llvm::errs() << "oldFunc: " << *gutils->oldFunc << "\n";
-      llvm::errs() << "newFunc: " << *gutils->newFunc << "\n";
-      llvm::errs() << "could not invert " << *BB;
-    }
-    assert(BB2);
-
-    if (BB2->getTerminator())
-      Builder2.SetInsertPoint(BB2->getTerminator());
-    else
-      Builder2.SetInsertPoint(BB2);
-    Builder2.SetCurrentDebugLocation(
-        gutils->getNewFromOriginal(Builder2.getCurrentDebugLocation()));
-    Builder2.setFastMathFlags(getFast());
+  void getReverseBuilder(IRBuilder<> &Builder2, bool original = true) {
+    ((GradientUtils *)gutils)->getReverseBuilder(Builder2, original);
+  }
+
+  void getForwardBuilder(IRBuilder<> &Builder2) {
+    ((GradientUtils *)gutils)->getForwardBuilder(Builder2);
   }
 
   Value *diffe(Value *val, IRBuilder<> &Builder) {
@@ -1398,19 +1433,7 @@ class AdjointGenerator
 
   void createBinaryOperatorDual(llvm::BinaryOperator &BO) {
     IRBuilder<> Builder2(&BO);
-
-    Instruction *nBO = gutils->getNewFromOriginal(&BO);
-
-    assert(nBO);
-    assert(nBO->getNextNode());
-
-    if (nBO->getNextNode()) {
-      Builder2.SetInsertPoint(nBO->getNextNode());
-    }
-
-    Builder2.SetCurrentDebugLocation(
-        gutils->getNewFromOriginal(Builder2.getCurrentDebugLocation()));
-    Builder2.setFastMathFlags(getFast());
+    getForwardBuilder(Builder2);
 
     Value *orig_op0 = BO.getOperand(0);
     Value *orig_op1 = BO.getOperand(1);

enzyme/Enzyme/EnzymeLogic.cpp

@@ -2659,16 +2659,26 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
 
   assert(!todiff->empty());
 
-  ReturnType retVal =
-      returnValue ? (dretPtr ? ReturnType::ArgsWithTwoReturns
-                             : ReturnType::ArgsWithReturn)
-                  : (dretPtr ? ReturnType::ArgsWithReturn : ReturnType::Args);
+  ReturnType retVal;
+  if (fwdMode) {
+    auto TR = TA.analyzeFunction(oldTypeInfo);
+    bool retActive = TR.getReturnAnalysis().Inner0().isFloat();
+
+    retVal = returnValue
+                 ? (retActive ? ReturnType::TwoReturns : ReturnType::Return)
+                 : (retActive ? ReturnType::Return : ReturnType::Void);
+  } else {
+    retVal = returnValue
+                 ? (dretPtr ? ReturnType::ArgsWithTwoReturns
+                            : ReturnType::ArgsWithReturn)
+                 : (dretPtr ? ReturnType::ArgsWithReturn : ReturnType::Args);
+  }
 
   bool diffeReturnArg = fwdMode ? false : retType == DIFFE_TYPE::OUT_DIFF;
 
   DiffeGradientUtils *gutils = DiffeGradientUtils::CreateFromClone(
       *this, topLevel, todiff, TLI, TA, retType, diffeReturnArg, constant_args,
-      fwdMode ? ReturnType::Return : retVal, additionalArg);
+      retVal, additionalArg);
 
   if (omp)
     gutils->setupOMPFor();

enzyme/Enzyme/GradientUtils.h

@@ -957,11 +957,10 @@ class GradientUtils : public CacheUtility {
         if (!TR.query(inst).Inner0().isPossiblePointer())
           continue;
 
-        Instruction *newi = getNewFromOriginal(inst);
-
         if (isa<LoadInst>(inst)) {
-          IRBuilder<> BuilderZ(getNextNonDebugInstruction(newi));
-          BuilderZ.setFastMathFlags(getFast());
+          IRBuilder<> BuilderZ(inst);
+          getForwardBuilder(BuilderZ);
+
           PHINode *anti = BuilderZ.CreatePHI(inst->getType(), 1,
                                              inst->getName() + "'il_phi");
           invertedPointers[inst] = anti;
@@ -987,8 +986,9 @@ class GradientUtils : public CacheUtility {
           continue;
         }
 
-        IRBuilder<> BuilderZ(getNextNonDebugInstruction(newi));
-        BuilderZ.setFastMathFlags(getFast());
+        IRBuilder<> BuilderZ(inst);
+        getForwardBuilder(BuilderZ);
+
         PHINode *anti =
             BuilderZ.CreatePHI(op->getType(), 1, op->getName() + "'ip_phi");
         invertedPointers[inst] = anti;
@@ -1181,6 +1181,39 @@ class GradientUtils : public CacheUtility {
                                            /*successor*/ BasicBlock *>>>
           &targetToPreds,
       const std::map<BasicBlock *, PHINode *> *replacePHIs = nullptr);
+
+  void getReverseBuilder(IRBuilder<> &Builder2, bool original = true) {
+    BasicBlock *BB = Builder2.GetInsertBlock();
+    if (original)
+      BB = getNewFromOriginal(BB);
+    BasicBlock *BB2 = reverseBlocks[BB].back();
+    if (!BB2) {
+      llvm::errs() << "oldFunc: " << oldFunc << "\n";
+      llvm::errs() << "newFunc: " << newFunc << "\n";
+      llvm::errs() << "could not invert " << *BB;
+    }
+    assert(BB2);
+
+    if (BB2->getTerminator())
+      Builder2.SetInsertPoint(BB2->getTerminator());
+    else
+      Builder2.SetInsertPoint(BB2);
+    Builder2.SetCurrentDebugLocation(
+        getNewFromOriginal(Builder2.getCurrentDebugLocation()));
+    Builder2.setFastMathFlags(getFast());
+  }
+
+  void getForwardBuilder(IRBuilder<> &Builder2) {
+    Instruction *insert = &*Builder2.GetInsertPoint();
+    Instruction *nInsert = getNewFromOriginal(insert);
+
+    assert(nInsert);
+
+    Builder2.SetInsertPoint(getNextNonDebugInstruction(nInsert));
+    Builder2.SetCurrentDebugLocation(
+        getNewFromOriginal(Builder2.getCurrentDebugLocation()));
+    Builder2.setFastMathFlags(getFast());
+  }
 };
 
 class DiffeGradientUtils : public GradientUtils {

enzyme/Enzyme/Utils.h

@@ -220,6 +220,7 @@ enum class ReturnType {
   Tape,
   TwoReturns,
   Return,
+  Void,
 };
 
 /// Potential differentiable argument classifications

enzyme/test/Enzyme/ForwardMode/addOneMem.ll

@@ -0,0 +1,50 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -inline -mem2reg -instsimplify -gvn -dse -dse -S | FileCheck %s
+
+; __attribute__((noinline))
+; void addOneMem(double *x) {
+;     *x += 1;
+; }
+; 
+; void test_derivative(double *x, double *xp) {
+;   __builtin_autodiff(addOneMem, x, xp);
+; }
+
+; Function Attrs: noinline norecurse nounwind uwtable
+define dso_local void @addOneMem(double* nocapture %x) {
+entry:
+  %0 = load double, double* %x, align 8, !tbaa !2
+  %add = fadd fast double %0, 1.000000e+00
+  store double %add, double* %x, align 8, !tbaa !2
+  ret void
+}
+
+; Function Attrs: nounwind uwtable
+define dso_local void @test_derivative(double* %x, double* %xp) local_unnamed_addr {
+entry:
+  %0 = tail call double (void (double*)*, ...) @__enzyme_fwddiff(void (double*)* nonnull @addOneMem, double* %x, double* %xp)
+  ret void
+}
+
+; Function Attrs: nounwind
+declare double @__enzyme_fwddiff(void (double*)*, ...)
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 7.1.0 "}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"double", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+
+
+; CHECK: define {{(dso_local )?}}void @test_derivative(double* %x, double* %xp) 
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = load double, double* %x, align 8, !tbaa !2
+; CHECK-NEXT:   %1 = load double, double* %xp
+; CHECK-NEXT:   %add.i = fadd fast double %0, 1.000000e+00
+; CHECK-NEXT:   store double %add.i, double* %x, align 8, !tbaa !2
+; CHECK-NEXT:   store double %1, double* %xp, align 8
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }

enzyme/test/Enzyme/ForwardMode/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_fwddiff(double (double)* nonnull @tester, double %x, double 1.0)
+  ret double %0
+}
+
+declare void @myprint(double %x) #0
+
+; Function Attrs: nounwind
+declare double @__enzyme_fwddiff(double (double)*, ...)
+
+attributes #0 = { "enzyme_inactive" }
+
+; CHECK: define internal {{(dso_local )?}}{ double } @diffetester(double %x, double %"x'")
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   tail call void @myprint(double %x)
+; CHECK-NEXT:   %0 = insertvalue { double } undef, double %"x'", 0
+; CHECK-NEXT:   ret { double } %0
+; CHECK-NEXT: }