[llvm] 83433d9 - [OpenMP][IRBuilder] Handle `target ... nowait` when codegen targets host (#124720)

[via llvm-commits]

Author: Kareem Ergawy
Date: 2025-01-28T12:57:40+01:00
New Revision: 83433d936195c612a51b54397f82ab0d97369d86

URL: https://github.com/llvm/llvm-project/commit/83433d936195c612a51b54397f82ab0d97369d86
DIFF: https://github.com/llvm/llvm-project/commit/83433d936195c612a51b54397f82ab0d97369d86.diff

LOG: [OpenMP][IRBuilder] Handle `target ... nowait` when codegen targets host (#124720)

Fixes https://github.com/llvm/llvm-project/issues/124578

Handles the `nowait` clause for `omp.target` ops when the actual target
is the host (i.e. there is no target device). Rather than only checking
for the `HasNoWait` boolean, we also check for the presence/absence of a
`DeviceID` value. We only emit the target task if both are present.

Added: 
    mlir/test/Target/LLVMIR/omptarget-nowait-host-only.mlir

Modified: 
    llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 3e74f4c6d8ccdb..6d234262e402fa 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -7258,10 +7258,12 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
     // If `HasNoWait == true`, we call  @__kmpc_omp_target_task_alloc to provide
     // the DeviceID to the deferred task and also since
     // @__kmpc_omp_target_task_alloc creates an untied/async task.
+    bool NeedsTargetTask = HasNoWait && DeviceID;
     Function *TaskAllocFn =
-        !HasNoWait ? getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_alloc)
-                   : getOrCreateRuntimeFunctionPtr(
-                         OMPRTL___kmpc_omp_target_task_alloc);
+        !NeedsTargetTask
+            ? getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_alloc)
+            : getOrCreateRuntimeFunctionPtr(
+                  OMPRTL___kmpc_omp_target_task_alloc);
 
     // Arguments - `loc_ref` (Ident) and `gtid` (ThreadID)
     // call.
@@ -7310,8 +7312,10 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
         /*sizeof_task=*/TaskSize, /*sizeof_shared=*/SharedsSize,
         /*task_func=*/ProxyFn};
 
-    if (HasNoWait)
+    if (NeedsTargetTask) {
+      assert(DeviceID && "Expected non-empty device ID.");
       TaskAllocArgs.push_back(DeviceID);
+    }
 
     TaskData = Builder.CreateCall(TaskAllocFn, TaskAllocArgs);
 
@@ -7333,7 +7337,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
     // ---------------------------------------------------------------
     // The above means that the lack of a nowait on the target construct
     // translates to '#pragma omp task if(0)'
-    if (!HasNoWait) {
+    if (!NeedsTargetTask) {
       if (DepArray) {
         Function *TaskWaitFn =
             getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_wait_deps);

diff  --git a/mlir/test/Target/LLVMIR/omptarget-nowait-host-only.mlir b/mlir/test/Target/LLVMIR/omptarget-nowait-host-only.mlir
new file mode 100644
index 00000000000000..6b634226a35686
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-nowait-host-only.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Tests `target ... nowait` when code gen targets the host rather than a
+// device.
+
+module attributes {omp.is_target_device = false} {
+  llvm.func @omp_target_nowait_() {
+    %0 = llvm.mlir.constant(1 : i64) : i64
+    %1 = llvm.alloca %0 x f32 {bindc_name = "x"} : (i64) -> !llvm.ptr
+    %3 = omp.map.info var_ptr(%1 : !llvm.ptr, f32) map_clauses(tofrom) capture(ByRef) -> !llvm.ptr {name = "x"}
+    omp.target nowait map_entries(%3 -> %arg0 : !llvm.ptr) {
+      %4 = llvm.mlir.constant(5.000000e+00 : f32) : f32
+      llvm.store %4, %arg0 : f32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+}
+
+// CHECK: define void @omp_target_nowait_()
+// CHECK-NOT: define {{.*}} @
+// CHECK-NOT: call ptr @__kmpc_omp_target_task_alloc({{.*}})
+// Verify that we directly emit a call to the "target" region's body from the
+// parent function of the the `omp.target` op.
+// CHECK: call void @__omp_offloading_[[DEV:.*]]_[[FIL:.*]]_omp_target_nowait__l[[LINE:.*]](ptr {{.*}})
+// CHECK-NEXT: ret void
+
+// CHECK: define internal void @__omp_offloading_[[DEV]]_[[FIL]]_omp_target_nowait__l[[LINE]](ptr %[[ADDR_X:.*]])
+// CHECK: store float 5{{.*}}, ptr %[[ADDR_X]], align 4