One shot all reduce & symm mem sync

stack-info: PR: #245, branch: joydddd/stack/12

Author: joydddd

examples/all_reduce.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+
+import helion
+
+
+@helion.jit(
+    config=helion.Config(
+        block_sizes=[24],
+        num_warps=32,
+        indexing="pointers",
+    ),
+    static_shapes=True,
+)
+def one_shot_all_reduce_kernel(
+    buffer_ptr_addrs,
+    signal_pad_ptrs,
+    output_ptr,
+    numel: tl.constexpr,
+    rank: tl.constexpr,
+    world_size: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    output = torch.empty_like(x)
+    ptx_utils.symm_mem_sync(
+        signal_pad_ptrs, None, rank, world_size, hasSubsequenceMemAccess=True
+    )
+
+    pid = tl.program_id(axis=0)
+    buffer_ptr_addrs = buffer_ptr_addrs.to(tl.pointer_type(tl.uint64))
+    output_ptr = output_ptr.to(tl.pointer_type(tl.bfloat16))
+    block_start = pid * BLOCK_SIZE
+
+    while block_start < numel:
+        # Each thread processes 128 bits.
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < numel
+
+        acc = tl.zeros((BLOCK_SIZE,), dtype=tl.bfloat16)
+        for i in range(world_size):
+            buffer_ptr = tl.load(buffer_ptr_addrs + i).to(tl.pointer_type(tl.bfloat16))
+            tl.multiple_of(buffer_ptr, 16)
+            x = tl.load(buffer_ptr + offsets, mask=mask)
+            acc += x
+        tl.store(output_ptr + offsets, acc, mask=mask)
+        block_start += tl.num_programs(axis=0) * BLOCK_SIZE
+
+    ptx_utils.symm_mem_sync(
+        signal_pad_ptrs, None, rank, world_size, hasPreviousMemAccess=True
+    )
+
+
+def one_shot_all_reduce(tensor: torch.Tensor, **kwargs) -> torch.Tensor:
+    config = {
+        "max_num_blocks": kwargs.get("max_num_blocks", 24),
+        "num_warps": kwargs.get("num_warps", 32),
+        "BLOCK_SIZE": kwargs.get("BLOCK_SIZE", 8192),
+    }
+
+    assert tensor.dtype == torch.bfloat16, "Only bfloat16 is supported for now."
+    assert tensor.numel() % 8 == 0, "The number of elements must be 128-bit aligned."
+    assert config["BLOCK_SIZE"] % (config["num_warps"] * 32) == 0, (
+        "BLOCK_SIZE must be a multiple of num_warps * 32"
+    )
+
+    num_blocks = min(
+        triton.cdiv(tensor.numel(), config["BLOCK_SIZE"]), config["max_num_blocks"]
+    )
+
+    symm_mem_hdl = symm_mem.rendezvous(tensor, group=dist.group.WORLD)
+    output = torch.empty_like(tensor)
+
+    signal_pads = tuple(
+        [
+            symm_mem_hdl.get_signal_pad(i, dtype=torch.int32)
+            for i in range(symm_mem_hdl.world_size)
+        ]
+    )
+
+    one_shot_all_reduce_kernel[(num_blocks, 1, 1)](
+        symm_mem_hdl.buffer_ptrs_dev,
+        signal_pads,
+        output,
+        numel=tensor.numel(),
+        rank=symm_mem_hdl.rank,
+        world_size=symm_mem_hdl.world_size,
+        BLOCK_SIZE=config["BLOCK_SIZE"],
+        num_warps=config["num_warps"],
+    )
+
+    return output

helion/_compiler/type_propagation.py

@@ -1203,6 +1203,18 @@ def populate_symbol_origins(self, origin: Origin) -> None:
         for i, subtype in enumerate(self.element_types):
             subtype.populate_symbol_origins(GetItemOrigin(origin, i))
 
+    def propagate_getitem(self, key: TypeInfo, origin: Origin) -> TypeInfo:
+        # Check if all elements have the same type
+        first_type = self.element_types[0]
+        for element_type in self.element_types[1:]:
+            if type(element_type) != type(first_type):
+                raise exc.TypeInferenceError(
+                    f"Sequence contains mixed types: cannot safely index. "
+                    f"Found {type(first_type).__name__} and {type(element_type).__name__}"
+                )
+
+        return first_type
+
     def merge(self, other: TypeInfo) -> TypeInfo:
         if isinstance(other, SequenceType):
             self_elements = self.element_types

test/test_distributed.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+import torch.distributed._symmetric_memory as symm_mem
+from torch.testing._internal.common_distributed import MultiProcessTestCase
+from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
+from torch.testing._internal.common_utils import instantiate_parametrized_tests
+from torch.testing._internal.common_utils import run_tests
+
+import helion
+from helion._testing import code_and_output
+import helion.language as hl
+
+
+@helion.jit
+def symm_mem_sync_kernel(
+    remote_signal_pad_ptrs: torch.Tensor,
+    local_signal_pad: torch.Tensor,
+    rank: hl.constexpr,
+) -> None:
+    N, world_size = local_signal_pad.size()
+    world_size = hl.specialize(world_size)
+
+    assert world_size == remote_signal_pad_ptrs.size(0)
+    for n in hl.grid(N):
+        for multicast_tile in hl.tile(world_size, block_size=world_size):
+            peer_bar_offset = (n * world_size + rank) * 4
+            hl.signal(
+                remote_signal_pad_ptrs[multicast_tile] + peer_bar_offset,
+                wait_for=0,
+                signal=1,
+                op="atomic_cas",
+                sem="relaxed",
+                scope="sys",
+                skip_sync=True,
+                as_ptrs=True,
+            )
+            hl.wait(
+                local_signal_pad,
+                [n, multicast_tile],
+                signal=1,
+                update=0,
+                scope="sys",
+                op="atomic_cas",
+            )
+
+
+@instantiate_parametrized_tests
+class SymmMemBarrier(MultiProcessTestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        self._spawn_processes()
+
+    @property
+    def world_size(self) -> int:
+        # world_size > 2 is needed to verify accumulation order
+        return 4
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device(f"cuda:{self.rank}")
+
+    def _init_process(self):
+        torch.cuda.set_device(self.device)
+        store = dist.FileStore(self.file_name, self.world_size)
+        dist.init_process_group(
+            backend="nccl",
+            world_size=self.world_size,
+            rank=self.rank,
+            store=store,
+        )
+        torch.manual_seed(42 + self.rank)
+
+    @skip_if_lt_x_gpu(4)
+    def test_symm_mem_barrier(self):
+        self._init_process()
+        t = symm_mem.empty(4096, device=self.device)
+        symm_mem_hdl = symm_mem.rendezvous(t, group=dist.group.WORLD)
+        local_signal_pad_t = symm_mem_hdl.get_signal_pad(
+            symm_mem_hdl.rank, (32, symm_mem_hdl.world_size), dtype=torch.int32
+        )
+        signa_pad_pointers_t = torch.as_tensor(
+            symm_mem_hdl.signal_pad_ptrs, dtype=torch.uint64
+        ).to(self.device)
+
+        code, result = code_and_output(
+            symm_mem_sync_kernel,
+            (
+                signa_pad_pointers_t,
+                local_signal_pad_t,
+                symm_mem_hdl.rank,
+            ),
+        )
+
+        signal_pad = symm_mem_hdl.get_signal_pad(symm_mem_hdl.rank)
+        assert signal_pad.eq(0).all().item()
+
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    run_tests()