Speculate

Author: wsmoses

enzyme/Enzyme/FunctionUtils.cpp

@@ -141,6 +141,9 @@ cl::opt<bool>
     EnzymeNameInstructions("enzyme-name-instructions", cl::init(false),
                            cl::Hidden,
                            cl::desc("Have enzyme name all instructions"));
+
+cl::opt<bool> EnzymeSelectOpt("enzyme-select-opt", cl::init(true), cl::Hidden,
+                              cl::desc("Run Enzyme select optimization"));
 }
 
 /// Is the use of value val as an argument of call CI potentially captured
@@ -1648,17 +1651,20 @@ void PreProcessCache::optimizeIntermediate(Function *F) {
   PromotePass().run(*F, FAM);
   GVN().run(*F, FAM);
   SROA().run(*F, FAM);
+
+  if (EnzymeSelectOpt) {
 #if LLVM_VERSION_MAJOR >= 12
-  SimplifyCFGOptions scfgo;
+    SimplifyCFGOptions scfgo;
 #else
-  SimplifyCFGOptions scfgo(
-      /*unsigned BonusThreshold=*/1, /*bool ForwardSwitchCond=*/false,
-      /*bool SwitchToLookup=*/false, /*bool CanonicalLoops=*/true,
-      /*bool SinkCommon=*/true, /*AssumptionCache *AssumpCache=*/nullptr);
+    SimplifyCFGOptions scfgo(
+        /*unsigned BonusThreshold=*/1, /*bool ForwardSwitchCond=*/false,
+        /*bool SwitchToLookup=*/false, /*bool CanonicalLoops=*/true,
+        /*bool SinkCommon=*/true, /*AssumptionCache *AssumpCache=*/nullptr);
 #endif
-  SimplifyCFGPass(scfgo).run(*F, FAM);
-  CorrelatedValuePropagationPass().run(*F, FAM);
-  SelectOptimization(F);
+    SimplifyCFGPass(scfgo).run(*F, FAM);
+    CorrelatedValuePropagationPass().run(*F, FAM);
+    SelectOptimization(F);
+  }
   // EarlyCSEPass(/*memoryssa*/ true).run(*F, FAM);
 
   for (Function &Impl : *F->getParent()) {

enzyme/Enzyme/GradientUtils.cpp

@@ -75,6 +75,9 @@ llvm::cl::opt<bool>
 llvm::cl::opt<bool>
     EnzymeRegisterReduce("enzyme-register-reduce", cl::init(false), cl::Hidden,
                          cl::desc("Reduce the amount of register reduce"));
+llvm::cl::opt<bool>
+    EnzymeSpeculatePHIs("enzyme-speculate-phis", cl::init(false), cl::Hidden,
+                        cl::desc("Speculatively execute phi computations"));
 }
 
 bool isPotentialLastLoopValue(Value *val, const BasicBlock *loc,
@@ -869,12 +872,20 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 
               unwrap_cache[blocks[i]] = unwrap_cache[oldB];
               lookup_cache[blocks[i]] = lookup_cache[oldB];
+              auto PB = *done[std::make_pair(parent, predBlocks[i])].begin();
+
+              if (auto inst = dyn_cast<Instruction>(
+                      phi->getIncomingValueForBlock(PB))) {
+                if (inst->mayReadFromMemory() || !EnzymeSpeculatePHIs)
+                  vals.push_back(
+                      getOpFull(B, phi->getIncomingValueForBlock(PB), PB));
+                else
+                  vals.push_back(getOpFull(
+                      BuilderM, phi->getIncomingValueForBlock(PB), PB));
+              } else
+                vals.push_back(
+                    getOpFull(BuilderM, phi->getIncomingValueForBlock(PB), PB));
 
-              vals.push_back(getOpFull(
-                  B,
-                  phi->getIncomingValueForBlock(
-                      *done[std::make_pair(parent, predBlocks[i])].begin()),
-                  *done[std::make_pair(parent, predBlocks[i])].begin()));
               if (!vals[i]) {
                 for (size_t j = 0; j < i; i++) {
                   reverseBlocks[fwd].erase(std::find(reverseBlocks[fwd].begin(),
@@ -1034,7 +1045,7 @@ Value *GradientUtils::unwrapM(Value *const val, IRBuilder<> &BuilderM,
 
         if (auto inst =
                 dyn_cast<Instruction>(phi->getIncomingValueForBlock(PB))) {
-          if (inst->mayReadFromMemory())
+          if (inst->mayReadFromMemory() || !EnzymeSpeculatePHIs)
             vals.push_back(getOpFull(B, phi->getIncomingValueForBlock(PB), PB));
           else
             vals.push_back(