diff options
| author | Crefeda Rodrigues <65665931+cfRod@users.noreply.github.com> | 2024-03-22 10:08:03 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-03-22 10:08:03 +0000 |
| commit | 465ea0bfa69aa48afef58666b084467a1c96c81b (patch) | |
| tree | d8d30be0cd5e83ddc22c42c0355724584bfd9f3e /mlir | |
| parent | 99d8c25b3104fc07f46532bd681515c5f3c71133 (diff) | |
| download | llvm-465ea0bfa69aa48afef58666b084467a1c96c81b.zip llvm-465ea0bfa69aa48afef58666b084467a1c96c81b.tar.gz llvm-465ea0bfa69aa48afef58666b084467a1c96c81b.tar.bz2 | |
[mlir][vector] Propagate scalability in TransferWriteNonPermutationLowering (#85632)
Updates `extendVectorRank` so that scalability in patterns
that use it (in particular, `TransferWriteNonPermutationLowering`),
is correctly propagated.
Closed related previous PR
https://github.com/llvm/llvm-project/pull/85270
---------
Signed-off-by: Crefeda Rodrigues <crefeda.rodrigues@arm.com>
Co-authored-by: Benjamin Maxwell <macdue@dueutil.tech>
Diffstat (limited to 'mlir')
| -rw-r--r-- | mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp | 8 | ||||
| -rw-r--r-- | mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir | 18 |
2 files changed, 24 insertions, 2 deletions
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp index 4a5e8fc..0693aa5 100644 --- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp @@ -41,8 +41,12 @@ static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec, SmallVector<int64_t> newShape(addedRank, 1); newShape.append(originalVecType.getShape().begin(), originalVecType.getShape().end()); - VectorType newVecType = - VectorType::get(newShape, originalVecType.getElementType()); + + SmallVector<bool> newScalableDims(addedRank, false); + newScalableDims.append(originalVecType.getScalableDims().begin(), + originalVecType.getScalableDims().end()); + VectorType newVecType = VectorType::get( + newShape, originalVecType.getElementType(), newScalableDims); return builder.create<vector::BroadcastOp>(loc, newVecType, vec); } diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir index 13e07f5..31bd19c 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir @@ -41,6 +41,24 @@ func.func @permutation_with_mask_scalable(%2: memref<?x?xf32>, %dim_1: index, %d return %1 : vector<8x[4]x2xf32> } +// CHECK: func.func @permutation_with_mask_transfer_write_scalable( +// CHECK-SAME: %[[ARG_0:.*]]: vector<4x[8]xi16>, +// CHECK-SAME: %[[ARG_1:.*]]: memref<1x4x?x1x1x1x1xi16>, +// CHECK-SAME: %[[MASK:.*]]: vector<4x[8]xi1>) { +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: %[[BCAST_1:.*]] = vector.broadcast %[[ARG_0]] : vector<4x[8]xi16> to vector<1x1x1x1x4x[8]xi16> +// CHECK: %[[BCAST_2:.*]] = vector.broadcast %[[MASK]] : vector<4x[8]xi1> to vector<1x1x1x1x4x[8]xi1> +// CHECK: %[[TRANSPOSE_1:.*]] = vector.transpose %[[BCAST_2]], [4, 5, 0, 1, 2, 3] : vector<1x1x1x1x4x[8]xi1> to vector<4x[8]x1x1x1x1xi1> +// CHECK: %[[TRANSPOSE_2:.*]] = vector.transpose %[[BCAST_1]], [4, 5, 0, 1, 2, 3] : vector<1x1x1x1x4x[8]xi16> to vector<4x[8]x1x1x1x1xi16> +// CHECK: vector.transfer_write %[[TRANSPOSE_2]], %[[ARG_1]]{{\[}}%[[C0]], %[[C0]], %[[C0]], %[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[TRANSPOSE_1]] {in_bounds = [true, true, true, true, true, true]} : vector<4x[8]x1x1x1x1xi16>, memref<1x4x?x1x1x1x1xi16> +// CHECK: return +func.func @permutation_with_mask_transfer_write_scalable(%arg0: vector<4x[8]xi16>, %arg1: memref<1x4x?x1x1x1x1xi16>, %mask: vector<4x[8]xi1>){ + %c0 = arith.constant 0 : index + vector.transfer_write %arg0, %arg1[%c0, %c0, %c0, %c0, %c0, %c0, %c0], %mask {in_bounds = [true, true], permutation_map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d2)> +} : vector<4x[8]xi16>, memref<1x4x?x1x1x1x1xi16> + + return +} module attributes {transform.with_named_sequence} { transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) { %f = transform.structured.match ops{["func.func"]} in %module_op |