[MLIR][Python][NOMERGE] demo finding live ops via nanobind internals

[makslevental]

This demo PR shows how one can use nanobind internals to find "live operations". Specifically we get the inst_c2p pointer map and iterate it to find all instances which are of type PyOperation. The result is

# CHECK: %subview_10 = memref.subview
# CHECK: %0 = arith.addi %c3_2, %c4_3 : index
print_live_ops()

Note, not for the faint of heart - performs #include "nb_internals.h" in a fairly cursed, but I believe robust, way...

[makslevental]

Currently, this is NOMERGE because it's very "dangerous" but I would not be opposed to landing some form of this (heavily guarded behind scary #defines and CMake vars) to help users like @teqdruid (see discussion starting here #155114 (comment)).

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

This demo PR shows how one can use nanobind internals to find "live operations". Specifically we get the inst_c2p pointer map and iterate it to find all instances which are of type PyOperation. The result is

# CHECK: %subview_10 = memref.subview
# CHECK: %0 = arith.addi %c3_2, %c4_3 : index
print_live_ops()

Note, not for the faint of heart - performs #include "nb_internals.h" in a fairly cursed, but I believe robust, way...

---

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

4 Files Affected:

• (modified) mlir/examples/standalone/python/CMakeLists.txt (+3)
• (modified) mlir/lib/Bindings/Python/IRCore.cpp (+32)
• (modified) mlir/python/CMakeLists.txt (+3)
• (modified) mlir/test/python/dialects/memref.py (+4)

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<PyLocation> &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<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 @@ 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

[teqdruid]

Something like this could work for me, but this makes me uncomfortable as all hell. Since I want to enable it for all of the CIRCT Python bindings I don't want things to break when someone uses a newer (or potentially older) version of nanobind.

[makslevental]

I don't want things to break when someone uses a newer

There's not large implementation details being used here (currently) and what is used is fundamental to how nanobind works - the inst_c2p map has always been there, even in pybind, and nanobind wouldn't work without it. Granted the internal API surface could change 🤷.

(or potentially older) version of nanobind.

We already have a minimum version requirement (which I recently bumped):

llvm-project/mlir/cmake/modules/MLIRDetectPythonEnv.cmake

Line 57 in efd96af

find_package(nanobind 2.9 CONFIG REQUIRED)