[Mlir-commits] [mlir] 0c46a91 - [MLIR] Infer locations for block argument conversion

[Nandor Licker]

Author: Nandor Licker
Date: 2023-08-18T19:44:49+03:00
New Revision: 0c46a9189c8bb559926c3391097b5c980d006f8e

URL: https://github.com/llvm/llvm-project/commit/0c46a9189c8bb559926c3391097b5c980d006f8e
DIFF: https://github.com/llvm/llvm-project/commit/0c46a9189c8bb559926c3391097b5c980d006f8e.diff

LOG: [MLIR] Infer locations for block argument conversion

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

Reviewed By: Mogball, springerm

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

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/DialectConversion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 78d7b47558b553..a345f8d5b5be49 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -724,9 +724,18 @@ Block *ArgConverter::applySignatureConversion(
   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);