[Mlir-commits] [mlir] 5f782d2 - Preserve argument locations when cloning a region.

[Mehdi Amini]

Author: Dominik Grewe
Date: 2022-01-16T21:17:23Z
New Revision: 5f782d25a742302d25ef3c8b84b54f7483c2deb9

URL: https://github.com/llvm/llvm-project/commit/5f782d25a742302d25ef3c8b84b54f7483c2deb9
DIFF: https://github.com/llvm/llvm-project/commit/5f782d25a742302d25ef3c8b84b54f7483c2deb9.diff

LOG: Preserve argument locations when cloning a region.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D117403

Added: 
    

Modified: 
    mlir/lib/IR/Region.cpp
    mlir/test/Transforms/inlining.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index 161470de8554a..d1783c482d47f 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -92,7 +92,7 @@ void Region::cloneInto(Region *dest, Region::iterator destPos,
     // argument to the cloned block.
     for (auto arg : block.getArguments())
       if (!mapper.contains(arg))
-        mapper.map(arg, newBlock->addArgument(arg.getType()));
+        mapper.map(arg, newBlock->addArgument(arg.getType(), arg.getLoc()));
 
     // Clone and remap the operations within this block.
     for (auto &op : block)

diff  --git a/mlir/test/Transforms/inlining.mlir b/mlir/test/Transforms/inlining.mlir
index a81050f16ecb0..4c894eb0a9e80 100644
--- a/mlir/test/Transforms/inlining.mlir
+++ b/mlir/test/Transforms/inlining.mlir
@@ -202,3 +202,26 @@ func @inline_gpu_ops() -> memref<1024xf32> {
   %m = call @gpu_alloc() : () -> memref<1024xf32>
   return %m : memref<1024xf32>
 }
+
+// Test block arguments location propagation.
+// Use two call-sites to force cloning.
+func @func_with_block_args_location(%arg0 : i32) {
+  br ^bb1(%arg0 : i32)
+^bb1(%x : i32 loc("foo")):
+  "test.foo" (%x) : (i32) -> () loc("bar")
+  return
+}
+
+// INLINE-LOC-LABEL: func @func_with_block_args_location_callee1
+// INLINE-LOC: br
+// INLINE-LOC: ^bb{{[0-9]+}}(%{{.*}}: i32 loc("foo")
+func @func_with_block_args_location_callee1(%arg0 : i32) {
+  call @func_with_block_args_location(%arg0) : (i32) -> ()
+  return
+}
+
+// CHECK-LABEL: func @func_with_block_args_location_callee2
+func @func_with_block_args_location_callee2(%arg0 : i32) {
+  call @func_with_block_args_location(%arg0) : (i32) -> ()
+  return
+}