update

.github/workflows/mscclpp-lang.yml

@@ -23,31 +23,13 @@ jobs:
       run: |
         CMAKE_ARGS="-DMSCCLPP_BYPASS_GPU_CHECK=ON -DMSCCLPP_USE_CUDA=ON" pip3 install .
 
-  # 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-tools and dependencies
-  #     run: |
-  #       pip install --upgrade pip
-  #       pip install -r requirements.txt
-  #   - name: Copy test script/config to temp directory
-  #     run: |
-  #       cp tests/generate_test_results.py $RUNNER_TEMP/
-  #       cp tests/configs/test-config.json $RUNNER_TEMP/
-  #   - name: generate outputs
-  #     run: |
-  #       python $RUNNER_TEMP/generate_test_results.py examples/mscclang/ $RUNNER_TEMP/test-config.json $RUNNER_TEMP/tests/pr-outputs/
+    - name: Copy test script/config to temp directory
+      run: |
+        cp python/test/generate_mscclpp_lang_test_result.py $RUNNER_TEMP/
+        cp python/test/configs/mscclpp_lang_test_config.json $RUNNER_TEMP/
+    - name: generate outputs
+      run: |
+        python3 $RUNNER_TEMP/generate_test_results.py examples/mscclang/ $RUNNER_TEMP/test-config.json $RUNNER_TEMP/tests/pr-outputs/
   #   - name: Checkout main branch
   #     uses: actions/checkout@v4
   #     if: github.event_name == 'pull_request' || github.event_name == 'push'

python/mscclpp/language/dag/optimizer.py

@@ -124,8 +124,7 @@ def try_fuse_with_put(self, op: Op, next_op: Op, tb: Threadblock, queue: list) -
         :return: True if operations are merged, False otherwise.
         """
         if (
-            next_op.inst == Instruction.put
-            or next_op.inst == Instruction.put_packet
+            (next_op.inst == Instruction.put or next_op.inst == Instruction.put_packet)
             and same_tb(op, next_op)
             and same_count(op, next_op)
             and buf_dst_src_match(op, next_op)

python/test/configs/msccllang_test_config.json

@@ -0,0 +1,30 @@
+[
+    {
+        "filename": "allreduce_allpairs_packet.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "allreduce_allpairs_get.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "allreduce_allpairs.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "allreduce_ring.py",
+        "args": ["8", "8"]
+    },
+    {
+        "filename": "send_recv_packet.py",
+        "args": ["2"]
+    },
+    {
+        "filename": "send_recv_proxy.py",
+        "args": ["2"]
+    },
+    {
+        "filename": "allreduce_nvls.py",
+        "args": ["8", "2"]
+    }
+]

python/test/generate_mscclpp_lang_test_result.py

@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import argparse
+import json
+from pathlib import Path
+import subprocess
+
+
+def run_examples(input_folder, configs, output_folder):
+    for config in configs:
+        file_name = config["filename"]
+        args = config["args"]
+
+        input_file_path = Path(input_folder) / file_name
+        # Strip the ".py" from the filename and add ".output"
+        base_file_name = file_name[:-3] if file_name.endswith(".py") else file_name
+        base_file_name = base_file_name.replace("/", "_")
+        output_file_path = Path(output_folder) / f"{base_file_name}.output"
+
+        # Construct the command to run the Python script
+        command = ["python3", str(input_file_path)] + args
+
+        # Run the command and capture output
+        with open(output_file_path, "w") as output_file:
+            result = subprocess.run(command, stdout=output_file, stderr=subprocess.STDOUT, text=True)
+
+        # Optional: Check the return code to handle errors
+        if result.returncode != 0:
+            print(f"Error running {file_name}. See {output_file_path} for details.")
+
+
+def main(input_folder, config_path, output_folder):
+    with open(config_path, "r") as f:
+        config = json.load(f)
+
+    Path(output_folder).mkdir(parents=True, exist_ok=True)
+    run_examples(input_folder, config, output_folder)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Process files according to a configuration and save the results.")
+    parser.add_argument("input_folder", type=str, help="Path to the folder containing the input files.")
+    parser.add_argument("config", type=str, help="Path to the configuration file.")
+    parser.add_argument("output_folder", type=str, help="Path to the folder where the processed files will be saved.")
+    args = parser.parse_args()
+    main(args.input_folder, args.config, args.output_folder)