[Mlir-commits] [mlir] de61875 - [MLIR][Vector] Generalize DropUnitDimFromElementwiseOps to non leading / trailing dimensions. (#98455)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 17 02:22:30 PDT 2024
Author: Hugo Trachino
Date: 2024-07-17T10:22:25+01:00
New Revision: de61875e9d276cdb1899c2bea818bca86025b35a
URL: https://github.com/llvm/llvm-project/commit/de61875e9d276cdb1899c2bea818bca86025b35a
DIFF: https://github.com/llvm/llvm-project/commit/de61875e9d276cdb1899c2bea818bca86025b35a.diff
LOG: [MLIR][Vector] Generalize DropUnitDimFromElementwiseOps to non leading / trailing dimensions. (#98455)
Generalizes DropUnitDimFromElementwiseOps to support inner unit
dimensions.
This change stems from improving lowering of contractionOps for Arm SME.
Where we end up with inner unit dimensions on MulOp, BroadcastOp and
TransposeOp, preventing the generation of outerproducts.
discussed
[here](https://discourse.llvm.org/t/on-improving-arm-sme-lowering-resilience-in-mlir/78543/17?u=nujaa).
Fix after : https://github.com/llvm/llvm-project/pull/97652 showed an
unhandled edge case when all dimensions are one. The generated target
VectorType would be `vector<f32>` which is apparently not supported by
the mulf.
In case all dimensions are dropped, the target vectorType is
vector<1xf32>
---------
Co-authored-by: Benjamin Maxwell <macdue at dueutil.tech>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index b69e9421384b0..2686277bba59d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1627,7 +1627,33 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
}
};
-/// For vectors with either leading or trailing unit dim, replaces:
+// Helper function dropping unit non-scalable dimension from a VectorType
+// keeping at least 1 dimension to avoid generating 0-D vectors. Scalable unit
+// dimensions are not dropped. Folding such dimensions would require "shifting"
+// the scalable flag onto some other fixed-width dim (e.g. vector<[1]x4xf32> ->
+// vector<[4]xf32>). This could be implemented in the future.
+static VectorType dropNonScalableUnitDimFromType(VectorType inVecTy) {
+ auto inVecShape = inVecTy.getShape();
+ SmallVector<int64_t> newShape;
+ SmallVector<bool> newScalableDims;
+ for (auto [dim, isScalable] :
+ llvm::zip_equal(inVecShape, inVecTy.getScalableDims())) {
+ if (dim == 1 && !isScalable)
+ continue;
+
+ newShape.push_back(dim);
+ newScalableDims.push_back(isScalable);
+ }
+ // All dims have been dropped, return vector<1xeType>.
+ if (newShape.empty()) {
+ newShape.push_back(1);
+ newScalableDims.push_back(false);
+ }
+
+ return VectorType::get(newShape, inVecTy.getElementType(), newScalableDims);
+}
+
+/// For vectors with at least one unit dim, replaces:
/// elementwise(a, b)
/// with:
/// sc_a = shape_cast(a)
@@ -1639,20 +1665,16 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
/// 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`.
@@ -1677,37 +1699,26 @@ struct DropUnitDimFromElementwiseOps final
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
- int64_t dim = hasLeadingDimUnitFixed ? 0 : sourceVectorType.getRank() - 1;
-
SmallVector<Value> newOperands;
auto loc = op->getLoc();
for (auto operand : op->getOperands()) {
auto opVectorType = cast<VectorType>(operand.getType());
- VectorType newVType = VectorType::Builder(opVectorType).dropDim(dim);
+ auto newVType = dropNonScalableUnitDimFromType(opVectorType);
+ if (newVType == opVectorType)
+ return rewriter.notifyMatchFailure(op, "No unit dimension to remove.");
+
auto opSC = rewriter.create<vector::ShapeCastOp>(loc, newVType, operand);
newOperands.push_back(opSC);
}
VectorType newResultVectorType =
- VectorType::Builder(resultVectorType).dropDim(dim);
- // Create an updated elementwise Op without leading/trailing unit dim
+ dropNonScalableUnitDimFromType(resultVectorType);
+ // Create an updated elementwise Op without unit dim.
Operation *elementwiseOp =
rewriter.create(loc, op->getName().getIdentifier(), newOperands,
newResultVectorType, op->getAttrs());
- // Restore the leading/trailing unit dim by applying vector.shape_cast
- // to the result
+ // Restore the unit dim by applying vector.shape_cast to the result.
rewriter.replaceOpWithNewOp<ShapeCastOp>(op, resultVectorType,
elementwiseOp->getResult(0));
diff --git a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
index 5fd3cbd54aa58..303f841e8a828 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-flatten.mlir
@@ -604,6 +604,57 @@ func.func @fold_unit_dims_entirely(%arg0 : vector<8xi32>,
// -----
+func.func @fold_inner_unit_dim(%arg0 : vector<8x1x3xf128>,
+ %arg1 : vector<1x8x3xf128>) -> vector<8x3xf128> {
+ %sc_arg1 = vector.shape_cast %arg1 : vector<1x8x3xf128> to vector<8x1x3xf128>
+ %mul = arith.mulf %arg0, %sc_arg1 : vector<8x1x3xf128>
+ %res = vector.shape_cast %mul : vector<8x1x3xf128> to vector<8x3xf128>
+ return %res : vector<8x3xf128>
+}
+
+// CHECK-LABEL: func.func @fold_inner_unit_dim(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1x3xf128>,
+// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x3xf128>) -> vector<8x3xf128> {
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x3xf128> to vector<8x3xf128>
+// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x3xf128> to vector<8x3xf128>
+// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x3xf128>
+// CHECK: return %[[VAL_4]] : vector<8x3xf128>
+
+// -----
+
+func.func @fold_inner_unit_dim_scalable(%arg0 : vector<8x1x[1]x3xf128>,
+ %arg1 : vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {
+ %sc_arg1 = vector.shape_cast %arg1 : vector<1x8x[1]x3xf128> to vector<8x1x[1]x3xf128>
+ %mul = arith.mulf %arg0, %sc_arg1 : vector<8x1x[1]x3xf128>
+ %res = vector.shape_cast %mul : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>
+ return %res : vector<8x[1]x3xf128>
+}
+
+// CHECK-LABEL: func.func @fold_inner_unit_dim_scalable(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1x[1]x3xf128>,
+// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>
+// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x[1]x3xf128> to vector<8x[1]x3xf128>
+// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x[1]x3xf128>
+// CHECK: return %[[VAL_4]] : vector<8x[1]x3xf128>
+
+// -----
+
+func.func @fold_all_unit_dims(%arg0: vector<1x1xf32>) -> vector<1xf32> {
+ %0 = arith.mulf %arg0, %arg0 : vector<1x1xf32>
+ %res = vector.shape_cast %0 : vector<1x1xf32> to vector<1xf32>
+ return %res : vector<1xf32>
+}
+
+// CHECK-LABEL: func.func @fold_all_unit_dims(
+// CHECK-SAME: %[[VAL_0:.*]]: vector<1x1xf32>) -> vector<1xf32>
+// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x1xf32> to vector<1xf32>
+// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x1xf32> to vector<1xf32>
+// CHECK: %[[VAL_3:.*]] = arith.mulf %[[VAL_1]], %[[VAL_2]] : vector<1xf32>
+// CHECK: return %[[VAL_3]] : vector<1xf32>
+
+// -----
+
func.func @negative_out_of_bound_transfer_read(
%arg : memref<?x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {
%c0 = arith.constant 0 : index
More information about the Mlir-commits
mailing list