sycl: honor GGML_HINT_SRC0_IS_HADAMARD (llama/27298)
Kernel is a port of `ggml-cuda/fwht.cu` (us/run, median): ``` m x n x k GEMM FWHT speedup 64 x 1 x 64 10.20 2.93 3.48x 64 x 2048 x 64 10.75 2.71 3.97x 128 x 1 x 128 10.33 2.88 3.59x 128 x 32 x 128 9.20 2.77 3.33x 128 x 2048 x 128 16.46 2.76 5.95x 256 x 1 x 256 10.19 2.77 3.68x 256 x 2048 x 256 16.69 3.41 4.89x 512 x 2048 x 512 54.16 12.89 4.20x ```
This commit is contained in:
parent
c6a2bd0973
commit
dbab353f03
|
|
@ -0,0 +1,119 @@
|
|||
#include "fwht.hpp"
|
||||
|
||||
#include <cmath>
|
||||
|
||||
template <int N>
|
||||
static void fwht_kernel(const float * __restrict__ src, float * __restrict__ dst, const int64_t n_rows,
|
||||
const float scale, const sycl::nd_item<2> & item) {
|
||||
const sycl::sub_group sg = item.get_sub_group();
|
||||
|
||||
const int64_t r = item.get_global_id(0);
|
||||
if (r >= n_rows) {
|
||||
return;
|
||||
}
|
||||
|
||||
src += r * N;
|
||||
dst += r * N;
|
||||
|
||||
constexpr int el_w = N / WARP_SIZE;
|
||||
static_assert(el_w >= 1 && N % WARP_SIZE == 0, "row must be a whole number of sub-group widths");
|
||||
|
||||
float reg[el_w];
|
||||
const int lane = sg.get_local_linear_id();
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < el_w; ++i) {
|
||||
reg[i] = src[i * WARP_SIZE + lane] * scale;
|
||||
}
|
||||
|
||||
// Butterflies inside the sub-group. The partner of a lane with bit h clear is the
|
||||
// lower index of the pair, so it takes the sum and the upper takes lower - upper.
|
||||
#pragma unroll
|
||||
for (int h = 1; h < WARP_SIZE; h *= 2) {
|
||||
#pragma unroll
|
||||
for (int j = 0; j < el_w; ++j) {
|
||||
const float val = reg[j];
|
||||
const float val2 = dpct::permute_sub_group_by_xor(sg, val, h, WARP_SIZE);
|
||||
|
||||
reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
|
||||
}
|
||||
}
|
||||
|
||||
// Butterflies across registers: h is a multiple of WARP_SIZE, so the partner of
|
||||
// element i*WARP_SIZE + lane lives in reg[i + h/WARP_SIZE] on the same lane.
|
||||
#pragma unroll
|
||||
for (int h = WARP_SIZE; h < N; h *= 2) {
|
||||
const int step = h / WARP_SIZE;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < el_w; j += 2 * step) {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < step; ++k) {
|
||||
const float x = reg[j + k];
|
||||
const float y = reg[j + k + step];
|
||||
|
||||
reg[j + k] = x + y;
|
||||
reg[j + k + step] = x - y;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < el_w; ++i) {
|
||||
dst[i * WARP_SIZE + lane] = reg[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <int N>
|
||||
static void launch_fwht(const float * src, float * dst, const int64_t n_rows, const float scale,
|
||||
dpct::queue_ptr stream) {
|
||||
constexpr int rows_per_block = 4;
|
||||
|
||||
const int64_t num_blocks = (n_rows + rows_per_block - 1) / rows_per_block;
|
||||
|
||||
// dim 1 is the fastest-varying, so a sub-group is exactly one row's WARP_SIZE lanes.
|
||||
const sycl::range<2> global(num_blocks * rows_per_block, WARP_SIZE);
|
||||
const sycl::range<2> local(rows_per_block, WARP_SIZE);
|
||||
|
||||
stream->parallel_for(sycl::nd_range<2>(global, local),
|
||||
[=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
|
||||
fwht_kernel<N>(src, dst, n_rows, scale, item);
|
||||
});
|
||||
}
|
||||
|
||||
bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
|
||||
if (src->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
|
||||
return false;
|
||||
}
|
||||
if (!ggml_are_same_shape(src, dst)) {
|
||||
return false;
|
||||
}
|
||||
if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const int n = (int) src->ne[0];
|
||||
const int64_t rows = ggml_nrows(src);
|
||||
|
||||
const float * src_d = (const float *) src->data;
|
||||
float * dst_d = (float *) dst->data;
|
||||
dpct::queue_ptr stream = ctx.stream();
|
||||
|
||||
const float scale = 1.0f / std::sqrt((float) n);
|
||||
|
||||
switch (n) {
|
||||
case 64:
|
||||
launch_fwht<64>(src_d, dst_d, rows, scale, stream);
|
||||
return true;
|
||||
case 128:
|
||||
launch_fwht<128>(src_d, dst_d, rows, scale, stream);
|
||||
return true;
|
||||
case 256:
|
||||
launch_fwht<256>(src_d, dst_d, rows, scale, stream);
|
||||
return true;
|
||||
case 512:
|
||||
launch_fwht<512>(src_d, dst_d, rows, scale, stream);
|
||||
return true;
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
#ifndef GGML_SYCL_FWHT_HPP
|
||||
#define GGML_SYCL_FWHT_HPP
|
||||
|
||||
#include "common.hpp"
|
||||
|
||||
// Fast Walsh-Hadamard transform, the fast path for a MUL_MAT whose src0 ggml has
|
||||
// tagged GGML_HINT_SRC0_IS_HADAMARD. src0 is not read at all. Returns false if the
|
||||
// shape is not one this can serve, in which case the caller must fall through to the
|
||||
// ordinary mat-mul dispatch.
|
||||
bool ggml_sycl_op_fwht(ggml_backend_sycl_context & ctx, const ggml_tensor * src, ggml_tensor * dst);
|
||||
|
||||
#endif // GGML_SYCL_FWHT_HPP
|
||||
|
|
@ -58,6 +58,7 @@
|
|||
#include "ggml-sycl/backend.hpp"
|
||||
#include "ggml-sycl/common.hpp"
|
||||
#include "ggml-sycl/element_wise.hpp"
|
||||
#include "ggml-sycl/fwht.hpp"
|
||||
#include "ggml-sycl/gemm.hpp"
|
||||
#include "ggml-sycl/getrows.hpp"
|
||||
#include "ggml-sycl/norm.hpp"
|
||||
|
|
@ -4473,6 +4474,18 @@ static bool can_use_mul_mat_vec_q(const ggml_tensor * src0, const ggml_tensor *
|
|||
|
||||
static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
|
||||
scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2);
|
||||
|
||||
// Handle HADAMARAD hint given from further up the pipeline and pass it to the correct
|
||||
// kernel.
|
||||
//
|
||||
// The op check is not redundant: this backend also routes MUL_MAT_ID through here with a
|
||||
// stack copy of dst, which carries MUL_MAT_ID's own op_params. ggml_mul_mat_set_hint()
|
||||
// asserts GGML_OP_MUL_MAT for the same reason.
|
||||
if (dst->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(dst, 1) == GGML_HINT_SRC0_IS_HADAMARD &&
|
||||
ggml_sycl_op_fwht(ctx, src1, dst)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const bool split = ggml_backend_buffer_is_sycl_split(src0->buffer);
|
||||
int64_t min_compute_capability = INT_MAX;
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue