[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/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/IR/IntrinsicInst.h"
+#include "llvm/IR/IntrinsicsAMDGPU.h"
+#include "llvm/IR/IntrinsicsNVPTX.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"
 
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 true;
+    if (auto *II = dyn_cast<IntrinsicInst>(Cmp->getOperand(0)))
+      if (II->getIntrinsicID() == Intrinsic::amdgcn_workitem_id_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
+define void @kernel() {
+  call void @__kmpc_kernel_init(i32 512, i16 1)
+  call void @nvptx()
+  call void @amdgcn()
+  ret void
+}
 
-; CHECK: [openmp-opt] Basic block @bar entry is executed by a single thread.
+; CHECK-NOT: [openmp-opt] Basic block @nvptx entry is executed by a single thread.
+; CHECK: [openmp-opt] Basic block @nvptx if.then is executed by a single thread.
+; CHECK-NOT: [openmp-opt] Basic block @nvptx if.end is executed by a single thread.
 ; Function Attrs: noinline nounwind uwtable
-define internal void @bar() {
+define dso_local void @nvptx() {
 entry:
+  %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
 }
 
-; 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.
+; CHECK-NOT: [openmp-opt] Basic block @amdgcn entry is executed by a single thread.
+; CHECK: [openmp-opt] Basic block @amdgcn if.then is executed by a single thread.
+; CHECK-NOT: [openmp-opt] Basic block @amdgcn if.end is executed by a single thread.
 ; Function Attrs: noinline nounwind uwtable
-define dso_local void @foo() {
+define dso_local void @amdgcn() {
 entry:
-  %call = call i32 @omp_get_thread_num()
+  %call = call i32 @llvm.amdgcn.workitem.id.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()
+; 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
+}
+
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+
+declare i32 @llvm.amdgcn.workitem.id.x()
+
+declare void @__kmpc_kernel_init(i32, i16)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4}
+!nvvm.annotations = !{!5}
 
-!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}
+!0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 12.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "single_threaded_execution.c", directory: "/tmp/single_threaded_execution.c")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{void ()* @kernel, !"kernel", i32 1}