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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Wed Mar 4 16:53:54 PST 2026


Author: Mehdi Amini
Date: 2026-03-05T00:53:49Z
New Revision: d07ab183b6695a6262c251d24bb86d4b977aff3d

URL: https://github.com/llvm/llvm-project/commit/d07ab183b6695a6262c251d24bb86d4b977aff3d
DIFF: https://github.com/llvm/llvm-project/commit/d07ab183b6695a6262c251d24bb86d4b977aff3d.diff

LOG: [mlir] Add cross-context checks to the IR verifier (#184627)

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

Added: 
    mlir/unittests/IR/VerifierTest.cpp

Modified: 
    mlir/lib/IR/Verifier.cpp
    mlir/unittests/IR/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index e19537a901d18..33da2cd867f42 100644
--- a/mlir/lib/IR/Verifier.cpp
+++ b/mlir/lib/IR/Verifier.cpp
@@ -110,9 +110,30 @@ static bool mayBeValidWithoutTerminator(Block *block) {
 }
 
 LogicalResult OperationVerifier::verifyOnEntrance(Block &block) {
-  for (auto arg : block.getArguments())
+  // Get the parent op and context for cross-context checks. Both are available
+  // whenever the block lives inside a region that has a parent operation.
+  Operation *parentOp = block.getParentOp();
+  MLIRContext *blockCtx = parentOp ? parentOp->getContext() : nullptr;
+
+  for (auto [idx, arg] : llvm::enumerate(block.getArguments())) {
     if (arg.getOwner() != &block)
       return emitError(arg.getLoc(), "block argument not owned by block");
+    if (blockCtx) {
+      // Check the location first; if it is wrong we must use the parent op's
+      // location to emit the error (the arg location would route to the wrong
+      // context's diagnostic handler).
+      if (arg.getLoc().getContext() != blockCtx)
+        return emitError(parentOp->getLoc(), "block argument #")
+               << idx
+               << " location from a 
diff erent MLIRContext than its "
+                  "parent operation";
+      if (arg.getType().getContext() != blockCtx)
+        return emitError(arg.getLoc(), "block argument #")
+               << idx
+               << " type from a 
diff erent MLIRContext than its "
+                  "parent operation";
+    }
+  }
 
   // Verify that this block has a terminator.
   if (block.empty()) {
@@ -129,6 +150,17 @@ LogicalResult OperationVerifier::verifyOnEntrance(Block &block) {
     if (op.getNumSuccessors() != 0 && &op != &block.back())
       return op.emitError(
           "operation with block successors must terminate its parent block");
+    // Check that each op's location (which defines its context) is from the
+    // same MLIRContext as the enclosing block. We cannot use op.emitError()
+    // here because op's context is the wrong one; emit via the parent op's
+    // location instead.
+    if (blockCtx && op.getContext() != blockCtx) {
+      emitError(parentOp->getLoc(), "operation '")
+          << op.getName()
+          << "' has a location from a 
diff erent MLIRContext than its "
+             "enclosing block";
+      return failure();
+    }
   }
 
   return success();
@@ -155,13 +187,38 @@ LogicalResult OperationVerifier::verifyOnExit(Block &block) {
 }
 
 LogicalResult OperationVerifier::verifyOnEntrance(Operation &op) {
-  // Check that operands are non-nil and structurally ok.
-  for (auto operand : op.getOperands())
+  // op.getContext() is defined as location->getContext(), so opCtx is the
+  // location's context by construction.  The OperationName, however, carries
+  // its own context reference and can independently point elsewhere.
+  MLIRContext *opCtx = op.getContext();
+  if (op.getName().getContext() != opCtx)
+    return op.emitError(
+        "operation name from a 
diff erent MLIRContext than this operation");
+
+  // Check result types are from the same context as the operation.
+  for (auto [i, result] : llvm::enumerate(op.getResults())) {
+    if (result.getType().getContext() != opCtx)
+      return op.emitOpError()
+             << "result #" << i
+             << " type from a 
diff erent 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 
diff erent 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 
diff erent 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..9b2dd3c73454a
--- /dev/null
+++ b/mlir/unittests/IR/VerifierTest.cpp
@@ -0,0 +1,177 @@
+//===- 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/Block.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 "mlir/IR/Region.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();
+}
+
+TEST(VerifierTest, CrossContextOperationName) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Location (and thus op.getContext()) is from ctxA; OperationName is from
+  // ctxB.  op.getContext() == op.getLoc().getContext() by definition, so the
+  // OperationName is the independent source of a cross-context violation here.
+  Operation *op =
+      Operation::create(UnknownLoc::get(&ctxA), OperationName("foo.bar", &ctxB),
+                        {}, {}, NamedAttrList(), nullptr, {}, 0);
+
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(op)));
+  op->destroy();
+}
+
+TEST(VerifierTest, CrossContextOperationLocation) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Create an outer op in ctxA with one region.
+  Operation *outerOp = Operation::create(
+      UnknownLoc::get(&ctxA), OperationName("foo.outer", &ctxA), {}, {},
+      NamedAttrList(), nullptr, {}, /*numRegions=*/1);
+
+  Block *block = new Block();
+  outerOp->getRegion(0).push_back(block);
+
+  // Create an inner op whose location (and thus context) comes from ctxB, and
+  // place it inside the ctxA block.  This is the cross-context violation:
+  // op.getContext() == ctxB but the enclosing block belongs to ctxA.
+  Operation *innerOp = Operation::create(UnknownLoc::get(&ctxB),
+                                         OperationName("foo.inner", &ctxB), {},
+                                         {}, NamedAttrList(), nullptr, {}, 0);
+  block->push_back(innerOp);
+
+  // The verifier emits via the parent op's location (ctxA), so only ctxA's
+  // handler needs to suppress it.
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(outerOp)));
+  outerOp->destroy();
+}
+
+TEST(VerifierTest, CrossContextBlockArgType) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Create an unregistered op in ctxA with one region.
+  Operation *op =
+      Operation::create(UnknownLoc::get(&ctxA), OperationName("foo.bar", &ctxA),
+                        {}, {}, NamedAttrList(), nullptr, {}, /*numRegions=*/1);
+
+  // Add a block with one argument whose type comes from ctxB.
+  Block *block = new Block();
+  op->getRegion(0).push_back(block);
+  Type i32B = IntegerType::get(&ctxB, 32);
+  block->addArgument(i32B, UnknownLoc::get(&ctxA));
+
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(op)));
+  op->destroy();
+}
+
+TEST(VerifierTest, CrossContextBlockArgLocation) {
+  MLIRContext ctxA, ctxB;
+  ctxA.allowUnregisteredDialects();
+  ctxB.allowUnregisteredDialects();
+
+  // Create an unregistered op in ctxA with one region.
+  Operation *op =
+      Operation::create(UnknownLoc::get(&ctxA), OperationName("foo.bar", &ctxA),
+                        {}, {}, NamedAttrList(), nullptr, {}, /*numRegions=*/1);
+
+  // Add a block with one argument whose location comes from ctxB.
+  Block *block = new Block();
+  op->getRegion(0).push_back(block);
+  Type i32A = IntegerType::get(&ctxA, 32);
+  block->addArgument(i32A, UnknownLoc::get(&ctxB));
+
+  // The verifier emits the diagnostic via the parent op's location (ctxA).
+  ScopedDiagnosticHandler suppress(&ctxA,
+                                   [](Diagnostic &) { return success(); });
+  EXPECT_TRUE(failed(verify(op)));
+  op->destroy();
+}
+
+} // namespace


        


More information about the Mlir-commits mailing list