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