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

[Cullen Rhodes]

================
@@ -209,28 +210,40 @@ struct AssignTileIDsPattern
     if (tileOp.getTileId())
       return failure();
 
+    auto func = tileOp->getParentOfType<FunctionOpInterface>();
+    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) {
+      rewriter.updateRootInPlace(tileOp, [&] {
+        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());
-
+    TileMask tilesInUse =
+        static_cast<TileMask>(getDiscardableIntAttr(kTilesInUseAttr));
     auto tileId = allocateTileId(*tileType, tilesInUse);
-    if (failed(tileId))
-      return tileOp.emitError("ran out of SME virtual tiles!");
-
-    rewriter.updateRootInPlace(func, [&]() {
-      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);
+    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).
----------------
c-rhodes wrote:

```suggestion
      // If we could not find a real tile, use a virtual tile ID (ID >= 16).
```

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