pytorch

1
#!/usr/bin/env python3
2

3

4
import argparse
5
import os
6
import sys
7

8
sys.path.append(
9
    os.path.realpath(
10
        os.path.join(
11
            __file__, os.path.pardir, os.path.pardir, os.path.pardir, "torch", "utils"
12
        )
13
    )
14
)
15

16
from hipify import hipify_python  # type: ignore[import]
17

18
parser = argparse.ArgumentParser(
19
    description="Top-level script for HIPifying, filling in most common parameters"
20
)
21
parser.add_argument(
22
    "--out-of-place-only",
23
    action="store_true",
24
    help="Whether to only run hipify out-of-place on source files",
25
)
26

27
parser.add_argument(
28
    "--project-directory",
29
    type=str,
30
    default="",
31
    help="The root of the project.",
32
    required=False,
33
)
34

35
parser.add_argument(
36
    "--output-directory",
37
    type=str,
38
    default="",
39
    help="The directory to store the hipified project",
40
    required=False,
41
)
42

43
parser.add_argument(
44
    "--extra-include-dir",
45
    type=str,
46
    default=[],
47
    nargs="+",
48
    help="The list of extra directories in caffe2 to hipify",
49
    required=False,
50
)
51

52
args = parser.parse_args()
53

54
amd_build_dir = os.path.dirname(os.path.realpath(__file__))
55
proj_dir = os.path.join(os.path.dirname(os.path.dirname(amd_build_dir)))
56

57
if args.project_directory:
58
    proj_dir = args.project_directory
59

60
out_dir = proj_dir
61
if args.output_directory:
62
    out_dir = args.output_directory
63

64
includes = [
65
    "caffe2/operators/*",
66
    "caffe2/sgd/*",
67
    "caffe2/image/*",
68
    "caffe2/transforms/*",
69
    "caffe2/video/*",
70
    "caffe2/distributed/*",
71
    "caffe2/queue/*",
72
    "caffe2/contrib/aten/*",
73
    "binaries/*",
74
    "caffe2/**/*_test*",
75
    "caffe2/core/*",
76
    "caffe2/db/*",
77
    "caffe2/utils/*",
78
    "caffe2/contrib/gloo/*",
79
    "caffe2/contrib/nccl/*",
80
    "c10/cuda/*",
81
    "c10/cuda/test/CMakeLists.txt",
82
    "modules/*",
83
    "third_party/nvfuser/*",
84
    # PyTorch paths
85
    # Keep this synchronized with is_pytorch_file in hipify_python.py
86
    "aten/src/ATen/cuda/*",
87
    "aten/src/ATen/native/cuda/*",
88
    "aten/src/ATen/native/cudnn/*",
89
    "aten/src/ATen/native/quantized/cudnn/*",
90
    "aten/src/ATen/native/nested/cuda/*",
91
    "aten/src/ATen/native/sparse/cuda/*",
92
    "aten/src/ATen/native/quantized/cuda/*",
93
    "aten/src/ATen/native/transformers/cuda/attention_backward.cu",
94
    "aten/src/ATen/native/transformers/cuda/attention.cu",
95
    "aten/src/ATen/native/transformers/cuda/sdp_utils.cpp",
96
    "aten/src/ATen/native/transformers/cuda/sdp_utils.h",
97
    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/debug_utils.h",
98
    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/gemm_kernel_utils.h",
99
    "aten/src/ATen/native/transformers/cuda/mem_eff_attention/pytorch_utils.h",
100
    "aten/src/ATen/native/transformers/cuda/flash_attn/flash_api.h",
101
    "aten/src/THC/*",
102
    "aten/src/ATen/test/*",
103
    # CMakeLists.txt isn't processed by default, but there are a few
104
    # we do want to handle, so explicitly specify them
105
    "aten/src/THC/CMakeLists.txt",
106
    "torch/*",
107
    "tools/autograd/templates/python_variable_methods.cpp",
108
]
109

110
includes = [os.path.join(proj_dir, include) for include in includes]
111

112
for new_dir in args.extra_include_dir:
113
    abs_new_dir = os.path.join(proj_dir, new_dir)
114
    if os.path.exists(abs_new_dir):
115
        abs_new_dir = os.path.join(abs_new_dir, "**/*")
116
        includes.append(abs_new_dir)
117

