colossalai
54 строки · 1.6 Кб
1import platform
2
3from ..cuda_extension import _CudaExtension
4from ..utils import append_nvcc_threads
5
6
7class CpuAdamX86Extension(_CudaExtension):
8def __init__(self):
9super().__init__(name="cpu_adam_x86")
10
11def is_hardware_available(self) -> bool:
12return platform.machine() == "x86_64" and super().is_hardware_available()
13
14def assert_hardware_compatible(self) -> None:
15arch = platform.machine()
16assert (
17arch == "x86_64"
18), f"[extension] The {self.name} kernel requires the CPU architecture to be x86_64 but got {arch}"
19super().assert_hardware_compatible()
20
21# necessary 4 functions
22def sources_files(self):
23ret = [
24self.csrc_abs_path("cuda/cpu_adam.cpp"),
25]
26return ret
27
28def include_dirs(self):
29return [self.csrc_abs_path("includes"), self.get_cuda_home_include()]
30
31def cxx_flags(self):
32extra_cxx_flags = [
33"-std=c++14",
34"-std=c++17",
35"-lcudart",
36"-lcublas",
37"-g",
38"-Wno-reorder",
39"-fopenmp",
40"-march=native",
41]
42return ["-O3"] + self.version_dependent_macros + extra_cxx_flags
43
44def nvcc_flags(self):
45extra_cuda_flags = [
46"-std=c++14",
47"-std=c++17",
48"-U__CUDA_NO_HALF_OPERATORS__",
49"-U__CUDA_NO_HALF_CONVERSIONS__",
50"-U__CUDA_NO_HALF2_OPERATORS__",
51"-DTHRUST_IGNORE_CUB_VERSION_CHECK",
52]
53ret = ["-O3", "--use_fast_math"] + self.version_dependent_macros + extra_cuda_flags
54return append_nvcc_threads(ret)
55