[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

158

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

1173

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 ;)

1263

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

1277–1279

UniqueReplVal should not be needed anywhere here.

1365

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

155

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

1170

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.

1173

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

1170

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

1420–1423

1437

Never reset Changed, see also below

1444

leftover

1490–1493

Never reset Changed

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1336

Here and below, we should cache this, right?

1351–1354

leftover.

1364

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

1367–1559

Caching.

[sstefan1]

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

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1364

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

945

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

1240

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

1252

What is this about?

1256

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.

1273

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

1458

This seems unfortunate but OK for now.

[sstefan1]

addressing comments

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1252

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.

1273

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,
LLVM_DEBUG(dbgs() << "[Attributor] Done with " << SCC.size()
                       << " functions, result: " << Changed << ".\n");
 
     return Changed == ChangeStatus::CHANGED;
   }
 
   /// Populate the Attributor with abstract attribute opportunities in the
   /// function.
   void registerAAs() {
+    if (SCC.empty())
+      return;
+
     for (Function *F : SCC) {
       if (F->isDeclaration())
         continue;
 
-      A.getOrCreateAAFor<AAICVTracker>(IRPosition::function(*F));
+      // Create CallSite AA for all Getters.
+      for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
+        auto ICVInfo = OMPInfoCache.ICVs[static_cast<InternalControlVar>(Idx)];
+
+        auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
+
+        auto CreateAA = [&](Use &U, Function &Caller) {
+          CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &GetterRFI);
+          if (!CI)
+            return false;
+
+          auto &CB = cast<CallBase>(*CI);
+
+          IRPosition CBPos = IRPosition::callsite_function(CB);
+          A.getOrCreateAAFor<AAICVTracker>(CBPos);
+          return false;
+        };
+
+        GetterRFI.foreachUse(CreateAA, F);
+      }
     }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
   {
bool OpenMPOpt::rewriteDeviceCodeStateMachine() {
   return Changed;
 }
 
 /// Abstract Attribute for tracking ICV values.
 struct AAICVTracker : public StateWrapper<BooleanState, AbstractAttribute> {
   using Base = StateWrapper<BooleanState, AbstractAttribute>;
   AAICVTracker(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
 
+  void initialize(Attributor &A) override {
+    Function *F = getAnchorScope();
+    if (!F || !A.isFunctionIPOAmendable(*F))
+      indicatePessimisticFixpoint();
+  }
+
   /// Returns true if value is assumed to be tracked.
   bool isAssumedTracked() const { return getAssumed(); }
 
   /// 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 {
+    return None;
+  }
+
+  /// 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 = 0;
+
+  // Currently only nthreads is being tracked.
+  // this array will only grow with time.
+  InternalControlVar TrackableICVs[1] = {ICV_nthreads};
 
   /// 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
   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) {}
 
   // FIXME: come up with better string.
-  const std::string getAsStr() const override { return "ICVTracker"; }
+  const std::string getAsStr() const override { return "ICVTrackerFunction"; }
 
   // FIXME: come up with some stats.
   void trackStatistics() const override {}
 
-  /// TODO: decide whether to deduplicate here, or use current
-  /// deduplicateRuntimeCalls function.
+  /// We don't manifest anything for this AA.
   ChangeStatus manifest(Attributor &A) override {
-    ChangeStatus Changed = ChangeStatus::UNCHANGED;
-
-    for (InternalControlVar &ICV : TrackableICVs)
-      if (deduplicateICVGetters(ICV, A))
-        Changed = ChangeStatus::CHANGED;
-
-    return Changed;
-  }
-
-  bool deduplicateICVGetters(InternalControlVar &ICV, Attributor &A) {
-    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
-    auto &ICVInfo = OMPInfoCache.ICVs[ICV];
-    auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
-
-    bool Changed = false;
-
-    auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
-      CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &GetterRFI);
-      Instruction *UserI = cast<Instruction>(U.getUser());
-      Value *ReplVal = getReplacementValue(ICV, UserI, A);
-
-      if (!ReplVal || !CI)
-        return false;
-
-      A.removeCallSite(CI);
-      CI->replaceAllUsesWith(ReplVal);
-      CI->eraseFromParent();
-      Changed = true;
-      return true;
-    };
-
-    GetterRFI.foreachUse(ReplaceAndDeleteCB, getAnchorScope());
-    return Changed;
+    return ChangeStatus::UNCHANGED;
   }
 
   // Map of ICV to their values at specific program point.
