github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-12-01 00:24:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

[setup] support pre-build and jit-build of cuda kernels (#2374)

Browse Source 
* [setup] support pre-build and jit-build of cuda kernels

* polish code

* polish code

* polish code

* polish code

* polish code

* polish code
This commit is contained in:
Frank Lee
2023-01-06 20:50:26 +08:00
committed by GitHub
parent 12c8bf38d7
commit 40d376c566
36 changed files with 414 additions and 390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
op_builder/cpu_adam.py
View File

@@ -6,24 +6,22 @@ from .utils import append_nvcc_threads
class CPUAdamBuilder(Builder):
NAME = "cpu_adam"
BASE_DIR = "cuda_native"
PREBUILT_IMPORT_PATH = "colossalai._C.cpu_adam"
def __init__(self):
self.name = CPUAdamBuilder.NAME
super().__init__()
super().__init__(name=CPUAdamBuilder.NAME, prebuilt_import_path=CPUAdamBuilder.PREBUILT_IMPORT_PATH)
self.version_dependent_macros = ['-DVERSION_GE_1_1', '-DVERSION_GE_1_3', '-DVERSION_GE_1_5']
# necessary 4 functions
def sources_files(self):
ret = [
os.path.join(CPUAdamBuilder.BASE_DIR, "csrc/cpu_adam.cpp"),
self.csrc_abs_path('cpu_adam.cpp'),
]
return [self.colossalai_src_path(path) for path in ret]
return ret
def include_dirs(self):
return [
self.colossalai_src_path(os.path.join(CPUAdamBuilder.BASE_DIR, "includes")),
self.csrc_abs_path("includes"),
self.get_cuda_home_include()
]
@@ -36,7 +34,5 @@ class CPUAdamBuilder(Builder):
'-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__',
'-U__CUDA_NO_HALF2_OPERATORS__', '-DTHRUST_IGNORE_CUB_VERSION_CHECK'
]
return append_nvcc_threads(['-O3', '--use_fast_math'] + self.version_dependent_macros + extra_cuda_flags)
# necessary 4 functions
ret = ['-O3', '--use_fast_math'] + self.version_dependent_macros + extra_cuda_flags
return append_nvcc_threads(ret)