[Mlir-commits] [mlir] c02d07f - [mlir][vector] Add pattern to drop unit dim from elementwise(a, b)) (#74817)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Dec 13 12:29:17 PST 2023
Author: Andrzej Warzyński
Date: 2023-12-13T20:29:12Z
New Revision: c02d07fdf007afc6b928cda0342751889cc2604b
URL: https://github.com/llvm/llvm-project/commit/c02d07fdf007afc6b928cda0342751889cc2604b
DIFF: https://github.com/llvm/llvm-project/commit/c02d07fdf007afc6b928cda0342751889cc2604b.diff
LOG: [mlir][vector] Add pattern to drop unit dim from elementwise(a, b)) (#74817)
For vectors with either leading or trailing unit dim, replaces:
elementwise(a, b)
with:
sc_a = shape_cast(a)
sc_b = shape_cast(b)
res = elementwise(sc_a, sc_b)
return shape_cast(res)
The newly inserted shape_cast Ops fold (before elementwise Op) and then
restore (after elementwise Op) the unit dim. Vectors `a` and `b` are
required to be rank > 1.
Example:
```mlir
%mul = arith.mulf %B_row, %A_row : vector<1x[4]xf32>
%cast = vector.shape_cast %mul : vector<1x[4]xf32> to vector<[4]xf32>
```
gets converted to:
```mlir
%B_row_sc = vector.shape_cast %B_row : vector<1x[4]xf32> to vector<[4]xf32>
%A_row_sc = vector.shape_cast %A_row : vector<1x[4]xf32> to vector<[4]xf32>
%mul = arith.mulf %B_row_sc, %A_row_sc : vector<[4]xf32>
%mul_sc = vector.shape_cast %mul : vector<[4]xf32> to vector<1x[4]xf32>
%cast = vector.shape_cast %mul_sc : vector<1x[4]xf32> to vector<[4]xf32>
```
In practice, the bottom 2 shape_cast(s) will be folded away.
Added:
Modified:
mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
index 08c08172d0531e..17173c01ab762a 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
@@ -294,6 +294,15 @@ void populateCastAwayVectorLeadingOneDimPatterns(RewritePatternSet &patterns,
void populateVectorTransferDropUnitDimsPatterns(RewritePatternSet &patterns,
PatternBenefit benefit = 1);
+/// Collect a set of patterns that use vector.shape_cast to help fold unit dims.
+///
+/// These patterns use vector.shape_cast to remove unit dims from e.g.
+/// arithmetic operations on Vectors. The newly inserted shape_casts will either
+/// cancel each other out or will be folded away when combined with other
+/// patterns.
+void populateDropUnitDimWithShapeCastPatterns(RewritePatternSet &patterns,
+ PatternBenefit benefit = 1);
+
/// Collect a set of patterns to flatten n-D vector transfers on contiguous
/// memref.
///
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index ed42e6508b4310..b761d1ed888973 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -922,4 +922,5 @@ void mlir::vector::populateFlattenVectorTransferPatterns(
FlattenContiguousRowMajorTransferWritePattern>(
patterns.getContext(), benefit);
populateShapeCastFoldingPatterns(patterns, benefit);
+ populateDropUnitDimWithShapeCastPatterns(patterns, benefit);
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 6e7fab293d3a1c..b468d874bcf306 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1446,6 +1446,90 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
}
};
+/// For vectors with either leading or trailing unit dim, replaces:
+/// elementwise(a, b)
+/// with:
+/// sc_a = shape_cast(a)
+/// sc_b = shape_cast(b)
+/// res = elementwise(sc_a, sc_b)
+/// return shape_cast(res)
+/// The newly inserted shape_cast Ops fold (before elementwise Op) and then
+/// restore (after elementwise Op) the unit dim. Vectors `a` and `b` are
+/// required to be rank > 1.
+///
+/// Ex:
+/// ```
+/// %mul = arith.mulf %B_row, %A_row : vector<1x[4]xf32>
+/// %cast = vector.shape_cast %mul : vector<1x[4]xf32> to vector<[4]xf32>
+/// ```
+///
+/// gets converted to:
+///
+/// ```
+/// %B_row_sc = vector.shape_cast %B_row : vector<1x[4]xf32> to vector<[4]xf32>
+/// %A_row_sc = vector.shape_cast %A_row : vector<1x[4]xf32> to vector<[4]xf32>
+/// %mul = arith.mulf %B_row_sc, %A_row_sc : vector<[4]xf32>
+/// %cast_new = vector.shape_cast %mul : vector<[4]xf32> to vector<1x[4]xf32>
+/// %cast = vector.shape_cast %cast_new : vector<1x[4]xf32> to vector<[4]xf32>
+/// ```
+///
+/// Patterns for folding shape_casts should instantly eliminate `%cast_new` and
+/// `%cast`.
+struct DropUnitDimFromElementwiseOps final
+ : public OpTraitRewritePattern<OpTrait::Elementwise> {
+ using OpTraitRewritePattern::OpTraitRewritePattern;
+ LogicalResult matchAndRewrite(Operation *op,
+ PatternRewriter &rewriter) const override {
+ if (op->getNumResults() != 1)
+ return failure();
+
+ // Check the pre-condiitions. For `Elementwise` Ops all operands
+ // are guaranteed to have identical shapes and it suffices to only check the
+ // first one.
+ auto op1 = op->getOperands()[0];
+ auto sourceVectorType = dyn_cast<VectorType>(op1.getType());
+ if (!sourceVectorType)
+ return failure();
+
+ if (sourceVectorType.getRank() < 2)
+ return failure();
+
+ bool hasTrailingDimUnitFixed =
+ ((sourceVectorType.getShape().back() == 1) &&
+ (!sourceVectorType.getScalableDims().back()));
+ bool hasLeadingDimUnitFixed =
+ ((sourceVectorType.getShape().front() == 1) &&
+ (!sourceVectorType.getScalableDims().front()));
+ if (!hasLeadingDimUnitFixed && !hasTrailingDimUnitFixed)
+ return failure();
+
+ // Drop leading/trailing unit dim by applying vector.shape_cast to all
+ // operands
+ auto elTy = sourceVectorType.getElementType();
+ int64_t dim = hasLeadingDimUnitFixed ? 0 : sourceVectorType.getRank() - 1;
+ VectorType newVType = VectorType::Builder(sourceVectorType).dropDim(dim);
+
+ SmallVector<Value> newOperands;
+ auto loc = op->getLoc();
+ for (auto operand : op->getOperands()) {
+ auto opSC = rewriter.create<vector::ShapeCastOp>(loc, newVType, operand);
+ newOperands.push_back(opSC);
+ }
+
+ // Create an updated elementwise Op without leading/trailing unit dim
+ Operation *elementwiseOp =
+ rewriter.create(loc, op->getName().getIdentifier(), newOperands,
+ newVType, op->getAttrs());
+
+ // Restore the leading/trailing unit dim by applying vector.shape_cast to
+ // the result
+ rewriter.replaceOpWithNewOp<ShapeCastOp>(op, sourceVectorType,
+ elementwiseOp->getResult(0));
+
+ return success();
+ }
+};
+
/// Pattern to eliminate redundant zero-constants added to reduction operands.
/// It's enough for there to be one initial zero value, so we can eliminate the
/// extra ones that feed into `vector.reduction <add>`. These get created by the
@@ -1514,6 +1598,12 @@ void mlir::vector::populateShapeCastFoldingPatterns(RewritePatternSet &patterns,
patterns.add<ShapeCastOpFolder>(patterns.getContext(), benefit);
}
+void mlir::vector::populateDropUnitDimWithShapeCastPatterns(
+ RewritePatternSet &patterns, PatternBenefit benefit) {
+ patterns.add<DropUnitDimFromElementwiseOps, ShapeCastOpFolder>(
+ patterns.getContext(), benefit);
+}
+
void mlir::vector::populateBubbleVectorBitCastOpPatterns(
RewritePatternSet &patterns, PatternBenefit benefit) {
patterns.add<BubbleDownVectorBitCastForExtract,
diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index ebec2274655e46..b81491b9c07404 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -254,3 +254,91 @@ func.func @transfer_read_flattenable_negative2(
// CHECK-LABEL: func @transfer_read_flattenable_negative2
// CHECK: vector.transfer_read {{.*}} vector<5x4x3x2xi8>
+
+// -----
+
+func.func @fold_unit_dim_add_basic(%arg0 : vector<1x8xi32>) -> vector<1x8xi32> {
+ %add = arith.addi %arg0, %arg0 : vector<1x8xi32>
+ return %add : vector<1x8xi32>
+}
+// CHECK-LABEL: func.func @fold_unit_dim_add_basic(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<1x8xi32>) -> vector<1x8xi32> {
+// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8xi32> to vector<8xi32>
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8xi32> to vector<8xi32>
+// CHECK: %[[VAL_3:.*]] = arith.addi %[[VAL_1]], %[[VAL_2]] : vector<8xi32>
+// CHECK: %[[VAL_4:.*]] = vector.shape_cast %[[VAL_3]] : vector<8xi32> to vector<1x8xi32>
+// CHECK: return %[[VAL_4]] : vector<1x8xi32>
+
+// -----
+
+func.func @fold_unit_dim_add_leading_and_trailing(%arg0 : vector<1x8x1xi32>) -> vector<1x8x1xi32> {
+ %add = arith.addi %arg0, %arg0 : vector<1x8x1xi32>
+ return %add : vector<1x8x1xi32>
+}
+// CHECK-LABEL: func.func @fold_unit_dim_add_leading_and_trailing(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<1x8x1xi32>) -> vector<1x8x1xi32> {
+// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8x1xi32> to vector<8xi32>
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8x1xi32> to vector<8xi32>
+// CHECK: %[[VAL_3:.*]] = arith.addi %[[VAL_1]], %[[VAL_2]] : vector<8xi32>
+// CHECK: %[[VAL_4:.*]] = vector.shape_cast %[[VAL_3]] : vector<8xi32> to vector<1x8x1xi32>
+// CHECK: return %[[VAL_4]] : vector<1x8x1xi32>
+
+// -----
+
+func.func @fold_unit_dim_add(%arg0 : vector<8x1xi32>,
+ %arg1 : vector<1x8xi32>) -> vector<8xi32> {
+ %sc_arg0 = vector.shape_cast %arg0 : vector<8x1xi32> to vector<1x8xi32>
+ %add = arith.addi %sc_arg0, %arg1 : vector<1x8xi32>
+ %res = vector.shape_cast %add : vector<1x8xi32> to vector<8xi32>
+ return %res : vector<8xi32>
+}
+
+// CHECK-LABEL: func.func @fold_unit_dim_add(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1xi32>,
+// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8xi32>) -> vector<8xi32> {
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1xi32> to vector<8xi32>
+// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8xi32> to vector<8xi32>
+// CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_2]], %[[VAL_3]] : vector<8xi32>
+// CHECK: return %[[VAL_4]] : vector<8xi32>
+
+// -----
+
+func.func @fold_unit_dim_mulf(%arg0 : vector<8x[2]x1xf32>,
+ %arg1 : vector<1x8x[2]xf32>) -> vector<8x[2]xf32> {
+ %sc_arg0 = vector.shape_cast %arg0 : vector<8x[2]x1xf32> to vector<1x8x[2]xf32>
+ %add = arith.mulf %sc_arg0, %arg1 : vector<1x8x[2]xf32>
+ %res = vector.shape_cast %add : vector<1x8x[2]xf32> to vector<8x[2]xf32>
+ return %res : vector<8x[2]xf32>
+}
+
+// CHECK-LABEL: func.func @fold_unit_dim_mulf(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8x[2]x1xf32>,
+// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x[2]xf32>) -> vector<8x[2]xf32> {
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x[2]x1xf32> to vector<8x[2]xf32>
+// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x[2]xf32> to vector<8x[2]xf32>
+// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x[2]xf32>
+// CHECK: return %[[VAL_4]] : vector<8x[2]xf32>
+
+// -----
+
+// All shape casts are folded away
+
+func.func @fold_unit_dims_entirely(%arg0 : vector<8xi32>,
+ %arg1 : vector<8xi32>,
+ %arg2 : vector<8xi32>) -> vector<8xi32> {
+ %sc_arg0 = vector.shape_cast %arg0 : vector<8xi32> to vector<1x8xi32>
+ %sc_arg1 = vector.shape_cast %arg1 : vector<8xi32> to vector<1x8xi32>
+ %sc_arg2 = vector.shape_cast %arg2 : vector<8xi32> to vector<1x8xi32>
+ %mul = arith.muli %sc_arg0, %sc_arg1 : vector<1x8xi32>
+ %add = arith.addi %mul, %sc_arg2 : vector<1x8xi32>
+ %res = vector.shape_cast %add : vector<1x8xi32> to vector<8xi32>
+ return %res : vector<8xi32>
+}
+
+// CHECK-LABEL: func.func @fold_unit_dims_entirely(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8xi32>, %[[VAL_1:.*]]: vector<8xi32>,
+// CHECK-SAME: %[[VAL_2:.*]]: vector<8xi32>) -> vector<8xi32> {
+// CHECK: %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_1]] : vector<8xi32>
+// CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_2]] : vector<8xi32>
+// CHECK: return %[[VAL_4]] : vector<8xi32>
+
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index e593c0defcd29e..a643343e9342ad 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -455,6 +455,7 @@ struct TestFlattenVectorTransferPatterns
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<memref::MemRefDialect>();
registry.insert<affine::AffineDialect>();
+ registry.insert<vector::VectorDialect>();
}
void runOnOperation() override {
RewritePatternSet patterns(&getContext());
More information about the Mlir-commits
mailing list