| 39 | |
| 40 | |
| 41 | def get_extensions(): |
| 42 | this_dir = path.dirname(path.abspath(__file__)) |
| 43 | extensions_dir = path.join(this_dir, "detectron2", "layers", "csrc") |
| 44 | |
| 45 | main_source = path.join(extensions_dir, "vision.cpp") |
| 46 | sources = glob.glob(path.join(extensions_dir, "**", "*.cpp")) |
| 47 | |
| 48 | is_rocm_pytorch = False |
| 49 | if torch_ver >= [1, 5]: |
| 50 | from torch.utils.cpp_extension import ROCM_HOME |
| 51 | |
| 52 | is_rocm_pytorch = ( |
| 53 | True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False |
| 54 | ) |
| 55 | |
| 56 | if is_rocm_pytorch: |
| 57 | hipify_python.hipify( |
| 58 | project_directory=this_dir, |
| 59 | output_directory=this_dir, |
| 60 | includes="/detectron2/layers/csrc/*", |
| 61 | show_detailed=True, |
| 62 | is_pytorch_extension=True, |
| 63 | ) |
| 64 | |
| 65 | # Current version of hipify function in pytorch creates an intermediate directory |
| 66 | # named "hip" at the same level of the path hierarchy if a "cuda" directory exists, |
| 67 | # or modifying the hierarchy, if it doesn't. Once pytorch supports |
| 68 | # "same directory" hipification (PR pendeing), the source_cuda will be set |
| 69 | # similarly in both cuda and hip paths, and the explicit header file copy |
| 70 | # (below) will not be needed. |
| 71 | source_cuda = glob.glob(path.join(extensions_dir, "**", "hip", "*.hip")) + glob.glob( |
| 72 | path.join(extensions_dir, "hip", "*.hip") |
| 73 | ) |
| 74 | |
| 75 | shutil.copy( |
| 76 | "detectron2/layers/csrc/box_iou_rotated/box_iou_rotated_utils.h", |
| 77 | "detectron2/layers/csrc/box_iou_rotated/hip/box_iou_rotated_utils.h", |
| 78 | ) |
| 79 | shutil.copy( |
| 80 | "detectron2/layers/csrc/deformable/deform_conv.h", |
| 81 | "detectron2/layers/csrc/deformable/hip/deform_conv.h", |
| 82 | ) |
| 83 | |
| 84 | else: |
| 85 | source_cuda = glob.glob(path.join(extensions_dir, "**", "*.cu")) + glob.glob( |
| 86 | path.join(extensions_dir, "*.cu") |
| 87 | ) |
| 88 | |
| 89 | sources = [main_source] + sources |
| 90 | |
| 91 | extension = CppExtension |
| 92 | |
| 93 | extra_compile_args = {"cxx": []} |
| 94 | define_macros = [] |
| 95 | |
| 96 | if (torch.cuda.is_available() and ((CUDA_HOME is not None) or is_rocm_pytorch)) or os.getenv( |
| 97 | "FORCE_CUDA", "0" |
| 98 | ) == "1": |