Support Fuse Operation on MSCCPP DSL

python/mscclpp/language/collectives.py

@@ -68,3 +68,22 @@ def init_buffers(self):
             }
             rank_buffers.append(buffers)
         return rank_buffers
+
+
+class AllReduce(Collective):
+    def __init__(self, num_ranks, chunk_factor, inplace):
+        Collective.__init__(self, num_ranks, chunk_factor, inplace)
+        self.name = "allreduce"
+
+    # Initializes input buffer for an allgather
+    def init_buffers(self):
+        rank_buffers = []
+        for rank in range(self.num_ranks):
+            input_buffer_size = self.num_ranks * self.chunk_factor
+            output_buffer_size = self.num_ranks * self.chunk_factor
+            buffers = {
+                BufferType.input: BaseBuffer(rank, BufferType.input, 0, input_buffer_size),
+                BufferType.output: BaseBuffer(rank, BufferType.output, 0, output_buffer_size),
+            }
+            rank_buffers.append(buffers)
+        return rank_buffers

python/mscclpp/language/general.py

@@ -2,4 +2,5 @@
 
 
 def JSON():
+    get_program().optimize_operations()
     return get_program().to_json()

python/mscclpp/language/internal/gpu.py

@@ -34,11 +34,11 @@ def setup_channel(self, tb: int, channel) -> int:
         return self.threadblocks[tb].add_channel(channel)
 
     def add_remote_buffer(self, tb: int, remote_buffer: RemoteBuffer) -> int:
-        if (remote_buffer.rank, remote_buffer.type) not in self.remote_buffers:
+        if (remote_buffer.remote_rank, remote_buffer.type) not in self.remote_buffers:
             remote_buffer.set_id()
-            self.remote_buffers[(remote_buffer.rank, remote_buffer.type)] = remote_buffer
+            self.remote_buffers[(remote_buffer.remote_rank, remote_buffer.type)] = remote_buffer
         else:
-            gpu_remote_buffer = self.remote_buffers[(remote_buffer.rank, remote_buffer.type)]
+            gpu_remote_buffer = self.remote_buffers[(remote_buffer.remote_rank, remote_buffer.type)]
             gpu_remote_buffer.channel_access.update(remote_buffer.channel_access)
             remote_buffer = gpu_remote_buffer
 
@@ -53,6 +53,10 @@ def add_operation(self, tb: int, operation: BaseOperation):
 
         self.threadblocks[tb].add_operation(operation)
 
+    def optimize_operations(self):
+        for tb in self.threadblocks:
+            tb.optimize_operations()
+
     def to_json(self) -> dict:
         return {
             "id": self.id,