| 1 | #!/usr/bin/env python3 |
| 2 | # SPDX-License-Identifier: Apache-2.0 |
| 3 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | |
| 5 | # |
| 6 | # A command line tool for running pytorch's hipify preprocessor on CUDA |
| 7 | # source files. |
| 8 | # |
| 9 | # See https://github.com/ROCm/hipify_torch |
| 10 | # and <torch install dir>/utils/hipify/hipify_python.py |
| 11 | # |
| 12 | |
| 13 | import argparse |
| 14 | import os |
| 15 | import shutil |
| 16 | |
| 17 | from torch.utils.hipify.hipify_python import get_hip_file_path, hipify |
| 18 | |
| 19 | |
| 20 | def _expected_hip_build_path(source_abs: str, output_directory: str) -> str: |
| 21 | """Match torch.utils.hipify.hipify_python.preprocessor fout_path naming.""" |
| 22 | rel = os.path.relpath(source_abs, output_directory) |
| 23 | return os.path.abspath( |
| 24 | os.path.join( |
| 25 | output_directory, get_hip_file_path(rel, is_pytorch_extension=True) |
| 26 | ) |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | if __name__ == "__main__": |
| 31 | parser = argparse.ArgumentParser() |
| 32 | |
| 33 | # Project directory where all the source + include files live. |
| 34 | parser.add_argument( |
| 35 | "-p", |
| 36 | "--project_dir", |
| 37 | help="The project directory.", |
| 38 | ) |
| 39 | |
| 40 | # Directory where hipified files are written. |
| 41 | parser.add_argument( |
| 42 | "-o", |
| 43 | "--output_dir", |
| 44 | help="The output directory.", |
| 45 | ) |
| 46 | |
| 47 | # Source files to convert. |
| 48 | parser.add_argument( |
| 49 | "sources", help="Source files to hipify.", nargs="*", default=[] |
| 50 | ) |
| 51 | |
| 52 | args = parser.parse_args() |
| 53 | |
| 54 | # Limit include scope to project_dir only |
| 55 | includes = [os.path.join(args.project_dir, "*")] |
| 56 | |
| 57 | # Get absolute path for all source files. |
| 58 | extra_files = [os.path.abspath(s) for s in args.sources] |
| 59 | |
| 60 | # Copy sources from project directory to output directory. |
| 61 | # The directory might already exist to hold object files so we ignore that. |
| 62 | shutil.copytree(args.project_dir, args.output_dir, dirs_exist_ok=True) |
| 63 | |
| 64 | hipify_result = hipify( |
| 65 | project_directory=args.project_dir, |
| 66 | output_directory=args.output_dir, |
| 67 | # Hipify resolves quoted includes next to the including file first; vLLM |
| 68 | # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h" |
| 69 | # from quantization/w8a8/fp8/*.cu). Without an include root here, those |
| 70 | # headers are never found and are not hipified or rewritten in dependents. |
| 71 | header_include_dirs=["."], |
| 72 | includes=includes, |
| 73 | extra_files=extra_files, |
| 74 | show_detailed=True, |
| 75 | is_pytorch_extension=True, |
| 76 | hipify_extra_files_only=True, |
| 77 | ) |
| 78 | |
| 79 | hipified_sources = [] |
| 80 | for source in args.sources: |
| 81 | s_abs = os.path.abspath(source) |
| 82 | if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: |
| 83 | path = hipify_result[s_abs].hipified_path |
| 84 | # PyTorch skips writing when is_pytorch_extension and text unchanged; |
| 85 | # hipified_path then stays *.cu. CMake expects *.hip under output_dir. |
| 86 | if s_abs.endswith(".cu") and path.endswith(".cu"): |
| 87 | dest = _expected_hip_build_path(s_abs, args.output_dir) |
| 88 | if os.path.normpath(path) != os.path.normpath(dest): |
| 89 | os.makedirs(os.path.dirname(dest), exist_ok=True) |
| 90 | shutil.copy2(path, dest) |
| 91 | hipified_s_abs = dest |
| 92 | else: |
| 93 | hipified_s_abs = path |
| 94 | else: |
| 95 | hipified_s_abs = s_abs |
| 96 | hipified_sources.append(hipified_s_abs) |
| 97 | |
| 98 | assert len(hipified_sources) == len(args.sources) |
| 99 | |
| 100 | # Print hipified source files. |
| 101 | print("\n".join(hipified_sources)) |