deepspeed

1
# Copyright (c) 2023 Habana Labs, Ltd. an Intel Company
2
# Copyright (c) Microsoft Corporation.
3
# SPDX-License-Identifier: Apache-2.0
4

5
# DeepSpeed Team
6

7
from .builder import CPUOpBuilder
8

9

10
class CPUAdamBuilder(CPUOpBuilder):
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
        return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp']
22

23
    def cxx_args(self):
24
        args = super().cxx_args()
25
        args += ['-DENABLE_BFLOAT16']
26
        return args
27

28
    def libraries_args(self):
29
        args = super().libraries_args()
30
        return args
31

32
    def include_paths(self):
33
        return ['csrc/includes']
34