[llvm] 205b0bd - [OpenMP][IRBuilder] Handle `target` directives with both `if` & `nowait` (#125029)

[via llvm-commits]

Author: Kareem Ergawy
Date: 2025-01-30T18:40:34+01:00
New Revision: 205b0bddcd8bdd19381e095a547302fffaf2750d

URL: https://github.com/llvm/llvm-project/commit/205b0bddcd8bdd19381e095a547302fffaf2750d
DIFF: https://github.com/llvm/llvm-project/commit/205b0bddcd8bdd19381e095a547302fffaf2750d.diff

LOG: [OpenMP][IRBuilder] Handle `target` directives with both `if` & `nowait` (#125029)

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.

Added: 
    mlir/test/Target/LLVMIR/omptarget-if-nowait.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 4c4a0d25906cbf..695b15ac31f380 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)
+      if (OutlinedFnID && DeviceID)
         return OMPBuilder.emitKernelLaunch(Builder, OutlinedFnID,
                                            EmitTargetCallFallbackCB, KArgs,
                                            DeviceID, RTLoc, TargetTaskAllocaIP);
+
+      // 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.
+      //
       // 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: }
+