-  EnumeratedArray<SmallSetVector<ICVValue, 4>, InternalControlVar,
+  EnumeratedArray<DenseMap<Instruction *, Value *>, InternalControlVar,
                   InternalControlVar::ICV___last>
-      ICVValuesMap;
-
-  // Currently only nthreads is being tracked.
-  // this array will only grow with time.
-  InternalControlVar TrackableICVs[1] = {ICV_nthreads};
+      ICVReplacementValuesMap;
 
   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 &ValuesMap = ICVReplacementValuesMap[ICV];
       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 (ValuesMap.insert(std::make_pair(CI, CI->getArgOperand(0))).second)
           HasChanged = ChangeStatus::CHANGED;
 
         return false;
       };
 
+      auto CallCheck = [&](Instruction &I) {
+        Optional<Value *> ReplVal = getValueForCall(A, &I, ICV);
+        if (ReplVal.hasValue()) {
+          ValuesMap.insert(std::make_pair(&I, *ReplVal));
+          HasChanged = ChangeStatus::CHANGED;
+        }
+        return true;
+      };
+
+      // Track all changes of an ICV.
       SetterRFI.foreachUse(TrackValues, F);
+
+      A.checkForAllInstructions(CallCheck, *this, {Instruction::Call},
+                                /* CheckBBLivenessOnly */ true);
+
+      Instruction *Entry = &F->getEntryBlock().front();
+      if (HasChanged == ChangeStatus::CHANGED && !ValuesMap.count(Entry))
+        ValuesMap.insert(std::make_pair(Entry, nullptr));
     }
 
     return HasChanged;
   }
 
