[OpenMP][IRBuilder] Handle target directives with both if & nowait

[ergawy]

This fixes a bug when a target directive has both an if and a nowait clauses. The bug happens because we tried to emitKernelLaunch for else branch of the if clause.

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-flang-openmp

@llvm/pr-subscribers-mlir-llvm

Author: Kareem Ergawy (ergawy)

Changes

This fixes a bug when a target directive has both an if and a nowait clauses. The bug happens because we tried to emitKernelLaunch for else branch of the if clause.

---

Full diff: https://github.com/llvm/llvm-project/pull/125029.diff

2 Files Affected:

• (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+6-1)
• (added) mlir/test/Target/LLVMIR/omptarget-if-nowait.mlir (+46)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 4c4a0d25906cbf..f618613b9ccad2 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -7438,10 +7438,15 @@ emitTargetCall(OpenMPIRBuilder &OMPBuilder, IRBuilderBase &Builder,
       // '@.omp_target_task_proxy_func' in the pseudo code above)
       // "@.omp_target_task_proxy_func' is generated by
       // emitTargetTaskProxyFunction.
-      if (OutlinedFnID)
+      //
+      // We only need to do the outlining if `DeviceID` is set to avoid calling
+      // `emitKernelLaunch` if we want to code-gen for the host; e.g. if we are
+      // generating the `else` branch of an `if` clause.
+      if (OutlinedFnID && DeviceID)
         return OMPBuilder.emitKernelLaunch(Builder, OutlinedFnID,
                                            EmitTargetCallFallbackCB, KArgs,
                                            DeviceID, RTLoc, TargetTaskAllocaIP);
+
       // When OutlinedFnID is set to nullptr, then it's not an offloading call.
       // In this case, we execute the host implementation directly.
       return EmitTargetCallFallbackCB(OMPBuilder.Builder.saveIP());
diff --git a/mlir/test/Target/LLVMIR/omptarget-if-nowait.mlir b/mlir/test/Target/LLVMIR/omptarget-if-nowait.mlir
new file mode 100644
index 00000000000000..6f8d938a4d5f26
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-if-nowait.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {omp.is_target_device = false, omp.target_triples = ["amdgcn-amd-amdhsa"]} {
+  llvm.func @target_if_nowait(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
+    %0 = llvm.mlir.constant(1 : i64) : i64
+    %3 = llvm.alloca %0 x i32 {bindc_name = "cond"} : (i64) -> !llvm.ptr
+    %6 = llvm.load %3 : !llvm.ptr -> i32
+    %7 = llvm.mlir.constant(0 : i64) : i32
+    %8 = llvm.icmp "ne" %6, %7 : i32
+    %9 = omp.map.info var_ptr(%3 : !llvm.ptr, i32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "cond"}
+    %10 = omp.map.info var_ptr(%arg0 : !llvm.ptr, f32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "var"}
+    %11 = omp.map.info var_ptr(%arg1 : !llvm.ptr, f32) map_clauses(implicit, exit_release_or_enter_alloc) capture(ByCopy) -> !llvm.ptr {name = "val"}
+    omp.target if(%8) nowait map_entries(%10 -> %arg3, %11 -> %arg4 : !llvm.ptr, !llvm.ptr) {
+      %12 = llvm.load %arg4 : !llvm.ptr -> f32
+      llvm.store %12, %arg3 : f32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+}
+
+// CHECK: define void @target_if_nowait{{.*}} {
+// CHECK: omp_if.then:
+// CHECK:   br label %[[TARGET_TASK_BB:.*]]
+
+// CHECK: [[TARGET_TASK_BB]]:
+// CHECK:   call ptr @__kmpc_omp_target_task_alloc
+// CHECK:   br label %[[OFFLOAD_CONT:.*]]
+
+// CHECK: [[OFFLOAD_CONT]]:
+// CHECK:   br label %omp_if.end
+
+// CHECK: omp_if.else:
+// CHECK:   br label %[[HOST_TASK_BB:.*]]
+
+// CHECK: [[HOST_TASK_BB]]:
+// CHECK:   call ptr @__kmpc_omp_task_alloc
+// CHECK:   br label %[[HOST_TASK_CONT:.*]]
+
+// CHECK: [[HOST_TASK_CONT]]:
+// CHECK:   br label %omp_if.end
+
+// CHECK: omp_if.end:
+// CHECK:   ret void
+// CHECK: }
+

[skatrak]

Thank you, LGTM.

[bhandarkar-pranav]

LGTM. Thank you for the fix.

[llvm-ci]

LLVM Buildbot has detected a new failure on builder llvm-clang-aarch64-darwin running on doug-worker-4 while building llvm,mlir at step 6 "test-build-unified-tree-check-all".

Full details are available at: https://lab.llvm.org/buildbot/#/builders/190/builds/13812

Here is the relevant piece of the build log for the reference

Step 6 (test-build-unified-tree-check-all) failure: test (failure)
******************** TEST 'LLVM :: ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll' FAILED ********************
Exit Code: 2

Command Output (stderr):
--
RUN: at line 1: /Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli -jit-kind=orc-lazy -compile-threads=2 -thread-entry hello /Users/buildbot/buildbot-root/aarch64-darwin/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll | /Users/buildbot/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/aarch64-darwin/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
+ /Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli -jit-kind=orc-lazy -compile-threads=2 -thread-entry hello /Users/buildbot/buildbot-root/aarch64-darwin/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
+ /Users/buildbot/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/aarch64-darwin/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
 #0 0x00000001012ed4e0 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100eb94e0)
 #1 0x00000001012eb564 llvm::sys::RunSignalHandlers() (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100eb7564)
 #2 0x00000001012edb9c SignalHandler(int) (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100eb9b9c)
 #3 0x0000000198d56584 (/usr/lib/system/libsystem_platform.dylib+0x18047a584)
 #4 0x0000000198d2521c (/usr/lib/system/libsystem_pthread.dylib+0x18044921c)
 #5 0x0000000198c4bad0 (/usr/lib/libc++.1.dylib+0x18036fad0)
 #6 0x0000000100e93c8c void llvm::detail::UniqueFunctionBase<void, llvm::Expected<llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef, llvm::DenseMapInfo<llvm::orc::SymbolStringPtr, void>, llvm::detail::DenseMapPair<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef>>>>::CallImpl<llvm::orc::Platform::lookupInitSymbols(llvm::orc::ExecutionSession&, llvm::DenseMap<llvm::orc::JITDylib*, llvm::orc::SymbolLookupSet, llvm::DenseMapInfo<llvm::orc::JITDylib*, void>, llvm::detail::DenseMapPair<llvm::orc::JITDylib*, llvm::orc::SymbolLookupSet>> const&)::$_45>(void*, llvm::Expected<llvm::DenseMap<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef, llvm::DenseMapInfo<llvm::orc::SymbolStringPtr, void>, llvm::detail::DenseMapPair<llvm::orc::SymbolStringPtr, llvm::orc::ExecutorSymbolDef>>>&) (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100a5fc8c)
 #7 0x0000000100e8fa7c llvm::orc::AsynchronousSymbolQuery::handleComplete(llvm::orc::ExecutionSession&)::RunQueryCompleteTask::run() (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100a5ba7c)
 #8 0x0000000100f4c2d8 void* std::__1::__thread_proxy[abi:un170006]<std::__1::tuple<std::__1::unique_ptr<std::__1::__thread_struct, std::__1::default_delete<std::__1::__thread_struct>>, llvm::orc::DynamicThreadPoolTaskDispatcher::dispatch(std::__1::unique_ptr<llvm::orc::Task, std::__1::default_delete<llvm::orc::Task>>)::$_0>>(void*) (/Users/buildbot/buildbot-root/aarch64-darwin/build/bin/lli+0x100b182d8)
 #9 0x0000000198d25f94 (/usr/lib/system/libsystem_pthread.dylib+0x180449f94)
#10 0x0000000198d20d34 (/usr/lib/system/libsystem_pthread.dylib+0x180444d34)
FileCheck error: '<stdin>' is empty.
FileCheck command line:  /Users/buildbot/buildbot-root/aarch64-darwin/build/bin/FileCheck /Users/buildbot/buildbot-root/aarch64-darwin/llvm-project/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll

--

********************