[PyCDE][Python] Fixing leaks reported by nanobind (#8029)

The ESI and MSFT Python modules had object memory leaks which upstream
nanobind detected. Fixing them.

Author: teqdruid

Diff for: frontends/PyCDE/src/pycde/esi.py

@@ -19,8 +19,16 @@
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, cast
 
+import atexit
+
 __dir__ = Path(__file__).parent
 
+
+@atexit.register
+def _cleanup():
+  raw_esi.cleanup()
+
+
 FlattenStructPorts = "esi.portFlattenStructs"
 PortInSuffix = "esi.portInSuffix"
 PortOutSuffix = "esi.portOutSuffix"

Diff for: frontends/PyCDE/test/test_compreg.py

@@ -4,8 +4,9 @@
 # RUN: FileCheck %s --input-file %t/hw/CompReg.tcl --check-prefix TCL
 
 import pycde
-from pycde import types, Module, AppID, Clock, Input, Output
+from pycde import Module, AppID, Clock, Input, Output
 from pycde.devicedb import LocationVector
+from pycde.types import Bits
 
 from pycde.module import generator
 
@@ -27,9 +28,9 @@
 
 class CompReg(Module):
   clk = Clock()
-  rst = Input(types.i1)
-  input = Input(types.i8)
-  output = Output(types.i8)
+  rst = Input(Bits(1))
+  input = Input(Bits(8))
+  output = Output(Bits(8))
 
   @generator
   def build(ports):
@@ -40,7 +41,6 @@ def build(ports):
 
 
 mod = pycde.System([CompReg], name="CompReg", output_directory=sys.argv[1])
-mod.print()
 mod.generate()
 top_inst = mod.get_instance(CompReg)
 mod.createdb()

Diff for: frontends/PyCDE/test/test_xrt.py

@@ -14,27 +14,24 @@
 # RUN: FileCheck %s --input-file %t/hw/XrtTop.sv --check-prefix=TOP
 
 import pycde
-from pycde import Clock, Input, Module, generator, types
+from pycde import Clock, Input, Module, generator
+from pycde.types import Bits
 from pycde.bsp import XrtBSP
 
 import sys
 
 
 class Main(Module):
-  clk = Clock(types.i1)
-  rst = Input(types.i1)
+  clk = Clock()
+  rst = Input(Bits(1))
 
   @generator
   def construct(ports):
     pass
 
 
 gendir = sys.argv[1]
-s = pycde.System(XrtBSP(Main),
-                 name="ESILoopback",
-                 output_directory=gendir,
-                 sw_api_langs=["python"])
-s.run_passes(debug=True)
+s = pycde.System(XrtBSP(Main), name="ESILoopback", output_directory=gendir)
 s.compile()
 s.package()
 

Diff for: lib/Bindings/Python/ESIModule.cpp

@@ -31,24 +31,39 @@ using namespace circt::esi;
 // The main entry point into the ESI Assembly API.
 //===----------------------------------------------------------------------===//
 
+/// Container for a Python function that will be called to generate a service.
+class ServiceGenFunc {
+public:
+  ServiceGenFunc(py::object genFunc) : genFunc(std::move(genFunc)) {}
+
+  MlirLogicalResult run(MlirOperation reqOp, MlirOperation declOp,
+                        MlirOperation recOp) {
+    py::gil_scoped_acquire acquire;
+    py::object rc = genFunc(reqOp, declOp, recOp);
+    return rc.cast<bool>() ? mlirLogicalResultSuccess()
+                           : mlirLogicalResultFailure();
+  }
+
+private:
+  py::object genFunc;
+};
+
 // Mapping from unique identifier to python callback. We use std::string
 // pointers since we also need to allocate memory for the string.
-llvm::DenseMap<std::string *, PyObject *> serviceGenFuncLookup;
+llvm::DenseMap<std::string *, ServiceGenFunc> serviceGenFuncLookup;
 static MlirLogicalResult serviceGenFunc(MlirOperation reqOp,
                                         MlirOperation declOp,
                                         MlirOperation recOp, void *userData) {
   std::string *name = static_cast<std::string *>(userData);
-  py::handle genFunc(serviceGenFuncLookup[name]);
-  py::gil_scoped_acquire();
-  py::object rc = genFunc(reqOp, declOp, recOp);
-  return rc.cast<bool>() ? mlirLogicalResultSuccess()
-                         : mlirLogicalResultFailure();
+  auto iter = serviceGenFuncLookup.find(name);
+  if (iter == serviceGenFuncLookup.end())
+    return mlirLogicalResultFailure();
+  return iter->getSecond().run(reqOp, declOp, recOp);
 }
 
 void registerServiceGenerator(std::string name, py::object genFunc) {
   std::string *n = new std::string(name);
-  genFunc.inc_ref();
-  serviceGenFuncLookup[n] = genFunc.ptr();
+  serviceGenFuncLookup.try_emplace(n, ServiceGenFunc(genFunc));
   circtESIRegisterGlobalServiceGenerator(wrap(*n), serviceGenFunc, n);
 }
 
@@ -81,6 +96,12 @@ void circt::python::populateDialectESISubmodule(py::module &m) {
   m.doc() = "ESI Python Native Extension";
   ::registerESIPasses();
 
+  // Clean up references when the module is unloaded.
+  auto cleanup = []() { serviceGenFuncLookup.clear(); };
+  m.def("cleanup", cleanup,
+        "Cleanup various references. Must be called before the module is "
+        "unloaded in order to not leak.");
+
   m.def("registerServiceGenerator", registerServiceGenerator,
         "Register a service generator for a given service name.",
         py::arg("impl_type"), py::arg("generator"));

Diff for: lib/Bindings/Python/MSFTModule.cpp

@@ -117,18 +117,18 @@ class PlacementDB {
 class PyLocationVecIterator {
 public:
   /// Get item at the specified position, translating a nullptr to None.
-  static py::handle getItem(MlirAttribute locVec, intptr_t pos) {
+  static std::optional<MlirAttribute> getItem(MlirAttribute locVec,
+                                              intptr_t pos) {
     MlirAttribute loc = circtMSFTLocationVectorAttrGetElement(locVec, pos);
     if (loc.ptr == nullptr)
-      return py::none();
-    return py::detail::type_caster<MlirAttribute>().cast(
-        loc, py::return_value_policy::automatic, py::handle());
+      return std::nullopt;
+    return loc;
   }
 
   PyLocationVecIterator(MlirAttribute attr) : attr(attr) {}
   PyLocationVecIterator &dunderIter() { return *this; }
 
-  py::handle dunderNext() {
+  std::optional<MlirAttribute> dunderNext() {
     if (nextIndex >= circtMSFTLocationVectorAttrGetNumElements(attr)) {
       throw py::stop_iteration();
     }