[OpenMP] Fold `__kmpc_is_generic_main_thread_id` if possible

[jhuber6]

Summary

This patch adds the ability to fold __kmpc_is_generic_main_thread_id if we
know for a fact that it is executed by the initial thread using
AAExecutionDomain. This combined with folding __kmpc_is_spmd_exec_mode will
allow us to fully fold __kmpc_is_generic_main_thread.

Depends on D106438 D106437

[jhuber6]

Fixing hang caused by folding call to an undefined value.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3632

llvm/test/Transforms/OpenMP/fold_generic_main_thread.ll

40

Shouldn't this be folded?

[jhuber6]

llvm/test/Transforms/OpenMP/fold_generic_main_thread.ll

40

It should be folded. I have CHECK-NOT because the only way that call can be removed is through the folding because it's only a declaration in the module. I couldn't decide if I should just use the automatic test script for this, personally I prefer tests like this to be narrow if possible.

[jdoerfert]

llvm/test/Transforms/OpenMP/fold_generic_main_thread.ll

40

check-not like above is really easy to break. Let's assume the type is changed, the name of the argument, etc.
If you want it narrow, check for %pred = icmp eq i8 1, 1 or br i1 true, label %if.then, label %if.end

[jhuber6]

Running test script on test and indicate pessimistic fixpoint.

[jdoerfert]

LG

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

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_is_generic_main_thread_id:
+      Changed |= foldIsGenericMainThread(A);
+      break;
     default:
       llvm_unreachable("Unhandled OpenMP runtime function!");
     }
 
     return Changed;
   }
 
   ChangeStatus manifest(Attributor &A) override {
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
 
     if (SimplifiedValue.hasValue() && SimplifiedValue.getValue()) {
       Instruction &CB = *getCtxI();
       A.changeValueAfterManifest(CB, **SimplifiedValue);
       A.deleteAfterManifest(CB);
+
+      LLVM_DEBUG(dbgs() << TAG << "Folding runtime call: " << CB << " with "
+                        << **SimplifiedValue << "\n");
+
       Changed = ChangeStatus::CHANGED;
     }
 
     return Changed;
   }
 
   ChangeStatus indicatePessimisticFixpoint() override {
     SimplifiedValue = nullptr;
if (KnownSPMDCount || AssumedSPMDCount) {
       // must be none.
       assert(!SimplifiedValue.hasValue() && "SimplifiedValue should be none");
     }
 
     return SimplifiedValue == SimplifiedValueBefore ? ChangeStatus::UNCHANGED
                                                     : ChangeStatus::CHANGED;
   }
 
+  /// Fold __kmpc_is_generic_main_thread_id into a constant if possible.
+  ChangeStatus foldIsGenericMainThread(Attributor &A) {
+    Optional<Value *> SimplifiedValueBefore = SimplifiedValue;
+
+    CallBase &CB = cast<CallBase>(getAssociatedValue());
+    Function *F = CB.getFunction();
+    const auto &ExecutionDomainAA = A.getAAFor<AAExecutionDomain>(
+        *this, IRPosition::function(*F), DepClassTy::REQUIRED);
+
+    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;
+  }
+
   /// 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;
 };
