[mlir][ArmSME] Implement tile allocation

[c-rhodes]

Summary

This patch adds a pass '-allocate-sme-tiles' to the ArmSME dialect that
implements allocation of SME ZA tiles.

It does this at the 'func.func' op level by replacing
'arm_sme.get_tile_id' ops with 'arith.constant' ops that represent the
tile number. The tiles in use in a given function are tracked by an
integer function attribute 'tilesInUse' that is a 16-bit tile mask with
a bit for each 128-bit element tile (ZA0.Q-ZA15.Q), the smallest ZA tile
granule. 'tilesInUse' is initialized on the first 'arm_sme.get_tile_id'
rewrite and updated on each subsequent rewrite. Mixing of different
element tile types is supported.

Section B2.3.2 of the SME spec [1] describes how the 128-bit element
tiles overlap with other element tiles.

Depends on D154941

[1] https://developer.arm.com/documentation/ddi0616/aa

[awarzynski]

Thanks for working on this! The overall logic makes a lot of sense to me.

My only concern is that "tile allocation" could be confused with "tiling" in e.g. Linalg (I'm sure there's more examples in MLIR). I suggest making it a bit clearer that this is "SME virtual tile allocation" (or "tile allocation for virtual SME tiles"). Even in the context of SME, we will be using Linalg tiling :) I made some specific suggestions inline. You could try even some bolder changes, e.g. "TileAllocationPass" --> "SMEVirtualTileAllocationPass" (or something similar). This is a nice-to-have.

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

43

Looks like we use arm- for pretty much everything Arm-related, so suggest adding it here for consistency.

46–48

Just suggested to refine/expand a bit.

I would drop "currently" as there are no timescales (yet) for relaxing that.

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

64

[nit] I would rename this to something a bit more specific. Also, isn't the convention to use underscores? So, e.g. arm_sme_tiles_in_use?

153

I suggest making this error a bit more specific so that it's easy to grep for. I appreciate that this is not ambiguous when only one pass is run. But this will eventually be embedded in some larger compilation pipelines that will include tiling :)

[c-rhodes]

Address comments

[c-rhodes]

In D154955#4493675, @awarzynski wrote:

Thanks for working on this! The overall logic makes a lot of sense to me.

My only concern is that "tile allocation" could be confused with "tiling" in e.g. Linalg (I'm sure there's more examples in MLIR). I suggest making it a bit clearer that this is "SME virtual tile allocation" (or "tile allocation for virtual SME tiles"). Even in the context of SME, we will be using Linalg tiling :) I made some specific suggestions inline. You could try even some bolder changes, e.g. "TileAllocationPass" --> "SMEVirtualTileAllocationPass" (or something similar). This is a nice-to-have.

Thanks for the comments! I'll keep it as TileAllocation for simplicity but we can rename it in the future if it's causing confusion within the dialect.

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

43

Looks like we use arm- for pretty much everything Arm-related, so suggest adding it here for consistency.

Good spot!

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

153

I suggest making this error a bit more specific so that it's easy to grep for. I appreciate that this is not ambiguous when only one pass is run. But this will eventually be embedded in some larger compilation pipelines that will include tiling :)

Updated, but just wanted to note this error is on the arm_sme.get_tile_id op so confusion with tiling seems unlikely.

[awarzynski]

Cheers for addressing my comments, LGTM!

[c-rhodes]

• Scalar int type of arith.constant should match that of the arm_get.get_tile_id it's replacing.
• Add -allocate-arm-sme-tiles to mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir.

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

 #define GEN_PASS_DECL
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 /// Pass to enable Armv9 Streaming SVE mode.
 std::unique_ptr<Pass>
 createEnableArmStreamingPass(const ArmStreaming mode = ArmStreaming::Default,
                              const bool enableZA = false);
 
