Per-module precomputed constants, improve enum casting perf. (#1184)

Creating runtime constants, particularly strings, has a significant cost on
Python. This commit adds a per-module stash of precomputed constants, allowing
us to avoid this cost and improve the performance of casting enums from Python
in the process.

Author: hpkfft

include/nanobind/nb_defs.h

@@ -28,15 +28,22 @@
 #  define NB_INLINE          __forceinline
 #  define NB_NOINLINE        __declspec(noinline)
 #  define NB_INLINE_LAMBDA
+#  define NB_NOUNROLL
 #else
 #  define NB_EXPORT          __attribute__ ((visibility("default")))
 #  define NB_IMPORT          NB_EXPORT
 #  define NB_INLINE          inline __attribute__((always_inline))
 #  define NB_NOINLINE        __attribute__((noinline))
 #  if defined(__clang__)
 #    define NB_INLINE_LAMBDA __attribute__((always_inline))
+#    define NB_NOUNROLL      _Pragma("nounroll")
 #  else
 #    define NB_INLINE_LAMBDA
+#    if defined(__GNUC__)
+#      define NB_NOUNROLL    _Pragma("GCC unroll 0")
+#    else
+#      define NB_NOUNROLL
+#    endif
 #  endif
 #endif
 
@@ -202,12 +209,14 @@
     X(const X &) = delete;                                                     \
     X &operator=(const X &) = delete;
 
+#define NB_MOD_STATE_SIZE 80
+
 // Helper macros to ensure macro arguments are expanded before token pasting/stringification
 #define NB_MODULE_IMPL(name, variable) NB_MODULE_IMPL2(name, variable)
 #define NB_MODULE_IMPL2(name, variable)                                        \
     static void nanobind_##name##_exec_impl(nanobind::module_);                \
     static int nanobind_##name##_exec(PyObject *m) {                           \
-        nanobind::detail::init(NB_DOMAIN_STR);                                 \
+        nanobind::detail::nb_module_exec(NB_DOMAIN_STR, m);                    \
         try {                                                                  \
             nanobind_##name##_exec_impl(                                       \
                 nanobind::borrow<nanobind::module_>(m));                       \
@@ -227,8 +236,9 @@
         NB_MODULE_SLOTS_2                                                      \
     };                                                                         \
     static struct PyModuleDef nanobind_##name##_module = {                     \
-        PyModuleDef_HEAD_INIT, #name, nullptr, 0, nullptr,                     \
-        nanobind_##name##_slots, nullptr, nullptr, nullptr                     \
+        PyModuleDef_HEAD_INIT, #name, nullptr, NB_MOD_STATE_SIZE, nullptr,     \
+        nanobind_##name##_slots, nanobind::detail::nb_module_traverse,         \
+        nanobind::detail::nb_module_clear, nanobind::detail::nb_module_free    \
     };                                                                         \
     extern "C" [[maybe_unused]] NB_EXPORT PyObject *PyInit_##name(void);       \
     extern "C" PyObject *PyInit_##name(void) {                                 \

include/nanobind/nb_lib.h

@@ -9,8 +9,17 @@
 
 NAMESPACE_BEGIN(NB_NAMESPACE)
 
+NAMESPACE_BEGIN(dlpack)
+
+// The version of DLPack that is supported by libnanobind
+static constexpr uint32_t major_version = 0;
+static constexpr uint32_t minor_version = 0;
+
 // Forward declarations for types in ndarray.h (1)
-namespace dlpack { struct dltensor; struct dtype; }
+struct dltensor;
+struct dtype;
+
+NAMESPACE_END(dlpack)
 
 NAMESPACE_BEGIN(detail)
 
@@ -107,7 +116,10 @@ NB_CORE void raise_next_overload_if_null(void *p);
 
 // ========================================================================
 
-NB_CORE void init(const char *domain);
+NB_CORE void nb_module_exec(const char *domain, PyObject *m);
+NB_CORE int  nb_module_traverse(PyObject *m, visitproc visit, void *arg);
+NB_CORE int  nb_module_clear(PyObject *m);
+NB_CORE void nb_module_free(void *m);
 
 // ========================================================================
 

