xref: /aosp_15_r20/external/pytorch/tools/amd_build/build_amd.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1#!/usr/bin/env python3
2
3
4import argparse
5import os
6import sys
7
8
9sys.path.append(
10    os.path.realpath(
11        os.path.join(
12            __file__, os.path.pardir, os.path.pardir, os.path.pardir, "torch", "utils"
13        )
14    )
15)
16
17from hipify import hipify_python  # type: ignore[import]
18
19
20parser = argparse.ArgumentParser(
21    description="Top-level script for HIPifying, filling in most common parameters"
22)
23parser.add_argument(
24    "--out-of-place-only",
25    action="store_true",
26    help="Whether to only run hipify out-of-place on source files",
27)
28
29parser.add_argument(
30    "--project-directory",
31    type=str,
32    default="",
33    help="The root of the project.",
34    required=False,
35)
36
37parser.add_argument(
38    "--output-directory",
39    type=str,
40    default="",
41    help="The directory to store the hipified project",
42    required=False,
43)
44
45parser.add_argument(
46    "--extra-include-dir",
47    type=str,
48    default=[],
49    nargs="+",
50    help="The list of extra directories in caffe2 to hipify",
51    required=False,
52)
53
54args = parser.parse_args()
55
56amd_build_dir = os.path.dirname(os.path.realpath(__file__))
57proj_dir = os.path.join(os.path.dirname(os.path.dirname(amd_build_dir)))
58
59if args.project_directory:
60    proj_dir = args.project_directory
61
62out_dir = proj_dir
63if args.output_directory:
64    out_dir = args.output_directory
65
66includes = [
67    "caffe2/operators/*",
68    "caffe2/sgd/*",
69    "caffe2/image/*",
70    "caffe2/transforms/*",
71    "caffe2/video/*",
72    "caffe2/distributed/*",
73    "caffe2/queue/*",
74    "caffe2/contrib/aten/*",
75    "binaries/*",
76    "caffe2/**/*_test*",
77    "caffe2/core/*",
78    "caffe2/db/*",
79    "caffe2/utils/*",
80    "caffe2/contrib/gloo/*",
81    "caffe2/contrib/nccl/*",
82    "c10/cuda/*",
83    "c10/cuda/test/CMakeLists.txt",
84    "modules/*",
85    "third_party/nvfuser/*",
86    # PyTorch paths
87    # Keep this synchronized with is_pytorch_file in hipify_python.py
88    "aten/src/ATen/cuda/*",
89    "aten/src/ATen/native/cuda/*",
90    "aten/src/ATen/native/cudnn/*",
91    "aten/src/ATen/native/quantized/cudnn/*",
92    "aten/src/ATen/native/nested/cuda/*",
93    "aten/src/ATen/native/sparse/cuda/*",
94    "aten/src/ATen/native/quantized/cuda/*",
95    "aten/src/ATen/native/transformers/cuda/attention_backward.cu",
96    "aten/src/ATen/native/transformers/cuda/attention.cu",
97    "aten/src/ATen/native/transformers/cuda/sdp_utils.cpp",
98    "aten/src/ATen/native/transformers/cuda/sdp_utils.h",
99    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/debug_utils.h",
100    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/gemm_kernel_utils.h",
101    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/pytorch_utils.h",
102    "aten/src/ATen/native/transformers/cuda/flash_attn/flash_api.h",
103    "aten/src/THC/*",
104    "aten/src/ATen/test/*",
105    # CMakeLists.txt isn't processed by default, but there are a few
106    # we do want to handle, so explicitly specify them
107    "aten/src/THC/CMakeLists.txt",
108    "torch/*",
109    "tools/autograd/templates/python_variable_methods.cpp",
110]
111
112includes = [os.path.join(proj_dir, include) for include in includes]
113
114for new_dir in args.extra_include_dir:
115    abs_new_dir = os.path.join(proj_dir, new_dir)
116    if os.path.exists(abs_new_dir):
117        abs_new_dir = os.path.join(abs_new_dir, "**/*")
118        includes.append(abs_new_dir)
119
120ignores = [
121    "caffe2/operators/depthwise_3x3_conv_op_cudnn.cu",
122    "caffe2/operators/pool_op_cudnn.cu",
123    "*/hip/*",
124    # These files are compatible with both cuda and hip
125    "aten/src/ATen/core/*",
126    # Correct path to generate HIPConfig.h:
127    #   CUDAConfig.h.in -> (amd_build) HIPConfig.h.in -> (cmake) HIPConfig.h
128    "aten/src/ATen/cuda/CUDAConfig.h",
129    "third_party/nvfuser/csrc/codegen.cpp",
130    "third_party/nvfuser/runtime/block_reduction.cu",
131    "third_party/nvfuser/runtime/block_sync_atomic.cu",
132    "third_party/nvfuser/runtime/block_sync_default_rocm.cu",
133    "third_party/nvfuser/runtime/broadcast.cu",
134    "third_party/nvfuser/runtime/grid_reduction.cu",
135    "third_party/nvfuser/runtime/helpers.cu",
136    "torch/csrc/jit/codegen/fuser/cuda/resource_strings.h",
137    "torch/csrc/jit/tensorexpr/ir_printer.cpp",
138    # generated files we shouldn't frob
139    "torch/lib/tmp_install/*",
140    "torch/include/*",
141]
142
143ignores = [os.path.join(proj_dir, ignore) for ignore in ignores]
144
145
146# Check if the compiler is hip-clang.
147def is_hip_clang() -> bool:
148    try:
149        hip_path = os.getenv("HIP_PATH", "/opt/rocm/hip")
150        with open(hip_path + "/lib/.hipInfo") as f:
151            return "HIP_COMPILER=clang" in f.read()
152    except OSError:
153        return False
154
155
156# TODO Remove once the following submodules are updated
157hip_platform_files = [
158    "third_party/fbgemm/fbgemm_gpu/CMakeLists.txt",
159    "third_party/fbgemm/fbgemm_gpu/cmake/Hip.cmake",
160    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_dense_host.cpp",
161    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_split_host_template.cpp",
162    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_split_template.cu",
163    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_forward_quantized_split_lookup.cu",
164    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/utils/cuda_prelude.cuh",
165    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/utils/stochastic_rounding.cuh",
166    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/utils/vec4.cuh",
167    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/utils/weight_row.cuh",
168    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.cuh",
169    "third_party/fbgemm/fbgemm_gpu/src/jagged_tensor_ops.cu",
170    "third_party/fbgemm/fbgemm_gpu/src/quantize_ops.cu",
171    "third_party/fbgemm/fbgemm_gpu/src/sparse_ops.cu",
172    "third_party/fbgemm/fbgemm_gpu/src/split_embeddings_cache_cuda.cu",
173    "third_party/fbgemm/fbgemm_gpu/src/topology_utils.cpp",
174    "third_party/fbgemm/src/EmbeddingSpMDM.cc",
175    "third_party/gloo/cmake/Dependencies.cmake",
176    "third_party/gloo/gloo/cuda.cu",
177    "third_party/kineto/libkineto/CMakeLists.txt",
178    "third_party/nvfuser/CMakeLists.txt",
179    "third_party/tensorpipe/cmake/Hip.cmake",
180]
181
182
183def remove_hcc(line: str) -> str:
184    line = line.replace("HIP_PLATFORM_HCC", "HIP_PLATFORM_AMD")
185    line = line.replace("HIP_HCC_FLAGS", "HIP_CLANG_FLAGS")
186    return line
187
188
189for hip_platform_file in hip_platform_files:
190    do_write = False
191    if os.path.exists(hip_platform_file):
192        with open(hip_platform_file) as sources:
193            lines = sources.readlines()
194        newlines = [remove_hcc(line) for line in lines]
195        if lines == newlines:
196            print(f"{hip_platform_file} skipped")
197        else:
198            with open(hip_platform_file, "w") as sources:
199                for line in newlines:
200                    sources.write(line)
201            print(f"{hip_platform_file} updated")
202
203hipify_python.hipify(
204    project_directory=proj_dir,
205    output_directory=out_dir,
206    includes=includes,
207    ignores=ignores,
208    extra_files=[
209        "torch/_inductor/codegen/cpp_wrapper_cpu.py",
210        "torch/_inductor/codegen/cpp_wrapper_cuda.py",
211        "torch/_inductor/codegen/wrapper.py",
212    ],
213    out_of_place_only=args.out_of_place_only,
214    hip_clang_launch=is_hip_clang(),
215)
216