colossalai

1
import platform
2

3
from ..cpp_extension import _CppExtension
4

5

6
class CpuAdamArmExtension(_CppExtension):
7
    def __init__(self):
8
        super().__init__(name="cpu_adam_arm")
9

10
    def is_hardware_available(self) -> bool:
11
        # only arm allowed
12
        return platform.machine() == "aarch64"
13

14
    def assert_hardware_compatible(self) -> None:
15
        arch = platform.machine()
16
        assert (
17
            arch == "aarch64"
18
        ), f"[extension] The {self.name} kernel requires the CPU architecture to be aarch64 but got {arch}"
19

20
    # necessary 4 functions
21
    def sources_files(self):
22
        ret = [
23
            self.csrc_abs_path("arm/cpu_adam_arm.cpp"),
24
        ]
25
        return ret
26

27
    def include_dirs(self):
28
        return []
29

30
    def cxx_flags(self):
31
        extra_cxx_flags = [
32
            "-std=c++14",
33
            "-std=c++17",
34
            "-g",
35
            "-Wno-reorder",
36
            "-fopenmp",
37
        ]
38
        return ["-O3"] + self.version_dependent_macros + extra_cxx_flags
39

40
    def nvcc_flags(self):
41
        return []
42