Preserve argument locations when cloning a region.

Reviewed By: mehdi_amini

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

Author: dominikgrewe

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)

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
+}