deepspeed

1
# Copyright (c) Microsoft Corporation.
2
# SPDX-License-Identifier: Apache-2.0
3

4
# DeepSpeed Team
5

6
import os
7
from .builder import TorchCPUOpBuilder
8

9

10
class CPUAdamBuilder(TorchCPUOpBuilder):
11
    BUILD_VAR = "DS_BUILD_CPU_ADAM"
12
    NAME = "cpu_adam"
13

14
    def __init__(self):
15
        super().__init__(name=self.NAME)
16

17
    def absolute_name(self):
18
        return f'deepspeed.ops.adam.{self.NAME}_op'
19

20
    def sources(self):
21
        if self.build_for_cpu:
22
            return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp']
23

24
        return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp', 'csrc/common/custom_cuda_kernel.cu']
25

26
    def libraries_args(self):
27
        args = super().libraries_args()
28
        if self.build_for_cpu:
29
            return args
30

31
        if not self.is_rocm_pytorch():
32
            args += ['curand']
33

34
        return args
35

36
    def include_paths(self):
37
        import torch
38
        if self.build_for_cpu:
39
            CUDA_INCLUDE = []
40
        elif not self.is_rocm_pytorch():
41
            CUDA_INCLUDE = [os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")]
42
        else:
43
            CUDA_INCLUDE = []
44
        return ['csrc/includes'] + CUDA_INCLUDE
45