[MLIR][ArmSME] Initial SME Dialect Implementation

[WanderAway]

Summary

This patch is the initial implementation of the ArmSME dialect, and contains the dialect itself, along with the following ops:

• zero
• load.tile
• store.tile
• mopa
• mops

Along with their corresponding LLVM intrinsics.

The next step would be to implement the legalization transformation pass so that it may be lowered to LLVM.

PATCH [1/?] for RFC: https://discourse.llvm.org/t/rfc-creating-a-armsme-dialect/67208

[WanderAway]

Update copy-paste description of the mopa op.

[tschuett]

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

26

Nit: you could use nested namespaces:
namespace mlir::arm_sme {
}

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

137

Nit: there

171

Nit: a or the?

[dcaballe]

Thanks, Frank!

Along with their corresponding LLVM intrinsics.

Where are these intrinsics defined? I can't find them.

[WanderAway]

In D152080#4403826, @dcaballe wrote:

Thanks, Frank!

Along with their corresponding LLVM intrinsics.

Where are these intrinsics defined? I can't find them.

They're in the tablegen file, with names being *IntrOp, inheriting from LLVM_IntrOpBase.

[c-rhodes]

Hi Frank, thanks for posting. As mentioned on the RFC there's a few of us working in this area now and there's a discussion scheduled for an upcoming MLIR ODM meeting but according to the agenda Public MLIR Open Meeting Agenda - Google Docs currently it's not until June 29th. As you've probably seen from upstream activity I've been focusing on Streaming SVE support in MLIR / IREE but more recently I've started looking at SME proper. To begin with focusing on a path from linalg.fill to SME as a simple example. In the spirit of openness I've shared a patch D152508 that adds a lowering to the SME zero instruction in the simplest case.

There's a lot going on in this patch, personally I'd find a few smaller patches that demonstrate the approach end-to-end for a simple example more useful, particularly in the context of the tile resource and how that could work.

[c-rhodes]

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

17–18

this has been mentioned a couple of times now but it would be good to clarify what it really means. The attribute exists in two places:

1. The EnableArmStreaming.cpp pass as a static constexpr char that adds it to func.func ops.
2. llvm.func ops which propagate it to the backend, this must stay as is.

I guess the idea here is 1. could be moved to a function as part of extraClassDeclaration under dialect def (example in GPU dialect), e.g.

def ArmSME_Dialect : Dialect {
  ...
  let extraClassDeclaration = [{
    /// Get the name of the attribute used to annotate functions to enable Arm
    /// streaming mode.
    static StringRef getArmStreamingAttrName() { return "arm_streaming"; }
    /// Get the name of the attribute used to annotate functions to enable Arm
    /// streaming mode locally.
    static StringRef getArmLocallyStreamingAttrName() {
      return "arm_locally_streaming";
    }
  }];
}

46–47

I don't think so, AFAIK there's no lowerings (from vector or linalg) to the custom ops in the SVE dialect anyway so there's no practical reason for this (yet).

[dcaballe]

Where are these intrinsics defined? I can't find them.

They're in the tablegen file, with names being *IntrOp, inheriting from LLVM_IntrOpBase.

Do you think you could create a separate patch only with the LLVM intrinsics? I think we could land that already!

[WanderAway]

In D152080#4416389, @dcaballe wrote:

Where are these intrinsics defined? I can't find them.

They're in the tablegen file, with names being *IntrOp, inheriting from LLVM_IntrOpBase.

Do you think you could create a separate patch only with the LLVM intrinsics? I think we could land that already!

Sure, I'll raise a separate patch when I get the time.

[awarzynski]

In D152080#4417351, @WanderAway wrote:

In D152080#4416389, @dcaballe wrote:

Where are these intrinsics defined? I can't find them.

They're in the tablegen file, with names being *IntrOp, inheriting from LLVM_IntrOpBase.

Do you think you could create a separate patch only with the LLVM intrinsics? I think we could land that already!

Sure, I'll raise a separate patch when I get the time.

