[OpenMPOpt] ICV tracking for calls

[sstefan1]

Summary

Introduce two new AAs. AAICVTrackerFunctionReturned which checks if a
function can have a unique ICV value after it is finished, and
AAICVCallSiteReturned which checks AAICVTrackerFunctionReturned for a
call site. This enables us to check the value of a call and if it
changes the ICV. This also changes the approach in
getReplacementValues() to a worklist-based approach so we can explore
all relevant BBs.

[sstefan1]

Note: I've added D85052 as a parent revision. However this can also be considered as a slight change in approach. If needed, those 2 can be merged together.

Things like ICVValue struct could be deleted if this approach turns out to be the one we go with. But I'd like to hear from others first.

[jdoerfert]

Cool! Can you describe in the commit message in more detail what is different in this approach, how it works in general. Maybe pros and cons compared to before?

llvm/test/Transforms/OpenMP/icv_tracking.ll

281

Can we have more negative test cases too?

1. this function with weak linkage, should be handled by the Attributor already but let's make sure we verify.
2. two returns with different ICV values, e.g. unknown and known, or two different known ones.
3. not a call but an invoke of a function with a call to a unknown one that may throw. I guess we can/should just not propagate over invokes now.

[sstefan1]

Better commit message.
Will add more tests tomorrow.

[sstefan1]

adding more tests

[jdoerfert]

We need a test with more basic blocks between the setter and getter.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1145

The comment of getUniqueValue is not helpful. It looks like this will compute the value of ICV after the function was called, correct? Whatever it does, write it here ;)

1339

If this is only for calls, why not take a CallBase &CB instead?

1353–1355

UniqueReplVal should not be needed anywhere here.

1447

We should have a worklist of instructions we need to look at. Initially it is the argument, if we go through the prdedecessor blocks we add their terminator to the worklist. I'm not sure we need the explorer at all in that case.

[sstefan1]

• worklist approach
• don't use MBEC explorer
• more tests

[sstefan1]

Unclear if we still want D85052 to go in before this, since most of the login is replaced with this patch.

Should I just merge the good parts with this patch?

[jdoerfert]

llvm/test/Transforms/OpenMP/icv_tracking.ll

369

this should not happen. The weak function should not be used in IPO. I guess we could make the AA actually work on all IRPositions. (call site) return for the value after a call, and function for the body. Then, I hope, the Attributor should not give you the Return value if you query the call site return value of a weak function.

[sstefan1]

don't track values in weak functions

[jdoerfert]

In D85544#2211355, @sstefan1 wrote:

Unclear if we still want D85052 to go in before this, since most of the login is replaced with this patch.

Should I just merge the good parts with this patch?

Yes I think so.

TBH, I'm thinking we should embrace the IRPositions and Attributor design fully, as it seems a bit messy right now (= we reason about returns, calls, and instructions all somewhat together).

• AAICVTrackerFunction for the function body, given an ICV and an instruction, what is the value by backwards lookup (as you have it now). When you hit a call, ask AAICVTrackerCallSiteReturn
• AAICVTrackerFunctionReturn which should forward queries to AAICVTrackerFunction if it is not invalid. The queries are done for all return inst.
• AAICVTrackerCallSiteReturn which should forward queries to AAICVTrackerFunctionReturn.

WDYT?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1142

I am unsure about the difference of these two functions from the declarations. If we need two member functions we should name them appropriately and elaborate in the description.

1145

The description is now better, we should also change the name to match it. Maybe getICVValueAfterOnReturn ?

[sstefan1]

At first, I thought it wouldn't be very useful to have AAs for multiple positions since we only interact with the function one.
Now that I actually tried to make it work, turns out it is quiet a nice clean up and makes much more sense.

I'll update shortly.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1142

With the new AAs for CSRet and Ret positions, one of those goes away and I think it will be clearer then.

[sstefan1]

Introduce 2 new AAs
Merge good parts of D85052

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1502–1505

1519

Never reset Changed, see also below

1526

leftover

1572–1575

Never reset Changed

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1348–1350

Here and below, we should cache this, right?

1358–1361

leftover.

1371

why do we need the hasTrackedValue call (and function)?
Caching again. And a TODO that we should look at call sites if possible.

1449–1454

Caching.

[sstefan1]

• Add AAICVTrackerCallSite
• Move manifest to AAICVTrackerCallSite
• Only register callsite AAs for getters.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1371

Turns out hasTrackedValue is actually not needed.

As for caching, values are kind of cached inside appropriate AAICVTrackerFunctionReturned, AAICVTrackerCallSiteReturned or AAICVTrackerCallSite

[sstefan1]

• Remove hasTrackedValue()

[jdoerfert]

Looks much better now. Mostly minor comments.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

923

Nit: We can now get rid of the outer loop and the conditional as well by using the SCC version of foreachUse.

1264

Don't we need to check if we already assumed this value?

1299

What is this about?

1302

I feel this update doesn't need to look at everything but only the things we have in the cache. Right now it will look at all calls for all ICVs, which is probably not necessary.

1347

Arguably, the value for the ICV is the getter call, so I.

