Code refactor for mscclang

.github/workflows/codeql.yml

@@ -19,12 +19,12 @@ jobs:
 
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v2
+      uses: actions/checkout@v4
 
     - name: Initialize CodeQL
-      uses: github/codeql-action/init@v1
+      uses: github/codeql-action/init@v3
       with:
         languages: python
 
     - name: Perform CodeQL Analysis
-      uses: github/codeql-action/analyze@v1
+      uses: github/codeql-action/analyze@v3

.github/workflows/tests.yaml

@@ -11,14 +11,14 @@ 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 }}
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install msccl and dependencies
@@ -28,3 +28,44 @@ jobs:
     - name: Run tests and check at least 90% coverage
       run: |
         pytest
+
+  compare_outputs:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.8', '3.9', '3.10']
+    name: Compare outputs with Python ${{ matrix.python-version }}
+
+    steps:
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Checkout current branch
+      uses: actions/checkout@v4
+    - name: Install msccl and dependencies
+      run: |
+        pip install --upgrade pip
+        pip install -r requirements.txt
+    - name: Copy test script/config to temp directory
+      run: |
+        cp tests/generate_example_results.py $RUNNER_TEMP/
+        cp tests/configs/example-config.json $RUNNER_TEMP/
+    - name: generate outputs
+      run: |
+        python $RUNNER_TEMP/generate_example_results.py examples/mscclang/ $RUNNER_TEMP/example-config.json $RUNNER_TEMP/tests/pr-outputs/
+    - name: Checkout main branch
+      uses: actions/checkout@v4
+      with:
+        ref: main
+    - name: Install msccl and dependencies
+      run: |
+        pip install --upgrade pip
+        pip install -r requirements.txt
+    - name: generate outputs
+      run: |
+        python $RUNNER_TEMP/generate_example_results.py examples/mscclang/ $RUNNER_TEMP/example-config.json $RUNNER_TEMP/tests/main-outputs/
+    - name: Compare outputs
+      run: |
+        diff -rw $RUNNER_TEMP/tests/main-outputs/ $RUNNER_TEMP/tests/pr-outputs/
+

.gitignore

@@ -131,3 +131,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# vscode
+.vscode/

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