[MLIR][Python] demo finding live ops via nanobind internals
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 2a4fd99..e3a76c2 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -152,3 +152,6 @@
add_dependencies(StandalonePythonModules "${_mlir_typestub_gen_target}")
endif()
add_dependencies(StandalonePythonModules "${_standaloneDialectsNanobind_typestub_gen_target}")
+
+target_include_directories(StandalonePythonModules.extension._mlir.dso PRIVATE ${NB_DIR}/ext/robin_map/include)
+target_include_directories(StandalonePythonModules.extension._mlir.dso PRIVATE ${NB_DIR}/src)
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757..5bef751 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2897,6 +2897,36 @@
// Populates the core exports of the 'ir' submodule.
//------------------------------------------------------------------------------
+#include "nb_internals.h"
+
+auto print_live_op = [](void *k, PyObject *v) {
+ nb::handle op_py_type = nb::type<PyOperation>();
+ nb::handle maybe_op_inst{v};
+ nb::str end{" "};
+ if (maybe_op_inst.type().is(op_py_type)) {
+ nb::print("found live operation:", end);
+ nb::print(maybe_op_inst);
+ }
+};
+
+void print_live_ops() 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_live_op(k, seq->inst);
+ }
+ } else {
+ print_live_op(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 +4505,6 @@
PyErr_SetObject(PyExc_Exception, obj.ptr());
}
});
+
+ m.def("print_live_ops", &print_live_ops);
}
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index d6686bb..9ac2e6b 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -974,3 +974,6 @@
if(MLIR_INCLUDE_TESTS)
add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}")
endif()
+message(STATUS "NB_DIR=${NB_DIR}")
+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 b91fdc3..7d95ef7 100644
--- a/mlir/test/python/dialects/memref.py
+++ b/mlir/test/python/dialects/memref.py
@@ -179,6 +179,10 @@
# 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)
+ # CHECK: %subview_10 = memref.subview
+ # CHECK: %0 = arith.addi %c3_2, %c4_3 : index
+ print_live_ops()
+
# CHECK-LABEL: TEST: testSubViewOpInferReturnTypeExtensiveSlicing
@run