+1

Thinking of the upcoming MLIR ODM, you could prepare some slides on this dialect and the approach that you are taking. Only if you are interested - I have no intention to delegate extra work to anyone :) I will prepare an introduction into SME and the work done so far (i.e. "status update").

Looking forward to discussing this :)

[WanderAway]

you could prepare some slides on this dialect and the approach that you are taking.

Yup that's the plan. Looking forward to it.

[WanderAway]

Here's the separate patch for just the intrinsic ops: https://reviews.llvm.org/D152878 @dcaballe @awarzynski

[WanderAway]

Dialect on hold until use case becomes clearer.

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"
+
+namespace mlir {
+namespace arm_sme {
+class SMETile : public SideEffects::Resource::Base<SMETile> {
+public:
+  StringRef getName() final { return "SMETile"; }
+};
+} // namespace arm_sme
+} // namespace mlir
+
+#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 basic operations for the ArmSME dialect.
+// Currently contains the following ops:
+// * zero       - allocates a tile initialized with zeros
+// * mopa/mops  - matrix outer product accumulate/subtract
+// * load.tile  - allocates a tile initialized with data from memory
+// * store.tile - stores tile to memory while deallocating it.
+//
+// There are a few outstanding TODOs:
+// * Move the streaming mode function attribute into here
+// * Think a bit more about how to share more functionality with vector dialect
+// * Save/restore functionality for cross-function tile usage
+// * Implement a more robust tile allocation scheme
+// * Implement lowering of more general vector ops on SME tiles, e.g. extract
+//
+//===----------------------------------------------------------------------===//
+
+#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 specific Arm SME
+    scalable matrix operations.
+
+    Source:
+    https://developer.arm.com/documentation/ddi0616/aa
+  }];
+  // FIXME: Is this necessary?
+  let dependentDialects = ["arm_sve::ArmSVEDialect"];
+}
+
+//===----------------------------------------------------------------------===//
+// ArmSME misc definitions
+//===----------------------------------------------------------------------===//
+
+// Effects
+def TileResource : Resource<"SMETile">;
+def TileAlloc : MemAlloc<TileResource>;
+def TileWrite : MemWrite<TileResource>;
+def TileRead : MemRead<TileResource>;
+def TileFree : MemFree<TileResource>;
+
+def Predicate : ScalableVectorOfLengthAndType<[16, 8, 4, 2], [I1]>;
+def SMEVector : ScalableVectorOfLengthAndType<
+                    [16, 8, 4, 2], [SI8, SI16, UI8, UI16, BF16, F16, F32, F64]>;
+def SMETile : AllOfType<[
+              ScalableVectorOf<[SI8, SI16, SI32, SI64, UI8, UI16, UI32, UI64,
+                                BF16, F16, F32, F64]>,
+              ScalableVectorOfRank<[2]>,
+              ScalableVectorOfLength<[256, 64, 16, 4]>]>;
+
+//===----------------------------------------------------------------------===//
+// ArmSME op definitions
+//===----------------------------------------------------------------------===//
+
+class ArmSME_Op<string mnemonic, list<Trait> traits = []> :
+  Op<ArmSME_Dialect, mnemonic, traits> {}
+
+def ZeroOp : ArmSME_Op<"zero"> {
+  let summary = "Returns a SME tile with its contents zeroed.";
+  let description = [{
+    Allocates a tile in the background. This will impose restrictions on the
+    returned vector (for now):
+    * The vector can only be read from before being written to. After which
+      reading from that vector will become invalid.
+    * Once overwritten (e.g. mopa), the new returned tile will represent the
+      same physical tile.
+  }];
+  let results = (outs Res<SMETile, "tile", [TileAlloc, TileWrite]>: $tile);
+  let assemblyFormat = "attr-dict `:` type($tile)";
+}
+
+class MOPOpBase<string mnemonic, bit accumulate>
+    : ArmSME_Op<mnemonic, [AllTypesMatch<["tile", "result"]>,
+                           AllShapesMatch<["lhsPred", "lhs"]>,
+                           AllShapesMatch<["rhsPred", "rhs"]>,
+                           AllElementTypesMatch<["lhs", "rhs"]>,
+                           AllElementCountsMatch<["lhs", "rhs"]>]> {
+  let arguments = (ins
+          Arg<SMETile, "tile", [TileRead]>:$tile,
+          Predicate:$lhsPred,
+          Predicate:$rhsPred,
+          SMEVector:$lhs,
+          SMEVector:$rhs
+  );
+  let results = (outs Res<SMETile, "", [TileWrite]>:$result);
+  let extraClassDeclaration = [{
+    bool isAccumulate() { return }] # accumulate # [{;
+    }
+    bool isSubtract() { return }] # !not(accumulate) # [{; }
+    bool isWidening() {
+      auto elTy = cast<VectorType>(this->getLhs().getType()).getElementType();
+      if (elTy.isF32() || elTy.isF64())
+        return false;
+      else
+        return true;
+    }
+  }];
+  let assemblyFormat =[{ $tile`,` $lhsPred`,` $rhsPred`,` $lhs`,` $rhs attr-dict
+                         `:` custom<MOPOperandTypes>(type($tile),
+                             type($lhsPred), type($rhsPred), type($lhs),
+                             type($rhs)) }];
+  let hasVerifier = 1;
+}
+
+def MopaOp : MOPOpBase<"mopa", /*accumulate=*/true> {
+  let summary = "Vector-vector outer product and accumulate op";
+  let description = [{
+    MOPA: Outer product accumulate.
+
+    This function maps to the *MOPA instructions, it takes scalable vector
+    operands which will be used to compute the outer product matrix. Two
+    masking predicate operands for each of the floating point operands will also
+    be provided, such that elements marked inactive by the predicate will not
+    update the corresponding row/column in the result matrix tile, specified by
+    the attribute.
+
+    Theere are two variations of MOPA instructions - widening and non-widening.
+
+    Non-widening MOPAs will take a 1D vector of f32 or f64 as input and
+    accumulate into 32b and 64b tiles respectively (za*s and za*d).
+
+    Widening MOPAs will pack two f16/bf16 or four (signed or unsigned) i8
+    elements into a single 32b lane of the vector and accumulate into 32b tiles
+    (za*s); Or it will pack four (signed or unsigned) i16 elements into a 64b
+    lane and accumulate into 64b tiles (za*d). Hence widening MOPAs will take
+    2D scalable vectors as input, i.e. `<[4x2]xf16>, <[2x4]xsi16>, <[4x4]xsi8>`
+
+    Example: Assume `vscale == 2`, `%lhs = %rhs = <1, 2, 3, 4> : <[2]xfp64>`,
+    `%lhsPred = %rhsPred = <true, true, false, true>`, then:
+    ```
+    arm_sme.zero za0d
+    arm_sme.fmopa za0d, %lhsPred, %rhsPred, %lhs, %rhs
+                  : vector<[2]xi1>, vector<[2]xf64>
+    ```
+
+    Would result in za0d containing:
+    ```
+    1  2  0  4
+    2  4  0  8
+    0  0  0  0
+    4  8  0 16
+    ```
+  }];
+}
+
+def MopsOp : MOPOpBase<"mops", /*accumulate=*/false> {
+  let summary = "Vector-vector outer product and subtract op";
+  let description = [{
+    MOPS: Outer product subtract.
+
+    Similar to the MOPA instruction, except a the outer product is subtracted
+    from the tile.
+  }];
+}
+
+def LoadTile : ArmSME_Op<"load.tile"> {
+  let summary = "Loads a 2D vector from memory into a SME tile";
+  let description = [{
+    This op will allocate a tile (similar to `arm_sme.zero`) in addition to
+    loading a tile from memory. This means the vector returned from this op
+    requires the same constraint as `arm_sme.zero`.
+  }];
+  let arguments = (ins Arg<AnyMemRef, "reference to load from", [MemRead]>:$base,
+                       Variadic<Index>:$indices);
+  let results = (outs Res<SMETile, "tile", [TileAlloc, TileWrite]>:$tile);
+  let assemblyFormat = "$base `[` $indices `]` attr-dict "
+                       "`:` type($base) `,` type($tile)";
+}
+
+def StoreTile : ArmSME_Op<"store.tile"> {
+  let summary = "Stores a 2D vector into memory from a SME tile";
+  let description = [{
+    This op will deallocate a tile (is this necessary?) in addition to
+    storing the tile.
+  }];
+  let arguments = (ins Arg<SMETile, "tile", [TileFree, TileRead]>:$tile,
+                       Arg<AnyMemRef, "reference to load from", [MemRead]>:$base,
+                       Variadic<Index>:$indices);
+  let assemblyFormat = "$tile `,` $base `[` $indices `]` attr-dict "
+                       "`:` type($base) `,` type($tile)";
+}
+
+//===----------------------------------------------------------------------===//
+// ArmSME Intrinsic op definitions
+//===----------------------------------------------------------------------===//
+
+class ArmSME_IntrOverloadedOp<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 ZeroIntrOp : ArmSME_IntrOverloadedOp<"zero", []>,
+                 Arguments<(ins Arg<I32, "Tile register ID", [TileAlloc]>)>;
+
+// MOP's
+class ArmSME_IntrMopOverloadedOp<string mnemonic>
+    : ArmSME_IntrOverloadedOp<mnemonic, [4]>,
+      Arguments<(ins Arg<I32, "Tile register ID", [TileWrite, TileRead]>,
+                 Arg<Predicate, "LHS predicate">,
+                 Arg<Predicate, "RHS predicate">,
+                 Arg<AnyScalableVector, "LHS vector operand">,
+                 Arg<AnyScalableVector, "RHS vector operand">)>;
+
+def FmopaIntrOp : ArmSME_IntrMopOverloadedOp<"mopa">;
+def FmopsIntrOp : ArmSME_IntrMopOverloadedOp<"mops">;
+def FmopaWidenIntrOp : ArmSME_IntrMopOverloadedOp<"mopa.wide">;
+def FmopsWidenIntrOp : ArmSME_IntrMopOverloadedOp<"mops.wide">;
+def SmopaIntrOp : ArmSME_IntrMopOverloadedOp<"smopa.wide">;
+def SmopsIntrOp : ArmSME_IntrMopOverloadedOp<"smops.wide">;
+def UmopaIntrOp : ArmSME_IntrMopOverloadedOp<"umopa.wide">;
+def UmopsIntrOp : ArmSME_IntrMopOverloadedOp<"umops.wide">;
+def SUmopaIntrOp : ArmSME_IntrMopOverloadedOp<"sumopa.wide">;
+def SUmopsIntrOp : ArmSME_IntrMopOverloadedOp<"sumops.wide">;
+def USmopaIntrOp : ArmSME_IntrMopOverloadedOp<"usmopa.wide">;
+def USmopsIntrOp : ArmSME_IntrMopOverloadedOp<"usmops.wide">;
+
+// Loads
+class ArmSME_IntrLoadOverloadedOp<string mnemonic>
+    : ArmSME_IntrOverloadedOp<mnemonic>,
+      Arguments<(ins Arg<Predicate, "Vector predicate">,
+                 Arg<LLVM_AnyPointer, "Load address", [MemRead]>,
+                 Arg<I32, "Tile register ID", [TileWrite]>,
+                 Arg<I32, "Vector number">)>;
+
+def LoadHorizontalWordsIntrOp : ArmSME_IntrLoadOverloadedOp<"ld1w.horiz">;
+def LoadHorizontalDoublesIntrOp : ArmSME_IntrLoadOverloadedOp<"ld1d.horiz">;
+def LoadVerticalWordsIntrOp : ArmSME_IntrLoadOverloadedOp<"ld1w.vert">;
+def LoadVerticalDoublesIntrOp : ArmSME_IntrLoadOverloadedOp<"ld1d.vert">;
+
+// Stores
+class ArmSME_IntrStoreOverloadedOp<string mnemonic>
+    : ArmSME_IntrOverloadedOp<mnemonic>,
+      Arguments<(ins Arg<Predicate, "Vector predicate">,
+                 Arg<LLVM_AnyPointer, "Store address", [MemWrite]>,
+                 Arg<I32, "Tile register ID", [TileRead]>,
+                 Arg<I32, "Vector number">)>;
+
+def StoreHorizontalWordsIntrOp : ArmSME_IntrStoreOverloadedOp<"st1w.horiz">;
+def StoreHorizontalDoublesIntrOp : ArmSME_IntrStoreOverloadedOp<"st1d.horiz">;
+def StoreVerticalWordsIntrOp : ArmSME_IntrStoreOverloadedOp<"st1w.vert">;
+def StoreVerticalDoublesIntrOp : ArmSME_IntrStoreOverloadedOp<"st1d.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/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/ArmSVE/ArmSVEDialect.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;
+
+//===----------------------------------------------------------------------===//
+// Custom Printer/Parser/Verifiers
+//===----------------------------------------------------------------------===//
+
+/// Custom parser for MOPA and MOPS types - Since the predicates are always
+/// going to be same shape as their respective vectors, we should be able to
+/// omit them.
+static ParseResult parseMOPOperandTypes(OpAsmParser &parser, Type &tileTy,
+                                        Type &lhsPredTy, Type &rhsPredTy,
+                                        Type &lhsTy, Type &rhsTy) {
+  MLIRContext *ctx = parser.getContext();
+  // Parse tile type
+  ParseResult result = parser.parseType(tileTy);
+  if (result)
+    return result;
+
+  result = parser.parseComma();
+  if (result)
+    return result;
+
+  // Parse LHS type
+  result = parser.parseType(lhsTy);
+  if (result)
+    return result;
+
+  // Construct lhsPredTy
+  VectorType vecTy = dyn_cast<VectorType>(lhsTy);
+  llvm::SMLoc loc = parser.getCurrentLocation();
+  StringRef expectVectorErr = "expected vector type";
+  if (!vecTy)
+    return parser.emitError(loc, expectVectorErr);
+  lhsPredTy = vecTy.clone(IntegerType::get(ctx, 1));
+
+  result = parser.parseComma();
+  if (result)
+    return result;
+
+  // Parse RHS type
+  result = parser.parseType(rhsTy);
+  if (result)
+    return result;
+
+  // Construct rhsPredTy
+  vecTy = dyn_cast<VectorType>(rhsTy);
+  loc = parser.getCurrentLocation();
+  if (!vecTy)
+    return parser.emitError(loc, expectVectorErr);
+  rhsPredTy = vecTy.clone(IntegerType::get(ctx, 1));
+
+  return result;
+}
+
+static void printMOPOperandTypes(OpAsmPrinter &printer, Operation *op,
+                                 Type tileTy, Type lhsPredTy, Type rhsPredTy,
+                                 Type lhsTy, Type rhsTy) {
+  (void)op;
+  (void)lhsPredTy;
+  (void)rhsPredTy;
+  printer << tileTy << ", " << lhsTy << ", " << rhsTy;
+}
+
+// Verifier for MOPA/MOPS, additions to constraints in ODS
+static LogicalResult verifyMOPOps(Operation *op, Value tile, Value lhs,
+                                  Value rhs, bool isWidening) {
+  VectorType tileTy = cast<VectorType>(tile.getType());
+  VectorType lhsTy = cast<VectorType>(lhs.getType());
+  VectorType rhsTy = cast<VectorType>(rhs.getType());
+  Type tileElTy = tileTy.getElementType();
+  Type vecElTy = lhsTy.getElementType();
+
+  // Verify tile type
+  unsigned tileElWidth = tileElTy.getIntOrFloatBitWidth();
+  if (tileElWidth != 32 && tileElWidth != 64)
+    return op->emitOpError("expected 32 or 64-bit output tiles for MOPA/MOPS");
+  if (tileTy.getShape()[0] != tileTy.getShape()[1])
+    return op->emitOpError("expecting square tiles for SME ops");
+  if (tileTy.getShape()[0] != lhsTy.getShape()[0] ||
+      tileTy.getShape()[1] != rhsTy.getShape().back())
+    return op->emitOpError("invalid shapes for outer product");
+
+  // Integer sign check
+  if (vecElTy.isSignedInteger() != tileElTy.isSignedInteger() ||
+      vecElTy.isUnsignedInteger() != tileElTy.isUnsignedInteger())
+    return op->emitOpError(
+        "expecting tile and vector operands to have same signedness");
+
+  // Scalable dimension check
+  if (lhsTy.getNumScalableDims() != lhsTy.getRank() ||
+      rhsTy.getNumScalableDims() != rhsTy.getRank() ||
+      tileTy.getNumScalableDims() != tileTy.getRank())
+    return op->emitOpError(
+        "expecting all dimensions of all operands to be scalable");
+
+  // Vector length check
+  if (lhsTy.getNumElements() * vecElTy.getIntOrFloatBitWidth() != 128)
+    return op->emitOpError(
+        "expecting input operands to have unit vector length of 128");
+
+  // For widening MOPA/MOPS
+  if (isWidening) {
+    // Input vectors should be of rank 2, with specific shapes
+    if (lhsTy.getRank() != 2 || rhsTy.getRank() != 2)
+      return op->emitOpError(
+          "widening sme outer product instructions expects 2d input vectors");
+
+    if (vecElTy.isInteger(8)) { // SI8/UI8 accumulates to SI32/UI32
+      if (!tileElTy.isInteger(32))
+        return op->emitOpError("expecting i8 to accumulate into i32");
+
+      // LHS and RHS should all be <[4x4]xi8>
+      if (lhsTy.getShape()[0] != 4 || lhsTy != rhsTy)
+        return op->emitOpError(
+            "expecting shape of [4x4] for 8-bit integer operands");
+
+    } else if (vecElTy.isInteger(16)) { // SI16/UI16 accumulates to SI64/UI64
+      if (!tileElTy.isInteger(64))
+        return op->emitOpError("expecting i16 to accumulate into i64");
+
+      // LHS should be [2x4], while RHS should be [4x2]
+      if (lhsTy.getShape()[0] != 2 || rhsTy.getShape()[0] != 4)
+        return op->emitOpError("expecting outer product of LHS [2x4] and RHS "
+                               "[4x2] for 16-bit integer operands");
+
+    } else { // BF16 and F16 should all accumulate to F32
+      if (!tileElTy.isF32())
+        return op->emitOpError("expecting f16/bf16 to accumulate into f32");
+
+      // LHS and RHS should be [4x2] and [2x4] respectively
+      if (lhsTy.getShape()[0] != 4 || rhsTy.getShape()[0] != 2)
+        return op->emitOpError("expecting outer product of LHS [4x2] and RHS "
+                               "[2x4] for 16-bit floating point operands");
+    }
+    return LogicalResult::success();
+  } // End verification of widening MOPA/MOPS
+
+  if (lhsTy.getRank() != 1 || rhsTy.getRank() != 1)
+    return op->emitOpError(
+        "expecting non-widening MOPA/MOPS to have vector operands of rank 1");
+
+  if (vecElTy != tileElTy)
+    return op->emitOpError("expecting same input and tile element types of "
+                           "non-widening MOPA/MOPS");
+  return LogicalResult::success();
+}
+
+LogicalResult MopaOp::verify() {
+  return verifyMOPOps(*this, getTile(), getLhs(), getRhs(), isWidening());
+}
+
+LogicalResult MopsOp::verify() {
+  return verifyMOPOps(*this, getTile(), getLhs(), getRhs(), isWidening());
+}
+
+//===----------------------------------------------------------------------===//
+// 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/test/Dialect/ArmSME/roundtrip.mlir

