wip

Author: caiomcbr

python/mscclpp/language/internal/operations.py

@@ -330,7 +330,8 @@ def to_json(self):
         result["dst_buff"] = []
         for chunk in self.dst_buff:
             result["dst_buff"].append(chunk.to_json())
-        result["channel_ids"] = self.channel_ids
+        if self.channel_ids == ChannelType.port:
+            result["channel_ids"] = self.channel_ids
         result["channel_type"] = self.channel_type.value
         return result
 
@@ -437,7 +438,7 @@ def __add__(self, other):
                 remote_dst_buff=self.remote_dst_buff + other.dst_buff,
                 channel_ids=self.channel_ids,
                 put_channel_ids=self.put_channel_ids + other.channel_ids,
-                channel_type=self.channel_type,
+                channel_type=other.channel_type,
                 reduce_operation=self.reduce_operation,
                 packet=self.packet,
             )
@@ -460,10 +461,10 @@ def to_json(self):
             for chunk in self.remote_dst_buff:
                 result["dst_buff"].append(chunk.to_json())
 
-        if len(self.channel_ids) > 0:
+        """ if len(self.channel_ids) > 0:
             result["channel_ids"] = self.channel_ids
         if len(self.put_channel_ids) > 0:
-            result["output_channel_ids"] = self.put_channel_ids
+            result["output_channel_ids"] = self.put_channel_ids """
         if self.channel_type != ChannelType.none:
             result["channel_type"] = self.channel_type.value
         result["reduce_op"] = self.reduce_operation.value

python/mscclpp/language/tests/allgather.py

@@ -44,7 +44,7 @@ def allgather_example(name, num_threads_per_block, min_message_size, max_message
                 # Skip sending from a rank to itself
                 if src_rank != dst_rank:
                     # Define a channel from src_rank → dst_rank using memory channel
-                    ch = Channel(dst_rank, src_rank, ChannelType.memory)
+                    ch = Channel(dst_rank, src_rank)
                     # Step 1: source signals to indicate it is ready to receive data
                     ch.signal(tb=0, sync=None, relaxed=True)
                     # Step 2: wait for the destination rank to be ready

python/mscclpp/language/tests/allgather_pkt.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def allgather_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size):
+    chunksperloop = 1
+    collective = AllGather(gpu_size, chunksperloop, True)
+    with MSCCLPPProgram(
+        name,
+        collective,
+        gpu_size,
+        protocol="LL",
+        num_threads_per_block=num_threads_per_block,
+        use_double_scratch_buffer=True,
+        min_message_size=min_message_size,
+        max_message_size=max_message_size,
+    ):
+        # Creating Scratch Buffers
+        scratch_buffer = []
+        for gpu in range(gpu_size):
+            scratch_buffer.append(Buffer(gpu, gpu_size))
+
+        # Putting packet in the remote scratch buffer
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            for peer in range(1, gpu_size):
+                dst_rank = (gpu + peer) % gpu_size
+                ch = Channel(dst_rank, gpu)
+                tb = peer - 1
+                ch.put_packet(scratch_buffer[dst_rank][gpu : gpu + 1], output_buffer[gpu : gpu + 1], tb)
+
+        # Copying packet from local scratch buffer to local buffer
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            output_buffer = rank.get_output_buffer()
+            for peer in range(1, gpu_size):
+                dst_rank = (gpu + peer) % gpu_size
+                rank.copy(output_buffer[dst_rank: dst_rank + 1], scratch_buffer[gpu][dst_rank: dst_rank + 1], tb=gpu_size + peer - 2, from_packet=True)
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--num_gpus", type=int, help="number of gpus")
+parser.add_argument("--num_threads_per_block", type=int, default=1024, help="number of threads per block")
+parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size")
+parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size")
+
+args = parser.parse_args()
+
+allgather_example(args.name, args.num_gpus, args.num_threads_per_block, args.min_message_size, args.max_message_size)

python/mscclpp/language/tests/allreduce.py

@@ -10,7 +10,7 @@
 
 
 def allreduce_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size):
