[AbstractAttributor] Fold __kmpc_parallel_level if possible

[tianshilei1992]

Summary

Similar to D105787, this patch tries to fold __kmpc_parallel_level if possible.
Note that __kmpc_parallel_level doesn't take activeness into consideration,
based on current deviceRTLs, its return value can be such as 0, 1, 2, instead
of 0, 129, 130, etc. that also indicate activeness.

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

444–446

These changes are covered by D106149, which will not be part of this patch eventually.

[tianshilei1992]

rebase

[tianshilei1992]

add missing feature and a trivial test case

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3296

The logic seems problematic here. Need to revise later.

llvm/test/Transforms/OpenMP/parallel_level_fold.ll

1

The enhancement of the test case is still WIP.

[tianshilei1992]

add assertion

[jdoerfert]

add noinline to the microtrask function to avoid us accidentally doing the wrong thing. I'll put a patch up that removes those noinline in the module pass.

[tianshilei1992]

more conservative

[tianshilei1992]

rebase

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3314

3659

3678

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3314

Seems unnecessary because this variable is only written in the function call but I can do it.

[tianshilei1992]

rebase and fixed comments

[jdoerfert]

LG, some minor comments below.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3296

mark __kmpc_parallel_51 as noinline to make this sound.

3671

3671–3672

[tianshilei1992]

fix comments

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3671

I changed this a little bit. If CallerKernelInfoAA.ParallelLevels is empty, we can only tell we currently cannot fold it, but it doesn't mean it cannot be folded in the future. If the size is more than 1, yeah, definitely it is pessimistic.

[tianshilei1992]

rebase

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3675

not nullptr but it shuld continue to be none

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3675

Aha, I recall it. Yeah, will do it.

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

update

[tianshilei1992]

add __attribute__((noinline)) to __kmpc_parallel_level

[jhuber6]

openmp/libomptarget/deviceRTLs/interface.h

444 ↗

