[OpenMPOpt] ICV Tracking

[sstefan1]

Summary

This is first part if ICV tracking support in OpenMPOpt. Things that are included in this pathc:

• ICV information definition (not complete) in OMPKinds.def.
• OpenMP-specific information cache based on Attributor's InformationCache. This should should make it easier to share information between them.
• possibility of running the Attributor through runAttributor().
• AAICVTracking AbstractAttribute with basic tracking.

Since there are a lot of changes to agree up on, I left out the deduplication part for now. There are 2 possible places where that can be done: current deduplication functions, or within AAICVTracking's manifest.

This can also be treated as a RFC.

[jdoerfert]

I left a first set of comments. I think the overall direction is good, some interesting opportunities are opening up (see below). We need tests and also to split the patch (some suggestions below).

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

367 ↗

(On Diff #270574)

Remove the comment for now. and also the clause stuff.

llvm/include/llvm/Transforms/IPO/Attributor.h

822

Why do we want a pointer here? We can derive from InformationCache elsewhere and still provide a reference, right?

877

I think everything else would just duplicate code. This can be public but the comment should mention that use *after* the Attributor was started to run is restricted to the Attributor itself. Initial seeding can be done via this function.

929

Just moving this to be public with a new comment LGTM. Feel free to do this right away. If you want input on the comment, let me know.

3004

I think the AA should live in OpenMP opt. TBH, similar to a comment below, I would like us to have a "simplification AA" here at some point. More specifically for this use case a "call simplifier". Those interfaces would live with the Attributor and be used by it, via some registration mechanism. It might even be useful to have the general logic in here eventually so that only a few functions need to be implemented by the user, e.g., bool interactsWithHiddenState(CallBase &CB), Optional<Value> getNewHiddenState(CallBase &CB), bool returnsHiddenState(CallBase &CB). The generic algorithm would track the "hidden state" (here ICV value) and provide all the logic. We can do that in a follow up or now. WDYT?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

66

As long as there is no other user we keep it here.

97

I think we need to do that eventually but for a first restricted deduction, see below, we might not need to. Just remove the comments for now then.

131

Why this change?

151

Documentation gone? We need some for both functions.

312

I think an RFC commit moving everything into a "cache" helper class should be provided seperatly. That seems trivial to accept and will significantly shrink this diff.

737

I feel the instance might be preferable so we can access AAs later, e.g., via A.lookupAAFor.

739

I think in the long run we want to register the ICV tracking attribute as another "simplification" attribute (for runtime calls). Then, the Attributor can ask all simplification attributes for a better value for a call, potentially resulting in better liveness and therefore better ICV tracking ;)

For now we could whitelist only ICV tracking (I guess).

859

Style: Add braces if you have a comment that makes the body multi-line.
You have to verify there is no cal in-between the two instructions too:

set_num_threads(2);
foo(); // resets num_threads
get_num_threads();

^ negative test we should keep.

We expand our tracking efforts in a follow up, for now we can be restrictive and correct with some positive and negative tests.

866

I don't think this works:

if (...)
  set_num_threads(2);
get_num_threads();

^ negative test we should keep.

For the first version go with the same basic block only

[sstefan1]

In D81788#2091453, @jdoerfert wrote:

I left a first set of comments. I think the overall direction is good, some interesting opportunities are opening up (see below). We need tests and also to split the patch (some suggestions below).

Maybe D81114 should be commited first, or should I just abandon that and put the tests here?

I'll split this with at least two more patches.

llvm/include/llvm/Transforms/IPO/Attributor.h

822

So that we can later cast it to OMPInformationCache and use openmp-specific information. This seemed like least amount of change.

877

I guess then it should also be allowed to be used after Attributor run is finished? Say we want to know ICV value after the Attributor run and all the deduplication, we still won't have querying attribute. Or should we also expose lookupAAFor?

929

You mean a NFC commit without a review?

3004

I think this should work, but I'll have to put more thought into (mainly the registration mechanism). this I'm in favor of getting this in first, since this is already proven to be rather large patch with lots of changes.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

131

