[mlir][Transforms] Fix use-after-free when accessing replaced block args

[matthias-springer]

This commit fixes a bug in a dialect conversion. Currently, when a block is replaced via a signature conversion, the block is erased during the "commit" phase. This is problematic because the block arguments may still be referenced internal data structures of the dialect conversion (mapping). Blocks should be treated same as ops: they should be erased during the "cleanup" phase.

Note: The test case fails without this fix when running with ASAN, but may pass when running without ASAN.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

This commit fixes a bug in a dialect conversion. Currently, when a block is replaced via a signature conversion, the block is erased during the "commit" phase. This is problematic because the block arguments may still be referenced internal data structures of the dialect conversion (mapping). Blocks should be treated same as ops: they should be erased during the "cleanup" phase.

Note: The test case fails without this fix when running with ASAN, but may pass when running without ASAN.

---

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

3 Files Affected:

• (modified) mlir/lib/Transforms/Utils/DialectConversion.cpp (+20-10)
• (modified) mlir/test/Transforms/test-legalizer.mlir (+12)
• (modified) mlir/test/lib/Dialect/Test/TestOps.td (+2-2)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 26899301eb742e..7846f1ab56811a 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -207,13 +207,13 @@ class IRRewrite {
   /// Roll back the rewrite. Operations may be erased during rollback.
   virtual void rollback() = 0;
 
-  /// Commit the rewrite. Operations may be unlinked from their blocks during
-  /// the commit phase, but they must not be erased yet. This is because
-  /// internal dialect conversion state (such as `mapping`) may still be using
-  /// them. Operations must be erased during cleanup.
+  /// Commit the rewrite. Operations/blocks may be unlinked during the commit
+  /// phase, but they must not be erased yet. This is because internal dialect
+  /// conversion state (such as `mapping`) may still be using them. Operations/
+  /// blocks must be erased during cleanup.
   virtual void commit() {}
 
-  /// Cleanup operations. Cleanup is called after commit.
+  /// Cleanup operations/blocks. Cleanup is called after commit.
   virtual void cleanup() {}
 
   Kind getKind() const { return kind; }
@@ -282,9 +282,9 @@ class CreateBlockRewrite : public BlockRewrite {
 };
 
 /// Erasure of a block. Block erasures are partially reflected in the IR. Erased
-/// blocks are immediately unlinked, but only erased when the rewrite is
-/// committed. This makes it easier to rollback a block erasure: the block is
-/// simply inserted into its original location.
+/// blocks are immediately unlinked, but only erased during cleanup. This makes
+/// it easier to rollback a block erasure: the block is simply inserted into its
+/// original location.
 class EraseBlockRewrite : public BlockRewrite {
 public:
   EraseBlockRewrite(ConversionPatternRewriterImpl &rewriterImpl, Block *block,
@@ -297,7 +297,8 @@ class EraseBlockRewrite : public BlockRewrite {
   }
 
   ~EraseBlockRewrite() override {
-    assert(!block && "rewrite was neither rolled back nor committed");
+    assert(!block &&
+           "rewrite was neither rolled back nor committed/cleaned up");
   }
 
   void rollback() override {
@@ -312,7 +313,7 @@ class EraseBlockRewrite : public BlockRewrite {
     block = nullptr;
   }
 
-  void commit() override {
+  void cleanup() override {
     // Erase the block.
     assert(block && "expected block");
     assert(block->empty() && "expected empty block");
@@ -440,6 +441,8 @@ class BlockTypeConversionRewrite : public BlockRewrite {
 
   void commit() override;
 
+  void cleanup() override;
+
   void rollback() override;
 
 private:
@@ -986,7 +989,9 @@ void BlockTypeConversionRewrite::commit() {
           rewriterImpl.mapping.lookupOrDefault(castValue, origArg.getType()));
     }
   }
+}
 
+void BlockTypeConversionRewrite::cleanup() {
   assert(origBlock->empty() && "expected empty block");
   origBlock->dropAllDefinedValueUses();
   delete origBlock;
@@ -1484,6 +1489,11 @@ void ConversionPatternRewriterImpl::notifyBlockInserted(
     Block *block, Region *previous, Region::iterator previousIt) {
   assert(!wasOpReplaced(block->getParentOp()) &&
          "attempting to insert into a region within a replaced/erased op");
+  LLVM_DEBUG({
+    logger.startLine() << "** Insert Block into : '"
+                       << block->getParentOp()->getName() << "'("
+                       << block->getParentOp() << ")\n";
+  });
 
   if (!previous) {
     // This is a newly created block.
diff --git a/mlir/test/Transforms/test-legalizer.mlir b/mlir/test/Transforms/test-legalizer.mlir
index 62d776cd7573ee..ccdc9fe78ea0d3 100644
--- a/mlir/test/Transforms/test-legalizer.mlir
+++ b/mlir/test/Transforms/test-legalizer.mlir
@@ -346,3 +346,15 @@ func.func @test_properties_rollback() {
       {modify_inplace}
   "test.return"() : () -> ()
 }
+
+// -----
+
+//      CHECK: func.func @use_of_replaced_bbarg(
+// CHECK-SAME:     %[[arg0:.*]]: f64)
+//      CHECK:   "test.valid"(%[[arg0]])
+func.func @use_of_replaced_bbarg(%arg0: i64) {
+  %0 = "test.op_with_region_fold"(%arg0) ({
+    "foo.op_with_region_terminator"() : () -> ()
+  }) : (i64) -> (i64)
+  "test.invalid"(%0) : (i64) -> ()
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 91ce0af9cd7fd5..dfd2f21a5ea249 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1270,8 +1270,8 @@ def TestOpWithRegionFoldNoMemoryEffect : TEST_Op<
 
 // Op for testing folding of outer op with inner ops.
 def TestOpWithRegionFold : TEST_Op<"op_with_region_fold"> {
-  let arguments = (ins I32:$operand);
-  let results = (outs I32);
+  let arguments = (ins AnyType:$operand);
+  let results = (outs AnyType);
   let regions = (region SizedRegion<1>:$region);
   let hasFolder = 1;
 }

[jpienaar]

Reviewed 3 of 3 files at r1, all commit messages.
Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @matthias-springer)