[MLIR][Transforms] Update block arg locations during inlining

[Dinistro]

This commit changes the inlining to also update the locations of block arguments. Not updating these locations leads to LLVM IR verification issues when exporting converted block arguments to phi nodes. This lack of location update was not visible due to ignoring the argument locations until recently.
Relevant change: #105534

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Christian Ulmann (Dinistro)

Changes

This commit changes the inlining to also update the locations of block arguments. Not updating these locations leads to LLVM IR verification issues when exporting converted block arguments to phi nodes. This lack of location update was not visible due to ignoring the argument locations until recently.
Relevant change: #105534

---

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

2 Files Affected:

• (modified) mlir/lib/Transforms/Utils/InliningUtils.cpp (+23-9)
• (modified) mlir/test/Transforms/inlining.mlir (+2-2)

diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index ba146920fae2e9..da49f79d069e76 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -25,22 +25,36 @@
 
 using namespace mlir;
 
-/// Remap locations from the inlined blocks with CallSiteLoc locations with the
-/// provided caller location.
+/// Remap all locations from the inlined blocks with CallSiteLoc locations with
+/// the provided caller location.
 static void
 remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
                       Location callerLoc) {
   DenseMap<Location, Location> mappedLocations;
-  auto remapOpLoc = [&](Operation *op) {
-    auto it = mappedLocations.find(op->getLoc());
+  auto remapLoc = [&](Location loc) {
+    auto it = mappedLocations.find(loc);
     if (it == mappedLocations.end()) {
-      auto newLoc = CallSiteLoc::get(op->getLoc(), callerLoc);
-      it = mappedLocations.try_emplace(op->getLoc(), newLoc).first;
+      auto newLoc = CallSiteLoc::get(loc, callerLoc);
+      it = mappedLocations.try_emplace(loc, newLoc).first;
     }
-    op->setLoc(it->second);
+    return it->second;
   };
-  for (auto &block : inlinedBlocks)
-    block.walk(remapOpLoc);
+
+  AttrTypeReplacer attrReplacer;
+  attrReplacer.addReplacement(
+      [&](LocationAttr loc) -> std::pair<LocationAttr, WalkResult> {
+        return {remapLoc(loc), WalkResult::skip()};
+      });
+
+  for (Block &block : inlinedBlocks) {
+    for (BlockArgument &arg : block.getArguments())
+      if (LocationAttr newLoc = remapLoc(arg.getLoc()))
+        arg.setLoc(newLoc);
+
+    for (Operation &op : block)
+      attrReplacer.recursivelyReplaceElementsIn(&op, /*replaceAttrs=*/false,
+                                                /*replaceLocs=*/true);
+  }
 }
 
 static void remapInlinedOperands(iterator_range<Region::iterator> inlinedBlocks,
diff --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index 2a08e625ba79e2..79a2936b104fa1 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -215,9 +215,9 @@ func.func @func_with_block_args_location(%arg0 : i32) {
 
 // INLINE-LOC-LABEL: func @func_with_block_args_location_callee1
 // INLINE-LOC: cf.br
-// INLINE-LOC: ^bb{{[0-9]+}}(%{{.*}}: i32 loc("foo")
+// INLINE-LOC: ^bb{{[0-9]+}}(%{{.*}}: i32 loc(callsite("foo" at "bar"))
 func.func @func_with_block_args_location_callee1(%arg0 : i32) {
-  call @func_with_block_args_location(%arg0) : (i32) -> ()
+  call @func_with_block_args_location(%arg0) : (i32) -> () loc("bar")
   return
 }
 

[gysit]

LGTM modulo nits.

[Dinistro]

FYI: @abidh , in case you observed some issues with this.