I think this comes from the fact that the same OMPInfoCache is shared between the Attributor and OpenMPOPT. Maybe we'll be able to avoid this after splitting up the patch, but I'm not sure.

151

Forgot to copy. Will do.

866

I agree. Not to proud of this part tbh. I'll leave this out for now.

[jdoerfert]

I think this needs a rebase once some other landed.

llvm/include/llvm/Transforms/IPO/Attributor.h

822

You can cast regardless. If there is no other reason, please keep it a reference :)

877

Iff you are careful, you can use this *after* a run, e.g. if you only use information from attributes at a fixpoint. Exposing lookup makes this less error prone and would allow us to have an assert even. The other way is to create and remember all interesting locations. That is usually what we did so far but if we have non AA users lookup might be useful as well.

929

Yes. The code movement parts.

3004

Agreed. We can even split them further and submit the NFC parts so it becomes clearer to me what changes :)

[sstefan1]

• clean up the patch
• basic deduplication

[jdoerfert]

Nit: Add documentation to functions and classes please. Describe what happens and why if applicable.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

755

I'm really happy one can add a local AA like this :)

807

Style: & in both above.

812

Nit: Is this necessary?

831

Nit: GetterRFI

833

We might want to "push" the information from setters instead. That would allow us to remove setters that are not useful, e.g. it is reset before it could have been read. We also need to look at other/implicit reads and writes later: https://godbolt.org/z/fdBqiF

I'm fine with starting here but let me know what you think about pushing instead?

838

Can we replace the pair with a struct that has proper member names and documentation.

873

Somehow broken.

881

Style: & and GetterRFI.

885

Make both conditionals early exits

llvm/test/Transforms/OpenMP/icv_tracking.ll

56–57

I thought this min was necessary but I doubt it reading the standard again. Maybe just remove the fixme

95–96

I guess this should be use(10) too, right? And tmp4 should be removed.

[sstefan1]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

833

I guess we would first need to be aware of 'other' reads and writes, since without that information we could for example remove a setter that was supposed to be read.

I think we can do that with this approach as well. If the value wasn't read when we reach new setter, we mark the last one for removal? Not really sure if one is better than other.

873

The comment is broken?

885

I guess this makes sense now, but as we improve this, more conditionals will be added and it won't make sense to exit early.

llvm/test/Transforms/OpenMP/icv_tracking.ll

95–96

Yes, but only once we add callsite information. Right now we don't know if use might have set a new value.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

833

Both directions conceptually work. Let's wait and see. The correctness should never be compromised, I mean if you see something that you don't know you have to assume it reads and writes the ICV anyway.

873

Yes, a stray \p , right?

885

It depends. Having 3 early exists is, for me at least, nicer than one big condition.

llvm/test/Transforms/OpenMP/icv_tracking.ll

95–96

I see.

[sstefan1]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

885

Just to make sure we're on the same page here. I'm not saying there will be one big condition. Rather, right now there is the case when both instructions are in the same BB. In the future we also want to cover the case when they are not. Does that make sense?

[sstefan1]

addressing comments

[jdoerfert]

Some nits to be addressed, otherwise LGTM.

Address the clang-tidy warnings where possible please.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

799

& and below. I assume you don't want to copy it all the time.

830

Note: Once we do proper optimistic initialization we need to compute the replacement value in the update. For now this is fine.

864

Nit: If you declare this as template<> struct DenseMapInfo<ICVValue> you don't need to specify it with the map later.

911

Wouldn't a const block do just fine below?

[lebedev.ri]

I've reverted this in rG1d542f0ca83fa1411d6501a8d088450d83abd5b8.

There appears to be some kind of memory corruption/use-after-free/etc
going on here. In particular, in OpenMPOpt::deleteParallelRegions(),
in DeleteCallCB(), CI is garbage.

Reduced reproducer: ./bin/opt -O3 /tmp/test.ll:

; ModuleID = '/home/lebedevri/CREDUCE/input.c'
source_filename = "/home/lebedevri/CREDUCE/input.c"
target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
target triple = "x86_64-pc-linux-gnu"

%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

