[mlir][ArmSME] Dialect and Intrinsic Op Definition

[WanderAway]

Summary

This patch creates the ArmSME dialect, and provides the intrinsic op definition necessary for lowering to LLVM IR.

This will cover most instructions interacting with the ZA tile register, not covering SME2 instructions.

Source: https://developer.arm.com/documentation/ddi0616/latest

[WanderAway]

Remove unneeded type constraints

[c-rhodes]

Thanks for the patch Frank this is great. I've left some comments, mostly minor but generally LGTM

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

10–11

nit: This file defines the ArmSME dialect and contains intrinsic ops to lower to LLVM IR.

44

nit: ArmSME_IntrOp is sufficient I think given overloading is optional.

56

I prefer the naming format used in the AMX dialect e.g. LLVM_x86_amx_tilezero. It makes it obvious these are LLVM intrinsics and if this were LLVM_aarch64_sme_zero it would be found when grepping on aarch64.sme.zero (alongside the LLVM intrinsic and tests). What do you think?

56

this is the default and can be removed.

57

mask

87

tile slice?

89–92

may as well add b/h/q intrinsics as well (and below for stores)

mlir/test/Target/LLVMIR/arm-sme.mlir

4

should we follow LLVM naming nxv2f64 given this is scalable?

19–54

missing tests for outer product and subtract intrinsics

20

nit: move all types to next line? E.g.

"arm_sme.intr.mopa"(%c0, %v2i1, %v2i1, %v2f64, %v2f64) :
  (i32, vector<[2]xi1>, vector<[2]xi1>, vector<[2]xf64>, vector<[2]xf64>) -> ()

[awarzynski]

Hi Frank, thanks for working on this! Also LGTM modulo some minor comments.

Does this cover all SME intrinsics? In particular, I don't see any SME 2. That's totally fine, but would be good to clarify the scope in the summary.

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

34

This spec doesn't cover all the SME instructions. It would be good to add a link to https://developer.arm.com/documentation/ddi0602/2023-03/SME-Instructions?lang=en as well.

57

The tiles in SME aren't registers. And also aren't real tiles :) I've seen people using the term "virtual tile", which IMHO makes a lot of sense.

65

Why aren't the input arguments limited to "supported" vector sizes like predicates?

def Predicate : ScalableVectorOfLengthAndType<[16, 8, 4, 2], [I1]>;

It's something that we can change later. I'm just curious about the rationale.

Also, what happens when invalid predicate/input size is used?

86

99–100

mlir/test/Target/LLVMIR/arm-sme.mlir

4

should we follow LLVM naming nxv2f64 given this is scalable?

+1

Also, could you split this into multiple functions? Perhaps a function per intrinsic group. I'd also further split into integer vs FP.

[WanderAway]

Addressed comments

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

10–11

Oops, fixed.

44

Makes sense.

56

Makes sense. Changed.

57

mask

Oops too much copy pasting

I've seen people using the term "virtual tile"

Sounds good.

65

Why aren't the input arguments limited to "supported" vector sizes like predicates?

Honestly I can't remember, there was a reason at some point, but right now I think it makes sense to add a constraint for it.

what happens when invalid predicate/input size is used?

If the predicate and vector sizes mismatch, then it violates the intrinsic signature, and emit a error during LLVM verifier. However there are currently no checks in place to make sure that the vector unit length is 128b (i.e. <vscalex2xf32> is a valid operand type). Right now there are no built-in constraints to check for something like this (although we could add one), and I suspect it will fail during ISel.

[c-rhodes]

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

32–33

Please could we keep the previous link to the spec as well, that contains general information about the SME architecture like ZA layout etc.

46

46–47

nit: fix indentation

50

nit: add whitespace around concat operator

51

and here

110–114

nit: please could we order from smallest to largest types (i.e define b first and q last) to be consistent with other parts of the codebase

mlir/test/Target/LLVMIR/arm-sme.mlir

1

add -split-input-file and separate each test with // ----- to run in parallel?

3

missing CHECK-LABELs

[WanderAway]

Fixed more comments.

[c-rhodes]

Just one last comment but otherwise LGTM, thanks!

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

16–21

I think we can drop these includes?

[WanderAway]

Removed unnecessary header files.

[awarzynski]

LGTM, thank you for addressing my comments!