(On Diff #361315)

Isn't there a NOINLINE definition in target_impl.h?

[tianshilei1992]

rebase

[tianshilei1992]

rebase and update

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

#include "llvm/Frontend/OpenMP/OMPKinds.def"
   /// Collection of known kernels (\see Kernel) in the module.
   SmallPtrSetImpl<Kernel> &Kernels;
 
   /// Collection of known OpenMP runtime functions..
   DenseSet<const Function *> RTLFunctions;
 };
 
 template <typename Ty, bool InsertInvalidates = true>
 struct BooleanStateWithSetVector : public BooleanState {
   bool contains(const Ty &Elem) const { return Set.contains(Elem); }
   bool insert(const Ty &Elem) {
     if (InsertInvalidates)
       BooleanState::indicatePessimisticFixpoint();
     return Set.insert(Elem);
   }
 
   const Ty &operator[](int Idx) const { return Set[Idx]; }
   bool operator==(const BooleanStateWithSetVector &RHS) const {
     return BooleanState::operator==(RHS) && Set == RHS.Set;
struct KernelInfoState : AbstractState {
   CallBase *KernelDeinitCB = nullptr;
 
   /// Flag to indicate if the associated function is a kernel entry.
   bool IsKernelEntry = false;
 
   /// State to track what kernel entries can reach the associated function.
   BooleanStateWithPtrSetVector<Function, false> ReachingKernelEntries;
 
+  /// State to track what parallel levels the associated function can be at.
+  BooleanStateWithSetVector<uint8_t, false> ParallelLevels;
+
   /// Abstract State interface
   ///{
 
   KernelInfoState() {}
   KernelInfoState(bool BestState) {
     if (!BestState)
       indicatePessimisticFixpoint();
   }
void initialize(Attributor &A) override {
 
     Function *Fn = getAnchorScope();
     if (!OMPInfoCache.Kernels.count(Fn))
       return;
 
     // Add itself to the reaching kernel and set IsKernelEntry.
     ReachingKernelEntries.insert(Fn);
     IsKernelEntry = true;
+    // Kernel entry is at level 0, which means not in parallel region.
+    ParallelLevels.insert(0);
 
     OMPInformationCache::RuntimeFunctionInfo &InitRFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_target_init];
     OMPInformationCache::RuntimeFunctionInfo &DeinitRFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_target_deinit];
 
     // For kernels we perform more initialization work, first we find the init
     // and deinit calls.
ChangeStatus updateImpl(Attributor &A) override {
     };
 
     bool UsedAssumedInformationInCheckRWInst = false;
     if (!SPMDCompatibilityTracker.isAtFixpoint())
       if (!A.checkForAllReadWriteInstructions(
               CheckRWInst, *this, UsedAssumedInformationInCheckRWInst))
         SPMDCompatibilityTracker.indicatePessimisticFixpoint();
 
-    if (!IsKernelEntry)
+    if (!IsKernelEntry) {
       updateReachingKernelEntries(A);
+      updateParallelLevels(A);
+    }
 
     // Callback to check a call instruction.
     bool AllSPMDStatesWereFixed = true;
     auto CheckCallInst = [&](Instruction &I) {
       auto &CB = cast<CallBase>(I);
       auto &CBAA = A.getAAFor<AAKernelInfo>(
           *this, IRPosition::callsite_function(CB), DepClassTy::OPTIONAL);
       getState() ^= CBAA.getState();
void updateReachingKernelEntries(Attributor &A) {
     };
 
     bool AllCallSitesKnown;
     if (!A.checkForAllCallSites(PredCallSite, *this,
                                 true /* RequireAllCallSites */,
                                 AllCallSitesKnown))
       ReachingKernelEntries.indicatePessimisticFixpoint();
   }
+
+  /// Update info regarding parallel levels.
+  void updateParallelLevels(Attributor &A) {
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    OMPInformationCache::RuntimeFunctionInfo &Parallel51RFI =
+        OMPInfoCache.RFIs[OMPRTL___kmpc_parallel_51];
+
+    auto PredCallSite = [&](AbstractCallSite ACS) {
+      Function *Caller = ACS.getInstruction()->getFunction();
+
+      assert(Caller && "Caller is nullptr");
+
+      auto &CAA =
+          A.getOrCreateAAFor<AAKernelInfo>(IRPosition::function(*Caller));
+      if (CAA.ParallelLevels.isValidState()) {
+        if (Caller != Parallel51RFI.Declaration)
+          ParallelLevels ^= CAA.ParallelLevels;
+        else
+          for (uint8_t L : CAA.ParallelLevels)
+            ParallelLevels.insert(L + 1);
+
+        if (ParallelLevels.size() > 1)
+          ParallelLevels.indicatePessimisticFixpoint();
+
+        return true;
+      }
+
+      // We lost track of the caller of the associated function, any kernel
+      // could reach now.
+      ParallelLevels.indicatePessimisticFixpoint();
+
+      return true;
+    };
+
+    bool AllCallSitesKnown;
+    if (!A.checkForAllCallSites(PredCallSite, *this,
+                                true /* RequireAllCallSites */,
+                                AllCallSitesKnown))
+      ParallelLevels.indicatePessimisticFixpoint();
+  }
 };
 
 /// The call site kernel info abstract attribute, basically, what can we say
 /// about a call site with regards to the KernelInfoState. For now this simply
 /// forwards the information from the callee.
 struct AAKernelInfoCallSite : AAKernelInfo {
   AAKernelInfoCallSite(const IRPosition &IRP, Attributor &A)
       : AAKernelInfo(IRP, A) {}
struct AAFoldRuntimeCallCallSiteReturned : AAFoldRuntimeCall {
 
   ChangeStatus updateImpl(Attributor &A) override {
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
 
     switch (RFKind) {
     case OMPRTL___kmpc_is_spmd_exec_mode:
       Changed |= foldIsSPMDExecMode(A);
       break;
+    case OMPRTL___kmpc_parallel_level:
+      Changed |= foldParallelLevel(A);
+      break;
     default:
       llvm_unreachable("Unhandled OpenMP runtime function!");
     }
 
     return Changed;
   }
 
   ChangeStatus manifest(Attributor &A) override {
if (KnownSPMDCount || AssumedSPMDCount) {
       // must be none.
       assert(!SimplifiedValue.hasValue() && "SimplifiedValue should be none");
     }
 
     return SimplifiedValue == SimplifiedValueBefore ? ChangeStatus::UNCHANGED
                                                     : ChangeStatus::CHANGED;
   }
 
+  /// Fold __kmpc_parallel_level into a constant if possible.
+  ChangeStatus foldParallelLevel(Attributor &A) {
+    Optional<Value *> SimplifiedValueBefore = SimplifiedValue;
+
+    auto &CallerKernelInfoAA = A.getAAFor<AAKernelInfo>(
+        *this, IRPosition::function(*getAnchorScope()), DepClassTy::REQUIRED);
+
+    if (!CallerKernelInfoAA.ParallelLevels.isValidState())
+      return indicatePessimisticFixpoint();
+
+    // If the caller is at multiple parallel levels, the associated function
+    // call cannot be folded anymore.
+    // TODO: I feel this will not happen as we will invalid
+    // CallerKernelInfoAA.ParallelLevels when updating AAKernelInfoFunction,
+    // therefore this case will be handled by checking valid state above
+    // directly.
+    if (CallerKernelInfoAA.ParallelLevels.size() > 1)
+      return indicatePessimisticFixpoint();
+
+    // If we cannot deduct parallel region of the caller, we set SimplifiedValue
+    // to nullptr to indicate the associated function call cannot be folded for
+    // now.
+    if (CallerKernelInfoAA.ParallelLevels.empty())
+      SimplifiedValue = nullptr;
+    else {
+      const uint8_t Level = CallerKernelInfoAA.ParallelLevels[0];
+      auto &Ctx = getAnchorValue().getContext();
+      SimplifiedValue = ConstantInt::get(Type::getInt8Ty(Ctx), Level);
+    }
+
+    return SimplifiedValue == SimplifiedValueBefore ? ChangeStatus::UNCHANGED
+                                                    : ChangeStatus::CHANGED;
+  }
+
   /// An optional value the associated value is assumed to fold to. That is, we
   /// assume the associated value (which is a call) can be replaced by this
   /// simplified value.
   Optional<Value *> SimplifiedValue;
 
   /// The runtime function kind of the callee of the associated call site.
   RuntimeFunction RFKind;
 };
 
 } // namespace
 
IsSPMDRFI.foreachUse(SCC, [&](Use &U, Function &) {
       if (!CI)
         return false;
       A.getOrCreateAAFor<AAFoldRuntimeCall>(
           IRPosition::callsite_returned(*CI), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,
           /* UpdateAfterInit */ false);
       return false;
     });
+
+    auto &ParallelLevelRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_parallel_level];
+    ParallelLevelRFI.foreachUse(SCC, [&](Use &U, Function &) {
+      CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &ParallelLevelRFI);
+      if (!CI)
+        return false;
+      A.getOrCreateAAFor<AAFoldRuntimeCall>(
+          IRPosition::callsite_returned(*CI), /* QueryingAA */ nullptr,
+          DepClassTy::NONE, /* ForceUpdate */ false,
+          /* UpdateAfterInit */ false);
+
+      return false;
+    });
   }
 
   // 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];
 

