diff --git a/scripts/generate_public_headers.py b/scripts/generate_public_headers.py index 5daa57e..516d4aa 100644 --- a/scripts/generate_public_headers.py +++ b/scripts/generate_public_headers.py @@ -15,6 +15,7 @@ ("nvidia", "data_type_.h", "native/cuda/nvidia/data_type_.h"), ("nvidia", "device_.h", "native/cuda/nvidia/device_.h"), ("nvidia", "runtime_.h", "native/cuda/nvidia/runtime_.h"), + ("nvidia", "driver_.h", "native/cuda/nvidia/driver_.h"), ), "iluvatar": ( ("iluvatar", "data_type_.h", "native/cuda/iluvatar/data_type_.h"), @@ -116,7 +117,7 @@ def _rewrite_detail_include(match): _DETAIL_INCLUDE_PATTERN = re.compile( - r'#include "((?:common|native)/[^"]+|data_type\.h|device\.h|dispatcher\.h|hash\.h|runtime\.h|tensor_view\.h)"' + r'#include "((?:common|native)/[^"]+|data_type\.h|device\.h|dispatcher\.h|driver\.h|hash\.h|runtime\.h|tensor_view\.h)"' ) @@ -148,6 +149,7 @@ def _write_detail_headers(include_root, source_root, devices): "data_type.h", "device.h", "dispatcher.h", + "driver.h", "hash.h", "runtime.h", "tensor_view.h", @@ -200,6 +202,7 @@ def _write_generated_header(include_root, source_root, devices): "#include ", f"#include {_detail_include('data_type.h')}", f"#include {_detail_include('device.h')}", + f"#include {_detail_include('driver.h')}", f"#include {_detail_include('hash.h')}", f"#include {_detail_include('runtime.h')}", f"#include {_detail_include('tensor_view.h')}", @@ -211,6 +214,10 @@ def _write_generated_header(include_root, source_root, devices): for device in devices: includes.append(f"#include ") + for device in devices: + if any(h == "driver_.h" for _, h, _ in _DEVICE_HEADERS[device]): + includes.append(f"#include ") + runtime_declarations = "\n\n".join( f"{function.signature()};" for function in public_runtime_functions ) diff --git a/src/driver.h b/src/driver.h new file mode 100644 index 0000000..9f32847 --- /dev/null +++ b/src/driver.h @@ -0,0 +1,34 @@ +#ifndef INFINI_RT_DRIVER_H_ +#define INFINI_RT_DRIVER_H_ + +#include + +#include "device.h" + +namespace infini::rt::driver { + +template +struct Driver; + +template +struct DriverBase { + static constexpr bool Validate() { + static_assert( + std::is_same_v, + Device::Type>, + "`Driver` must define `static constexpr Device::Type kDeviceType`."); + static_assert(sizeof(typename Derived::Result) > 0, + "`Driver` must define a `Result` type alias."); + static_assert(std::is_same_v, + typename Derived::Result>, + "`Driver` must define `static constexpr Result kSuccess`."); + return true; + } +}; + +template +using DeviceDriver = DriverBase; + +} // namespace infini::rt::driver + +#endif diff --git a/src/native/cuda/driver_.h b/src/native/cuda/driver_.h new file mode 100644 index 0000000..cd2cee2 --- /dev/null +++ b/src/native/cuda/driver_.h @@ -0,0 +1,47 @@ +#ifndef INFINI_RT_CUDA_DRIVER__H_ +#define INFINI_RT_CUDA_DRIVER__H_ + +#include + +#include "driver.h" + +namespace infini::rt::driver { + +template +struct CudaDriver : DeviceDriver { + static constexpr bool Validate() { + DeviceDriver::Validate(); + static_assert(sizeof(typename Derived::Function) > 0, + "`Driver` must define a `Function` type alias."); + static_assert(sizeof(typename Derived::Module) > 0, + "`Driver` must define a `Module` type alias."); + static_assert(sizeof(typename Derived::Stream) > 0, + "`Driver` must define a `Stream` type alias."); + static_assert(std::is_invocable_v, + "`Driver::ModuleLoadData` must be callable with `(Module*, " + "const void*)`."); + static_assert(std::is_invocable_v, + "`Driver::ModuleGetFunction` must be callable with " + "`(Function*, Module, const char*)`."); + static_assert(std::is_invocable_v, + "`Driver::ModuleUnload` must be callable with `(Module)`."); + static_assert( + std::is_invocable_v, + "`Driver::LaunchKernel` must be callable with " + "`(Function, unsigned int, unsigned int, unsigned int, unsigned int, " + "unsigned int, unsigned int, unsigned int, Stream, void**, void**)`."); + return true; + } +}; + +} // namespace infini::rt::driver + +#endif diff --git a/src/native/cuda/nvidia/driver_.h b/src/native/cuda/nvidia/driver_.h new file mode 100644 index 0000000..2459653 --- /dev/null +++ b/src/native/cuda/nvidia/driver_.h @@ -0,0 +1,55 @@ +#ifndef INFINI_RT_NVIDIA_DRIVER__H_ +#define INFINI_RT_NVIDIA_DRIVER__H_ + +// clang-format off +#include +// clang-format on + +#include "native/cuda/driver_.h" +#include "native/cuda/nvidia/device_.h" + +namespace infini::rt::driver { + +template <> +struct Driver + : CudaDriver> { + using Result = CUresult; + + using Module = CUmodule; + + using Function = CUfunction; + + using Stream = CUstream; + + static constexpr Device::Type kDeviceType = Device::Type::kNvidia; + + static constexpr Result kSuccess = CUDA_SUCCESS; + + static constexpr auto ModuleLoadData = cuModuleLoadData; + + static constexpr auto ModuleGetFunction = cuModuleGetFunction; + + static constexpr auto ModuleUnload = cuModuleUnload; + + static constexpr auto FuncGetAttribute = cuFuncGetAttribute; + + static constexpr auto kFuncAttributeSharedSizeBytes = + CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES; + + static constexpr auto FuncSetCacheConfig = cuFuncSetCacheConfig; + + static constexpr auto kFuncCachePreferShared = CU_FUNC_CACHE_PREFER_SHARED; + + static constexpr auto FuncSetAttribute = cuFuncSetAttribute; + + static constexpr auto kFuncAttributeMaxDynamicSharedSizeBytes = + CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES; + + static constexpr auto LaunchKernel = cuLaunchKernel; +}; + +static_assert(Driver::Validate()); + +} // namespace infini::rt::driver + +#endif