[mlir][ArmSME] Add custom get_tile_id and cast ops

[c-rhodes]

Summary

This patch adds three new custom ops to the ArmSME dialect:

• arm_sme.get_tile_id - returns a scalar integer representing an SME "virtual tile" that is not in use.
• arm_sme.cast_tile_to_vector - casts from a tile id to a 2-d scalable vector type, which represents an SME "virtual tile".
• arm_sme.cast_vector_to_tile - casts from a 2-d scalable vector type, which represents an SME "virtual tile", to a tile id.

The 'arm_sme.get_tile_id' op currently only supports tile 0, a follow-up
patch will implement proper tile allocation. A further follow-up patch
will demonstrate load/store to/from ZA using these ops.

See the op descriptions for further details and examples.

Thanks to @paulwalker-arm and @awarzynski for helping drive this.

[WanderAway]

LGTM % a few minor details. Great work!

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

54

What about f32 and f64?

188

It may be useful to also indicate what is the intended live range of the allocated tile?

[c-rhodes]

Add FP types.

[c-rhodes]

In D154941#4489283, @WanderAway wrote:

LGTM % a few minor details. Great work!

thanks for taking a look!

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

54

What about f32 and f64?

Thanks for pointing that out, I've also added F16/BF16.

188

It may be useful to also indicate what is the intended live range of the allocated tile?

Sorry could clarify? Not sure exactly what you mean

[WanderAway]

It may be useful to also indicate what is the intended live range of the allocated tile?

Sorry could clarify? Not sure exactly what you mean

