[mlir][ArmSME] Add basic lowering of vector.transfer write to zero

[c-rhodes]

Summary

This patch adds support for lowering a vector.transfer_write of zeroes
and type vector<[16x16]xi8> to the SME zero {za} instruction [1],
which zeroes the entire accumulator, and then writing it out to memory
with the str instruction [2].

This contributes to supporting a path from linalg.fill to SME.

[1] https://developer.arm.com/documentation/ddi0602/2022-06/SME-Instructions/ZERO--Zero-a-list-of-64-bit-element-ZA-tiles-
[2] https://developer.arm.com/documentation/ddi0602/2022-06/SME-Instructions/STR--Store-vector-from-ZA-array-

[awarzynski]

Hi @c-rhodes , thanks for working on this :)

I've made a few suggestions inline. Mostly asking for more documentation and wondering whether this could be trimmed a bit more - there are some references to ArmSME Ops, but none are defined, so perhaps some code is not needed?

Sadly, it doesn't build for me ATM (tried ToT: 15a16ef8e06e):

/llvm-project/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEDialect.h:21:10: fatal error: 'mlir/Dialect/ArmSME/ArmSMEDialect.h.inc' file not found
#include "mlir/Dialect/ArmSME/ArmSMEDialect.h.inc"
         ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1 error generated.

mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEDialect.h

22

Is this needed?

25

If there are no SME Ops then why would this header be needed?

mlir/include/mlir/Dialect/ArmSME/Transforms/Transforms.h

25

1. There are no ArmSME ops :)
2. There are no patterns :)

Do we need this hook?

28

[nit] There are no ArmSME ops :)

mlir/lib/Dialect/ArmSME/Transforms/LowerVectorOps.cpp

9

Could you add a Doxygen note to document what is this file intended for?

19

We are missing some documentation :)

Also, it would be good to document that only the i8 case is supported ATM:

• element size: i8
• number of tiles: 1
• tile size: [16x16]xi8

And that this will be extended shortly :)

36

Could you replace 255 with some constant? Otherwise it's a magic number and it's unclear what it means.

[c-rhodes]

Fix build error and address comments.

[c-rhodes]

In D152508#4412985, @awarzynski wrote:

Hi @c-rhodes , thanks for working on this :)

I've made a few suggestions inline. Mostly asking for more documentation and wondering whether this could be trimmed a bit more - there are some references to ArmSME Ops, but none are defined, so perhaps some code is not needed?

Sadly, it doesn't build for me ATM (tried ToT: 15a16ef8e06e):

/llvm-project/mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEDialect.h:21:10: fatal error: 'mlir/Dialect/ArmSME/ArmSMEDialect.h.inc' file not found
#include "mlir/Dialect/ArmSME/ArmSMEDialect.h.inc"
         ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1 error generated.

thanks for reviewing! Apologies for the build error, path was wrong for that .inc (and a few others), it's in the IR directory. Should build now.

mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEDialect.h

22

Is this needed?

It is yeah that .inc file is generated by tablegen and declares the dialect.

25

If there are no SME Ops then why would this header be needed?

there's no custom ops but the intrinsic definition LLVM_aarch64_sme_zero is still an op.

mlir/include/mlir/Dialect/ArmSME/Transforms/Transforms.h

25

1. There are no ArmSME ops :)
2. There are no patterns :)

Do we need this hook?

there's no custom ops so you're right this isn't needed, removed it.

28

[nit] There are no ArmSME ops :)

The LLVM_aarch64_sme_zero intrinsic definition is as op, it's marked legal in this function.

[dcaballe]

LGTM % ongoing comments. Thanks!

mlir/lib/Dialect/ArmSME/Transforms/LowerVectorOps.cpp

46

Just a direct and simple translation from the Vector dialect... That's great!

mlir/test/Dialect/ArmSME/vector_ops.mlir

17

You can add --split-input-file and add a line with // ----- between each test for them to run independently and in parallel

[c-rhodes]

• Added missing LLVMIR test mlir/test/Target/LLVMIR/arm-sme.mlir.
• Use -split-input-file.

[c-rhodes]

In D152508#4416400, @dcaballe wrote:

LGTM % ongoing comments. Thanks!

