Fix merge issue

msccl/language/instruction_dag.py

@@ -53,6 +53,16 @@ def circular_dep_after_merge(op: Op, other_op: Op):
             frontier.append(n)
         frontier = frontier[1:]
 
+"""
+For case: op2.prev = [op1, op3]. op1.next = [op2]. op3.next = [op2]. And op1 and op2 are satisfied to merge.
+We only apply the merge if all previous ops of op2 are visited after the merge.
+"""
+def all_prevs_visited_after_merge(op: Op, other_op: Op):
+    step = op.step
+    for prev in other_op.prev:
+        if prev.step > step:
+            return False
+    return True
 
 def same_tb(op1: Op, op2: Op):
     return op1.tb == op2.tb and op1.channel == op2.channel

msccl/language/mscclpp/instruction_optimizer.py

@@ -10,6 +10,7 @@
     same_count,
     same_buf_dst,
     same_buf_src,
+    all_prevs_visited_after_merge,
 )
 from msccl.language.types import ChunkRef, ChannelType, MscclppInstruction as Instruction, Op, Threadblock
 
@@ -41,6 +42,7 @@ def try_merge_same_instructions(
             and same_count(op, next_op)
             and same_chan_type(op, next_op)
             and not circular_dep_after_merge(op, next_op)
+            and all_prevs_visited_after_merge(op, next_op)
         ):
             # Append the source chunks from next_op
             op.srcs.append(
@@ -85,6 +87,7 @@ def try_compact_instructions(
                 and same_chan_type(op, seq_op)
                 and same_count(op, seq_op)
                 and not circular_dep_after_merge(op, seq_op)
+                and all_prevs_visited_after_merge(op, seq_op)
             ):
                 # Append the source and destination chunks from seq_op
                 op.dsts.append(
@@ -124,6 +127,7 @@ def try_fuse_with_put(self, op: Op, next_op: Op, tb: Threadblock, queue: list) -
             and next_op.channel_type == ChannelType.sm
             and (op.channel_type == ChannelType.none or op.channel_type == ChannelType.sm)
             and not circular_dep_after_merge(op, next_op)
+            and all_prevs_visited_after_merge(op, next_op)
         ):
             if len(op.dsts) > 0 and op.dsts[0][0].buffer != next_op.dst.buffer:
                 return False
@@ -170,6 +174,7 @@ def try_fuse_instructions_using_proxy_channel(
             and same_chan_type(op, next_op)
             and op.channel_type == ChannelType.proxy
             and not circular_dep_after_merge(op, next_op)
+            and all_prevs_visited_after_merge(op, next_op)
         ):
             if op.inst == Instruction.put and next_op.inst == Instruction.signal:
                 op.inst = Instruction.put_with_signal