diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 67f8a93248d1..9debf7312c01 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -34,6 +34,7 @@ #include "ggml-cuda/im2col.cuh" #include "ggml-cuda/mmf.cuh" #include "ggml-cuda/mmq.cuh" +#include "ggml-cuda/mul_mat_f32_tall_skinny.cuh" #include "ggml-cuda/mmvf.cuh" #include "ggml-cuda/mmvq.cuh" #include "ggml-cuda/mmvdq.cuh" @@ -1787,6 +1788,11 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); return; } + // f32 tall-skinny GEMM for gfx1151 (RDNA3.5): K=2048, N=32. + if (ggml_cuda_should_use_mul_mat_f32_tall_skinny(src0, src1, dst, cc)) { + ggml_cuda_mul_mat_f32_tall_skinny(ctx, src0, src1, dst); + return; + } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); return; diff --git a/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cu b/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cu new file mode 100644 index 000000000000..d83b44fe6eb7 --- /dev/null +++ b/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cu @@ -0,0 +1,233 @@ +#include "ggml.h" +#include "common.cuh" +#include "mul_mat_f32_tall_skinny.cuh" + +#if defined(GGML_USE_HIP) + +// Wave32 sum reduction, result in lane 0. On RDNA3.5 (the only arch this kernel runs on) +// use row-shift DPP to tree-add within each 16-lane row, then one permlanex16 to add the +// two rows. Those builtins need gfx10+ instructions, so other archs (compiled but never +// launched) fall back to a portable shuffle reduction. +#if defined(RDNA3_5) +#define DPP_ROW_SHIFT_ADD(v, ctrl) (v) += __int_as_float(__builtin_amdgcn_mov_dpp(__float_as_int(v), ctrl, 0xf, 0xf, 1)) +static __device__ __forceinline__ float permute_xor16(float v) { + return __int_as_float(__builtin_amdgcn_permlanex16(__float_as_int(v), __float_as_int(v), + 0x76543210u, 0xFEDCBA98u, false, false)); +} +static __device__ __forceinline__ float warp_reduce_sum(float v) { + DPP_ROW_SHIFT_ADD(v, 0x108); DPP_ROW_SHIFT_ADD(v, 0x104); DPP_ROW_SHIFT_ADD(v, 0x102); DPP_ROW_SHIFT_ADD(v, 0x101); + v += permute_xor16(v); + return v; +} +#else +static __device__ __forceinline__ float warp_reduce_sum(float v) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) v += __shfl_xor_sync(0xffffffff, v, offset, 32); + return v; +} +#endif // defined(RDNA3_5) + +// Software-pipelined weight stream shared by both kernels. Two alternating register +// sets (cur/nxt) keep DEPTH*WROWS b128 loads in flight across the loop backedge; the +// sched_barriers pin the load-then-compute order so the counter does not drain to zero. +// Each lane reduces a K/32 stripe into acc[WROWS][ACOLS]; act_cols_lds holds ACOLS +// activation columns (float4 per K step). STEP indices count 32-wide K steps. +#define TSN_PIPE_STAGE(cur, nxt, load_step, use_step) do { \ + _Pragma("unroll") \ + for (int d = 0; d < DEPTH; ++d) { \ + const int load_k = lane + ((load_step) + d) * 32; \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) nxt[d][c] = weight_row[c][load_k]; \ + __builtin_amdgcn_sched_barrier(0); \ + const int use_k = lane + ((use_step) + d) * 32; \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) { const float4 w = cur[d][c]; \ + _Pragma("unroll") \ + for (int e = 0; e < ACOLS; ++e) { const float4 a = act_cols_lds[e * k4 + use_k]; \ + acc[c][e] += w.x*a.x + w.y*a.y + w.z*a.z + w.w*a.w; } } \ + __builtin_amdgcn_sched_barrier(0); \ + } } while (0) +#define TSN_PIPE_DRAIN(cur, use_step) do { \ + _Pragma("unroll") \ + for (int d = 0; d < DEPTH; ++d) { \ + const int use_k = lane + ((use_step) + d) * 32; \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) { const float4 w = cur[d][c]; \ + _Pragma("unroll") \ + for (int e = 0; e < ACOLS; ++e) { const float4 a = act_cols_lds[e * k4 + use_k]; \ + acc[c][e] += w.x*a.x + w.y*a.y + w.z*a.z + w.w*a.w; } } \ + } } while (0) + +// Runs the DEPTH-deep prologue, the 2*DEPTH-unrolled steady loop and the epilogue over +// the K dimension for the WROWS x ACOLS output tile owned by this warp, then writes it. +// Requires n_ksteps % (2*DEPTH) == 0. +#define TSN_STREAM_AND_STORE(n_lo, n_hi) do { \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) { const int row = min(nb + c, N - 1); \ + weight_row[c] = (const float4 *) __builtin_assume_aligned(x + (int64_t) row * ldx, 16); }\ + float acc[WROWS][ACOLS]; \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) _Pragma("unroll") for (int e = 0; e < ACOLS; ++e) acc[c][e] = 0.f; \ + float4 w_cur[DEPTH][WROWS], w_nxt[DEPTH][WROWS]; \ + _Pragma("unroll") \ + for (int d = 0; d < DEPTH; ++d) _Pragma("unroll") for (int c = 0; c < WROWS; ++c) w_cur[d][c] = weight_row[c][lane + d * 32]; \ + int step = DEPTH; \ + for (; step + 2 * DEPTH <= n_ksteps; step += 2 * DEPTH) { \ + TSN_PIPE_STAGE(w_cur, w_nxt, step, step - DEPTH); \ + TSN_PIPE_STAGE(w_nxt, w_cur, step + DEPTH, step); \ + } \ + TSN_PIPE_STAGE(w_cur, w_nxt, step, step - DEPTH); \ + TSN_PIPE_DRAIN(w_nxt, step); \ + _Pragma("unroll") \ + for (int c = 0; c < WROWS; ++c) _Pragma("unroll") for (int e = 0; e < ACOLS; ++e) { \ + const float sum = warp_reduce_sum(acc[c][e]); \ + if (lane == 0 && (nb + c) < (n_hi) && (mb + e) < M) dst[(int64_t)(mb + e) * N + (nb + c)] = sum; } \ + } while (0) + +// dst[N,M] = weight[K,N]^T . activation[K,M]. The activation is staged resident in LDS +// and the weight is streamed from global memory. Persistent over M tiles: while the +// weight stream of the current M tile is computed, the next tile's activation is +// prefetched into registers and stored into the other of two LDS buffers, overlapping +// the activation load with the weight compute. vec_per_thread = (K/4)/blockDim.x. +template +static __global__ void __launch_bounds__(512) +mul_mat_f32_tall_skinny_double_buffered(const float * __restrict__ x, const float * __restrict__ y, + float * __restrict__ dst, const int K, const int N, const int M, const int ldx) { + const int block_size = blockDim.x, k4 = K / 4, n_ksteps = k4 / 32; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, nwarps = block_size >> 5; + const int m_stride = gridDim.x * ACOLS; + extern __shared__ float smem[]; + float4 * act_buf0 = (float4 *) smem; + float4 * act_buf1 = act_buf0 + (int64_t) ACOLS * k4; + float4 prefetch[ACOLS][VEC_PER_THREAD > 0 ? VEC_PER_THREAD : 1]; + { const int mb = blockIdx.x * ACOLS; + #pragma unroll + for (int e = 0; e < ACOLS; ++e) { const int col = min(mb + e, M - 1); + const float4 * yr = (const float4 *) __builtin_assume_aligned(y + (int64_t) col * K, 16); + #pragma unroll + for (int p = 0; p < VEC_PER_THREAD; ++p) { const int j = tid + p * block_size; + prefetch[e][p] = mb < M ? yr[j] : make_float4(0, 0, 0, 0); } } } + int buf = 0; + for (int mb = blockIdx.x * ACOLS; mb < M; mb += m_stride) { + float4 * act_cols_lds = buf ? act_buf1 : act_buf0; + #pragma unroll + for (int e = 0; e < ACOLS; ++e) + #pragma unroll + for (int p = 0; p < VEC_PER_THREAD; ++p) { const int j = tid + p * block_size; act_cols_lds[e * k4 + j] = prefetch[e][p]; } + __syncthreads(); + const int next_mb = mb + m_stride; + #pragma unroll + for (int e = 0; e < ACOLS; ++e) { const int col = min(next_mb + e, M - 1); + const float4 * yr = (const float4 *) __builtin_assume_aligned(y + (int64_t) col * K, 16); + #pragma unroll + for (int p = 0; p < VEC_PER_THREAD; ++p) { const int j = tid + p * block_size; + prefetch[e][p] = next_mb < M ? yr[j] : make_float4(0, 0, 0, 0); } } + for (int nb = warp * WROWS; nb < N; nb += nwarps * WROWS) { + const float4 * weight_row[WROWS]; + TSN_STREAM_AND_STORE(0, N); + } + __syncthreads(); + buf ^= 1; + } +} + +// Same weight-streamed kernel, but the N (weight-row) dimension is split across gridDim.y +// so more blocks run concurrently (higher occupancy) without extra accumulators or a +// cross-block reduction, since each split writes a disjoint N range of dst. +template +static __global__ void __launch_bounds__(512, 1) +mul_mat_f32_tall_skinny_split_n(const float * __restrict__ x, const float * __restrict__ y, + float * __restrict__ dst, const int K, const int N, const int M, const int ldx) { + const int block_size = blockDim.x, k4 = K / 4, n_ksteps = k4 / 32; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, nwarps = block_size >> 5; + const int n_splits = gridDim.y, split = blockIdx.y; + const int n_lo = (int) (((int64_t) N * split) / n_splits); + const int n_hi = (int) (((int64_t) N * (split + 1)) / n_splits); + extern __shared__ float smem[]; + float4 * act_cols_lds = (float4 *) smem; + for (int mb = blockIdx.x * ACOLS; mb < M; mb += gridDim.x * ACOLS) { + #pragma unroll + for (int e = 0; e < ACOLS; ++e) { const int col = min(mb + e, M - 1); + const float4 * yr = (const float4 *) __builtin_assume_aligned(y + (int64_t) col * K, 16); + for (int p = tid; p < k4; p += block_size) act_cols_lds[e * k4 + p] = yr[p]; } + __syncthreads(); + for (int nb = n_lo + warp * WROWS; nb < n_hi; nb += nwarps * WROWS) { + const float4 * weight_row[WROWS]; + TSN_STREAM_AND_STORE(n_lo, n_hi); + } + __syncthreads(); + } +} + +#undef TSN_STREAM_AND_STORE +#undef TSN_PIPE_STAGE +#undef TSN_PIPE_DRAIN +#undef DPP_ROW_SHIFT_ADD + +#endif // defined(GGML_USE_HIP) + +bool ggml_cuda_should_use_mul_mat_f32_tall_skinny(const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * dst, int cc) { +#if defined(GGML_USE_HIP) + if (!GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return false; + } + if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { + return false; + } + if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { + return false; + } + const int64_t K = src0->ne[0], N = src0->ne[1]; + return K == 2048 && N == 32 && src1->ne[0] == K; +#else + GGML_UNUSED(src0); GGML_UNUSED(src1); GGML_UNUSED(dst); GGML_UNUSED(cc); + return false; +#endif +} + +void ggml_cuda_mul_mat_f32_tall_skinny(ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { +#if defined(GGML_USE_HIP) + const int K = (int) src0->ne[0]; + const int N = (int) src0->ne[1]; + const int M = (int) src1->ne[1]; + const int ldx = (int) (src0->nb[1] / sizeof(float)); + const float * x = (const float *) src0->data; + const float * y = (const float *) src1->data; + float * d = (float *) dst->data; + cudaStream_t stream = ctx.stream(); + + constexpr int block_size = 256; + constexpr int vec_per_thread = (2048 / 4) / block_size; + // Per-M configuration: small batches use the double-buffered kernel, mid batches the + // split-N kernel (128 blocks up to M=256, then 64), large batches the double-buffered + // kernel again with a wider activation tile. + if (M <= 128) { + constexpr int act_cols = 2; + const dim3 grid(64, 1, 1), block(block_size, 1, 1); + const size_t nbytes_shared = (size_t) 2 * act_cols * K * sizeof(float); + mul_mat_f32_tall_skinny_double_buffered<4, act_cols, 2, vec_per_thread> + <<>>(x, y, d, K, N, M, ldx); + } else if (M <= 1024) { + constexpr int act_cols = 4; + const int grid_x = M <= 256 ? 128 : 64; + const dim3 grid(grid_x, 1, 1), block(block_size, 1, 1); + const size_t nbytes_shared = (size_t) act_cols * K * sizeof(float); + mul_mat_f32_tall_skinny_split_n<4, act_cols, 1> + <<>>(x, y, d, K, N, M, ldx); + } else { + constexpr int act_cols = 4; + const dim3 grid(40, 1, 1), block(block_size, 1, 1); + const size_t nbytes_shared = (size_t) 2 * act_cols * K * sizeof(float); + mul_mat_f32_tall_skinny_double_buffered<4, act_cols, 1, vec_per_thread> + <<>>(x, y, d, K, N, M, ldx); + } +#else + GGML_UNUSED(ctx); GGML_UNUSED(src0); GGML_UNUSED(src1); GGML_UNUSED(dst); + GGML_ABORT("ggml_cuda_mul_mat_f32_tall_skinny is HIP-only"); +#endif +} diff --git a/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cuh b/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cuh new file mode 100644 index 000000000000..5ba330e62f0a --- /dev/null +++ b/ggml/src/ggml-cuda/mul_mat_f32_tall_skinny.cuh @@ -0,0 +1,14 @@ +#include "common.cuh" + +// f32 tall-skinny GEMM for gfx1151 / RDNA3.5: the weight [K, N] is tall (large K) +// and skinny (small N). Computes ggml_mul_mat(src0=weight[K,N], src1=activation[K,M]) +// -> dst[N,M] for K=2048, N=32. The activation is staged resident in LDS while the +// weight is streamed from global memory with a backedge software pipeline; the K +// reduction uses a wave32 DPP tree. AMD-only: on non-HIP builds the predicate returns +// false and this path is never taken. + +bool ggml_cuda_should_use_mul_mat_f32_tall_skinny( + const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int cc); + +void ggml_cuda_mul_mat_f32_tall_skinny(ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cac4ac6b4140..b5dc173af1df 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8717,6 +8717,10 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1056, 1, 67, {1, 1}, {4, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 16, 32, 32, { 1, 1}, {1, 1}, {0, 1, 2, 3}, 64, 3)); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 64, 77, 77, {12,1}, {1,1})); + // f32 tall-skinny GEMM (mul_mat_f32_tall_skinny): K=2048, weight rows=32, spanning M buckets. + for (int64_t act : {8, 64, 128, 512, 2048, 4096}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, act, 2048, {1, 1}, {1, 1})); + } test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1, 1}, {1, 1})); @@ -9564,6 +9568,10 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16416, 1, 128, {8, 1}, {4, 1}, {0, 2, 1, 3})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 1, 16416, {8, 1}, {4, 1}, {0, 1, 2, 3}, 2*16416)); + // f32 tall-skinny GEMM (mul_mat_f32_tall_skinny): K=2048, weight rows=32, spanning M buckets. + for (int64_t act : {8, 64, 128, 512, 2048, 4096}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, act, 2048, {1, 1}, {1, 1})); + } // FWHT tests test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 1, 128));