llvm/test/Transforms/OpenMP/parallel_level_fold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s
+target triple = "nvptx64"
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@no_parallel_exec_mode = weak constant i8 0
+@G = external global i8
[e-mail] = appending global [1 x i8*] [i8* @no_parallel_exec_mode], section "llvm.metadata"
+
+;.
+; CHECK: @[[NO_PARALLEL_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 0
+; CHECK: @[[G:[a-zA-Z0-9_$"\\.-]+]] = external global i8
+; CHECK: @[[LLVM_COMPILER_USED:[a-zA-Z0-9_$"\\.-]+]] = appending global [1 x i8*] [i8* @no_parallel_exec_mode], section "llvm.metadata"
+;.
+define weak void @no_parallel() {
+; CHECK-LABEL: define {{[^@]+}}@no_parallel() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @foo()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define internal void @foo() {
+; CHECK-LABEL: define {{[^@]+}}@foo() {
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    ret void
+;
+  call void @bar()
+  ret void
+}
+
+define internal void @bar() {
+; CHECK-LABEL: define {{[^@]+}}@bar() {
+; CHECK-NEXT:    store i8 0, i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %level = call i8 @__kmpc_parallel_level()
+  store i8 %level, i8* @G
+  ret void
+}
+
+declare i8 @__kmpc_parallel_level()
+declare i32 @__kmpc_target_init(%struct.ident_t*, i1 zeroext, i1 zeroext, i1 zeroext) #1
+declare void @__kmpc_target_deinit(%struct.ident_t* nocapture readnone, i1 zeroext, i1 zeroext) #1
+
+!llvm.module.flags = !{!0, !1}
+!nvvm.annotations = !{!2}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+!2 = !{void ()* @no_parallel, !"kernel", i32 1}
+;.
+; CHECK: [[META0:![0-9]+]] = !{i32 7, !"openmp", i32 50}
+; CHECK: [[META1:![0-9]+]] = !{i32 7, !"openmp-device", i32 50}
+; CHECK: [[META2:![0-9]+]] = !{void ()* @no_parallel, !"kernel", i32 1}
+;.