[Mlir-commits] [mlir] [mlir][ArmSME] Add rudimentary support for tile spills to the stack (PR #76086)

[Andrzej Warzyński]

================
@@ -200,36 +202,52 @@ static void findDependantOps(Value rootValue,
         });
   }
 }
-
 struct AssignTileIDsPattern
     : public OpInterfaceRewritePattern<ArmSMETileOpInterface> {
   using OpInterfaceRewritePattern::OpInterfaceRewritePattern;
   LogicalResult matchAndRewrite(ArmSMETileOpInterface tileOp,
                                 PatternRewriter &rewriter) const override {
+    auto func = tileOp->getParentOfType<FunctionOpInterface>();
     if (tileOp.getTileId())
       return failure();
 
+    auto getDiscardableIntAttr = [&](StringRef name, unsigned defaultVal = 0) {
+      if (auto attr = llvm::dyn_cast_or_null<IntegerAttr>(
+              func->getDiscardableAttr(name)))
+        return unsigned(attr.getInt());
+      return defaultVal;
+    };
+
+    auto setDiscardableIntAttr = [&](StringRef name, auto value) {
+      func->setDiscardableAttr(name,
+                               rewriter.getI32IntegerAttr((unsigned)value));
+    };
+
     std::optional<ArmSMETileType> tileType = tileOp.getAllocatedTileType();
     if (!tileType)
       return rewriter.notifyMatchFailure(tileOp, "op does not allocate a tile");
 
-    auto func = tileOp->getParentOfType<FunctionOpInterface>();
-    TileMask tilesInUse = TileMask::kNone;
-    if (auto tilesInUseAttr = llvm::dyn_cast_or_null<IntegerAttr>(
-            func->getDiscardableAttr(kTilesInUseAttr)))
-      tilesInUse = static_cast<TileMask>(tilesInUseAttr.getInt());
-
-    auto tileId = allocateTileId(*tileType, tilesInUse);
-    if (failed(tileId))
-      return tileOp.emitError("ran out of SME virtual tiles!");
-
-    func->setDiscardableAttr(kTilesInUseAttr,
-                             rewriter.getI32IntegerAttr((unsigned)tilesInUse));
-
     // Find all the ops that (transitively) depend on this tile.
     SetVector<Operation *> dependantOps;
     findDependantOps(tileOp->getResult(0), dependantOps);
 
+    TileMask tilesInUse =
+        static_cast<TileMask>(getDiscardableIntAttr(kTilesInUseAttr));
+    auto tileId = allocateTileId(*tileType, tilesInUse);
+    bool tileIsInMemory = failed(tileId);
+    if (!tileIsInMemory)
+      setDiscardableIntAttr(kTilesInUseAttr, tilesInUse);
+    else {
+      // If we could not find a real tile, set use a virtual tile ID (ID >= 16).
+      // A later pass will insert the necessary spills and reloads.
+      tileId = getDiscardableIntAttr(kNextTileMemoryIndex, kInMemoryTileIdBase);
+      setDiscardableIntAttr(kNextTileMemoryIndex, *tileId + 1);
+      tileOp->emitWarning(
+          "failed to allocate physical tile to operation, all tile "
----------------
banach-space wrote:

[nit] "physical tile" is a bit misleading - there are no actual tiles in hardware ("virtual tiles" are more like "views" into ZA). In any case, "SME virtual tile" would be more consistent with the terminology we have been using so far.

https://github.com/llvm/llvm-project/pull/76086