[flang] Extend symbol update to ArrayAttrext in `external-name-interop` (#150061)

In the `external-name-interop` pass, when a symbol is changed, all the
uses of the renamed symbols should also be updated. The update was only
applied to the `SymbolRefAttr` type. With this change, the update will
be applied to `ArrayAttr` containing elements of type `SymbolRefAttr`.

---------

Co-authored-by: Delaram Talaashrafi <dtalaashrafi@nvidia.com>

2 files changed, 89 insertions, 11 deletions

diff --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
index 3d84eaa..2fcff87 100644
--- a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
+++ b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
@@ -41,6 +41,23 @@ mangleExternalName(const std::pair<fir::NameUniquer::NameKind,
                                                   appendUnderscore);
 }
 
+/// Process a symbol reference and return the updated symbol reference if
+/// needed.
+std::optional<mlir::SymbolRefAttr>
+processSymbolRef(mlir::SymbolRefAttr symRef, mlir::Operation *nestedOp,
+                 const llvm::DenseMap<mlir::StringAttr, mlir::FlatSymbolRefAttr>
+                     &remappings) {
+  if (auto remap = remappings.find(symRef.getLeafReference());
+      remap != remappings.end()) {
+    mlir::SymbolRefAttr symAttr = mlir::FlatSymbolRefAttr(remap->second);
+    if (mlir::isa<mlir::gpu::LaunchFuncOp>(nestedOp))
+      symAttr = mlir::SymbolRefAttr::get(
+          symRef.getRootReference(), {mlir::FlatSymbolRefAttr(remap->second)});
+    return symAttr;
+  }
+  return std::nullopt;
+}
+
 namespace {
 
 class ExternalNameConversionPass
@@ -97,21 +114,40 @@ void ExternalNameConversionPass::runOnOperation() {
 
   // Update all uses of the functions and globals that have been renamed.
   op.walk([&remappings](mlir::Operation *nestedOp) {
-    llvm::SmallVector<std::pair<mlir::StringAttr, mlir::SymbolRefAttr>> updates;
+    llvm::SmallVector<std::pair<mlir::StringAttr, mlir::SymbolRefAttr>>
+        symRefUpdates;
+    llvm::SmallVector<std::pair<mlir::StringAttr, mlir::ArrayAttr>>
+        arrayUpdates;
     for (const mlir::NamedAttribute &attr : nestedOp->getAttrDictionary())
       if (auto symRef = llvm::dyn_cast<mlir::SymbolRefAttr>(attr.getValue())) {
-        if (auto remap = remappings.find(symRef.getLeafReference());
-            remap != remappings.end()) {
-          mlir::SymbolRefAttr symAttr = mlir::FlatSymbolRefAttr(remap->second);
-          if (mlir::isa<mlir::gpu::LaunchFuncOp>(nestedOp))
-            symAttr = mlir::SymbolRefAttr::get(
-                symRef.getRootReference(),
-                {mlir::FlatSymbolRefAttr(remap->second)});
-          updates.emplace_back(std::pair<mlir::StringAttr, mlir::SymbolRefAttr>{
-              attr.getName(), symAttr});
+        if (auto newSymRef = processSymbolRef(symRef, nestedOp, remappings))
+          symRefUpdates.emplace_back(
+              std::pair<mlir::StringAttr, mlir::SymbolRefAttr>{attr.getName(),
+                                                               *newSymRef});
+      } else if (auto arrayAttr =
+                     llvm::dyn_cast<mlir::ArrayAttr>(attr.getValue())) {
+        llvm::SmallVector<mlir::Attribute> symbolRefs;
+        for (auto element : arrayAttr) {
+          if (!element) {
+            symbolRefs.push_back(element);
+            continue;
+          }
+          auto symRef = llvm::dyn_cast<mlir::SymbolRefAttr>(element);
+          std::optional<mlir::SymbolRefAttr> updatedSymRef;
+          if (symRef)
+            updatedSymRef = processSymbolRef(symRef, nestedOp, remappings);
+          if (!symRef || !updatedSymRef)
+            symbolRefs.push_back(element);
+          else
+            symbolRefs.push_back(*updatedSymRef);
         }
+        arrayUpdates.push_back(std::make_pair(
+            attr.getName(),
+            mlir::ArrayAttr::get(nestedOp->getContext(), symbolRefs)));
       }
-    for (auto update : updates)
+    for (auto update : symRefUpdates)
+      nestedOp->setAttr(update.first, update.second);
+    for (auto update : arrayUpdates)
       nestedOp->setAttr(update.first, update.second);
   });
 }
diff --git a/flang/test/Transforms/external-name-interop-symref-array.fir b/flang/test/Transforms/external-name-interop-symref-array.fir
new file mode 100644
index 0000000..ce14f17
--- /dev/null
+++ b/flang/test/Transforms/external-name-interop-symref-array.fir
@@ -0,0 +1,42 @@
+// Test fir.do_concurrent.loop operation with array of symbol reference attributes
+// This test demonstrates operations that have ArrayAttr containing SymbolRefAttr elements
+
+// RUN: fir-opt %s --external-name-interop | fir-opt | FileCheck %s
+
+// Define reduction operations that will be referenced in the symbol array
+func.func @_QPadd_reduction_i32_init(%arg0: i32, %arg1: !fir.ref<i32>) {
+  %0 = arith.constant 0 : i32
+  fir.store %0 to %arg1 : !fir.ref<i32>
+  return
+}
+
+func.func @_QPadd_reduction_i32_combiner(%arg0: i32, %arg1: i32) -> i32 {
+  %0 = arith.addi %arg0, %arg1 : i32
+  return %0 : i32
+}
+
+// Define a local privatizer that will be referenced in local_syms
+func.func @_QPlocal_var_privatizer(%arg0: !fir.ref<i32>) -> !fir.ref<i32> {
+  return %arg0 : !fir.ref<i32>
+}
+
+// Test function demonstrating both local_syms and reduce_syms arrays
+func.func @_QPtest_symbol_arrays(%i_lb: index, %i_ub: index, %i_st: index) {
+  %local_var = fir.alloca i32
+  %sum = fir.alloca i32
+  
+  fir.do_concurrent {
+    %i = fir.alloca i32
+    fir.do_concurrent.loop (%i_iv) = (%i_lb) to (%i_ub) step (%i_st)
+      local(@_QPlocal_var_privatizer %local_var -> %local_arg : !fir.ref<i32>)
+      reduce(@_QPadd_reduction_i32_init #fir.reduce_attr<add> %sum -> %sum_arg : !fir.ref<i32>) {
+      %0 = fir.convert %i_iv : (index) -> i32
+      fir.store %0 to %i : !fir.ref<i32>
+    }
+  }
+  return
+}
+
+// CHECK: local(@local_var_privatizer_
+// CHECK: reduce(@add_reduction_i32_init_
+