[Mlir-commits] [llvm] [mlir] [mlir][OpenMP] - Honor dependencies in code-generation of the if clause in `omp.task` correctly (PR #90891)

[Pranav Bhandarkar]

https://github.com/bhandarkar-pranav updated https://github.com/llvm/llvm-project/pull/90891

>From 70c818b7259372f24a198605eaf709a3d43f4a31 Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar at amd.com>
Date: Thu, 2 May 2024 13:50:43 -0500
Subject: [PATCH 1/2] [mlir][OpenMP] - Honor dependencies in code-generation of
 the if clause in omp.task correctly

This patch fixes code generation of the if clause specifically when the condition
evaluates to false and when the task directive has the depend clause on it.
When the if clause of a task construct evaluates to false, then the task is an
undeferred task. This undeferred task still has to honor dependencies. Previously,
the OpenMPIRbuilder didn't honor dependencies. This patch fixes that.

Fixes https://github.com/llvm/llvm-project/issues/90869
---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     |  9 ++++++++
 mlir/test/Target/LLVMIR/omptask_if_false.mlir | 23 +++++++++++++++++++
 2 files changed, 32 insertions(+)
 create mode 100644 mlir/test/Target/LLVMIR/omptask_if_false.mlir

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 4d2d352f7520b2..695308419e0d20 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1887,6 +1887,15 @@ OpenMPIRBuilder::createTask(const LocationDescription &Loc,
       SplitBlockAndInsertIfThenElse(IfCondition, IfTerminator, &ThenTI,
                                     &ElseTI);
       Builder.SetInsertPoint(ElseTI);
+
+      if (Dependencies.size()) {
+        Function *TaskWaitFn =
+            getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_taskwait_deps_51);
+        Builder.CreateCall(TaskWaitFn, {Ident, ThreadID, Builder.getInt32(Dependencies.size()),
+                                        DepArray, ConstantInt::get(Builder.getInt32Ty(), 0),
+                                        ConstantPointerNull::get(PointerType::getUnqual(M.getContext())),
+                                        ConstantInt::get(Builder.getInt32Ty(), false)});
+      }
       Function *TaskBeginFn =
           getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_begin_if0);
       Function *TaskCompleteFn =
diff --git a/mlir/test/Target/LLVMIR/omptask_if_false.mlir b/mlir/test/Target/LLVMIR/omptask_if_false.mlir
new file mode 100644
index 00000000000000..c7fd5eb62a9772
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptask_if_false.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<"dlti.endianness", "little">, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>>, fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = "", llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128", llvm.target_triple = "x86_64-unknown-linux-gnu", omp.is_gpu = false, omp.is_target_device = false, omp.version = #omp.version<version = 11>}
+{
+  llvm.func @foo_(%arg0: !llvm.ptr {fir.bindc_name = "n"}, %arg1: !llvm.ptr {fir.bindc_name = "r"}) attributes {fir.internal_name = "_QPfoo"} {
+    %0 = llvm.mlir.constant(false) : i1
+    omp.task if(%0) depend(taskdependin -> %arg0 : !llvm.ptr) {
+      %1 = llvm.load %arg0 : !llvm.ptr -> i32
+      llvm.store %1, %arg1 : i32, !llvm.ptr
+      omp.terminator
+    }
+    llvm.return
+  }
+  llvm.func @llvm.stacksave.p0() -> !llvm.ptr attributes {sym_visibility = "private"}
+  llvm.func @llvm.stackrestore.p0(!llvm.ptr) attributes {sym_visibility = "private"}
+}
+
+
+// CHECK: call void @__kmpc_omp_taskwait_deps_51
+// CHECK-NEXT: call void @__kmpc_omp_task_begin_if0
+// CHECK-NEXT: call void @foo_..omp_par
+// CHECK-NEXT: call void @__kmpc_omp_task_complete_if0
+

>From 09c87f091a77a8331e50904828c3e123e3a0ed2b Mon Sep 17 00:00:00 2001
From: Pranav Bhandarkar <pranav.bhandarkar at amd.com>
Date: Thu, 2 May 2024 14:00:28 -0500
Subject: [PATCH 2/2] Fix indentation

---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 695308419e0d20..7444778c783301 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1891,10 +1891,12 @@ OpenMPIRBuilder::createTask(const LocationDescription &Loc,
       if (Dependencies.size()) {
         Function *TaskWaitFn =
             getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_taskwait_deps_51);
-        Builder.CreateCall(TaskWaitFn, {Ident, ThreadID, Builder.getInt32(Dependencies.size()),
-                                        DepArray, ConstantInt::get(Builder.getInt32Ty(), 0),
-                                        ConstantPointerNull::get(PointerType::getUnqual(M.getContext())),
-                                        ConstantInt::get(Builder.getInt32Ty(), false)});
+        Builder.CreateCall(
+            TaskWaitFn,
+            {Ident, ThreadID, Builder.getInt32(Dependencies.size()), DepArray,
+             ConstantInt::get(Builder.getInt32Ty(), 0),
+             ConstantPointerNull::get(PointerType::getUnqual(M.getContext())),
+             ConstantInt::get(Builder.getInt32Ty(), false)});
       }
       Function *TaskBeginFn =
           getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_begin_if0);