[Mlir-commits] [mlir] [MLIR][Transforms] Update block arg locations during inlining (PR #106064)

[Christian Ulmann]

https://github.com/Dinistro updated https://github.com/llvm/llvm-project/pull/106064

>From 55b386e8564f266c833b7c828a684a0b2fc138e6 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Mon, 26 Aug 2024 12:15:43 +0000
Subject: [PATCH 1/3] [MLIR][Transforms] Update block arg locations during
 inlining

This commit changes the inlining to also update the locations of block
arguments.
---
 mlir/lib/Transforms/Utils/InliningUtils.cpp | 32 +++++++++++++++------
 mlir/test/Transforms/inlining.mlir          |  4 +--
 2 files changed, 25 insertions(+), 11 deletions(-)

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
 }
 

>From cc9420c0c4f2a7f6f33d9e3f8a20a5d35dd73653 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Mon, 26 Aug 2024 12:34:30 +0000
Subject: [PATCH 2/3] address review comments

---
 mlir/lib/Transforms/Utils/InliningUtils.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index da49f79d069e76..0c831969339165 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -25,8 +25,8 @@
 
 using namespace mlir;
 
-/// Remap all locations from the inlined blocks with CallSiteLoc locations with
-/// the provided caller location.
+/// Remap all locations reachable from the inlined blocks with CallSiteLoc
+/// locations with the provided caller location.
 static void
 remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
                       Location callerLoc) {

>From 264012570541c4f563e5dde25fe5b46b288def30 Mon Sep 17 00:00:00 2001
From: Christian Ulmann <christian.ulmann at nextsilicon.com>
Date: Mon, 26 Aug 2024 12:45:22 +0000
Subject: [PATCH 3/3] address second review comment

---
 mlir/lib/Transforms/Utils/InliningUtils.cpp | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 0c831969339165..0db097d14cd3c7 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -30,12 +30,13 @@ using namespace mlir;
 static void
 remapInlinedLocations(iterator_range<Region::iterator> inlinedBlocks,
                       Location callerLoc) {
-  DenseMap<Location, Location> mappedLocations;
+  DenseMap<Location, LocationAttr> mappedLocations;
   auto remapLoc = [&](Location loc) {
-    auto it = mappedLocations.find(loc);
-    if (it == mappedLocations.end()) {
+    auto [it, inserted] = mappedLocations.try_emplace(loc);
+    // Only query the attribute uniquer once per callsite attribute.
+    if (inserted) {
       auto newLoc = CallSiteLoc::get(loc, callerLoc);
-      it = mappedLocations.try_emplace(loc, newLoc).first;
+      it->getSecond() = newLoc;
     }
     return it->second;
   };