I wasn't entirely sure when the tile allocated by this will be "deallocated" since I misread the description of D154955 earlier, but upon a closer look I realized the tile allocated is valid in the scope of the function. I thought it'd be a good idea to make this clear in the description (and also how one probably wouldn't be able to pass this as an argument to another function call)

[c-rhodes]

In D154941#4489650, @WanderAway wrote:

It may be useful to also indicate what is the intended live range of the allocated tile?

Sorry could clarify? Not sure exactly what you mean

I wasn't entirely sure when the tile allocated by this will be "deallocated" since I misread the description of D154955 earlier, but upon a closer look I realized the tile allocated is valid in the scope of the function. I thought it'd be a good idea to make this clear in the description (and also how one probably wouldn't be able to pass this as an argument to another function call)

Yes that's right the scope is a function, we could add a deallocation op to free up tiles within functions and also support spilling/filling to memory instead of throwing a "ran out of tiles" error in the future. I'll fix the description to capture what you raised. Cheers.

[awarzynski]

Great work @c-rhodes , thank you! I've actually immediately rebased https://reviews.llvm.org/D154867 on top of this change and that immediately solved the "data flow" issue 🙏🏻

Overall this looks solid to me. I've left a few minor suggestions - mostly to clarify the documentation.

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

46–47

[nit] Perhaps extract the preconditions to dedicated definition?

Also, when would this be triggered:

CPred<"::llvm::cast<::mlir::VectorType>($_self).getShape() == ArrayRef<int64_t>({" # !interleave(dims, ", ") # "})"

It feels like a "VectorType verifier" that could be safely skipped (i.e. nothing SME specific).

75–76

Perhaps:

• "This is used in conjunction with `cast_vector_to_tile" --> "This would normally be used in conjunction with "virtual tile load" operations to model the output of such Ops. This is required to preserve data-flow as SME intrinsics do not return values."

Basically, this Op and CastVectorToTile complement each other, right? And I guess that's what we want to say here? But IMHO, this description should focus on CastTileToVector.

81–96

This example is a bit busy.I would focus on the Op that's defined here (i.e. CastTileToVector), so that this description is self-contained (try to avoid references to CastVectorToTile). My suggestion:

EXAMPLE:
Input:
```lang=cpp
    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
```

After lowering `vector.load`:
```lang=cpp
    %tile_id = arith.constant 0 : i32
    scf.for %vnum = %c0 to %num_vectors step %c1 {
      // ...
      "arm_sme.intr.ld1w.horiz"(%pg, %ptr, %tile_id, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
    }
    %tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[4]xi32>
    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
```

Another question - are vector.load and vector.store the right Ops here? We don't really lower from these ATM.

105–108

This comment refers to CastVectorToTile

110–113

This comment refers to "these casts", but this is just one cast ;-)

130–161

This example is a bit busy.I would focus on the Op that's defined here (i.e. CastVectorToTile), so that this description is self-contained (try to avoid references to CastTileToVector). My suggestion:

EXAMPLE:
Input:
```mlir

%tile = vector.load %mem1[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
```
 
Output after lowering `vector.store`:
```mlir
%tile = vector.load %mem1[%c0] : memref<?xi32>, vector<[4]x[4]xi32>

%tile_id = arm_sme.cast_vector_to_tile %tile : : (vector<[4]x[4]xi32>) -> i32
scf.for %vnum = %c0 to %num_vectors step %c1 {
  // ...
  "arm_sme.intr.st1w.horiz"(%pg, %ptr, %tile_id, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
}
```

Additionally, canonicalization will look through `cast_vector_to_tile` Ops and fold the
cast ops away if they come from `cast_tile_to_vector`.

178

[nit] Is there any "tile allocation" really taking place? Perhaps "Allocate and return a "virtual tile" ID"?

mlir/test/Dialect/ArmSME/canonicalize.mlir

10–11

What about "the other way round"?

%tile_id = arm_sme.cast_vector_to_tile %tile_1 : vector<[16]x[16]xi8> to i8  
%tile_2 = arm_sme.cast_tile_to_vector %tile_id : i8 to vector<[16]x[16]xi8>

mlir/test/Dialect/ArmSME/invalid.mlir

5

How about:

func.func @arm_sme_cast_tile_to_vector__bad_vector_type(%tile_id : i8) -> vector<[16]x16xi8>

and other combinations? For example:

• vector<[16]x[16]xi4>
• vector<16x[16]xi8>

mlir/test/Dialect/ArmSME/roundtrip.mlir

8

Could you add one more other element type? For example, vector<[1]x[1]xi128> (i.e. the other extreme).

[c-rhodes]

Address comments

[c-rhodes]

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

46–47

Also, when would this be triggered:

CPred<"::llvm::cast<::mlir::VectorType>($_self).getShape() == ArrayRef<int64_t>({" # !interleave(dims, ", ") # "})"

It feels like a "VectorType verifier" that could be safely skipped (i.e. nothing SME specific).

Please could you clarify, not sure what you mean? This verifies the shape, i.e. vector<[16]x[16]xi8> is (16, 16).

75–76

Perhaps:

• "This is used in conjunction with `cast_vector_to_tile" --> "This would normally be used in conjunction with "virtual tile load" operations to model the output of such Ops. This is required to preserve data-flow as SME intrinsics do not return values."

Basically, this Op and CastVectorToTile complement each other, right? And I guess that's what we want to say here? But IMHO, this description should focus on CastTileToVector.

Thanks for the suggestion this has cleaned it up nicely

178

[nit] Is there any "tile allocation" really taking place? Perhaps "Allocate and return a "virtual tile" ID"?

There isn't from the perspective of the op I suppose, it's the pass that does that. Updated the comment.

mlir/test/Dialect/ArmSME/canonicalize.mlir

10–11

What about "the other way round"?

%tile_id = arm_sme.cast_vector_to_tile %tile_1 : vector<[16]x[16]xi8> to i8  
%tile_2 = arm_sme.cast_tile_to_vector %tile_id : i8 to vector<[16]x[16]xi8>

Good spot!

mlir/test/Dialect/ArmSME/roundtrip.mlir

8

Could you add one more other element type? For example, vector<[1]x[1]xi128> (i.e. the other extreme).

I've added tests for all element types

[dcaballe]

LGTM! I like the abstraction! Awesome to see this moving forward!

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

44

Is there a construct to make this a VectorType instead of a ShapedType? I guess the subsequent predicates constraint the shaped type a bit more but it would be great if this could be an vector type directly.

118

This abstraction sounds really great!

131

I guess you also considered introducing a single cast op that could cast both ways depending on the order of the operand/types. I think having two makes sense since this cast is kind of crossing two domains...

166

nit: do the quotes imply anything on virtual tiles? I think I don't get what it is :)

[c-rhodes]

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

44

Is there a construct to make this a VectorType instead of a ShapedType? I guess the subsequent predicates constraint the shaped type a bit more but it would be great if this could be an vector type directly.

I don't believe there is an existing one, the vector ones in mlir/include/mlir/IR/OpBase.td I used for reference here also use this, but perhaps ShapedContainerType could be copied for a VectorType.

131

I guess you also considered introducing a single cast op that could cast both ways depending on the order of the operand/types. I think having two makes sense since this cast is kind of crossing two domains...

That didn't cross my mind actually, it's a good point I think of these casts as being similar to builtin.unrealized_conversion_cast and that does similar but with a single cast op like you say, perhaps this could be a single cast as well.

166

nit: do the quotes imply anything on virtual tiles? I think I don't get what it is :)

To be honest I would prefer we just use tile, but the rationale is that these are not real tiles but merely “views” into ZA.

[awarzynski]

LGTM, thanks for addressing my comments :)

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