+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+func.func @ui8mopa(%tile : vector<[4x4]xui32>, %pred : vector<[4x4]xi1>,
+                %operand : vector<[4x4]xui8>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[4x4]xui32>
+  %out = arm_sme.mopa %tile, %pred, %pred, %operand, %operand :
+              vector<[4x4]xui32>, vector<[4x4]xui8>, vector<[4x4]xui8>
+  return
+}
+
+func.func @ui16mopa(%tile : vector<[2x2]xui64>, %lpred : vector<[2x4]xi1>,
+                    %rpred : vector<[4x2]xi1>, %lhs : vector<[2x4]xui16>,
+                    %rhs : vector<[4x2]xui16>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[2x2]xui64>
+  %out = arm_sme.mopa %tile, %lpred, %rpred, %lhs, %rhs :
+              vector<[2x2]xui64>, vector<[2x4]xui16>, vector<[4x2]xui16>
+  return
+}
+
+func.func @si8mopa(%tile : vector<[4x4]xsi32>, %pred : vector<[4x4]xi1>,
+                %operand : vector<[4x4]xsi8>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[4x4]xsi32>
+  %out = arm_sme.mopa %tile, %pred, %pred, %operand, %operand :
+              vector<[4x4]xsi32>, vector<[4x4]xsi8>, vector<[4x4]xsi8>
+  return
+}
+
+func.func @si16mopa(%tile : vector<[2x2]xsi64>, %lpred : vector<[2x4]xi1>,
+                    %rpred : vector<[4x2]xi1>, %lhs : vector<[2x4]xsi16>,
+                    %rhs : vector<[4x2]xsi16>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[2x2]xsi64>
+  %out = arm_sme.mopa %tile, %lpred, %rpred, %lhs, %rhs :
+              vector<[2x2]xsi64>, vector<[2x4]xsi16>, vector<[4x2]xsi16>
+  return
+}
+
+func.func @bf16mopa(%tile : vector<[4x4]xf32>, %lpred : vector<[4x2]xi1>,
+                    %rpred : vector<[2x4]xi1>, %lhs : vector<[4x2]xbf16>,
+                    %rhs : vector<[2x4]xbf16>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[4x4]xf32>
+  %out = arm_sme.mopa %tile, %lpred, %rpred, %lhs, %rhs :
+              vector<[4x4]xf32>, vector<[4x2]xbf16>, vector<[2x4]xbf16>
+  return
+}
+
+func.func @f16mopa(%tile : vector<[4x4]xf32>, %lpred : vector<[4x2]xi1>,
+                    %rpred : vector<[2x4]xi1>, %lhs : vector<[4x2]xf16>,
+                    %rhs : vector<[2x4]xf16>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[4x4]xf32>
+  %out = arm_sme.mopa %tile, %lpred, %rpred, %lhs, %rhs :
+              vector<[4x4]xf32>, vector<[4x2]xf16>, vector<[2x4]xf16>
+  return
+}
+
+func.func @f32mopa(%tile : vector<[4x4]xf32>, %pred : vector<[4]xi1>,
+                %operand : vector<[4]xf32>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[4x4]xf32>
+  %out = arm_sme.mopa %tile, %pred, %pred, %operand, %operand :
+              vector<[4x4]xf32>, vector<[4]xf32>, vector<[4]xf32>
+  return
+}
+
+func.func @f64mopa(%tile : vector<[2x2]xf64>, %pred : vector<[2]xi1>,
+                %operand : vector<[2]xf64>) {
+// CHECK: arm_sme.mopa
+// CHECK-SAME: <[2x2]xf64>
+  %out = arm_sme.mopa %tile, %pred, %pred, %operand, %operand :
+              vector<[2x2]xf64>, vector<[2]xf64>, vector<[2]xf64>
+  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