[mlir][Arm] Fix invalid rewrite pattern API violations

[matthias-springer]

This commit fixes rewrite pattern API violations:

• Rewrite pattern must return "failure" if the IR was not modified.
• In-place op modifications must be communicated to the rewriter (updateRootInPlace).

This commit fixes test/Dialect/ArmSVE/legalize-vector-storage.mlir, test/Dialect/ArmSME/vector-ops-to-llvm.mlir,
test/Dialect/ArmSME/tile-allocation-invalid.mlir, test/Conversion/ArmSMEToLLVM/arm-sme-to-llvm.mlir, test/Conversion/ArmSMEToLLVM/tile-spills-and-fills.mlir, test/Conversion/ArmSMEToLLVM/unsupported.mlir when running with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

[llvmbot]

@llvm/pr-subscribers-mlir-sme
@llvm/pr-subscribers-mlir-sve

@llvm/pr-subscribers-mlir

Author: Matthias Springer (matthias-springer)

Changes

This commit fixes rewrite pattern API violations:

• Rewrite pattern must return "failure" if the IR was not modified.
• In-place op modifications must be communicated to the rewriter (updateRootInPlace).

This commit fixes test/Dialect/ArmSVE/legalize-vector-storage.mlir, test/Dialect/ArmSME/vector-ops-to-llvm.mlir,
test/Dialect/ArmSME/tile-allocation-invalid.mlir, test/Conversion/ArmSMEToLLVM/arm-sme-to-llvm.mlir, test/Conversion/ArmSMEToLLVM/tile-spills-and-fills.mlir, test/Conversion/ArmSMEToLLVM/unsupported.mlir when running with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

---

Full diff: https://github.com/llvm/llvm-project/pull/78246.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp (+15-6)
• (modified) mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp (+2-1)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 49ea6bb5f8614e7..b4630c834ff2428 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -232,14 +232,11 @@ struct AssignTileIDsPattern
         static_cast<TileMask>(getDiscardableIntAttr(kTilesInUseAttr));
     auto tileId = allocateTileId(*tileType, tilesInUse);
     bool tileIsInMemory = failed(tileId);
-    if (!tileIsInMemory)
-      setDiscardableIntAttr(kTilesInUseAttr, tilesInUse);
-    else {
+    if (tileIsInMemory) {
       // If we could not find a real tile ID, use an in-memory tile ID (ID >=
       // 16). A later pass will insert the necessary spills and reloads.
       tileId =
           getDiscardableIntAttr(kNextInMemoryTileIdAttr, kInMemoryTileIdBase);
-      setDiscardableIntAttr(kNextInMemoryTileIdAttr, *tileId + 1);
       tileOp->emitWarning(
           "failed to allocate SME virtual tile to operation, all tile "
           "operations will go through memory, expect degraded performance");
@@ -263,14 +260,26 @@ struct AssignTileIDsPattern
     SetVector<Operation *> dependantOps;
     findDependantOps(tileOp->getResult(0), dependantOps);
     auto tileIDAttr = rewriter.getI32IntegerAttr(*tileId);
-    rewriter.updateRootInPlace(tileOp, [&] { tileOp.setTileId(tileIDAttr); });
     for (auto *op : dependantOps) {
       if (auto dependantTileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op)) {
         auto currentTileId = dependantTileOp.getTileId();
         if (currentTileId && unsigned(currentTileId.getInt()) != tileId)
           return dependantTileOp.emitOpError(
               "already assigned different SME virtual tile!");
-        dependantTileOp.setTileId(tileIDAttr);
+      }
+    }
+
+    // Rewrite IR.
+    if (!tileIsInMemory)
+      setDiscardableIntAttr(kTilesInUseAttr, tilesInUse);
+    else {
+      setDiscardableIntAttr(kNextInMemoryTileIdAttr, *tileId + 1);
+    }
+    rewriter.updateRootInPlace(tileOp, [&] { tileOp.setTileId(tileIDAttr); });
+    for (auto *op : dependantOps) {
+      if (auto dependantTileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op)) {
+        rewriter.updateRootInPlace(
+            dependantTileOp, [&]() { dependantTileOp.setTileId(tileIDAttr); });
       }
     }
 
diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
index bee1f3659753b78..ed695db0372f28f 100644
--- a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
+++ b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp
@@ -106,7 +106,8 @@ struct RelaxScalableVectorAllocaAlignment
 
     // Set alignment based on the defaults for SVE vectors and predicates.
     unsigned aligment = vectorType.getElementType().isInteger(1) ? 2 : 16;
-    allocaOp.setAlignment(aligment);
+    rewriter.updateRootInPlace(allocaOp,
+                               [&]() { allocaOp.setAlignment(aligment); });
 
     return success();
   }

[MacDue]

Just a clarification: This has to be done for an update of any operation within a rewrite? If so, it's not immediately clear to me what "root" refers to in the API (and why it's not just updateOpInPlace)?

[matthias-springer]

Yes, that's right. We should indeed change the name.

[matthias-springer]

Here's an attempt to improve the situation: #78260

[MacDue]

LGTM, thanks for the fixes!