[mlir][ArmSME] Fix invalid rewriter API usage

[matthias-springer]

When operations are modified in-place, the rewriter must be notified. This commit fixes mlir/test/Conversion/ArmSMEToLLVM/unsupported.mlir, mlir/test/Dialect/ArmSME/tile-zero-masks.mlir and mlir/test/Dialect/ArmSME/vector-ops-to-llvm.mlir when running with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS enabled.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sme

Author: Matthias Springer (matthias-springer)

Changes

When operations are modified in-place, the rewriter must be notified. This commit fixes mlir/test/Conversion/ArmSMEToLLVM/unsupported.mlir and mlir/test/Dialect/ArmSME/vector-ops-to-llvm.mlir when running with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS enabled.

---

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

1 Files Affected:

• (modified) mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp (+7-4)

diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
index 597846e31e218e..8aa51f352f822d 100644
--- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
+++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp
@@ -223,8 +223,10 @@ struct AssignTileIDsPattern
     if (failed(tileId))
       return tileOp.emitError("ran out of SME virtual tiles!");
 
-    func->setDiscardableAttr(kTilesInUseAttr,
-                             rewriter.getI32IntegerAttr((unsigned)tilesInUse));
+    rewriter.updateRootInPlace(func, [&]() {
+      func->setDiscardableAttr(
+          kTilesInUseAttr, rewriter.getI32IntegerAttr((unsigned)tilesInUse));
+    });
 
     // Find all the ops that (transitively) depend on this tile.
     SetVector<Operation *> dependantOps;
@@ -245,14 +247,15 @@ struct AssignTileIDsPattern
     // scf.if, and moving the contents of %tileA or %tileB to result tile (based
     // on the %some_cond).
     auto tileIDAttr = rewriter.getI32IntegerAttr(*tileId);
-    tileOp.setTileId(tileIDAttr);
+    rewriter.updateRootInPlace(tileOp, [&]() { tileOp.setTileId(tileIDAttr); });
     for (auto *op : dependantOps) {
       if (auto tileOp = llvm::dyn_cast<ArmSMETileOpInterface>(op)) {
         auto currentTileId = tileOp.getTileId();
         if (currentTileId && unsigned(currentTileId.getInt()) != tileId)
           return tileOp.emitOpError(
               "already assigned different SME virtual tile!");
-        tileOp.setTileId(tileIDAttr);
+        rewriter.updateRootInPlace(tileOp,
+                                   [&]() { tileOp.setTileId(tileIDAttr); });
       }
     }
 

[c-rhodes]

LGTM, thanks for the fixes! Wasn't aware of MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS, good to know.

[banach-space]

Thanks for the fix Matthias!