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

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Delaram Talaashrafi (delaram-talaashrafi)

<details>
<summary>Changes</summary>

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`.

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


1 Files Affected:

- (modified) flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp (+47-11) 


``````````diff
diff --git a/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp b/flang/lib/Optimizer/Transforms/ExternalNameConversion.cpp
index 3d84eaa4c1595..2fcff87fdc393 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);
   });
 }

``````````

</details>


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