src/nb_enum.cpp

@@ -190,7 +190,8 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
         return false;
 
     if ((t->flags & (uint32_t) enum_flags::is_flag) != 0 && Py_TYPE(o) == t->type_py) {
-        PyObject *value_o = PyObject_GetAttrString(o, "value");
+        PyObject *value_o =
+                PyObject_GetAttr(o, static_pyobjects[pyobj_name::value_str]);
         if (value_o == nullptr) {
             PyErr_Clear();
             return false;

src/nb_internals.cpp

@@ -164,6 +164,71 @@ void default_exception_translator(const std::exception_ptr &p, void *) {
 nb_internals *internals = nullptr;
 PyTypeObject *nb_meta_cache = nullptr;
 
+
+static const char* interned_c_strs[pyobj_name::string_count] {
+    "value",
+    "copy",
+    "from_dlpack",
+    "__dlpack__",
+    "max_version",
+    "dl_device",
+};
+
+PyObject **static_pyobjects = nullptr;
+
+static bool init_pyobjects(PyObject* m) {
+    PyObject** pyobjects = (PyObject**) PyModule_GetState(m);
+    if (!pyobjects)
+        return false;
+
+    NB_NOUNROLL
+    for (int i = 0; i < pyobj_name::string_count; ++i)
+        pyobjects[i] = PyUnicode_InternFromString(interned_c_strs[i]);
+
+    pyobjects[pyobj_name::copy_tpl] =
+            PyTuple_Pack(1, pyobjects[pyobj_name::copy_str]);
+    pyobjects[pyobj_name::max_version_tpl] =
+            PyTuple_Pack(1, pyobjects[pyobj_name::max_version_str]);
+
+    PyObject* one = PyLong_FromLong(1);
+    PyObject* zero = PyLong_FromLong(0);
+    pyobjects[pyobj_name::dl_cpu_tpl] = PyTuple_Pack(2, one, zero);
+    Py_DECREF(zero);
+    Py_DECREF(one);
+
+    PyObject* major = PyLong_FromLong(dlpack::major_version);
+    PyObject* minor = PyLong_FromLong(dlpack::minor_version);
+    pyobjects[pyobj_name::dl_version_tpl] = PyTuple_Pack(2, major, minor);
+    Py_DECREF(minor);
+    Py_DECREF(major);
+
+    static_pyobjects = pyobjects;
+
+    return true;
+}
+
+NB_NOINLINE int nb_module_traverse(PyObject *m, visitproc visit, void *arg) {
+    PyObject** pyobjects = (PyObject**) PyModule_GetState(m);
+    NB_NOUNROLL
+    for (int i = 0; i < pyobj_name::total_count; ++i)
+        Py_VISIT(pyobjects[i]);
+    return 0;
+}
+
+NB_NOINLINE int nb_module_clear(PyObject *m) {
+    PyObject** pyobjects = (PyObject**) PyModule_GetState(m);
+    NB_NOUNROLL
+    for (int i = 0; i < pyobj_name::total_count; ++i)
+        Py_CLEAR(pyobjects[i]);
+    return 0;
+}
+
+void nb_module_free(void *m) {
+    // Allow nanobind_##name##_exec to omit calling nb_module_clear on error.
+    (void) nb_module_clear((PyObject *) m);
+}
+
+
 static bool is_alive_value = false;
 static bool *is_alive_ptr = &is_alive_value;
 bool is_alive() noexcept { return *is_alive_ptr; }
@@ -317,29 +382,34 @@ static void internals_cleanup() {
 #endif
 }
 
-NB_NOINLINE void init(const char *name) {
+NB_NOINLINE void nb_module_exec(const char *name, PyObject *m) {
     if (internals)
         return;
 
+    check(init_pyobjects(m), "nanobind::detail::nb_module_exec(): "
+                             "could not initialize module state!");
+
 #if defined(PYPY_VERSION)
     PyObject *dict = PyEval_GetBuiltins();
 #elif PY_VERSION_HEX < 0x03090000
     PyObject *dict = PyInterpreterState_GetDict(_PyInterpreterState_Get());
 #else
     PyObject *dict = PyInterpreterState_GetDict(PyInterpreterState_Get());
 #endif
-    check(dict, "nanobind::detail::init(): could not access internals dictionary!");
+    check(dict, "nanobind::detail::nb_module_exec(): "
+                "could not access internals dictionary!");
 
     PyObject *key = PyUnicode_FromFormat("__nb_internals_%s_%s__",
                                          abi_tag(), name ? name : "");
-    check(key, "nanobind::detail::init(): could not create dictionary key!");
+    check(key, "nanobind::detail::nb_module_exec(): "
+               "could not create dictionary key!");
 
     PyObject *capsule = dict_get_item_ref_or_fail(dict, key);
     if (capsule) {
         Py_DECREF(key);
         internals = (nb_internals *) PyCapsule_GetPointer(capsule, "nb_internals");
-        check(internals,
-              "nanobind::detail::internals_fetch(): capsule pointer is NULL!");
+        check(internals, "nanobind::detail::nb_module_exec(): "
+                         "capsule pointer is NULL!");
         nb_meta_cache = internals->nb_meta;
         is_alive_ptr = internals->is_alive_ptr;
         Py_DECREF(capsule);
@@ -381,7 +451,7 @@ NB_NOINLINE void init(const char *name) {
 
     check(p->nb_module && p->nb_meta && p->nb_type_dict && p->nb_func &&
               p->nb_method && p->nb_bound_method,
-          "nanobind::detail::init(): initialization failed!");
+          "nanobind::detail::nb_module_exec(): initialization failed!");
 
 #if PY_VERSION_HEX < 0x03090000
     p->nb_func->tp_flags |= NB_HAVE_VECTORCALL;
@@ -476,7 +546,7 @@ NB_NOINLINE void init(const char *name) {
     capsule = PyCapsule_New(p, "nb_internals", nullptr);
     int rv = PyDict_SetItem(dict, key, capsule);
     check(!rv && capsule,
-          "nanobind::detail::init(): capsule creation failed!");
+          "nanobind::detail::nb_module_exec(): capsule creation failed!");
     Py_DECREF(capsule);
     Py_DECREF(key);
     internals = p;

src/nb_internals.h

@@ -420,6 +420,30 @@ struct nb_internals {
     size_t shard_count = 1;
 };
 
+// Names for the PyObject* entries in the per-module state array.
+// These names are scoped, but will implicitly convert to int.
+struct pyobj_name {
+    enum : int {
+        value_str = 0,      // string "value"
+        copy_str,           // string "copy"
+        from_dlpack_str,    // string "from_dlpack"
+        dunder_dlpack_str,  // string "__dlpack__"
+        max_version_str,    // string "max_version"
+        dl_device_str,      // string "dl_device"
+        string_count,
+
+        copy_tpl = string_count,  // tuple ("copy")
+        max_version_tpl, // tuple ("max_version")
+        dl_cpu_tpl,      // tuple (1, 0), which corresponds to nb::device::cpu
+        dl_version_tpl,  // tuple (dlpack::major_version, dlpack::minor_version)
+        total_count
+    };
+};
+
+static_assert(pyobj_name::total_count * sizeof(PyObject*) == NB_MOD_STATE_SIZE);
+
+extern PyObject **static_pyobjects;
+
 /// Convenience macro to potentially access cached functions
 #if defined(Py_LIMITED_API)
 #  define NB_SLOT(type, name) internals->type##_##name

tests/inter_module.cpp

@@ -4,6 +4,10 @@ Shared create_shared() {
     return { 123 };
 }
 
-bool check_shared(const Shared &shared) {
-    return shared.value == 123;
+bool check_shared(const Shared &shared, int expected) {
+    return shared.value == expected;
+}
+
+void increment_shared(Shared &shared) {
+    ++shared.value;
 }

tests/inter_module.h

@@ -11,4 +11,5 @@ struct EXPORT_SHARED Shared {
 };
 
 extern EXPORT_SHARED Shared create_shared();
-extern EXPORT_SHARED bool check_shared(const Shared &shared);
+extern EXPORT_SHARED bool check_shared(const Shared &shared, int expected);
+extern EXPORT_SHARED void increment_shared(Shared &shared);

tests/test_inter_module.py

@@ -4,10 +4,86 @@
 import pytest
 from common import xfail_on_pypy_darwin
 
+try:
+    from concurrent import interpreters  # Added in Python 3.14
+    def needs_interpreters(x):
+        return x
+except:
+    needs_interpreters = pytest.mark.skip(reason="interpreters required")
+
+
 @xfail_on_pypy_darwin
 def test01_inter_module():
     s = t1.create_shared()
-    assert t2.check_shared(s)
+    assert t2.check_shared(s, 123)
+    t2.increment_shared(s)
+    assert t2.check_shared(s, 124)
     with pytest.raises(TypeError) as excinfo:
         assert t3.check_shared(s)
     assert 'incompatible function arguments' in str(excinfo.value)
+
+
+@xfail_on_pypy_darwin
+def test02_reload_module():
+    s1 = t1.create_shared()
+    s2 = t1.create_shared()
+    assert s2 is not s1
+    assert type(s2) is type(s1)
+    t2.increment_shared(s2)
+    import importlib
+    new_t1 = importlib.reload(t1)
+    assert new_t1 is t1
+    s3 = new_t1.create_shared()
+    assert type(s3) is type(s1)
+    new_t2 = importlib.reload(t2)
+    assert new_t2 is t2
+    s4 = new_t1.create_shared()
+    assert type(s4) is type(s1)
+    assert new_t2.check_shared(s2, 124)
+
+
+@xfail_on_pypy_darwin
+def test03_reimport_module():
+    s1 = t1.create_shared()
+    s2 = t1.create_shared()
+    t2.increment_shared(s2)
+    import sys
+    del sys.modules['test_inter_module_1_ext']
+    import test_inter_module_1_ext as new_t1
+    assert new_t1 is not t1
+    s3 = new_t1.create_shared()
+    assert type(s3) is type(s1)
+    del sys.modules['test_inter_module_2_ext']
+    with pytest.warns(RuntimeWarning, match="'Shared' was already registered"):
+        import test_inter_module_2_ext as new_t2
+    assert new_t2 is not t2
+    s4 = new_t1.create_shared()
+    assert type(s4) is type(s1)
+    assert new_t2.check_shared(s2, 124)
+
+
+def run():
+    import sys
+    if 'tests' not in sys.path[0]:
+        import os
+        builddir = sys.path[0]
+        sys.path.insert(0, os.path.join(builddir, 'tests', 'Release'))
+        sys.path.insert(0, os.path.join(builddir, 'tests', 'Debug'))
+        sys.path.insert(0, os.path.join(builddir, 'tests'))
+    import test_inter_module_1_ext as new_t1
+    import test_inter_module_2_ext as new_t2
+    success = True
+    s = new_t1.create_shared()
+    success &= new_t2.check_shared(s, 123)
+    new_t2.increment_shared(s)
+    success &= new_t2.check_shared(s, 124)
+    return success
+
+@needs_interpreters
+def test04_subinterpreters():
+    assert run()
+    interp = interpreters.create()
+    with pytest.raises(interpreters.ExecutionFailed) as excinfo:
+        assert interp.call(run)
+    assert 'does not support loading in subinterpreters' in str(excinfo.value)
+    interp.close()

tests/test_inter_module_2.cpp

@@ -6,4 +6,5 @@ namespace nb = nanobind;
 NB_MODULE(test_inter_module_2_ext, m) {
     nb::class_<Shared>(m, "Shared");
     m.def("check_shared", &check_shared);
+    m.def("increment_shared", &increment_shared);
 }