Use driver API

This commit is contained in:
Chenggang Zhao
2025-08-19 11:29:15 +08:00
parent 51d1e9cdd3
commit ea9c5d9270
2 changed files with 4 additions and 4 deletions

View File

@@ -8,7 +8,7 @@
namespace deep_gemm {
#if CUDART_VERSION >= 12080 and not defined(DG_JIT_USE_DRIVER_API)
#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API)
// Use CUDA runtime API
using LibraryHandle = cudaLibrary_t;

View File

@@ -28,9 +28,9 @@ third_party_include_dirs = [
'third-party/cutlass/include/cutlass',
]
# Use driver API for older CUDA compatibility
if int(os.environ.get('DG_JIT_USE_DRIVER_API', '0')):
cxx_flags.append('-DDG_JIT_USE_DRIVER_API')
# Use runtime API
if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')):
cxx_flags.append('-DDG_JIT_USE_RUNTIME_API')
class CustomBuildPy(build_py):