Thanks for reviewing!

mlir/test/Dialect/ArmSME/vector_ops.mlir

17

You can add --split-input-file and add a line with // ----- between each test for them to run independently and in parallel

Done, thank you!

[awarzynski]

LGTM % ongoing comments. Thanks!

Same. LGTM, thanks Cullen!

[c-rhodes]

Apologies I shouldn't have committed this until there's consensus at the ODM, reverted.

[awarzynski]

In D152508#4420362, @c-rhodes wrote:

Apologies I shouldn't have committed this until there's consensus at the ODM, reverted.

Thank you and sorry for the confusion, I should've left a clearer message when approving this.

In general, we do have +1 from Arm (myself), +1 from Google (Diego), but we should also make sure that this works for Huawei (Frank). Or, at least, that this wouldn't block Frank from pursuing the approach taken in https://reviews.llvm.org/D152080 (if that's still the preference after the ODM). In the meantime, we could merge https://reviews.llvm.org/D152878 as that change will be required regardless of whether going via the Vector dialect or not.

Thanks again for all the effort working on this!

[WanderAway]

we should also make sure that this works for Huawei (Frank)

Thanks for the consideration. I don't think this is in direct conflict with what we want to do, but again, we can discuss further during the ODM.

[c-rhodes]

Rebased now D152878 has landed.

[c-rhodes]

Rebase again now D153050 has landed.

[awarzynski]

In D152508#4421516, @WanderAway wrote:

we should also make sure that this works for Huawei (Frank)

Thanks for the consideration. I don't think this is in direct conflict with what we want to do, but again, we can discuss further during the ODM.

@WanderAway Could you confirm that there are no objections from your side after the ODM? Thanks!

[WanderAway]

@WanderAway Could you confirm that there are no objections from your side after the ODM? Thanks!

Yup, no objections here.

[c-rhodes]

This patch was pretty basic so I've made some improvements/fixes:

• Write ZA out to memory after zero {za}.
• Add integration test (runs on QEMU).
• Check the vector.transfer_write value is a dense arith.constant of zeroes.
• Simplify tests in mlir/test/Dialect/ArmSME/vector_ops.mlir to pass memref as argument rather than set it up in each function. Also added 3 more tests that check lowering doesn't happen for:
  • non-memref types.
  • non zero values.
  • vector.transfer_write value op where defining value isn't visible, previously crashed if it was passed as argument.

I was a bit unsure whether to abandon this and post a new patch given this has already been approved and there's a few changes, happy to do that if people prefer.

[awarzynski]

Thanks for the updates, I've left some comments inline.

I was a bit unsure whether to abandon this and post a new patch given this has already been approved and there's a few changes, happy to do that if people prefer.

IMO we can continue here. You are basically refining the initial design rather than proposing something completely new.

mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.td

123

Could you update "mlir/test/Target/LLVMIR/arm-sme.mlir" as well?

mlir/test/Dialect/ArmSME/vector_ops.mlir

24

Doesn't this store the same array vector on each iteration? IIUC, the only thing that's changing is the destination.

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector_ops.mlir

11 ↗

