From e7cff899ceba65873ded66b69a0e0efda2bdea82 Mon Sep 17 00:00:00 2001 From: Caio Rocha <164253795+caiomcbr@users.noreply.github.com> Date: Fri, 7 Feb 2025 11:24:27 -0800 Subject: [PATCH] Adjusting BFS to seek circular dependencies in the msccl-tools DAG (#459) Co-authored-by: Binyang Li --- python/mscclpp/language/ir.py | 4 +++- python/mscclpp/language/utils.py | 5 ++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/mscclpp/language/ir.py b/python/mscclpp/language/ir.py index 16a5b67f..3a1408ac 100644 --- a/python/mscclpp/language/ir.py +++ b/python/mscclpp/language/ir.py @@ -331,7 +331,9 @@ class _NopConverter(_OpConverter): def to_json(self, op: Op, tb_channel_dict: dict) -> _JsonInstruction: return _JsonInstruction( name=op.inst.value, - deps=list(map(lambda dep: {"tb": dep.tb, "step": dep.step}, op.depends)), + deps=sorted( + list(map(lambda dep: {"tb": dep.tb, "step": dep.step}, op.depends)), key=lambda x: (x["tb"], x["step"]) + ), ) diff --git a/python/mscclpp/language/utils.py b/python/mscclpp/language/utils.py index c97af881..02432b70 100644 --- a/python/mscclpp/language/utils.py +++ b/python/mscclpp/language/utils.py @@ -37,6 +37,7 @@ def merge_op(op: Op, other_op: Op): def circular_dep_after_merge(op: Op, other_op: Op): root = set([op, other_op]) frontier = set(op.next) + visited = set() if other_op in frontier: frontier.remove(other_op) frontier = list(frontier.union(other_op.next)) @@ -46,7 +47,9 @@ def circular_dep_after_merge(op: Op, other_op: Op): # The root node will be visited again if there is a circular dependency if n in root: return True - frontier.append(n) + if n not in visited: + frontier.append(n) + visited.add(n) frontier = frontier[1:]