44

How about ScalableVectorOf? Could that be applicable here?

46–47

I am just thinking that every vector that you create like this will satisfy this condition and to me this check feels redundant. But I am probably just failing to understand the underlying rationale. No harm in keeping this.

166

I've been suggesting "virtual tiles" as different people mean different things when referring to tiles. "SME virtual tiles" is just a way to highlight that:

• We mean the tiles in the context of the Arm SME extension (as opposed to e.g. tiles when tiling a matmul).
• These tiles are actually "views" into ZA rather than "tiles". A "tile" to me suggests that it's something "square" and so "ZA tile" could, incorrectly, imply "a square section of ZA".

It's a name that one of our architects at Arm has been using and I feel that's very fitting. Naming is hard!

[c-rhodes]

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

44

How about ScalableVectorOf? Could that be applicable here?

All of the existing scalable vector predicates only check is scalable (i.e. any dim) not all dims scalable

46–47

I am just thinking that every vector that you create like this will satisfy this condition and to me this check feels redundant.

And this verifies that :)

166

In the context of SME I think it's clear what a tile is, but I have no strong feelings either way.

[c-rhodes]

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

46–47

I am just thinking that every vector that you create like this will satisfy this condition and to me this check feels redundant.

And this verifies that :)

To clarify, without this check:

%tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[8]xi32>

would be valid

[c-rhodes]

Minor update to add scalar int type to name of functions for get_tile_id tests

[awarzynski]

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

46–47

To clarify, without this check:

%tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[8]xi32>

would be valid

Could you double-check? This works fine:

module {
  func.func @arm_sme_cast_tile_to_vector_i8(%arg0: i8) -> vector<[4]x[16]xi8> {
    %0 = arm_sme.cast_tile_to_vector %arg0 : i8 to vector<[4]x[16]xi8>
    return %0 : vector<[4]x[16]xi8>
  }
}

You will need to replace SMETile with AnyVectorOfAnyRank in the definition of CastTileToVector, but IsVectorOfShape should trigger in both cases, right?

[c-rhodes]

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

46–47

To clarify, without this check:

%tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[8]xi32>

would be valid

Could you double-check? This works fine:

module {
  func.func @arm_sme_cast_tile_to_vector_i8(%arg0: i8) -> vector<[4]x[16]xi8> {
    %0 = arm_sme.cast_tile_to_vector %arg0 : i8 to vector<[4]x[16]xi8>
    return %0 : vector<[4]x[16]xi8>
  }
}

this fails for me (as expected):

build/bin/mlir-opt foo.mlir
foo.mlir:4:8: error: 'arm_sme.cast_tile_to_vector' op result #0 must be vector<[16]x[16]xi8> of 8-bit signless integer values or vector<[8]x[8]xi16> of 16-bit signless integer values or vector<[4]x[4]xi32> of 32-bit signless integer values or vector<[2]x[2]xi64> of 64-bit signless integer values or vector<[1]x[1]xi128> of 128-bit signless integer values or vector<[8]x[8]xf16> of 16-bit float values or vector<[8]x[8]xbf16> of bfloat16 type values or vector<[4]x[4]xf32> of 32-bit float values or vector<[2]x[2]xf64> of 64-bit float values, but got 'vector<[4]x[16]xi8>'
  %0 = arm_sme.cast_tile_to_vector %tile_id : i8 to vector<[4]x[16]xi8>
       ^
foo.mlir:4:8: note: see current operation: %0 = "arm_sme.cast_tile_to_vector"(%arg0) : (i8) -> vector<[4]x[16]xi8>

and doesn't if I remove IsVectorOfShape<dims> check.

You will need to replace SMETile with AnyVectorOfAnyRank in the definition of CastTileToVector, but IsVectorOfShape should trigger in both cases, right?

