Intial impl of copy

Author: astroC86

iris/iris.py

@@ -407,11 +407,7 @@ def get(from_ptr, to_ptr, from_rank, to_rank, heap_bases, mask=None):
     Returns:
         None
     """
-    translated_from_ptr = __translate(from_ptr, from_rank, to_rank, heap_bases)
-
-    data = tl.load(translated_from_ptr, mask=mask)
-
-    tl.store(to_ptr, data, mask=mask)
+    copy(from_ptr, to_ptr, from_rank, to_rank , heap_bases, mask)
 
 
 @triton.jit
@@ -434,15 +430,29 @@ def put(from_ptr, to_ptr, from_rank, to_rank, heap_bases, mask=None):
     Returns:
         None
     """
-    translated_to_ptr = __translate(to_ptr, from_rank, to_rank, heap_bases)
+    copy(from_ptr, to_ptr, from_rank, to_rank, heap_bases, mask)
 
-    data = tl.load(from_ptr, mask=mask)
 
-    tl.store(translated_to_ptr, data, mask=mask)
+@triton.jit
+def copy(src_ptr, dst_ptr, from_rank, to_rank, heap_bases, mask=None):
+    """
+    Copies data from the specified rank's memory into the destination rank's memory.
+    This function performs the transfer by translating src_ptr from the from_rank's address
+    space to the to_rank's address space, performing a masked load from the translated
+    source, and storing the loaded data to dst_ptr in the to_rank memory location.
+    If from_rank and to_rank are the same, this function performs a local copy operation.
 
+    Args:
+        src_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the from_rank's local memory from which to read data.
+        dst_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the to_rank's local memory where the data will be written.
+        from_rank (int): The rank ID that owns src_ptr (source rank).
+        to_rank (int): The rank ID that will receive the data (destination rank).
+        heap_bases (triton.PointerType): Array containing the heap base addresses for all ranks.
+        mask (Block of triton.int1, optional): If mask[idx] is false, do not load from the translated src_ptr[idx] and do not store to dst_ptr[idx]. Defaults to None.
 
-@triton.jit
-def copy(dst_ptr, src_ptr, from_rank, to_rank, heap_bases, mask=None):
+    Returns:
+        None
+    """
     translated_src = __translate(src_ptr, from_rank, to_rank, heap_bases)
     data = tl.load(translated_src, mask=mask)
     tl.store(dst_ptr, data, mask=mask)

tests/unittests/test_copy.py

@@ -26,8 +26,8 @@ def copy_kernel(
         src_data  = data    + BLOCK_SIZE * cur_rank
         dest_data = results + BLOCK_SIZE * target_rank
         iris.copy(
-            dest_data + offsets,
             src_data  + offsets,
+            dest_data + offsets,
             cur_rank,
             target_rank,
             heap_bases,
@@ -77,15 +77,13 @@ def test_copy_get_semantics(dtype, BLOCK_SIZE):
     shmem.barrier()
 
     expected  = shmem.zeros((num_ranks, BLOCK_SIZE), dtype=dtype)
-    expected_2 = torch.zeros((num_ranks, BLOCK_SIZE), dtype=dtype, device="cuda")
     for rank_id in range(num_ranks):
-        expected[rank_id, :] = 999999
-        expected_2[rank_id, :] = (rank_id + num_ranks) * (cur_rank + 1)
+        expected[rank_id, :] = (rank_id + num_ranks) * (cur_rank + 1)
 
     try:
         torch.testing.assert_close(results, expected, rtol=0, atol=0)
     except AssertionError as e:
         print(e)
-        print("Expected:", expected_2)
+        print("Expected:", expected)
         print("Actual:", results)
         raise