Memory refactor #1205

Andy-Jost · 2025-10-31T22:41:52Z

Major refactoring of the memory package.

Overview
This PR refactors the _memory.pyx module into a dedicated package (_memory/) to address its growing size and complexity, which were hindering further development. The primary goals are to physically separate the code into more manageable submodules, simplify the internal logic, and enhance the overall structure, including the addition of .pxd headers for better Cython integration.

Major Changes

Split _memory.pyx into submodules, the major ones being the following:
- Buffers: _buffer.*
- Device memory resources: _dmr.*
- IPC (Inter-Process Communication): _ipc.*
- Virtual memory management: _vmm.*
Introduced Cython headers (.pxd) for public definitions to improve modularity and type safety.
Refactored DeviceMemoryResource to isolate IPC-related code, reducing coupling.
Simplified IPC implementation by adding an IPCData class to encapsulate relevant data members and eliminating a redundant uuid field.
Streamlined the class hierarchy by removing unnecessary classes.
Simplified the Cython interface for memory allocation and deallocation operations.

Minor Improvements

Added __all__ lists to modules for explicit control over exports.
Extracted long implementation functions from class definitions to make classes more concise and readable.
Renamed various private attributes and methods for consistency (e.g., _handle instead of _mempool_handle).
Consolidated and alphabetized property definitions for better organization.
Converted additional classes and functions to Cython for performance gains.

copy-pr-bot · 2025-10-31T22:41:55Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Andy-Jost · 2025-11-03T17:33:00Z

/ok to test f13a44e

github-actions · 2025-11-03T17:43:00Z

Doc Preview CI
🚀 View preview at https://nvidia.github.io/cuda-python/pr-preview/pr-1205/
https://nvidia.github.io/cuda-python/pr-preview/pr-1205/cuda-core/
https://nvidia.github.io/cuda-python/pr-preview/pr-1205/cuda-bindings/
https://nvidia.github.io/cuda-python/pr-preview/pr-1205/cuda-pathfinder/
Preview will be ready when the GitHub Pages deployment is complete.

Andy-Jost · 2025-11-03T22:07:18Z

/ok to test 7c97d22

rparolin · 2025-11-04T01:06:14Z

@Andy-Jost To frame the code review, can you fill in the PR description with more details about what the goals of the refactor were.

Andy-Jost

This one is a beast. I tried to leave some helpful comments.

Andy-Jost · 2025-11-04T01:56:08Z

cuda_core/cuda/core/experimental/_memory/__init__.py

This file assembles the _memory package by combining the public elements of each submodule. This should match the public interface of the old _memory.pyx module.

Andy-Jost · 2025-11-04T01:57:44Z

cuda_core/cuda/core/experimental/_memory/_buffer.pxd

This file contains Cython declarations related to Buffer. I also put the declaration of MemoryResource here because I couldn't find a better place.

Classes _cyBuffer and _cyMemoryResource were eliminated.

Andy-Jost · 2025-11-04T01:58:00Z

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

This file contains the implementation of Buffer.

Andy-Jost · 2025-11-04T01:58:44Z

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

+    def _clear(self):
+        self._ptr = 0
+        self._size = 0
+        self._mr = None
+        self._ptr_obj = None
+        self._alloc_stream = None


This is used by VMM. I did not try to change the logic of that code.

Andy-Jost · 2025-11-04T01:59:27Z

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

+        stream: Stream = None
+    ) -> Buffer:
+        """Import a buffer that was exported from another process."""
+        return _ipc.Buffer_from_ipc_descriptor(cls, mr, ipc_buffer, stream)


Implementation of IPC functions has been moved to the _ipc module.

Andy-Jost · 2025-11-04T02:17:20Z

cuda_core/cuda/core/experimental/_memory/_vmm.py

+        raise_if_driver_error(res2)
+
+        # Invalidate the old buffer so its destructor won't try to free again
+        buf._clear()


This was the only change to the VMM code. Cf. _memory.pyx:1432-5

Andy-Jost · 2025-11-04T02:19:04Z

cuda_core/cuda/core/experimental/_device.pyx

                        )
                    )
                if attr == 1:
+                    from cuda.core.experimental._memory import DeviceMemoryResource


I believe this import should be delayed (like this) to avoid a circular dependency.

Andy-Jost · 2025-11-04T02:19:12Z

cuda_core/cuda/core/experimental/_device.pyx

+                    from cuda.core.experimental._memory import DeviceMemoryResource
                    device._mr = DeviceMemoryResource(dev_id)
                else:
+                    from cuda.core.experimental._memory import _SynchronousMemoryResource


Andy-Jost · 2025-11-04T02:19:35Z

cuda_core/cuda/core/experimental/_device.pyx


    @memory_resource.setter
    def memory_resource(self, mr):
+        from cuda.core.experimental._memory import MemoryResource


Andy-Jost · 2025-11-04T02:21:47Z

cuda_core/tests/memory_ipc/test_serialize.py

        # Receive the memory resource.
        handle = mp.reduction.recv_handle(conn)
        mr = DeviceMemoryResource.from_allocation_handle(device, handle)
-        os.close(handle)


A small functional change.

Fixed in 0d5f08b

Andy-Jost · 2025-11-04T15:21:19Z

