Skip to content

Lazily load nvmath subpackages and bindings wrappers - #69

Draft
juenglin wants to merge 1 commit into
NVIDIA:mainfrom
juenglin:nvmath-python-lazy-bindings
Draft

Lazily load nvmath subpackages and bindings wrappers#69
juenglin wants to merge 1 commit into
NVIDIA:mainfrom
juenglin:nvmath-python-lazy-bindings

Conversation

@juenglin

Copy link
Copy Markdown

Summary

Importing any nvmath.bindings Cython submodule (for example CuPy's from nvmath.bindings.cycurand cimport ...) previously executed nvmath/__init__.py, which imported fft / linalg / sparse / tensor and every library wrapper, including optional *Mp and NVSHMEM modules. That mapped ~130 extra .so files (~350 MB RSS) into the process and slowed Python object construction process-wide (visible on cupy.empty_like and even on bytearray(256)).

This PR loads those subpackages and wrappers on first access (PEP 562). import nvmath.bindings.cycurand now maps cycurand and its own deps only. Public names, __all__, and "optional modules are None when unavailable" are unchanged. Native-library preload for host APIs still runs when nvmath.fft / linalg / sparse / tensor are actually imported.

Test plan

  • pytest tests/nvmath_tests/test_bindings_init.py
  • pytest tests/nvmath_tests/linalg/generic/matmul/test_qualifier.py
  • pytest tests/nvmath_tests/linalg/test_layout.py
  • pytest tests/nvmath_tests/ndbuffer/test_no_cupy.py
  • A/B: import nvmath.bindings.cycurand — 130 → 7 extra .so, 347 MB → 8.6 MB RSS

…ibrary wrapper and host API into the process.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant