Fixups

Author: MacDue

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

@@ -24,7 +24,9 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+namespace mlir::arm_sme {
 #include "mlir/Dialect/ArmSME/IR/ArmSMEOpInterfaces.h.inc"
+}
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/ArmSME/IR/ArmSMEAttrDefs.h.inc"

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

@@ -40,7 +40,7 @@ def ArmSMETileType : I32EnumAttr<"ArmSMETileType", "Arm SME tile type",
 def ArmSMETileOpInterface : OpInterface<"ArmSMETileOpInterface"> {
   let description = [{
     An interface for operations that use or allocate Arm SME tiles. These
-    operations need to be assigned a tile ID an i32 attribute, which specifies
+    operations need to be assigned a tile ID, an i32 attribute, which specifies
     which virtual tile within the ZA storage to use. The number of tiles
     available depends on the type of the tile. This is summarized below:
 
@@ -52,7 +52,7 @@ def ArmSMETileOpInterface : OpInterface<"ArmSMETileOpInterface"> {
     | `vector<[2]x[2]xi64>` or `vector<[2]x[2]xf64>`                          | 0 to 7 (inclusive)  |
     | `vector<[1]x[1]xi128>`                                                  | 0 to 15 (inclusive) |
 
-    Operations that allocate a new tiles (such as arm_sme.get_tile), are used as
+    Operations that allocate a new tile (such as arm_sme.get_tile), are used as
     the roots for tile allocation, with all operations that (transitively)
     depend on a root being assigned the same tile ID.
   }];
@@ -71,7 +71,10 @@ def ArmSMETileOpInterface : OpInterface<"ArmSMETileOpInterface"> {
       }]
     >,
     InterfaceMethod<
-      "Returns the (possibly null) tile ID assigned to this operation.",
+      [{
+        Returns the tile ID assigned to this operation. This will be null before
+        tile allocation.
+      }],
       /*returnType=*/"mlir::IntegerAttr",
       /*methodName=*/"getTileId",
       /*arguments=*/(ins),
@@ -82,13 +85,16 @@ def ArmSMETileOpInterface : OpInterface<"ArmSMETileOpInterface"> {
       }]
     >,
     InterfaceMethod<
-      "The type of tile this operation allocates (or none)",
+      [{
+        The type of tile this operation allocates. Returns none (std::nullopt)
+        if this operation does not allocate a tile.
+      }],
       /*returnType=*/"std::optional<::mlir::arm_sme::ArmSMETileType>",
       /*methodName=*/"getAllocatedTileType",
       /*arguments=*/(ins),
       /*methodBody=*/[{}],
       /*defaultImpl=*/ [{
-        // Do not allocate a new tile.
+        // This operation does not allocate a tile.
         return std::nullopt;
       }]
     >
@@ -104,14 +110,16 @@ def ArmSMETileOpInterface : OpInterface<"ArmSMETileOpInterface"> {
       return op;
     }
 
-    // A helper to replace this operation and forward any tile ID.
+    // A helper to replace this operation and forward its tile ID (if present).
     template<typename T, typename... Args>
     T replaceWithAndForwardTileId(::mlir::RewriterBase& rewriter, Args &&...args) {
       auto newOp = createOpAndForwardTileId<T>(rewriter, $_op.getLoc(), std::forward<Args>(args)...);
       rewriter.replaceOp($_op, newOp);
       return newOp;
     }
   }];