+/// Pass that replaces 'arm_sme.get_tile_id' ops with actual tiles.
+std::unique_ptr<Pass> createTileAllocationPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
 
 } // namespace arm_sme
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/ArmSME/Transforms/Passes.td

Option<"mode", "mode", "mlir::arm_sme::ArmStreaming",
                "manages PSTATE.SM on entry/exit.")
           )}]>,
     Option<"enableZA", "enable-za", "bool", /*default=*/"false",
            "Enable ZA storage array.">,
   ];
   let dependentDialects = ["func::FuncDialect"];
 }
 
+def TileAllocation
+    : Pass<"allocate-arm-sme-tiles", "mlir::func::FuncOp"> {
+  let summary = "Allocate SME tiles";
+  let description = [{
+    This pass does tile allocation for SME "virtual tiles". It is run at the
+    'func.func' op level, replacing 'arm_sme.get_tile_id' ops with (i32) tile
+    ids. An error will be emitted when there's no tiles left.
+  }];
+  let constructor = "mlir::arm_sme::createTileAllocationPass()";
+  let dependentDialects = ["func::FuncDialect"];
+}
+
 #endif // MLIR_DIALECT_ARMSME_TRANSFORMS_PASSES_TD

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

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

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

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
 
 /// Lower 'arm_sme.zero'. Use 'arm_sme.cast_tile_to_vector' to model the return
 /// value. The latter is a nop, which should be folded away (e.g. during
 /// canonicalisation).
 ///
 ///  BEFORE:
 ///  ```mlir

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

