[flang-commits] [flang] [mlir] [Flang][OpenMP] Properly bind arguments of composite operations (PR #113682)

[Sergio Afonso via flang-commits]

https://github.com/skatrak updated https://github.com/llvm/llvm-project/pull/113682

>From 58f02e5294f2b998e49d2f7f130dc97dc6545c34 Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Fri, 25 Oct 2024 11:45:19 +0100
Subject: [PATCH] [Flang][OpenMP] Properly bind arguments of composite
 operations

When composite constructs are lowered, clauses for each leaf construct are
lowered before creating the set of loop wrapper operations, using these outside
values to populate their operand lists. Then, when the loop nest associated to
that composite construct is lowered, the binding of Fortran symbols to the
entry block arguments defined by these loop wrappers is performed, resulting in
the creation of `hlfir.declare` operations in the entry block of the
`omp.loop_nest`.

This approach prevents `hlfir.declare` operations related to the binding and
other operations resulting from the evaluation of the clauses from being
inserted between loop wrapper operations, which would be an illegal MLIR
representation. However, this introduces the problem of entry block arguments
defined by a wrapper that then should be used by one of its nested wrappers,
because the corresponding Fortran symbol would still be mapped to an outside
value at the time of gathering the list of operands for the nested wrapper.

This patch adds operand re-mapping logic to update wrappers without changing
when clauses are evaluated or where the `hlfir.declare` creation is performed.
---
 flang/lib/Lower/OpenMP/OpenMP.cpp             | 21 +++++++++++++++++--
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      |  1 +
 mlir/test/Target/LLVMIR/openmp-todo.mlir      | 14 +++++++++++++
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 84985b880b1ec2..329cbf3d7539f5 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -589,10 +589,27 @@ static void genLoopVars(
   llvm::SmallVector<mlir::Location> locs(args.size(), loc);
   firOpBuilder.createBlock(&region, {}, tiv, locs);
 
+  // Update nested wrapper operands if parent wrappers have mapped these values
+  // to block arguments.
+  //
+  // Binding these values earlier would take care of this, but we cannot rely on
+  // that approach because binding in between the creation of a wrapper and the
+  // next one would result in 'hlfir.declare' operations being introduced inside
+  // of a wrapper, which is illegal.
+  mlir::IRMapping mapper;
+  for (auto [argGeneratingOp, blockArgs] : wrapperArgs) {
+    for (mlir::OpOperand &operand : argGeneratingOp->getOpOperands())
+      operand.set(mapper.lookupOrDefault(operand.get()));
+
+    for (const auto [arg, var] : llvm::zip_equal(
+             argGeneratingOp->getRegion(0).getArguments(), blockArgs.getVars()))
+      mapper.map(var, arg);
+  }
+
   // Bind the entry block arguments of parent wrappers to the corresponding
   // symbols.
-  for (auto [argGeneratingOp, args] : wrapperArgs)
-    bindEntryBlockArgs(converter, argGeneratingOp, args);
+  for (auto [argGeneratingOp, blockArgs] : wrapperArgs)
+    bindEntryBlockArgs(converter, argGeneratingOp, blockArgs);
 
   // The argument is not currently in memory, so make a temporary for the
   // argument, and store it there, then bind that location to the argument.
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index b47b2cfb29eb97..dca29f55661b0c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -536,6 +536,7 @@ convertIgnoredWrapper(omp::LoopWrapperInterface &opInst,
         forwardArgs(blockArgIface.getPrivateBlockArgs(), op.getPrivateVars());
         forwardArgs(blockArgIface.getReductionBlockArgs(),
                     op.getReductionVars());
+        op.emitWarning() << "simd information on composite construct discarded";
         return success();
       })
       .Default([&](Operation *op) {
diff --git a/mlir/test/Target/LLVMIR/openmp-todo.mlir b/mlir/test/Target/LLVMIR/openmp-todo.mlir
index 5a7d8193bdeca1..3c9bd9031c3e85 100644
--- a/mlir/test/Target/LLVMIR/openmp-todo.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-todo.mlir
@@ -52,6 +52,20 @@ llvm.func @cancellation_point() {
 
 // -----
 
+llvm.func @do_simd(%lb : i32, %ub : i32, %step : i32) {
+  omp.wsloop {
+    // expected-warning at below {{simd information on composite construct discarded}}
+    omp.simd {
+      omp.loop_nest (%iv) : i32 = (%lb) to (%ub) step (%step) {
+        omp.yield
+      }
+    } {omp.composite}
+  } {omp.composite}
+  llvm.return
+}
+
+// -----
+
 llvm.func @distribute(%lb : i32, %ub : i32, %step : i32) {
   // expected-error at below {{unsupported OpenMP operation: omp.distribute}}
   // expected-error at below {{LLVM Translation failed for operation: omp.distribute}}