[Mlir-commits] [mlir] [MLIR] Fix use-after-free in Remark by owning string data (PR #179889)

[Guray Ozen]

https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/179889

>From 94458c2df4453c04f8a9d497489237701cbc939b Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Thu, 5 Feb 2026 10:01:16 +0100
Subject: [PATCH] [MLIR] Fix use-after-free in Remark by owning string data

Change Remark's StringRef members to std::string to ensure remarks own
their data, preventing dangling pointers when used with RemarkEmittingPolicyFinal.
---
 mlir/include/mlir/IR/Remarks.h   | 27 ++++++-----
 mlir/unittests/IR/RemarkTest.cpp | 79 ++++++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+), 11 deletions(-)

diff --git a/mlir/include/mlir/IR/Remarks.h b/mlir/include/mlir/IR/Remarks.h
index 3102542731b33..0a787f20c5704 100644
--- a/mlir/include/mlir/IR/Remarks.h
+++ b/mlir/include/mlir/IR/Remarks.h
@@ -89,9 +89,10 @@ class Remark {
 public:
   Remark(RemarkKind remarkKind, DiagnosticSeverity severity, Location loc,
          RemarkOpts opts)
-      : remarkKind(remarkKind), functionName(opts.functionName), loc(loc),
-        categoryName(opts.categoryName), subCategoryName(opts.subCategoryName),
-        remarkName(opts.remarkName) {
+      : remarkKind(remarkKind), functionName(opts.functionName.str()), loc(loc),
+        categoryName(opts.categoryName.str()),
+        subCategoryName(opts.subCategoryName.str()),
+        remarkName(opts.remarkName.str()) {
     if (!categoryName.empty() && !subCategoryName.empty()) {
       (llvm::Twine(categoryName) + ":" + subCategoryName)
           .toStringRef(fullCategoryName);
@@ -183,21 +184,25 @@ class Remark {
   /// Keeps the MLIR diagnostic kind, which is used to determine the
   /// diagnostic kind in the LLVM remark streamer.
   RemarkKind remarkKind;
-  /// Name of the convering function like interface
-  StringRef functionName;
+  /// Name of the covering function like interface.
+  /// Stored as std::string to ensure the Remark owns its data.
+  std::string functionName;
 
   Location loc;
-  /// Sub category passname e.g., "Unroll" or "UnrollAndJam"
-  StringRef categoryName;
+  /// Category name e.g., "Unroll" or "UnrollAndJam".
+  /// Stored as std::string to ensure the Remark owns its data.
+  std::string categoryName;
 
-  /// Sub category name "Loop Optimizer"
-  StringRef subCategoryName;
+  /// Sub category name e.g., "Loop Optimizer".
+  /// Stored as std::string to ensure the Remark owns its data.
+  std::string subCategoryName;
 
   /// Combined name for category and sub-category
   SmallString<64> fullCategoryName;
 
-  /// Remark identifier
-  StringRef remarkName;
+  /// Remark identifier.
+  /// Stored as std::string to ensure the Remark owns its data.
+  std::string remarkName;
 
   /// Args collected via the streaming interface.
   SmallVector<Arg, 4> args;
diff --git a/mlir/unittests/IR/RemarkTest.cpp b/mlir/unittests/IR/RemarkTest.cpp
index dca86632071d4..df8e8c8cc066b 100644
--- a/mlir/unittests/IR/RemarkTest.cpp
+++ b/mlir/unittests/IR/RemarkTest.cpp
@@ -409,4 +409,83 @@ TEST(Remark, TestArgWithAttribute) {
   EXPECT_FALSE(argWithoutAttr.getAttribute()); // Returns null Attribute
   EXPECT_EQ(argWithoutAttr.val, "Value");
 }
+
+// Test that Remark correctly owns its string data and doesn't have
+// use-after-free issues when the original strings go out of scope.
+// This is particularly important for RemarkEmittingPolicyFinal which
+// stores remarks and emits them later during finalize().
+TEST(Remark, TestRemarkOwnsStringData) {
+  testing::internal::CaptureStderr();
+
+  // These are the expected values we'll check for in the output.
+  // They must match what we create in the inner scope below.
+  const char *expectedCategory = "DynamicCategory";
+  const char *expectedName = "DynamicRemarkName";
+  const char *expectedFunction = "dynamicFunction";
+  const char *expectedMessage = "Dynamic message content";
+
+  {
+    MLIRContext context;
+    Location loc = FileLineColLoc::get(&context, "test.cpp", 42, 10);
+
+    // Setup with RemarkEmittingPolicyFinal - this stores remarks and emits
+    // them only when the engine is destroyed (during finalize).
+    // Note: The 'passed' filter must be set for remark::passed() to emit.
+    mlir::remark::RemarkCategories cats{
+        /*all=*/std::nullopt,
+        /*passed=*/expectedCategory, // Enable passed remarks for this category
+        /*missed=*/std::nullopt,
+        /*analysis=*/std::nullopt,
+        /*failed=*/std::nullopt};
+
+    std::unique_ptr<remark::RemarkEmittingPolicyFinal> policy =
+        std::make_unique<remark::RemarkEmittingPolicyFinal>();
+    LogicalResult isEnabled = remark::enableOptimizationRemarks(
+        context, std::make_unique<MyCustomStreamer>(), std::move(policy), cats,
+        /*printAsEmitRemarks=*/true);
+    ASSERT_TRUE(succeeded(isEnabled)) << "Failed to enable remark engine";
+
+    // Create dynamic strings in an inner scope that will go out of scope
+    // BEFORE the RemarkEngine is destroyed and finalize() is called.
+    {
+      std::string dynamicCategory(expectedCategory);
+      std::string dynamicName(expectedName);
+      std::string dynamicFunction(expectedFunction);
+      std::string dynamicSubCategory("DynamicSubCategory");
+      std::string dynamicMessage(expectedMessage);
+
+      // Emit a remark with all dynamic strings
+      remark::passed(loc, remark::RemarkOpts::name(dynamicName)
+                              .category(dynamicCategory)
+                              .subCategory(dynamicSubCategory)
+                              .function(dynamicFunction))
+          << dynamicMessage;
+
+      // dynamicCategory, dynamicName, dynamicFunction, dynamicSubCategory,
+      // and dynamicMessage all go out of scope here!
+    }
+
+    // At this point, all the dynamic strings have been destroyed.
+    // The Remark stored in RemarkEmittingPolicyFinal must have its own
+    // copies of the string data, otherwise we'd have dangling pointers.
+
+    // Context destruction triggers RemarkEngine destruction, which calls
+    // finalize() on the policy, which then emits the stored remarks.
+    // If Remark doesn't own its strings, this would crash or produce garbage.
+  }
+
+  llvm::errs().flush();
+  std::string errOut = ::testing::internal::GetCapturedStderr();
+
+  // Verify the output contains our expected strings - this proves the
+  // Remark correctly copied and owns the string data.
+  EXPECT_NE(errOut.find(expectedCategory), std::string::npos)
+      << "Expected category not found in output. Got: " << errOut;
+  EXPECT_NE(errOut.find(expectedName), std::string::npos)
+      << "Expected name not found in output. Got: " << errOut;
+  EXPECT_NE(errOut.find(expectedFunction), std::string::npos)
+      << "Expected function not found in output. Got: " << errOut;
+  EXPECT_NE(errOut.find(expectedMessage), std::string::npos)
+      << "Expected message not found in output. Got: " << errOut;
+}
 } // namespace