118
ignores = [
119
    "caffe2/operators/depthwise_3x3_conv_op_cudnn.cu",
120
    "caffe2/operators/pool_op_cudnn.cu",
121
    "*/hip/*",
122
    # These files are compatible with both cuda and hip
123
    "aten/src/ATen/core/*",
124
    # Correct path to generate HIPConfig.h:
125
    #   CUDAConfig.h.in -> (amd_build) HIPConfig.h.in -> (cmake) HIPConfig.h
126
    "aten/src/ATen/cuda/CUDAConfig.h",
127
    "third_party/nvfuser/csrc/codegen.cpp",
128
    "third_party/nvfuser/runtime/block_reduction.cu",
129
    "third_party/nvfuser/runtime/block_sync_atomic.cu",
130
    "third_party/nvfuser/runtime/block_sync_default_rocm.cu",
131
    "third_party/nvfuser/runtime/broadcast.cu",
132
    "third_party/nvfuser/runtime/grid_reduction.cu",
133
    "third_party/nvfuser/runtime/helpers.cu",
134
    "torch/csrc/jit/codegen/fuser/cuda/resource_strings.h",
135
    "torch/csrc/jit/tensorexpr/ir_printer.cpp",
136
    # generated files we shouldn't frob
137
    "torch/lib/tmp_install/*",
138
    "torch/include/*",
139
]
140

141
ignores = [os.path.join(proj_dir, ignore) for ignore in ignores]
142

143

144
# Check if the compiler is hip-clang.
145
def is_hip_clang() -> bool:
146
    try:
147
        hip_path = os.getenv("HIP_PATH", "/opt/rocm/hip")
148
        with open(hip_path + "/lib/.hipInfo") as f:
149
            return "HIP_COMPILER=clang" in f.read()
150
    except OSError:
151
        return False
152

153

154
# TODO Remove once the following submodules are updated
155
hip_platform_files = [
156
    "third_party/fbgemm/fbgemm_gpu/CMakeLists.txt",
157
    "third_party/fbgemm/fbgemm_gpu/cmake/Hip.cmake",
158
    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_dense_host.cpp",
159
    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_split_host_template.cpp",
160
    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_backward_split_template.cu",
161
    "third_party/fbgemm/fbgemm_gpu/codegen/embedding_forward_quantized_split_lookup.cu",
162
    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/fbgemm_cuda_utils.cuh",
163
    "third_party/fbgemm/fbgemm_gpu/include/fbgemm_gpu/sparse_ops.cuh",
164
    "third_party/fbgemm/fbgemm_gpu/src/jagged_tensor_ops.cu",
165
    "third_party/fbgemm/fbgemm_gpu/src/quantize_ops.cu",
166
    "third_party/fbgemm/fbgemm_gpu/src/sparse_ops.cu",
167
    "third_party/fbgemm/fbgemm_gpu/src/split_embeddings_cache_cuda.cu",
168
    "third_party/fbgemm/fbgemm_gpu/src/topology_utils.cpp",
169
    "third_party/fbgemm/src/EmbeddingSpMDM.cc",
170
    "third_party/gloo/cmake/Dependencies.cmake",
171
    "third_party/gloo/gloo/cuda.cu",
172
    "third_party/kineto/libkineto/CMakeLists.txt",
173
    "third_party/nvfuser/CMakeLists.txt",
174
    "third_party/tensorpipe/cmake/Hip.cmake",
175
]
176

177

178
def remove_hcc(line: str) -> str:
179
    line = line.replace("HIP_PLATFORM_HCC", "HIP_PLATFORM_AMD")
180
    line = line.replace("HIP_HCC_FLAGS", "HIP_CLANG_FLAGS")
181
    return line
182

183

184
for hip_platform_file in hip_platform_files:
185
    do_write = False
186
    if os.path.exists(hip_platform_file):
187
        with open(hip_platform_file) as sources:
188
            lines = sources.readlines()
189
        newlines = [remove_hcc(line) for line in lines]
190
        if lines == newlines:
191
            print(f"{hip_platform_file} skipped")
192
        else:
193
            with open(hip_platform_file, "w") as sources:
194
                for line in newlines:
195
                    sources.write(line)
196
            print(f"{hip_platform_file} updated")
197

198
hipify_python.hipify(
199
    project_directory=proj_dir,
200
    output_directory=out_dir,
201
    includes=includes,
202
    ignores=ignores,
203
    extra_files=[
204
        "torch/_inductor/codegen/cpp_wrapper_cpu.py",
205
        "torch/_inductor/codegen/cpp_wrapper_cuda.py",
206
        "torch/_inductor/codegen/wrapper.py",
207
    ],
208
    out_of_place_only=args.out_of_place_only,
209
    hip_clang_launch=is_hip_clang(),
210
)
211