; Function Attrs: nounwind uwtable
define dso_local i32 @b() #0 {
  %1 = alloca i32, align 4
  %2 = call i32 @a()
  %3 = load i32, i32* %1, align 4
  ret i32 %3
}

; Function Attrs: nounwind uwtable
define internal i32 @a() #0 {
  %1 = alloca i32, align 4
  %2 = call i32 @b()
  call void (%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) @__kmpc_fork_call(%struct.ident_t* @0, i32 0, void (i32*, i32*, ...)* bitcast (void (i32*, i32*)* @.omp_outlined. to void (i32*, i32*, ...)*))
  %3 = load i32, i32* %1, align 4
  ret i32 %3
}

; Function Attrs: norecurse nounwind uwtable
define internal void @.omp_outlined.(i32* noalias %0, i32* noalias %1) #1 {
  %3 = alloca i32*, align 8
  %4 = alloca i32*, align 8
  store i32* %0, i32** %3, align 8, !tbaa !2
  store i32* %1, i32** %4, align 8, !tbaa !2
  ret void
}

; Function Attrs: nounwind
declare !callback !6 void @__kmpc_fork_call(%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...) #2

attributes #0 = { nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
attributes #1 = { norecurse nounwind uwtable "correctly-rounded-divide-sqrt-fp-math"="false" "disable-tail-calls"="false" "frame-pointer"="none" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "target-cpu"="x86-64" "target-features"="+cx8,+fxsr,+mmx,+sse,+sse2,+x87" "unsafe-fp-math"="false" "use-soft-float"="false" }
attributes #2 = { nounwind }

!llvm.module.flags = !{!0}
!llvm.ident = !{!1}

!0 = !{i32 1, !"wchar_size", i32 4}
!1 = !{!"Debian clang version 11.0.0-++20200709100646+c92a8c0a0f6-1~exp1~20200709201313.3348"}
!2 = !{!3, !3, i64 0}
!3 = !{!"any pointer", !4, i64 0}
!4 = !{!"omnipotent char", !5, i64 0}
!5 = !{!"Simple C/C++ TBAA"}
!6 = !{!7}
!7 = !{i64 2, i64 -1, i64 -1, i1 true}

[sstefan1]

fixing the issue with dead uses.

[jdoerfert]

Add the reproducer as test please. With that one and the change below, LGTM.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

306

-Size
+RFIs.size()

just add the function to the EnumeratedArray please.

[jdoerfert]

llvm/include/llvm/ADT/EnumeratedArray.h

41 ↗

