[Mlir-commits] [mlir] [mlir][py] avoid crashing on None contexts in custom `get`s (PR #171140)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Oleksandr "Alex" Zinenko (ftynse)

<details>
<summary>Changes</summary>

Following a series of refactorings, MLIR Python bindings would crash if a
dialect object requiring a context defined using mlir_attribute/type_subclass
was constructed outside of the `ir.Context` context manager. The type caster
for `MlirContext` would try using `ir.Context.current` when the default `None`
value was provided to the `get`, which would also just return `None`. The
caster would then attempt to obtain the MLIR capsule for that `None`, fail,
but access it anyway without checking, leading to a C++ assertion failure or
segfault.

Guard against this case in nanobind adaptors. Also emit a warning to the user
to clarify expectations, as the default message confusingly says that `None` is
accepted as context and then fails with a type error. Using Python C API is
currently recommended by nanobind in this case since the surrounding function
must be marked `noexcept`.

The corresponding test is in the PDL dialect since it is where I first observed
the behavior. Core types are not using the `mlir_type_subclass` mechanism and
are immune to the problem, so cannot be used for checking.

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


2 Files Affected:

- (modified) mlir/include/mlir/Bindings/Python/NanobindAdaptors.h (+14-6) 
- (modified) mlir/test/python/dialects/pdl_types.py (+13) 


``````````diff
diff --git a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
index 847951ab5fd46..6594670abaaa7 100644
--- a/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
+++ b/mlir/include/mlir/Bindings/Python/NanobindAdaptors.h
@@ -181,14 +181,22 @@ struct type_caster<MlirContext> {
   bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept {
     if (src.is_none()) {
       // Gets the current thread-bound context.
-      // TODO: This raises an error of "No current context" currently.
-      // Update the implementation to pretty-print the helpful error that the
-      // core implementations print in this case.
       src = mlir::python::irModule().attr("Context").attr("current");
     }
-    std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src);
-    value = mlirPythonCapsuleToContext(capsule->ptr());
-    return !mlirContextIsNull(value);
+    // If there is no context, including thread-bound, emit a warning (since
+    // this function is not allowed to throw) and fail to cast.
+    if (src.is_none()) {
+      PyErr_Warn(
+          PyExc_RuntimeWarning,
+          "Passing None as MLIR Context is only allowed inside "
+          "the " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") " context manager.");
+      return false;
+    }
+    if (std::optional<nanobind::object> capsule = mlirApiObjectToCapsule(src)) {
+      value = mlirPythonCapsuleToContext(capsule->ptr());
+      return !mlirContextIsNull(value);
+    }
+    return false;
   }
 };
 
diff --git a/mlir/test/python/dialects/pdl_types.py b/mlir/test/python/dialects/pdl_types.py
index 16a41e2a4c1ce..bbc6007511721 100644
--- a/mlir/test/python/dialects/pdl_types.py
+++ b/mlir/test/python/dialects/pdl_types.py
@@ -148,3 +148,16 @@ def test_value_type():
     print(parsedType)
     # CHECK: !pdl.value
     print(constructedType)
+
+
+# CHECK-LABEL: TEST: test_type_without_context
+ at run
+def test_type_without_context():
+  # Constructing a type without the surrounding ir.Context context manager
+  # should raise an exception but not crash.
+  try:
+    constructedType = pdl.ValueType.get()
+  except TypeError:
+    pass
+  else:
+    assert False, "Expected TypeError to be raised."

``````````

</details>


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