From 43cfc9176fb37e59a30195064ba54ad095ce4284 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 16 Jan 2024 09:44:11 +0000 Subject: [PATCH 1/4] [mlir][Arm] Fix invalid rewrite pattern API violations 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`. --- .../ArmSME/Transforms/TileAllocation.cpp | 21 +++++++++++++------ .../Transforms/LegalizeVectorStorage.cpp | 3 ++- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp index 49ea6bb5f8614..b4630c834ff24 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(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 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(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(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 bee1f3659753b..ed695db0372f2 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(); } From 18c5ef0f0abbc785310a0a48c6cc5b1f33f3f9d6 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 16 Jan 2024 12:06:02 +0100 Subject: [PATCH 2/4] Update mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp Co-authored-by: Benjamin Maxwell --- mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp index b4630c834ff24..30df0f0db8028 100644 --- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp +++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp @@ -272,9 +272,8 @@ struct AssignTileIDsPattern // Rewrite IR. if (!tileIsInMemory) setDiscardableIntAttr(kTilesInUseAttr, tilesInUse); - else { + else setDiscardableIntAttr(kNextInMemoryTileIdAttr, *tileId + 1); - } rewriter.updateRootInPlace(tileOp, [&] { tileOp.setTileId(tileIDAttr); }); for (auto *op : dependantOps) { if (auto dependantTileOp = llvm::dyn_cast(op)) { From 711f2bdb925f58301c51891947ac11de1f4fa555 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 16 Jan 2024 12:06:15 +0100 Subject: [PATCH 3/4] Update mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp Co-authored-by: Benjamin Maxwell --- mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp index 30df0f0db8028..4d49efecbe05c 100644 --- a/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp +++ b/mlir/lib/Dialect/ArmSME/Transforms/TileAllocation.cpp @@ -278,7 +278,7 @@ struct AssignTileIDsPattern for (auto *op : dependantOps) { if (auto dependantTileOp = llvm::dyn_cast(op)) { rewriter.updateRootInPlace( - dependantTileOp, [&]() { dependantTileOp.setTileId(tileIDAttr); }); + dependantTileOp, [&] { dependantTileOp.setTileId(tileIDAttr); }); } } From 436a16ed02c71d905e1be9858224a796ff496902 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 16 Jan 2024 12:06:21 +0100 Subject: [PATCH 4/4] Update mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp Co-authored-by: Benjamin Maxwell --- mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp index ed695db0372f2..bf627d95ae557 100644 --- a/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp +++ b/mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp @@ -107,7 +107,7 @@ struct RelaxScalableVectorAllocaAlignment // Set alignment based on the defaults for SVE vectors and predicates. unsigned aligment = vectorType.getElementType().isInteger(1) ? 2 : 16; rewriter.updateRootInPlace(allocaOp, - [&]() { allocaOp.setAlignment(aligment); }); + [&] { allocaOp.setAlignment(aligment); }); return success(); }