-  /// Return the value with which \p I can be replaced for specific \p ICV.
-  Value *getReplacementValue(InternalControlVar ICV, const Instruction *I,
-                             Attributor &A) override {
-    const BasicBlock *CurrBB = I->getParent();
+  /// Hepler to check if \p I is a call and get the value for it if it is
+  /// unique.
+  Optional<Value *> getValueForCall(Attributor &A, const Instruction *I,
+                                    InternalControlVar &ICV) const {
+
+    const auto *CB = dyn_cast<CallBase>(I);
+    if (!CB)
+      return None;
 
-    auto &ValuesSet = ICVValuesMap[ICV];
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
+    auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
+    Function *CalledFunction = CB->getCalledFunction();
 
-    for (const auto &ICVVal : ValuesSet) {
-      if (CurrBB == ICVVal.Inst->getParent()) {
-        if (!ICVVal.Inst->comesBefore(I))
-          continue;
+    if (CalledFunction == GetterRFI.Declaration)
+      return None;
+    if (CalledFunction == SetterRFI.Declaration) {
+      if (ICVReplacementValuesMap[ICV].count(I))
+        return ICVReplacementValuesMap[ICV].lookup(I);
+
+      return nullptr;
+    }
 
-        // both instructions are in the same BB and at \p I we know the ICV
-        // value.
-        while (I != ICVVal.Inst) {
-          // we don't yet know if a call might update an ICV.
-          // TODO: check callsite AA for value.
-          if (const auto *CB = dyn_cast<CallBase>(I))
-            if (CB->getCalledFunction() != GetterRFI.Declaration)
-              return nullptr;
+    // Since we don't know, assume it changes the ICV.
+    if (CalledFunction->isDeclaration())
 
-          I = I->getPrevNode();
-        }
+    const auto &ICVTrackingAA =
+        A.getAAFor<AAICVTracker>(*this, IRPosition::callsite_returned(*CB));
+
+    if (ICVTrackingAA.isAssumedTracked())
+      return ICVTrackingAA.getUniqueReplacementValue(ICV);
+
+    // If we don't know, assume it changes.
+    return nullptr;
 
-        // No call in between, return the value.
-        return ICVVal.TrackedValue;
-      }
+  // We don't check unique value for a function, so return None.
+  Optional<Value *>
+  getUniqueReplacementValue(InternalControlVar ICV) const override {
+    return None;
+
+  /// 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;
           }
 
-    // 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;
+
+          break;
+        }
+
+        Optional<Value *> NewReplVal = getValueForCall(A, CurrInst, ICV);
+        if (!NewReplVal.hasValue())
+          continue;
+
+        // Unknown value, track new.
+        if (!ReplVal.hasValue()) {
+          ReplVal = NewReplVal;
+          break;
+        }
+
+        // if (NewReplVal.hasValue())
+        // We found a new value, we can't know the icv value anymore.
+        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;
+
+      // 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);
+    }
+
+    return ReplVal;
+  }
+};
+
+struct AAICVTrackerFunctionReturned : AAICVTracker {
+  AAICVTrackerFunctionReturned(const IRPosition &IRP, Attributor &A)
+      : AAICVTracker(IRP, A) {}
+
+  // FIXME: come up with better string.
+  const std::string getAsStr() const override {
+    return "ICVTrackerFunctionReturned";
+  }
+
+  // FIXME: come up with some stats.
+  void trackStatistics() const override {}
+
+  /// We don't manifest anything for this AA.
+  ChangeStatus manifest(Attributor &A) override {
+    return ChangeStatus::UNCHANGED;
+  }
+
+  // Map of ICV to their values at specific program point.
+  EnumeratedArray<Optional<Value *>, InternalControlVar,
+                  InternalControlVar::ICV___last>
+      ICVReplacementValuesMap;
+
+  /// Return the value with which \p I can be replaced for specific \p ICV.
+  Optional<Value *>
+  getUniqueReplacementValue(InternalControlVar ICV) const override {
+    return ICVReplacementValuesMap[ICV];
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+    const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
+        *this, IRPosition::function(*getAnchorScope()));
+
+    if (!ICVTrackingAA.isAssumedTracked())
+      return indicatePessimisticFixpoint();
+
+    for (InternalControlVar ICV : TrackableICVs) {
+      Optional<Value *> &ReplVal = ICVReplacementValuesMap[ICV];
+      Optional<Value *> UniqueICVValue;
+
+      auto CheckReturnInst = [&](Instruction &I) {
+        Optional<Value *> NewReplVal =
+            ICVTrackingAA.getReplacementValue(ICV, &I, A);
+
+        // If we found a second ICV value there is no unique returned value.
+        if (UniqueICVValue.hasValue() && UniqueICVValue != NewReplVal)
+          return false;
+
+        UniqueICVValue = NewReplVal;
+
+        return true;
+      };
+
+      if (!A.checkForAllInstructions(CheckReturnInst, *this, {Instruction::Ret},
+                                     /* CheckBBLivenessOnly */ true))
+        UniqueICVValue = nullptr;
+
+      if (UniqueICVValue == ReplVal)
+        continue;
+
+      ReplVal = UniqueICVValue;
+      Changed = ChangeStatus::CHANGED;
+    }
+
+    return Changed;
+  }
+};
+
+struct AAICVTrackerCallSite : AAICVTracker {
+  AAICVTrackerCallSite(const IRPosition &IRP, Attributor &A)
+      : AAICVTracker(IRP, A) {}
+
+  void initialize(Attributor &A) override {
+    Function *F = getAnchorScope();
+    if (!F || !A.isFunctionIPOAmendable(*F))
+      indicatePessimisticFixpoint();
+
+    // We only initialize this AA for getters, so we need to know which ICV it
+    // gets.
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    for (InternalControlVar ICV : TrackableICVs) {
+      auto ICVInfo = OMPInfoCache.ICVs[ICV];
+      auto &Getter = OMPInfoCache.RFIs[ICVInfo.Getter];
+      if (Getter.Declaration == getAssociatedFunction()) {
+        AssociatedICV = ICVInfo.Kind;
+        return;
+      }
+    }
+
+    /// Unknown ICV.
+    indicatePessimisticFixpoint();
+  }
+
+  ChangeStatus manifest(Attributor &A) override {
+    if (!ReplVal.hasValue() || !ReplVal.getValue())
+      return ChangeStatus::UNCHANGED;
+
+    A.changeValueAfterManifest(*getCtxI(), **ReplVal);
+    A.deleteAfterManifest(*getCtxI());
+
+    return ChangeStatus::CHANGED;
+  }
+
+  // FIXME: come up with better string.
+  const std::string getAsStr() const override { return "ICVTrackerCallSite"; }
+
+  // FIXME: come up with some stats.
+  void trackStatistics() const override {}
+
+  InternalControlVar AssociatedICV;
+  Optional<Value *> ReplVal;
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
+        *this, IRPosition::function(*getAnchorScope()));
+
+    // We don't have any information, so we assume it changes the ICV.
+    if (!ICVTrackingAA.isAssumedTracked())
+      return indicatePessimisticFixpoint();
+
+    Optional<Value *> NewReplVal =
+        ICVTrackingAA.getReplacementValue(AssociatedICV, getCtxI(), A);
+
+    if (ReplVal == NewReplVal)
+      return ChangeStatus::UNCHANGED;
+
+    ReplVal = NewReplVal;
+    return ChangeStatus::CHANGED;
+  }
+
+  // Return the value with which associated value can be replaced for specific
+  // \p ICV.
+  Optional<Value *>
+  getUniqueReplacementValue(InternalControlVar ICV) const override {
+    return ReplVal;
+  }
+};
+
+struct AAICVTrackerCallSiteReturned : AAICVTracker {
+  AAICVTrackerCallSiteReturned(const IRPosition &IRP, Attributor &A)
+      : AAICVTracker(IRP, A) {}
+
+  // FIXME: come up with better string.
+  const std::string getAsStr() const override {
+    return "ICVTrackerCallSiteReturned";
+  }
+
+  // FIXME: come up with some stats.
+  void trackStatistics() const override {}
+
+  /// We don't manifest anything for this AA.
+  ChangeStatus manifest(Attributor &A) override {
+    return ChangeStatus::UNCHANGED;
+  }
+
+  // Map of ICV to their values at specific program point.
+  EnumeratedArray<Optional<Value *>, InternalControlVar,
+                  InternalControlVar::ICV___last>
+      ICVReplacementValuesMap;
+
+  /// Return the value with which associated value can be replaced for specific
+  /// \p ICV.
+  Optional<Value *>
+  getUniqueReplacementValue(InternalControlVar ICV) const override {
+    return ICVReplacementValuesMap[ICV];
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+    const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
+        *this, IRPosition::returned(*getAssociatedFunction()));
+
+    // We don't have any information, so we assume it changes the ICV.
+    if (!ICVTrackingAA.isAssumedTracked())
+      return indicatePessimisticFixpoint();
+
+    for (InternalControlVar ICV : TrackableICVs) {
+      Optional<Value *> &ReplVal = ICVReplacementValuesMap[ICV];
+      Optional<Value *> NewReplVal =
+          ICVTrackingAA.getUniqueReplacementValue(ICV);
+
+      if (ReplVal == NewReplVal)
+        continue;
+
+      ReplVal = NewReplVal;
+      Changed = ChangeStatus::CHANGED;
+    }
+    return Changed;
   }
 };
 } // namespace
 
 const char AAICVTracker::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
   case IRPosition::IRP_ARGUMENT:
-  case IRPosition::IRP_RETURNED:
-  case IRPosition::IRP_CALL_SITE_RETURNED:
   case IRPosition::IRP_CALL_SITE_ARGUMENT:
-  case IRPosition::IRP_CALL_SITE:
-    llvm_unreachable("ICVTracker can only be created for function position!");
+    llvm_unreachable("ICVTracker: invalid IRPosition!");
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAICVTrackerFunction(IRP, A);
     break;
+  case IRPosition::IRP_RETURNED:
+    AA = new (A.Allocator) AAICVTrackerFunctionReturned(IRP, A);
+    break;
+  case IRPosition::IRP_CALL_SITE_RETURNED:
+    AA = new (A.Allocator) AAICVTrackerCallSiteReturned(IRP, A);
+    break;
+  case IRPosition::IRP_CALL_SITE:
+    AA = new (A.Allocator) AAICVTrackerCallSite(IRP, A);
+    break;
   }
 
   return *AA;
 }
 
 PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
                                      CGSCCAnalysisManager &AM,
                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
   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());
 
-  Attributor A(Functions, InfoCache, CGUpdater);
+  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();
bool runOnSCC(CallGraphSCC &CGSCC) override {
 
     AnalysisGetter AG;
     SetVector<Function *> Functions(SCC.begin(), SCC.end());
     BumpPtrAllocator Allocator;
     OMPInformationCache InfoCache(
         *(Functions.back()->getParent()), AG, Allocator,
         /*CGSCC*/ Functions, OMPInModule.getKernels());
 
-    Attributor A(Functions, InfoCache, CGUpdater);
+    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);
     return OMPOpt.run();
   }
 
   bool doFinalization(CallGraph &CG) override { return CGUpdater.finalize(); }
 };

