diff --git a/ggml/src/ggml-sycl/backend.hpp b/ggml/src/ggml-sycl/backend.hpp index f299bcf62..51ab6f930 100644 --- a/ggml/src/ggml-sycl/backend.hpp +++ b/ggml/src/ggml-sycl/backend.hpp @@ -26,6 +26,7 @@ #include "dmmv.hpp" #include "element_wise.hpp" #include "fattn.hpp" +#include "fusion.hpp" #include "gated_delta_net.hpp" #include "gla.hpp" #include "im2col.hpp" diff --git a/ggml/src/ggml-sycl/fusion.cpp b/ggml/src/ggml-sycl/fusion.cpp new file mode 100644 index 000000000..4a6027f39 --- /dev/null +++ b/ggml/src/ggml-sycl/fusion.cpp @@ -0,0 +1,44 @@ +#include "fusion.hpp" + +bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list ops) { + if (!g_ggml_sycl_enable_fusion) { + return false; + } + + if (!ggml_can_fuse(cgraph, node_idx, ops)) { + return false; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + const ggml_tensor * rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + + GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); + + if (mul->src[0]->type != GGML_TYPE_F32 || + mul->src[1]->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32) { + return false; + } + + // if rms norm is the B operand, then we don't handle broadcast + if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { + return false; + } + + const ggml_tensor * mul_w = (mul->src[0] == rms_norm) ? mul->src[1] : mul->src[0]; + // the fused kernel indexes the weight as mul[col], so it must span ncols contiguously + if (mul_w->ne[0] != rms_norm->ne[0] || mul_w->nb[0] != ggml_type_size(mul_w->type)) { + return false; + } + + if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { + return false; + } + + return true; + } + + return false; +} diff --git a/ggml/src/ggml-sycl/fusion.hpp b/ggml/src/ggml-sycl/fusion.hpp new file mode 100644 index 000000000..7d7c79e02 --- /dev/null +++ b/ggml/src/ggml-sycl/fusion.hpp @@ -0,0 +1,15 @@ +#ifndef GGML_SYCL_FUSION_HPP +#define GGML_SYCL_FUSION_HPP + +#include + +#include "common.hpp" + +// Backend-side fusability test. `ops` names a candidate op sequence starting at cgraph node +// `node_idx`; the result is true only if ggml considers that subgraph fusable *and* the SYCL +// kernel which would service it accepts the tensors involved (types, shapes, contiguity). +// +// Lives in its own translation unit because it grows a branch per supported op sequence. +bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list ops); + +#endif // GGML_SYCL_FUSION_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 3b807c7cb..49bdc97d2 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5398,6 +5398,13 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc } } #endif + if (node->op == GGML_OP_RMS_NORM && + ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { + ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]); + i++; + continue; + } + bool ok = ggml_sycl_compute_forward(*sycl_ctx, node); if (!ok) { GGML_LOG_ERROR("%s: error: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); diff --git a/ggml/src/ggml-sycl/norm.cpp b/ggml/src/ggml-sycl/norm.cpp index c4472e4bd..682a9f51e 100644 --- a/ggml/src/ggml-sycl/norm.cpp +++ b/ggml/src/ggml-sycl/norm.cpp @@ -147,10 +147,13 @@ static void group_norm_f32(const float* x, float* dst, const int group_size, con } } +template static void rms_norm_f32(const float* x, float* dst, const int ncols, const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, - const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size) { + const float eps, const sycl::nd_item<3>& item_ct1, float* s_sum, int block_size, + const float* mul = nullptr, const int64_t mul_stride_row = 0, const int64_t mul_stride_channel = 0, + const int64_t mul_stride_sample = 0, const int mul_nrows = 0, const int mul_nchannels = 0, const int mul_nsamples = 0) { const int nrows = item_ct1.get_group_range(2); const int nchannels = item_ct1.get_group_range(1); @@ -170,6 +173,12 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, x += src_offset; dst += dst_offset; + if constexpr (do_multiply) { + const int mul_row = row % mul_nrows; + const int mul_channel = channel % mul_nchannels; + const int mul_sample = sample % mul_nsamples; + mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row; + } float tmp = 0.0f; // partial sum for thread in warp @@ -202,7 +211,11 @@ static void rms_norm_f32(const float* x, float* dst, const int ncols, const float scale = sycl::rsqrt(mean + eps); for (int col = tid; col < ncols; col += block_size) { - dst[col * dst_stride_col] = scale * x[col * src_stride_col]; + if constexpr (do_multiply) { + dst[col * dst_stride_col] = scale * x[col * src_stride_col] * mul[col]; + } else { + dst[col * dst_stride_col] = scale * x[col * src_stride_col]; + } } } @@ -376,6 +389,49 @@ static void rms_norm_f32_sycl(const float* x, float* dst, const int ncols, const } } +static void rms_norm_mul_f32_sycl(const float* x, const float* mul, float* dst, const int ncols, const int nrows, + const int nchannels, const int nsamples, + const int64_t src_stride_col, const int64_t src_stride_row, const int64_t src_stride_channel, const int64_t src_stride_sample, + const int64_t dst_stride_col, const int64_t dst_stride_row, const int64_t dst_stride_channel, const int64_t dst_stride_sample, + const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, + const int mul_nrows, const int mul_nchannels, const int mul_nsamples, + const float eps, queue_ptr stream, int device) { + const sycl::range<3> global_dims(nsamples, nchannels, nrows); + if (ncols < 1024) { + const sycl::range<3> block_dims(1, 1, WARP_SIZE); + stream->submit([&](sycl::handler& cgh) { + cgh.parallel_for( + sycl::nd_range<3>(global_dims * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) + [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, nullptr, WARP_SIZE, + mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples); + }); + }); + } + else { + const int work_group_size = ggml_sycl_info().max_work_group_sizes[device]; + assert(work_group_size % (WARP_SIZE * WARP_SIZE) == 0); + const sycl::range<3> block_dims(1, 1, work_group_size); + stream->submit([&](sycl::handler& cgh) { + sycl::local_accessor s_sum_acc_ct1(sycl::range<1>(work_group_size / WARP_SIZE), cgh); + cgh.parallel_for( + sycl::nd_range<3>(global_dims * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) + [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + rms_norm_f32(x, dst, ncols, + src_stride_col, src_stride_row, src_stride_channel, src_stride_sample, + dst_stride_col, dst_stride_row, dst_stride_channel, dst_stride_sample, + eps, item_ct1, get_pointer(s_sum_acc_ct1), work_group_size, + mul, mul_stride_row, mul_stride_channel, mul_stride_sample, mul_nrows, mul_nchannels, mul_nsamples); + }); + }); + } +} + template static void l2_norm_f32_sycl(const float * x, float * dst, @@ -518,6 +574,66 @@ void ggml_sycl_op_rms_norm(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { ss0, ss1, ss2, ss3, ds0, ds1, ds2, ds3, eps, main_stream, ctx.device); } +void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { + const ggml_tensor * rms_norm_src = dst->src[0]; + float eps = 0.0f; + memcpy(&eps, dst->op_params, sizeof(float)); + + const float * src0_dd = static_cast(rms_norm_src->data); + const float * mul_dd = nullptr; + const ggml_tensor * mul_src = nullptr; + if (mul_tensor->src[0] == dst) { + mul_dd = static_cast(mul_tensor->src[1]->data); + mul_src = mul_tensor->src[1]; + } else if (mul_tensor->src[1] == dst) { + mul_dd = static_cast(mul_tensor->src[0]->data); + mul_src = mul_tensor->src[0]; + } else { + GGML_ASSERT(false); + } + float * dst_dd = static_cast(mul_tensor->data); + + dpct::queue_ptr main_stream = ctx.stream(); + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + GGML_ASSERT(rms_norm_src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(mul_tensor->type == GGML_TYPE_F32); + GGML_ASSERT(eps >= 0.0f); + + const int64_t ne00 = rms_norm_src->ne[0]; + const int64_t ne01 = rms_norm_src->ne[1]; + const int64_t ne02 = rms_norm_src->ne[2]; + const int64_t ne03 = rms_norm_src->ne[3]; + + const size_t ts0 = ggml_type_size(rms_norm_src->type); + GGML_ASSERT(rms_norm_src->nb[0] == ts0); + const int64_t s00 = rms_norm_src->nb[0] / ts0; + const int64_t s01 = rms_norm_src->nb[1] / ts0; + const int64_t s02 = rms_norm_src->nb[2] / ts0; + const int64_t s03 = rms_norm_src->nb[3] / ts0; + + const size_t tdst = ggml_type_size(mul_tensor->type); + GGML_ASSERT(mul_tensor->nb[0] == tdst); + const int64_t d00 = mul_tensor->nb[0] / tdst; + const int64_t d01 = mul_tensor->nb[1] / tdst; + const int64_t d02 = mul_tensor->nb[2] / tdst; + const int64_t d03 = mul_tensor->nb[3] / tdst; + + const size_t ts_mul = ggml_type_size(mul_src->type); + GGML_ASSERT(mul_src->nb[0] == ts_mul); + const int64_t mul_s01 = mul_src->nb[1] / ts_mul; + const int64_t mul_s02 = mul_src->nb[2] / ts_mul; + const int64_t mul_s03 = mul_src->nb[3] / ts_mul; + const int mul_nrows = mul_src->ne[1]; + const int mul_nchannels = mul_src->ne[2]; + const int mul_nsamples = mul_src->ne[3]; + + rms_norm_mul_f32_sycl(src0_dd, mul_dd, dst_dd, ne00, ne01, ne02, ne03, + s00, s01, s02, s03, d00, d01, d02, d03, + mul_s01, mul_s02, mul_s03, mul_nrows, mul_nchannels, mul_nsamples, eps, main_stream, ctx.device); +} + void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); diff --git a/ggml/src/ggml-sycl/norm.hpp b/ggml/src/ggml-sycl/norm.hpp index 8cb885eb2..51217c421 100644 --- a/ggml/src/ggml-sycl/norm.hpp +++ b/ggml/src/ggml-sycl/norm.hpp @@ -19,6 +19,8 @@ void ggml_sycl_op_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); void ggml_sycl_op_rms_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst); +void ggml_sycl_op_rms_norm_fused(ggml_backend_sycl_context& ctx, ggml_tensor* dst, ggml_tensor* mul); + void ggml_sycl_op_rms_norm_back(ggml_backend_sycl_context& ctx, ggml_tensor* dst); void ggml_sycl_op_group_norm(ggml_backend_sycl_context& ctx, ggml_tensor* dst);