diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt index 2a4fd99d243e0..e3a76c21156fb 100644 --- a/mlir/examples/standalone/python/CMakeLists.txt +++ b/mlir/examples/standalone/python/CMakeLists.txt @@ -152,3 +152,6 @@ if(NOT EXTERNAL_PROJECT_BUILD) 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 83a8757bb72c7..5bef751924aa9 100644 --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -2897,6 +2897,36 @@ maybeGetTracebackLocation(const std::optional &location) { // 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(); + 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 @@ void mlir::python::populateIRCore(nb::module_ &m) { 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 d6686bb89ce4e..9ac2e6b6b53b2 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -974,3 +974,6 @@ add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}") 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 b91fdc367cf30..7d95ef7f86d9b 100644 --- a/mlir/test/python/dialects/memref.py +++ b/mlir/test/python/dialects/memref.py @@ -179,6 +179,10 @@ 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) + # CHECK: %subview_10 = memref.subview + # CHECK: %0 = arith.addi %c3_2, %c4_3 : index + print_live_ops() + # CHECK-LABEL: TEST: testSubViewOpInferReturnTypeExtensiveSlicing @run