diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp index da5954b70a2ec..e344ea6b64f40 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp @@ -1622,7 +1622,33 @@ struct ChainedReduction final : OpRewritePattern { } }; -/// 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 newShape; + SmallVector 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) @@ -1634,20 +1660,16 @@ struct ChainedReduction final : OpRewritePattern { /// 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`. @@ -1672,37 +1694,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 newOperands; auto loc = op->getLoc(); for (auto operand : op->getOperands()) { auto opVectorType = cast(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(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(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>) -> vector<5x4x3x2xi8> { %c0 = arith.constant 0 : index