| #include "cuda_driver_loader.h"
|
| #include "rowm_cuda_validation.h"
|
|
|
| extern const unsigned char gemm_ptx[];
|
| extern const unsigned int gemm_ptx_len;
|
|
|
| enum {
|
| SOV_GEMM_ERR_INVALID_ARGUMENT = -1,
|
| SOV_GEMM_ERR_MODULE_LOAD = -2,
|
| SOV_GEMM_ERR_TENSOR_LOOKUP = -3,
|
| SOV_GEMM_ERR_SCALAR_LOOKUP = -4,
|
| SOV_GEMM_ERR_LAUNCH = -5,
|
| SOV_GEMM_ERR_GRID_TOO_LARGE = -6
|
| };
|
|
|
| static CUmodule g_gemm_module;
|
| static CUfunction g_gemm_tensor_kernel;
|
| static CUfunction g_gemm_scalar_kernel;
|
| static uint64_t g_gemm_context_generation;
|
|
|
| static int checked_mul_u64(uint64_t left, uint64_t right, uint64_t* product) {
|
| if (!product) return -1;
|
| if (left != 0 && right > UINT64_MAX / left) return -1;
|
| *product = left * right;
|
| return 0;
|
| }
|
|
|
| #define UINT64_MAX ((uint64_t)0xFFFFFFFFFFFFFFFFu)
|
|
|
| static int matrix_end(CUdeviceptr base, uint64_t rows, uint64_t cols,
|
| uint64_t ld, CUdeviceptr* end) {
|
| uint64_t last_row, elements, bytes;
|
| if (!base || !rows || !cols || !end) return -1;
|
| if (checked_mul_u64(rows - 1u, ld, &last_row) != 0
|
| || last_row > UINT64_MAX - cols) return -1;
|
| elements = last_row + cols;
|
| if (checked_mul_u64(elements, 2u, &bytes) != 0
|
| || base > UINT64_MAX - bytes) return -1;
|
| *end = base + bytes;
|
| return 0;
|
| }
|
|
|
| static int ranges_overlap(CUdeviceptr lb, CUdeviceptr le,
|
| CUdeviceptr rb, CUdeviceptr re) {
|
| return lb < re && rb < le;
|
| }
|
|
|
| int sov_cuda_gemm_init(void) {
|
| CUresult result;
|
| uint64_t ctx_gen;
|
|
|
| if (gemm_ptx_len == 0u || gemm_ptx[gemm_ptx_len - 1u] != 0u)
|
| return SOV_GEMM_ERR_MODULE_LOAD;
|
| if (!sov_cuda_is_initialized() && sov_cuda_init() != 0)
|
| return SOV_GEMM_ERR_MODULE_LOAD;
|
|
|
| ctx_gen = sov_cuda_context_generation();
|
| if (ctx_gen == 0u) return SOV_GEMM_ERR_MODULE_LOAD;
|
|
|
| if (g_gemm_tensor_kernel && g_gemm_scalar_kernel
|
| && g_gemm_context_generation == ctx_gen) return 0;
|
|
|
| if (g_gemm_context_generation != ctx_gen) {
|
| g_gemm_module = 0;
|
| g_gemm_tensor_kernel = 0;
|
| g_gemm_scalar_kernel = 0;
|
| g_gemm_context_generation = 0;
|
| }
|
|
|
| result = sov_cuda_module_load_data(&g_gemm_module, (const void*)gemm_ptx);
|
| if (result != CUDA_SUCCESS) { g_gemm_module = 0; return SOV_GEMM_ERR_MODULE_LOAD; }
|
|
|
| result = sov_cuda_module_get_function(&g_gemm_tensor_kernel, g_gemm_module,
|
| "gemm_f16_f32_accum");
|
| if (result != CUDA_SUCCESS) {
|
| sov_cuda_module_unload(g_gemm_module);
|
| g_gemm_module = 0; g_gemm_tensor_kernel = 0;
|
| return SOV_GEMM_ERR_TENSOR_LOOKUP;
|
| }
|
|
|
| result = sov_cuda_module_get_function(&g_gemm_scalar_kernel, g_gemm_module,
|
| "gemm_f16_f32_accum_scalar");
|
| if (result != CUDA_SUCCESS) {
|
| sov_cuda_module_unload(g_gemm_module);
|
| g_gemm_module = 0; g_gemm_tensor_kernel = 0; g_gemm_scalar_kernel = 0;
|
| return SOV_GEMM_ERR_SCALAR_LOOKUP;
|
| }
|
|
|
| g_gemm_context_generation = ctx_gen;
|
| return 0;
|
| }
|
|
|
| int sov_cuda_gemm_ex(CUdeviceptr A, CUdeviceptr B, CUdeviceptr C,
|
| int M, int N, int K,
|
| int lda, int ldb, int ldc) {
|
| CUfunction kernel;
|
| CUresult result;
|
| unsigned int grid_x, grid_y, block_x, block_y;
|
| int power_state, use_tensor, prepare;
|
| CUdeviceptr a_end, b_end, c_end;
|
| void* args[10];
|
|
|
| if (!A || !B || !C
|
| || M < 0 || N < 0 || K < 0
|
| || lda < K || ldb < N || ldc < N
|
| || lda < 0 || ldb < 0 || ldc < 0)
|
| return SOV_GEMM_ERR_INVALID_ARGUMENT;
|
|
|
| if (M == 0 || N == 0) return 0;
|
|
|
| if (matrix_end(C, (uint64_t)M, (uint64_t)N, (uint64_t)ldc, &c_end) != 0)
|
| return SOV_GEMM_ERR_INVALID_ARGUMENT;
|
|
|
| if (K > 0) {
|
| if (matrix_end(A, (uint64_t)M, (uint64_t)K, (uint64_t)lda, &a_end) != 0
|
| || matrix_end(B, (uint64_t)K, (uint64_t)N, (uint64_t)ldb, &b_end) != 0
|
| || ranges_overlap(A, a_end, C, c_end)
|
| || ranges_overlap(B, b_end, C, c_end))
|
| return SOV_GEMM_ERR_INVALID_ARGUMENT;
|
| }
|
|
|
| use_tensor = (M & 15) == 0 && (N & 7) == 0 && (K & 15) == 0;
|
| if (use_tensor) {
|
| grid_x = (unsigned int)M / 16u;
|
| grid_y = (unsigned int)N / 8u;
|
| block_x = 32u; block_y = 1u;
|
| } else {
|
| grid_x = ((unsigned int)N + 15u) / 16u;
|
| grid_y = ((unsigned int)M + 15u) / 16u;
|
| block_x = 16u; block_y = 16u;
|
| }
|
| if (grid_y > 65535u) return SOV_GEMM_ERR_GRID_TOO_LARGE;
|
|
|
| prepare = sov_cuda_gemm_init();
|
| if (prepare != 0) return prepare;
|
|
|
| prepare = sov_cuda_validation_require_authorized(sov_cuda_context_generation());
|
| if (prepare != SOV_CUDA_ROWM_OK) return prepare;
|
|
|
| kernel = use_tensor ? g_gemm_tensor_kernel : g_gemm_scalar_kernel;
|
| power_state = 0;
|
|
|
| args[0] = &A; args[1] = &B; args[2] = &C;
|
| args[3] = &M; args[4] = &N; args[5] = &K;
|
| args[6] = &lda; args[7] = &ldb; args[8] = &ldc;
|
| args[9] = &power_state;
|
|
|
| result = sov_cuda_launch_kernel(kernel,
|
| grid_x, grid_y, 1u,
|
| block_x, block_y, 1u,
|
| 0u, (CUstream)0, args, 0);
|
| return result == CUDA_SUCCESS ? 0 : SOV_GEMM_ERR_LAUNCH;
|
| }
|
|
|
| int sov_cuda_gemm(CUdeviceptr A, CUdeviceptr B, CUdeviceptr C,
|
| int M, int N, int K) {
|
| return sov_cuda_gemm_ex(A, B, C, M, N, K, K, N, N);
|
| }
|
|
|
| void sov_cuda_gemm_shutdown(void) {
|
| if (g_gemm_module
|
| && sov_cuda_is_initialized()
|
| && g_gemm_context_generation == sov_cuda_context_generation())
|
| sov_cuda_module_unload(g_gemm_module);
|
| g_gemm_module = 0;
|
| g_gemm_tensor_kernel = 0;
|
| g_gemm_scalar_kernel = 0;
|
| g_gemm_context_generation = 0;
|
| }
|
|
|
