Feat: Impl copy

Author: astroC86

iris/__init__.py

@@ -33,6 +33,8 @@
     load,
     store,
     copy,
+    get,
+    put,
     atomic_add,
     atomic_cas,
     atomic_xchg,
@@ -86,6 +88,8 @@
     "load",
     "store",
     "copy",
+    "get",
+    "put",
     "atomic_add",
     "atomic_cas",
     "atomic_xchg",
@@ -102,4 +106,4 @@
     "INFO",
     "WARNING",
     "ERROR",
-]
+]

iris/iris.py

@@ -22,6 +22,7 @@
     >>> tensor = ctx.zeros(1024, 1024, dtype=torch.float32)
 """
 
+from iris.util import trap_if
 import triton
 import triton.language as tl
 
@@ -1532,7 +1533,7 @@ def store(pointer, value, from_rank, to_rank, heap_bases, mask=None):
 
 
 @triton.jit
-def copy(src_ptr, dst_ptr, from_rank, to_rank, heap_bases, mask=None):
+def copy(src_ptr, dst_ptr, from_rank, to_rank, cur_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
@@ -1548,6 +1549,51 @@ def copy(src_ptr, dst_ptr, from_rank, to_rank, heap_bases, mask=None):
         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.
 
+    Returns:
+        None
+    """
+
+    trap_if((cur_rank != from_rank) and (cur_rank != to_rank))
+
+    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)
+
+
+@triton.jit
+def get(from_ptr, to_ptr, from_rank, to_rank, heap_bases, mask=None):
+    """
+    Copies data from the specified rank's memory to the current rank's local memory.
+
+    This function performs a memory read operation by translating the from_ptr
+    from the current rank's address space to the from_rank's address space, loading data
+    from the from_rank memory location, and storing it to the local to_ptr.
+    If the from_rank is the same as the current rank, this function performs a local copy operation.
+
+    Args:
+        from_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's address space that will be translated to the from_rank's address space. Must be the current rank where the pointer is local.
+        to_ptr (triton.PointerType, or block of dtype=triton.PointerType): Pointer in the current rank's local memory where the data will be stored.
+        from_rank (int): The from_rank ID from which to read the data.
+        to_rank (int): The current rank ID where the data will be stored.
+        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 the data at address from_ptr[idx] and do not store to to_ptr[idx]. Defaults to None.
+
     Returns:
         None
 

iris/util.py

@@ -142,3 +142,24 @@ def do_bench(
 
     times = [s.elapsed_time(e) for s, e in zip(start_event, end_event)]
     return _summarize_statistics(times, quantiles, return_mode)
+
+
+@triton.jit
+def trap_if(cond):
+    drv = tl.zeros([1], dtype=tl.uint32)
+    cond_u32 = tl.where(cond, drv, drv + 1)
+    if tl.program_id(0) == 0:
+        tl.inline_asm_elementwise(
+            asm="""
+                    s_cmp_lg_u32 $1, 0
+                    s_cbranch_scc1 0f
+                    s_trap 2
+                0:
+                """,
+                constraints="=v,s",
+                args=[cond_u32],
+                dtype=tl.uint32,
+                is_pure=False,
+                pack=1,
+            )
+        

tests/unittests/test_copy.py

@@ -25,7 +25,7 @@ def copy_kernel(
     for target_rank in range(num_ranks):
         src_data = data + BLOCK_SIZE * cur_rank
         dest_data = results + BLOCK_SIZE * target_rank
-        iris.copy(src_data + offsets, dest_data + offsets, cur_rank, target_rank, heap_bases, mask)
+        iris.copy(src_data + offsets, dest_data + offsets, target_rank, cur_rank, cur_rank,  heap_bases, mask)
 
 
 @pytest.mark.parametrize(

tests/unittests/test_get.py

@@ -31,7 +31,7 @@ def get_kernel(
     # Loop over all ranks, get the stored data.
     # load to local register, accumulate.
     for target_rank in range(num_ranks):
-        iris.copy(data + offsets, results + offsets, cur_rank, target_rank, heap_bases, mask=mask)
+        iris.get(data + offsets, results + offsets, cur_rank, target_rank, heap_bases, mask=mask)
         acc += tl.load(results + offsets, mask=mask)
 
     # Store the accumulated value back to the output.
@@ -81,4 +81,4 @@ def test_get_api(dtype, BLOCK_SIZE):
         print(e)
         print("Expected:", expected)
         print("Actual:", results)
-        raise
+        raise

tests/unittests/test_put.py

@@ -29,7 +29,7 @@ def put_kernel(
     # Put data in all ranks
     # Doesn't matter which rank stores at the end, the data should all be the same at the end.
     for target_rank in range(num_ranks):
-        iris.copy(data + offsets, results + offsets, cur_rank, target_rank, heap_bases, mask=mask)
+        iris.put(data + offsets, results + offsets, cur_rank, target_rank, heap_bases, mask=mask)
 
 
 @pytest.mark.parametrize(
@@ -75,4 +75,4 @@ def test_put_api(dtype, BLOCK_SIZE):
         print(e)
         print("Expected:", expected)
         print("Actual:", results)
-        raise
+        raise