[Mlir-commits] [mlir] db8a119 - [mlir][ArmSME] Fix invalid rewriter API usage (#76123)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2023-12-21T17:39:36+09:00
New Revision: db8a119e8f04dfccd40ab7675c62ada81423e5c2

URL: https://github.com/llvm/llvm-project/commit/db8a119e8f04dfccd40ab7675c62ada81423e5c2
DIFF: https://github.com/llvm/llvm-project/commit/db8a119e8f04dfccd40ab7675c62ada81423e5c2.diff

LOG: [mlir][ArmSME] Fix invalid rewriter API usage (#76123)

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.

Added: 
    

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

Removed: 
    


################################################################################
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 
diff erent SME virtual tile!");
-        tileOp.setTileId(tileIDAttr);
+        rewriter.updateRootInPlace(tileOp,
+                                   [&]() { tileOp.setTileId(tileIDAttr); });
       }
     }