/ok to test 0fac800

…a file descriptor (caller closes the fd).

Andy-Jost · 2025-11-04T19:35:16Z

/ok to test 0d5f08b

Andy-Jost · 2025-11-04T21:08:21Z

/ok to test 567ea2c

rparolin · 2025-11-07T19:27:48Z

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

+    def _clear(self):
+        self._ptr = 0
+        self._size = 0
+        self._mr = None


Nit: Consider renaming _mr -> _memory_resource or mem_resource.

rparolin · 2025-11-07T19:42:57Z

cuda_core/cuda/core/experimental/_memory/_buffer.pyx

+        stream: Stream | None = None
+    ):
+        cdef Buffer self = Buffer.__new__(cls)
+        self._ptr = <intptr_t>(int(ptr))


Shouldn't this be a uintptr_t or a uint64_t?

… uintptr_t.

Andy-Jost · 2025-11-07T20:43:52Z

/ok to test cf4dc9d

Andy-Jost · 2025-11-10T20:38:35Z

/ok to test 19e4b8f

Andy-Jost added 25 commits October 28, 2025 15:43

Resolve a Cython build warning.

37849a3

Make memory module into a package.

ac8a69c

Rename cyStream to _cyStream for consistency.

123aa24

Move defs to memory.pxd header

fe4b67e

Separate VMM.

ce77d44

Weaken dependencies from device to memory module.

c5179bc

Move LegacyPinnedMemoryResource to a submodule.

e192748

Move _SynchronousMemoryResource into a submodule.

729c900

Partly separates the IPC implementation.

8735455

Move IPC registry to ipc module.

b2517f6

Collect and reorder DeviceMemoryResource properties.

a61317a

Move more IPC implementation out of DeviceMemoryResource.

0e2d1d8

Minor refactoring.

5387629

Move Buffer IPC implementation.

7fa38ca

Simplify the class hierarchy (remove _cyBuffer and _cyMemoryResource).

f357abd

Refactor to shrink Cython interface.

89057f9

Simplify Buffer close.

00b60eb

Refactor DeviceMemoryResource.__init__.

ecc9405

Move Buffer into a separate module.

228936b

Refactors DeviceMemoryResource IPC implementation.

9a86bde

Removes superfluous _uuid member of DeviceMemoryResource.

c7f6cde

Adds __all__ lists.

216b4fb

Prepend underscore to submodules, add a test for package contents.

6a30a39

Refactor IPC data of DMR into IPCData class.

229ddc6

General clean up.

0fd3ca9

Andy-Jost requested review from cpcloud, leofang and mdboom and removed request for cpcloud October 31, 2025 22:41

Andy-Jost force-pushed the memory-refactor branch 2 times, most recently from feda70e to 52164c0 Compare November 3, 2025 17:30

Andy-Jost force-pushed the memory-refactor branch from f13a44e to d3dc347 Compare November 3, 2025 21:52

Touch-ups

44f7587

Andy-Jost force-pushed the memory-refactor branch from d3dc347 to 44f7587 Compare November 3, 2025 22:06

Merge branch 'main' into memory-refactor

7c97d22

Cythonize DeviceMemoryResourceAttributes.

0fac800

Andy-Jost commented Nov 4, 2025

View reviewed changes

Restore previous behavior for DMR.from_allocation_handle when passed …

0d5f08b

…a file descriptor (caller closes the fd).

Andy-Jost mentioned this pull request Nov 4, 2025

Implement non-pooling memory allocation in DeviceMemoryResource Andy-Jost/cuda-python#1

Draft

Merge branch 'main' into memory-refactor

567ea2c

Andy-Jost added the cuda.core Everything related to the cuda.core module label Nov 4, 2025

rparolin reviewed Nov 7, 2025

View reviewed changes

Andy-Jost added 2 commits November 7, 2025 12:42

Rename _mr to _memory_resource. Change pointer types from intptr_t to…

7315e28

… uintptr_t.

Merge branch 'main' into memory-refactor

cf4dc9d

leofang assigned Andy-Jost Nov 10, 2025

leofang added this to the cuda.core beta 9 milestone Nov 10, 2025

leofang added enhancement Any code-related improvements P0 High priority - Must do! labels Nov 10, 2025

Merge branch 'main' into memory-refactor

19e4b8f

Memory refactor #1205

Are you sure you want to change the base?

Memory refactor #1205

Conversation

Andy-Jost commented Oct 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Major refactoring of the memory package.

Uh oh!

copy-pr-bot bot commented Oct 31, 2025

Uh oh!

Andy-Jost commented Nov 3, 2025

Uh oh!

github-actions bot commented Nov 3, 2025

Preview will be ready when the GitHub Pages deployment is complete.

Uh oh!

Andy-Jost commented Nov 3, 2025

Uh oh!

rparolin commented Nov 4, 2025

Uh oh!

Andy-Jost left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andy-Jost commented Nov 4, 2025

Uh oh!

Andy-Jost commented Nov 4, 2025

Uh oh!

Andy-Jost commented Nov 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andy-Jost commented Nov 7, 2025

Uh oh!

Andy-Jost commented Nov 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Andy-Jost commented Oct 31, 2025 •

edited

Loading