[Attributor] Change AAExecutionDomain to only accept intrinsics

[jhuber6]

Summary

The OpenMP runtime functions don't always provide unique thread ID's to
determine if a basic block is truly single-threaded. Change the implementation
to only check NVPTX intrinsics for now.

[jdoerfert]

LG, can you please add the AMDGPU intrinsic as well (with test)

[jhuber6]

Adding AMDGCN intrinsic check.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #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"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsNVPTX.h"
 
 using namespace llvm;
 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
struct AAExecutionDomainFunction : public AAExecutionDomain {
   /// Set of basic blocks that are executed by a single thread.
   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()),
auto IsSingleThreadOnly = [&](BranchInst *Edge, BasicBlock *SuccessorBB) {
     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;
-      }
-    }
+    if (auto *II = dyn_cast<IntrinsicInst>(Cmp->getOperand(0)))
+      if (II->getIntrinsicID() == Intrinsic::nvvm_read_ptx_sreg_tid_x)
 
     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))

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
 ; REQUIRES: asserts
 ; 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()
+  %dummy = call i32 @omp_get_thread_num()
+  %call = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
   %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()
 
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+
 !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}