[Mlir-commits] [mlir] [MLIR][Python] demo finding live ops via nanobind internals (PR #160632)

[Maksim Levental]

https://github.com/makslevental created https://github.com/llvm/llvm-project/pull/160632

None

>From cc1df73748cb875e4df4534af453a9c0345c6b66 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Wed, 24 Sep 2025 20:02:12 -0700
Subject: [PATCH] [MLIR][Python] demo finding live ops via nanobind internals

---
 mlir/lib/Bindings/Python/IRCore.cpp | 27 +++++++++++++++++++++++++++
 mlir/python/CMakeLists.txt          |  2 ++
 mlir/test/python/dialects/memref.py |  2 ++
 3 files changed, 31 insertions(+)

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..12f511f6f8eca 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2897,6 +2897,31 @@ maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
+#include "nb_internals.h"
+
+auto print_leak = [](void *k, PyObject *v) {
+  nanobind::detail::type_data *tp = nanobind::detail::nb_type_data(Py_TYPE(v));
+  fprintf(stderr, " - leaked instance %p of type \"%s\"\n", k, tp->name);
+};
+
+void print_leaks() noexcept {
+  nanobind::detail::nb_internals *p = nanobind::detail::internals;
+  for (size_t i = 0; i < p->shard_count; ++i) {
+    nanobind::detail::nb_shard &s = p->shards[i];
+    nanobind::detail::lock_shard lock(s);
+    for (auto [k, v] : s.inst_c2p) {
+      if (NB_UNLIKELY(nanobind::detail::nb_is_seq(v))) {
+        nanobind::detail::nb_inst_seq *seq = nanobind::detail::nb_get_seq(v);
+        for (; seq != nullptr; seq = seq->next) {
+          print_leak(k, seq->inst);
+        }
+      } else {
+        print_leak(k, (PyObject *)v);
+      }
+    }
+  }
+}
+
 void mlir::python::populateIRCore(nb::module_ &m) {
   // disable leak warnings which tend to be false positives.
   nb::set_leak_warnings(false);
@@ -4475,4 +4500,6 @@ void mlir::python::populateIRCore(nb::module_ &m) {
       PyErr_SetObject(PyExc_Exception, obj.ptr());
     }
   });
+
+  m.def("print_leaks", &print_leaks);
 }
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index d6686bb89ce4e..91ab8cda4ecb7 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -974,3 +974,5 @@ add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}")
 if(MLIR_INCLUDE_TESTS)
   add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")
 endif()
+target_include_directories(MLIRPythonModules.extension._mlir.dso PRIVATE ${NB_DIR}/ext/robin_map/include)
+target_include_directories(MLIRPythonModules.extension._mlir.dso PRIVATE ${NB_DIR}/src)
diff --git a/mlir/test/python/dialects/memref.py b/mlir/test/python/dialects/memref.py
index b91fdc367cf30..4efc64abcb9d9 100644
--- a/mlir/test/python/dialects/memref.py
+++ b/mlir/test/python/dialects/memref.py
@@ -179,6 +179,8 @@ def testSubViewOpInferReturnTypeSemantics():
             # CHECK: %{{.*}} = memref.subview %[[DYNAMICALLOC]][1, 1] [3, 3] [1, 1] : memref<10x10xi32, strided<[10, 1], offset: ?>> to memref<3x3xi32, strided<[10, 1], offset: ?>>
             print(y.owner)
 
+            print_leaks()
+
 
 # CHECK-LABEL: TEST: testSubViewOpInferReturnTypeExtensiveSlicing
 @run