perf: intern cached str objects

metaopt · Nov 30, 2023 · 16112fb · 16112fb
1 parent abb03ab
commit 16112fb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/include/utils.h b/include/utils.h
@@ -53,16 +53,22 @@ inline void HashCombine(py::ssize_t& seed, const T& v) {  // NOLINT[runtime/refe
 constexpr bool NONE_IS_LEAF = true;
 constexpr bool NONE_IS_NODE = false;
 
-// NOTE: Use raw pointers to leak the memory intentionally to avoid py::object deallocation and
-//       garbage collection.
-#define Py_Declare_ID(name)                                 \
-    inline PyObject* Py_ID_##name() {                       \
-        static PyObject* ptr = (new py::str{#name})->ptr(); \
-        return ptr;                                         \
+#define Py_Declare_ID(name)                                    \
+    inline PyObject* Py_ID_##name() {                          \
+        static PyObject* ptr = []() {                          \
+            PyObject* ptr = PyUnicode_InternFromString(#name); \
+            if (ptr == nullptr) [[unlikely]] {                 \
+                throw py::error_already_set();                 \
+            }                                                  \
+            Py_INCREF(ptr); /* leak a reference on purpose */  \
+            return ptr;                                        \
+        }();                                                   \
+        return ptr;                                            \
     }
 
 #define Py_Get_ID(name) (Py_ID_##name())
 
+Py_Declare_ID(optree);
 Py_Declare_ID(__module__);         // type.__module__
 Py_Declare_ID(__qualname__);       // type.__qualname__
 Py_Declare_ID(__name__);           // type.__name__

diff --git a/src/optree.cpp b/src/optree.cpp
@@ -128,12 +128,12 @@ void BuildModule(py::module_& mod) {  // NOLINT[runtime/references]
             .value("DEQUE", PyTreeKind::Deque, "A collections.deque.")
             .value("STRUCTSEQUENCE", PyTreeKind::StructSequence, "A PyStructSequence.");
     reinterpret_cast<PyTypeObject*>(PyTreeKindTypeObject.ptr())->tp_name = "optree.PyTreeKind";
-    py::setattr(PyTreeKindTypeObject.ptr(), "__module__", py::str("optree"));
+    py::setattr(PyTreeKindTypeObject.ptr(), Py_Get_ID(__module__), Py_Get_ID(optree));
 
     auto PyTreeSpecTypeObject =
         py::class_<PyTreeSpec>(mod, "PyTreeSpec", "Representing the structure of the pytree.");
     reinterpret_cast<PyTypeObject*>(PyTreeSpecTypeObject.ptr())->tp_name = "optree.PyTreeSpec";
-    py::setattr(PyTreeSpecTypeObject.ptr(), "__module__", py::str("optree"));
+    py::setattr(PyTreeSpecTypeObject.ptr(), Py_Get_ID(__module__), Py_Get_ID(optree));
 
     PyTreeSpecTypeObject
         .def_property_readonly(