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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Mar 20 01:45:00 PDT 2025


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-openmp

Author: Kareem Ergawy (ergawy)

<details>
<summary>Changes</summary>

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         |


---
Full diff: https://github.com/llvm/llvm-project/pull/132165.diff


2 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp (+39-1) 
- (modified) mlir/test/Target/LLVMIR/openmp-target-nesting-in-host-ops.mlir (+21) 


``````````diff
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 722107c7ec6d7..cfc4e98ad3341 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -5323,6 +5323,43 @@ convertTargetOpsInNest(Operation *op, llvm::IRBuilderBase &builder,
             if (auto blockArgsIface =
                     dyn_cast<omp::BlockArgOpenMPOpInterface>(oper))
               forwardArgs(moduleTranslation, blockArgsIface);
+            else {
+              // 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.
+              //
+              // clang-format off
+              // =============================================================================
+              // | 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         |
+              // =============================================================================
+              // clang-format on
+              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() &&
@@ -5340,9 +5377,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:       }
 }

``````````

</details>


https://github.com/llvm/llvm-project/pull/132165


More information about the Mlir-commits mailing list