diff --git a/ggml/src/ggml-cuda/solve_tri.cu b/ggml/src/ggml-cuda/solve_tri.cu index 177ffc268..72ff808a4 100644 --- a/ggml/src/ggml-cuda/solve_tri.cu +++ b/ggml/src/ggml-cuda/solve_tri.cu @@ -88,7 +88,7 @@ static void solve_tri_f32_cublas(ggml_backend_cuda_context & ctx, # pragma clang diagnostic push # pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ -template +template static __global__ void solve_tri_f32_fast(const float * __restrict__ A, const float * __restrict__ B, float * __restrict__ X, @@ -100,15 +100,19 @@ static __global__ void solve_tri_f32_fast(const float * __restrict__ A, const size_t nb2, const size_t nb3, const int n_arg, - const int k_arg) { - const int n = n_template == 0 ? n_arg : n_template; - const int k = k_template == 0 ? k_arg : k_template; + const int ncols_arg, + const int ld_arg, + const int col0_arg) { + const int n = n_template == 0 ? n_arg : n_template; + const int ncols = ncols_template == 0 ? ncols_arg : ncols_template; + const int ld = ld_arg; + const int col0 = col0_arg; const int batch_idx = blockIdx.x; const int lane = threadIdx.x; const int col_idx = threadIdx.y; - if (col_idx >= k) { + if (col_idx >= ncols) { return; } @@ -125,7 +129,7 @@ static __global__ void solve_tri_f32_fast(const float * __restrict__ A, const int offset = threadIdx.x + threadIdx.y * blockDim.x; #pragma unroll - for (int i = 0; i < n * n; i += k * WARP_SIZE) { + for (int i = 0; i < n * n; i += ncols * WARP_SIZE) { const int i0 = i + offset; if (i0 < n * n) { sA[i0] = A_batch[i0]; @@ -134,8 +138,8 @@ static __global__ void solve_tri_f32_fast(const float * __restrict__ A, __syncthreads(); - float x_low = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f; - float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f; + float x_low = (lane < n) ? B_batch[lane * ld + col0 + col_idx] : 0.0f; + float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * ld + col0 + col_idx] : 0.0f; const int half = WARP_SIZE; const int nrows_low = (n < half) ? n : half; @@ -172,7 +176,7 @@ static __global__ void solve_tri_f32_fast(const float * __restrict__ A, const int row = rr * WARP_SIZE + lane; if (row < n) { const float val = (row < half) ? x_low : x_high; - X_batch[row * k + col_idx] = val; + X_batch[row * ld + col0 + col_idx] = val; } } } @@ -180,72 +184,99 @@ static __global__ void solve_tri_f32_fast(const float * __restrict__ A, # pragma clang diagnostic pop #endif // __clang__ -static void solve_tri_f32_cuda(const float * A, - const float * B, - float * X, - int n, - int k, - int64_t ne02, - int64_t ne03, - size_t nb02, - size_t nb03, - size_t nb12, - size_t nb13, - size_t nb2, - size_t nb3, - cudaStream_t stream) { +static void solve_tri_f32_cuda_tile(const float * A, + const float * B, + float * X, + int n, + int ncols, + int ld, + int col0, + int64_t ne02, + int64_t ne03, + size_t nb02, + size_t nb03, + size_t nb12, + size_t nb13, + size_t nb2, + size_t nb3, + cudaStream_t stream) { const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02); - dim3 threads(WARP_SIZE, k); + dim3 threads(WARP_SIZE, ncols); dim3 grid(ne02 * ne03); if (n == 64) { - switch (k) { + switch (ncols) { case 32: solve_tri_f32_fast<64, 32> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 16: solve_tri_f32_fast<64, 16> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 14: solve_tri_f32_fast<64, 14> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 12: solve_tri_f32_fast<64, 12> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 10: solve_tri_f32_fast<64, 10> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 8: solve_tri_f32_fast<64, 8> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 6: solve_tri_f32_fast<64, 6> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 4: solve_tri_f32_fast<64, 4> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 2: solve_tri_f32_fast<64, 2> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; case 1: solve_tri_f32_fast<64, 1> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0, ld, col0); break; default: solve_tri_f32_fast<0, 0> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, n, k); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, n, ncols, ld, col0); } } else { // run general case solve_tri_f32_fast<0, 0> - <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, n, k); + <<>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, n, ncols, ld, col0); + } +} + +static void solve_tri_f32_cuda(const float * A, + const float * B, + float * X, + int n, + int k, + int64_t ne02, + int64_t ne03, + size_t nb02, + size_t nb03, + size_t nb12, + size_t nb13, + size_t nb2, + size_t nb3, + cudaStream_t stream) { + if (k <= MAX_K_FAST) { + solve_tri_f32_cuda_tile(A, B, X, n, k, k, 0, ne02, ne03, nb02, nb03, nb12, nb13, nb2, nb3, stream); + return; + } + + for (int col0 = 0; col0 < k; col0 += MAX_K_FAST) { + const int ncols = std::min(MAX_K_FAST, k - col0); + solve_tri_f32_cuda_tile(A, B, X, n, ncols, k, col0, ne02, ne03, nb02, nb03, nb12, nb13, nb2, nb3, stream); } } @@ -261,7 +292,11 @@ void ggml_cuda_op_solve_tri(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const int64_t ne02 = src0->ne[2]; const int64_t ne03 = src0->ne[3]; - if (n <= MAX_N_FAST && k <= MAX_K_FAST) { + if (n <= MAX_N_FAST && (k <= MAX_K_FAST +#if defined(GGML_USE_HIP) + || n == 64 +#endif + )) { solve_tri_f32_cuda((const float *) src0->data, (const float *) src1->data, (float *) dst->data, n, k, src0->ne[2], src0->ne[3], src0->nb[2] / sizeof(float), src0->nb[3] / sizeof(float), src1->nb[2] / sizeof(float), src1->nb[3] / sizeof(float), dst->nb[2] / sizeof(float),