[Mlir-commits] [mlir] b983783 - [MLIR][linalg] Preserve location during elementwise fusion

[Geoffrey Martin-Noble]

Author: Geoffrey Martin-Noble
Date: 2021-10-05T09:43:53-07:00
New Revision: b983783d2e58742b03f7e9d495c50eed4001ec66

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

LOG: [MLIR][linalg] Preserve location during elementwise fusion

This otherwise loses a lot of debugging info and results in a painful
debugging experience.

Reviewed By: mravishankar, stellaraccident

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index befff7b4795e..5f721e2098ee 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -1214,13 +1214,17 @@ class FoldConstants : public OpRewritePattern<GenericOp> {
       // values at the constant index dropped.
       SmallVector<AffineMap> fusedIndexMaps;
       SmallVector<Value> fusedOperands;
+      SmallVector<Location> fusedLocs{genericOp.getLoc()};
       fusedIndexMaps.reserve(genericOp.getNumInputsAndOutputs());
       fusedOperands.reserve(genericOp.getNumInputs());
+      fusedLocs.reserve(fusedLocs.size() + genericOp.getNumInputs());
       for (OpOperand *inputOperand : genericOp.getInputOperands()) {
         if (inputOperand == opOperand)
           continue;
+        Value inputValue = inputOperand->get();
         fusedIndexMaps.push_back(genericOp.getTiedIndexingMap(inputOperand));
-        fusedOperands.push_back(inputOperand->get());
+        fusedOperands.push_back(inputValue);
+        fusedLocs.push_back(inputValue.getLoc());
       }
       for (OpOperand *outputOperand : genericOp.getOutputOperands())
         fusedIndexMaps.push_back(genericOp.getTiedIndexingMap(outputOperand));
@@ -1237,7 +1241,7 @@ class FoldConstants : public OpRewritePattern<GenericOp> {
 
       SmallVector<Value> outputOperands = genericOp.getOutputOperands();
       auto fusedOp = rewriter.create<GenericOp>(
-          rewriter.getUnknownLoc(), genericOp->getResultTypes(),
+          rewriter.getFusedLoc(fusedLocs), genericOp->getResultTypes(),
           /*inputs=*/fusedOperands,
           /*outputs=*/outputOperands,
           rewriter.getAffineMapArrayAttr(fusedIndexMaps),