(On Diff #277182)

Nit: const

llvm/include/llvm/Transforms/IPO/Attributor.h

struct Attributor {
   ///                  the abstract attributes.
   /// \param CGUpdater Helper to update an underlying call graph.
   /// \param Allowed If not null, a set limiting the attribute opportunities.
   Attributor(SetVector<Function *> &Functions, InformationCache &InfoCache,
              CallGraphUpdater &CGUpdater,
              DenseSet<const char *> *Allowed = nullptr)
       : Allocator(InfoCache.Allocator), Functions(Functions),
         InfoCache(InfoCache), CGUpdater(CGUpdater), Allowed(Allowed) {}
 
   ~Attributor();
 
   /// Run the analyses until a fixpoint is reached or enforced (timeout).
   ///
   /// The attributes registered with this Attributor can be used after as long
   /// as the Attributor is not destroyed (it owns the attributes now).
   ///
   /// \Returns CHANGED if the IR was changed, otherwise UNCHANGED.
const AAType &getAndUpdateAAFor(const AbstractAttribute &QueryingAA,
                                   bool TrackDependence = true,
                                   DepClassTy DepClass = DepClassTy::REQUIRED) {
     return getOrCreateAAFor<AAType>(IRP, &QueryingAA, TrackDependence, DepClass,
                                     /* ForceUpdate */ true);
   }
 
   /// The version of getAAFor that allows to omit a querying abstract
   /// attribute. Using this after Attributor started running is restricted to
   /// only the Attributor itself. Initial seeding of AAs can be done via this
   /// function.
   template <typename AAType>
   const AAType &getOrCreateAAFor(const IRPosition &IRP,
                                  const AbstractAttribute *QueryingAA = nullptr,
                                  bool TrackDependence = false,
                                  DepClassTy DepClass = DepClassTy::OPTIONAL,
                                  bool ForceUpdate = false) {
     if (AAType *AAPtr = lookupAAFor<AAType>(IRP, QueryingAA, TrackDependence)) {
const AAType &getOrCreateAAFor(const IRPosition &IRP,
     updateAA(AA);
 
     if (TrackDependence && AA.getState().isValidState())
       recordDependence(AA, const_cast<AbstractAttribute &>(*QueryingAA),
                        DepClass);
     return AA;
   }
 
   /// Return the attribute of \p AAType for \p IRP if existing. This also allows
   /// non-AA users lookup.
   template <typename AAType>
   AAType *lookupAAFor(const IRPosition &IRP,
                       const AbstractAttribute *QueryingAA = nullptr,
                       bool TrackDependence = false,
                       DepClassTy DepClass = DepClassTy::OPTIONAL) {
     static_assert(std::is_base_of<AbstractAttribute, AAType>::value,
                   "Cannot query an attribute with a type not derived from "
struct Attributor {
   /// \p F.
   void markLiveInternalFunction(const Function &F) {
     assert(F.hasLocalLinkage() &&
            "Only local linkage is assumed dead initially.");
 
     identifyDefaultAbstractAttributes(const_cast<Function &>(F));
   }
 
+  /// Helper function to remove callsite.
+  void removeCallSite(CallInst *CI) {
+    if (!CI)
+      return;
+
+    CGUpdater.removeCallSite(*CI);
+  }
+
   /// Record that \p U is to be replaces with \p NV after information was
   /// manifested. This also triggers deletion of trivially dead istructions.
   bool changeUseAfterManifest(Use &U, Value &NV) {
     Value *&V = ToBeChangedUses[&U];
     if (V && (V->stripPointerCasts() == NV.stripPointerCasts() ||
               isa_and_nonnull<UndefValue>(V)))
       return false;
     assert((!V || V == &NV || isa<UndefValue>(NV)) &&
enum AttributorRunOption {
   NONE = 0,
   MODULE = 1 << 0,
   CGSCC = 1 << 1,
   ALL = MODULE | CGSCC
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_IPO_FUNCTIONATTRS_H

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

STATISTIC(NumOpenMPRuntimeFunctionUsesIdentified,
           "Number of OpenMP runtime function uses identified");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
 #endif
 
 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,
                       BumpPtrAllocator &Allocator, SetVector<Function *> *CGSCC,
                       SmallPtrSetImpl<Function *> &ModuleSlice)
       : InformationCache(M, AG, Allocator, CGSCC), ModuleSlice(ModuleSlice),
         OMPBuilder(M) {
     initializeTypes(M);
     initializeRuntimeFunctions();
     initializeInternalControlVars();
 
     OMPBuilder.initialize();
   }
struct InternalControlVarInfo {
 
     /// Initial value.
     ConstantInt *InitValue;
 
     /// Setter RTL function associated with this ICV.
     RuntimeFunction Setter;
 
     /// Getter RTL function associated with this ICV.
     RuntimeFunction Getter;
 
     /// RTL Function corresponding to the override clause of this ICV
     RuntimeFunction Clause;
   };
 
   /// Generic information that describes a runtime function
   struct RuntimeFunctionInfo {
 
struct RuntimeFunctionInfo {
     /// The declaration if available.
     Function *Declaration = nullptr;
 
     /// Uses of this runtime function per function containing the use.
     using UseVector = SmallVector<Use *, 16>;
 
     /// Return the vector of uses in function \p F.
     UseVector &getOrCreateUseVector(Function *F) {
-      std::unique_ptr<UseVector> &UV = UsesMap[F];
+      std::shared_ptr<UseVector> &UV = UsesMap[F];
       if (!UV)
-        UV = std::make_unique<UseVector>();
+        UV = std::make_shared<UseVector>();
       return *UV;
     }
 
     /// Return the vector of uses in function \p F or `nullptr` if there are
     /// none.
     const UseVector *getUseVector(Function &F) const {
       auto I = UsesMap.find(&F);
       if (I != UsesMap.end())
         return I->second.get();
       return nullptr;
     }
 
     /// Return how many functions contain uses of this runtime function.
     size_t getNumFunctionsWithUses() const { return UsesMap.size(); }
 
     /// Return the number of arguments (or the minimal number for variadic
     /// functions).
     size_t getNumArgs() const { return ArgumentTypes.size(); }
 
     /// Run the callback \p CB on each use and forget the use if the result is
     /// true. The callback will be fed the function in which the use was
     /// encountered as second argument.
     void foreachUse(function_ref<bool(Use &, Function &)> CB) {
       for (auto &It : UsesMap)
         foreachUse(CB, It.first, It.second.get());
     }
 
     /// Run the callback \p CB on each use within the function \p F and forget
     /// the use if the result is true.
     void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F,
void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F,
         UV[Idx] = UV.back();
         UV.pop_back();
       }
     }
 
   private:
     /// Map from functions to all uses of this runtime function contained in
     /// them.
-    DenseMap<Function *, std::unique_ptr<UseVector>> UsesMap;
+    DenseMap<Function *, std::shared_ptr<UseVector>> UsesMap;
   };
 
   /// The slice of the module we are allowed to look at.
   SmallPtrSetImpl<Function *> &ModuleSlice;
 
   /// An OpenMP-IR-Builder instance
   OpenMPIRBuilder OMPBuilder;
 
void initializeRuntimeFunctions() {
     Module &M = *((*ModuleSlice.begin())->getParent());
 
 #define OMP_RTL(_Enum, _Name, _IsVarArg, _ReturnType, ...)                     \
   {                                                                            \
     SmallVector<Type *, 8> ArgsTypes({__VA_ARGS__});                           \
     Function *F = M.getFunction(_Name);                                        \
     if (declMatchesRTFTypes(F, _ReturnType, ArgsTypes)) {                      \
       auto &RFI = RFIs[_Enum];                                                 \
       RFI.Kind = _Enum;                                                        \
       RFI.Name = _Name;                                                        \
       RFI.IsVarArg = _IsVarArg;                                                \
       RFI.ReturnType = _ReturnType;                                            \
       RFI.ArgumentTypes = std::move(ArgsTypes);                                \
       RFI.Declaration = F;                                                     \
       unsigned NumUses = CollectUses(RFI);                                     \
       (void)NumUses;                                                           \
       LLVM_DEBUG({                                                             \
         dbgs() << TAG << RFI.Name << (RFI.Declaration ? "" : " not")           \
                << " found\n";                                                  \
         if (RFI.Declaration)                                                   \
           dbgs() << TAG << "-> got " << NumUses << " uses in "                 \
                  << RFI.getNumFunctionsWithUses()                              \
                  << " different functions.\n";                                 \
       });                                                                      \
     }                                                                          \
   }
 #include "llvm/Frontend/OpenMP/OMPKinds.def"
 
     // TODO: We should attach the attributes defined in OMPKinds.def.
   }
 };
 
 struct OpenMPOpt {
 
   using OptimizationRemarkGetter =
       function_ref<OptimizationRemarkEmitter &(Function *)>;
 
   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
             OptimizationRemarkGetter OREGetter,
-            OMPInformationCache &OMPInfoCache)
+            OMPInformationCache &OMPInfoCache, Attributor &A)
       : M(*(*SCC.begin())->getParent()), SCC(SCC), CGUpdater(CGUpdater),
-        OREGetter(OREGetter), OMPInfoCache(OMPInfoCache) {}
+        OREGetter(OREGetter), OMPInfoCache(OMPInfoCache), A(A) {}
 
   /// Run all OpenMP optimizations on the underlying SCC/ModuleSlice.
   bool run() {
     bool Changed = false;
 
     LLVM_DEBUG(dbgs() << TAG << "Run on SCC with " << SCC.size()
                       << " functions in a slice with "
                       << OMPInfoCache.ModuleSlice.size() << " functions\n");
if (PrintICVValues) {
                               : "IMPLEMENTATION_DEFINED");
           };
 
           emitRemarkOnFunction(F, "OpenMPICVTracker", Remark);
         }
       }
     }
 
+    Changed |= runAttributor();
     Changed |= deduplicateRuntimeCalls();
     Changed |= deleteParallelRegions();
 
     return Changed;
   }
 
   /// Return the call if \p U is a callee use in a regular call. If \p RFI is
   /// given it has to be the callee or a nullptr is returned.
private:
   /// the second an optional replacement call.
   CallGraphUpdater &CGUpdater;
 
   /// Callback to get an OptimizationRemarkEmitter from a Function *
   OptimizationRemarkGetter OREGetter;
 
   /// OpenMP-specific information cache. Also Used for Attributor runs.
   OMPInformationCache &OMPInfoCache;
+
+  /// Attributor instance.
+  Attributor &A;
+
+  /// Helper function to run Attributor on SCC.
+  bool runAttributor() {
+    if (SCC.empty())
+      return false;
+
+    registerAAs();
+
+    ChangeStatus Changed = A.run();
+
+    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() {
+    for (Function *F : SCC) {
+      if (F->isDeclaration())
+        continue;
+
+      A.getOrCreateAAFor<AAICVTracker>(IRPosition::function(*F));
+    }
+  }
+};
+
+/// 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) {}
+
+  /// 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;
+
+  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"; }
+
+  // FIXME: come up with some stats.
+  void trackStatistics() const override {}
+
+  /// TODO: decide whether to deduplicate here, or use current
+  /// deduplicateRuntimeCalls function.
+  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);
+    return Changed;
+  }
+
+  /// 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) {}
+  };
+
+  // Provide DenseMapInfo for ICVValue
+  struct DenseMapICVValueKeyInfo {
+    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);
+    }
+  };
+
+  using SetVectorTy =
+      SetVector<ICVValue, SmallVector<ICVValue, 4>,
+                SmallDenseSet<ICVValue, 4, DenseMapICVValueKeyInfo>>;
+
+  // Map of ICV to their values at specific program point.
+  EnumeratedArray<SetVectorTy, InternalControlVar,
+                  InternalControlVar::ICV___last>
+      ICVValuesMap;
+
+  // 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 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))))
+          HasChanged = ChangeStatus::CHANGED;
+
+        return false;
+      };
+
+      SetterRFI.foreachUse(TrackValues, F);
+    }
+
+    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 {
+    BasicBlock *CurrBB = const_cast<BasicBlock *>(I->getParent());
+
+    auto &ValuesSet = ICVValuesMap[ICV];
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
+
+    for (auto &ICVVal : ValuesSet) {
+      if (CurrBB == ICVVal.Inst->getParent()) {
+        if (!ICVVal.Inst->comesBefore(I))
+          continue;
+
+        // 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 (auto *CB = dyn_cast<CallBase>(I))
+            if (CB->getCalledFunction() != GetterRFI.Declaration)
+              return nullptr;
+
+          I = I->getPrevNode();
+        }
+
+        // No call in between, return the value.
+        return ICVVal.TrackedValue;
+      }
+    }
+
+    // No value was tracked.
+    return nullptr;
+  }
 };
 } // 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!");
