[Mlir-commits] [mlir] [mlir] Add cross-context checks to the IR verifier (PR #184627)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

Detect IR where result types, operand types, or discardable attribute values come from a different MLIRContext than the operation itself. Mixing contexts is a latent use-after-free hazard when one context is destroyed before the other; the verifier now reports a clear error instead of silently allowing the invalid IR through.

Three new unit tests in MLIRIRTests cover each scenario.

Fix #<!-- -->61569

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


3 Files Affected:

- (modified) mlir/lib/IR/Verifier.cpp (+20-2) 
- (modified) mlir/unittests/IR/CMakeLists.txt (+1) 
- (added) mlir/unittests/IR/VerifierTest.cpp (+83) 


``````````diff
diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index e19537a901d18..13b0ce2644ca7 100644
--- a/mlir/lib/IR/Verifier.cpp
+++ b/mlir/lib/IR/Verifier.cpp
@@ -155,13 +155,31 @@ LogicalResult OperationVerifier::verifyOnExit(Block &block) {
 }
 
 LogicalResult OperationVerifier::verifyOnEntrance(Operation &op) {
-  // Check that operands are non-nil and structurally ok.
-  for (auto operand : op.getOperands())
+  MLIRContext *opCtx = op.getContext();
+
+  // Check result types are from the same context as the operation.
+  for (OpResult result : op.getResults()) {
+    if (result.getType().getContext() != opCtx)
+      return op.emitError(
+          "result type from a different MLIRContext than this operation");
+  }
+
+  // Check that operands are non-nil and their types are from the same context.
+  for (auto [i, operand] : llvm::enumerate(op.getOperands())) {
     if (!operand)
       return op.emitError("null operand found");
+    if (operand.getType().getContext() != opCtx)
+      return op.emitOpError()
+             << "operand #" << i
+             << " type from a different MLIRContext than this operation";
+  }
 
   /// Verify that all of the attributes are okay.
   for (auto attr : op.getDiscardableAttrDictionary()) {
+    if (attr.getValue().getContext() != opCtx)
+      return op.emitOpError()
+             << "discardable attribute '" << attr.getName()
+             << "' value from a different MLIRContext than this operation";
     // Check for any optional dialect specific attributes.
     if (auto *dialect = attr.getNameDialect())
       if (failed(dialect->verifyOperationAttribute(&op, attr)))
diff --git a/mlir/unittests/IR/CMakeLists.txt b/mlir/unittests/IR/CMakeLists.txt
index dd3b110dcd295..49eee0c93f02e 100644
--- a/mlir/unittests/IR/CMakeLists.txt
+++ b/mlir/unittests/IR/CMakeLists.txt
@@ -21,6 +21,7 @@ add_mlir_unittest(MLIRIRTests
   TypeAttrNamesTest.cpp
   OpPropertiesTest.cpp
   ValueTest.cpp
+  VerifierTest.cpp
   BlobManagerTest.cpp
 
   DEPENDS
diff --git a/mlir/unittests/IR/VerifierTest.cpp b/mlir/unittests/IR/VerifierTest.cpp
new file mode 100644
index 0000000000000..ad542919876c0
--- /dev/null
+++ b/mlir/unittests/IR/VerifierTest.cpp
@@ -0,0 +1,83 @@
+//===- mlir/unittest/IR/VerifierTest.cpp - Verifier unit tests ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Verifier.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+
+TEST(VerifierTest, CrossContextResultType) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Result type comes from ctxB but the op lives in ctxA.
+  Type i32B = IntegerType::get(&ctxB, 32);
+  Operation *op =
+      Operation::create(UnknownLoc::get(&ctxA), OperationName("foo.bar", &ctxA),
+                        {i32B}, {}, NamedAttrList(), nullptr, {}, 0);
+
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(op)));
+  op->destroy();
+}
+
+TEST(VerifierTest, CrossContextDiscardableAttr) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Attribute value comes from ctxB but the op lives in ctxA.
+  IntegerAttr attrB = IntegerAttr::get(IntegerType::get(&ctxB, 32), 42);
+  NamedAttrList attrs;
+  attrs.append(StringAttr::get(&ctxA, "my_attr"), attrB);
+
+  Operation *op =
+      Operation::create(UnknownLoc::get(&ctxA), OperationName("foo.bar", &ctxA),
+                        {}, {}, std::move(attrs), nullptr, {}, 0);
+
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(op)));
+  op->destroy();
+}
+
+TEST(VerifierTest, CrossContextOperand) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Producer op lives in ctxB; its result type is also from ctxB.
+  Type i32B = IntegerType::get(&ctxB, 32);
+  Operation *producerOp = Operation::create(
+      UnknownLoc::get(&ctxB), OperationName("foo.producer", &ctxB), {i32B}, {},
+      NamedAttrList(), nullptr, {}, 0);
+  Value valFromCtxB = producerOp->getResult(0);
+
+  // Consumer op lives in ctxA but uses the value (whose type is in ctxB).
+  Operation *consumerOp = Operation::create(
+      UnknownLoc::get(&ctxA), OperationName("foo.consumer", &ctxA), {},
+      {valFromCtxB}, NamedAttrList(), nullptr, {}, 0);
+
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(consumerOp)));
+  consumerOp->destroy();
+  producerOp->destroy();
+}
+
+} // namespace

``````````

</details>


https://github.com/llvm/llvm-project/pull/184627