Author: Nirvedh
Date: 2022-03-15T18:42:43Z
New Revision: b8d211fc317ffefaed1d65b226cda6c464f7d216
URL:
https://github.com/llvm/llvm-project/commit/b8d211fc317ffefaed1d65b226cda6c464f7d216
DIFF:
https://github.com/llvm/llvm-project/commit/b8d211fc317ffefaed1d65b226cda6c464f7d216.diff
LOG: [MLIR][Linalg] Canonicalization patterns for linalg.generic.
Fold linalg.fill into linalg.generic.
Remove dead arguments used in linalg.generic.
Reviewed By: ThomasRaoux
Differential Revision: https://reviews.llvm.org/D121535
Added:
Modified:
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
mlir/test/Dialect/Linalg/canonicalize.mlir
mlir/test/Dialect/Linalg/fusion-indexed.mlir
Removed:
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 331a8b91bd330..2c9b1e0c53553 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -105,6 +105,33 @@ static LogicalResult foldMemRefCast(Operation *op) {
return success(folded);
}
+/// Helper function to find if there is atleast one dimension in an AffineMap
+/// testMap that is contained in `testMapLocation` of `maps` but not in any
+/// other locations
+static bool hasaUniqueDim(ArrayRef maps, unsigned testMapLocation) {
+ AffineMap testMap = maps[testMapLocation];
+ llvm::SmallDenseSet dimsToCheck;
+ for (auto result : testMap.getResults()) {
+auto expr = result.dyn_cast();
+if (expr != nullptr)
+ dimsToCheck.insert(expr.getPosition());
+ }
+ for (auto It : llvm::enumerate(maps)) {
+if (It.index() == testMapLocation)
+ continue;
+auto map = It.value();
+for (auto result : map.getResults()) {
+ auto expr = result.dyn_cast();
+ if (expr != nullptr) {
+dimsToCheck.erase(expr.getPosition());
+ }
+ if (dimsToCheck.empty())
+return false;
+}
+ }
+ return true;
+}
+
//===--===//
// Region builder helper.
// TODO: Move this to a utility library.
@@ -826,11 +853,95 @@ struct EraseIdentityGenericOp : public
OpRewritePattern {
return success();
}
};
+
+/// Drop dead args of a linalg generic op.
+/// An arg is dead if it has zero uses in the op region.
+struct DeadArgsGenericOpInputs : public OpRewritePattern {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(GenericOp genericOp,
+PatternRewriter &rewriter) const override {
+SmallVector oldIndexingMaps = genericOp.getIndexingMaps();
+// Maps must be projected permutations.
+if (llvm::any_of(genericOp.getIndexingMaps(), [](AffineMap map) {
+ return !map.isProjectedPermutation();
+}))
+ return failure();
+Block &payload = genericOp.region().front();
+SmallVector newInputOperands;
+SmallVector newIndexingMaps;
+bool deadArgFound = false;
+int inputSize = genericOp.getInputOperands().size();
+for (int i = inputSize - 1; i >= 0; i--) {
+ OpOperand *opOperand = genericOp.getInputOperand(i);
+ // Iterate in reverse, so that we erase later args first, preventing the
+ // argument list from shifting unexpectedly and invalidating all our
+ // indices.
+ if (payload.getArgument(i).use_empty() &&
+ !hasaUniqueDim(oldIndexingMaps, i)) {
+payload.eraseArgument(i);
+deadArgFound = true;
+// remove this indexing map out of consideration for hasaUniqueDim
check
+oldIndexingMaps.erase(oldIndexingMaps.begin() + i);
+ } else {
+newInputOperands.insert(newInputOperands.begin(), opOperand->get());
+newIndexingMaps.insert(newIndexingMaps.begin(),
+ genericOp.getTiedIndexingMap(opOperand));
+ }
+}
+// Bail out if there are no dead args.
+if (!deadArgFound)
+ return failure();
+for (OpOperand *opOperand : genericOp.getOutputOperands())
+ newIndexingMaps.push_back(genericOp.getTiedIndexingMap(opOperand));
+SmallVector outputOperands = genericOp.getOutputOperands();
+
+auto newOp = rewriter.create(
+genericOp.getLoc(), genericOp->getResultTypes(), newInputOperands,
+outputOperands, rewriter.getAffineMapArrayAttr(newIndexingMaps),
+genericOp.iterator_types(), genericOp.docAttr(),
+genericOp.library_callAttr());
+// Copy over unknown attributes. They might be load bearing for some flow.
+ArrayRef odsAttrs = genericOp.getAttributeNames();
+for (NamedAttribute kv : genericOp->getAttrs()) {
+ if (!llvm::is_contained(odsAttrs, kv.getName().getValue())) {
+newOp->setAttr(kv.getName(), kv.getValue());
+ }
+}
+rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
+newOp.region().begin());
+rewriter.replaceOp(genericOp, newOp->ge
