[Attributor] Add AAExecutionDomainInfo interface to OpenMPOpt

[jhuber6]

Summary

Add the AAExecutionDomainInfo attributor instance to OpenMPOpt.
This will infer information relating to domain information that an
instruction might be expecting in. Right now this only includes a very
crude check for instructions that will be executed by the master thread
by comparing a thread-id function with a constant zero.

[jdoerfert]

Looks like a reasonable start, some comments:

Hook it up to the globalization to make it easier to test it. We also need negative tests and such, swapped conditions, swapped edges, ...

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

3851 ↗

(On Diff #341661)

Comments are copy&paste but I don't assume we need these functions at all.
The interface we want is below/

3866 ↗

(On Diff #341661)

Nit: take references

llvm/lib/Transforms/IPO/AttributorAttributes.cpp

4347 ↗

(On Diff #341661)

Not needed

8150 ↗

(On Diff #341661)

8177 ↗

(On Diff #341661)

-Master
+MainThreadOnly or FirstThreadOnly

---

I don't get the TrueBB part, you don't know which successor the block you are looking at is from the code below. Make it "SuccBB" and check the condition first, then check if this is the successor under which the condition is tid == 0.

---

Given that we have all the openmp specific functions in here we should move this into OpenMPOpt.

8185 ↗

(On Diff #341661)

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1166

I assume this is only for testing, right?

[jhuber6]

Moving to OpenMPOpt.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2281

return how many blocks are single threaded of how many, as string: "3/7 BBs thread 0 only"

2312

You could use a set, during initialize put all blocks in the set, remove one if you cannot argue it's single threaded. you know something changed if the size at the end of update is not the size at the beginning. That should make things a lot simpler, whenever you fail to show single threaded, remove it from the set right away, no need to communicate booleans all over the place and check the set all the time to update "changed".

2333

also check if it is an equality, this would allow tid >= 0, I think.

2369

Move down to the usage, also IsSingleThreadOnly

2380–2382

2444

unreachable for all but function. See the other createForPosition

[jhuber6]

Making suggested changes.

[jdoerfert]

LG, two minor nits.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

2298

Here and elsewhere, if the type is trivial, use it. 4 vs 9 chars but better readability.
E is not a good variable name for a BB.

2339–2341

[jhuber6]

Fixing nits.

[sebastian-ne]

Hi,
The added test llvm/test/Transforms/OpenMP/single_threaded_execution.ll fails in release builds when assertions are disabled.
You probably need to add ; REQUIRES: asserts like it is done in llvm/test/CodeGen/Mips/brind-tailcall.ll.

[jhuber6]

In D101578#2735986, @sebastian-ne wrote:

Hi,
The added test llvm/test/Transforms/OpenMP/single_threaded_execution.ll fails in release builds when assertions are disabled.
You probably need to add ; REQUIRES: asserts like it is done in llvm/test/CodeGen/Mips/brind-tailcall.ll.

Thanks, seems like someone already fixed it in 05146fe5171076eaec3ef6a9eeeab4c8a33acf15.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 //
 // - Deduplication of runtime calls, e.g., omp_get_thread_num.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/OpenMPOpt.h"
 
 #include "llvm/ADT/EnumeratedArray.h"
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/InitializePasses.h"
STATISTIC(NumOpenMPParallelRegionsMerged,
           "Number of OpenMP parallel regions merged");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
 #endif
 
 namespace {
 
+struct AAExecutionDomain
+    : public StateWrapper<BooleanState, AbstractAttribute> {
+  using Base = StateWrapper<BooleanState, AbstractAttribute>;
+  AAExecutionDomain(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAExecutionDomain &createForPosition(const IRPosition &IRP,
+                                              Attributor &A);
+
+  /// See AbstractAttribute::getName().
+  const std::string getName() const override { return "AAExecutionDomain"; }
+
+  /// See AbstractAttribute::getIdAddr().
+  const char *getIdAddr() const override { return &ID; }
+
+  /// Check if an instruction is executed by a single thread.
+  virtual bool isSingleThreadExeuction(const Instruction &) const = 0;
+
+  virtual bool isSingleThreadExeuction(const BasicBlock &) const = 0;
+
+  /// This function should return true if the type of the \p AA is
+  /// AAExecutionDomain.
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  /// Unique ID (due to the unique address)
+  static const char ID;
+};
+
 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<Kernel> &Kernels)
for (const auto GlobalizationCallID : GlobalizationRuntimeIDs) {
         if (CallInst *CI = getCallIfRegularCall(U, &RFI)) {
           auto Remark = [&](OptimizationRemarkAnalysis ORA) {
             return ORA
                    << "Found thread data sharing on the GPU. "
                    << "Expect degraded performance due to data globalization.";
           };
           emitRemark<OptimizationRemarkAnalysis>(CI, "OpenMPGlobalization",
                                                  Remark);
         }
 
         return false;
       };
 
       RFI.foreachUse(SCC, CheckGlobalization);
     }
   }
 
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
 
         IRPosition CBPos = IRPosition::callsite_function(CB);
         A.getOrCreateAAFor<AAICVTracker>(CBPos);
         return false;
       };
 
       GetterRFI.foreachUse(SCC, CreateAA);
     }
+
+    for (auto &F : M) {
+      if (!F.isDeclaration())
+        A.getOrCreateAAFor<AAExecutionDomain>(IRPosition::function(F));
+    }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
for (InternalControlVar ICV : TrackableICVs) {
         continue;
 
       ReplVal = NewReplVal;
       Changed = ChangeStatus::CHANGED;
     }
     return Changed;
   }
 };
+
+struct AAExecutionDomainFunction : public AAExecutionDomain {
+  AAExecutionDomainFunction(const IRPosition &IRP, Attributor &A)
+      : AAExecutionDomain(IRP, A) {}
+
+  const std::string getAsStr() const override {
+    return "[AAExecutionDomain] " + std::to_string(SingleThreadedBBs.size()) +
+           "/" + std::to_string(NumBBs) + " BBs thread 0 only.";
+  }
+
+  /// See AbstractAttribute::trackStatistics().
+  void trackStatistics() const override {}
+
+  void initialize(Attributor &A) override {
+    Function *F = getAnchorScope();
+    for (const auto &BB : *F)
+      SingleThreadedBBs.insert(&BB);
+    NumBBs = SingleThreadedBBs.size();
+  }
+
+  ChangeStatus manifest(Attributor &A) override {
+    LLVM_DEBUG({
+      for (const auto &E : SingleThreadedBBs)
+        dbgs() << TAG << " Basic block @" << getAnchorScope()->getName() << " "
+               << E->getName() << " is executed by a single thread.\n";
+    });
+    return ChangeStatus::UNCHANGED;
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override;
+
+  /// Check if an instruction is executed by a single thread.
+  bool isSingleThreadExeuction(const Instruction &I) const override {
+    return isSingleThreadExeuction(*I.getParent());
+  }
+
+  bool isSingleThreadExeuction(const BasicBlock &BB) const override {
+    return SingleThreadedBBs.contains(&BB);
+  }
+
+  /// Map between basic blocks and their execution domain.
+  DenseSet<const BasicBlock *> SingleThreadedBBs;
+
+  /// Total number of basic blocks in this function.
+  long unsigned NumBBs;
+};
+
+ChangeStatus AAExecutionDomainFunction::updateImpl(Attributor &A) {
+  auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+  Function *F = getAnchorScope();
+  ReversePostOrderTraversal<Function *> RPOT(F);
+  auto NumSingleThreadedBBs = SingleThreadedBBs.size();
+
+  bool AllCallSitesKnown;
+  auto PredForCallSite = [&](AbstractCallSite ACS) {
+    const auto &ExecutionDomainAA = A.getAAFor<AAExecutionDomain>(
+        *this, IRPosition::function(*ACS.getInstruction()->getFunction()),
+        DepClassTy::REQUIRED);
+    return ExecutionDomainAA.isSingleThreadExeuction(*ACS.getInstruction());
+  };
+
+  if (!A.checkForAllCallSites(PredForCallSite, *this,
+                              /* RequiresAllCallSites */ true,
+                              AllCallSitesKnown)) {
+    SingleThreadedBBs.erase(&F->getEntryBlock());
+  }
+
+  // Check if the edge into the successor block compares a thread-id function to
+  // a constant zero.
+  // TODO: Use AAValueSimplify to simplify and propogate constants.
+  // TODO: Check more than a single use for thread ID's.
+  auto IsSingleThreadOnly = [&](BranchInst *Edge, BasicBlock *SuccessorBB) {
+    if (!Edge || !Edge->isConditional())
+      return false;
+    if (Edge->getSuccessor(0) != SuccessorBB)
+      return false;
+
+    auto *Cmp = dyn_cast<CmpInst>(Edge->getCondition());
+    if (!Cmp || !Cmp->isTrueWhenEqual() || !Cmp->isEquality())
+      return false;
+
+    ConstantInt *C = dyn_cast<ConstantInt>(Cmp->getOperand(1));
+    if (!C || !C->isZero())
+      return false;
+
+    if (auto *CB = dyn_cast<CallBase>(Cmp->getOperand(0))) {
+      RuntimeFunction ThreadNumRuntimeIDs[] = {OMPRTL_omp_get_thread_num,
+                                               OMPRTL___kmpc_master,
+                                               OMPRTL___kmpc_global_thread_num};
+
+      for (const auto ThreadNumRuntimeID : ThreadNumRuntimeIDs) {
+        auto &RFI = OMPInfoCache.RFIs[ThreadNumRuntimeID];
+        if (CB->getCalledFunction() == RFI.Declaration)
+          return true;
+      }
+    }
+
+    return false;
+  };
+
+  // Merge all the predecessor states into the current basic block. A basic
+  // block is executed by a single thread if all of its predecessors are.
+  auto MergePredecessorStates = [&](BasicBlock *BB) {
+    if (pred_begin(BB) == pred_end(BB))
+      return SingleThreadedBBs.contains(BB);
+
+    bool IsSingleThreaded = true;
+    for (auto PredBB = pred_begin(BB), PredEndBB = pred_end(BB);
+         PredBB != PredEndBB; ++PredBB) {
+      if (!IsSingleThreadOnly(dyn_cast<BranchInst>((*PredBB)->getTerminator()),
+                              BB))
+        IsSingleThreaded &= SingleThreadedBBs.contains(*PredBB);
+    }
+
+    return IsSingleThreaded;
+  };
+
+  for (auto *BB : RPOT) {
+    if (!MergePredecessorStates(BB))
+      SingleThreadedBBs.erase(BB);
+  }
+
+  return (NumSingleThreadedBBs == SingleThreadedBBs.size())
+             ? ChangeStatus::UNCHANGED
+             : ChangeStatus::CHANGED;
+}
+
 } // namespace
 
 const char AAICVTracker::ID = 0;
+const char AAExecutionDomain::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:
AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAICVTrackerFunction(IRP, A);
     break;
   }
 
   return *AA;
 }
 
