Skip to content

[MLIR][AArch64] Lower vector.contract to SVE FEAT_BF16 operations #147052

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.

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
4 changes: 4 additions & 0 deletions mlir/include/mlir/Conversion/Passes.td
Original file line number Diff line number Diff line change
Expand Up @@ -1437,6 +1437,10 @@ def ConvertVectorToLLVMPass : Pass<"convert-vector-to-llvm"> {
"bool", /*default=*/"false",
"Enables the use of Arm FEAT_I8MM instructions while lowering "
"the vector dialect.">,
Option<"armBF16", "enable-arm-bf16",
"bool", /*default=*/"false",
"Enables the use of Arm FEAT_BF16 instructions while lowering "
"the vector dialect.">,
Option<"x86Vector", "enable-x86vector",
"bool", /*default=*/"false",
"Enables the use of X86Vector dialect while lowering the vector "
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,25 @@ include "mlir/Dialect/Transform/IR/TransformAttrs.td"
include "mlir/Dialect/Transform/IR/TransformDialect.td"
include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"

def ApplyArmSVELowerContractionPatternsOp
def ApplyArmSVELowerContractionToI8MMPatternsOp
: Op<Transform_Dialect, "apply_patterns.arm_sve.vector_contract_to_i8mm",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Indicates that vector contraction-like operations should be lowered to
finer-grained vector primitives using the ArmSVE dialect.
Indicates that vector contract operations should be lowered to
to ArmSVE dialect operations mapping to instructions from FEAT_I8MM.
}];

let assemblyFormat = "attr-dict";
}

def ApplyArmSVELowerContractionToBFMMLAPatternsOp
: Op<Transform_Dialect, "apply_patterns.arm_sve.vector_contract_to_bfmmla",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
let description = [{
Indicates that vector contract operations should be lowered to
ArmSVE dialect operations mapping to instructions from FEAT_BF16.
}];

let assemblyFormat = "attr-dict";
}
#endif // ARMSVE_VECTOR_TRANSFORM_OPS
2 changes: 2 additions & 0 deletions mlir/include/mlir/Dialect/ArmSVE/Transforms/Transforms.h
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ void populateArmSVELegalizeForLLVMExportPatterns(
void populateLowerContractionToSVEI8MMPatternPatterns(
RewritePatternSet &patterns);

void populateLowerContractionToSVEBFMMLAPatterns(RewritePatternSet &patterns);

/// Configure the target to support lowering ArmSVE ops to ops that map to LLVM
/// intrinsics.
void configureArmSVELegalizeForExportTarget(LLVMConversionTarget &target);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,9 @@ void ConvertVectorToLLVMPass::runOnOperation() {
if (armSVE)
populateLowerContractionToSVEI8MMPatternPatterns(patterns);
}
if (armBF16)
populateLowerContractionToSVEBFMMLAPatterns(patterns);

(void)applyPatternsGreedily(getOperation(), std::move(patterns));
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
// TODO: There may be opportunities to unify this with a similar pattern
// for SVE. See:
// https://github.com/llvm/llvm-project/issues/145559
// LowerContractionToSVEI8MMPattern.cpp
// LowerContractToSVEPatterns.cpp
//
//===----------------------------------------------------------------------===//

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,16 @@ using namespace mlir;
// Apply...PatternsOp
//===----------------------------------------------------------------------===//

void transform::ApplyArmSVELowerContractionPatternsOp::populatePatterns(
void transform::ApplyArmSVELowerContractionToI8MMPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
mlir::populateLowerContractionToSVEI8MMPatternPatterns(patterns);
}

void transform::ApplyArmSVELowerContractionToBFMMLAPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
mlir::populateLowerContractionToSVEBFMMLAPatterns(patterns);
}

//===----------------------------------------------------------------------===//
// Transform op registration
//===----------------------------------------------------------------------===//
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Dialect/ArmSVE/Transforms/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
add_mlir_dialect_library(MLIRArmSVETransforms
LegalizeForLLVMExport.cpp
LegalizeVectorStorage.cpp
LowerContractionToSVEI8MMPattern.cpp
LowerContractToSVEPatterns.cpp

DEPENDS
MLIRArmSVEConversionsIncGen
Expand Down
Loading
Loading