Skip to content

[mlir][vector] Refine Vector to LLVM lowering options #159553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Sep 23, 2025
Merged

[mlir][vector] Refine Vector to LLVM lowering options #159553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
060f098
[mlir][vector] Refine Vector to LLVM lowering options
banach-space Sep 17, 2025
9b36aca
Fix Python test
banach-space Sep 18, 2025
1269750
Addressing comments from Jakub and Diego
banach-space Sep 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions mlir/include/mlir/Conversion/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1489,8 +1489,8 @@ def ConvertVectorToLLVMPass : Pass<"convert-vector-to-llvm"> {
VectorContractLoweringAttr.summary, [{::llvm::cl::values(
clEnumValN(::mlir::vector::VectorContractLowering::Dot, "dot",
"Progressively lower to finer grained `vector.contract` and dot-products. (default)"),
clEnumValN(::mlir::vector::VectorContractLowering::Matmul, "matmul",
"Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics."),
clEnumValN(::mlir::vector::VectorContractLowering::LLVMIntr, "llvmintr",
"Lower directly to `llvm.intr.matrix.multiply`."),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

llvm only could be misleading as that's not the only option we have to lower to llvm. What about ~llvm-matmul?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point!

Since this is about lowering to "LLVM intrinsics", and the intrinsics are:

  • llvm.intr.matrix.multiply (i.e. llvm.intr.<...>)
  • llvm.intr.matrix.transpose (i.e. llvm.intr.<...>),

let me rename this as llvmintr.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LLVM intrinsic sounds like the best name for me

clEnumValN(::mlir::vector::VectorContractLowering::OuterProduct, "outerproduct",
"Lower to `vector.outerproduct`."),
clEnumValN(::mlir::vector::VectorContractLowering::ParallelArith, "parallelarith",
Expand All @@ -1502,8 +1502,8 @@ def ConvertVectorToLLVMPass : Pass<"convert-vector-to-llvm"> {
VectorTransposeLoweringAttr.summary, [{::llvm::cl::values(
clEnumValN(::mlir::vector::VectorTransposeLowering::EltWise, "eltwise",
"Lower transpose into element-wise extract and inserts (default)"),
clEnumValN(::mlir::vector::VectorTransposeLowering::Flat, "flat",
"Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix intrinsics"),
clEnumValN(::mlir::vector::VectorTransposeLowering::LLVMIntr, "llvmintr",
"Lower 2-D transpose directly to `llvm.intr.matrix.transpose`"),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same, in general, throughout the file

clEnumValN(::mlir::vector::VectorTransposeLowering::Shuffle1D, "shuffle1d",
"Lower 2-D transpose to `vector.shuffle` on 1-D vector."),
clEnumValN(::mlir::vector::VectorTransposeLowering::Shuffle16x16, "shuffle16x16",
Expand Down
17 changes: 8 additions & 9 deletions mlir/include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,9 @@ include "mlir/IR/EnumAttr.td"
// Lower transpose into element-wise extract and inserts.
def VectorTransposeLowering_Elementwise:
I32EnumAttrCase<"EltWise", 0, "eltwise">;
// Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix
// intrinsics.
def VectorTransposeLowering_FlatTranspose:
I32EnumAttrCase<"Flat", 1, "flat_transpose">;
// Lower directly to LLVM matrix intrinsics.
def VectorTransposeLowering_LLVMIntr:
I32EnumAttrCase<"LLVMIntr", 1, "llvmintr">;
// Lower 2-D transpose to `vector.shuffle` on 1-D vector.
def VectorTransposeLowering_Shuffle1D:
I32EnumAttrCase<"Shuffle1D", 2, "shuffle_1d">;
Expand All @@ -27,7 +26,7 @@ def VectorTransposeLowering_Shuffle16x16:
def VectorTransposeLoweringAttr : I32EnumAttr<
"VectorTransposeLowering",
"control the lowering of `vector.transpose` operations.",
[VectorTransposeLowering_Elementwise, VectorTransposeLowering_FlatTranspose,
[VectorTransposeLowering_Elementwise, VectorTransposeLowering_LLVMIntr,
VectorTransposeLowering_Shuffle1D, VectorTransposeLowering_Shuffle16x16]> {
let cppNamespace = "::mlir::vector";
}
Expand All @@ -48,9 +47,9 @@ def VectorMultiReductionLoweringAttr: I32EnumAttr<

// Progressively lower to finer grained `vector.contract` and dot-products.
def VectorContractLowering_Dot: I32EnumAttrCase<"Dot", 0, "dot">;
// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
def VectorContractLowering_Matmul:
I32EnumAttrCase<"Matmul", 1, "matmulintrinsics">;
// Lower directly to LLVM intrinsics.
def VectorContractLowering_LLVMIntr:
I32EnumAttrCase<"LLVMIntr", 1, "llvmintr">;
// Lower to `vector.outerproduct`.
def VectorContractLowering_OuterProduct:
I32EnumAttrCase<"OuterProduct", 2, "outerproduct">;
Expand All @@ -61,7 +60,7 @@ def VectorContractLowering_ParallelArith:
def VectorContractLoweringAttr: I32EnumAttr<
"VectorContractLowering",
"control the lowering of `vector.contract` operations.",
[VectorContractLowering_Dot, VectorContractLowering_Matmul,
[VectorContractLowering_Dot, VectorContractLowering_LLVMIntr,
VectorContractLowering_OuterProduct, VectorContractLowering_ParallelArith]> {
let cppNamespace = "::mlir::vector";
}
Expand Down
49 changes: 28 additions & 21 deletions mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1987,41 +1987,36 @@ struct VectorScalableStepOpLowering
/// %e = add %c, %d
/// ```
/// `vector.matrix_multiply` later lowers to `llvm.matrix.multiply`.
//
/// This only kicks in when vectorContractLowering is set to Matmul and
/// the vector.contract op is a row-major matrix multiply.
class ContractionOpToMatmulOpLowering
: public vector::MaskableOpRewritePattern<vector::ContractionOp> {
public:
using MaskableOpRewritePattern::MaskableOpRewritePattern;

ContractionOpToMatmulOpLowering(
vector::VectorContractLowering vectorContractLowering,
MLIRContext *context, PatternBenefit benefit = 100)
ContractionOpToMatmulOpLowering(MLIRContext *context,
PatternBenefit benefit = 100)
: MaskableOpRewritePattern<vector::ContractionOp>(context, benefit) {}

FailureOr<Value>
matchAndRewriteMaskableOp(vector::ContractionOp op, MaskingOpInterface maskOp,
PatternRewriter &rewriter) const override;
};

/// Progressively lower a `vector.contract %a, %b, %c` with row-major matmul
/// semantics to:
/// Lower a qualifying `vector.contract %a, %b, %c` (with row-major matmul
/// semantics directly into `llvm.intr.matrix.multiply`:
/// BEFORE:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit/optional: you can tag code blocks with mlir, e.g.:

```mlir
  %res = ...
```

also below

/// ```mlir
/// %res = vector.contract #matmat_trait %lhs, %rhs, %acc
/// : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32>
/// ```
/// %mta = maybe_transpose
/// %mtb = maybe_transpose
/// %flattened_a = vector.shape_cast %mta
/// %flattened_b = vector.shape_cast %mtb
/// %flattened_d = llvm.intr.matrix.multiply %flattened_a, %flattened_b
/// %mtd = vector.shape_cast %flattened_d
/// %d = maybe_untranspose %mtd
/// %e = add %c, %d
///
/// AFTER:
/// ```mlir
/// %lhs = vector.shape_cast %arg0 : vector<2x4xf32> to vector<8xf32>
/// %rhs = vector.shape_cast %arg1 : vector<4x3xf32> to vector<12xf32>
/// %matmul = llvm.intr.matrix.multiply %lhs, %rhs
/// %res = arith.addf %acc, %matmul : vector<2x3xf32>
/// ```
//
/// This only kicks in when vectorContractLowering is set to `Matmul`.
/// vector.transpose operations are inserted if the vector.contract op is not a
/// row-major matrix multiply.
///
/// Scalable vectors are not supported.
FailureOr<Value> ContractionOpToMatmulOpLowering::matchAndRewriteMaskableOp(
vector::ContractionOp op, MaskingOpInterface maskOp,
Expand Down Expand Up @@ -2116,7 +2111,19 @@ FailureOr<Value> ContractionOpToMatmulOpLowering::matchAndRewriteMaskableOp(
return res;
}

/// Lowers vector.transpose to llvm.intr.matrix.transpose
/// Lowers vector.transpose directly to llvm.intr.matrix.transpose
///
/// BEFORE:
/// ```mlir
/// %tr = vector.transpose %vec, [1, 0] : vector<2x4xf32> to vector<4x2xf32>
/// ```
/// AFTER:
/// ```mlir
/// %vec_cs = vector.shape_cast %vec : vector<2x4xf32> to vector<8xf32>
/// %tr = llvm.intr.matrix.transpose %vec_sc
/// {columns = 2 : i32, rows = 4 : i32} : vector<8xf32> into vector<8xf32>
/// %res = vector.shape_cast %tr : vector<8xf32> to vector<4x2xf32>
/// ```
class TransposeOpToMatrixTransposeOpLowering
: public OpRewritePattern<vector::TransposeOp> {
public:
Expand Down
4 changes: 2 additions & 2 deletions mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ void ConvertVectorToLLVMPass::runOnOperation() {
populateVectorBitCastLoweringPatterns(patterns);
populateVectorBroadcastLoweringPatterns(patterns);
populateVectorContractLoweringPatterns(patterns, vectorContractLowering);
if (vectorContractLowering == vector::VectorContractLowering::Matmul) {
if (vectorContractLowering == vector::VectorContractLowering::LLVMIntr) {
// This pattern creates a dependency on the LLVM dialect, hence we don't
// include it in `populateVectorContractLoweringPatterns` that is part of
// the Vector dialect (and should not depend on LLVM).
Expand All @@ -80,7 +80,7 @@ void ConvertVectorToLLVMPass::runOnOperation() {
populateVectorShapeCastLoweringPatterns(patterns);
populateVectorInterleaveLoweringPatterns(patterns);
populateVectorTransposeLoweringPatterns(patterns, vectorTransposeLowering);
if (vectorTransposeLowering == vector::VectorTransposeLowering::Flat) {
if (vectorTransposeLowering == vector::VectorTransposeLowering::LLVMIntr) {
// This pattern creates a dependency on the LLVM dialect, hence we don't
// include it in `populateVectorTransposeLoweringPatterns` that is part of
// the Vector dialect (and should not depend on LLVM).
Expand Down
6 changes: 3 additions & 3 deletions mlir/test/Conversion/VectorToLLVM/pass-option-serialization.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

// RUN: mlir-opt --convert-vector-to-llvm --dump-pass-pipeline %s 2>&1 | FileCheck %s --check-prefix=DEFAULT

// RUN: mlir-opt --convert-vector-to-llvm='vector-contract-lowering=matmul vector-transpose-lowering=flat' \
// RUN: mlir-opt --convert-vector-to-llvm='vector-contract-lowering=llvmintr vector-transpose-lowering=llvmintr' \
// RUN: --dump-pass-pipeline %s 2>&1 | FileCheck %s --check-prefix=NON-DEFAULT

// CHECK: builtin.module(
Expand All @@ -26,5 +26,5 @@
// CHECK-SAME: reassociate-fp-reductions={{[aA-zZ0-9]+}}
// DEFAULT: vector-contract-lowering=dot
// DEFAULT: vector-transpose-lowering=eltwise
// NON-DEFAULT: vector-contract-lowering=matmul
// NON-DEFAULT: vector-transpose-lowering=flat
// NON-DEFAULT: vector-contract-lowering=llvm
// NON-DEFAULT: vector-transpose-lowering=llvm
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: mlir-opt %s --convert-vector-to-llvm='vector-contract-lowering=matmul' | FileCheck %s
// RUN: mlir-opt %s --convert-vector-to-llvm='vector-contract-lowering=llvmintr' | FileCheck %s

#matmat_accesses = [
affine_map<(i, j, k) -> (i, k)>,
Expand Down
2 changes: 1 addition & 1 deletion mlir/test/Dialect/Vector/vector-transpose-to-matrix-intrinsics-transform.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// RUN: mlir-opt %s --convert-vector-to-llvm='vector-transpose-lowering=flat' --split-input-file | FileCheck %s
// RUN: mlir-opt %s --convert-vector-to-llvm='vector-transpose-lowering=llvmintr' --split-input-file | FileCheck %s

// CHECK-LABEL: func @transpose(
func.func @transpose(%arg0: vector<2x4xf32>) -> vector<4x2xf32> {
Expand Down
12 changes: 6 additions & 6 deletions mlir/test/python/dialects/transform_vector_ext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,9 +74,9 @@ def enum_configurable_patterns():
# CHECK: transform.apply_patterns.vector.lower_contraction
vector.ApplyLowerContractionPatternsOp()
# CHECK: transform.apply_patterns.vector.lower_contraction
# CHECK-SAME: lowering_strategy = matmulintrinsics
# CHECK-SAME: lowering_strategy = llvmintr
vector.ApplyLowerContractionPatternsOp(
lowering_strategy=vector.VectorContractLowering.Matmul
lowering_strategy=vector.VectorContractLowering.LLVMIntr
)
# CHECK: transform.apply_patterns.vector.lower_contraction
# CHECK-SAME: lowering_strategy = parallelarith
Expand Down Expand Up @@ -105,9 +105,9 @@ def enum_configurable_patterns():
lowering_strategy=vector.VectorTransposeLowering.EltWise
)
# CHECK: transform.apply_patterns.vector.lower_transpose
# CHECK-SAME: lowering_strategy = flat_transpose
# CHECK-SAME: lowering_strategy = llvmintr
vector.ApplyLowerTransposePatternsOp(
lowering_strategy=vector.VectorTransposeLowering.Flat
lowering_strategy=vector.VectorTransposeLowering.LLVMIntr
)
# CHECK: transform.apply_patterns.vector.lower_transpose
# CHECK-SAME: lowering_strategy = shuffle_1d
Expand All @@ -120,10 +120,10 @@ def enum_configurable_patterns():
lowering_strategy=vector.VectorTransposeLowering.Shuffle16x16
)
# CHECK: transform.apply_patterns.vector.lower_transpose
# CHECK-SAME: lowering_strategy = flat_transpose
# CHECK-SAME: lowering_strategy = llvmintr
# CHECK-SAME: avx2_lowering_strategy = true
vector.ApplyLowerTransposePatternsOp(
lowering_strategy=vector.VectorTransposeLowering.Flat,
lowering_strategy=vector.VectorTransposeLowering.LLVMIntr,
avx2_lowering_strategy=True,
)

Expand Down