+AAExecutionDomain &AAExecutionDomain::createForPosition(const IRPosition &IRP,
+                                                        Attributor &A) {
+  AAExecutionDomainFunction *AA = nullptr;
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_INVALID:
+  case IRPosition::IRP_FLOAT:
+  case IRPosition::IRP_ARGUMENT:
+  case IRPosition::IRP_CALL_SITE_ARGUMENT:
+  case IRPosition::IRP_RETURNED:
+  case IRPosition::IRP_CALL_SITE_RETURNED:
+  case IRPosition::IRP_CALL_SITE:
+    llvm_unreachable(
+        "AAExecutionDomain can only be created for function position!");
+  case IRPosition::IRP_FUNCTION:
+    AA = new (A.Allocator) AAExecutionDomainFunction(IRP, A);
+    break;
+  }
+
+  return *AA;
+}
+
 PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   if (!containsOpenMP(M, OMPInModule))
     return PreservedAnalyses::all();
 
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   // Look at every function definition in the Module.

llvm/test/Transforms/OpenMP/single_threaded_execution.ll

+; RUN: opt -passes=openmp-opt-cgscc -debug-only=openmp-opt -disable-output < %s 2>&1 | FileCheck %s
+; ModuleID = 'single_threaded_exeuction.c'
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@.str = private unnamed_addr constant [4 x i8] c"%d\0A\00", align 1
+@0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
+@1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 0, i8* getelementptr inbounds ([23 x i8], [23 x i8]* @0, i32 0, i32 0) }, align 8
+
+; CHECK: [openmp-opt] Basic block @bar entry is executed by a single thread.
+; Function Attrs: noinline nounwind uwtable
+define internal void @bar() {
+entry:
+  ret void
+}
+
+; CHECK-NOT: [openmp-opt] Basic block @foo entry is executed by a single thread.
+; CHECK: [openmp-opt] Basic block @foo if.then is executed by a single thread.
+; CHECK-NOT: [openmp-opt] Basic block @foo if.end is executed by a single thread.
+; Function Attrs: noinline nounwind uwtable
+define dso_local void @foo() {
+entry:
+  %call = call i32 @omp_get_thread_num()
+  %cmp = icmp eq i32 %call, 0
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:
+  call void @bar()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare dso_local i32 @omp_get_thread_num()
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 13.0.0"}
+!2 = !{!3}
+!3 = !{i64 2, i64 -1, i64 -1, i1 true}