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

Fixes a bug introduced by
#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 |

Author: ergawy

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();
 

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:       }
 }