I'm not sure I follow, please could you clarify?

[awarzynski]

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

46–47

You are right and I am wrong, sorry.

I've just checked the generated CPP code and it's this:

((::llvm::cast<::mlir::VectorType>(type).getShape() == ArrayRef<int64_t>({16, 16})))
// other similar checks

So RHS is taken from:

def SMETile : AnyTypeOf<[nxnxv16i8, nxnxv8i16, nxnxv4i32, nxnxv2i64, nxnxv1i128,
                         nxnxv8f16, nxnxv8bf16, nxnxv4f32, nxnxv2f64]>;

I thought that for this example (vector<[4]x[8]xi32>) it would check the following instead:

((::llvm::cast<::mlir::VectorType>(type).getShape() == ArrayRef<int64_t>({4, 8})))

i.e. take the RHS from the input (vector<[4]x[8]xi32>). Hence the confusion.

[c-rhodes]

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

46–47

no worries, thanks for checking

[c-rhodes]

I'll land this tomorrow unless there's any further comments by then.

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

let description = [{
     Sources:
     https://developer.arm.com/documentation/ddi0616
     https://developer.arm.com/documentation/ddi0602/2023-03/SME-Instructions
   }];
   let dependentDialects = ["scf::SCFDialect"];
 }
 
 //===----------------------------------------------------------------------===//