+  case IRPosition::IRP_FUNCTION:
+    AA = new (A.Allocator) AAICVTrackerFunction(IRP, A);
+    break;
+  }
+
+  return *AA;
+}
+
 PreservedAnalyses OpenMPOptPass::run(LazyCallGraph::SCC &C,
                                      CGSCCAnalysisManager &AM,
                                      LazyCallGraph &CG, CGSCCUpdateResult &UR) {
   if (!containsOpenMP(*C.begin()->getFunction().getParent(), OMPInModule))
     return PreservedAnalyses::all();
 
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
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, ModuleSlice);
 
+  Attributor A(Functions, InfoCache, CGUpdater);
+
   // TODO: Compute the module slice we are allowed to look at.
-  OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache);
+  OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
   bool Changed = OMPOpt.run();
   (void)Changed;
   return PreservedAnalyses::all();
 }
 
 namespace {
 
 struct OpenMPOptLegacyPass : public CallGraphSCCPass {
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, ModuleSlice);
 
+    Attributor A(Functions, InfoCache, CGUpdater);
+
     // TODO: Compute the module slice we are allowed to look at.
-    OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache);
+    OpenMPOpt OMPOpt(SCC, CGUpdater, OREGetter, InfoCache, A);
     return OMPOpt.run();
   }
 
   bool doFinalization(CallGraph &CG) override { return CGUpdater.finalize(); }
 };
 
 } // end anonymous namespace
 

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
 
 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:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 [[TMP1]])
-; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    [[TMP6:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    tail call void @use(i32 [[TMP4]])
-; CHECK-NEXT:    tail call void @use(i32 [[TMP5]])
+; 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*, ...)*))
-; CHECK-NEXT:    [[TMP7:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    tail call void @use(i32 [[TMP7]])
+; CHECK-NEXT:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
+; CHECK-NEXT:    tail call void @use(i32 [[TMP3]])
 ; CHECK-NEXT:    ret i32 0
 ;
   tail call void @omp_set_num_threads(i32 %0)
   %3 = tail call i32 @omp_get_max_threads()
 ; FIXME: this value should be tracked and the rest of the getters deduplicated and replaced with it.
   tail call void @omp_set_num_threads(i32 %1)
   %4 = tail call i32 @omp_get_max_threads()
   %5 = tail call i32 @omp_get_max_threads()

 
 define internal void @.omp_outlined.(i32* %0, i32* %1) {
 ; CHECK-LABEL: define {{[^@]+}}@.omp_outlined.
 ; CHECK-SAME: (i32* [[TMP0:%.*]], i32* [[TMP1:%.*]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP4]])
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 10)
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    tail call void @use(i32 [[TMP5]])
+; CHECK-NEXT:    tail call void @use(i32 10)
 ; CHECK-NEXT:    ret void
 ;
 ; FIXME: this value should be tracked and the rest of the getters deduplicated and replaced with it.
   %3 = tail call i32 @omp_get_max_threads()
   %4 = tail call i32 @omp_get_max_threads()
   tail call void @use(i32 %4)
-; FIXME: this value ( min(%3, 10) ) should be tracked and the rest of the getters deduplicated and replaced with it.
   tail call void @omp_set_num_threads(i32 10)
   %5 = tail call i32 @omp_get_max_threads()
   tail call void @use(i32 %5)
   ret void
 }
 
 declare !callback !0 void @__kmpc_fork_call(%struct.ident_t*, i32, void (i32*, i32*, ...)*, ...)
 
 define dso_local i32 @bar(i32 %0, i32 %1) {
 ; CHECK-LABEL: define {{[^@]+}}@bar
 ; 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:    [[TMP5:%.*]] = tail call i32 @omp_get_max_threads()
 ; 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..1 to void (i32*, i32*, ...)*))
-; CHECK-NEXT:    [[TMP6:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    tail call void @use(i32 [[TMP6]])
+; 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
 ; FIXME: getters can be replaced with %4
   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..1 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..1(i32* %0, i32*  %1) {
 ; CHECK-LABEL: define {{[^@]+}}@.omp_outlined..1
 ; CHECK-SAME: (i32* [[TMP0:%.*]], i32* [[TMP1:%.*]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP3]])
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 10)
+; CHECK-NEXT:    tail call void @use(i32 10)
 ; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP4]])
-; CHECK-NEXT:    [[TMP5:%.*]] = tail call i32 @omp_get_max_threads()
-; CHECK-NEXT:    tail call void @use(i32 [[TMP5]])
 ; CHECK-NEXT:    ret void
 ;
   %3 = tail call i32 @omp_get_max_threads()
   tail call void @use(i32 %3)
   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
 }
 
 !0 = !{!1}
 !1 = !{i64 2, i64 -1, i64 -1, i1 true}