llvm/test/Transforms/OpenMP/icv_tracking.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
 ; 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
 
+; doesn't modify any ICVs.
+define i32 @icv_free_use(i32 %0) {
+; CHECK-LABEL: define {{[^@]+}}@icv_free_use
+; CHECK-SAME: (i32 [[TMP0:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  %2 = add nsw i32 %0, 1
+  ret i32 %2
+}
+
+define i32 @bad_use(i32 %0) {
+; CHECK-LABEL: define {{[^@]+}}@bad_use
+; CHECK-SAME: (i32 [[TMP0:%.*]])
+; CHECK-NEXT:    tail call void @use(i32 [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  tail call void @use(i32 %0)
+  %2 = add nsw i32 %0, 1
+  ret i32 %2
+}
+
 define dso_local i32 @foo(i32 %0, i32 %1) {
 ; CHECK-LABEL: define {{[^@]+}}@foo
 ; CHECK-SAME: (i32 [[TMP0:%.*]], i32 [[TMP1:%.*]])
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 [[TMP0]])
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 [[TMP1]])
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP1]])
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP1]])
 ; CHECK-NEXT:    tail call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* nonnull @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined. to void (i32*, i32*, ...)*))
;
   tail call void @omp_set_num_threads(i32 10)
   %4 = tail call i32 @omp_get_max_threads()
   tail call void @use(i32 %4)
   %5 = tail call i32 @omp_get_max_threads()
   tail call void @use(i32 %5)
   ret void
 }
 
+define dso_local i32 @bar1(i32 %0, i32 %1) {
+; CHECK-LABEL: define {{[^@]+}}@bar1
+; CHECK-SAME: (i32 [[TMP0:%.*]], i32 [[TMP1:%.*]])
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP0]], [[TMP1]]
+; CHECK-NEXT:    [[TMP4:%.*]] = select i1 [[TMP3]], i32 [[TMP0]], i32 [[TMP1]]
+; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 [[TMP4]])
+; CHECK-NEXT:    tail call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* nonnull @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined..2 to void (i32*, i32*, ...)*))
+; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @omp_get_max_threads()
+; CHECK-NEXT:    tail call void @use(i32 [[TMP5]])
+; CHECK-NEXT:    ret i32 0
+;
+  %3 = icmp sgt i32 %0, %1
+  %4 = select i1 %3, i32 %0, i32 %1
+  tail call void @omp_set_num_threads(i32 %4)
+  %5 = tail call i32 @omp_get_max_threads()
+  tail call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* nonnull @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined..2 to void (i32*, i32*, ...)*))
+  %6 = tail call i32 @omp_get_max_threads()
+  tail call void @use(i32 %6)
+  ret i32 0
+}
+
+define internal void @.omp_outlined..2(i32* %0, i32*  %1) {
+; CHECK-LABEL: define {{[^@]+}}@.omp_outlined..2
+; CHECK-SAME: (i32* [[TMP0:%.*]], i32* [[TMP1:%.*]])
+; CHECK-NEXT:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @icv_free_use(i32 [[TMP3]])
+; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 10)
+; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    [[TMP6:%.*]] = tail call i32 @icv_free_use(i32 10)
+; CHECK-NEXT:    ret void
+;
+  %3 = tail call i32 @omp_get_max_threads()
+  %4 = tail call i32 @icv_free_use(i32 %3)
+  tail call void @omp_set_num_threads(i32 10)
+  %5 = tail call i32 @omp_get_max_threads()
+  %6 = tail call i32 @icv_free_use(i32 %5)
+  %7 = tail call i32 @omp_get_max_threads()
+  %8 = tail call i32 @icv_free_use(i32 %7)
+  ret void
+}
+define void @test(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@test
+; 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 [[TMP4:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    call void @use(i32 10)
+; 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
+;
+  call void @omp_set_num_threads(i32 2)
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %4, label %3
+
+3:                                                ; preds = %1
+  call void @use(i32 10)
+  br label %4
+
+4:                                                ; preds = %3, %1
+  %5 = call i32 @omp_get_max_threads()
+  call void @use(i32 %5)
+  ret void
+}
+
+define void @test1(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@test1
+; 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:    call void @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()
+  call void @use(i32 %6)
+  ret void
+}
+
+define void @bad_use_test(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@bad_use_test
+; 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 @bad_use(i32 10)
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    [[TMP6:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    call void @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 @bad_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 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}