+
+  let verify = [{ return ::mlir::arm_sme::verifyOperationHasValidTileId($_op); }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -222,11 +230,11 @@ def ArmSME_CombiningKindAttr : EnumAttr<ArmSME_Dialect, CombiningKind,
 class ArmSME_Op<string mnemonic, list<Trait> traits = []> :
   Op<ArmSME_Dialect, mnemonic, traits> {}
 
-def GetTile : ArmSME_Op<"get_tile", [ArmSMETileOpInterface]> {
+def GetTileOp : ArmSME_Op<"get_tile", [ArmSMETileOpInterface]> {
   let summary = "Returns a SME virtual tile";
   let description = [{
     Allocates a new SME "virtual tile" within a function. The contents of the
-    tile returned from this operation undefined.
+    tile returned from this operation are undefined.
 
     Example 1:
 
@@ -264,12 +272,12 @@ def GetTile : ArmSME_Op<"get_tile", [ArmSMETileOpInterface]> {
   }];
 }
 
-def MaterializeSSATile : ArmSME_Op<"materialize_ssa_tile", [Pure]> {
+def MaterializeSSATileOp : ArmSME_Op<"materialize_ssa_tile", [Pure]> {
   let summary = "SME tile placeholder";
   let description = [{
     A placeholder to preserve dataflow while lowering to SME intrinsics (which
-    do not take or return tile values). This operation is intended to be DCE'd
-    once all ArmSME operations have been lowered.
+    do not take or return SME virtual tile values). This operation is intended
+    to be DCE'd once all ArmSME operations have been lowered.
 
     This operation is not intended to be used outside of the ArmSME -> LLVM
     conversion.

mlir/include/mlir/Dialect/ArmSME/Utils/Utils.h

@@ -35,8 +35,12 @@ bool isValidSMETileElementType(Type type);
 /// otherwise.
 bool isValidSMETileVectorType(VectorType vType);
 
+/// Returns the type of SME tile this vector type corresponds to or none.
 std::optional<ArmSMETileType> getSMETileType(VectorType);
 
+/// Verifies the tile ID (if set) on this tile operation is valid.
+LogicalResult verifyOperationHasValidTileId(Operation *);
+
 } // namespace mlir::arm_sme
 
 #endif // MLIR_DIALECT_ARMSME_UTILS_UTILS_H_

mlir/lib/Conversion/ArmSMEToLLVM/ArmSMEToLLVM.cpp

@@ -32,21 +32,21 @@ using namespace mlir;
 
 namespace {
 
-IntegerAttr getTileIdOrError(ArmSMETileOpInterface op) {
+IntegerAttr getTileIdOrError(arm_sme::ArmSMETileOpInterface op) {
   auto tileId = op.getTileId();
   if (!tileId)
     op.emitOpError(
         "expected tile ID to be allocated before conversion to LLVM");
   return tileId;
 }
 
-struct GetTileConversion : public ConvertOpToLLVMPattern<arm_sme::GetTile> {
-  using ConvertOpToLLVMPattern<arm_sme::GetTile>::ConvertOpToLLVMPattern;
+struct GetTileConversion : public ConvertOpToLLVMPattern<arm_sme::GetTileOp> {
+  using ConvertOpToLLVMPattern<arm_sme::GetTileOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
-  matchAndRewrite(arm_sme::GetTile getTile, OpAdaptor,
+  matchAndRewrite(arm_sme::GetTileOp getTile, OpAdaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<arm_sme::MaterializeSSATile>(
+    rewriter.replaceOpWithNewOp<arm_sme::MaterializeSSATileOp>(
         getTile, getTile.getTileType());
     return success();
   }
@@ -140,7 +140,7 @@ struct ZeroOpConversion : public ConvertOpToLLVMPattern<arm_sme::ZeroOp> {
         loc, rewriter.getI32IntegerAttr(zeroMask));
 
     // Create a placeholder op to preserve dataflow.
-    rewriter.replaceOpWithNewOp<arm_sme::MaterializeSSATile>(
+    rewriter.replaceOpWithNewOp<arm_sme::MaterializeSSATileOp>(
         zero, zero.getVectorType());
 
     return success();
@@ -558,7 +558,7 @@ struct ConvertArmSMEToLLVMPass
 void mlir::configureArmSMEToLLVMConversionLegality(ConversionTarget &target) {
   target.addIllegalDialect<arm_sme::ArmSMEDialect>();
   target.addLegalOp<
-      arm_sme::MaterializeSSATile, arm_sme::aarch64_sme_zero,
+      arm_sme::MaterializeSSATileOp, arm_sme::aarch64_sme_zero,
       arm_sme::aarch64_sme_str, arm_sme::aarch64_sme_ld1b_horiz,
       arm_sme::aarch64_sme_ld1h_horiz, arm_sme::aarch64_sme_ld1w_horiz,
       arm_sme::aarch64_sme_ld1d_horiz, arm_sme::aarch64_sme_ld1q_horiz,

mlir/lib/Conversion/ArmSMEToLLVM/CMakeLists.txt

@@ -10,7 +10,6 @@ add_mlir_conversion_library(MLIRArmSMEToLLVM
   LINK_LIBS PUBLIC
   MLIRArmSMETransforms
   MLIRArmSMEDialect
-  MLIRArmSMEUtils
   MLIRTransforms
   MLIRLLVMCommonConversion
   MLIRLLVMDialect)

mlir/lib/Conversion/ArmSMEToSCF/ArmSMEToSCF.cpp

@@ -89,7 +89,7 @@ struct TileLoadOpConversion : public OpRewritePattern<arm_sme::TileLoadOp> {
     auto tileElementType = tileType.getElementType();
 
     // Allocate a new SME tile.
-    auto tile = tileLoadOp.createOpAndForwardTileId<arm_sme::GetTile>(
+    auto tile = tileLoadOp.createOpAndForwardTileId<arm_sme::GetTileOp>(
         rewriter, loc, tileType);
 
     // Create a loop that loads each ZA tile slice from memory.
@@ -299,7 +299,7 @@ struct TileLoadOpWithMaskAndPadNonZeroConversion
         loc, rewriter.getI32Type(), numCols);
 
     // Allocate a new SME tile.
-    auto tile = tileLoadOp.createOpAndForwardTileId<arm_sme::GetTile>(
+    auto tile = tileLoadOp.createOpAndForwardTileId<arm_sme::GetTileOp>(
         rewriter, loc, tileType);
 
     // Create a loop that loads each ZA tile slice from memory.

mlir/lib/Conversion/ArmSMEToSCF/CMakeLists.txt

@@ -9,6 +9,5 @@ add_mlir_conversion_library(MLIRArmSMEToSCF
 
   LINK_LIBS PUBLIC
   MLIRArmSMEDialect
-  MLIRArmSMEUtils
   MLIRTransforms
   )

mlir/lib/Conversion/VectorToArmSME/CMakeLists.txt

@@ -10,6 +10,5 @@ add_mlir_conversion_library(MLIRVectorToArmSME
 
   LINK_LIBS PUBLIC
   MLIRArmSMEDialect
-  MLIRArmSMEUtils
   MLIRLLVMCommonConversion
   )

mlir/lib/Conversion/VectorToArmSME/VectorToArmSME.cpp

@@ -253,7 +253,7 @@ struct ConstantOpToArmSMELowering : public OpRewritePattern<arith::ConstantOp> {
         tileSliceType, denseAttr.getSplatValue<Attribute>());
     auto constantOp1D = rewriter.create<arith::ConstantOp>(loc, denseAttr1D);
 
-    auto tile = rewriter.create<arm_sme::GetTile>(loc, tileType);
+    auto tile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
 
     auto forOp = getLoopOverTileSlices(rewriter, loc, tileElementType);
     auto tileSliceIndex = forOp.getInductionVar();
@@ -315,7 +315,7 @@ struct BroadcastOpToArmSMELowering
     else
       return failure();
 
-    auto tile = rewriter.create<arm_sme::GetTile>(loc, tileType);
+    auto tile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
 
     // Create a loop over ZA tile slices.
     auto forOp = getLoopOverTileSlices(rewriter, loc, tileElementType);
@@ -371,7 +371,7 @@ struct SplatOpToArmSMELowering : public OpRewritePattern<vector::SplatOp> {
     Value broadcastOp1D = rewriter.create<vector::BroadcastOp>(
         loc, tileSliceType, splatOp.getInput());
 
-    auto tile = rewriter.create<arm_sme::GetTile>(loc, tileType);
+    auto tile = rewriter.create<arm_sme::GetTileOp>(loc, tileType);
 
     // Next, create a loop over ZA tile slices and "move" the generated 1-d
     // vector to each slice.

mlir/lib/Dialect/ArmSME/CMakeLists.txt

@@ -1,3 +1,2 @@
 add_subdirectory(IR)
 add_subdirectory(Transforms)
-add_subdirectory(Utils)

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

@@ -1,5 +1,6 @@
 add_mlir_dialect_library(MLIRArmSMEDialect
   ArmSME.cpp
+  Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/ArmSME
@@ -15,5 +16,4 @@ add_mlir_dialect_library(MLIRArmSMEDialect
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRVectorDialect
-  MLIRArmSMEUtils
 )

mlir/lib/Dialect/ArmSME/Utils/Utils.cpp renamed to mlir/lib/Dialect/ArmSME/IR/Utils.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/ArmSME/Utils/Utils.h"
+#include "mlir/Dialect/ArmSME/IR/ArmSME.h"
 
 namespace mlir::arm_sme {
 
@@ -59,4 +60,17 @@ std::optional<ArmSMETileType> getSMETileType(VectorType type) {
   }
 }
 
+LogicalResult verifyOperationHasValidTileId(Operation *op) {
+  auto tileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op);
+  if (!tileOp)
+    return success(); // Not a tile op (no need to check).
+  auto tileId = tileOp.getTileId();
+  if (!tileId)
+    return success(); // Not having a tile ID (yet) is okay.
+  if (!tileId.getType().isSignlessInteger(32))
+    return tileOp.emitOpError("tile ID should be a 32-bit signless integer");
+  // TODO: Verify value of tile ID is in range.
+  return success();
+}
+
 } // namespace mlir::arm_sme

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

@@ -11,7 +11,6 @@ add_mlir_dialect_library(MLIRArmSMETransforms
 
   LINK_LIBS PUBLIC
   MLIRArmSMEDialect
-  MLIRArmSMEUtils
   MLIRFuncDialect
   MLIRLLVMCommonConversion
   MLIRVectorDialect

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

@@ -136,9 +136,10 @@ static ArrayRef<TileMask> getMasks(ArmSMETileType type) {
   }
 }
 
-/// Allocates a tile to 'tileId' or returns an error if there are no tiles left.
-static FailureOr<unsigned> getTile(ArmSMETileType tileType,
-                                   TileMask &tilesInUse) {
+/// Allocates and returns a tile ID. Returns an error if there are no tiles
+/// left.
+static FailureOr<unsigned> allocateTileId(ArmSMETileType tileType,
+                                          TileMask &tilesInUse) {
   auto masks = getMasks(tileType);
   for (auto [tileId, tileMask] : llvm::enumerate(masks)) {
     if ((tilesInUse & tileMask) == TileMask::kNone) {
@@ -168,7 +169,7 @@ struct AssignTileIDsPattern
     else
       tilesInUse = TileMask::kNone;
 
-    auto tileId = getTile(*tileType, tilesInUse);
+    auto tileId = allocateTileId(*tileType, tilesInUse);
     if (failed(tileId))
       return tileOp.emitError("ran out of SME virtual tiles!");
 

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

@@ -21,9 +21,9 @@
 using namespace mlir;
 
 // This is needed because these matchers are defined as overloaded functions.
-using HasOpAttrName = mlir::detail::AttrOpMatcher(StringRef);
-using HasOpName = mlir::detail::NameOpMatcher(StringRef);
-using IsConstantOp = mlir::detail::constant_op_matcher();
+using HasOpAttrName = detail::AttrOpMatcher(StringRef);
+using HasOpName = detail::NameOpMatcher(StringRef);
+using IsConstantOp = detail::constant_op_matcher();
 
 namespace test {
 #ifdef MLIR_INCLUDE_TESTS