[Mlir-commits] [mlir] [MLIR][Vector] Add scalable lowering for trailing transposes (PR #237937)
Muneeb Khan
llvmlistbot at llvm.org
Mon Nov 2 04:08:43 PST 2026
https://github.com/muneebkhan85 created https://github.com/llvm/llvm-project/pull/237937
Lower fixed-size vector.transpose operations that swap the final two dimensions to loops that process runtime-sized scalable-vector chunks. Treat all leading dimensions as batches and use masked loads and stores for tails, so the lowering remains vectorized instead of operating on scalar elements.
vector.interleave and vector.deinterleave operate on two streams, so one of the final two dimensions must contain exactly two elements. The transformation is opt-in because it introduces stack allocations and scf.for loops.
Expose the minimum target vector width through the Transform dialect. Test rank-2 and batched rank-N transposes in both directions, tails, unsupported shapes and permutations, scalable inputs, and configurable target widths.
>From 3747c85043fdfb14263eccaaaaa4f51098064ac0 Mon Sep 17 00:00:00 2001
From: Muneeb Khan <muneeb.khan at huawei.com>
Date: Mon, 2 Nov 2026 18:54:06 +0800
Subject: [PATCH] [MLIR][Vector] Add scalable lowering for trailing transposes
Lower fixed-size vector.transpose operations that swap the final two
dimensions to loops that process runtime-sized scalable-vector chunks. Treat
all leading dimensions as batches and use masked loads and stores for tails,
so the lowering remains vectorized instead of operating on scalar elements.
vector.interleave and vector.deinterleave operate on two streams, so one of
the final two dimensions must contain exactly two elements. The transformation
is opt-in because it introduces stack allocations and scf.for loops.
Expose the minimum target vector width through the Transform dialect. Test
rank-2 and batched rank-N transposes in both directions, tails, unsupported
shapes and permutations, scalable inputs, and configurable target widths.
---
.../Vector/TransformOps/VectorTransformOps.td | 22 +++
.../Vector/Transforms/LoweringPatterns.h | 9 +
.../Vector/Transforms/VectorTransforms.h | 14 ++
.../TransformOps/VectorTransformOps.cpp | 7 +
.../Transforms/LowerVectorTranspose.cpp | 187 ++++++++++++++++++
.../vector-transpose-scalable-lowering.mlir | 147 ++++++++++++++
.../python/dialects/transform_vector_ext.py | 4 +
7 files changed, 390 insertions(+)
create mode 100644 mlir/test/Dialect/Vector/vector-transpose-scalable-lowering.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index af690211b894de..9f1660c45a2f27 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -490,6 +490,28 @@ def ApplyLowerTransposePatternsOp : Op<Transform_Dialect,
}];
}
+def ApplyLowerTransposeToScalablePatternsOp : Op<Transform_Dialect,
+ "apply_patterns.vector.lower_transpose_to_scalable",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collects patterns that lower fixed-size `vector.transpose` operations
+ swapping the final two dimensions to runtime-sized scalable vector
+ operations. Any leading dimensions are treated as batch dimensions.
+
+ The lowering processes the fixed-size input in chunks whose runtime size
+ is derived from `vector.vscale` and `target_vector_bitwidth`. It uses
+ `vector.interleave` and `vector.deinterleave`, so one of the transposed
+ dimensions must be exactly 2.
+
+ This is an opt-in transformation because its profitability depends on the
+ target and it introduces stack allocations and an `scf.for` loop.
+ }];
+
+ let arguments = (ins
+ ConfinedAttr<I64Attr, [IntPositive]>:$target_vector_bitwidth);
+ let assemblyFormat = "`target_vector_bitwidth` `=` $target_vector_bitwidth attr-dict";
+}
+
def ApplyLowerInterleavePatternsOp : Op<Transform_Dialect,
"apply_patterns.vector.lower_interleave",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
index 556979ad3e9be0..ba52cfa4be0c20 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h
@@ -155,6 +155,15 @@ void populateVectorTransposeLoweringPatterns(
VectorTransposeLowering vectorTransposeLowering,
PatternBenefit benefit = 1);
+/// Populate patterns that lower fixed-size `vector.transpose` operations that
+/// swap the final two dimensions to runtime-sized scalable vector operations.
+/// The pattern uses binary `vector.interleave` and `vector.deinterleave`
+/// operations, so one of the two transposed dimensions must be exactly 2.
+void populateVectorTransposeToScalableLoweringPatterns(
+ RewritePatternSet &patterns,
+ const VectorTransposeToScalableLoweringOptions &options,
+ PatternBenefit benefit = 1);
+
/// Populate the pattern set with the following patterns:
///
/// [TransferReadToVectorLoadLowering]
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
index f98e0a47f39eea..27c13c78da5c3f 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h
@@ -62,6 +62,20 @@ struct VectorTransformsOptions {
}
};
+/// Options controlling the lowering of fixed-size transposes of the final two
+/// dimensions to scalable vector operations.
+struct VectorTransposeToScalableLoweringOptions {
+ /// Minimum size of a target scalable vector in bits. The lowering uses this
+ /// to derive the minimum number of lanes in its scalable vector types.
+ unsigned targetVectorBitWidth = 0;
+
+ VectorTransposeToScalableLoweringOptions &
+ setTargetVectorBitWidth(unsigned bitWidth) {
+ targetVectorBitWidth = bitWidth;
+ return *this;
+ }
+};
+
//===----------------------------------------------------------------------===//
// Standalone transformations and helpers.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index 44d461e4195c4f..660741a3589496 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -284,6 +284,13 @@ void transform::ApplyLowerTransposePatternsOp::populatePatterns(
}
}
+void transform::ApplyLowerTransposeToScalablePatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ vector::VectorTransposeToScalableLoweringOptions options;
+ options.setTargetVectorBitWidth(getTargetVectorBitwidth());
+ vector::populateVectorTransposeToScalableLoweringPatterns(patterns, options);
+}
+
void transform::ApplyLowerInterleavePatternsOp::populatePatterns(
RewritePatternSet &patterns) {
vector::populateVectorInterleaveLoweringPatterns(patterns);
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
index c3f7de0ac3c4e8..dcedf93bb5f5dc 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
@@ -11,7 +11,9 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
@@ -489,6 +491,183 @@ class TransposeOp2DToShuffleLowering
/// Options to control the vector patterns.
vector::VectorTransposeLowering vectorTransposeLowering;
};
+
+/// Lower a fixed-size transpose of the final two dimensions using scalable
+/// vector operations.
+///
+/// The input and output are staged through scoped stack allocations because a
+/// fixed-size vector cannot be directly cast to a runtime-sized scalable
+/// vector. Each loop iteration processes one scalable chunk and masks its
+/// tail.
+///
+/// This lowering requires one dimension to be exactly 2. That restriction is
+/// inherent in its use of the binary `vector.interleave` and
+/// `vector.deinterleave` constructs: the dimension of size 2 identifies the
+/// two streams that are joined or split by those operations.
+class TransposeOpToScalableLowering
+ : public OpRewritePattern<vector::TransposeOp> {
+public:
+ TransposeOpToScalableLowering(
+ MLIRContext *context,
+ const vector::VectorTransposeToScalableLoweringOptions &options,
+ PatternBenefit benefit = 1)
+ : OpRewritePattern<vector::TransposeOp>(context, benefit),
+ targetVectorBitWidth(options.targetVectorBitWidth) {}
+
+ LogicalResult matchAndRewrite(vector::TransposeOp op,
+ PatternRewriter &rewriter) const override {
+ VectorType srcType = op.getSourceVectorType();
+ int64_t rank = srcType.getRank();
+ if (targetVectorBitWidth == 0 || srcType.isScalable() || rank < 2)
+ return failure();
+
+ ArrayRef<int64_t> permutation = op.getPermutation();
+ for (int64_t dim = 0; dim < rank - 2; ++dim) {
+ if (permutation[dim] != dim)
+ return rewriter.notifyMatchFailure(
+ op, "expected only the final two dimensions to be transposed");
+ }
+ if (permutation[rank - 2] != rank - 1 || permutation[rank - 1] != rank - 2)
+ return rewriter.notifyMatchFailure(
+ op, "expected the final two dimensions to be transposed");
+
+ ArrayRef<int64_t> shape = srcType.getShape();
+ // Interleave and deinterleave operate on exactly two streams. Therefore,
+ // this lowering is applicable only when one of the final two dimensions
+ // is 2. Any preceding dimensions are handled as independent batches.
+ bool rowsToColumns = shape[rank - 1] == 2;
+ bool columnsToRows = shape[rank - 2] == 2;
+ if (!rowsToColumns && !columnsToRows)
+ return rewriter.notifyMatchFailure(
+ op, "expected one of the final two dimensions to have size 2");
+
+ Type elementType = srcType.getElementType();
+ if (!elementType.isIntOrFloat())
+ return rewriter.notifyMatchFailure(op,
+ "expected integer or float elements");
+ unsigned elementBitWidth = elementType.getIntOrFloatBitWidth();
+ if (elementBitWidth == 0 || targetVectorBitWidth % elementBitWidth != 0)
+ return rewriter.notifyMatchFailure(
+ op, "target vector bit width must be divisible by element bit width");
+
+ Location loc = op.getLoc();
+ VectorType dstType = op.getResultVectorType();
+ int64_t streamLength = rowsToColumns ? shape[rank - 2] : shape[rank - 1];
+ int64_t numElements = srcType.getNumElements();
+ int64_t sliceSize = 2 * streamLength;
+ int64_t batchCount = numElements / sliceSize;
+ int64_t minLanes = targetVectorBitWidth / elementBitWidth;
+ auto flatType = VectorType::get({numElements}, elementType);
+ auto bufferType = MemRefType::get({numElements}, elementType);
+ auto laneType = VectorType::get({minLanes}, elementType, {true});
+ auto pairType = VectorType::get({2 * minLanes}, elementType, {true});
+ auto laneMaskType = laneType.cloneWith(std::nullopt, rewriter.getI1Type());
+ auto pairMaskType = pairType.cloneWith(std::nullopt, rewriter.getI1Type());
+
+ auto scope =
+ memref::AllocaScopeOp::create(rewriter, loc, TypeRange{dstType});
+ Block *scopeBody = rewriter.createBlock(&scope.getBodyRegion());
+ rewriter.setInsertionPointToStart(scopeBody);
+
+ Value input = memref::AllocaOp::create(rewriter, loc, bufferType);
+ Value output = memref::AllocaOp::create(rewriter, loc, bufferType);
+ Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ Value two = arith::ConstantIndexOp::create(rewriter, loc, 2);
+ Value upper = arith::ConstantIndexOp::create(rewriter, loc, streamLength);
+ Value minLaneCount =
+ arith::ConstantIndexOp::create(rewriter, loc, minLanes);
+ Value vscale = vector::VectorScaleOp::create(rewriter, loc);
+ Value step = arith::MulIOp::create(rewriter, loc, minLaneCount, vscale);
+ Value flat =
+ vector::ShapeCastOp::create(rewriter, loc, flatType, op.getVector());
+ vector::StoreOp::create(rewriter, loc, flat, input, ValueRange{zero});
+
+ Value lanePassthru = arith::ConstantOp::create(
+ rewriter, loc,
+ SplatElementsAttr::get(laneType, rewriter.getZeroAttr(elementType)));
+ Value pairPassthru = arith::ConstantOp::create(
+ rewriter, loc,
+ SplatElementsAttr::get(pairType, rewriter.getZeroAttr(elementType)));
+
+ auto createTransposeLoop = [&](Value sliceBase) {
+ auto loop = scf::ForOp::create(rewriter, loc, zero, upper, step);
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(loop.getBody());
+ Value index = loop.getInductionVar();
+ Value remaining = arith::SubIOp::create(rewriter, loc, upper, index);
+ Value laneMask = vector::CreateMaskOp::create(
+ rewriter, loc, laneMaskType, ValueRange{remaining});
+ Value pairCount = arith::MulIOp::create(rewriter, loc, remaining, two);
+ Value pairMask = vector::CreateMaskOp::create(
+ rewriter, loc, pairMaskType, ValueRange{pairCount});
+ Value pairOffset = arith::MulIOp::create(rewriter, loc, index, two);
+ Value firstIndex =
+ arith::AddIOp::create(rewriter, loc, sliceBase, index);
+ Value secondOffset = arith::AddIOp::create(rewriter, loc, upper, index);
+ Value secondIndex =
+ arith::AddIOp::create(rewriter, loc, sliceBase, secondOffset);
+
+ if (rowsToColumns) {
+ Value pairIndex =
+ arith::AddIOp::create(rewriter, loc, sliceBase, pairOffset);
+ Value interleaved = vector::MaskedLoadOp::create(
+ rewriter, loc, pairType, input, ValueRange{pairIndex}, pairMask,
+ pairPassthru);
+ auto deinterleaved =
+ vector::DeinterleaveOp::create(rewriter, loc, interleaved);
+ vector::MaskedStoreOp::create(rewriter, loc, output,
+ ValueRange{firstIndex}, laneMask,
+ deinterleaved.getRes1());
+ vector::MaskedStoreOp::create(rewriter, loc, output,
+ ValueRange{secondIndex}, laneMask,
+ deinterleaved.getRes2());
+ } else {
+ Value lhs = vector::MaskedLoadOp::create(
+ rewriter, loc, laneType, input, ValueRange{firstIndex}, laneMask,
+ lanePassthru);
+ Value rhs = vector::MaskedLoadOp::create(
+ rewriter, loc, laneType, input, ValueRange{secondIndex}, laneMask,
+ lanePassthru);
+ Value interleaved =
+ vector::InterleaveOp::create(rewriter, loc, lhs, rhs);
+ Value pairIndex =
+ arith::AddIOp::create(rewriter, loc, sliceBase, pairOffset);
+ vector::MaskedStoreOp::create(rewriter, loc, output,
+ ValueRange{pairIndex}, pairMask,
+ interleaved);
+ }
+ }
+ return loop;
+ };
+
+ if (batchCount == 1) {
+ createTransposeLoop(zero);
+ } else {
+ Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
+ Value batchUpper =
+ arith::ConstantIndexOp::create(rewriter, loc, batchCount);
+ Value sliceSizeValue =
+ arith::ConstantIndexOp::create(rewriter, loc, sliceSize);
+ auto batchLoop = scf::ForOp::create(rewriter, loc, zero, batchUpper, one);
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointToStart(batchLoop.getBody());
+ Value sliceBase = arith::MulIOp::create(
+ rewriter, loc, batchLoop.getInductionVar(), sliceSizeValue);
+ createTransposeLoop(sliceBase);
+ }
+
+ Value result = vector::LoadOp::create(rewriter, loc, flatType, output,
+ ValueRange{zero});
+ result = vector::ShapeCastOp::create(rewriter, loc, dstType, result);
+ memref::AllocaScopeReturnOp::create(rewriter, loc, ValueRange{result});
+ rewriter.replaceOp(op, scope.getResults());
+ return success();
+ }
+
+private:
+ unsigned targetVectorBitWidth;
+};
} // namespace
void mlir::vector::populateVectorTransposeLoweringPatterns(
@@ -499,3 +678,11 @@ void mlir::vector::populateVectorTransposeLoweringPatterns(
patterns.add<TransposeOpLowering, TransposeOp2DToShuffleLowering>(
vectorTransposeLowering, patterns.getContext(), benefit);
}
+
+void mlir::vector::populateVectorTransposeToScalableLoweringPatterns(
+ RewritePatternSet &patterns,
+ const VectorTransposeToScalableLoweringOptions &options,
+ PatternBenefit benefit) {
+ patterns.add<TransposeOpToScalableLowering>(patterns.getContext(), options,
+ benefit);
+}
diff --git a/mlir/test/Dialect/Vector/vector-transpose-scalable-lowering.mlir b/mlir/test/Dialect/Vector/vector-transpose-scalable-lowering.mlir
new file mode 100644
index 00000000000000..220078e4d855f0
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-transpose-scalable-lowering.mlir
@@ -0,0 +1,147 @@
+// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s
+
+// Check the deinterleave direction, including a tail because 37 is not a
+// multiple of the minimum four f32 lanes in a 128-bit scalable vector.
+// CHECK-LABEL: func.func @rows_to_columns(
+// CHECK-NOT: vector.transpose
+// CHECK: %[[RESULT:.*]] = memref.alloca_scope -> (vector<2x37xf32>) {
+// CHECK: %[[INPUT:.*]] = memref.alloca() : memref<74xf32>
+// CHECK: %[[OUTPUT:.*]] = memref.alloca() : memref<74xf32>
+// CHECK: %[[VSCALE:.*]] = vector.vscale
+// CHECK: %[[STEP:.*]] = arith.muli %[[VSCALE]], %{{.*}} : index
+// CHECK: scf.for %[[IV:.*]] = %{{.*}} to %{{.*}} step %[[STEP]] {
+// CHECK: %[[PAIR:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<74xf32>, vector<[8]xi1>, vector<[8]xf32>
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[PAIR]] : vector<[8]xf32> -> vector<[4]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[EVEN]] : memref<74xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[ODD]] : memref<74xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: }
+// CHECK: memref.alloca_scope.return %{{.*}} : vector<2x37xf32>
+// CHECK: }
+// CHECK: return %[[RESULT]] : vector<2x37xf32>
+func.func @rows_to_columns(%arg: vector<37x2xf32>) -> vector<2x37xf32> {
+ %result = vector.transpose %arg, [1, 0] : vector<37x2xf32> to vector<2x37xf32>
+ return %result : vector<2x37xf32>
+}
+
+// Check the inverse interleave direction.
+// CHECK-LABEL: func.func @columns_to_rows(
+// CHECK-NOT: vector.transpose
+// CHECK: memref.alloca_scope -> (vector<37x2xf32>) {
+// CHECK: %[[INPUT:.*]] = memref.alloca() : memref<74xf32>
+// CHECK: %[[OUTPUT:.*]] = memref.alloca() : memref<74xf32>
+// CHECK: scf.for
+// CHECK: %[[LHS:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<74xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: %[[RHS:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<74xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: %[[PAIR:.*]] = vector.interleave %[[LHS]], %[[RHS]] : vector<[4]xf32> -> vector<[8]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[PAIR]] : memref<74xf32>, vector<[8]xi1>, vector<[8]xf32>
+func.func @columns_to_rows(%arg: vector<2x37xf32>) -> vector<37x2xf32> {
+ %result = vector.transpose %arg, [1, 0] : vector<2x37xf32> to vector<37x2xf32>
+ return %result : vector<37x2xf32>
+}
+
+// Leading dimensions are flattened into a batch loop. Each batch still uses
+// the same scalable deinterleave kernel for the final two dimensions.
+// CHECK-LABEL: func.func @batched_rows_to_columns(
+// CHECK-NOT: vector.transpose
+// CHECK: memref.alloca_scope -> (vector<2x5x2x37xf32>) {
+// CHECK: %[[INPUT:.*]] = memref.alloca() : memref<740xf32>
+// CHECK: %[[OUTPUT:.*]] = memref.alloca() : memref<740xf32>
+// CHECK: scf.for %[[BATCH:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK: %[[BASE:.*]] = arith.muli %[[BATCH]], %{{.*}} : index
+// CHECK: scf.for
+// CHECK: %[[PAIR:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<740xf32>, vector<[8]xi1>, vector<[8]xf32>
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[PAIR]] : vector<[8]xf32> -> vector<[4]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[EVEN]] : memref<740xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[ODD]] : memref<740xf32>, vector<[4]xi1>, vector<[4]xf32>
+func.func @batched_rows_to_columns(
+ %arg: vector<2x5x37x2xf32>) -> vector<2x5x2x37xf32> {
+ %result = vector.transpose %arg, [0, 1, 3, 2]
+ : vector<2x5x37x2xf32> to vector<2x5x2x37xf32>
+ return %result : vector<2x5x2x37xf32>
+}
+
+// Check the inverse batched interleave direction.
+// CHECK-LABEL: func.func @batched_columns_to_rows(
+// CHECK-NOT: vector.transpose
+// CHECK: memref.alloca_scope -> (vector<3x37x2xf32>) {
+// CHECK: %[[INPUT:.*]] = memref.alloca() : memref<222xf32>
+// CHECK: %[[OUTPUT:.*]] = memref.alloca() : memref<222xf32>
+// CHECK: scf.for
+// CHECK: scf.for
+// CHECK: %[[LHS:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<222xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: %[[RHS:.*]] = vector.maskedload %[[INPUT]][%{{.*}}], %{{.*}}, %{{.*}} : memref<222xf32>, vector<[4]xi1>, vector<[4]xf32>
+// CHECK: %[[PAIR:.*]] = vector.interleave %[[LHS]], %[[RHS]] : vector<[4]xf32> -> vector<[8]xf32>
+// CHECK: vector.maskedstore %[[OUTPUT]][%{{.*}}], %{{.*}}, %[[PAIR]] : memref<222xf32>, vector<[8]xi1>, vector<[8]xf32>
+func.func @batched_columns_to_rows(
+ %arg: vector<3x2x37xf32>) -> vector<3x37x2xf32> {
+ %result = vector.transpose %arg, [0, 2, 1]
+ : vector<3x2x37xf32> to vector<3x37x2xf32>
+ return %result : vector<3x37x2xf32>
+}
+
+// The binary interleave/deinterleave implementation requires one dimension to
+// be exactly 2, so a 7x3 transpose must remain unchanged.
+// CHECK-LABEL: func.func @no_dimension_of_size_two(
+// CHECK: vector.transpose %{{.*}}, [1, 0] : vector<7x3xf32> to vector<3x7xf32>
+func.func @no_dimension_of_size_two(%arg: vector<7x3xf32>) -> vector<3x7xf32> {
+ %result = vector.transpose %arg, [1, 0] : vector<7x3xf32> to vector<3x7xf32>
+ return %result : vector<3x7xf32>
+}
+
+// A dimension of size 2 is insufficient when dimensions other than the final
+// two are transposed.
+// CHECK-LABEL: func.func @non_trailing_transpose(
+// CHECK: vector.transpose %{{.*}}, [1, 0, 2] : vector<2x3x4xf32> to vector<3x2x4xf32>
+func.func @non_trailing_transpose(
+ %arg: vector<2x3x4xf32>) -> vector<3x2x4xf32> {
+ %result = vector.transpose %arg, [1, 0, 2]
+ : vector<2x3x4xf32> to vector<3x2x4xf32>
+ return %result : vector<3x2x4xf32>
+}
+
+// Already-scalable inputs are outside the scope of this fixed-to-scalable
+// lowering.
+// CHECK-LABEL: func.func @already_scalable(
+// CHECK: vector.transpose %{{.*}}, [1, 0] : vector<2x[4]xf32> to vector<[4]x2xf32>
+func.func @already_scalable(%arg: vector<2x[4]xf32>) -> vector<[4]x2xf32> {
+ %result = vector.transpose %arg, [1, 0] : vector<2x[4]xf32> to vector<[4]x2xf32>
+ return %result : vector<[4]x2xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root : !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.vector.lower_transpose_to_scalable
+ target_vector_bitwidth = 128
+ } : !transform.op<"func.func">
+ transform.yield
+ }
+}
+
+// -----
+
+// Check that the target vector width is configurable rather than implicitly
+// assuming SVE's 128-bit minimum. With 256 bits, f32 chunks have eight lanes.
+// CHECK-LABEL: func.func @target_width_256(
+// CHECK: vector.maskedload {{.*}} : memref<38xf32>, vector<[16]xi1>, vector<[16]xf32>
+// CHECK: vector.deinterleave {{.*}} : vector<[16]xf32> -> vector<[8]xf32>
+func.func @target_width_256(%arg: vector<19x2xf32>) -> vector<2x19xf32> {
+ %result = vector.transpose %arg, [1, 0] : vector<19x2xf32> to vector<2x19xf32>
+ return %result : vector<2x19xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(
+ %root : !transform.any_op {transform.readonly}) {
+ %funcs = transform.structured.match ops{["func.func"]} in %root
+ : (!transform.any_op) -> !transform.op<"func.func">
+ transform.apply_patterns to %funcs {
+ transform.apply_patterns.vector.lower_transpose_to_scalable
+ target_vector_bitwidth = 256
+ } : !transform.op<"func.func">
+ transform.yield
+ }
+}
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index 26225e014ed7b2..22d36bf4dc9493 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -148,6 +148,10 @@ def enum_configurable_patterns():
avx2_lowering_strategy=True,
)
+ # CHECK: transform.apply_patterns.vector.lower_transpose_to_scalable
+ # CHECK-SAME: target_vector_bitwidth = 128
+ vector.ApplyLowerTransposeToScalablePatternsOp(target_vector_bitwidth=128)
+
# CHECK: transform.apply_patterns.vector.split_transfer_full_partial
vector.ApplySplitTransferFullPartialPatternsOp()
# CHECK: transform.apply_patterns.vector.split_transfer_full_partial
More information about the Mlir-commits
mailing list