Preserve argument locations when cloning a region.

[dominikgrewe]

Summary

Preserve argument locations when cloning a region.

[jpienaar]

Good catch! Could you add a simple test for this?

[mehdi_amini]

Thanks!

Can you add a test?
This should do it:

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

in mlir/test/Transforms/inlining.mlir

[dominikgrewe]

Added test.

[dominikgrewe]

Thanks for the neat test! I've added it to the CL.

[mehdi_amini]

LG, I assume I need to land this for you!

[dominikgrewe]

Thanks Mehdi!

mlir/lib/IR/Region.cpp

for (Block &block : *this) {
     Block *newBlock = new Block();
     mapper.map(&block, newBlock);
 
     // Clone the block arguments. The user might be deleting arguments to the
     // block by specifying them in the mapper. If so, we don't add the
     // 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)
       newBlock->push_back(op.clone(mapper));
 
     dest->getBlocks().insert(destPos, newBlock);
   }
 

mlir/test/Transforms/inlining.mlir

 }
 
 // CHECK-LABEL: func @inline_gpu_ops
 func @inline_gpu_ops() -> memref<1024xf32> {
   // CHECK-NEXT: gpu.alloc
   %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
+}