[MLIR] Infer locations for block argument conversion

[nand]

Summary

To enable signature conversions to be used in CIRCT, locations should no longer be dropped from block arguments.

[nand]

Fixed conversion for inserted arguments

[nand]

remove unnecessary loop

[springerm]

mlir/lib/Transforms/Utils/DialectConversion.cpp

730

< is more idiomatic

[Mogball]

Thank you! I just noticed the conversion was spitting out unknown locations today was about to try to fix it!

[nand]

apply nit

mlir/lib/Transforms/Utils/DialectConversion.cpp

Block *ArgConverter::applySignatureConversion(
   if (origArgCount == 0 && convertedTypes.empty())
     return block;
 
   // Split the block at the beginning to get a new block to use for the updated
   // signature.
   Block *newBlock = block->splitBlock(block->begin());
   block->replaceAllUsesWith(newBlock);
 
-  // FIXME: We should map the new arguments to proper locations.
+  // Map all new arguments to the location of the argument they originate from.
   SmallVector<Location> newLocs(convertedTypes.size(),
                                 rewriter.getUnknownLoc());
+  for (unsigned i = 0; i < origArgCount; ++i) {
+    auto inputMap = signatureConversion.getInputMapping(i);
+    if (!inputMap || inputMap->replacementValue)
+      continue;
+    Location origLoc = block->getArgument(i).getLoc();
+    for (unsigned j = 0; j < inputMap->size; ++j)
+      newLocs[inputMap->inputNo + j] = origLoc;
+  }
+
   SmallVector<Value, 4> newArgRange(
       newBlock->addArguments(convertedTypes, newLocs));
   ArrayRef<Value> newArgs(newArgRange);
 
   // Remap each of the original arguments as determined by the signature
   // conversion.
   ConvertedBlockInfo info(block, converter);
   info.argInfo.resize(origArgCount);