[Mlir-commits] [mlir] f23a6ef - [flang][OpenMP] Process `omp.atomic.update` while translating scopes for target device (#132165)

[llvmlistbot at llvm.org]

Author: Kareem Ergawy
Date: 2025-03-20T16:21:09-05:00
New Revision: f23a6ef54c104b357e140e7782fd66248d9382f1

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

LOG: [flang][OpenMP] Process `omp.atomic.update` while translating scopes for target device (#132165)

Fixes a bug introduced by
https://github.com/llvm/llvm-project/pull/130078.

For non-BlockArgOpenMPOpInterface ops, we also want to map their entry
block arguments to their operands, if any. For the current support in
the OpenMP dialect, the table below lists all ops that have arguments
(SSA operands and/or attributes) and not target-related. Of all these
ops, we need to only process `omp.atomic.update` since it is the only op
that has SSA operands & an attached region. Therefore, the region's
entry block arguments must be mapped to the op's operands in case they
are referenced inside the region.


| op | operands? | region(s)? | parent is func? | processed? |

|--------------|-------------|------------|------------------|-------------|
| atomic.read | yes | no | yes | no |
| atomic.write | yes | no | yes | no |
| atomic.update | yes | yes | yes | yes |
| critical | no | no | yes | no |
| declare_mapper | no | yes | no | no |
| declare_reduction | no | yes | no | no |
| flush | yes | no | yes | no |
| private | no | yes | yes | no |
| threadprivate | yes | no | yes | no |
| yield | yes | no | yes | no |

Added: 
    

Modified: 
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
    mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index a745726bd51d5..d41489921bd13 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5320,6 +5320,20 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
             if (auto blockArgsIface =
                     dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
               forwardArgs(moduleTranslation, blockArgsIface);
+            else {
+              // Here we map entry block arguments of
+              // non-BlockArgOpenMPOpInterface ops if they can be encountered
+              // inside of a function and they define any of these arguments.
+              if (isa<mlir::omp::AtomicUpdateOp>(oper))
+                for (auto [operand, arg] :
+                     llvm::zip_equal(oper->getOperands(),
+                                     oper->getRegion(0).getArguments())) {
+                  moduleTranslation.mapValue(
+                      arg, builder.CreateLoad(
+                               moduleTranslation.convertType(arg.getType()),
+                               moduleTranslation.lookupValue(operand)));
+                }
+            }
 
             if (auto loopNest = dyn_cast<omp::LoopNestOp>(oper)) {
               assert(builder.GetInsertBlock() &&
@@ -5337,9 +5351,10 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
               // translation of the OpenMP construct being converted (e.g. no
               // OpenMP runtime calls will be generated). We just need this to
               // prepare the kernel invocation args.
+              SmallVector<llvm::PHINode *> phis;
               auto result = convertOmpOpRegions(
                   region, oper->getName().getStringRef().str() + ".fake.region",
-                  builder, moduleTranslation);
+                  builder, moduleTranslation, &phis);
               if (failed(handleError(result, *oper)))
                 return WalkResult::interrupt();
 

diff  --git a/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir b/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
index 71a4c29eaf0aa..c618b68d52aaf 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir
@@ -97,6 +97,20 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
     llvm.return
   }
 
+  llvm.func @test_target_and_atomic_update(%x: !llvm.ptr, %expr : i32) {
+    omp.target {
+      omp.terminator
+    }
+
+    omp.atomic.update %x : !llvm.ptr {
+    ^bb0(%xval: i32):
+      %newval = llvm.add %xval, %expr : i32
+      omp.yield(%newval : i32)
+    }
+
+    llvm.return
+  }
+
 // CHECK-LABEL: define void @test_nested_target_in_parallel_with_private({{.*}}) {
 // CHECK:        br label %omp.parallel.fake.region
 // CHECK:       omp.parallel.fake.region:
@@ -132,4 +146,11 @@ module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true,
 // CHECK:         call void @__kmpc_target_deinit()
 // CHECK:         ret void
 // CHECK:       }
+
+// CHECK-LABEL: define {{.*}} amdgpu_kernel void @__omp_offloading_{{.*}}_test_target_and_atomic_update_{{.*}} {
+// CHECK:         call i32 @__kmpc_target_init
+// CHECK:       user_code.entry:
+// CHECK:         call void @__kmpc_target_deinit()
+// CHECK:         ret void
+// CHECK:       }
 }