[Mlir-commits] [mlir] [MLIR][Vector] Add scalable lowering for trailing transposes (PR #237937)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Nov 2 04:09:32 PST 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Muneeb Khan (muneebkhan85)
<details>
<summary>Changes</summary>
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.
---
Patch is 23.50 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/237937.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td (+22)
- (modified) mlir/include/mlir/Dialect/Vector/Transforms/LoweringPatterns.h (+9)
- (modified) mlir/include/mlir/Dialect/Vector/Transforms/VectorTransforms.h (+14)
- (modified) mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp (+7)
- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp (+187)
- (added) mlir/test/Dialect/Vector/vector-transpose-scalable-lowering.mlir (+147)
- (modified) mlir/test/python/dialects/transform_vector_ext.py (+4)
``````````diff
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.
+// C...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/237937
More information about the Mlir-commits
mailing list