Fix DMA-BUF Export/Import with PyTorch Caching Allocator Offsets

[mawad-amd]

Motivation

DMA-BUF export/import was failing when tensors were suballocated from PyTorch's caching allocator. hipMemGetHandleForAddressRange exports a handle to the entire allocation buffer, not just the specific tensor's memory. Without offset correction, imports would map to the wrong location, causing data corruption.

This became apparent after importing iris.ops (which loads tritonBLAS), causing subsequent tensors to be suballocated with non-zero offsets.

Technical Details

Changes to iris/hip.py:

• export_dmabuf_handle() - Now returns (fd, base_ptr, base_size) tuple. Uses hipMemGetAddressRange() to query allocation base and size.
• import_dmabuf_handle() - Now accepts original_ptr and base_ptr parameters. Calculates offset and returns mapped_base + offset.

Changes to iris/allocators/torch_allocator.py:

• Updated peer memory exchange to transmit metadata (base_ptr, base_size, heap_ptr) along with FD using struct.pack('QQQ', ...).

Changes to tests/unittests/test_dmabuf_apis.py:

• Updated all tests to use new API.
• Added test_dmabuf_with_offset() to explicitly validate offset handling.

Breaking Change: API changed from returning single fd to tuple (fd, base_ptr, base_size).

Test Plan

• All 6 DMA-BUF unit tests pass (including new offset test)
• Multi-rank tests verified with 1 and 2 ranks
• test_dmabuf_with_offset() explicitly tests the fix by allocating two tensors (forcing offset > 0) and verifying correct data import

Test Result

✅ All 6 DMA-BUF tests pass:

test_dmabuf_export                    PASSED [ 16%]
test_dmabuf_import                    PASSED [ 33%]
test_dmabuf_export_import_roundtrip   PASSED [ 50%]
test_iris_symmetric_heap_creation     PASSED [ 66%]
test_dmabuf_with_offset               PASSED [ 83%]
test_dmabuf_multirank_exchange        PASSED [100%]

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This pull request fixes a critical bug in DMA-BUF export/import functionality when working with PyTorch's caching allocator. The issue arose because hipMemGetHandleForAddressRange exports handles to entire allocation buffers rather than specific sub-allocations, causing data corruption when tensors had non-zero offsets within those buffers.

Changes:

• Modified export_dmabuf_handle() to return a tuple (fd, base_ptr, base_size) instead of just fd, using hipMemGetAddressRange to query the base allocation information
• Updated import_dmabuf_handle() to accept offset parameters and correctly map to the intended memory location
• Updated TorchAllocator to pack/unpack the additional metadata when exchanging handles between ranks
• Added comprehensive tests including a dedicated offset test to validate the fix

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 6 comments.

File	Description
iris/hip.py	Added hipMemGetAddressRange call to query base allocation; modified return type to tuple; added offset calculation in import function
iris/allocators/torch_allocator.py	Updated handle exchange to pack/unpack metadata (base_ptr, base_size, heap_ptr); updated import call to use offset parameters
tests/unittests/test_dmabuf_apis.py	Updated all existing tests for new API; added test_dmabuf_with_offset() to explicitly test offset handling