diff --git a/.github/workflows/mscclpp-lang.yml b/.github/workflows/mscclpp-lang.yml index 250bc1164..94b4ad0ab 100644 --- a/.github/workflows/mscclpp-lang.yml +++ b/.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' diff --git a/python/mscclpp/language/dag/optimizer.py b/python/mscclpp/language/dag/optimizer.py index 44f7f664b..62fc0f5e8 100644 --- a/python/mscclpp/language/dag/optimizer.py +++ b/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) diff --git a/python/test/configs/msccllang_test_config.json b/python/test/configs/msccllang_test_config.json new file mode 100644 index 000000000..6aa0ca9b3 --- /dev/null +++ b/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"] + } +] diff --git a/python/test/generate_mscclpp_lang_test_result.py b/python/test/generate_mscclpp_lang_test_result.py new file mode 100644 index 000000000..036c71c5f --- /dev/null +++ b/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)