[Mlir-commits] [mlir] [WIP][flang][OpenMP] Translate OpenMP scopes when compiling for target device (PR #130078)

[Kareem Ergawy]

https://github.com/ergawy created https://github.com/llvm/llvm-project/pull/130078

**Todo**:
- [ ] Add tests
- [ ] Handle all ops with block arguments.

With these changes, we can handle inputs like:
```fortran
subroutine test_sections_target_with_allocatable
implicit none
integer,parameter :: N=100, rt=8
real(kind=rt), allocatable, dimension(:) :: a
integer :: i

allocate(a(1:N))

!$omp parallel
  !$omp parallel
    !$omp sections
      !$omp section
        !$omp target
          a(1) = 1.0_rt
        !$omp end target 
    !$omp end sections
  !$omp end parallel
!$omp end parallel


!$omp parallel
  !$omp parallel
    !$omp do
    do i=1,10
      !$omp target
        a(i) = 1.0_rt
      !$omp end target 
    end do
    !$omp end do
  !$omp end parallel
!$omp end parallel

end subroutine
```


>From 43a03621b20b7a79ba281755acdc7525c7cf6d9f Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 6 Mar 2025 03:16:59 -0600
Subject: [PATCH] [WIP][flang][OpenMP] Translate OpenMP scopes when compiling
 for target device

---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 62 ++++++++++++++++---
 1 file changed, 53 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 32c7c501d03c3..6af8856618f27 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -542,6 +542,13 @@ static llvm::omp::ProcBindKind getProcBindKind(omp::ClauseProcBindKind kind) {
   llvm_unreachable("Unknown ClauseProcBindKind kind");
 }
 
+static void forwardArgs(LLVM::ModuleTranslation &moduleTranslation,
+                        llvm::ArrayRef<BlockArgument> blockArgs,
+                        OperandRange operands) {
+  for (auto [arg, var] : llvm::zip_equal(blockArgs, operands))
+    moduleTranslation.mapValue(arg, moduleTranslation.lookupValue(var));
+}
+
 /// Helper function to map block arguments defined by ignored loop wrappers to
 /// LLVM values and prevent any uses of those from triggering null pointer
 /// dereferences.
@@ -554,18 +561,12 @@ convertIgnoredWrapper(omp::LoopWrapperInterface opInst,
   // Map block arguments directly to the LLVM value associated to the
   // corresponding operand. This is semantically equivalent to this wrapper not
   // being present.
-  auto forwardArgs =
-      [&moduleTranslation](llvm::ArrayRef<BlockArgument> blockArgs,
-                           OperandRange operands) {
-        for (auto [arg, var] : llvm::zip_equal(blockArgs, operands))
-          moduleTranslation.mapValue(arg, moduleTranslation.lookupValue(var));
-      };
-
   return llvm::TypeSwitch<Operation *, LogicalResult>(opInst)
       .Case([&](omp::SimdOp op) {
         auto blockArgIface = cast<omp::BlockArgOpenMPOpInterface>(*op);
-        forwardArgs(blockArgIface.getPrivateBlockArgs(), op.getPrivateVars());
-        forwardArgs(blockArgIface.getReductionBlockArgs(),
+        forwardArgs(moduleTranslation, blockArgIface.getPrivateBlockArgs(),
+                    op.getPrivateVars());
+        forwardArgs(moduleTranslation, blockArgIface.getReductionBlockArgs(),
                     op.getReductionVars());
         op.emitWarning() << "simd information on composite construct discarded";
         return success();
@@ -5255,6 +5256,49 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
               return WalkResult::interrupt();
             return WalkResult::skip();
           }
+
+          // Not-target ops might nest target-related ops, therefore, we
+          // translate them as non-OpenMP scopes. Translating them is needed by
+          // nested target-related ops since they might LLVM values defined in
+          // their parent non-target ops.
+          if (isa<omp::OpenMPDialect>(oper->getDialect()) &&
+              oper->getParentOfType<LLVM::LLVMFuncOp>() &&
+              !oper->getRegions().empty()) {
+
+            // TODO Handle other ops with entry block args.
+            if (auto wsloop = dyn_cast<omp::WsloopOp>(oper)) {
+              auto blockArgIface = cast<omp::BlockArgOpenMPOpInterface>(oper);
+              if (blockArgIface) {
+                forwardArgs(moduleTranslation,
+                            blockArgIface.getPrivateBlockArgs(),
+                            wsloop.getPrivateVars());
+                forwardArgs(moduleTranslation,
+                            blockArgIface.getReductionBlockArgs(),
+                            wsloop.getReductionVars());
+              }
+            }
+
+            if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
+              for (auto iv : loopNest.getIVs()) {
+                // Create fake allocas just to maintain IR validity.
+                moduleTranslation.mapValue(
+                    iv, builder.CreateAlloca(
+                            moduleTranslation.convertType(iv.getType())));
+              }
+            }
+
+            for (Region &region : oper->getRegions()) {
+              auto result = convertOmpOpRegions(
+                  region, oper->getName().getStringRef().str() + ".fake.region",
+                  builder, moduleTranslation);
+              if (result.takeError())
+                return WalkResult::interrupt();
+              builder.SetInsertPoint(result.get(), result.get()->end());
+            }
+
+            return WalkResult::skip();
+          }
+
           return WalkResult::advance();
         }).wasInterrupted();
   return failure(interrupted);