Implements copy function

[astroC86]

Motivation

Closes #98
Implements copy. I keep the get and put as wrappers to the copy function so that the tests pass

Technical Details

Test Plan

Test Result

(.iris_dev) root@2-6-0-gpu-mi300x1-192gb-devcloud-atl1:~/iris# mpirun -np 2 pytest ./tests/unittests/test_get.py 
============================================================================================== test session starts ==============================================================================================
platform linux -- Python 3.12.3, pytest-8.4.1, pluggy-1.6.0
============================================================================================== test session starts ==============================================================================================
platform linux -- Python 3.12.3, pytest-8.4.1, pluggy-1.6.0
rootdir: /root/iris
configfile: pyproject.toml
plugins: mpi-0.6
collecting ... rootdir: /root/iris
configfile: pyproject.toml
plugins: mpi-0.6
collected 16 items                                                                                                                                                                                              

collected 16 items                                                                                                                                                                                              

tests/unittests/test_get.py ................................                                                                                                                                                              [100%]                                                                                                                                                              [100%]



============================================================================================== 16 passed in 13.63s ==============================================================================================
============================================================================================== 16 passed in 13.63s ==============================================================================================
(.iris_dev) root@2-6-0-gpu-mi300x1-192gb-devcloud-atl1:~/iris# mpirun -np 2 pytest ./tests/unittests/test_put.py 
============================================================================================== test session starts ==============================================================================================
platform linux -- Python 3.12.3, pytest-8.4.1, pluggy-1.6.0
rootdir: /root/iris
configfile: pyproject.toml
plugins: mpi-0.6
collecting ... ============================================================================================== test session starts ==============================================================================================
platform linux -- Python 3.12.3, pytest-8.4.1, pluggy-1.6.0
rootdir: /root/iris
configfile: pyproject.toml
plugins: mpi-0.6
collected 16 items                                                                                                                                                                                              
collected 16 items                                                                                                                                                                                              

tests/unittests/test_put.py 
tests/unittests/test_put.py ................................                                                                                                                                                              [100%]                                                                                                                                                              [100%]

============================================================================================== 16 passed in 13.52s ==============================================================================================


============================================================================================== 16 passed in 13.51s ==============================================================================================

Submission Checklist

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

[neoblizz]

This PR will have to update examples that use put or get, there should be a few.

[mawad-amd]

Yeah, there are a few other places where we use iris.put. But, I think there something about the semantics of the copy is not correct.

[mawad-amd]

Shouldn't this be:

Suggested change

	iris.copy(data + offsets, results + offsets, cur_rank, target_rank, heap_bases, mask=mask)
	iris.copy(data + offsets, results + offsets, target_rank, cur_rank, heap_bases, mask=mask)

[mawad-amd]

The code I suggest fails the test btw but it shouldn't according to the docstring.

[astroC86]

I assume here the from_rank is always the current rank. I think if we want to allow for interchanging then perhaps this would more appropriate ? let me know what you think

@triton.jit
def copy(src_ptr, dst_ptr, from_rank, to_rank, cur_rank, heap_bases, mask=None):
    assert cur_rank == from_rank or cur_rank == to_rank, "Cannot copy between two arbitrary ranks"
    
    cur_base  = tl.load(heap_bases + cur_rank)

    from_base = tl.load(heap_bases + from_rank)
    to_base   = tl.load(heap_bases + to_rank)

    src_ptr_int = tl.cast(src_ptr, tl.uint64)
    src_offset = src_ptr_int - cur_base

    dst_ptr_int = tl.cast(dst_ptr, tl.uint64)
    dst_offset  = dst_ptr_int - cur_base

    from_base_byte = tl.cast(from_base, tl.pointer_type(tl.int8))
    to_base_byte   = tl.cast(to_base  , tl.pointer_type(tl.int8))

    translated_src = tl.cast(from_base_byte + src_offset, src_ptr.dtype)
    translated_dst = tl.cast(to_base_byte   + dst_offset, src_ptr.dtype)

    data = tl.load(translated_src, mask=mask)
    tl.store(translated_dst, data, mask=mask)

[mawad-amd]

The solution you proposed here is good but adds that additional overhead of the two translates. I have been thinking about this and I am not sure if there is away to resolve this cleanly.

I don’t really like the put/get names but maybe we will just stick to them for now. Let’s keep this PR open for now and we can come back to it later if we get better ideas. Thanks for your time looking into this and sorry this feature was not very well thought through.

[astroC86]

Hi no worries at all! Thanks a lot for taking the time to review my solution!