-    chunksperloop = 1
+    chunksperloop = gpu_size
     collective = AllReduce(gpu_size, chunksperloop, True)
     with MSCCLPPProgram(
         name,
@@ -25,7 +25,7 @@ def allreduce_example(name, gpu_size, num_threads_per_block, min_message_size, m
         # Creating Scratch Buffers
         scratch_buffer = []
         for gpu in range(gpu_size):
-            scratch_buffer.append(Buffer(gpu, 2 * gpu_size))
+            scratch_buffer.append(Buffer(gpu, 2 * (gpu_size - 1) * gpu_size))
 
         # Creating Channels
         channels = {}
@@ -40,33 +40,47 @@ def allreduce_example(name, gpu_size, num_threads_per_block, min_message_size, m
             input_buffer = rank.get_input_buffer()
             for peer in range(gpu_size):
                 if peer != gpu:
+                    scratch_index = gpu if gpu < peer else gpu - 1
+                    scratch_index *= gpu_size
+                    input_index = peer * gpu_size
+                    tb = peer if peer < gpu else peer - 1
                     channels[(peer, gpu)].put_packet(
-                        scratch_buffer[peer][gpu : gpu + 1], input_buffer[peer : peer + 1], 0
+                        scratch_buffer[peer][scratch_index : scratch_index + gpu_size], input_buffer[input_index : input_index + gpu_size], tb
                     )
 
         # Each rank performs a local reduction on the nth chunk
         for gpu in range(gpu_size):
-            chunks = []
-            for peer in range(gpu_size):
-                if peer != gpu:
-                    chunks.append(scratch_buffer[gpu][peer : peer + 1])
-            rank = Rank(gpu)
-            input_buffer = rank.get_input_buffer()
-            rank.reduce(input_buffer[gpu : gpu + 1], chunks, 0, packet=True)
-            for peer in range(gpu_size):
-                if peer != gpu:
-                    channels[(peer, gpu)].put_packet(
-                        scratch_buffer[peer][gpu_size + gpu : gpu_size + gpu + 1], input_buffer[gpu : gpu + 1], 0
-                    )
+            for index in range(gpu_size):
+                chunks = []
+                for peer in range(gpu_size):
+                    if peer != gpu:
+                        scratch_index = peer if peer < gpu else peer - 1
+                        scratch_index = scratch_index * gpu_size + index
+                        chunks.append(scratch_buffer[gpu][scratch_index : scratch_index + 1])
+                rank = Rank(gpu)
+                input_buffer = rank.get_input_buffer()
+                input_index = gpu * gpu_size + index
+                rank.reduce(input_buffer[input_index : input_index + 1], chunks, index, packet=True)
+                
+                for peer in range(gpu_size):
+                    if peer != gpu:
+                        scratch_index = gpu if gpu < peer else gpu - 1
+                        scratch_index = gpu_size * (gpu_size - 1) + scratch_index * gpu_size + index
+                        channels[(peer, gpu)].put_packet(
+                            scratch_buffer[peer][scratch_index : scratch_index + 1], input_buffer[input_index : input_index + 1], index
+                        )
 
         # Each rank get final result from scratch space
         for gpu in range(gpu_size):
             rank = Rank(gpu)
             input_buffer = rank.get_input_buffer()
             for peer in range(gpu_size):
                 if peer != gpu:
+                    input_index = peer * gpu_size
+                    scratch_index = peer if peer < gpu else peer - 1
+                    scratch_index = gpu_size * (gpu_size - 1) + scratch_index * gpu_size
                     rank.copy(
-                        input_buffer[peer : peer + 1], scratch_buffer[gpu][gpu_size + peer : gpu_size + peer + 1], 0
+                        input_buffer[input_index : input_index + gpu_size], scratch_buffer[gpu][scratch_index : scratch_index + gpu_size], peer, from_packet=True
                     )
 
         print(JSON())

python/mscclpp/language/tests/allreduce_naivy.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+from mscclpp.language.channel import *
+from mscclpp.language.rank import *
+from mscclpp.language.general import *
+from mscclpp.language.program import *
+from mscclpp.language.collectives import *
+
+
+def allreduce_example(name, gpu_size, num_threads_per_block, min_message_size, max_message_size):
+    chunksperloop = 1
+    collective = AllReduce(gpu_size, chunksperloop, True)
+    with MSCCLPPProgram(
+        name,
+        collective,
+        gpu_size,
+        protocol="LL",
+        num_threads_per_block=num_threads_per_block,
+        use_double_scratch_buffer=False,
+        min_message_size=min_message_size,
+        max_message_size=max_message_size,
+    ):
+        # Creating Scratch Buffers
+        scratch_buffer = []
+        for gpu in range(gpu_size):
+            scratch_buffer.append(Buffer(gpu, 2 * gpu_size))
+
+        # Creating Channels
+        channels = {}
+        for gpu in range(gpu_size):
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    channels[(peer, gpu)] = Channel(peer, gpu)
+
+        # Each rank sends the nth chunk to the nth rank into scratch space
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            input_buffer = rank.get_input_buffer()
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    channels[(peer, gpu)].put_packet(
+                        scratch_buffer[peer][gpu : gpu + 1], input_buffer[peer : peer + 1], 0
+                    )
+
+        # Each rank performs a local reduction on the nth chunk
+        for gpu in range(gpu_size):
+            chunks = []
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    chunks.append(scratch_buffer[gpu][peer : peer + 1])
+            rank = Rank(gpu)
+            input_buffer = rank.get_input_buffer()
+            rank.reduce(input_buffer[gpu : gpu + 1], chunks, 0, packet=True)
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    channels[(peer, gpu)].put_packet(
+                        scratch_buffer[peer][gpu_size + gpu : gpu_size + gpu + 1], input_buffer[gpu : gpu + 1], 0
+                    )
+
+        # Each rank get final result from scratch space
+        for gpu in range(gpu_size):
+            rank = Rank(gpu)
+            input_buffer = rank.get_input_buffer()
+            for peer in range(gpu_size):
+                if peer != gpu:
+                    rank.copy(
+                        input_buffer[peer : peer + 1], scratch_buffer[gpu][gpu_size + peer : gpu_size + peer + 1], 0, from_packet=True
+                    )
+
+        print(JSON())
+
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--name", type=str, help="name of the program")
+parser.add_argument("--num_gpus", type=int, help="number of gpus")
+parser.add_argument("--num_threads_per_block", type=int, default=1024, help="number of threads per block")
+parser.add_argument("--min_message_size", type=int, default=0, help="minimum message size")
+parser.add_argument("--max_message_size", type=int, default=2**64 - 1, help="maximum message size")
+
+args = parser.parse_args()
+
+allreduce_example(args.name, args.num_gpus, args.num_threads_per_block, args.min_message_size, args.max_message_size)