if (IsModulePass) {
     // callbacks before any other AA has the chance to create an AAValueSimplify
     // or similar.
     for (Function *Kernel : OMPInfoCache.Kernels)
       A.getOrCreateAAFor<AAKernelInfo>(
           IRPosition::function(*Kernel), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,
           /* UpdateAfterInit */ false);
 
+    auto &IsMainRFI =
+        OMPInfoCache.RFIs[OMPRTL___kmpc_is_generic_main_thread_id];
+    IsMainRFI.foreachUse(SCC, [&](Use &U, Function &F) {
+      CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &IsMainRFI);
+      if (!CI)
+        return false;
+      A.getOrCreateAAFor<AAFoldRuntimeCall>(
+          IRPosition::callsite_returned(*CI), /* QueryingAA */ nullptr,
+          DepClassTy::NONE, /* ForceUpdate */ false,
+          /* UpdateAfterInit */ false);
+      return false;
+    });
+
     auto &IsSPMDRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_is_spmd_exec_mode];
     IsSPMDRFI.foreachUse(SCC, [&](Use &U, Function &) {
       CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &IsSPMDRFI);
       if (!CI)
         return false;
       A.getOrCreateAAFor<AAFoldRuntimeCall>(
           IRPosition::callsite_returned(*CI), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,

llvm/test/Transforms/OpenMP/fold_generic_main_thread.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; RUN: opt -S -passes='openmp-opt' < %s | FileCheck %s
+; ModuleID = 'single_threaded_exeuction.c'
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@0 = private unnamed_addr constant [1 x i8] c"\00", align 1
+@1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 0, i8* getelementptr inbounds ([1 x i8], [1 x i8]* @0, i32 0, i32 0) }, align 8
+
+define void @kernel() {
+; CHECK-LABEL: define {{[^@]+}}@kernel() {
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* nonnull @[[GLOB1:[0-9]+]], i1 false, i1 false, i1 false)
+; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[CALL]], -1
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    br label [[IF_END:%.*]]
+; CHECK:       if.else:
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 true)
+; CHECK-NEXT:    ret void
+;
+  %call = call i32 @__kmpc_target_init(%struct.ident_t* nonnull @1, i1 false, i1 false, i1 false)
+  %cmp = icmp eq i32 %call, -1
+  br i1 %cmp, label %if.then, label %if.else
+if.then:
+  call void @foo()
+  br label %if.end
+if.else:
+  call void @bar()
+  br label %if.end
+if.end:
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 true)
+  ret void
+}
+
+; Function Attrs: noinline
+define internal void @foo() {
+; CHECK-LABEL: define {{[^@]+}}@foo() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TID:%.*]] = call i32 @__kmpc_get_hardware_thread_id()
+; CHECK-NEXT:    br label [[IF_THEN:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    call void @baz()
+; CHECK-NEXT:    br label [[IF_END:%.*]]
+; CHECK:       if.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tid = call i32 @__kmpc_get_hardware_thread_id();
+  %ismain = call signext i8 @__kmpc_is_generic_main_thread_id(i32 %tid)
+  %pred = icmp eq i8 %ismain, 1
+  br i1 %pred, label %if.then, label %if.end
+
+if.then:
+  call void @baz()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+; Function Attrs: noinline
+define internal void @bar() {
+; CHECK-LABEL: define {{[^@]+}}@bar() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TID:%.*]] = call i32 @__kmpc_get_hardware_thread_id()
+; CHECK-NEXT:    [[ISMAIN:%.*]] = call signext i8 @__kmpc_is_generic_main_thread_id(i32 [[TID]])
+; CHECK-NEXT:    [[PRED:%.*]] = icmp eq i8 [[ISMAIN]], 1
+; CHECK-NEXT:    br i1 [[PRED]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    call void @baz()
+; CHECK-NEXT:    br label [[IF_END]]
+; CHECK:       if.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %tid = call i32 @__kmpc_get_hardware_thread_id();
+  %ismain = call signext i8 @__kmpc_is_generic_main_thread_id(i32 %tid)
+  %pred = icmp eq i8 %ismain, 1
+  br i1 %pred, label %if.then, label %if.end
+
+if.then:
+  call void @baz()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare void @baz()
+
+declare i8 @__kmpc_is_generic_main_thread_id(i32)
+
+declare i32 @__kmpc_get_hardware_thread_id()
+
+declare i32 @__kmpc_target_init(%struct.ident_t*, i1, i1, i1)
+
+declare void @__kmpc_target_deinit(%struct.ident_t*, i1, i1)
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4, !5, !6}
+!nvvm.annotations = !{!7}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 13.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "fold_generic_main_thread.c", directory: "/tmp/fold_generic_main_thread.c")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 7, !"openmp", i32 50}
+!6 = !{i32 7, !"openmp-device", i32 50}
+!7 = !{void ()* @kernel, !"kernel", i32 1}