(On Diff #534515)

[nit] %i1 (variable) is easy to confuse with i1 (type). I would use c1 instead. And if you need different types, c1_idx and c1_i32.

17 ↗

(On Diff #534515)

[nit] It's worth elaborating what svl stands for in this test. And how do we know that it's going to be "streaming vector length" rather than "vector length"?

33 ↗

(On Diff #534515)

[nit] Did you mean init_1 instead? Similar comment for addition below.

34 ↗

(On Diff #534515)

[nit] In this case the induction variable has a very specific meaning. Also, it can be confusing that there 16 bytes being loaded, but the induction variable is only increased by 1.

36 ↗

(On Diff #534515)

38 ↗

(On Diff #534515)

[nit] Same point as for the previous scf.for

40 ↗

(On Diff #534515)

50 ↗

(On Diff #534515)

It would be good to also verify that the result is != 1 when the elements in the matrix are different. Would it be possible to set one element to 123 and verify that the result is 123?

74 ↗

(On Diff #534515)

Shouldn't this be ... ?

82 ↗

(On Diff #534515)

It would be good to also verify that the result is != 0 when the elements in the matrix are different. Would it be possible to set one element to 321 and see what happens?

[awarzynski]

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector_ops.mlir

21–24 ↗

(On Diff #534515)

I am fine with loops, but not sure about the comment 🤔 .

vector.store - scalable - https://github.com/llvm/llvm-project/blob/79c83e12c8884fa46f2f2594836af93474f6ca5a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-sve.mlir#L9-L20

vector.transfer_write - scalable - https://github.com/llvm/llvm-project/blob/79c83e12c8884fa46f2f2594836af93474f6ca5a/mlir/test/Dialect/Linalg/vectorization-scalable.mlir#L27

Like I said, loops are fine (I really like the simplicity). I am just curious what exactly is missing :) But we can investigate that independently of this patch.

[c-rhodes]

Rebase and address comments.

[c-rhodes]

mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.td

123

Could you update "mlir/test/Target/LLVMIR/arm-sme.mlir" as well?

Good spot, cheers.

mlir/test/Dialect/ArmSME/vector_ops.mlir

24

Doesn't this store the same array vector on each iteration? IIUC, the only thing that's changing is the destination.

Doh! It does yeah good spot, fixed.

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector_ops.mlir

17 ↗

(On Diff #534515)

... And how do we know that it's going to be "streaming vector length" rather than "vector length"?

streaming-mode is enabled by the -enable-arm-streaming pass.

21–24 ↗

(On Diff #534515)

I am fine with loops, but not sure about the comment 🤔 .

vector.store - scalable - https://github.com/llvm/llvm-project/blob/79c83e12c8884fa46f2f2594836af93474f6ca5a/mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/test-sve.mlir#L9-L20

vector.transfer_write - scalable - https://github.com/llvm/llvm-project/blob/79c83e12c8884fa46f2f2594836af93474f6ca5a/mlir/test/Dialect/Linalg/vectorization-scalable.mlir#L27

Like I said, loops are fine (I really like the simplicity). I am just curious what exactly is missing :) But we can investigate that independently of this patch.

Apologies, I've updated the comment to reflect the actual problem.

33 ↗

(On Diff #534515)

[nit] Did you mean init_1 instead? Similar comment for addition below.

It was intended as first init / second init but I can see how that's confusing, updated to your suggestion.

38 ↗

(On Diff #534515)

[nit] Same point as for the previous scf.for

not sure col is applicable here, changed it to (row) offset.

74 ↗

(On Diff #534515)

Shouldn't this be ... ?

it should. Good spot.

[awarzynski]

Thanks for addressing my comments! I've tested locally and can confirm that the integration test runs correctly 🎉. Great job, LGTM!

Just a few final nits that you can either ignore or address when merging.

Btw, most test files use hyphen "-" rather than "_" underscore: "vector_ops.mlir" --> "vector-ops.mlir"?

As this remains within the scope of the original submission, I think it's fine to merge without waiting for the other reviewers to confirm (I'm just being mindful that this change has evolved since being originally OK'ed).

mlir/lib/Dialect/ArmSME/Transforms/LowerVectorOps.cpp

42

Just a nit. I feel that it would be good to make it clear, consistently, that these are "virtual" SME tiles.

53

Replace 16 with a constant (it's repeated a few times).

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector_ops.mlir

37–38 ↗

(On Diff #536170)

IMHO, readability would be better without wrapping lines like this one. And 80-char limit is rarely observed in tests. This is a nit ;-)

83 ↗

(On Diff #536170)

A comment might make it easier for our future selves to figure out where the "magic" 60 comes from :) This is a nit.

[c-rhodes]

In D152508#4467879, @awarzynski wrote:

Thanks for addressing my comments! I've tested locally and can confirm that the integration test runs correctly 🎉. Great job, LGTM!

Just a few final nits that you can either ignore or address when merging.

Btw, most test files use hyphen "-" rather than "_" underscore: "vector_ops.mlir" --> "vector-ops.mlir"?

As this remains within the scope of the original submission, I think it's fine to merge without waiting for the other reviewers to confirm (I'm just being mindful that this change has evolved since being originally OK'ed).

Thanks for reviewing again! Addressed all your comments before committing. Cheers.

mlir/include/mlir/Conversion/Passes.td

let options = [
     Option<"armNeon", "enable-arm-neon",
            "bool", /*default=*/"false",
            "Enables the use of ArmNeon dialect while lowering the vector "
      "dialect.">,
     Option<"armSVE", "enable-arm-sve",
            "bool", /*default=*/"false",
            "Enables the use of ArmSVE dialect while lowering the vector "
        "dialect.">,
+    Option<"armSME", "enable-arm-sme",
+           "bool", /*default=*/"false",
+           "Enables the use of ArmSME dialect while lowering the vector "
+       "dialect.">,
     Option<"x86Vector", "enable-x86vector",
            "bool", /*default=*/"false",
            "Enables the use of X86Vector dialect while lowering the vector "
      "dialect.">,
     Option<"useOpaquePointers", "use-opaque-pointers", "bool",
                /*default=*/"true", "Generate LLVM IR using opaque pointers "
                "instead of typed pointers">
   ];

mlir/include/mlir/Dialect/ArmSME/CMakeLists.txt

+add_subdirectory(IR)
 add_subdirectory(Transforms)

mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.td

+//===-- ArmSME.td - ArmSME dialect operation definitions ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the basic operations for the ArmSME dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ArmSME
+#define ArmSME
+
+include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+
+//===----------------------------------------------------------------------===//
+// ArmSME dialect definition.
+//===----------------------------------------------------------------------===//
+
+def ArmSME_Dialect : Dialect {
+  let name = "arm_sme";
+  let cppNamespace = "::mlir::arm_sme";
+  let summary = "Dialect to target the Armv9 Scalable Matrix Extension (SME)";
+  let description = [{
+    This dialect contains the definitions necessary to target specific Arm SME
+    operations.
+
+    For more details on the architecture, see the Arm documentation:
+    https://developer.arm.com/documentation/ddi0616
+  }];
+  let usePropertiesForAttributes = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// LLVMIR Intrinsics
+//===----------------------------------------------------------------------===//
+
+class ArmSME_IntrOp<string mnemonic, int numResults = 1,
+                    list<Trait> traits = []> :
+  LLVM_IntrOpBase<ArmSME_Dialect, "intr." # mnemonic,
+                  "aarch64_sme_" # !subst(".", "_", mnemonic), [], [], traits,
+                  numResults>;
+
+/// Create a call to aarch64_sme_zero intrinsic.
+def LLVM_aarch64_sme_zero
+    : ArmSME_IntrOp<"zero", 0>, Arguments<(ins I32:$imm)>;
+
+#endif // ArmSME

mlir/include/mlir/Dialect/ArmSME/IR/ArmSMEDialect.h

+//===- ArmSMEDialect.h - MLIR Dialect for Arm SME ---------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the Target dialect for ArmSME in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARMSME_ARMSMEDIALECT_H
+#define MLIR_DIALECT_ARMSME_ARMSMEDIALECT_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h.inc"
+
+#endif // MLIR_DIALECT_ARMSME_ARMSMEDIALECT_H

mlir/include/mlir/Dialect/ArmSME/IR/CMakeLists.txt

+add_mlir_dialect(ArmSME arm_sme ArmSME)
+add_mlir_doc(ArmSME ArmSME Dialects/ -gen-dialect-doc -dialect=arm_sme)
+
+set(LLVM_TARGET_DEFINITIONS ArmSME.td)
+mlir_tablegen(ArmSMEConversions.inc -gen-llvmir-conversions)
+add_public_tablegen_target(MLIRArmSMEConversionsIncGen)

mlir/include/mlir/Dialect/ArmSME/Transforms/Transforms.h

+//===- Transforms.h - ArmSME Dialect Transformation Entrypoints -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARMSME_TRANSFORMS_H
+#define MLIR_DIALECT_ARMSME_TRANSFORMS_H
+
+namespace mlir {
+
+class LLVMConversionTarget;
+class RewritePatternSet;
+
+namespace arm_sme {
+void populateVectorTransferLoweringPatterns(RewritePatternSet &patterns);
+} // namespace arm_sme
+
+/// Configure the target to support lowering ArmSME ops to ops that map to LLVM
+/// intrinsics.
+void configureArmSMELegalizeForExportTarget(LLVMConversionTarget &target);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARMSME_TRANSFORMS_H

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
registry.insert<acc::OpenACCDialect,
                   nvgpu::NVGPUDialect,
                   scf::SCFDialect,
                   omp::OpenMPDialect,
                   pdl::PDLDialect,
                   pdl_interp::PDLInterpDialect,
                   quant::QuantizationDialect,
                   spirv::SPIRVDialect,
                   arm_sve::ArmSVEDialect,
+                  arm_sme::ArmSMEDialect,
                   vector::VectorDialect,
                   NVVM::NVVMDialect,
                   ROCDL::ROCDLDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   transform::TransformDialect,
                   tosa::TosaDialect,

mlir/include/mlir/Target/LLVMIR/Dialect/All.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 #define MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 
 #include "mlir/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.h"
 
 namespace mlir {
 class DialectRegistry;
 
 /// Registers all dialects that can be translated to LLVM IR and the
 /// corresponding translation interfaces.
 static inline void registerAllToLLVMIRTranslations(DialectRegistry &registry) {
   registerArmNeonDialectTranslation(registry);
   registerAMXDialectTranslation(registry);
   registerArmSVEDialectTranslation(registry);
+  registerArmSMEDialectTranslation(registry);
   registerBuiltinDialectTranslation(registry);
   registerGPUDialectTranslation(registry);
   registerLLVMDialectTranslation(registry);
   registerNVVMDialectTranslation(registry);
   registerOpenACCDialectTranslation(registry);
   registerOpenMPDialectTranslation(registry);
   registerROCDLDialectTranslation(registry);
   registerX86VectorDialectTranslation(registry);

mlir/include/mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h

+//=======- ArmSMEToLLVMIRTranslation.h - ArmSME to LLVM IR --*- C++ -*-=======//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This provides registration calls for ArmSME dialect to LLVM IR translation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_LLVMIR_DIALECT_ARMSME_ARMSMETOLLVMIRTRANSLATION_H
+#define MLIR_TARGET_LLVMIR_DIALECT_ARMSME_ARMSMETOLLVMIRTRANSLATION_H
+
+namespace mlir {
+
+class DialectRegistry;
+class MLIRContext;
+
+/// Register the ArmSME dialect and the translation from it to the LLVM IR in
+/// the given registry;
+void registerArmSMEDialectTranslation(DialectRegistry &registry);
+
+/// Register the ArmSME dialect and the translation from it in the registry
+/// associated with the given context.
+void registerArmSMEDialectTranslation(MLIRContext &context);
+
+} // namespace mlir
+
+#endif // MLIR_TARGET_LLVMIR_DIALECT_ARMSME_ARMSMETOLLVMIRTRANSLATION_H

mlir/lib/Conversion/VectorToLLVM/CMakeLists.txt

add_mlir_conversion_library(MLIRVectorToLLVM
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRArmNeonDialect
   MLIRArmSVEDialect
   MLIRArmSVETransforms
+  MLIRArmSMEDialect
+  MLIRArmSMETransforms
   MLIRAMXDialect
   MLIRAMXTransforms
   MLIRLLVMCommonConversion
   MLIRLLVMDialect
   MLIRMemRefDialect
   MLIRTargetLLVMIRExport
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorTransforms
   MLIRX86VectorDialect
   MLIRX86VectorTransforms
   )

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVMPass.cpp

 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/AMX/Transforms.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
+#include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/ArmSVE/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/X86Vector/Transforms.h"
struct LowerVectorToLLVMPass
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<LLVM::LLVMDialect>();
     registry.insert<arith::ArithDialect>();
     registry.insert<memref::MemRefDialect>();
     if (armNeon)
       registry.insert<arm_neon::ArmNeonDialect>();
     if (armSVE)
       registry.insert<arm_sve::ArmSVEDialect>();
+    if (armSME)
+      registry.insert<arm_sme::ArmSMEDialect>();
     if (amx)
       registry.insert<amx::AMXDialect>();
     if (x86Vector)
       registry.insert<x86vector::X86VectorDialect>();
   }
   void runOnOperation() override;
 };
 } // namespace
if (armNeon) {
     // LLVM-compatible operations here. So far, all operations in the dialect
     // can be translated to LLVM IR so there is no conversion necessary.
     target.addLegalDialect<arm_neon::ArmNeonDialect>();
   }
   if (armSVE) {
     configureArmSVELegalizeForExportTarget(target);
     populateArmSVELegalizeForLLVMExportPatterns(converter, patterns);
   }
+  if (armSME) {
+    configureArmSMELegalizeForExportTarget(target);
+    arm_sme::populateVectorTransferLoweringPatterns(patterns);
+  }
   if (amx) {
     configureAMXLegalizeForExportTarget(target);
     populateAMXLegalizeForLLVMExportPatterns(converter, patterns);
   }
   if (x86Vector) {
     configureX86VectorLegalizeForExportTarget(target);
     populateX86VectorLegalizeForLLVMExportPatterns(converter, patterns);
   }
 
   if (failed(
           applyPartialConversion(getOperation(), target, std::move(patterns))))
     signalPassFailure();
 }

mlir/lib/Dialect/ArmSME/CMakeLists.txt

+add_subdirectory(IR)
 add_subdirectory(Transforms)

mlir/lib/Dialect/ArmSME/IR/ArmSMEDialect.cpp

+//===- ArmSMEDialect.cpp - MLIR ArmSME dialect implementation -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the ArmSME dialect and its operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+//===----------------------------------------------------------------------===//
+// Tablegen Definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.cpp.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/ArmSME/IR/ArmSME.cpp.inc"
+
+void ArmSMEDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/ArmSME/IR/ArmSME.cpp.inc"
+      >();
+}

mlir/lib/Dialect/ArmSME/IR/CMakeLists.txt

+add_mlir_dialect_library(MLIRArmSMEDialect
+  ArmSMEDialect.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME
+
+  DEPENDS
+  MLIRArmSMEIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLLVMDialect
+  )

mlir/lib/Dialect/ArmSME/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRArmSMETransforms
   EnableArmStreaming.cpp
+  LegalizeForLLVMExport.cpp
+  LowerVectorOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME/Transforms
 
   DEPENDS
   MLIRArmSMETransformsIncGen
 
   LINK_LIBS PUBLIC
+  MLIRArmSMEDialect
   MLIRFuncDialect
+  MLIRLLVMDialect
+  MLIRVectorDialect
+  MLIRLLVMCommonConversion
+  MLIRIR
   MLIRPass
   )

mlir/lib/Dialect/ArmSME/Transforms/LegalizeForLLVMExport.cpp

+//===- LegalizeForLLVMExport.cpp - Prepare ArmSME for LLVM translation ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
+#include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+namespace {} // namespace
+
+void mlir::configureArmSMELegalizeForExportTarget(
+    LLVMConversionTarget &target) {
+  target.addLegalOp<aarch64_sme_zero>();
+}

mlir/lib/Dialect/ArmSME/Transforms/LowerVectorOps.cpp

+//===- LowerVectorOps.cpp - Lower vector ops to SME -----------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements rewrite patterns to lower vector dialect ops to ArmSME.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
+#include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+static constexpr unsigned kZeroZAMask = 255;
+
+namespace {
+/// Lower `vector.transfer_write` op to `arm_sme.intr.zero` op. Currently only
+/// supports 2d scalable vector type `vector<[16x16]xi8>` that maps to the ZA0.B
+/// SME tile. This will be extended to support more element types.
+struct TransferWriteToArmSMEZeroLowering
+    : public OpRewritePattern<vector::TransferWriteOp> {
+  TransferWriteToArmSMEZeroLowering(MLIRContext *context)
+      : OpRewritePattern<vector::TransferWriteOp>(context) {}
+
+  LogicalResult matchAndRewrite(vector::TransferWriteOp write,
+                                PatternRewriter &rewriter) const override {
+    auto vType = write.getVectorType();
+    if (vType.getRank() != 2)
+      return failure();
+    if (vType.getShape() != ArrayRef<int64_t>({16, 16}))
+      return failure();
+    if (vType.getElementType() != rewriter.getI8Type())
+      return failure();
+    if (vType.getNumScalableDims() != 2)
+      return failure();
+    auto tile = rewriter.create<arith::ConstantOp>(
+        write.getLoc(), rewriter.getI32Type(),
+        rewriter.getI32IntegerAttr(kZeroZAMask));
+    rewriter.replaceOpWithNewOp<arm_sme::aarch64_sme_zero>(write, tile);
+    return success();
+  }
+};
+} // namespace
+
+void mlir::arm_sme::populateVectorTransferLoweringPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<TransferWriteToArmSMEZeroLowering>(patterns.getContext());
+}

mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRVectorTransforms
   MLIRMemRefDialect
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorInterfaces
   MLIRVectorUtils
+  MLIRArmSMEDialect
   )

mlir/lib/Target/LLVMIR/CMakeLists.txt

add_mlir_translation_library(MLIRTargetLLVMIRExport
   )
 
 add_mlir_translation_library(MLIRToLLVMIRTranslationRegistration
   ConvertToLLVMIR.cpp
 
   LINK_LIBS PUBLIC
   MLIRArmNeonToLLVMIRTranslation
   MLIRArmSVEToLLVMIRTranslation
+  MLIRArmSMEToLLVMIRTranslation
   MLIRAMXToLLVMIRTranslation
   MLIRBuiltinToLLVMIRTranslation
   MLIRGPUToLLVMIRTranslation
   MLIRX86VectorToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRNVVMToLLVMIRTranslation
   MLIROpenACCToLLVMIRTranslation
   MLIROpenMPToLLVMIRTranslation

mlir/lib/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.cpp

+//======- ArmSMEToLLVMIRTranslation.cpp - Translate ArmSME to LLVM IR -=======//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a translation between the ArmSME dialect and LLVM IR.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSMEDialect.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Target/LLVMIR/ModuleTranslation.h"
+
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicsAArch64.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+
+namespace {
+/// Implementation of the dialect interface that converts operations belonging
+/// to the ArmSME dialect to LLVM IR.
+class ArmSMEDialectLLVMIRTranslationInterface
+    : public LLVMTranslationDialectInterface {
+public:
+  using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
+
+  /// Translates the given operation to LLVM IR using the provided IR builder
+  /// and saving the state in `moduleTranslation`.
+  LogicalResult
+  convertOperation(Operation *op, llvm::IRBuilderBase &builder,
+                   LLVM::ModuleTranslation &moduleTranslation) const final {
+    Operation &opInst = *op;
+#include "mlir/Dialect/ArmSME/IR/ArmSMEConversions.inc"
+
+    return failure();
+  }
+};
+} // namespace
+
+void mlir::registerArmSMEDialectTranslation(DialectRegistry &registry) {
+  registry.insert<arm_sme::ArmSMEDialect>();
+  registry.addExtension(+[](MLIRContext *ctx, arm_sme::ArmSMEDialect *dialect) {
+    dialect->addInterfaces<ArmSMEDialectLLVMIRTranslationInterface>();
+  });
+}
+
+void mlir::registerArmSMEDialectTranslation(MLIRContext &context) {
+  DialectRegistry registry;
+  registerArmSMEDialectTranslation(registry);
+  context.appendDialectRegistry(registry);
+}

mlir/lib/Target/LLVMIR/Dialect/ArmSME/CMakeLists.txt

+add_mlir_translation_library(MLIRArmSMEToLLVMIRTranslation
+  ArmSMEToLLVMIRTranslation.cpp
+
+  DEPENDS
+  MLIRArmSMEConversionsIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRArmSMEDialect
+  MLIRLLVMDialect
+  MLIRSupport
+  MLIRTargetLLVMIRExport
+  )

mlir/lib/Target/LLVMIR/Dialect/CMakeLists.txt

 add_subdirectory(ArmNeon)
 add_subdirectory(ArmSVE)
+add_subdirectory(ArmSME)
 add_subdirectory(AMX)
 add_subdirectory(Builtin)
 add_subdirectory(GPU)
 add_subdirectory(LLVMIR)
 add_subdirectory(NVVM)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenMP)
 add_subdirectory(ROCDL)
 add_subdirectory(X86Vector)

mlir/test/Dialect/ArmSME/vector_ops.mlir

+// RUN: mlir-opt %s -convert-vector-to-llvm="enable-arm-sme" | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: @transfer_write_2d_zero_i8
+// CHECK: %[[C255:.*]] = arith.constant 255 : i32
+// CHECK: "arm_sme.intr.zero"(%[[C255]]) : (i32) -> ()
+func.func @transfer_write_2d_zero_i8() {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %vscale = vector.vscale
+  %dim = arith.muli %c16, %vscale : index
+  %0 = memref.alloc(%dim, %dim) : memref<?x?xi8>
+  %cst = arith.constant dense<0> : vector<[16x16]xi8>
+  vector.transfer_write %cst, %0[%c0, %c0] {in_bounds = [true, true]} : vector<[16x16]xi8>, memref<?x?xi8>
+  memref.dealloc %0 : memref<?x?xi8>
+  return
+}
+
+// The following tests check the 'vector.transfer_write' -> 'arm_sme.intr.zero'
+// lowering only occurs for vector types of correct rank, shape, element size
+// and number of scalable dims.
+
+// CHECK-LABEL: @transfer_write_2d_zero__bad_type
+// CHECK: vector.transfer_write
+// CHECK-NOT: arm_sme.intr.zero
+func.func @transfer_write_2d_zero__bad_type() {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %vscale = vector.vscale
+  %dim = arith.muli %c16, %vscale : index
+  %0 = memref.alloc(%dim, %dim) : memref<?x?xi4>
+  %cst = arith.constant dense<0> : vector<[16x16]xi4>
+  vector.transfer_write %cst, %0[%c0, %c0] {in_bounds = [true, true]} : vector<[16x16]xi4>, memref<?x?xi4>
+  memref.dealloc %0 : memref<?x?xi4>
+  return
+}
+
+// CHECK-LABEL: @transfer_write_2d_zero__bad_shape
+// CHECK: vector.transfer_write
+// CHECK-NOT: arm_sme.intr.zero
+func.func @transfer_write_2d_zero__bad_shape() {
+  %c0 = arith.constant 0 : index
+  %c8 = arith.constant 8 : index
+  %vscale = vector.vscale
+  %dim = arith.muli %c8, %vscale : index
+  %0 = memref.alloc(%dim, %dim) : memref<?x?xi8>
+  %cst = arith.constant dense<0> : vector<[8x8]xi8>
+  vector.transfer_write %cst, %0[%c0, %c0] {in_bounds = [true, true]} : vector<[8x8]xi8>, memref<?x?xi8>
+  memref.dealloc %0 : memref<?x?xi8>
+  return
+}
+
+// CHECK-LABEL: @transfer_write_2d_zero__bad_rank
+// CHECK: vector.transfer_write
+// CHECK-NOT: arm_sme.intr.zero
+func.func @transfer_write_2d_zero__bad_rank() {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %vscale = vector.vscale
+  %dim = arith.muli %c16, %vscale : index
+  %0 = memref.alloc(%dim, %dim, %dim) : memref<?x?x?xi8>
+  %cst = arith.constant dense<0> : vector<[16x16x16]xi8>
+  vector.transfer_write %cst, %0[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<[16x16x16]xi8>, memref<?x?x?xi8>
+  memref.dealloc %0 : memref<?x?x?xi8>
+  return
+}
+
+// CHECK-LABEL: @transfer_write_2d_zero__bad_num_scalable_dims
+// CHECK: vector.transfer_write
+// CHECK-NOT: arm_sme.intr.zero
+func.func @transfer_write_2d_zero__bad_num_scalable_dims() {
+  %c0 = arith.constant 0 : index
+  %c16 = arith.constant 16 : index
+  %vscale = vector.vscale
+  %dim = arith.muli %c16, %vscale : index
+  %0 = memref.alloc(%dim) : memref<16x?xi8>
+  %cst = arith.constant dense<0> : vector<16x[16]xi8>
+  vector.transfer_write %cst, %0[%c0, %c0] {in_bounds = [true, true]} : vector<16x[16]xi8>, memref<16x?xi8>
+  memref.dealloc %0 : memref<16x?xi8>
+  return
+}