Implement fast block lookup and phi recomputation (rust-lang#234)

wsmoses · web-flow · commit d4d3ae9c5dd8 · 2021-07-18T05:19:15.000-04:00
* Implement fast block lookup and phi recomputation

* Attempt debug ci

* Fix cache lookup

* try debug build

* Fix pending unwrap bug
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -3013,8 +3013,13 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
         for (auto V : unwrapToOrig[newi]) {
           ValueToValueMapTy empty;
           IRBuilder<> lb(cast<Instruction>(V));
+          // This must disallow caching here as otherwise performing the loop in
+          // the wrong order may result in first replacing the later unwrapped
+          // value, caching it, then attempting to reuse it for an earlier
+          // replacement.
           V->replaceAllUsesWith(
-              gutils->unwrapM(nexti, lb, empty, UnwrapMode::LegalFullUnwrap));
+              gutils->unwrapM(nexti, lb, empty, UnwrapMode::LegalFullUnwrap,
+                              /*scope*/ nullptr, /*permitCache*/ false));
           cast<Instruction>(V)->eraseFromParent();
         }
       }
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -176,15 +176,13 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
   // assert(!val->getName().startswith("$tapeload"));
   if (permitCache && unwrap_cache[BuilderM.GetInsertBlock()].find(idx) !=
                          unwrap_cache[BuilderM.GetInsertBlock()].end()) {
-    if (unwrap_cache[BuilderM.GetInsertBlock()][idx]->getType() !=
-        val->getType()) {
+    auto cachedValue = unwrap_cache[BuilderM.GetInsertBlock()][idx];
+    if (cachedValue->getType() != val->getType()) {
       llvm::errs() << "val: " << *val << "\n";
-      llvm::errs() << "unwrap_cache[cidx]: "
-                   << *unwrap_cache[BuilderM.GetInsertBlock()][idx] << "\n";
+      llvm::errs() << "unwrap_cache[cidx]: " << *cachedValue << "\n";
     }
-    assert(unwrap_cache[BuilderM.GetInsertBlock()][idx]->getType() ==
-           val->getType());
-    return unwrap_cache[BuilderM.GetInsertBlock()][idx];
+    assert(cachedValue->getType() == val->getType());
+    return cachedValue;
   }
 
 #define getOpFullest(Builder, vtmp, frominst, check)                           \
@@ -361,8 +359,9 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
     auto toreturn = BuilderM.CreateBinOp(op->getOpcode(), op0, op1,
                                          op->getName() + "_unwrap");
     unwrappedLoads[toreturn] = val;
-    if (auto newi = dyn_cast<Instruction>(toreturn))
+    if (auto newi = dyn_cast<Instruction>(toreturn)) {
       newi->copyIRFlags(op);
+    }
     if (permitCache)
       unwrap_cache[BuilderM.GetInsertBlock()][idx] = toreturn;
     assert(val->getType() == toreturn->getType());
@@ -834,17 +833,10 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
             if (BuilderM.GetInsertPoint() != oldB->end())
               goto endCheck;
 
-            // todo speed this up
-            BasicBlock *fwd = nullptr;
-            for (const auto &pair : reverseBlocks) {
-              const std::vector<BasicBlock *> &vec = pair.second;
-              if (std::find(vec.begin(), vec.end(), oldB) != vec.end()) {
-                fwd = pair.first;
-                break;
-              }
-            }
-            if (!fwd)
+            auto found = reverseBlockToPrimal.find(oldB);
+            if (found == reverseBlockToPrimal.end())
               goto endCheck;
+            BasicBlock *fwd = found->second;
 
             SmallVector<BasicBlock *, 2> predBlocks;
             predBlocks.push_back(bi2->getSuccessor(0));
@@ -876,6 +868,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
               blocks[i]->moveAfter(last);
               last = blocks[i];
               reverseBlocks[fwd].push_back(blocks[i]);
+              reverseBlockToPrimal[blocks[i]] = fwd;
               IRBuilder<> B(blocks[i]);
 
               unwrap_cache[blocks[i]] = unwrap_cache[oldB];
@@ -884,7 +877,17 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 
               if (auto inst = dyn_cast<Instruction>(
                       phi->getIncomingValueForBlock(PB))) {
-                if (inst->mayReadFromMemory() || !EnzymeSpeculatePHIs)
+                // Recompute the phi computation with the conditional if:
+                // 1) the instruction may reat from memory AND does not
+                //    dominate the current insertion point (thereby
+                //    potentially making such recomputation without the
+                //    condition illegal)
+                // 2) the value is a call or load and option is set to not
+                //    speculatively recompute values within a phi
+                if ((inst->mayReadFromMemory() &&
+                     !DT.dominates(inst->getParent(), phi->getParent())) ||
+                    (!EnzymeSpeculatePHIs &&
+                     (isa<CallInst>(inst) || isa<LoadInst>(inst))))
                   vals.push_back(
                       getOpFull(B, phi->getIncomingValueForBlock(PB), PB));
                 else
@@ -895,10 +898,11 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
                     getOpFull(BuilderM, phi->getIncomingValueForBlock(PB), PB));
 
               if (!vals[i]) {
-                for (size_t j = 0; j < i; i++) {
+                for (size_t j = 0; j < i; j++) {
                   reverseBlocks[fwd].erase(std::find(reverseBlocks[fwd].begin(),
                                                      reverseBlocks[fwd].end(),
                                                      blocks[j]));
+                  reverseBlockToPrimal.erase(blocks[j]);
                   unwrap_cache.erase(blocks[j]);
                   lookup_cache.erase(blocks[j]);
                   SmallVector<Instruction *, 4> toErase;
@@ -910,7 +914,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
                   }
                 }
                 bret->eraseFromParent();
-                for (size_t j = 0; j < i; i++) {
+                for (size_t j = 0; j < i; j++) {
                   blocks[j]->eraseFromParent();
                 };
                 goto endCheck;
@@ -939,6 +943,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 
             BuilderM.SetInsertPoint(bret);
             reverseBlocks[fwd].push_back(bret);
+            reverseBlockToPrimal[bret] = fwd;
             auto toret = BuilderM.CreatePHI(val->getType(), vals.size());
             for (size_t i = 0; i < vals.size(); i++)
               toret->addIncoming(vals[i], endingBlocks[i]);
@@ -992,17 +997,10 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
       if (BuilderM.GetInsertPoint() != oldB->end())
         goto endCheck;
 
-      // todo speed this up
-      BasicBlock *fwd = nullptr;
-      for (const auto &pair : reverseBlocks) {
-        const std::vector<BasicBlock *> &vec = pair.second;
-        if (std::find(vec.begin(), vec.end(), oldB) != vec.end()) {
-          fwd = pair.first;
-          break;
-        }
-      }
-      if (!fwd)
+      auto found = reverseBlockToPrimal.find(oldB);
+      if (found == reverseBlockToPrimal.end())
         goto endCheck;
+      BasicBlock *fwd = found->second;
 
       SmallVector<BasicBlock *, 2> predBlocks;
       Value *cond = nullptr;
@@ -1046,14 +1044,24 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
         blocks[i]->moveAfter(last);
         last = blocks[i];
         reverseBlocks[fwd].push_back(blocks[i]);
+        reverseBlockToPrimal[blocks[i]] = fwd;
         IRBuilder<> B(blocks[i]);
 
         unwrap_cache[blocks[i]] = unwrap_cache[oldB];
         lookup_cache[blocks[i]] = lookup_cache[oldB];
 
         if (auto inst =
                 dyn_cast<Instruction>(phi->getIncomingValueForBlock(PB))) {
-          if (inst->mayReadFromMemory() || !EnzymeSpeculatePHIs)
+          // Recompute the phi computation with the conditional if:
+          // 1) the instruction may reat from memory AND does not dominate
+          //    the current insertion point (thereby potentially making such
+          //    recomputation without the condition illegal)
+          // 2) the value is a call or load and option is set to not
+          //    speculatively recompute values within a phi
+          if ((inst->mayReadFromMemory() &&
+               !DT.dominates(inst->getParent(), phi->getParent())) ||
+              (!EnzymeSpeculatePHIs &&
+               (isa<CallInst>(inst) || isa<LoadInst>(inst))))
             vals.push_back(getOpFull(B, phi->getIncomingValueForBlock(PB), PB));
           else
             vals.push_back(
@@ -1063,10 +1071,11 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
               getOpFull(BuilderM, phi->getIncomingValueForBlock(PB), PB));
 
         if (!vals[i]) {
-          for (size_t j = 0; j < i; i++) {
+          for (size_t j = 0; j < i; j++) {
             reverseBlocks[fwd].erase(std::find(reverseBlocks[fwd].begin(),
                                                reverseBlocks[fwd].end(),
                                                blocks[j]));
+            reverseBlockToPrimal.erase(blocks[j]);
             unwrap_cache.erase(blocks[j]);
             lookup_cache.erase(blocks[j]);
             SmallVector<Instruction *, 4> toErase;
@@ -1078,7 +1087,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
             }
           }
           bret->eraseFromParent();
-          for (size_t j = 0; j < i; i++) {
+          for (size_t j = 0; j < i; j++) {
             blocks[j]->eraseFromParent();
           };
           goto endCheck;
@@ -1088,6 +1097,38 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
         endingBlocks.push_back(B.GetInsertBlock());
       }
 
+      // Fast path to not make a split block if no additional instructions
+      // were made in the two blocks
+      if (isa<BranchInst>(equivalentTerminator) && blocks[0]->size() == 1 &&
+          blocks[1]->size() == 1) {
+        for (size_t j = 0; j < blocks.size(); j++) {
+          reverseBlocks[fwd].erase(std::find(
+              reverseBlocks[fwd].begin(), reverseBlocks[fwd].end(), blocks[j]));
+          reverseBlockToPrimal.erase(blocks[j]);
+          unwrap_cache.erase(blocks[j]);
+          lookup_cache.erase(blocks[j]);
+          SmallVector<Instruction *, 4> toErase;
+          for (auto &I : *blocks[j]) {
+            toErase.push_back(&I);
+          }
+          for (auto I : toErase) {
+            erase(I);
+          }
+        }
+        bret->eraseFromParent();
+        for (size_t j = 0; j < blocks.size(); j++) {
+          blocks[j]->eraseFromParent();
+        };
+        Value *toret = BuilderM.CreateSelect(cond, vals[0], vals[1],
+                                             phi->getName() + "_unwrap");
+        if (permitCache) {
+          unwrap_cache[BuilderM.GetInsertBlock()][idx] = toret;
+        }
+        if (auto instRet = dyn_cast<Instruction>(toret))
+          unwrappedLoads[instRet] = val;
+        return toret;
+      }
+
       bret->moveAfter(last);
       if (isa<BranchInst>(equivalentTerminator)) {
         BuilderM.CreateCondBr(cond, blocks[0], blocks[1]);
@@ -1102,6 +1143,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
       }
       BuilderM.SetInsertPoint(bret);
       reverseBlocks[fwd].push_back(bret);
+      reverseBlockToPrimal[bret] = fwd;
       auto toret = BuilderM.CreatePHI(val->getType(), vals.size());
       for (size_t i = 0; i < vals.size(); i++)
         toret->addIncoming(vals[i], endingBlocks[i]);
@@ -1328,9 +1370,10 @@ Value *GradientUtils::cacheForReverse(IRBuilder<> &BuilderQ, Value *malloc,
           for (auto u : users) {
             if (auto li = dyn_cast<LoadInst>(u)) {
               IRBuilder<> lb(li);
-              ValueToValueMapTy empty;
-              li->replaceAllUsesWith(
-                  unwrapM(ret, lb, empty, UnwrapMode::LegalFullUnwrap));
+              auto replacewith =
+                  (idx < 0) ? tape
+                            : lb.CreateExtractValue(tape, {(unsigned)idx});
+              li->replaceAllUsesWith(replacewith);
               erase(li);
             } else {
               llvm::errs() << "newFunc: " << *newFunc << "\n";
@@ -1768,13 +1811,10 @@ bool GradientUtils::legalRecompute(const Value *val,
       if (BuilderM) {
         fwdBlockIfReverse = BuilderM->GetInsertBlock();
         if (!reverse) {
-          for (auto pair : reverseBlocks) {
-            if (std::find(pair.second.begin(), pair.second.end(),
-                          BuilderM->GetInsertBlock()) != pair.second.end()) {
-              fwdBlockIfReverse = pair.first;
-              reverse = true;
-              break;
-            }
+          auto found = reverseBlockToPrimal.find(BuilderM->GetInsertBlock());
+          if (found != reverseBlockToPrimal.end()) {
+            fwdBlockIfReverse = found->second;
+            reverse = true;
           }
         }
         if (fwdBlockIfReverse->getParent() != oldFunc)
diff --git a/enzyme/Enzyme/GradientUtils.h b/enzyme/Enzyme/GradientUtils.h
@@ -106,7 +106,12 @@ class GradientUtils : public CacheUtility {
   ScalarEvolution &OrigSE;
   std::shared_ptr<ActivityAnalyzer> ATA;
   SmallVector<BasicBlock *, 12> originalBlocks;
+
+  // Map of primal block to corresponding block(s) in reverse
   std::map<BasicBlock *, std::vector<BasicBlock *>> reverseBlocks;
+  // Map of block in reverse to corresponding primal block
+  std::map<BasicBlock *, BasicBlock *> reverseBlockToPrimal;
+
   SmallPtrSet<PHINode *, 4> fictiousPHIs;
   ValueToValueMapTy originalToNewFn;
   std::vector<CallInst *> originalCalls;
@@ -441,25 +446,23 @@ class GradientUtils : public CacheUtility {
   BasicBlock *addReverseBlock(BasicBlock *currentBlock, Twine name,
                               bool forkCache = true) {
     assert(reverseBlocks.size());
-
-    // todo speed this up
-    for (auto &pair : reverseBlocks) {
-      std::vector<BasicBlock *> &vec = pair.second;
-      if (vec.back() == currentBlock) {
-
-        BasicBlock *rev =
-            BasicBlock::Create(currentBlock->getContext(), name, newFunc);
-        rev->moveAfter(currentBlock);
-        vec.push_back(rev);
-        if (forkCache) {
-          unwrap_cache[rev] = unwrap_cache[currentBlock];
-          lookup_cache[rev] = lookup_cache[currentBlock];
-        }
-        return rev;
-      }
-    }
-    assert(0 && "cannot find reverse location to add into");
-    llvm_unreachable("cannot find reverse location to add into");
+    auto found = reverseBlockToPrimal.find(currentBlock);
+    assert(found != reverseBlockToPrimal.end());
+
+    std::vector<BasicBlock *> &vec = reverseBlocks[found->second];
+    assert(vec.size());
+    assert(vec.back() == currentBlock);
+
+    BasicBlock *rev =
+        BasicBlock::Create(currentBlock->getContext(), name, newFunc);
+    rev->moveAfter(currentBlock);
+    vec.push_back(rev);
+    reverseBlockToPrimal[rev] = found->second;
+    if (forkCache) {
+      unwrap_cache[rev] = unwrap_cache[currentBlock];
+      lookup_cache[rev] = lookup_cache[currentBlock];
+    }
+    return rev;
   }
 
 public:
@@ -831,19 +834,13 @@ class GradientUtils : public CacheUtility {
 
 private:
   BasicBlock *originalForReverseBlock(BasicBlock &BB2) const {
-    assert(reverseBlocks.size() != 0);
-    for (auto BB : originalBlocks) {
-      auto it = reverseBlocks.find(BB);
-      assert(it != reverseBlocks.end());
-      if (std::find(it->second.begin(), it->second.end(), &BB2) !=
-          it->second.end()) {
-        return BB;
-      }
+    auto found = reverseBlockToPrimal.find(&BB2);
+    if (found == reverseBlockToPrimal.end()) {
+      llvm::errs() << "newFunc: " << *newFunc << "\n";
+      llvm::errs() << BB2 << "\n";
     }
-    llvm::errs() << *newFunc << "\n";
-    llvm::errs() << BB2 << "\n";
-    assert(0 && "could not find original block for given reverse block");
-    report_fatal_error("could not find original block for given reverse block");
+    assert(found != reverseBlockToPrimal.end());
+    return found->second;
   }
 
 public:
@@ -1264,8 +1261,10 @@ class DiffeGradientUtils : public GradientUtils {
     for (BasicBlock *BB : originalBlocks) {
       if (BB == inversionAllocs)
         continue;
-      reverseBlocks[BB].push_back(BasicBlock::Create(
-          BB->getContext(), "invert" + BB->getName(), newFunc));
+      BasicBlock *RBB = BasicBlock::Create(BB->getContext(),
+                                           "invert" + BB->getName(), newFunc);
+      reverseBlocks[BB].push_back(RBB);
+      reverseBlockToPrimal[RBB] = BB;
     }
     assert(reverseBlocks.size() != 0);
   }
diff --git a/enzyme/test/Enzyme/ReverseMode/reorderrep.ll b/enzyme/test/Enzyme/ReverseMode/reorderrep.ll