Feel free to ignore my final nit.

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

43–45

[nit] I believe that we will be using these for Ops other than outer product as well. I would use a more generic name instead (e.g. SVEPredicate/SVEVector or SSVEPredicate/SSVEVector to be more precise).

[WanderAway]

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

43–45

I believe that we will be using these for Ops other than outer product as well. I would use a more generic name instead (e.g. SVEPredicate/SVEVector or SSVEPredicate/SSVEVector to be more precise).

The only issue I see here is that different instructions may have different types supported (such as load/stores and MOPs defined here), so it may be too much of a generalization to rename these to SVE* or SSVE* at the current moment. I think it may be a good idea to rename these to something appropriate whenever we add more to this file.

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

+add_subdirectory(IR)
 add_subdirectory(Transforms)

mlir/include/mlir/Dialect/ArmSME/IR/ArmSME.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_IR_ARMSME_H
+#define MLIR_DIALECT_ARMSME_IR_ARMSME_H
+
+#include "mlir/Bytecode/BytecodeOpInterface.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_IR_ARMSME_H

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 ArmSME dialect and contains intrinsic ops to lower to
+// LLVM IR.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ARMSME_OPS
+#define ARMSME_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+
+//===----------------------------------------------------------------------===//
+// ArmSME dialect definition
+//===----------------------------------------------------------------------===//
+
+def ArmSME_Dialect : Dialect {
+  let name = "arm_sme";
+  let cppNamespace = "::mlir::arm_sme";
+  let summary = "Basic dialect to target Arm SME architectures";
+  let description = [{
+    This dialect contains the definitions necessary to target Arm SME
+    scalable matrix operations.
+
+    Source:
+    https://developer.arm.com/documentation/ddi0602/2023-03/SME-Instructions
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// ArmSME Intrinsic op definitions
+//===----------------------------------------------------------------------===//
+
+def MOPPredicate : ScalableVectorOfLengthAndType<[16, 8, 4, 2], [I1]>;
+def MOPVector : ScalableVectorOfLengthAndType<[16, 8, 4, 2],
+                                              [I8, I16, BF16, F16, F32, F64]>;
+def LDSTPredicate : ScalableVectorOfLengthAndType<[16, 8, 4, 2, 1], [I1]>;
+
+class ArmSME_IntrOp<string mnemonic, list<int> overloadOperands = [],
+                              list<Trait> traits = []>
+    : LLVM_IntrOpBase<
+          /*Dialect dialect=*/ArmSME_Dialect,
+          /*string opName=*/"intr." #mnemonic,
+          /*string enumName=*/"aarch64_sme_" #!subst(".", "_", mnemonic),
+          /*list<int> overloadedResults=*/[],
+          /*list<int> overloadedOperands=*/overloadOperands,
+          /*list<Trait> traits=*/traits,
+          /*int numResults=*/0>;
+
+// Zero
+def LLVM_aarch64_sme_zero : ArmSME_IntrOp<"zero">,
+                            Arguments<(ins Arg<I32, "Tile mask">)>;
+
+// MOP's
+class ArmSME_IntrMopOverloadedOp<string mnemonic>
+    : ArmSME_IntrOp<mnemonic, [4]>,
+      Arguments<(ins Arg<I32, "Virtual tile ID">,
+                 Arg<MOPPredicate, "LHS predicate">,
+                 Arg<MOPPredicate, "RHS predicate">,
+                 Arg<MOPVector, "LHS vector operand">,
+                 Arg<MOPVector, "RHS vector operand">)>;
+
+def LLVM_aarch64_sme_mopa : ArmSME_IntrMopOverloadedOp<"mopa">;
+def LLVM_aarch64_sme_mops : ArmSME_IntrMopOverloadedOp<"mops">;
+def LLVM_aarch64_sme_mopa_wide : ArmSME_IntrMopOverloadedOp<"mopa.wide">;
+def LLVM_aarch64_sme_mops_wide : ArmSME_IntrMopOverloadedOp<"mops.wide">;
+def LLVM_aarch64_sme_smopa_wide : ArmSME_IntrMopOverloadedOp<"smopa.wide">;
+def LLVM_aarch64_sme_smops_wide : ArmSME_IntrMopOverloadedOp<"smops.wide">;
+def LLVM_aarch64_sme_umopa_wide : ArmSME_IntrMopOverloadedOp<"umopa.wide">;
+def LLVM_aarch64_sme_umops_wide : ArmSME_IntrMopOverloadedOp<"umops.wide">;
+def LLVM_aarch64_sme_sumopa_wide : ArmSME_IntrMopOverloadedOp<"sumopa.wide">;
+def LLVM_aarch64_sme_sumops_wide : ArmSME_IntrMopOverloadedOp<"sumops.wide">;
+def LLVM_aarch64_sme_usmopa_wide : ArmSME_IntrMopOverloadedOp<"usmopa.wide">;
+def LLVM_aarch64_sme_usmops_wide : ArmSME_IntrMopOverloadedOp<"usmops.wide">;
+
+// Loads
+class ArmSME_IntrLoadOp<string mnemonic>
+    : ArmSME_IntrOp<mnemonic>,
+      Arguments<(ins Arg<LDSTPredicate, "Vector predicate">,
+                 Arg<LLVM_AnyPointer, "Load address", [MemRead]>,
+                 Arg<I32, "Virtual tile ID">,
+                 Arg<I32, "Tile slice">)>;
+
+def LLVM_aarch64_sme_ld1q_horiz : ArmSME_IntrLoadOp<"ld1q.horiz">;
+def LLVM_aarch64_sme_ld1w_horiz : ArmSME_IntrLoadOp<"ld1w.horiz">;
+def LLVM_aarch64_sme_ld1d_horiz : ArmSME_IntrLoadOp<"ld1d.horiz">;
+def LLVM_aarch64_sme_ld1h_horiz : ArmSME_IntrLoadOp<"ld1h.horiz">;
+def LLVM_aarch64_sme_ld1b_horiz : ArmSME_IntrLoadOp<"ld1b.horiz">;
+def LLVM_aarch64_sme_ld1q_vert : ArmSME_IntrLoadOp<"ld1q.vert">;
+def LLVM_aarch64_sme_ld1w_vert : ArmSME_IntrLoadOp<"ld1w.vert">;
+def LLVM_aarch64_sme_ld1d_vert : ArmSME_IntrLoadOp<"ld1d.vert">;
+def LLVM_aarch64_sme_ld1h_vert : ArmSME_IntrLoadOp<"ld1h.vert">;
+def LLVM_aarch64_sme_ld1b_vert : ArmSME_IntrLoadOp<"ld1b.vert">;
+
+// Stores
+class ArmSME_IntrStoreOp<string mnemonic>
+    : ArmSME_IntrOp<mnemonic>,
+      Arguments<(ins Arg<LDSTPredicate, "Vector predicate">,
+                 Arg<LLVM_AnyPointer, "Store address", [MemWrite]>,
+                 Arg<I32, "Virtual tile ID">,
+                 Arg<I32, "Tile slice">)>;
+
+def LLVM_aarch64_sme_st1q_horiz : ArmSME_IntrStoreOp<"st1q.horiz">;
+def LLVM_aarch64_sme_st1w_horiz : ArmSME_IntrStoreOp<"st1w.horiz">;
+def LLVM_aarch64_sme_st1d_horiz : ArmSME_IntrStoreOp<"st1d.horiz">;
+def LLVM_aarch64_sme_st1h_horiz : ArmSME_IntrStoreOp<"st1h.horiz">;
+def LLVM_aarch64_sme_st1b_horiz : ArmSME_IntrStoreOp<"st1b.horiz">;
+def LLVM_aarch64_sme_st1q_vert : ArmSME_IntrStoreOp<"st1q.vert">;
+def LLVM_aarch64_sme_st1w_vert : ArmSME_IntrStoreOp<"st1w.vert">;
+def LLVM_aarch64_sme_st1d_vert : ArmSME_IntrStoreOp<"st1d.vert">;
+def LLVM_aarch64_sme_st1h_vert : ArmSME_IntrStoreOp<"st1h.vert">;
+def LLVM_aarch64_sme_st1b_vert : ArmSME_IntrStoreOp<"st1b.vert">;
+
+#endif // ARMSME_OPS

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/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/ArmSME.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,
                   ml_program::MLProgramDialect,
                   nvgpu::NVGPUDialect,
                   scf::SCFDialect,
                   omp::OpenMPDialect,
                   pdl::PDLDialect,
                   pdl_interp::PDLInterpDialect,
                   quant::QuantizationDialect,
                   spirv::SPIRVDialect,
+                  arm_sme::ArmSMEDialect,
                   arm_sve::ArmSVEDialect,
                   vector::VectorDialect,
                   NVVM::NVVMDialect,
                   ROCDL::ROCDLDialect,
                   shape::ShapeDialect,
                   sparse_tensor::SparseTensorDialect,
                   tensor::TensorDialect,
                   transform::TransformDialect,

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);
+  registerArmSMEDialectTranslation(registry);
   registerArmSVEDialectTranslation(registry);
   registerBuiltinDialectTranslation(registry);
   registerGPUDialectTranslation(registry);
   registerLLVMDialectTranslation(registry);
   registerNVVMDialectTranslation(registry);
   registerOpenACCDialectTranslation(registry);
   registerOpenMPDialectTranslation(registry);
   registerROCDLDialectTranslation(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/Dialect/ArmSME/CMakeLists.txt

+add_subdirectory(IR)
 add_subdirectory(Transforms)

mlir/lib/Dialect/ArmSME/IR/ArmSME.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/ArmSME.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/TypeSwitch.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"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/ArmSME/IR/ArmSMETypes.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
+  ArmSME.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME
+
+  DEPENDS
+  MLIRArmSMEIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLLVMDialect
+  MLIRSideEffectInterfaces
+)

mlir/lib/Target/LLVMIR/CMakeLists.txt

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

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/ArmSME.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(ArmSME)
 add_subdirectory(ArmSVE)
 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/Target/LLVMIR/arm-sme.mlir

+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @arm_sme_zero() {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.zero(i32 0)
+  "arm_sme.intr.zero"(%c0) : (i32) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_fmopa(%nxv2f64 : vector<[2]xf64>,
+                         %nxv4f32 : vector<[4]xf32>,
+                         %nxv8f16 : vector<[8]xf16>,
+                         %nxv8bf16: vector<[8]xbf16>,
+                         %nxv2i1  : vector<[2]xi1>,
+                         %nxv4i1  : vector<[4]xi1>,
+                         %nxv8i1  : vector<[8]xi1>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.mopa.nxv2f64
+  "arm_sme.intr.mopa"(%c0, %nxv2i1, %nxv2i1, %nxv2f64, %nxv2f64) :
+    (i32, vector<[2]xi1>, vector<[2]xi1>, vector<[2]xf64>, vector<[2]xf64>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mopa.nxv4f32
+  "arm_sme.intr.mopa"(%c0, %nxv4i1, %nxv4i1, %nxv4f32, %nxv4f32) :
+    (i32, vector<[4]xi1>, vector<[4]xi1>, vector<[4]xf32>, vector<[4]xf32>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mopa.wide.nxv8f16
+  "arm_sme.intr.mopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8f16, %nxv8f16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xf16>, vector<[8]xf16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mopa.wide.nxv8bf16
+  "arm_sme.intr.mopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8bf16, %nxv8bf16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xbf16>, vector<[8]xbf16>) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_imopa(%nxv8i16 : vector<[8]xi16>,
+                         %nxv16i8 : vector<[16]xi8>,
+                         %nxv8i1  : vector<[8]xi1>,
+                         %nxv16i1 : vector<[16]xi1>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.smopa.wide.nxv8i16
+  "arm_sme.intr.smopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.umopa.wide.nxv8i16
+  "arm_sme.intr.umopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.sumopa.wide.nxv8i16
+  "arm_sme.intr.sumopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.usmopa.wide.nxv8i16
+  "arm_sme.intr.usmopa.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.smopa.wide.nxv16i8
+  "arm_sme.intr.smopa.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.umopa.wide.nxv16i8
+  "arm_sme.intr.umopa.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.sumopa.wide.nxv16i8
+  "arm_sme.intr.sumopa.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.usmopa.wide.nxv16i8
+  "arm_sme.intr.usmopa.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_fmops(%nxv2f64 : vector<[2]xf64>,
+                         %nxv4f32 : vector<[4]xf32>,
+                         %nxv8f16 : vector<[8]xf16>,
+                         %nxv8bf16: vector<[8]xbf16>,
+                         %nxv2i1  : vector<[2]xi1>,
+                         %nxv4i1  : vector<[4]xi1>,
+                         %nxv8i1  : vector<[8]xi1>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.mops.nxv2f64
+  "arm_sme.intr.mops"(%c0, %nxv2i1, %nxv2i1, %nxv2f64, %nxv2f64) :
+    (i32, vector<[2]xi1>, vector<[2]xi1>, vector<[2]xf64>, vector<[2]xf64>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mops.nxv4f32
+  "arm_sme.intr.mops"(%c0, %nxv4i1, %nxv4i1, %nxv4f32, %nxv4f32) :
+    (i32, vector<[4]xi1>, vector<[4]xi1>, vector<[4]xf32>, vector<[4]xf32>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mops.wide.nxv8f16
+  "arm_sme.intr.mops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8f16, %nxv8f16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xf16>, vector<[8]xf16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.mops.wide.nxv8bf16
+  "arm_sme.intr.mops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8bf16, %nxv8bf16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xbf16>, vector<[8]xbf16>) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_imops(%nxv8i16 : vector<[8]xi16>,
+                         %nxv16i8 : vector<[16]xi8>,
+                         %nxv8i1  : vector<[8]xi1>,
+                         %nxv16i1 : vector<[16]xi1>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.smops.wide.nxv8i16
+  "arm_sme.intr.smops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.umops.wide.nxv8i16
+  "arm_sme.intr.umops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.sumops.wide.nxv8i16
+  "arm_sme.intr.sumops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.usmops.wide.nxv8i16
+  "arm_sme.intr.usmops.wide"(%c0, %nxv8i1, %nxv8i1, %nxv8i16, %nxv8i16) :
+    (i32, vector<[8]xi1>, vector<[8]xi1>, vector<[8]xi16>, vector<[8]xi16>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.smops.wide.nxv16i8
+  "arm_sme.intr.smops.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.umops.wide.nxv16i8
+  "arm_sme.intr.umops.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.sumops.wide.nxv16i8
+  "arm_sme.intr.sumops.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  // CHECK: call void @llvm.aarch64.sme.usmops.wide.nxv16i8
+  "arm_sme.intr.usmops.wide"(%c0, %nxv16i1, %nxv16i1, %nxv16i8, %nxv16i8) :
+    (i32, vector<[16]xi1>, vector<[16]xi1>, vector<[16]xi8>, vector<[16]xi8>) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_load(%nxv1i1  : vector<[1]xi1>,
+                        %nxv2i1  : vector<[2]xi1>,
+                        %nxv4i1  : vector<[4]xi1>,
+                        %nxv8i1  : vector<[8]xi1>,
+                        %nxv16i1 : vector<[16]xi1>,
+                        %p8      : !llvm.ptr<i8>,
+                        %p16     : !llvm.ptr<i16>,
+                        %p32     : !llvm.ptr<i32>,
+                        %p64     : !llvm.ptr<i64>,
+                        %p128    : !llvm.ptr<i128>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.ld1q.horiz
+  "arm_sme.intr.ld1q.horiz"(%nxv1i1, %p128, %c0, %c0) :
+              (vector<[1]xi1>, !llvm.ptr<i128>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1d.horiz
+  "arm_sme.intr.ld1d.horiz"(%nxv2i1, %p64, %c0, %c0) :
+              (vector<[2]xi1>, !llvm.ptr<i64>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1w.horiz
+  "arm_sme.intr.ld1w.horiz"(%nxv4i1, %p32, %c0, %c0) :
+              (vector<[4]xi1>, !llvm.ptr<i32>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1h.horiz
+  "arm_sme.intr.ld1h.horiz"(%nxv8i1, %p16, %c0, %c0) :
+              (vector<[8]xi1>, !llvm.ptr<i16>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1b.horiz
+  "arm_sme.intr.ld1b.horiz"(%nxv16i1, %p8, %c0, %c0) :
+              (vector<[16]xi1>, !llvm.ptr<i8>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1q.vert
+  "arm_sme.intr.ld1q.vert"(%nxv1i1, %p128, %c0, %c0) :
+              (vector<[1]xi1>, !llvm.ptr<i128>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1d.vert
+  "arm_sme.intr.ld1d.vert"(%nxv2i1, %p64, %c0, %c0) :
+              (vector<[2]xi1>, !llvm.ptr<i64>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1w.vert
+  "arm_sme.intr.ld1w.vert"(%nxv4i1, %p32, %c0, %c0) :
+              (vector<[4]xi1>, !llvm.ptr<i32>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1h.vert
+  "arm_sme.intr.ld1h.vert"(%nxv8i1, %p16, %c0, %c0) :
+              (vector<[8]xi1>, !llvm.ptr<i16>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.ld1b.vert
+  "arm_sme.intr.ld1b.vert"(%nxv16i1, %p8, %c0, %c0) :
+              (vector<[16]xi1>, !llvm.ptr<i8>, i32, i32) -> ()
+  llvm.return
+}
+
+llvm.func @arm_sme_store(%nxv1i1  : vector<[1]xi1>,
+                         %nxv2i1  : vector<[2]xi1>,
+                         %nxv4i1  : vector<[4]xi1>,
+                         %nxv8i1  : vector<[8]xi1>,
+                         %nxv16i1 : vector<[16]xi1>,
+                         %p8      : !llvm.ptr<i8>,
+                         %p16     : !llvm.ptr<i16>,
+                         %p32     : !llvm.ptr<i32>,
+                         %p64     : !llvm.ptr<i64>,
+                         %p128    : !llvm.ptr<i128>) {
+  %c0 = llvm.mlir.constant(0 : index) : i32
+  // CHECK: call void @llvm.aarch64.sme.st1q.horiz
+  "arm_sme.intr.st1q.horiz"(%nxv1i1, %p128, %c0, %c0) :
+              (vector<[1]xi1>, !llvm.ptr<i128>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1d.horiz
+  "arm_sme.intr.st1d.horiz"(%nxv2i1, %p64, %c0, %c0) :
+              (vector<[2]xi1>, !llvm.ptr<i64>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1w.horiz
+  "arm_sme.intr.st1w.horiz"(%nxv4i1, %p32, %c0, %c0) :
+              (vector<[4]xi1>, !llvm.ptr<i32>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1h.horiz
+  "arm_sme.intr.st1h.horiz"(%nxv8i1, %p16, %c0, %c0) :
+              (vector<[8]xi1>, !llvm.ptr<i16>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1b.horiz
+  "arm_sme.intr.st1b.horiz"(%nxv16i1, %p8, %c0, %c0) :
+              (vector<[16]xi1>, !llvm.ptr<i8>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1q.vert
+  "arm_sme.intr.st1q.vert"(%nxv1i1, %p128, %c0, %c0) :
+              (vector<[1]xi1>, !llvm.ptr<i128>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1d.vert
+  "arm_sme.intr.st1d.vert"(%nxv2i1, %p64, %c0, %c0) :
+              (vector<[2]xi1>, !llvm.ptr<i64>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1w.vert
+  "arm_sme.intr.st1w.vert"(%nxv4i1, %p32, %c0, %c0) :
+              (vector<[4]xi1>, !llvm.ptr<i32>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1h.vert
+  "arm_sme.intr.st1h.vert"(%nxv8i1, %p16, %c0, %c0) :
+              (vector<[8]xi1>, !llvm.ptr<i16>, i32, i32) -> ()
+  // CHECK: call void @llvm.aarch64.sme.st1b.vert
+  "arm_sme.intr.st1b.vert"(%nxv16i1, %p8, %c0, %c0) :
+              (vector<[16]xi1>, !llvm.ptr<i8>, i32, i32) -> ()
+  llvm.return
+}

mlir/test/mlir-opt/commandline.mlir

 // RUN: echo "" | mlir-opt --show-dialects | FileCheck %s
 // CHECK: Available Dialects:
 // CHECK-SAME: acc
 // CHECK-SAME: affine
 // CHECK-SAME: amdgpu
 // CHECK-SAME: amx
 // CHECK-SAME: arith
 // CHECK-SAME: arm_neon
+// CHECK-SAME: arm_sme
 // CHECK-SAME: arm_sve
 // CHECK-SAME: async
 // CHECK-SAME: bufferization
 // CHECK-SAME: builtin
 // CHECK-SAME: cf
 // CHECK-SAME: complex
 // CHECK-SAME: dlti
 // CHECK-SAME: emitc