Skip to content

[ROCm][VMM] Layer 1/4: RocmRawMemoryAllocation (hipMemCreate/hipMemRelease RAII)#928

Closed
phambinhfin wants to merge 1 commit into
mainfrom
phambinh/vmm-layer1-raw-memory
Closed

[ROCm][VMM] Layer 1/4: RocmRawMemoryAllocation (hipMemCreate/hipMemRelease RAII)#928
phambinhfin wants to merge 1 commit into
mainfrom
phambinh/vmm-layer1-raw-memory

Conversation

@phambinhfin

Copy link
Copy Markdown
Collaborator

Series (one PR per VMM layer)

This is the first in a series of four that splits the ROCm VMM allocator into small, reviewable layers (the combined work was previously one large branch, proposed upstream as openxla#40236). Splitting per layer keeps each PR well under the "large CL" threshold and lets each layer be reviewed independently, bottom-up.

  • PR 1/4: RocmRawMemoryAllocationhipMemCreate / hipMemRelease RAII ⟸ This PR
  • PR 2/4: RocmMemoryReservation — hipMemAddressReserve / hipMemMap / hipMemSetAccess / hipMemUnmap RAII
  • PR 3/4: RocmVmmAllocator — all-in-one allocator (create + reserve + map + setAccess)
  • PR 4/4: RocmDeviceAddressVmmAllocator — GPU timeline-based deferred deallocation + PJRT wiring

(The selective copy-vs-remap optimization continues to live in #926, which stacks on top of this series.)

What this PR adds

The lowest layer of the VMM stack: a standalone RAII handle around a single physical VMM allocation.

  • RocmRawMemoryAllocation owns the hipMemGenericAllocationHandle_t returned by hipMemCreate and releases it via hipMemRelease on destruction (move-only, no copy).
  • No dependency on any higher layer, so it builds and is testable on its own.

Tests

rocm_raw_memory_allocation_test.cc — allocate/release and move-semantics coverage.

Notes

  • ROCm-only target (tags = ["gpu", "rocm-only"]); no CUDA/NVIDIA impact.
  • +280 lines total.

…lease RAII)

First in a series that splits the ROCm VMM allocator (currently one large
branch, proposed upstream as openxla#40236) into one reviewable PR per
layer:

  Layer 1/4 - RocmRawMemoryAllocation: RAII wrapper for hipMemCreate /
              hipMemRelease  <= THIS PR
  Layer 2/4 - RocmMemoryReservation: RAII wrapper for hipMemAddressReserve /
              hipMemMap / hipMemSetAccess / hipMemUnmap
  Layer 3/4 - RocmVmmAllocator: all-in-one allocator (create + reserve + map
              + setAccess)
  Layer 4/4 - RocmDeviceAddressVmmAllocator: GPU timeline-based deferred
              deallocation + PJRT wiring

This layer adds a standalone RAII handle around a single physical VMM
allocation: it owns the hipMemGenericAllocationHandle_t returned by
hipMemCreate and releases it via hipMemRelease on destruction. It has no
dependency on the higher layers, so it can be reviewed and merged on its own.

Includes rocm_raw_memory_allocation_test.cc covering allocate/release and
move semantics.
@phambinhfin

Copy link
Copy Markdown
Collaborator Author

Retargeting this series to openxla/xla:main (the VMM allocator is proposed upstream as openxla#40236). Reopening Layer 1/4 as an upstream PR; closing this fork-internal one.

@phambinhfin phambinhfin closed this Jun 9, 2026
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