[Mlir-commits] [mlir] 12b4951 - [mlir][vector] Add missing support for scalable vectors
Andrzej Warzynski
llvmlistbot at llvm.org
Thu Aug 10 02:08:59 PDT 2023
Author: Andrzej Warzynski
Date: 2023-08-10T09:08:30Z
New Revision: 12b4951866622cfc38342629808a1a6542624689
URL: https://github.com/llvm/llvm-project/commit/12b4951866622cfc38342629808a1a6542624689
DIFF: https://github.com/llvm/llvm-project/commit/12b4951866622cfc38342629808a1a6542624689.diff
LOG: [mlir][vector] Add missing support for scalable vectors
This patch adds the missing logic so that the
`TransferReadPermutationLowering` can be used for scalable vectors. To
this end:
* TransferOp custom C++ builder is updated to support scalable
vectors,
* `TransferOpReduceRank` is also updated to support scalable vectors.
This pattern is relevant when lowering `linalg.matmul` via
`vector_multi_reduction` for scalable vectors.
I've also updated relevant code in `TransferOpReduceRank` not to use
`llvm::to_vector` for constructing `SmallVector` from `ArrayRef`. That
hook doesn't work for `ArraryRef<bool>` (*), so for consistency I
switched to an explicit constructor (so that both `newShape` and
`newScalableDim` are constructed in a similar fashion).
(*) IIUC, that's due how implicit narrowing conversions between `bool`
and `*bool` work. Note that these narrowing conversions change when
using initializer lists, see
* https://en.cppreference.com/w/cpp/language/list_initialization.
Depends on D157092
Differential Revision: https://reviews.llvm.org/D157268
Added:
Modified:
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 546f1b9c872ca6..9ff35593a79419 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -4949,11 +4949,15 @@ void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,
Value vector, ArrayRef<int64_t> transp) {
VectorType vt = llvm::cast<VectorType>(vector.getType());
SmallVector<int64_t, 4> transposedShape(vt.getRank());
- for (unsigned i = 0; i < transp.size(); ++i)
+ SmallVector<bool, 4> transposedScalableDims(vt.getRank());
+ for (unsigned i = 0; i < transp.size(); ++i) {
transposedShape[i] = vt.getShape()[transp[i]];
+ transposedScalableDims[i] = vt.getScalableDims()[transp[i]];
+ }
result.addOperands(vector);
- result.addTypes(VectorType::get(transposedShape, vt.getElementType()));
+ result.addTypes(VectorType::get(transposedShape, vt.getElementType(),
+ transposedScalableDims));
result.addAttribute(TransposeOp::getTranspAttrName(result.name),
builder.getI64ArrayAttr(transp));
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
index 9589482cd8f8d2..68160bcf59e667 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
@@ -115,8 +115,11 @@ struct TransferReadPermutationLowering
// Apply the reverse transpose to deduce the type of the transfer_read.
ArrayRef<int64_t> originalShape = op.getVectorType().getShape();
SmallVector<int64_t> newVectorShape(originalShape.size());
+ ArrayRef<bool> originalScalableDims = op.getVectorType().getScalableDims();
+ SmallVector<bool> newScalableDims(originalShape.size());
for (const auto &pos : llvm::enumerate(permutation)) {
newVectorShape[pos.value()] = originalShape[pos.index()];
+ newScalableDims[pos.value()] = originalScalableDims[pos.index()];
}
// Transpose in_bounds attribute.
@@ -126,8 +129,8 @@ struct TransferReadPermutationLowering
: ArrayAttr();
// Generate new transfer_read operation.
- VectorType newReadType =
- VectorType::get(newVectorShape, op.getVectorType().getElementType());
+ VectorType newReadType = VectorType::get(
+ newVectorShape, op.getVectorType().getElementType(), newScalableDims);
Value newRead = rewriter.create<vector::TransferReadOp>(
op.getLoc(), newReadType, op.getSource(), op.getIndices(),
AffineMapAttr::get(newMap), op.getPadding(), op.getMask(),
@@ -345,14 +348,16 @@ struct TransferOpReduceRank : public OpRewritePattern<vector::TransferReadOp> {
return success();
}
- SmallVector<int64_t> newShape = llvm::to_vector<4>(
+ SmallVector<int64_t> newShape(
originalVecType.getShape().take_back(reducedShapeRank));
+ SmallVector<bool> newScalableDims(
+ originalVecType.getScalableDims().take_back(reducedShapeRank));
// Vector rank cannot be zero. Handled by TransferReadToVectorLoadLowering.
if (newShape.empty())
return rewriter.notifyMatchFailure(op, "rank-reduced vector is 0-d");
- VectorType newReadType =
- VectorType::get(newShape, originalVecType.getElementType());
+ VectorType newReadType = VectorType::get(
+ newShape, originalVecType.getElementType(), newScalableDims);
ArrayAttr newInBoundsAttr =
op.getInBounds()
? rewriter.getArrayAttr(
diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
index 6ea53aa3f41b07..0b738cbfcd7cad 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
@@ -1,12 +1,12 @@
// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
-// CHECK-LABEL: func @lower_permutation_with_mask(
+// CHECK-LABEL: func @lower_permutation_with_mask_fixed_width(
// CHECK: %[[vec:.*]] = arith.constant dense<-2.000000e+00> : vector<7x1xf32>
// CHECK: %[[mask:.*]] = arith.constant dense<[true, false, true, false, true, true, true]> : vector<7xi1>
// CHECK: %[[b:.*]] = vector.broadcast %[[mask]] : vector<7xi1> to vector<1x7xi1>
// CHECK: %[[tp:.*]] = vector.transpose %[[b]], [1, 0] : vector<1x7xi1> to vector<7x1xi1>
// CHECK: vector.transfer_write %[[vec]], %{{.*}}[%{{.*}}, %{{.*}}], %[[tp]] {in_bounds = [false, true]} : vector<7x1xf32>, memref<?x?xf32>
-func.func @lower_permutation_with_mask(%A : memref<?x?xf32>, %base1 : index,
+func.func @lower_permutation_with_mask_fixed_width(%A : memref<?x?xf32>, %base1 : index,
%base2 : index) {
%fn1 = arith.constant -2.0 : f32
%vf0 = vector.splat %fn1 : vector<7xf32>
@@ -17,6 +17,30 @@ func.func @lower_permutation_with_mask(%A : memref<?x?xf32>, %base1 : index,
return
}
+// CHECK-LABEL: func.func @permutation_with_mask_scalable(
+// CHECK-SAME: %[[ARG_0:.*]]: memref<?x?xf32>,
+// CHECK-SAME: %[[IDX_1:.*]]: index,
+// CHECK-SAME: %[[IDX_2:.*]]: index) -> vector<8x[4]x2xf32> {
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[PASS_THROUGH:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[IDX_2]], %[[IDX_1]] : vector<2x[4]xi1>
+// CHECK: %[[T_READ:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0]], %[[C0]]], %[[PASS_THROUGH]], %[[MASK]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x[4]xf32>
+// CHECK: %[[BCAST:.*]] = vector.broadcast %[[T_READ]] : vector<2x[4]xf32> to vector<8x2x[4]xf32>
+// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[BCAST]], [0, 2, 1] : vector<8x2x[4]xf32> to vector<8x[4]x2xf32>
+// CHECK: return %[[TRANSPOSE]] : vector<8x[4]x2xf32>
+// CHECK: }
+func.func @permutation_with_mask_scalable(%2: memref<?x?xf32>, %dim_1: index, %dim_2: index) -> (vector<8x[4]x2xf32>) {
+
+ %c0 = arith.constant 0 : index
+ %cst_0 = arith.constant 0.000000e+00 : f32
+
+ %mask = vector.create_mask %dim_2, %dim_1 : vector<2x[4]xi1>
+ %1 = vector.transfer_read %2[%c0, %c0], %cst_0, %mask
+ {in_bounds = [true, true, true], permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>}
+ : memref<?x?xf32>, vector<8x[4]x2xf32>
+ return %1 : vector<8x[4]x2xf32>
+}
+
transform.sequence failures(propagate) {
^bb1(%module_op: !transform.any_op):
%f = transform.structured.match ops{["func.func"]} in %module_op
More information about the Mlir-commits
mailing list