Class AsyncIOBuilder

op_builder/async_io.py:13–113

class AsyncIOBuilder(TorchCPUOpBuilder):
    BUILD_VAR = "DS_BUILD_AIO"
    NAME = "async_io"

    def __init__(self):
        super().__init__(name=self.NAME)

    def absolute_name(self):
        return f'deepspeed.ops.aio.{self.NAME}_op'

    def lib_sources(self):
        src_list = [
            'csrc/aio/py_lib/deepspeed_py_io_handle.cpp', 'csrc/aio/py_lib/deepspeed_py_aio.cpp',
            'csrc/aio/py_lib/deepspeed_py_aio_handle.cpp', 'csrc/aio/py_lib/deepspeed_aio_thread.cpp',
            'csrc/aio/common/deepspeed_aio_utils.cpp', 'csrc/aio/common/deepspeed_aio_common.cpp',
            'csrc/aio/common/deepspeed_aio_types.cpp', 'csrc/aio/py_lib/deepspeed_cpu_op.cpp',
            'csrc/aio/py_lib/deepspeed_aio_op_desc.cpp', 'csrc/aio/py_lib/deepspeed_py_copy.cpp',
            'csrc/aio/py_lib/deepspeed_pin_tensor.cpp'
        ]
        return src_list

    def sources(self):
        return self.lib_sources() + ['csrc/aio/py_lib/py_ds_aio.cpp']

    def include_paths(self):
        import torch
        if self.build_for_cpu:
            CUDA_INCLUDE = []
        elif not self.is_rocm_pytorch():
            CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
        else:
            CUDA_INCLUDE = [
                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include"),
                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "rocrand"),
                os.path.join(torch.utils.cpp_extension.ROCM_HOME, "include", "hiprand"),
            ]
        return ['csrc/aio/py_lib', 'csrc/aio/common'] + CUDA_INCLUDE

    def cxx_args(self):
        # -O0 for improved debugging, since performance is bound by I/O
        args = super().cxx_args()
        import torch
        TORCH_MAJOR, TORCH_MINOR = map(int, torch.__version__.split('.')[0:2])
        if not (TORCH_MAJOR >= 2 and TORCH_MINOR >= 1):
            args.remove('-std=c++17')
            args.append('-std=c++14')
        args += ['-Wall', '-O0', '-shared', '-fPIC', '-Wno-reorder']
        return args

    def extra_ldflags(self):
        if self.build_for_cpu:
            return ['-fopenmp']

        import torch.utils.cpp_extension
        CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
        if CUDA_HOME is None:
            ldflags = ['-laio']  # the ROCM case
        else: