WIp

Azure · Aug 30, 2024 · a5ed7fe · a5ed7fe
1 parent 57c44ba
commit a5ed7fe
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/msccl/language/mscclpp/__init__.py b/msccl/language/mscclpp/__init__.py
@@ -186,7 +186,9 @@ def _get_buffer_index(self, remote_rank, buffer, index):
             return buffer, self.prog.buffers[remote_rank][buffer].instance_size()
         return buffer, index
 
-    def _put(self, dst, buffer=None, index=-1, sendtb=-1, chan_type=ChannelType.sm, use_packet=False, trans_to_packet=True):
+    def _put(
+        self, dst, buffer=None, index=-1, sendtb=-1, chan_type=ChannelType.sm, use_packet=False, trans_to_packet=False
+    ):
         self.prog.check_buffer_exists(dst, buffer)
         assert self.rank != dst, "Cannot put to the same rank"
         buffer, index = self._get_buffer_index(dst, buffer, index)
@@ -196,10 +198,7 @@ def _put(self, dst, buffer=None, index=-1, sendtb=-1, chan_type=ChannelType.sm,
         dst_chunkref = self.prog.get_ref(dst, buffer, index, self.size)
         self.prog.apply_send(self.rank, self.buffer, self.index, dst, buffer, index, self.size)
         if use_packet:
-            if trans_to_packet:
-                self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, chan_type, trans_to_packet)
-            else:
-                self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, chan_type)
+            self.prog.instr_dag.add_put(self.rank, self, dst_chunkref, sendtb, chan_type, trans_to_packet)
             self.prog.instr_dag.add_signal(self.rank, self, dst_chunkref, -1, ChannelType.none)
             self.prog.instr_dag.add_wait(dst, dst_chunkref, self, -1, ChannelType.none)
         else:

diff --git a/tests/configs/test-config.json b/tests/configs/test-config.json
@@ -62,5 +62,25 @@
     {
         "filename": "pipeline_a100_ring.py",
         "args": ["8", "4", "2"]
+    },
+    {
+        "filename": "mscclpp/allreduce_a100_allpairs_packet_mscclpp.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "mscclpp/allreduce_a100_allpairs_sm_mscclpp_get.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "mscclpp/allreduce_a100_allpairs_sm_mscclpp.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "mscclpp/allreduce_a100_ring_mscclpp.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "mscclpp/send_recv_a100_packet.mscclpp.py",
+        "args": ["2"]
     }
 ]