WIP

Azure · Sep 12, 2024 · ac1dfec · ac1dfec
1 parent 8006716
commit ac1dfec
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 14 deletions.
diff --git a/msccl/language/mscclpp/instruction_dag.py b/msccl/language/mscclpp/instruction_dag.py
@@ -348,7 +348,7 @@ def _optimize_get_put(self):
                     fused = False
                     if op.inst in campactable_inst:
                         if len(queue) > 1:
-                            fused = optimizer.try_compact_instruction(op, tb, queue, op.inst, same_src_dst_buffer_type)
+                            fused = optimizer.try_compact_instructions(op, tb, queue, op.inst, same_src_dst_buffer_type)
 
                     if fused:
                         continue
@@ -367,18 +367,8 @@ def _compact_signal_flush_wait(self):
                 while len(queue) > 0:
                     op = queue[0]
                     fused = False
-                    if op.inst == Instruction.signal:
-                        fused = optimizer.try_compact_instruction(
-                            op, tb, queue, Instruction.signal, same_src_dst_buffer_type
-                        )
-                    elif op.inst == Instruction.flush:
-                        fused = optimizer.try_compact_instruction(
-                            op, tb, queue, Instruction.flush, same_src_dst_buffer_type
-                        )
-                    elif op.inst == Instruction.wait:
-                        fused = optimizer.try_compact_instruction(
-                            op, tb, queue, Instruction.wait, same_src_dst_buffer_type
-                        )
+                    if op.inst == Instruction.signal or op.inst == Instruction.wait or op.inst == Instruction.flush:
+                        fused = optimizer.try_compact_instructions(op, tb, queue, op.inst, same_src_dst_buffer_type)
                     if fused:
                         continue
                     queue = queue[1:]

diff --git a/msccl/language/mscclpp/instruction_optimizer.py b/msccl/language/mscclpp/instruction_optimizer.py
@@ -64,7 +64,7 @@ def try_merge_same_instructions(
             return True
         return False
 
-    def try_compact_instruction(
+    def try_compact_instructions(
         self, op: Op, tb: Threadblock, queue: list, inst_type: Instruction, same_src_dst_func: callable
     ) -> bool:
         """