+//===- TileAllocation.cpp - Allocate SME ZA tiles -------------------------===//
+//
+// 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 pass allocates SME tiles at the 'func.func' op level for
+// 'arm_sme.get_tile_id' ops. It does this using a 16-bit tile mask that has a
+// bit for each 128-bit element tile (ZA0.Q-ZA15.Q), the smallest ZA tile
+// granule.
+//
+// The 128-bit tiles overlap with other element tiles as follows (see section
+// B2.3.2 of SME spec [1]):
+//
+//   Tile    Overlaps
+//   ---------------------------------------------------------------------------
+//   ZA0.B   ZA0.Q, ZA1.Q, ZA2.Q, ZA3.Q, ZA4.Q, ZA5.Q, ZA6.Q, ZA7.Q, ZA8.Q,
+//           ZA9.Q, ZA10.Q, ZA11.Q, ZA12.Q, ZA13.Q, ZA14.Q, ZA15.Q
+//   ZA0.H   ZA0.Q, ZA2.Q, ZA4.Q, ZA6.Q, ZA8.Q, ZA10.Q, ZA12.Q, ZA14.Q
+//   ZA1.H   ZA1.Q, ZA3.Q, ZA5.Q, ZA7.Q, ZA9.Q, ZA11.Q, ZA13.Q, ZA15.Q
+//   ZA0.S   ZA0.Q, ZA4.Q, ZA8.Q, ZA12.Q
+//   ZA1.S   ZA1.Q, ZA5.Q, ZA9.Q, ZA13.Q
+//   ZA2.S   ZA2.Q, ZA6.Q, ZA10.Q, ZA14.Q
+//   ZA3.S   ZA3.Q, ZA7.Q, ZA11.Q, ZA15.Q
+//   ZA0.D   ZA0.Q, ZA8.Q
+//   ZA1.D   ZA1.Q, ZA9.Q
+//   ZA2.D   ZA2.Q, ZA10.Q
+//   ZA3.D   ZA3.Q, ZA11.Q
+//   ZA4.D   ZA4.Q, ZA12.Q
+//   ZA5.D   ZA5.Q, ZA13.Q
+//   ZA6.D   ZA6.Q, ZA14.Q
+//   ZA7.D   ZA7.Q, ZA15.Q
+//
+// The tiles in use are tracked via a function attribute 'arm_sme.tiles_in_use'
+// that is initalized during the first 'arm_sme.get_tile_id' rewrite and
+// updated on each subsequent rewrite.
+//
+// [1] https://developer.arm.com/documentation/ddi0616/aa
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
+#include "mlir/Dialect/ArmSME/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#define DEBUG_TYPE "allocate-arm-sme-tiles"
+
+namespace mlir {
+namespace arm_sme {
+#define GEN_PASS_DEF_TILEALLOCATION
+#include "mlir/Dialect/ArmSME/Transforms/Passes.h.inc"
+} // namespace arm_sme
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::arm_sme;
+
+namespace {
+
+static constexpr char kTilesInUseAttr[] = "arm_sme.tiles_in_use";
+
+enum class TileMask : unsigned {
+  // clang-format off
+  kZA0B  = 0xffff, // 1111 1111 1111 1111
+
+  kZA0H  = 0xaaaa, // 1010 1010 1010 1010
+  kZA1H  = 0x5555, // 0101 0101 0101 0101
+
+  kZA0S  = 0x8888, // 1000 1000 1000 1000
+  kZA1S  = 0x4444, // 0100 0100 0100 0100
+  kZA2S  = 0x2222, // 0010 0010 0010 0010
+  kZA3S  = 0x1111, // 0001 0001 0001 0001
+
+  kZA0D  = 0x8080, // 1000 0000 1000 0000
+  kZA1D  = 0x4040, // 0100 0000 0100 0000
+  kZA2D  = 0x2020, // 0010 0000 0010 0000
+  kZA3D  = 0x1010, // 0001 0000 0001 0000
+  kZA4D  = 0x808,  // 0000 1000 0000 1000
+  kZA5D  = 0x404,  // 0000 0100 0000 0100
+  kZA6D  = 0x202,  // 0000 0010 0000 0010
+  kZA7D  = 0x101,  // 0000 0001 0000 0001
+
+  kZA0Q  = 0x8000, // 1000 0000 0000 0000
+  kZA1Q  = 0x4000, // 0100 0000 0000 0000
+  kZA2Q  = 0x2000, // 0010 0000 0000 0000
+  kZA3Q  = 0x1000, // 0001 0000 0000 0000
+  kZA4Q  = 0x800,  // 0000 1000 0000 0000
+  kZA5Q  = 0x400,  // 0000 0100 0000 0000
+  kZA6Q  = 0x200,  // 0000 0010 0000 0000
+  kZA7Q  = 0x100,  // 0000 0001 0000 0000
+  kZA8Q  = 0x80,   // 0000 0000 1000 0000
+  kZA9Q  = 0x40,   // 0000 0000 0100 0000
+  kZA10Q = 0x20,   // 0000 0000 0010 0000
+  kZA11Q = 0x10,   // 0000 0000 0001 0000
+  kZA12Q = 0x8,    // 0000 0000 0000 1000
+  kZA13Q = 0x4,    // 0000 0000 0000 0100
+  kZA14Q = 0x2,    // 0000 0000 0000 0010
+  kZA15Q = 0x1,    // 0000 0000 0000 0001
+
+  kNone = 0x0,     // 0000 0000 0000 0000
+  // clang-format on
+
+  LLVM_MARK_AS_BITMASK_ENUM(kZA0B)
+};
+
+/// Returns the set of masks relevant for the given type.
+static ArrayRef<TileMask> getMasks(Type type) {
+  static const SmallVector<TileMask> ZA_B_MASKS = {TileMask::kZA0B};
+  static const SmallVector<TileMask> ZA_H_MASKS = {TileMask::kZA0H,
+                                                   TileMask::kZA1H};
+  static const SmallVector<TileMask> ZA_S_MASKS = {
+      TileMask::kZA0S, TileMask::kZA1S, TileMask::kZA2S, TileMask::kZA3S};
+  static const SmallVector<TileMask> ZA_D_MASKS = {
+      TileMask::kZA0D, TileMask::kZA1D, TileMask::kZA2D, TileMask::kZA3D,
+      TileMask::kZA4D, TileMask::kZA5D, TileMask::kZA6D, TileMask::kZA7D};
+  static const SmallVector<TileMask> ZA_Q_MASKS = {
+      TileMask::kZA0Q,  TileMask::kZA1Q,  TileMask::kZA2Q,  TileMask::kZA3Q,
+      TileMask::kZA4Q,  TileMask::kZA5Q,  TileMask::kZA6Q,  TileMask::kZA7Q,
+      TileMask::kZA8Q,  TileMask::kZA9Q,  TileMask::kZA10Q, TileMask::kZA11Q,
+      TileMask::kZA12Q, TileMask::kZA13Q, TileMask::kZA14Q, TileMask::kZA15Q};
+  switch (cast<IntegerType>(type).getWidth()) {
+  default:
+    llvm_unreachable("unexpected type!");
+  case 8:
+    return ZA_B_MASKS;
+  case 16:
+    return ZA_H_MASKS;
+  case 32:
+    return ZA_S_MASKS;
+  case 64:
+    return ZA_D_MASKS;
+  case 128:
+    return ZA_Q_MASKS;
+  }
+}
+
+/// Allocates a tile to 'tileID' or returns an error if there are no tiles left.
+static LogicalResult getTile(GetTileID tileIDOp, TileMask &tilesInUse,
+                             unsigned &tileID) {
+  auto masks = getMasks(tileIDOp.getType());
+  for (const auto &it : llvm::enumerate(masks)) {
+    const auto tileMask = it.value();
+    if ((tilesInUse & tileMask) == TileMask::kNone) {
+      tilesInUse |= tileMask;
+      tileID = it.index();
+      return success();
+    }
+  }
+  return tileIDOp.emitError("ran out of SME virtual tiles!");
+}
+
+struct GetTileIDConversion : public OpRewritePattern<GetTileID> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(GetTileID tileIDOp,
+                                PatternRewriter &rewriter) const override {
+    auto funcOp = tileIDOp->getParentOfType<func::FuncOp>();
+    TileMask tilesInUse;
+    if (auto tilesInUseAttr =
+            funcOp->getAttrOfType<IntegerAttr>(kTilesInUseAttr))
+      tilesInUse = static_cast<TileMask>(tilesInUseAttr.getInt());
+    else
+      tilesInUse = TileMask::kNone;
+
+    unsigned tileID;
+    if (failed(getTile(tileIDOp, tilesInUse, tileID)))
+      return failure();
+
+    funcOp->setAttr(kTilesInUseAttr,
+                    rewriter.getI32IntegerAttr((unsigned)tilesInUse));
+
+    auto tileType = tileIDOp.getType();
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(
+        tileIDOp, tileType, rewriter.getIntegerAttr(tileType, tileID));
+    return success();
+  }
+};
+
+struct TileAllocationPass
+    : public arm_sme::impl::TileAllocationBase<TileAllocationPass> {
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    ConversionTarget target(getContext());
+    patterns.add<GetTileIDConversion>(patterns.getContext());
+    target.addLegalOp<arith::ConstantOp>();
+    target.addIllegalOp<GetTileID>();
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::arm_sme::createTileAllocationPass() {
+  return std::make_unique<TileAllocationPass>();
+}

mlir/test/Dialect/ArmSME/tile-allocation.mlir

+// RUN: mlir-opt %s -allocate-arm-sme-tiles -split-input-file -verify-diagnostics | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: mixed_tiles
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65534 : i32}
+func.func @mixed_tiles() {
+  // ZA0.Q, ZA2.Q, ZA4.Q, ZA6.Q, ZA8.Q, ZA10.Q, ZA12.Q, ZA14.Q
+  // CHECK-NEXT: arith.constant 0
+  %za0_h = arm_sme.get_tile_id : i16
+  // ZA1.Q, ZA5.Q, ZA9.Q, ZA13.Q
+  // CHECK-NEXT: arith.constant 1
+  %za1_s = arm_sme.get_tile_id : i32
+  // ZA3.D ZA3.Q, ZA11.Q
+  // CHECK-NEXT: arith.constant 3
+  %za3_d = arm_sme.get_tile_id : i64
+  // ZA7.Q
+  // CHECK-NEXT: arith.constant 7
+  %za7_q = arm_sme.get_tile_id : i128
+  // ZA15.Q is still free.
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_b
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_b() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_b = arm_sme.get_tile_id : i8
+  return
+}
+
+// -----
+
+func.func @za_b__out_of_tiles() {
+  %za0_b = arm_sme.get_tile_id : i8
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i8
+  return
+}
+
+// -----
+
+func.func @za_b_overlapping_za_q() {
+  %za0_b = arm_sme.get_tile_id : i8
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za0_h
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 43690 : i32}
+func.func @za0_h() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_h = arm_sme.get_tile_id : i16
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_h
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_h() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_h = arm_sme.get_tile_id : i16
+  // CHECK-NEXT: arith.constant 1
+  %za1_h = arm_sme.get_tile_id : i16
+  return
+}
+
+// -----
+
+func.func @za_h__out_of_tiles() {
+  %za0_h = arm_sme.get_tile_id : i16
+  %za1_h = arm_sme.get_tile_id : i16
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i16
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_h_overlapping_za_s
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_h_overlapping_za_s() {
+  // ZA0.Q, ZA2.Q, ZA4.Q, ZA6.Q, ZA8.Q, ZA10.Q, ZA12.Q, ZA14.Q
+  // CHECK-NEXT: arith.constant 0
+  %za0_h = arm_sme.get_tile_id : i16
+  // ZA1.Q, ZA5.Q, ZA9.Q, ZA13.Q
+  // CHECK-NEXT: arith.constant 1
+  %za1_s = arm_sme.get_tile_id : i32
+  // ZA3.Q, ZA7.Q, ZA11.Q, ZA15.Q
+  // CHECK-NEXT: arith.constant 3
+  %za3_s = arm_sme.get_tile_id : i32
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_h_overlapping_za_d
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_h_overlapping_za_d() {
+  // ZA0.Q, ZA2.Q, ZA4.Q, ZA6.Q, ZA8.Q, ZA10.Q, ZA12.Q, ZA14.Q
+  // CHECK-NEXT: arith.constant 0
+  %za0_h = arm_sme.get_tile_id : i16
+  // ZA1.Q, ZA9.Q
+  // CHECK-NEXT: arith.constant 1
+  %za1_d = arm_sme.get_tile_id : i64
+  // ZA3.Q, ZA11.Q
+  // CHECK-NEXT: arith.constant 3
+  %za3_d = arm_sme.get_tile_id : i64
+  // ZA5.Q, ZA13.Q
+  // CHECK-NEXT: arith.constant 5
+  %za5_d = arm_sme.get_tile_id : i64
+  // ZA7.Q, ZA15.Q
+  // CHECK-NEXT: arith.constant 7
+  %za7_d = arm_sme.get_tile_id : i64
+  return
+}
+
+// -----
+
+func.func @za_h_overlapping_za_q() {
+  %za0_h = arm_sme.get_tile_id : i16
+  %za0_q = arm_sme.get_tile_id : i128
+  %za2_q = arm_sme.get_tile_id : i128
+  %za4_q = arm_sme.get_tile_id : i128
+  %za6_q = arm_sme.get_tile_id : i128
+  %za8_q = arm_sme.get_tile_id : i128
+  %za10_q = arm_sme.get_tile_id : i128
+  %za12_q = arm_sme.get_tile_id : i128
+  %za14_q = arm_sme.get_tile_id : i128
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za0_s
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 34952 : i32}
+func.func @za0_s() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_s = arm_sme.get_tile_id : i32
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_s
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_s() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_s = arm_sme.get_tile_id : i32
+  // CHECK-NEXT: arith.constant 1
+  %za1_s = arm_sme.get_tile_id : i32
+  // CHECK-NEXT: arith.constant 2
+  %za2_s = arm_sme.get_tile_id : i32
+  // CHECK-NEXT: arith.constant 3
+  %za3_s = arm_sme.get_tile_id : i32
+  return
+}
+
+// -----
+
+func.func @za_s__out_of_tiles() {
+  %za0_s = arm_sme.get_tile_id : i32
+  %za1_s = arm_sme.get_tile_id : i32
+  %za2_s = arm_sme.get_tile_id : i32
+  %za3_s = arm_sme.get_tile_id : i32
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i32
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_s_overlapping_za_d
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_s_overlapping_za_d() {
+  // ZA0.Q, ZA4.Q, ZA8.Q, ZA12.Q
+  // CHECK-NEXT: arith.constant 0
+  %za0_s = arm_sme.get_tile_id : i32
+  // ZA1.Q, ZA5.Q, ZA9.Q, ZA13.Q
+  // CHECK-NEXT: arith.constant 1
+  %za1_s = arm_sme.get_tile_id : i32
+  // ZA2.Q, ZA6.Q, ZA10.Q, ZA14.Q
+  // CHECK-NEXT: arith.constant 2
+  %za2_s = arm_sme.get_tile_id : i32
+  // ZA3.Q, ZA11.Q
+  // CHECK-NEXT: arith.constant 3
+  %za3_d = arm_sme.get_tile_id : i64
+  // ZA7.Q, ZA15.Q
+  // CHECK-NEXT: arith.constant 7
+  %za7_d = arm_sme.get_tile_id : i64
+  return
+}
+
+// -----
+
+func.func @za_s_overlapping_za_q() {
+  %za0_s = arm_sme.get_tile_id : i32
+  %za1_q = arm_sme.get_tile_id : i128
+  %za2_q = arm_sme.get_tile_id : i128
+  %za3_q = arm_sme.get_tile_id : i128
+  %za5_q = arm_sme.get_tile_id : i128
+  %za6_q = arm_sme.get_tile_id : i128
+  %za7_q = arm_sme.get_tile_id : i128
+  %za9_q = arm_sme.get_tile_id : i128
+  %za10_q = arm_sme.get_tile_id : i128
+  %za11_q = arm_sme.get_tile_id : i128
+  %za13_q = arm_sme.get_tile_id : i128
+  %za14_q = arm_sme.get_tile_id : i128
+  %za15_q = arm_sme.get_tile_id : i128
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za0_d
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 32896 : i32}
+func.func @za0_d() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_d = arm_sme.get_tile_id : i64
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_d
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_d() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 1
+  %za1_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 2
+  %za2_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 3
+  %za3_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 4
+  %za4_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 5
+  %za5_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 6
+  %za6_d = arm_sme.get_tile_id : i64
+  // CHECK-NEXT: arith.constant 7
+  %za7_d = arm_sme.get_tile_id : i64
+  return
+}
+
+// -----
+
+func.func @za_d__out_of_tiles() {
+  %za0_d = arm_sme.get_tile_id : i64
+  %za1_d = arm_sme.get_tile_id : i64
+  %za2_d = arm_sme.get_tile_id : i64
+  %za3_d = arm_sme.get_tile_id : i64
+  %za4_d = arm_sme.get_tile_id : i64
+  %za5_d = arm_sme.get_tile_id : i64
+  %za6_d = arm_sme.get_tile_id : i64
+  %za7_d = arm_sme.get_tile_id : i64
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i64
+  return
+}
+
+// -----
+
+func.func @za_d_overlapping_za_q() {
+  %za0_d = arm_sme.get_tile_id : i64
+  %za1_q = arm_sme.get_tile_id : i128
+  %za2_q = arm_sme.get_tile_id : i128
+  %za3_q = arm_sme.get_tile_id : i128
+  %za4_q = arm_sme.get_tile_id : i128
+  %za5_q = arm_sme.get_tile_id : i128
+  %za6_q = arm_sme.get_tile_id : i128
+  %za7_q = arm_sme.get_tile_id : i128
+  %za9_q = arm_sme.get_tile_id : i128
+  %za10_q = arm_sme.get_tile_id : i128
+  %za11_q = arm_sme.get_tile_id : i128
+  %za12_q = arm_sme.get_tile_id : i128
+  %za13_q = arm_sme.get_tile_id : i128
+  %za14_q = arm_sme.get_tile_id : i128
+  %za15_q = arm_sme.get_tile_id : i128
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za0_q
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 32768 : i32}
+func.func @za0_q() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_q = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+// CHECK-LABEL: za_q
+// CHECK-SAME: attributes {arm_sme.tiles_in_use = 65535 : i32}
+func.func @za_q() {
+  // CHECK-NEXT: arith.constant 0
+  %za0_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 1
+  %za1_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 2
+  %za2_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 3
+  %za3_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 4
+  %za4_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 5
+  %za5_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 6
+  %za6_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 7
+  %za7_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 8
+  %za8_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 9
+  %za9_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 10
+  %za10_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 11
+  %za11_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 12
+  %za12_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 13
+  %za13_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 14
+  %za14_q = arm_sme.get_tile_id : i128
+  // CHECK-NEXT: arith.constant 15
+  %za15_q = arm_sme.get_tile_id : i128
+  return
+}
+
+// -----
+
+func.func @za_q__out_of_tiles() {
+  %za0_q = arm_sme.get_tile_id : i128
+  %za1_q = arm_sme.get_tile_id : i128
+  %za2_q = arm_sme.get_tile_id : i128
+  %za3_q = arm_sme.get_tile_id : i128
+  %za4_q = arm_sme.get_tile_id : i128
+  %za5_q = arm_sme.get_tile_id : i128
+  %za6_q = arm_sme.get_tile_id : i128
+  %za7_q = arm_sme.get_tile_id : i128
+  %za8_q = arm_sme.get_tile_id : i128
+  %za9_q = arm_sme.get_tile_id : i128
+  %za10_q = arm_sme.get_tile_id : i128
+  %za11_q = arm_sme.get_tile_id : i128
+  %za12_q = arm_sme.get_tile_id : i128
+  %za13_q = arm_sme.get_tile_id : i128
+  %za14_q = arm_sme.get_tile_id : i128
+  %za15_q = arm_sme.get_tile_id : i128
+  // expected-error@+2 {{failed to legalize operation 'arm_sme.get_tile_id' that was explicitly marked illegal}}
+  // expected-error@+1 {{ran out of SME virtual tiles!}}
+  %next_tile = arm_sme.get_tile_id : i128
+  return
+}

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/vector-ops.mlir

-// RUN: mlir-opt %s -convert-vector-to-arm-sme -enable-arm-streaming="mode=locally enable-za" \
-// RUN:   -convert-vector-to-llvm="enable-arm-sme" -test-lower-to-llvm | \
+// RUN: mlir-opt %s -enable-arm-streaming="mode=locally enable-za" \
+// RUN:   -convert-vector-to-arm-sme -convert-vector-to-llvm="enable-arm-sme" \
+// RUN:   -allocate-arm-sme-tiles -test-lower-to-llvm | \
 // RUN: mlir-translate -mlir-to-llvmir | \
 // RUN: %lli_aarch64_cmd --march=aarch64 --mattr="+sve,+sme" \
 // RUN:   --entry-function=entry \
 // RUN:   --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 func.func @entry() -> i32 {
   %c0 = arith.constant 0 : index