+// ArmSME type definitions
+//===----------------------------------------------------------------------===//
+
+class SMETileType<Type datatype, list<int> dims, string description>
+  : ShapedContainerType<[datatype],
+      And<[IsVectorOfRankPred<[2]>,
+           CPred<[{::llvm::cast<::mlir::VectorType>($_self).allDimsScalable()}]>,
+           CPred<"::llvm::cast<::mlir::VectorType>($_self).getShape() == ArrayRef<int64_t>({" # !interleave(dims, ", ") # "})">]>,
+  description>;
+
+def nxnxv16i8  : SMETileType<I8,   [16, 16], "vector<[16]x[16]xi8>">;
+def nxnxv8i16  : SMETileType<I16,  [8,  8 ], "vector<[8]x[8]xi16>">;
+def nxnxv4i32  : SMETileType<I32,  [4,  4 ], "vector<[4]x[4]xi32>">;
+def nxnxv2i64  : SMETileType<I64,  [2,  2 ], "vector<[2]x[2]xi64>">;
+def nxnxv1i128 : SMETileType<I128, [1,  1 ], "vector<[1]x[1]xi128>">;
+
+def nxnxv8f16  : SMETileType<F16,  [8,  8 ], "vector<[8]x[8]xf16>">;
+def nxnxv8bf16 : SMETileType<BF16, [8,  8 ], "vector<[8]x[8]xbf16>">;
+def nxnxv4f32  : SMETileType<F32,  [4,  4 ], "vector<[4]x[4]xf32>">;
+def nxnxv2f64  : SMETileType<F64,  [2,  2 ], "vector<[2]x[2]xf64>">;
+
+def SMETile : AnyTypeOf<[nxnxv16i8, nxnxv8i16, nxnxv4i32, nxnxv2i64, nxnxv1i128,
+                         nxnxv8f16, nxnxv8bf16, nxnxv4f32, nxnxv2f64]>;
+
+//===----------------------------------------------------------------------===//
+// ArmSME op definitions
+//===----------------------------------------------------------------------===//
+
+class ArmSME_Op<string mnemonic, list<Trait> traits = []> :
+  Op<ArmSME_Dialect, mnemonic, traits> {}
+
+def CastTileToVector : ArmSME_Op<"cast_tile_to_vector", [Pure]> {
+  let summary = "Cast from tile id to 2-d scalable vector type";
+  let description = [{
+    A `cast_tile_to_vector` operation does a cast from a tile id to a 2-d
+    scalable vector type, which represents an SME "virtual tile". This is used
+    in conjunction with `cast_vector_to_tile` to preserve dataflow and type
+    legality when lowering vector ops that have both inputs and outputs, to SME
+    intrinsics that have only inputs.
+
+    Example:
+    ```mlir
+
+    // input
+    %tile = vector.load %mem1[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+
+    // lower vector.load -> SME intrinsics
+
+    %tile_id = arm_sme.get_tile_id : i32
+    scf.for %vnum = %c0 to %num_vectors step %c1 {
+      // ...
+      "arm_sme.intr.ld1w.horiz"(%pg, %ptr, %tile_id, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
+    }
+    %tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[4]xi32>
+    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+    ```
+
+    In the example above, the `vector.load` can't be replaced with an SME
+    intrinsic that has no outputs since it is used by the `vector.store`.
+    However, by inserting a `cast_tile_to_vector` op after the load intrinsics
+    the `vector.load` can be replaced. This enables "local" rewrites on
+    individual vector ops, rather than "global" rewrites that would have to
+    look at the vector op uses and also lower them.
+
+    The opposite is true for the `vector.store`, when lowered to intrinsics
+    they would be preceded by a `cast_vector_to_tile` op. Once the lowering is
+    complete the canonicalizer will fold the casts away. The
+    `cast_vector_to_tile` op example shows the other half of the lowering.
+
+    These casts are expected to be folded, but may persist if there's an
+    incomplete lowering where a vector op has been lowered to SME but the uses
+    haven't, much like if `-reconcile-unrealized-casts` fails. Currently these
+    cast ops cannot be lowered to LLVM, but may be in the future.
+  }];
+  let arguments = (ins AnyTypeOf<[I8, I16, I32, I64, I128]>:$tile_id);
+  let results = (outs SMETile:$vector);
+  let assemblyFormat =
+    "$tile_id attr-dict `:` type($tile_id) `to` type($vector)";
+}
+
+def CastVectorToTile : ArmSME_Op<"cast_vector_to_tile", [Pure]> {
+  let summary = "Cast from 2-d scalable vector type to tile id";
+  let description = [{
+    A `cast_vector_to_tile` operation does a cast from a 2-d scalable vector
+    type, which represents an SME "virtual tile", to a tile id. This is used in
+    conjunction with `cast_tile_to_vector` to preserve dataflow and type
+    legality when lowering vector ops that have both inputs and outputs, to SME
+    intrinsics that have only inputs.
+
+    Example:
+    ```mlir
+
+    // input
+    %tile = vector.load %mem1[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+
+    // lower vector.load -> SME intrinsics
+    %tile_id = arm_sme.get_tile_id : i32
+    scf.for %vnum = %c0 to %num_vectors step %c1 {
+      // ...
+      "arm_sme.intr.ld1w.horiz"(%pg, %ptr, %tile_id, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
+    }
+    %tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[4]xi32>
+    vector.store %tile, %mem2[%c0] : memref<?xi32>, vector<[4]x[4]xi32>
+
+    // lower vector.store -> SME intrinsics
+    %tile_id_0 = arm_sme.get_tile_id : i32
+    scf.for %vnum = %c0 to %num_vectors step %c1 {
+      // ...
+      "arm_sme.intr.ld1w.horiz"(%pg, %ptr, %tile_id_0, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
+    }
+    %tile = arm_sme.cast_tile_to_vector %tile_id : i32 to vector<[4]x[4]xi32>
+    scf.for %vnum = %c0 to %num_vectors step %c1 {
+      // ...
+      %tile_id_1 = arm_sme.cast_vector_to_tile %tile : : (vector<[4]x[4]xi32>) -> i32
+      "arm_sme.intr.st1w.horiz"(%pg, %ptr, %tile_id_1, %vnum) : (vector<[4]xi1>, !llvm.ptr, i32, i32) -> ()
+    }
+
+    // canonicalization will look through `cast_vector_to_tile` and fold the
+    // cast ops away.
+    ```
+  }];
+  let arguments = (ins SMETile:$vector);
+  let results = (outs AnyTypeOf<[I8, I16, I32, I64, I128]>:$tile_id);
+  let assemblyFormat =
+    "$vector attr-dict `:` type($vector) `to` type($tile_id)";
+  let hasCanonicalizeMethod = 1;
+}
+
+def GetTileID : ArmSME_Op<"get_tile_id", [Pure]> {
+  let summary = "Returns an SME \"virtual tile\" id that is not in use";
+  let description = [{
+    A `get_tile_id` operation returns a scalar integer of given type
+    representing a tile id of an SME "virtual tile" that is not in use.
+
+    Example:
+    ```mlir
+    // Allocate an 8-bit element ZA tile
+    %za0_b = arm_sme.get_tile_id : i8
+    ```
+
+    Example:
+    ```
+    // Allocate two 16-bit element ZA tiles
+    %za0_h = arm_sme.get_tile_id : i16
+    %za1_h = arm_sme.get_tile_id : i16
+    ```
+
+    Example:
+    ```
+    // Allocate a 128-bit element ZA tile
+    %za0_q = arm_sme.get_tile_id : i128
+    ```
+  }];
+
+  let results = (outs AnyTypeOf<[I8, I16, I32, I64, I128]>:$tile_id);
+  let assemblyFormat = "attr-dict `:` type($tile_id)";
+}
+
+//===----------------------------------------------------------------------===//
 // 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]>;
 

mlir/include/mlir/IR/OpBase.td

class SignlessIntOfWidths<list<int> widths> :
     AnyTypeOf<!foreach(w, widths, I<w>),
               !interleave(widths, "/") # "-bit signless integer">;
 
 def I1  : I<1>;
 def I8  : I<8>;
 def I16 : I<16>;
 def I32 : I<32>;
 def I64 : I<64>;
+def I128 : I<128>;
 
 // Any signed integer type irrespective of its width.
 def AnySignedInteger : Type<
   CPred<"$_self.isSignedInteger()">, "signed integer">;
 
 // Signed integer type of a specific width.
 class SI<int width>
     : Type<CPred<"$_self.isSignedInteger(" # width # ")">,

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

 #include "mlir/Dialect/ArmSME/IR/ArmSMETypes.cpp.inc"
 
 void ArmSMEDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/ArmSME/IR/ArmSME.cpp.inc"
       >();
 }
