[Mlir-commits] [mlir] [mlir][vector] Expose vector linearize as transforms (PR #231412)
Adam Siemieniuk
llvmlistbot at llvm.org
Tue Oct 13 04:14:04 PDT 2026
https://github.com/adam-smnk updated https://github.com/llvm/llvm-project/pull/231412
>From 664a37003f73e90c9b36b14ab54169885587ee45 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Mon, 12 Oct 2026 12:05:19 -0700
Subject: [PATCH 1/2] [mlir][vector] Expose vector linearize as transforms
Adds vector linearize patterns as a transform op.
The transforms design follows existings tests which are also converted
to use transform schedules instead of dedicated test passes.
Assisted-by: Claude
---
.../Vector/TransformOps/VectorTransformOps.td | 47 +++++-
.../Vector/Transforms/VectorRewritePatterns.h | 6 +
.../TransformOps/VectorTransformOps.cpp | 43 +++++
.../Vector/Transforms/VectorLinearize.cpp | 67 ++++++++
.../Vector/linearize-subject-to-bitwidth.mlir | 8 +-
mlir/test/Dialect/Vector/linearize.mlir | 4 +-
mlir/test/Dialect/Vector/td/linearize.mlir | 32 ++++
.../Dialect/Vector/TestVectorTransforms.cpp | 159 ------------------
.../python/dialects/transform_vector_ext.py | 36 ++++
9 files changed, 239 insertions(+), 163 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/td/linearize.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
index 3bd53bb69c576..0106e6392947c 100644
--- a/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
+++ b/mlir/include/mlir/Dialect/Vector/TransformOps/VectorTransformOps.td
@@ -34,7 +34,52 @@ def ApplyVectorToLLVMConversionPatternsOp : Op<Transform_Dialect,
DefaultValuedAttr<BoolAttr, "false">:$reassociate_fp_reductions,
DefaultValuedAttr<BoolAttr, "true">:$force_32bit_vector_indices,
DefaultValuedAttr<BoolAttr, "false">:$use_vector_alignment);
- let assemblyFormat = "oilist(`reassociate_fp_reductions` `=` $reassociate_fp_reductions | `force_32bit_vector_indices` `=` $force_32bit_vector_indices | `use_vector_alignment` `=` $use_vector_alignment) attr-dict";
+ let assemblyFormat = [{
+ oilist (`reassociate_fp_reductions` `=` $reassociate_fp_reductions
+ | `force_32bit_vector_indices` `=` $force_32bit_vector_indices
+ | `use_vector_alignment` `=` $use_vector_alignment)
+ attr-dict
+ }];
+}
+
+def ApplyVectorLinearizeConversionPatternsOp : Op<Transform_Dialect,
+ "apply_conversion_patterns.vector.linearize",
+ [DeclareOpInterfaceMethods<ConversionPatternDescriptorOpInterface,
+ ["populateConversionTargetRules",
+ "verifyTypeConverter"]>]> {
+ let description = [{
+ Collects patterns for ND vector (n >= 2) vector linearization.
+
+ These patterns require the type converter built by the corresponding
+ linearize type converter op.
+
+ The patterns can be customized as follows:
+ - `shuffle_like_ops`: Linearizes ND vector ops into 1D vector shuffles.
+ - `target_vector_bitwidth`: If set, ops whose inner-most vector dimension
+ is at least this many bits wide are not linearized.
+ }];
+
+ let arguments = (ins
+ DefaultValuedAttr<BoolAttr, "false">:$shuffle_like_ops,
+ OptionalAttr<ConfinedAttr<I64Attr, [IntNonNegative]>>:$target_vector_bitwidth);
+ let assemblyFormat = [{
+ oilist (`shuffle_like_ops` `=` $shuffle_like_ops
+ | `target_vector_bitwidth` `=` $target_vector_bitwidth)
+ attr-dict
+ }];
+}
+
+def VectorLinearizeTypeConverterOp : Op<Transform_Dialect,
+ "apply_conversion_patterns.vector.linearize_type_converter",
+ [DeclareOpInterfaceMethods<TypeConverterBuilderOpInterface,
+ ["getTypeConverter",
+ "getTypeConverterType"]>]> {
+ let description = [{
+ Builds a type converter that controls which ND (n >= 2) vector types are
+ linearized into 1D forms, and materializes conversions.
+ }];
+
+ let assemblyFormat = "attr-dict";
}
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
index 6363279f6826b..8ecc8d7c3ef07 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
@@ -444,6 +444,12 @@ void populateVectorTransposeNarrowTypeRewritePatterns(
void populateForVectorLinearize(TypeConverter &typeConverter,
ConversionTarget &conversionTarget);
+/// Extends `conversionTarget` so that ops whose inner-most vector dimension is
+/// at least `targetBitWidth` bits wide (or has `index` elements) stay legal.
+/// Must be called after `populateForVectorLinearize`.
+void populateVectorLinearizeBitWidthTargetRules(
+ ConversionTarget &conversionTarget, unsigned targetBitWidth);
+
/// Populates `patterns` for ND vector (N >= 2) linearization. This currently
/// contains patterns for converting ConstantLike, Vectorizable, and
/// vector::BitCast ops.
diff --git a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
index c4c930738d2c9..b66d894ffcac3 100644
--- a/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
+++ b/mlir/lib/Dialect/Vector/TransformOps/VectorTransformOps.cpp
@@ -11,6 +11,7 @@
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialectDecl.h"
#include "mlir/Dialect/Transform/IR/TransformDialect.h"
#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
@@ -44,6 +45,47 @@ transform::ApplyVectorToLLVMConversionPatternsOp::verifyTypeConverter(
return success();
}
+void transform::ApplyVectorLinearizeConversionPatternsOp::populatePatterns(
+ TypeConverter &typeConverter, RewritePatternSet &patterns) {
+ ConversionTarget target(*getContext());
+ vector::populateVectorLinearizeBasePatterns(typeConverter, target, patterns);
+ if (getShuffleLikeOps())
+ vector::populateVectorLinearizeShuffleLikeOpsPatterns(typeConverter, target,
+ patterns);
+}
+
+void transform::ApplyVectorLinearizeConversionPatternsOp::
+ populateConversionTargetRules(const TypeConverter &typeConverter,
+ ConversionTarget &conversionTarget) {
+ // The legality rules need a mutable converter; they keep their own copy.
+ TypeConverter linearizer;
+ vector::populateForVectorLinearize(linearizer, conversionTarget);
+ if (std::optional<int64_t> bitwidth = getTargetVectorBitwidth())
+ vector::populateVectorLinearizeBitWidthTargetRules(conversionTarget,
+ *bitwidth);
+}
+
+LogicalResult
+transform::ApplyVectorLinearizeConversionPatternsOp::verifyTypeConverter(
+ transform::TypeConverterBuilderOpInterface builder) {
+ if (builder.getTypeConverterType() !=
+ VectorLinearizeTypeConverterOp::getTypeConverterType())
+ return emitOpError("expected VectorLinearizeTypeConverter");
+ return success();
+}
+
+std::unique_ptr<TypeConverter>
+transform::VectorLinearizeTypeConverterOp::getTypeConverter() {
+ auto typeConverter = std::make_unique<TypeConverter>();
+ ConversionTarget target(*getContext());
+ vector::populateForVectorLinearize(*typeConverter, target);
+ return typeConverter;
+}
+
+StringRef transform::VectorLinearizeTypeConverterOp::getTypeConverterType() {
+ return "VectorLinearizeTypeConverter";
+}
+
//===----------------------------------------------------------------------===//
// Apply...PatternsOp
//===----------------------------------------------------------------------===//
@@ -271,6 +313,7 @@ class VectorTransformDialectExtension
VectorTransformDialectExtension() {
declareGeneratedDialect<vector::VectorDialect>();
+ declareGeneratedDialect<arith::ArithDialect>();
declareGeneratedDialect<LLVM::LLVMDialect>();
registerTransformOps<
#define GET_OP_LIST
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 9ea70baa01bbe..a43f25c34def6 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -22,6 +22,7 @@
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/ArrayRef.h"
#include <cstdint>
+#include <limits>
#include <numeric>
#include <optional>
@@ -1007,6 +1008,72 @@ void mlir::vector::populateForVectorLinearize(TypeConverter &typeConverter,
});
}
+/// Get the set of operand/result types to check for sufficiently
+/// small inner-most dimension size.
+static SmallVector<std::pair<Type, unsigned>>
+getTypeBitWidthBoundPairs(Operation *op, unsigned targetBitWidth) {
+
+ if (auto insertOp = dyn_cast<vector::InsertOp>(op)) {
+ unsigned w = targetBitWidth < std::numeric_limits<unsigned>::max()
+ ? targetBitWidth + 1
+ : targetBitWidth;
+ return {{insertOp.getValueToStoreType(), w}};
+ }
+
+ auto resultTypes = op->getResultTypes();
+ SmallVector<std::pair<Type, unsigned>> resultsWithBitWidth;
+ resultsWithBitWidth.reserve(resultTypes.size());
+ for (Type type : resultTypes) {
+ resultsWithBitWidth.push_back({type, targetBitWidth});
+ }
+ return resultsWithBitWidth;
+}
+
+/// If `type` is VectorType with trailing dimension of (bit) size greater than
+/// or equal to `targetBitWidth`, its defining op is considered legal.
+static bool
+isNotLinearizableBecauseLargeInnerDimension(Type type,
+ unsigned targetBitWidth) {
+
+ VectorType vecType = dyn_cast<VectorType>(type);
+
+ // Not linearizable for reasons other than what this function checks.
+ if (!vecType || vecType.getRank() == 0)
+ return false;
+
+ // The width of the type 'index' is unbounded (and therefore potentially above
+ // the target width).
+ if (vecType.getElementType().isIndex())
+ return true;
+
+ unsigned finalDimSize = vecType.getShape().back();
+ unsigned nbBitsPerElm = vecType.getElementTypeBitWidth();
+ unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm;
+ return trailingVecDimBitWidth >= targetBitWidth;
+}
+
+static bool
+isNotLinearizableBecauseLargeInnerDimension(Operation *op,
+ unsigned targetBitWidth) {
+ // Check on bitwidths.
+ SmallVector<std::pair<Type, unsigned>> toCheck =
+ getTypeBitWidthBoundPairs(op, targetBitWidth);
+ return llvm::any_of(toCheck, [&](std::pair<Type, unsigned> typeWidth) {
+ return isNotLinearizableBecauseLargeInnerDimension(typeWidth.first,
+ typeWidth.second);
+ });
+}
+
+void mlir::vector::populateVectorLinearizeBitWidthTargetRules(
+ ConversionTarget &target, unsigned targetBitWidth) {
+ target.markUnknownOpDynamicallyLegal(
+ [=](Operation *op) -> std::optional<bool> {
+ if (isNotLinearizableBecauseLargeInnerDimension(op, targetBitWidth))
+ return true;
+ return std::nullopt;
+ });
+}
+
void mlir::vector::populateVectorLinearizeBasePatterns(
const TypeConverter &typeConverter, const ConversionTarget &target,
RewritePatternSet &patterns) {
diff --git a/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir b/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir
index 739fb2fb8b68b..18bdf762d3364 100644
--- a/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir
+++ b/mlir/test/Dialect/Vector/linearize-subject-to-bitwidth.mlir
@@ -1,5 +1,9 @@
-// RUN: mlir-opt %s -split-input-file -test-bit-width-constrained-vector-linearize=target-vector-bitwidth=128 | FileCheck %s --check-prefixes=ALL,BW-128
-// RUN: mlir-opt %s -split-input-file -test-bit-width-constrained-vector-linearize=target-vector-bitwidth=0 | FileCheck %s --check-prefixes=ALL,BW-0
+// RUN: mlir-opt %s -split-input-file \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/linearize.mlir' \
+// RUN: -transform-interpreter=entry-point=linearize_bitwidth_128 | FileCheck %s --check-prefixes=ALL,BW-128
+// RUN: mlir-opt %s -split-input-file \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/linearize.mlir' \
+// RUN: -transform-interpreter=entry-point=linearize_bitwidth_0 | FileCheck %s --check-prefixes=ALL,BW-0
// A vector<2x2xf32> has inner-most dimension with 64-bits. Check that at
// bitwidth threshold 128 (>= 64), operations are linearized, and at
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 3860e50a4924d..b0eaf7861e11e 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -1,4 +1,6 @@
-// RUN: mlir-opt %s -split-input-file -test-vector-linearize -verify-diagnostics | FileCheck %s
+// RUN: mlir-opt %s -split-input-file \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/linearize.mlir' \
+// RUN: -transform-interpreter=entry-point=linearize -verify-diagnostics | FileCheck %s
// CHECK-LABEL: test_linearize
// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)
diff --git a/mlir/test/Dialect/Vector/td/linearize.mlir b/mlir/test/Dialect/Vector/td/linearize.mlir
new file mode 100644
index 0000000000000..dfe9c7eefadc0
--- /dev/null
+++ b/mlir/test/Dialect/Vector/td/linearize.mlir
@@ -0,0 +1,32 @@
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @linearize(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_conversion_patterns to %func {
+ transform.apply_conversion_patterns.vector.linearize shuffle_like_ops = true
+ transform.apply_conversion_patterns.scf.structural_conversions
+ } with type_converter {
+ transform.apply_conversion_patterns.vector.linearize_type_converter
+ } <partial_conversion> : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @linearize_bitwidth_128(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_conversion_patterns to %func {
+ transform.apply_conversion_patterns.vector.linearize shuffle_like_ops = true target_vector_bitwidth = 128
+ } with type_converter {
+ transform.apply_conversion_patterns.vector.linearize_type_converter
+ } <partial_conversion> : !transform.any_op
+ transform.yield
+ }
+
+ transform.named_sequence @linearize_bitwidth_0(%module: !transform.any_op {transform.readonly}) {
+ %func = transform.structured.match ops{["func.func"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.apply_conversion_patterns to %func {
+ transform.apply_conversion_patterns.vector.linearize shuffle_like_ops = true target_vector_bitwidth = 0
+ } with type_converter {
+ transform.apply_conversion_patterns.vector.linearize_type_converter
+ } <partial_conversion> : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
index ab4386167f7ab..ab7bec74bb8fc 100644
--- a/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp
@@ -17,7 +17,6 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialectDecl.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
-#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
@@ -886,160 +885,6 @@ struct TestVectorEmulateMaskedLoadStore final
}
};
-/// Get the set of operand/result types to check for sufficiently
-/// small inner-most dimension size.
-static SmallVector<std::pair<Type, unsigned>>
-getTypeBitWidthBoundPairs(Operation *op, unsigned targetBitWidth) {
-
- if (auto insertOp = dyn_cast<vector::InsertOp>(op)) {
- unsigned w = targetBitWidth < std::numeric_limits<unsigned>::max()
- ? targetBitWidth + 1
- : targetBitWidth;
- return {{insertOp.getValueToStoreType(), w}};
- }
-
- auto resultTypes = op->getResultTypes();
- SmallVector<std::pair<Type, unsigned>> resultsWithBitWidth;
- resultsWithBitWidth.reserve(resultTypes.size());
- for (Type type : resultTypes) {
- resultsWithBitWidth.push_back({type, targetBitWidth});
- }
- return resultsWithBitWidth;
-}
-
-/// If `type` is VectorType with trailing dimension of (bit) size greater than
-/// or equal to `targetBitWidth`, its defining op is considered legal.
-static bool
-isNotLinearizableBecauseLargeInnerDimension(Type type,
- unsigned targetBitWidth) {
-
- VectorType vecType = dyn_cast<VectorType>(type);
-
- // Not linearizable for reasons other than what this function checks.
- if (!vecType || vecType.getRank() == 0)
- return false;
-
- // The width of the type 'index' is unbounded (and therefore potentially above
- // the target width).
- if (vecType.getElementType().isIndex())
- return true;
-
- unsigned finalDimSize = vecType.getShape().back();
- unsigned nbBitsPerElm = vecType.getElementTypeBitWidth();
- unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm;
- return trailingVecDimBitWidth >= targetBitWidth;
-}
-
-static bool
-isNotLinearizableBecauseLargeInnerDimension(Operation *op,
- unsigned targetBitWidth) {
- // Check on bitwidths.
- SmallVector<std::pair<Type, unsigned>> toCheck =
- getTypeBitWidthBoundPairs(op, targetBitWidth);
- return llvm::any_of(toCheck, [&](std::pair<Type, unsigned> typeWidth) {
- return isNotLinearizableBecauseLargeInnerDimension(typeWidth.first,
- typeWidth.second);
- });
-}
-
-void populateWithBitWidthConstraints(TypeConverter &typeConverter,
- ConversionTarget &target,
- unsigned targetBitWidth) {
-
- // The general purpose definition of what ops are legal must come first.
- populateForVectorLinearize(typeConverter, target);
-
- // Extend the set of legal ops to include those with large inner-most
- // dimensions on selected operands/results.
- target.markUnknownOpDynamicallyLegal(
- [=](Operation *op) -> std::optional<bool> {
- if (isNotLinearizableBecauseLargeInnerDimension(op, targetBitWidth)) {
- return true;
- }
- return {};
- });
-}
-
-struct TestVectorBitWidthLinearize final
- : public PassWrapper<TestVectorBitWidthLinearize, OperationPass<>> {
- MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorBitWidthLinearize)
-
- TestVectorBitWidthLinearize() = default;
- TestVectorBitWidthLinearize(const TestVectorBitWidthLinearize &pass)
- : PassWrapper(pass) {}
-
- StringRef getArgument() const override {
- return "test-bit-width-constrained-vector-linearize";
- }
- StringRef getDescription() const override {
- return "Linearizes ND vectors for N >= 2 into 1D vectors, with constraints "
- "in inner-most dimension's bit width.";
- }
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<vector::VectorDialect>();
- }
-
- Option<unsigned> targetVectorBitwidth{
- *this, "target-vector-bitwidth",
- llvm::cl::desc(
- "Minimum vector bitwidth to enable the flattening transformation"),
- llvm::cl::init(std::numeric_limits<unsigned>::max())};
- void runOnOperation() override {
- auto *context = &getContext();
-
- TypeConverter typeConverter;
- RewritePatternSet patterns(context);
- ConversionTarget target(*context);
-
- populateWithBitWidthConstraints(typeConverter, target,
- targetVectorBitwidth);
-
- vector::populateVectorLinearizeBasePatterns(typeConverter, target,
- patterns);
-
- vector::populateVectorLinearizeShuffleLikeOpsPatterns(typeConverter, target,
- patterns);
-
- if (failed(applyPartialConversion(getOperation(), target,
- std::move(patterns))))
- return signalPassFailure();
- }
-};
-
-struct TestVectorLinearize final
- : public PassWrapper<TestVectorLinearize, OperationPass<>> {
- MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestVectorLinearize)
-
- TestVectorLinearize() = default;
-
- StringRef getArgument() const override { return "test-vector-linearize"; }
- StringRef getDescription() const override {
- return "Linearizes ND vectors for N >= 2 into 1D vectors";
- }
- void getDependentDialects(DialectRegistry ®istry) const override {
- registry.insert<vector::VectorDialect, arith::ArithDialect>();
- }
-
- void runOnOperation() override {
- MLIRContext &context = getContext();
- TypeConverter converter;
- RewritePatternSet patterns(&context);
- ConversionTarget target(context);
-
- vector::populateForVectorLinearize(converter, target);
-
- vector::populateVectorLinearizeBasePatterns(converter, target, patterns);
- vector::populateVectorLinearizeShuffleLikeOpsPatterns(converter, target,
- patterns);
- mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
- converter, patterns, target);
-
- if (failed(applyPartialConversion(getOperation(), target,
- std::move(patterns))))
- return signalPassFailure();
- }
-};
-
struct TestVectorShuffleLowering
: public PassWrapper<TestVectorShuffleLowering,
OperationPass<func::FuncOp>> {
@@ -1099,10 +944,6 @@ void registerTestVectorLowerings() {
PassRegistration<TestFoldArithExtensionIntoVectorContractPatterns>();
PassRegistration<TestVectorEmulateMaskedLoadStore>();
-
- PassRegistration<TestVectorLinearize>();
-
- PassRegistration<TestVectorBitWidthLinearize>();
}
} // namespace test
} // namespace mlir
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index a3c53a45048b2..26225e014ed7b 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -5,6 +5,13 @@
from mlir.dialects.transform import vector
+def run(f):
+ with Context(), Location.unknown():
+ print("\nTEST:", f.__name__)
+ f()
+ return f
+
+
def run_apply_patterns(f):
with Context(), Location.unknown():
module = Module.create()
@@ -163,3 +170,32 @@ def enum_configurable_patterns():
vector.ApplySplitTransferFullPartialPatternsOp(
split_transfer_strategy=vector.VectorTransferSplit.ForceInBounds
)
+
+
+# CHECK-LABEL: TEST: linearize_conversion_patterns
+ at run
+def linearize_conversion_patterns():
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.AnyOpType.get(),
+ )
+ with InsertionPoint(sequence.body):
+ apply = transform.ApplyConversionPatternsOp(
+ sequence.bodyTarget, 1, partial_conversion=True
+ )
+ with InsertionPoint(apply.patterns.blocks.append()):
+ vector.ApplyVectorLinearizeConversionPatternsOp(
+ shuffle_like_ops=True, target_vector_bitwidth=128
+ )
+ with InsertionPoint(apply.default_type_converter_region[0].blocks.append()):
+ vector.VectorLinearizeTypeConverterOp()
+ transform.YieldOp()
+ # CHECK: apply_conversion_patterns to
+ # CHECK-NEXT: transform.apply_conversion_patterns.vector.linearize
+ # CHECK-SAME: shuffle_like_ops = true
+ # CHECK-SAME: target_vector_bitwidth = 128
+ # CHECK-NEXT: } with type_converter {
+ # CHECK-NEXT: transform.apply_conversion_patterns.vector.linearize_type_converter
+ # CHECK-NEXT: } <partial_conversion>
+ print(sequence)
>From 6241db5cf7eea5706bb1bf21c96ab2982e8a2854 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Tue, 13 Oct 2026 04:13:08 -0700
Subject: [PATCH 2/2] Refactor helpers
---
.../Vector/Transforms/VectorLinearize.cpp | 51 ++++++++-----------
1 file changed, 22 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index a43f25c34def6..6673c07132380 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -1029,38 +1029,31 @@ getTypeBitWidthBoundPairs(Operation *op, unsigned targetBitWidth) {
return resultsWithBitWidth;
}
-/// If `type` is VectorType with trailing dimension of (bit) size greater than
-/// or equal to `targetBitWidth`, its defining op is considered legal.
-static bool
-isNotLinearizableBecauseLargeInnerDimension(Type type,
- unsigned targetBitWidth) {
-
- VectorType vecType = dyn_cast<VectorType>(type);
-
- // Not linearizable for reasons other than what this function checks.
- if (!vecType || vecType.getRank() == 0)
- return false;
-
- // The width of the type 'index' is unbounded (and therefore potentially above
- // the target width).
- if (vecType.getElementType().isIndex())
- return true;
-
- unsigned finalDimSize = vecType.getShape().back();
- unsigned nbBitsPerElm = vecType.getElementTypeBitWidth();
- unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm;
- return trailingVecDimBitWidth >= targetBitWidth;
-}
-
-static bool
-isNotLinearizableBecauseLargeInnerDimension(Operation *op,
- unsigned targetBitWidth) {
+/// Returns true if any checked vector type of op's inner-most dimension
+/// size is equal to or exceeds target bit width.
+static bool hasWideInnerDim(Operation *op, unsigned targetBitWidth) {
// Check on bitwidths.
SmallVector<std::pair<Type, unsigned>> toCheck =
getTypeBitWidthBoundPairs(op, targetBitWidth);
+
+ // If `type` is VectorType with trailing dimension of (bit) size greater than
+ // or equal to `targetBitWidth`, its defining op is considered legal.
return llvm::any_of(toCheck, [&](std::pair<Type, unsigned> typeWidth) {
- return isNotLinearizableBecauseLargeInnerDimension(typeWidth.first,
- typeWidth.second);
+ VectorType vecType = dyn_cast<VectorType>(typeWidth.first);
+
+ // Not linearizable for reasons other than what this function checks.
+ if (!vecType || vecType.getRank() == 0)
+ return false;
+
+ // The width of the type 'index' is unbounded (and therefore potentially
+ // above the target width).
+ if (vecType.getElementType().isIndex())
+ return true;
+
+ unsigned finalDimSize = vecType.getShape().back();
+ unsigned nbBitsPerElm = vecType.getElementTypeBitWidth();
+ unsigned trailingVecDimBitWidth = finalDimSize * nbBitsPerElm;
+ return trailingVecDimBitWidth >= typeWidth.second;
});
}
@@ -1068,7 +1061,7 @@ void mlir::vector::populateVectorLinearizeBitWidthTargetRules(
ConversionTarget &target, unsigned targetBitWidth) {
target.markUnknownOpDynamicallyLegal(
[=](Operation *op) -> std::optional<bool> {
- if (isNotLinearizableBecauseLargeInnerDimension(op, targetBitWidth))
+ if (hasWideInnerDim(op, targetBitWidth))
return true;
return std::nullopt;
});
More information about the Mlir-commits
mailing list