[Mlir-commits] [mlir] [mlir][Transforms] Fix use-after-free in #82474 (PR #82504)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/82504

When a `ModifyOperationRewrite` is committed, the operation may already have been erased, so `OperationName` must be cached in the rewrite object.

Note: This will no longer be needed with #81757, which adds a "cleanup" method to `IRRewrite`.

>From 648cc44da386bdcbb5abd87f79af6a274c6f6de8 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 21 Feb 2024 16:24:38 +0000
Subject: [PATCH] [mlir][Transforms] Fix use-after-free in #82474

When a `ModifyOperationRewrite` is committed, the operation may already
have been erased, so `OperationName` must be cached in the rewrite
object.

Note: This will no longer be needed with #81757, which adds a "cleanup"
method to `IRRewrite`.
---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 88709bb2618744..4989ddc3ec94fb 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -965,14 +965,14 @@ class ModifyOperationRewrite : public OperationRewrite {
   ModifyOperationRewrite(ConversionPatternRewriterImpl &rewriterImpl,
                          Operation *op)
       : OperationRewrite(Kind::ModifyOperation, rewriterImpl, op),
-        loc(op->getLoc()), attrs(op->getAttrDictionary()),
+        name(op->getName()), loc(op->getLoc()), attrs(op->getAttrDictionary()),
         operands(op->operand_begin(), op->operand_end()),
         successors(op->successor_begin(), op->successor_end()) {
     if (OpaqueProperties prop = op->getPropertiesStorage()) {
       // Make a copy of the properties.
       propertiesStorage = operator new(op->getPropertiesStorageSize());
       OpaqueProperties propCopy(propertiesStorage);
-      op->getName().initOpProperties(propCopy, /*init=*/prop);
+      name.initOpProperties(propCopy, /*init=*/prop);
     }
   }
 
@@ -988,7 +988,9 @@ class ModifyOperationRewrite : public OperationRewrite {
   void commit() override {
     if (propertiesStorage) {
       OpaqueProperties propCopy(propertiesStorage);
-      op->getName().destroyOpProperties(propCopy);
+      // Note: The operation may have been erased in the mean time, so
+      // OperationName must be stored in this object.
+      name.destroyOpProperties(propCopy);
       operator delete(propertiesStorage);
       propertiesStorage = nullptr;
     }
@@ -1003,13 +1005,14 @@ class ModifyOperationRewrite : public OperationRewrite {
     if (propertiesStorage) {
       OpaqueProperties propCopy(propertiesStorage);
       op->copyProperties(propCopy);
-      op->getName().destroyOpProperties(propCopy);
+      name.destroyOpProperties(propCopy);
       operator delete(propertiesStorage);
       propertiesStorage = nullptr;
     }
   }
 
 private:
+  OperationName name;
   LocationAttr loc;
   DictionaryAttr attrs;
   SmallVector<Value, 8> operands;