1540

This seems unfortunate but OK for now.

[sstefan1]

addressing comments

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1299

As per IRC discussion, added TODO to avoid this. Probably in a follow up, I'll try to introduce new state and possibly make other AAs not depend on function AA as much.

This applies to the comment bellow as well.

1347

The reason we return None here is because getters don't change the ICV.
Maybe a different name for the function would be the best solution?

[jdoerfert]

LGTM.

[jdoerfert]

I'll revert this one in a second as it breaks some codes (at least on power).
The reason is that this extended the Attributor scope to the module slice. The Attributor thinks all of this is the SCC and when it tries to remove a function we see an assertion (as the function is actually not in the SCC).
The fix is to teach the Attributor about the module slice properly. Maybe in a pre-commit. For now we just revert this to unblock the build.

Here is the reproducer we should include in the OpenMP tests as well:

define internal fastcc void @"_omp$reduction$reduction_func14"() unnamed_addr {
  %call = call i8 @_ZStplIdESt7complexIT_ERKS2_S4_()
  ret void
}

define linkonce_odr hidden i8 @_ZStplIdESt7complexIT_ERKS2_S4_() local_unnamed_addr {
  ret i8 undef
}

declare void @__omp_offloading_2b_4010cad__ZN11qmcplusplus7ompBLAS17gemv_batched_implIfEEiRiciiPKT_PKS5_iS7_iS5_PKPS3_ii_l148(i64, i64, i64, float**, float**, i64, float**, float*, float*, i64) local_unnamed_addr

declare dso_local fastcc void @__kmpc_for_static_init_8u() unnamed_addr

!nvvm.annotations = !{!0}

!0 = !{void (i64, i64, i64, float**, float**, i64, float**, float*, float*, i64)* @__omp_offloading_2b_4010cad__ZN11qmcplusplus7ompBLAS17gemv_batched_implIfEEiRiciiPKT_PKS5_iS7_iS5_PKPS3_ii_l148, !"kernel", i32 1}

llvm/lib/Transforms/IPO/Attributor.cpp

bool Attributor::checkForAllInstructions(function_ref<bool(Instruction &)> Pred,
   const IRPosition &IRP = QueryingAA.getIRPosition();
   // Since we need to provide instructions we have to have an exact definition.
   const Function *AssociatedFunction = IRP.getAssociatedFunction();
   if (!AssociatedFunction)
     return false;
 
   // TODO: use the function scope once we have call site AAReturnedValues.
   const IRPosition &QueryIRP = IRPosition::function(*AssociatedFunction);
-  const auto &LivenessAA =
-      getAAFor<AAIsDead>(QueryingAA, QueryIRP, /* TrackDependence */ false);
+  const auto *LivenessAA =
+      CheckBBLivenessOnly ? nullptr
+                          : &(getAAFor<AAIsDead>(QueryingAA, QueryIRP,
+                                                 /* TrackDependence */ false));
 
   auto &OpcodeInstMap =
       InfoCache.getOpcodeInstMapForFunction(*AssociatedFunction);
   if (!checkForAllInstructionsImpl(this, OpcodeInstMap, Pred, &QueryingAA,
-                                   &LivenessAA, Opcodes, CheckBBLivenessOnly))
+                                   LivenessAA, Opcodes, CheckBBLivenessOnly))
     return false;
 
   return true;
 }
 
 bool Attributor::checkForAllReadWriteInstructions(
     function_ref<bool(Instruction &)> Pred, AbstractAttribute &QueryingAA) {
 

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

if (LookThroughConstantExprUses && isa<ConstantExpr>(U.getUser())) {
         Worklist.push_back(&CEU);
       continue;
     }
 
     CB(U);
   }
 }
 