+
+// cast_vector_to_tile(cast_tile_to_vector(tile_id)) -> tile_id
+LogicalResult CastVectorToTile::canonicalize(CastVectorToTile op,
+                                             PatternRewriter &rewriter) {
+  if (auto vectorOp =
+          dyn_cast<CastTileToVector>(op.getVector().getDefiningOp())) {
+    op.replaceAllUsesWith(vectorOp.getTileId());
+    return success();
+  }
+  return failure();
+}

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/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Dialect/ArmSME/IR/ArmSME.h"
 #include "mlir/Dialect/ArmSME/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 
 using namespace mlir;
 using namespace mlir::arm_sme;
 
 namespace {
 /// Insert 'llvm.aarch64.sme.za.enable' intrinsic at the start of 'func.func'
 /// ops to enable the ZA storage array.
LogicalResult matchAndRewrite(func::ReturnOp op,
                                 PatternRewriter &rewriter) const final {
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(op);
     rewriter.create<arm_sme::aarch64_sme_za_disable>(op->getLoc());
     rewriter.updateRootInPlace(op, [] {});
     return success();
   }
 };
+
+struct GetTileIDConversion : public ConvertOpToLLVMPattern<GetTileID> {
+  using ConvertOpToLLVMPattern<GetTileID>::ConvertOpToLLVMPattern;
+  LogicalResult
+  matchAndRewrite(GetTileID op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // TODO: implement tile allocation, currently only tile 0 is supported.
+    rewriter.replaceOpWithNewOp<LLVM::ConstantOp>(op, rewriter.getI32Type(), 0);
+    return success();
+  }
+};
 } // namespace
 
 void mlir::populateArmSMELegalizeForLLVMExportPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns) {
   patterns.add<EnableZAPattern, DisableZAPattern>(patterns.getContext());
 }
 
 void mlir::configureArmSMELegalizeForExportTarget(
     LLVMConversionTarget &target) {
-  target.addLegalOp<scf::ForOp, scf::YieldOp, arm_sme::aarch64_sme_zero,
+  target.addLegalOp<scf::ForOp, scf::YieldOp, arm_sme::CastTileToVector,
+                    arm_sme::CastVectorToTile, arm_sme::aarch64_sme_zero,
                     arm_sme::aarch64_sme_str, arm_sme::aarch64_sme_za_enable,
                     arm_sme::aarch64_sme_za_disable>();
+  target.addLegalOp<GetTileID>();
 
   // Mark 'func.func' ops as legal if either:
   //   1. no 'arm_za' function attribute is present.
   //   2. the 'arm_za' function attribute is present and the first op in the
   //      function is an 'arm_sme::aarch64_sme_za_enable' intrinsic.
   target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp funcOp) {
     if (funcOp.isDeclaration())
       return true;

mlir/test/Dialect/ArmSME/canonicalize.mlir

+// RUN: mlir-opt -canonicalize -split-input-file -verify-diagnostics %s | mlir-opt | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: @canonicalize_casts
+// CHECK-SAME: %[[TILE_ID:.*]]: i8
+func.func @canonicalize_casts(%tile_id_0 : i8) -> i8 {
+  // CHECK-NOT: arm_sme.cast_tile_to_vector
+  // CHECK-NOT: arm_sme.cast_vector_to_tile
+  %tile = arm_sme.cast_tile_to_vector %tile_id_0 : i8 to vector<[16]x[16]xi8>
+  %tile_id_1 = arm_sme.cast_vector_to_tile %tile : vector<[16]x[16]xi8> to i8
+  // CHECK-NEXT: return %[[TILE_ID]] : i8
+  return %tile_id_1 : i8
+}

mlir/test/Dialect/ArmSME/invalid.mlir

+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+// -----
+
+func.func @arm_sme_cast_tile_to_vector__bad_vector_type(%tile_id : i8) -> vector<[16]xi8> {
+  // expected-error@+1 {{op result #0 must be vector<[16]x[16]xi8> of 8-bit signless integer values or vector<[8]x[8]xi16> of 16-bit signless integer values or vector<[4]x[4]xi32> of 32-bit signless integer values or vector<[2]x[2]xi64> of 64-bit signless integer values or vector<[1]x[1]xi128> of 128-bit signless integer values or vector<[8]x[8]xf16> of 16-bit float values or vector<[8]x[8]xbf16> of bfloat16 type values or vector<[4]x[4]xf32> of 32-bit float values or vector<[2]x[2]xf64> of 64-bit float values, but got 'vector<[16]xi8>'}}
+  %0 = arm_sme.cast_tile_to_vector %tile_id : i8 to vector<[16]xi8>
+  return %0 : vector<[16]xi8>
+}
+
+// -----
+
+func.func @arm_sme_cast_vector_to_tile__bad_rank_1d(%vector : vector<[16]xi8>) -> i8 {
+  // expected-error@+1 {{op operand #0 must be vector<[16]x[16]xi8> of 8-bit signless integer values or vector<[8]x[8]xi16> of 16-bit signless integer values or vector<[4]x[4]xi32> of 32-bit signless integer values or vector<[2]x[2]xi64> of 64-bit signless integer values or vector<[1]x[1]xi128> of 128-bit signless integer values or vector<[8]x[8]xf16> of 16-bit float values or vector<[8]x[8]xbf16> of bfloat16 type values or vector<[4]x[4]xf32> of 32-bit float values or vector<[2]x[2]xf64> of 64-bit float values, but got 'vector<[16]xi8>'}}
+  %0 = arm_sme.cast_vector_to_tile %vector : vector<[16]xi8> to i8
+  return %0 : i8
+}
+
+// -----
+
+func.func @arm_sme_get_tile_id__bad_type() -> i1 {
+  // expected-error@+1 {{op result #0 must be 8-bit signless integer or 16-bit signless integer or 32-bit signless integer or 64-bit signless integer or 128-bit signless integer}}
+  %0 = arm_sme.get_tile_id : i1
+  return %0 : i1
+}

mlir/test/Dialect/ArmSME/roundtrip.mlir

+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | mlir-opt | FileCheck %s
+
+// -----
+
+func.func @arm_sme_cast_tile_to_vector(%tile_id : i8) -> vector<[16]x[16]xi8> {
+  // CHECK: arm_sme.cast_tile_to_vector {{.*}} : i8 to vector<[16]x[16]xi8>
+  %0 = arm_sme.cast_tile_to_vector %tile_id : i8 to vector<[16]x[16]xi8>
+  return %0 : vector<[16]x[16]xi8>
+}
+
+// -----
+
+func.func @arm_sme_cast_vector_to_tile(%vector : vector<[16]x[16]xi8>) -> i8 {
+  // CHECK: arm_sme.cast_vector_to_tile {{.*}} : vector<[16]x[16]xi8> to i8
+  %0 = arm_sme.cast_vector_to_tile %vector : vector<[16]x[16]xi8> to i8
+  return %0 : i8
+}
+
+// -----
+
+func.func @arm_sme_get_tile_id() -> i32 {
+  // CHECK: arm_sme.get_tile_id : i32
+  %0 = arm_sme.get_tile_id : i32
+  return %0 : i32
+}