[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 created https://github.com/llvm/llvm-project/pull/90891

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

>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] [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
+