-/// Helper struct to store tracked ICV values at specif instructions.
-struct ICVValue {
-  Instruction *Inst;
-  Value *TrackedValue;
-
-  ICVValue(Instruction *I, Value *Val) : Inst(I), TrackedValue(Val) {}
-};
-
-namespace llvm {
-
-// Provide DenseMapInfo for ICVValue
-template <> struct DenseMapInfo<ICVValue> {
-  using InstInfo = DenseMapInfo<Instruction *>;
-  using ValueInfo = DenseMapInfo<Value *>;
-
-  static inline ICVValue getEmptyKey() {
-    return ICVValue(InstInfo::getEmptyKey(), ValueInfo::getEmptyKey());
-  };
-
-  static inline ICVValue getTombstoneKey() {
-    return ICVValue(InstInfo::getTombstoneKey(), ValueInfo::getTombstoneKey());
-  };
-
-  static unsigned getHashValue(const ICVValue &ICVVal) {
-    return detail::combineHashValue(
-        InstInfo::getHashValue(ICVVal.Inst),
-        ValueInfo::getHashValue(ICVVal.TrackedValue));
-  }
-
-  static bool isEqual(const ICVValue &LHS, const ICVValue &RHS) {
-    return InstInfo::isEqual(LHS.Inst, RHS.Inst) &&
-           ValueInfo::isEqual(LHS.TrackedValue, RHS.TrackedValue);
-  }
-};
-
-} // end namespace llvm
-
 namespace {
 
 struct AAICVTracker;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
 struct OMPInformationCache : public InformationCache {
   OMPInformationCache(Module &M, AnalysisGetter &AG,
private:
   /// function.
   void registerAAs() {
     for (Function *F : SCC) {
       if (F->isDeclaration())
         continue;
 
       A.getOrCreateAAFor<AAICVTracker>(IRPosition::function(*F));
     }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
   {
struct AAICVTracker : public StateWrapper<BooleanState, AbstractAttribute> {
 
   /// Returns true if value is known to be tracked.
   bool isKnownTracked() const { return getAssumed(); }
 
   /// Create an abstract attribute biew for the position \p IRP.
   static AAICVTracker &createForPosition(const IRPosition &IRP, Attributor &A);
 
   /// Return the value with which \p I can be replaced for specific \p ICV.
-  virtual Value *getReplacementValue(InternalControlVar ICV,
-                                     const Instruction *I, Attributor &A) = 0;
+  virtual Optional<Value *> getReplacementValue(InternalControlVar ICV,
+                                                const Instruction *I,
+                                                Attributor &A) const = 0;
+
+  /// Return an assumed unique ICV value if a single candidate is found. If
+  /// there cannot be one, return a nullptr. If it is not clear yet, return the
+  /// Optional::NoneType.
+  virtual Optional<Value *> getUniqueReplacementValue(InternalControlVar ICV,
+                                                      const Instruction *I,
+                                                      Attributor &A) const = 0;
+
+  /// Compute the unique ICV value, if possible, after the function is called.
+  virtual Optional<Value *> getUniqueValue(InternalControlVar ICV,
+                                           Attributor &A) const = 0;
 
   /// Check if any value was tracked.
   virtual bool hasTrackedValue(InternalControlVar &ICV) const { return false; }
 
-  /// Check if any call potentially changes the ICV.
-  virtual bool hasUnknownCall(InternalControlVar &ICV, Attributor &A) const {
-    return false;
-  }
-
   /// See AbstractAttribute::getName()
   const std::string getName() const override { return "AAICVTracker"; }
 
   /// See AbstractAttribute::getIdAddr()
   const char *getIdAddr() const override { return &ID; }
 
-  /// This function should return true if the type of the \p AA is AAICVTracker
+  /// This function should return true if the type of the \p AA is
+  /// AAICVTracker
   static bool classof(const AbstractAttribute *AA) {
     return (AA->getIdAddr() == &ID);
   }
 
   static const char ID;
 };
 
 struct AAICVTrackerFunction : public AAICVTracker {
   AAICVTrackerFunction(const IRPosition &IRP, Attributor &A)
       : AAICVTracker(IRP, A) {}
 
+  void initialize(Attributor &A) override {
+    Function *F = getAnchorScope();
+    if (!F || !A.isFunctionIPOAmendable(*F))
+      indicatePessimisticFixpoint();
+  }
+
   // FIXME: come up with better string.
   const std::string getAsStr() const override { return "ICVTracker"; }
 
   // FIXME: come up with some stats.
   void trackStatistics() const override {}
 
   /// TODO: decide whether to deduplicate here, or use current
   /// deduplicateRuntimeCalls function.
auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
       return true;
     };
 
     GetterRFI.foreachUse(ReplaceAndDeleteCB, getAnchorScope());
     return Changed;
   }
 
   // Map of ICV to their values at specific program point.
-  EnumeratedArray<SmallSetVector<ICVValue, 4>, InternalControlVar,
-                  InternalControlVar::ICV___last>
-      ICVValuesMap;
-
-  // Map of ICV to their values at specific program point.
   EnumeratedArray<DenseMap<Instruction *, Value *>, InternalControlVar,
                   InternalControlVar::ICV___last>
       ICVReplacementValuesMap;
 
   bool hasTrackedValue(InternalControlVar &ICV) const override {
-    return !ICVValuesMap[ICV].empty();
-  }
-
-  bool hasUnknownCall(InternalControlVar &ICV, Attributor &A) const override {
-    for (BasicBlock &BB : *getAnchorScope())
-      for (Instruction &I : BB)
-        if (callChangesICV(A, &I, ICV))
-          return true;
-    return false;
+    return !ICVReplacementValuesMap[ICV].empty();
   }
 
   // Currently only nthreads is being tracked.
   // this array will only grow with time.
   InternalControlVar TrackableICVs[1] = {ICV_nthreads};
 
   ChangeStatus updateImpl(Attributor &A) override {
     ChangeStatus HasChanged = ChangeStatus::UNCHANGED;
 
     Function *F = getAnchorScope();
 
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
 
     for (InternalControlVar ICV : TrackableICVs) {
       auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
       auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
 
       auto TrackValues = [&](Use &U, Function &) {
         CallInst *CI = OpenMPOpt::getCallIfRegularCall(U);
         if (!CI)
           return false;
 
         // FIXME: handle setters with more that 1 arguments.
         /// Track new value.
-        if (ICVValuesMap[ICV].insert(ICVValue(CI, CI->getArgOperand(0))))
+        if (ICVReplacementValuesMap[ICV]
+                .insert(std::make_pair(CI, CI->getArgOperand(0)))
+                .second)
           HasChanged = ChangeStatus::CHANGED;
 
         return false;
       };
 
       // Map replacement values to Getters and their uses.
       auto MapReplacementValues = [&](Use &U, Function &) {
         CallInst *CI = OpenMPOpt::getCallIfRegularCall(U);
         if (!CI)
           return false;
 
-        Value *ReplVal = getReplacementValue(ICV, CI, A);
-        if (ICVReplacementValuesMap[ICV]
-                .insert(std::make_pair(CI, ReplVal))
-                .second)
+        Optional<Value *> ReplVal = getReplacementValue(ICV, CI, A);
+        if (ReplVal.hasValue() &&
+            ICVReplacementValuesMap[ICV]
+                .insert(std::make_pair(CI, ReplVal.getValue()))
+                .second) {
           HasChanged = ChangeStatus::CHANGED;
+        }
 
         assert((!ICVReplacementValuesMap[ICV].lookup(CI) ||
                 ICVReplacementValuesMap[ICV].lookup(CI) == ReplVal) &&
                "Getter should be either not mapped or mapped to ReplVal");
 
         return false;
       };
 
+      auto CheckReturnInst = [&](Instruction &I) {
+        Optional<Value *> ReplVal = getReplacementValue(ICV, &I, A);
+        if (ReplVal.hasValue() &&
+            ICVReplacementValuesMap[ICV]
+                .insert(std::make_pair(&I, ReplVal.getValue()))
+                .second)
+          HasChanged = ChangeStatus::CHANGED;
+
+        return true;
+      };
+
+      // Track all changes of an ICV.
       SetterRFI.foreachUse(TrackValues, F);
+      // Map all getters to a value if possible.
       GetterRFI.foreachUse(MapReplacementValues, F);
+      // Get ICV value at every return instruction. If all return instruction
+      // have the same value, that is the ICV value after this function call.
+      A.checkForAllInstructions(CheckReturnInst, *this, {Instruction::Ret},
+                                /* CheckBBLivenessOnly */ true);
     }
 
     return HasChanged;
   }
 
   /// Check if \p I is a call and whether it changes the ICV.
   bool callChangesICV(Attributor &A, const Instruction *I,
                       InternalControlVar &ICV) const {
     const auto *CB = dyn_cast<CallBase>(I);
     if (!CB)
       return false;
 
bool callChangesICV(Attributor &A, const Instruction *I,
     // Since we don't know, assume it changes the ICV.
     if (CalledFunction->isDeclaration())
       return true;
 
     const auto &ICVTrackingAA =
         A.getAAFor<AAICVTracker>(*this, IRPosition::function(*CalledFunction));
 
     if (ICVTrackingAA.isAssumedTracked())
-      return !ICVTrackingAA.hasTrackedValue(ICV) &&
-             ICVTrackingAA.hasUnknownCall(ICV, A);
+      return ICVTrackingAA.hasTrackedValue(ICV);
 
     return true;
   }
 
-  /// Return the value with which \p I can be replaced for specific \p ICV.
-  Value *getReplacementValue(InternalControlVar ICV, const Instruction *I,
-                             Attributor &A) override {
-    if (ICVReplacementValuesMap[ICV].count(I))
-      return ICVReplacementValuesMap[ICV].lookup(I);
-
-    const BasicBlock *CurrBB = I->getParent();
-
-    auto &ValuesSet = ICVValuesMap[ICV];
-
-    for (const auto &ICVVal : ValuesSet) {
-      if (CurrBB == ICVVal.Inst->getParent()) {
-        if (!ICVVal.Inst->comesBefore(I))
-          continue;
+  Optional<Value *> getUniqueReplacementValue(InternalControlVar ICV,
+                                              const Instruction *I,
+                                              Attributor &A) const override {
+    const auto &CB = cast<CallBase>(*I);
+    Function *CalledFunction = CB.getCalledFunction();
 
-        // both instructions are in the same BB and at \p I we know the ICV
-        // value.
-        const Instruction *CurrInst = I;
-        while (CurrInst != ICVVal.Inst) {
-          // we don't yet know if a call might update an ICV.
-          // TODO: check callsite AA for value.
-          if (callChangesICV(A, CurrInst, ICV))
-            return nullptr;
+    if (CalledFunction->isDeclaration())
 
-          CurrInst = CurrInst->getPrevNode();
-        }
+    const auto &ICVTrackingAA =
+        A.getAAFor<AAICVTracker>(*this, IRPosition::function(*CalledFunction));
+
+    if (ICVTrackingAA.isAssumedTracked())
+      return ICVTrackingAA.getUniqueValue(ICV, A);
+
+    return nullptr;
 
-        // No call in between, return the value.
-        return ICVVal.TrackedValue;
+  Optional<Value *> getUniqueValue(InternalControlVar ICV,
+                                   Attributor &A) const override {
+    Optional<Value *> UniqueICVValue;
+    auto CheckReturnInst = [&](Instruction &I) {
+      Optional<Value *> ReplVal = getReplacementValue(ICV, &I, A);
+
+      // If we found a second ICV value there is no unique returned value.
+      if (UniqueICVValue.hasValue() && UniqueICVValue != ReplVal) {
+        UniqueICVValue = nullptr;
+        return false;
       }
 
-      auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
-      auto &Explorer = OMPInfoCache.getMustBeExecutedContextExplorer();
-      for (const BasicBlock *Pred : predecessors(CurrBB)) {
-        if (ICVVal.Inst->getParent() == Pred) {
-          if (!Explorer.findInContextOf(ICVVal.Inst, I))
-            return nullptr;
-          Instruction *CurrInst = ICVVal.Inst->getNextNode();
-          while (CurrInst != Pred->getTerminator()) {
-            // If any of the calls after the tracked value, might change the
-            // ICV, we don't know the value.
-            if (callChangesICV(A, CurrInst, ICV))
-              return nullptr;
-            CurrInst = CurrInst->getNextNode();
-          }
-          return ICVVal.TrackedValue;
+      UniqueICVValue = ReplVal;
+
+      return true;
+    };
+
+    if (!A.checkForAllInstructions(CheckReturnInst, *this, {Instruction::Ret},
+                                   /* CheckBBLivenessOnly */ true))
+      UniqueICVValue = nullptr;
+
+    return UniqueICVValue;
+
+  /// Return the value with which \p I can be replaced for specific \p ICV.
+  Optional<Value *> getReplacementValue(InternalControlVar ICV,
+                                        const Instruction *I,
+                                        Attributor &A) const override {
+    const auto &ValuesMap = ICVReplacementValuesMap[ICV];
+    if (ValuesMap.count(I))
+      return ValuesMap.lookup(I);
+
+    SmallVector<const Instruction *, 16> Worklist;
+    SmallPtrSet<const Instruction *, 16> Visited;
+    Worklist.push_back(I);
+
+    Optional<Value *> ReplVal;
+
+    while (!Worklist.empty()) {
+      const Instruction *CurrInst = Worklist.pop_back_val();
+      if (!Visited.insert(CurrInst).second)
+        continue;
+
+      const BasicBlock *CurrBB = CurrInst->getParent();
+
+      // Go up and look for all potential setters/calls that might change the
+      // ICV.
+      while ((CurrInst = CurrInst->getPrevNode())) {
+        if (ValuesMap.count(CurrInst)) {
+          Optional<Value *> NewReplVal = ValuesMap.lookup(CurrInst);
+          // Unknown value, track new.
+          if (!ReplVal.hasValue()) {
+            ReplVal = NewReplVal;
+            break;
           }
 
-        // if any of the calls in a block that doesn't have tracked value might
-        // change the ICV, we don't know the value.
-        for (const Instruction &Inst : *Pred)
-          if (callChangesICV(A, &Inst, ICV))
+          // If we found a new value, we can't know the icv value anymore.
+          if (NewReplVal.hasValue())
+            if (ReplVal != NewReplVal)
               return nullptr;
+
+          break;
         }
+
+        if (!callChangesICV(A, CurrInst, ICV))
+          continue;
+
+        Optional<Value *> NewReplVal =
+            getUniqueReplacementValue(ICV, CurrInst, A);
+
+        // Unknown value, track new.
+        if (!ReplVal.hasValue()) {
+          ReplVal = NewReplVal;
+          break;
         }
 
-    // No value was tracked.
+        // If we found a new value, we can't know the icv value anymore.
+        if (NewReplVal.hasValue())
+          if (ReplVal != NewReplVal)
             return nullptr;
       }
+
+      // If we are in the same BB and we have a value, we are done.
+      if (CurrBB == I->getParent() && ReplVal.hasValue())
+        return ReplVal;
+
+      // If we are in the Entry BB and we don't have a value, we are done.
+      if (CurrBB == &(CurrBB->getParent()->getEntryBlock()))
+        if (!ReplVal.hasValue() && hasTrackedValue(ICV))
+          ReplVal = nullptr;
+
+      // Go through all predecessors and add terminators for analysis.
+      for (const BasicBlock *Pred : predecessors(CurrBB))
+        if (const Instruction *Terminator = Pred->getTerminator())
+          Worklist.push_back(Terminator);
+    }
+
+    // Couldn't find a replacement value, ICV value unknown.
+    if (!ReplVal.hasValue())
+      return None;
+
+    return ReplVal.getValue();
+  }
 };
 } // namespace
 
 const char AAICVTracker::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
for (LazyCallGraph::Node &N : C) {
     SCC.push_back(Fn);
 
     // Do we already know that the SCC contains kernels,
     // or that OpenMP functions are called from this SCC?
     if (SCCIsInteresting)
       continue;
     // If not, let's check that.
     SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
   }
 
   if (!SCCIsInteresting || SCC.empty())
     return PreservedAnalyses::all();
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
 
   AnalysisGetter AG(FAM);
 
   auto OREGetter = [&FAM](Function *F) -> OptimizationRemarkEmitter & {
     return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*F);
   };
 
   CallGraphUpdater CGUpdater;
   CGUpdater.initialize(CG, C, AM, UR);
 
   SetVector<Function *> Functions(SCC.begin(), SCC.end());
   BumpPtrAllocator Allocator;
   OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
                                 /*CGSCC*/ Functions, OMPInModule.getKernels());
 
   SetVector<Function *> ModuleSlice(InfoCache.ModuleSlice.begin(),
                                     InfoCache.ModuleSlice.end());
   Attributor A(ModuleSlice, InfoCache, CGUpdater);
 
   // TODO: Compute the module slice we are allowed to look at.
   OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run();
   if (Changed)
     return PreservedAnalyses::none();
 
   return PreservedAnalyses::all();
 }
 
 namespace {
 
 struct OpenMPOptLegacyPass : public CallGraphSCCPass {
   CallGraphUpdater CGUpdater;
   OpenMPInModule OMPInModule;
   static char ID;
 
   OpenMPOptLegacyPass() : CallGraphSCCPass(ID) {
     initializeOpenMPOptLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
bool runOnSCC(CallGraphSCC &CGSCC) override {
     SmallVector<Function *, 16> SCC;
     // If there are kernels in the module, we have to run on all SCC's.
     bool SCCIsInteresting = !OMPInModule.getKernels().empty();
     for (CallGraphNode *CGN : CGSCC) {
       Function *Fn = CGN->getFunction();
       if (!Fn || Fn->isDeclaration())
         continue;
       SCC.push_back(Fn);
 
       // Do we already know that the SCC contains kernels,
       // or that OpenMP functions are called from this SCC?
       if (SCCIsInteresting)
         continue;
       // If not, let's check that.
       SCCIsInteresting |= OMPInModule.containsOMPRuntimeCalls(Fn);
     }
 
     if (!SCCIsInteresting || SCC.empty())
       return false;
 

llvm/test/Transforms/OpenMP/icv_tracking.ll

-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --scrub-attributes --check-attributes
 ; RUN: opt -S -openmpopt < %s | FileCheck %s
 ; RUN: opt -S -passes=openmpopt < %s | FileCheck %s
 
 %struct.ident_t = type { i32, i32, i32, i32, i8* }
 
 @.str = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
 @0 = private unnamed_addr global %struct.ident_t { i32 0, i32 2, i32 0, i32 0, i8* getelementptr inbounds ([23 x i8], [23 x i8]* @.str, i32 0, i32 0) }, align 8
 
3:                                                ; preds = %1
   br label %5
 
 5:                                                ; preds = %3, %1
   %6 = call i32 @omp_get_max_threads()
   call void @use(i32 %6)
   ret void
 }
 
+define weak void @weak_known_unique_icv(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@weak_known_unique_icv
+; CHECK-SAME: (i1 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 2)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP5:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[TMP7:%.*]] = call i32 @icv_free_use(i32 [[TMP6]])
+; CHECK-NEXT:    ret void
+;
+  call void @omp_set_num_threads(i32 2)
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = call i32 @icv_free_use(i32 10)
+  br label %5
+
+5:                                                ; preds = %3, %1
+  %6 = call i32 @omp_get_max_threads()
+  %7 = call i32 @icv_free_use(i32 %6)
+  ret void
+}
+
+define void @known_unique_icv(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@known_unique_icv
+; CHECK-SAME: (i1 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 2)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP5:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @icv_free_use(i32 2)
+; CHECK-NEXT:    ret void
+;
+  call void @omp_set_num_threads(i32 2)
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = call i32 @icv_free_use(i32 10)
+  br label %5
+
+5:                                                ; preds = %3, %1
+  %6 = call i32 @omp_get_max_threads()
+  %7 = call i32 @icv_free_use(i32 %6)
+  ret void
+}
+
+define i32 @no_unique_icv(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@no_unique_icv
+; CHECK-SAME: (i1 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
+; CHECK-NEXT:    br i1 [[TMP0]], label [[TMP3:%.*]], label [[TMP2:%.*]]
+; CHECK:       2:
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 2)
+; CHECK-NEXT:    br label [[TMP3]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    ret i32 [[TMP4]]
+;
+  call void @omp_set_num_threads(i32 4)
+  br i1 %0, label %3, label %2
+
+2:                                                ; preds = %1
+  call void @omp_set_num_threads(i32 2)
+  br label %3
+
+3:                                                ; preds = %1, %2
+  %4 = call i32 @omp_get_max_threads()
+  ret i32 %4
+}
+
 define void @test2(i1 %0) {
 ; CHECK-LABEL: define {{[^@]+}}@test2
 ; CHECK-SAME: (i1 [[TMP0:%.*]])
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
 ; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP4:%.*]], label [[TMP3:%.*]]
 ; CHECK:       3:
 ; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
 ; CHECK-NEXT:    br label [[TMP4]]
 ; CHECK:       4:
 ; CHECK-NEXT:    [[TMP5:%.*]] = call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    call void @use(i32 [[TMP5]])
 ; CHECK-NEXT:    ret void
 ;
-  %2 = icmp eq i1 %0, 0
   br i1 %2, label %4, label %3
 
 3:                                                ; preds = %1
   call void @omp_set_num_threads(i32 4)
   br label %4
 
 4:                                                ; preds = %3, %1
   %5 = call i32 @omp_get_max_threads()
   call void @use(i32 %5)
   ret void
 }
 
+define void @test3(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@test3
+; CHECK-SAME: (i1 [[TMP0:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP4:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
+; CHECK-NEXT:    br label [[TMP4]]
+; CHECK:       4:
+; CHECK-NEXT:    call void @weak_known_unique_icv(i1 [[TMP0]])
+; CHECK-NEXT:    [[TMP5:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @icv_free_use(i32 [[TMP5]])
+; CHECK-NEXT:    ret void
+;
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %4, label %3
+
+3:                                                ; preds = %1
+  call void @omp_set_num_threads(i32 4)
+  br label %4
+
+4:                                                ; preds = %3, %1
+  call void @weak_known_unique_icv(i1 %0)
+  %5 = call i32 @omp_get_max_threads()
+  %6 = call i32 @icv_free_use(i32 %5)
+  ret void
+}
+
+declare void @__cxa_rethrow()
+
+define i32 @maybe_throw(i1 zeroext %0) {
+; CHECK-LABEL: define {{[^@]+}}@maybe_throw
+; CHECK-SAME: (i1 zeroext [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
+; CHECK-NEXT:    br i1 [[TMP0]], label [[TMP2:%.*]], label [[TMP3:%.*]]
+; CHECK:       2:
+; CHECK-NEXT:    tail call void @__cxa_rethrow()
+; CHECK-NEXT:    unreachable
+; CHECK:       3:
+; CHECK-NEXT:    ret i32 -1
+;
+  call void @omp_set_num_threads(i32 4)
+  br i1 %0, label %2, label %3
+
+2:                                                ; preds = %1
+  tail call void @__cxa_rethrow() #1
+  unreachable
+
+3:                                                ; preds = %1
+  ret i32 -1
+}
+
+define void @test4(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@test4
+; CHECK-SAME: (i1 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @known_unique_icv(i1 [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP4:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[VAL:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP4]]
+; CHECK:       4:
+; CHECK-NEXT:    call void @use(i32 2)
+; CHECK-NEXT:    [[TMP5:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @no_unique_icv(i1 [[TMP0]])
+; CHECK-NEXT:    call void @use(i32 [[TMP5]])
+; CHECK-NEXT:    ret void
+;
+  call void @known_unique_icv(i1 %0)
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %4, label %3
+
+3:                                                ; preds = %1
+  %val = call i32 @icv_free_use(i32 10)
+  br label %4
+
+4:                                                ; preds = %3, %1
+  %5 = call i32 @omp_get_max_threads()
+  call void @use(i32 %5)
+  %6 = call i32 @omp_get_max_threads()
+  call i32 @no_unique_icv(i1 %0)
+  call void @use(i32 %6)
+  ret void
+}
+
+define void @test4_invoke(i1 %0) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: define {{[^@]+}}@test4_invoke
+; CHECK-SAME: (i1 [[TMP0:%.*]]) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*)
+; CHECK-NEXT:    call void @known_unique_icv(i1 [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = invoke i32 @maybe_throw(i1 [[TMP0]])
+; CHECK-NEXT:    to label [[CONT:%.*]] unwind label [[EXC:%.*]]
+; CHECK:       cont:
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP3]], label [[TMP5:%.*]], label [[TMP4:%.*]]
+; CHECK:       exc:
+; CHECK-NEXT:    [[LP:%.*]] = landingpad { i8*, i32 }
+; CHECK-NEXT:    filter [0 x i8*] zeroinitializer
+; CHECK-NEXT:    unreachable
+; CHECK:       4:
+; CHECK-NEXT:    [[VAL:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    call void @use(i32 2)
+; CHECK-NEXT:    ret void
+;
+  call void @known_unique_icv(i1 %0)
+  invoke i32 @maybe_throw(i1 %0)
+  to label %cont unwind label %exc
+
+cont:
+  %3 = icmp eq i1 %0, 0
+  br i1 %3, label %5, label %4
+
+exc:
+  %lp = landingpad { i8*, i32 }
+  filter [0 x i8*] zeroinitializer
+  unreachable
+
+4:                                                ; preds = %1
+  %val = call i32 @icv_free_use(i32 10)
+  br label %5
+
+5:                                                ; preds = %3, %1
+  %6 = call i32 @omp_get_max_threads()
+  call void @use(i32 %6)
+  ret void
+}
+
+define i32 @test5(i32 %0)  #0 {
+; CHECK-LABEL: define {{[^@]+}}@test5
+; CHECK-SAME: (i32 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i32 [[TMP0]], 3
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP4:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    call void @use(i32 4)
+; CHECK-NEXT:    br label [[TMP12:%.*]]
+; CHECK:       4:
+; CHECK-NEXT:    [[TMP5:%.*]] = icmp sgt i32 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TMP5]], label [[TMP6:%.*]], label [[TMP8:%.*]]
+; CHECK:       6:
+; CHECK-NEXT:    [[TMP7:%.*]] = call i32 @icv_free_use(i32 [[TMP0]])
+; CHECK-NEXT:    br label [[TMP15:%.*]]
+; CHECK:       8:
+; CHECK-NEXT:    [[TMP9:%.*]] = icmp eq i32 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TMP9]], label [[TMP10:%.*]], label [[TMP12]]
+; CHECK:       10:
+; CHECK-NEXT:    [[TMP11:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP15]]
+; CHECK:       12:
+; CHECK-NEXT:    [[TMP13:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    [[TMP14:%.*]] = call i32 @icv_free_use(i32 [[TMP13]])
+; CHECK-NEXT:    br label [[TMP15]]
+; CHECK:       15:
+; CHECK-NEXT:    [[TMP16:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[TMP17:%.*]] = call i32 @icv_free_use(i32 [[TMP16]])
+; CHECK-NEXT:    ret i32 [[TMP17]]
+;
+  call void @omp_set_num_threads(i32 4)
+  %2 = icmp sgt i32 %0, 3
+  br i1 %2, label %3, label %5
+
+3:
+  %4 = call i32 @omp_get_max_threads()
+  call void @use(i32 %4)
+  br label %13
+
+5:
+  %6 = icmp sgt i32 %0, 0
+  br i1 %6, label %7, label %9
+
+7:
+  %8 = call i32 @icv_free_use(i32 %0)
+  br label %16
+
+9:
+  %10 = icmp eq i32 %0, 0
+  br i1 %10, label %11, label %13
+
+11:
+  %12 = call i32 @icv_free_use(i32 10)
+  br label %16
+
+13:
+  %14 = add nsw i32 %0, 1
+  %15 = call i32 @icv_free_use(i32 %14)
+  br label %16
+
+16:
+  %17 = call i32 @omp_get_max_threads()
+  %18 = call i32 @icv_free_use(i32 %17)
+  ret i32 %18
+}
+
+define i32 @test6(i32 %0) {
+; CHECK-LABEL: define {{[^@]+}}@test6
+; CHECK-SAME: (i32 [[TMP0:%.*]])
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i32 [[TMP0]], 3
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP3:%.*]], label [[TMP5:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    br label [[TMP16:%.*]]
+; CHECK:       5:
+; CHECK-NEXT:    [[TMP6:%.*]] = icmp sgt i32 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TMP6]], label [[TMP7:%.*]], label [[TMP9:%.*]]
+; CHECK:       7:
+; CHECK-NEXT:    [[TMP8:%.*]] = call i32 @icv_free_use(i32 [[TMP0]])
+; CHECK-NEXT:    br label [[TMP16]]
+; CHECK:       9:
+; CHECK-NEXT:    [[TMP10:%.*]] = icmp eq i32 [[TMP0]], 0
+; CHECK-NEXT:    br i1 [[TMP10]], label [[TMP11:%.*]], label [[TMP13:%.*]]
+; CHECK:       11:
+; CHECK-NEXT:    [[TMP12:%.*]] = call i32 @icv_free_use(i32 5)
+; CHECK-NEXT:    br label [[TMP16]]
+; CHECK:       13:
+; CHECK-NEXT:    [[TMP14:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    [[TMP15:%.*]] = call i32 @icv_free_use(i32 [[TMP14]])
+; CHECK-NEXT:    br label [[TMP16]]
+; CHECK:       16:
+; CHECK-NEXT:    [[TMP17:%.*]] = call i32 @icv_free_use(i32 4)
+; CHECK-NEXT:    ret i32 [[TMP17]]
+;
+  call void @omp_set_num_threads(i32 4)
+  %2 = icmp sgt i32 %0, 3
+  br i1 %2, label %3, label %5
+
+3:                                                ; preds = %1
+  %4 = call i32 @icv_free_use(i32 10)
+  br label %16
+
+5:                                                ; preds = %1
+  %6 = icmp sgt i32 %0, 0
+  br i1 %6, label %7, label %9
+
+7:                                                ; preds = %5
+  %8 = call i32 @icv_free_use(i32 %0)
+  br label %16
+
+9:                                                ; preds = %5
+  %10 = icmp eq i32 %0, 0
+  br i1 %10, label %11, label %13
+
+11:                                               ; preds = %9
+  %12 = call i32 @icv_free_use(i32 5)
+  br label %16
+
+13:                                               ; preds = %9
+  %14 = add nsw i32 %0, 1
+  %15 = call i32 @icv_free_use(i32 %14)
+  br label %16
+
+16:                                               ; preds = %7, %13, %11, %3
+  %17 = call i32 @omp_get_max_threads()
+  %18 = call i32 @icv_free_use(i32 %17)
+  ret i32 %18
+}
+
+declare i32 @__gxx_personality_v0(...)
+
 !0 = !{!1}
 !1 = !{i64 2, i64 -1, i64 -1, i1 true}