WIP

.github/workflows/tests.yaml

@@ -11,7 +11,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: [3.8, 3.9, 3.10]
 
     name: Test with Python ${{ matrix.python-version }}
 

msccl/autosynth/ndv4_plans.py

@@ -7,7 +7,7 @@
 from msccl.programs.alltoall_a100_yifan import alltoall_hierarchical
 from msccl.programs.alltoall_a100_8kp1 import alltoall_three_step
 from msccl.topologies import fully_connected
-from msccl.language.ir import ThreadblockPolicy
+from msccl.language.types import ThreadblockPolicy
 
 def register_ndv4_plans():
 
@@ -47,4 +47,4 @@ def ndv4_alltoall_three_step(prog, nodes):
     def ndv4_alltoall_hierarchical_config2(prog, nodes):
         alltoall_hierarchical(num_nodes=nodes, gpus_per_node=8)
 
-        
+

msccl/autosynth/registry.py

@@ -9,7 +9,7 @@
 import humanfriendly
 
 from msccl.language import MSCCLProgram, ir_to_xml
-from msccl.language.ir import ThreadblockPolicy
+from msccl.language.types import ThreadblockPolicy
 import msccl.language.collectives as lang_collectives
 from msccl.topologies import distributed_fully_connected
 
@@ -62,7 +62,7 @@ def wrapped(machines):
     return decorator
 
 
-def register_msccl_program(local_topology, collective, machine_type, machines=lambda x: True, sizes=None, protocol='Simple', 
+def register_msccl_program(local_topology, collective, machine_type, machines=lambda x: True, sizes=None, protocol='Simple',
     chunk_factor=1, priority=0, collective_obj=None, instances=1, inplace=False, threadblock_policy=ThreadblockPolicy.auto,
     interleaved_replication=True, dependence_nop=False):
     def decorator(fun):
@@ -81,7 +81,7 @@ def wrapped(machines):
                     co = lang_collectives.ReduceScatter(topology.num_nodes(), chunk_factor, inplace)
                 else:
                     raise RuntimeError(f'No collective_obj in msccl.language.collectives known for "{collective}"')
-            prog = MSCCLProgram(name, topology, co, instances, protocol, threadblock_policy=threadblock_policy, 
+            prog = MSCCLProgram(name, topology, co, instances, protocol, threadblock_policy=threadblock_policy,
                 interleaved_replication=interleaved_replication, dependence_nop=dependence_nop)
             with prog:
                 fun(prog, machines)
@@ -96,4 +96,4 @@ def wrapped(machines):
                              machine_type, machines, sizes, protocol, priority)
         # Return the original function to not break other usage
         return fun
-    return decorator
+    return decorator

msccl/language/__init__.py

@@ -6,7 +6,7 @@
 from msccl.language.tb_assignment import *
 from msccl.language.chunk import *
 from msccl.language.buffer import *
-from msccl.language.rank_dag import *
+from msccl.language.instruction_dag import *
 import msccl.language.msccl as msccl_lang
 import msccl.language.mscclpp as mscclpp
 from msccl.language.mscclpp import *