Fix recursive global (rust-lang#704)

wsmoses · web-flow · commit 3d28f8354167 · 2022-06-24T13:40:54.000-04:00
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -10864,8 +10864,30 @@ class AdjointGenerator
           (orig->mayWriteToMemory() ||
            !gutils->legalRecompute(orig, ValueToValueMapTy(), nullptr))) {
         if (!gutils->unnecessaryIntermediates.count(orig)) {
-          gutils->cacheForReverse(BuilderZ, newCall,
-                                  getIndex(orig, CacheType::Self));
+
+          std::map<UsageKey, bool> Seen;
+          bool primalNeededInReverse = false;
+          for (auto pair : gutils->knownRecomputeHeuristic)
+            if (!pair.second) {
+              if (pair.first == orig) {
+                primalNeededInReverse = true;
+                break;
+              } else {
+                Seen[UsageKey(pair.first, ValueType::Primal)] = false;
+              }
+            }
+          if (!primalNeededInReverse) {
+
+            auto minCutMode = (Mode == DerivativeMode::ReverseModePrimal)
+                                  ? DerivativeMode::ReverseModeGradient
+                                  : Mode;
+            primalNeededInReverse =
+                is_value_needed_in_reverse<ValueType::Primal>(
+                    gutils, orig, minCutMode, Seen, oldUnreachable);
+          }
+          if (primalNeededInReverse)
+            gutils->cacheForReverse(BuilderZ, newCall,
+                                    getIndex(orig, CacheType::Self));
         }
         eraseIfUnused(*orig);
         return;
@@ -11460,11 +11482,31 @@ class AdjointGenerator
           }
 
           if (Mode == DerivativeMode::ReverseModePrimal &&
-              is_value_needed_in_reverse<ValueType::Primal>(gutils, orig, Mode,
-                                                            oldUnreachable) &&
               !gutils->unnecessaryIntermediates.count(orig)) {
-            gutils->cacheForReverse(BuilderZ, dcall,
-                                    getIndex(orig, CacheType::Self));
+
+            std::map<UsageKey, bool> Seen;
+            bool primalNeededInReverse = false;
+            for (auto pair : gutils->knownRecomputeHeuristic)
+              if (!pair.second) {
+                if (pair.first == orig) {
+                  primalNeededInReverse = true;
+                  break;
+                } else {
+                  Seen[UsageKey(pair.first, ValueType::Primal)] = false;
+                }
+              }
+            if (!primalNeededInReverse) {
+
+              auto minCutMode = (Mode == DerivativeMode::ReverseModePrimal)
+                                    ? DerivativeMode::ReverseModeGradient
+                                    : Mode;
+              primalNeededInReverse =
+                  is_value_needed_in_reverse<ValueType::Primal>(
+                      gutils, orig, minCutMode, Seen, oldUnreachable);
+            }
+            if (primalNeededInReverse)
+              gutils->cacheForReverse(BuilderZ, dcall,
+                                      getIndex(orig, CacheType::Self));
           }
           BuilderZ.SetInsertPoint(newCall->getNextNode());
           gutils->erase(newCall);
diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.h b/enzyme/Enzyme/DifferentialUseAnalysis.h
@@ -275,23 +275,33 @@ static inline bool is_value_needed_in_reverse(
         return seen[idx] = true;
 
       if (auto SI = dyn_cast<StoreInst>(user)) {
-        // storing an active pointer into a location
-        // doesn't require the shadow pointer for the
-        // reverse pass
-        if (SI->getValueOperand() == inst &&
-            (mode == DerivativeMode::ReverseModeGradient ||
-             mode == DerivativeMode::ForwardModeSplit)) {
-          // Unless the store is into a backwards store, which would
-          // would then be performed in the reverse if the stored value was
-          // a possible pointer.
+        if (mode == DerivativeMode::ReverseModeGradient ||
+            mode == DerivativeMode::ForwardModeSplit) {
+
           bool rematerialized = false;
           for (auto pair : gutils->backwardsOnlyShadows)
             if (pair.second.stores.count(SI)) {
               rematerialized = true;
               break;
             }
-          if (!rematerialized)
-            goto endShadow;
+
+          if (SI->getValueOperand() == inst) {
+            // storing an active pointer into a location
+            // doesn't require the shadow pointer for the
+            // reverse pass
+            // Unless the store is into a backwards store, which would
+            // would then be performed in the reverse if the stored value was
+            // a possible pointer.
+            if (!rematerialized)
+              goto endShadow;
+          } else {
+            // Likewise, if not rematerializing in reverse pass, you
+            // don't need to keep the pointer operand for known pointers
+            if (!rematerialized &&
+                TR.query(const_cast<Value *>(SI->getValueOperand()))[{-1}] ==
+                    BaseType::Pointer)
+              goto endShadow;
+          }
         }
 
         if (!gutils->isConstantValue(
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -2472,8 +2472,36 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
   }
 
   SmallVector<CallInst *, 4> fnusers;
+  SmallVector<std::pair<GlobalVariable *, DerivativeMode>, 1> gfnusers;
   for (auto user : AugmentedCachedFunctions.find(tup)->second.fn->users()) {
-    fnusers.push_back(cast<CallInst>(user));
+    if (auto CI = dyn_cast<CallInst>(user)) {
+      fnusers.push_back(CI);
+    } else {
+      if (auto CS = dyn_cast<ConstantStruct>(user)) {
+        for (auto cuser : CS->users()) {
+          if (auto G = dyn_cast<GlobalVariable>(cuser)) {
+            if (("_enzyme_reverse_" + todiff->getName() + "'").str() ==
+                G->getName()) {
+              gfnusers.emplace_back(G, DerivativeMode::ReverseModeGradient);
+              continue;
+            }
+            if (("_enzyme_forwardsplit_" + todiff->getName() + "'").str() ==
+                G->getName()) {
+              gfnusers.emplace_back(G, DerivativeMode::ForwardModeSplit);
+              continue;
+            }
+          }
+          llvm::errs() << *gutils->newFunc->getParent() << "\n";
+          llvm::errs() << *cuser << "\n";
+          llvm::errs() << *user << "\n";
+          llvm_unreachable("Bad cuser of staging augmented forward fn");
+        }
+        continue;
+      }
+      llvm::errs() << *gutils->newFunc->getParent() << "\n";
+      llvm::errs() << *user << "\n";
+      llvm_unreachable("Bad user of staging augmented forward fn");
+    }
   }
   for (auto user : fnusers) {
     if (removeStruct) {
@@ -2514,6 +2542,25 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
     AugmentedCachedFunctions.find(tup)->second.tapeType = tapeType;
   insert_or_assign(AugmentedCachedFinished, tup, true);
 
+  for (auto pair : gfnusers) {
+    auto GV = pair.first;
+    GV->setName("_tmp");
+    auto R = gutils->GetOrCreateShadowFunction(
+        *this, TLI, TA, todiff, pair.second, width, gutils->AtomicAdd);
+    SmallVector<ConstantExpr *, 1> users;
+    for (auto U : GV->users()) {
+      if (auto CE = dyn_cast<ConstantExpr>(U)) {
+        if (CE->isCast()) {
+          users.push_back(CE);
+        }
+      }
+    }
+    for (auto U : users) {
+      U->replaceAllUsesWith(ConstantExpr::getPointerCast(R, U->getType()));
+    }
+    GV->eraseFromParent();
+  }
+
   {
     PreservedAnalyses PA;
     PPC.FAM.invalidate(*gutils->newFunc, PA);
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3582,13 +3582,9 @@ Constant *GradientUtils::GetOrCreateShadowConstant(
       Type *type = arg->getType()->getPointerElementType();
       auto shadow = new GlobalVariable(
           *arg->getParent(), type, arg->isConstant(), arg->getLinkage(),
-          arg->getInitializer()
-              ? GetOrCreateShadowConstant(Logic, TLI, TA,
-                                          cast<Constant>(arg->getOperand(0)),
-                                          mode, width, AtomicAdd)
-              : Constant::getNullValue(type),
-          arg->getName() + "_shadow", arg, arg->getThreadLocalMode(),
-          arg->getType()->getAddressSpace(), arg->isExternallyInitialized());
+          Constant::getNullValue(type), arg->getName() + "_shadow", arg,
+          arg->getThreadLocalMode(), arg->getType()->getAddressSpace(),
+          arg->isExternallyInitialized());
       arg->setMetadata("enzyme_shadow",
                        MDTuple::get(shadow->getContext(),
                                     {ConstantAsMetadata::get(shadow)}));
@@ -3598,6 +3594,10 @@ Constant *GradientUtils::GetOrCreateShadowConstant(
       shadow->setAlignment(arg->getAlignment());
 #endif
       shadow->setUnnamedAddr(arg->getUnnamedAddr());
+      if (arg->getInitializer())
+        shadow->setInitializer(GetOrCreateShadowConstant(
+            Logic, TLI, TA, cast<Constant>(arg->getOperand(0)), mode, width,
+            AtomicAdd));
       return shadow;
     }
   }
@@ -4103,14 +4103,11 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
         Type *elemTy = arg->getType()->getPointerElementType();
         Type *type = getShadowType(elemTy);
         IRBuilder<> B(inversionAllocs);
-        auto ip = arg->getInitializer() ? invertPointerM(arg->getInitializer(),
-                                                         B, /*nullShadow*/ true)
-                                        : Constant::getNullValue(type);
 
-        auto rule = [&](Value *ip) {
+        auto rule = [&]() {
           auto shadow = new GlobalVariable(
               *arg->getParent(), elemTy, arg->isConstant(), arg->getLinkage(),
-              cast<Constant>(ip), arg->getName() + "_shadow", arg,
+              Constant::getNullValue(type), arg->getName() + "_shadow", arg,
               arg->getThreadLocalMode(), arg->getType()->getAddressSpace(),
               arg->isExternallyInitialized());
           arg->setMetadata("enzyme_shadow",
@@ -4126,7 +4123,18 @@ Value *GradientUtils::invertPointerM(Value *const oval, IRBuilder<> &BuilderM,
           return shadow;
         };
 
-        Value *shadow = applyChainRule(oval->getType(), BuilderM, rule, ip);
+        Value *shadow = applyChainRule(oval->getType(), BuilderM, rule);
+
+        if (arg->getInitializer()) {
+          applyChainRule(
+              BuilderM,
+              [&](Value *shadow, Value *ip) {
+                cast<GlobalVariable>(shadow)->setInitializer(
+                    cast<Constant>(ip));
+              },
+              shadow,
+              invertPointerM(arg->getInitializer(), B, /*nullShadow*/ true));
+        }
 
         invertedPointers.insert(std::make_pair(
             (const Value *)oval, InvertedPointerVH(this, shadow)));
diff --git a/enzyme/test/Enzyme/ReverseMode/recurglobfn.ll b/enzyme/test/Enzyme/ReverseMode/recurglobfn.ll
@@ -0,0 +1,44 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -S | FileCheck %s
+
+@g = constant i8* bitcast (void (i8***)* @impl to i8*), align 8
+
+declare i32 @offset()
+
+define void @impl(i8*** %i) {
+  %o = call i32 @offset()
+  %g = getelementptr inbounds i8**, i8*** %i, i32 %o
+  store i8** @g, i8*** %g, align 8
+  ret void
+}
+
+define i8** @caller() {
+  %i6 = call i8** (...) @__enzyme_virtualreverse(i8** @g)
+  ret i8** %i6
+}
+
+declare i8** @__enzyme_virtualreverse(...)
+
+; CHECK: @g_shadow = constant i8* bitcast ({ i8* (i8***, i8***)*, void (i8***, i8***, i8*)* }* @"_enzyme_reverse_impl'" to i8*), align 8
+; CHECK: @g = constant i8* bitcast (void (i8***)* @impl to i8*), align 8, !enzyme_shadow !0
+; CHECK: @"_enzyme_reverse_impl'" = internal constant { i8* (i8***, i8***)*, void (i8***, i8***, i8*)* } { i8* (i8***, i8***)* @augmented_impl, void (i8***, i8***, i8*)* @diffeimpl }
+
+; CHECK: define i8** @caller() {
+; CHECK-NEXT:   ret i8** @g_shadow
+; CHECK-NEXT: }
+
+; CHECK: define internal i8* @augmented_impl(i8*** %i, i8*** %"i'")
+; CHECK-NEXT:   %o = call i32 @offset() 
+; CHECK-NEXT:   %"g'ipg" = getelementptr inbounds i8**, i8*** %"i'", i32 %o
+; CHECK-NEXT:   %g = getelementptr inbounds i8**, i8*** %i, i32 %o
+; CHECK-NEXT:   store i8** @g_shadow, i8*** %"g'ipg", align 8
+; CHECK-NEXT:   store i8** @g, i8*** %g, align 8
+; CHECK-NEXT:   ret i8* null
+; CHECK-NEXT: }
+
+; CHECK: define internal void @diffeimpl(i8*** %i, i8*** %"i'", i8* %tapeArg)
+; CHECK-NEXT: invert:
+; CHECK-NEXT:   tail call void @free(i8* nonnull %tapeArg)
+; CHECK-NEXT:   ret void
+; CHECK-NEXT: }
+
+; CHECK: !0 = !{i8** @g_shadow}
diff --git a/enzyme/test/Enzyme/ReverseMode/unnecessaryalloc.ll b/enzyme/test/Enzyme/ReverseMode/unnecessaryalloc.ll
@@ -51,20 +51,18 @@ declare void @free(i8*)
 !5 = distinct !{}
 
 
-; CHECK: define internal { i64, double } @augmented_diffemy_sin2(double %x, double %differeturn)
+; CHECK: define internal double @augmented_diffemy_sin2(double %x, double %differeturn)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %i = add i64 14, 1
 ; CHECK-NEXT:   %i1 = add nuw i64 %i, 1
 ; CHECK-NEXT:   %mallocsize = mul nuw nsw i64 %i1, 8
 ; CHECK-NEXT:   %malloccall = tail call noalias nonnull i8* @malloc(i64 %mallocsize)
 ; CHECK-NEXT:   tail call void @free(i8* nonnull %malloccall)
 ; CHECK-NEXT:   %0 = insertvalue { i8*, double } undef, double %x, 1
-; CHECK-NEXT:   %.fca.0.insert = insertvalue { i64, double } {{(undef|poison)}}, i64 14, 0
-; CHECK-NEXT:   %.fca.1.insert = insertvalue { i64, double } %.fca.0.insert, double %x, 1
-; CHECK-NEXT:   ret { i64, double } %.fca.1.insert
+; CHECK-NEXT:   ret double %x
 ; CHECK-NEXT: }
 
-; CHECK: define internal { double } @diffediffemy_sin2(double %x, double %differeturn1, double %differeturn, i64 %tapeArg)
+; CHECK: define internal { double } @diffediffemy_sin2(double %x, double %differeturn1, double %differeturn)
 ; CHECK-NEXT: entry:
 ; CHECK-NEXT:   %0 = insertvalue { double } undef, double %differeturn, 0
 ; CHECK-NEXT:   ret { double } %0
