[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

450–452

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

3326

The logic seems problematic here. Need to revise later.

llvm/test/Transforms/OpenMP/parallel_level_fold.ll

2

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

3344

3699

3718

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3344

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

3326

mark __kmpc_parallel_51 as noinline to make this sound.

3711

3733–3734

[tianshilei1992]

fix comments

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3711

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

3737

not nullptr but it shuld continue to be none

[tianshilei1992]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3737

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 {
     A.registerSimplificationCallback(
         IRPosition::callsite_argument(*KernelDeinitCB,
                                       DeinitRequiresFullRuntimeArgNo),
         IsGenericModeSimplifyCB);
 
     // Check if we know we are in SPMD-mode already.
     ConstantInt *IsSPMDArg =
         dyn_cast<ConstantInt>(KernelInitCB->getArgOperand(InitIsSPMDArgNo));
-    if (IsSPMDArg && !IsSPMDArg->isZero())
+    if (IsSPMDArg && !IsSPMDArg->isZero()) {
       SPMDCompatibilityTracker.indicateOptimisticFixpoint();
+      // Kernel entry is at level 1 if in SPMD mode.
+      // NOTE: This is quite dangerous because we assume there is no direct or
+      // indirect function call to `__kmpc_parallel_level` before we update the
+      // parallel level in `__kmpc_spmd_kernel_init`. However, we have to do it
+      // in this way to make the fold right. Alternatively, we could only fold
+      // to 0 if we don't use any assumed information.
+      ParallelLevels.insert(1);
+    } else {
+      // Kernel entry is at level 0 if in none-SPMD mode.
+      ParallelLevels.insert(0);
+    }
   }
 
   /// Modify the IR based on the KernelInfoState as the fixpoint iteration is
   /// finished now.
   ChangeStatus manifest(Attributor &A) override {
     // If we are not looking at a kernel with __kmpc_target_init and
     // __kmpc_target_deinit call we cannot actually manifest the information.
     if (!KernelInitCB || !KernelDeinitCB)
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()) {
+        // Any function that is called by `__kmpc_parallel_51` will not be
+        // folded as the parallel level in the function is updated. In order to
+        // get it right, all the analysis would depend on the implentation. That
+        // said, if in the future any change to the implementation, the analysis
+        // could be wrong. As a consequence, we are just conservative here.
+        if (Caller == Parallel51RFI.Declaration) {
+          ParallelLevels.indicatePessimisticFixpoint();
+          return true;
+        }
+
+        ParallelLevels ^= CAA.ParallelLevels;
+
+        return true;
+      }
+
+      // We lost track of the caller of the associated function, any kernel
+      // could reach now.
+      ParallelLevels.indicatePessimisticFixpoint();
+
+      return true;
+    };
+
+    bool AllCallSitesKnown = true;
+    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) {}
ChangeStatus updateImpl(Attributor &A) override {
 
     switch (RFKind) {
     case OMPRTL___kmpc_is_spmd_exec_mode:
       Changed |= foldIsSPMDExecMode(A);
       break;
     case OMPRTL___kmpc_is_generic_main_thread_id:
       Changed |= foldIsGenericMainThread(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 {
ChangeStatus foldIsGenericMainThread(Attributor &A) {
     if (!ExecutionDomainAA.isValidState())
       return indicatePessimisticFixpoint();
 
     auto &Ctx = getAnchorValue().getContext();
     if (ExecutionDomainAA.isExecutedByInitialThreadOnly(CB))
       SimplifiedValue = ConstantInt::get(Type::getInt8Ty(Ctx), true);
     else
       return indicatePessimisticFixpoint();
 
     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 (!CallerKernelInfoAA.ReachingKernelEntries.isValidState())
+      return indicatePessimisticFixpoint();
+
+    unsigned SPMDCount = 0;
+    for (Kernel K : CallerKernelInfoAA.ReachingKernelEntries) {
+      auto &AA = A.getAAFor<AAKernelInfo>(*this, IRPosition::function(*K),
+                                          DepClassTy::REQUIRED);
+      if (!AA.SPMDCompatibilityTracker.isValidState())
+        return indicatePessimisticFixpoint();
+
+      if (AA.SPMDCompatibilityTracker.isAssumed())
+        ++SPMDCount;
+    }
+
+    // If we cannot deduct parallel region of the caller, or it can be at
+    // multiple levels, we set SimplifiedValue to nullptr to indicate the
+    // associated function call cannot be folded.
+    if (CallerKernelInfoAA.ParallelLevels.empty() ||
+        CallerKernelInfoAA.ParallelLevels.size() > 1)
+      SimplifiedValue = nullptr;
+    else {
+      const uint8_t Level = CallerKernelInfoAA.ParallelLevels[0];
+      // If the caller can be reached by a SPMD kernel entry, the parallel level
+      // is 1. As a result, if the detected level is not 1, the caller can be at
+      // multiple levels, and we cannot fold it.
+      if (Level != 1 && SPMDCount)
+        return indicatePessimisticFixpoint();
+
+      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;
 };
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_spmd_exec_mode = weak constant i8 1
+@spmd_exec_mode = weak constant i8 0
+@parallel_exec_mode = weak constant i8 0
+@G = external global i8
[e-mail] = appending global [3 x i8*] [i8* @no_spmd_exec_mode, i8* @spmd_exec_mode, i8* @parallel_exec_mode], section "llvm.metadata"
+
+;.
+; CHECK: @[[NO_SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 1
+; CHECK: @[[SPMD_EXEC_MODE:[a-zA-Z0-9_$"\\.-]+]] = weak constant i8 0
+; CHECK: @[[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 [3 x i8*] [i8* @no_spmd_exec_mode, i8* @spmd_exec_mode, i8* @parallel_exec_mode], section "llvm.metadata"
+;.
+define weak void @none_spmd() {
+; CHECK-LABEL: define {{[^@]+}}@none_spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+; CHECK-NEXT:    call void @none_spmd_helper()
+; CHECK-NEXT:    call void @mixed_helper()
+; 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 @none_spmd_helper()
+  call void @mixed_helper()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define weak void @spmd() {
+; CHECK-LABEL: define {{[^@]+}}@spmd() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @spmd_helper()
+; CHECK-NEXT:    call void @mixed_helper()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+  call void @spmd_helper()
+  call void @mixed_helper()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+  ret void
+}
+
+define weak void @parallel() {
+; CHECK-LABEL: define {{[^@]+}}@parallel() {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* align 536870912 null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @spmd_helper()
+; CHECK-NEXT:    call void @__kmpc_parallel_51(%struct.ident_t* noalias noundef align 536870912 null, i32 noundef 0, i32 noundef 0, i32 noundef 0, i32 noundef 0, i8* noalias noundef align 536870912 null, i8* noalias noundef align 536870912 null, i8** noalias noundef align 536870912 null, i64 noundef 0)
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+  call void @spmd_helper()
+  call void @__kmpc_parallel_51(%struct.ident_t* null, i32 0, i32 0, i32 0, i32 0, i8* null, i8* null, i8** null, i64 0)
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+  ret void
+}
+
+define internal void @mixed_helper() {
+; CHECK-LABEL: define {{[^@]+}}@mixed_helper() {
+; CHECK-NEXT:    [[LEVEL:%.*]] = call i8 @__kmpc_parallel_level()
+; CHECK-NEXT:    store i8 [[LEVEL]], i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %level = call i8 @__kmpc_parallel_level()
+  store i8 %level, i8* @G
+  ret void
+}
+
+define internal void @none_spmd_helper() {
+; CHECK-LABEL: define {{[^@]+}}@none_spmd_helper() {
+; 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
+}
+
+define internal void @spmd_helper() {
+; CHECK-LABEL: define {{[^@]+}}@spmd_helper() {
+; CHECK-NEXT:    store i8 1, i8* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %level = call i8 @__kmpc_parallel_level()
+  store i8 %level, i8* @G
+  ret void
+}
+
+define internal void @__kmpc_parallel_51(%struct.ident_t*, i32, i32, i32, i32, i8*, i8*, i8**, i64) {
+; CHECK-LABEL: define {{[^@]+}}@__kmpc_parallel_51
+; CHECK-SAME: (%struct.ident_t* noalias nocapture nofree readnone align 536870912 [[TMP0:%.*]], i32 [[TMP1:%.*]], i32 [[TMP2:%.*]], i32 [[TMP3:%.*]], i32 [[TMP4:%.*]], i8* noalias nocapture nofree readnone align 536870912 [[TMP5:%.*]], i8* noalias nocapture nofree readnone align 536870912 [[TMP6:%.*]], i8** noalias nocapture nofree readnone align 536870912 [[TMP7:%.*]], i64 [[TMP8:%.*]]) {
+; CHECK-NEXT:    call void @parallel_helper()
+; CHECK-NEXT:    ret void
+;
+  call void @parallel_helper()
+  ret void
+}
+
+define internal void @parallel_helper() {
+; CHECK-LABEL: define {{[^@]+}}@parallel_helper() {
+; CHECK-NEXT:    [[LEVEL:%.*]] = call i8 @__kmpc_parallel_level()
+; CHECK-NEXT:    store i8 [[LEVEL]], 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, !3, !4}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+!2 = !{void ()* @none_spmd, !"kernel", i32 1}
+!3 = !{void ()* @spmd, !"kernel", i32 1}
+!4 = !{void ()* @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 ()* @none_spmd, !"kernel", i32 1}
+; CHECK: [[META3:![0-9]+]] = !{void ()* @spmd, !"kernel", i32 1}
+; CHECK: [[META4:![0-9]+]] = !{void ()* @parallel, !"kernel", i32 1}
+;.