[mlir] Make StripDebugInfo strip out block arguments locs

[ThomasRaoux]

Summary

Now that block arguments can have debug info we also need to strip those.

mlir/lib/Transforms/StripDebugInfo.cpp

 
 namespace {
 struct StripDebugInfo : public StripDebugInfoBase<StripDebugInfo> {
   void runOnOperation() override;
 };
 } // end anonymous namespace
 
 void StripDebugInfo::runOnOperation() {
-  // Strip the debug info from all operations.
   auto unknownLoc = UnknownLoc::get(&getContext());
-  getOperation()->walk([&](Operation *op) { op->setLoc(unknownLoc); });
+
+  // Strip the debug info from all operations.
+  getOperation()->walk([&](Operation *op) {
+    op->setLoc(unknownLoc);
+    // Strip block arguments debug info.
+    for (Region &region : op->getRegions()) {
+      for (Block &block : region.getBlocks()) {
+        for (BlockArgument &arg : block.getArguments()) {
+          arg.setLoc(unknownLoc);
+        }
+      }
+    }
+  });
 }
 
 /// Creates a pass to strip debug information from a function.
 std::unique_ptr<Pass> mlir::createStripDebugInfoPass() {
   return std::make_unique<StripDebugInfo>();
 }

mlir/test/Transforms/strip-debuginfo.mlir

func @inline_notation() -> i32 {
   affine.for %i0 = 0 to 8 {
   } loc(fused["foo", "mysource.cc":10:8])
 
   // CHECK: } loc(unknown)
   %2 = constant 4 : index
   affine.if #set0(%2) {
   } loc(fused<"myPass">["foo", "foo2"])
 
+  "foo.region"() ({
+  // CHECK: ^bb0(%{{.*}}: i32 loc(unknown), %{{.*}}: i32 loc(unknown)):
+ ^bb0(%a0: i32 loc("argloc"), %z: i32 loc("argloc2")):
+    %s = addi %a0, %a0 : i32
+    "foo.yield"(%s) : (i32) -> ()
+  }) : () -> ()
+
   // CHECK: return %0 : i32 loc(unknown)
   return %1 : i32 loc("bar")
 }