From 55edb2de442b50be0a29c2ed2ec88488560a96c5 Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Tue, 7 Jul 2026 15:48:50 +0800 Subject: [PATCH] [SYCL] support OP cross_entropy_loss, cross_entropy_loss_back (#25236) * support OP cross_entropy_loss, cross_entropy_loss_back * correct format issue --- docs/ops.md | 4 +- docs/ops/SYCL.csv | 9 +- ggml/src/ggml-sycl/cross_entropy_loss.cpp | 255 ++++++++++++++++++++++ ggml/src/ggml-sycl/cross_entropy_loss.hpp | 7 + ggml/src/ggml-sycl/ggml-sycl.cpp | 9 + 5 files changed, 278 insertions(+), 6 deletions(-) create mode 100644 ggml/src/ggml-sycl/cross_entropy_loss.cpp create mode 100644 ggml/src/ggml-sycl/cross_entropy_loss.hpp diff --git a/docs/ops.md b/docs/ops.md index 48f4e33cfb..c66b13c8d5 100644 --- a/docs/ops.md +++ b/docs/ops.md @@ -35,8 +35,8 @@ Legend: | COS | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ | | COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | -| CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | DIAG | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/docs/ops/SYCL.csv b/docs/ops/SYCL.csv index 59e2ab32dd..1bd777b2da 100644 --- a/docs/ops/SYCL.csv +++ b/docs/ops/SYCL.csv @@ -407,6 +407,7 @@ "SYCL0","GET_ROWS","type=i32,n=256,m=5,r=4,be1=7,be2=1,v=0","support","1","yes","SYCL" "SYCL0","GET_ROWS","type=i32,n=256,m=5,r=4,be1=7,be2=1,v=1","support","1","yes","SYCL" "SYCL0","GET_ROWS_BACK","type=f32,n=1,m=8,r=2,b=1,v=0","support","0","no","SYCL" +"SYCL0","GET_ROWS_BACK","type=f32,n=1,m=70000,r=4,b=1,v=0","support","0","no","SYCL" "SYCL0","GET_ROWS_BACK","type=f32,n=256,m=5,r=4,b=1,v=0","support","0","no","SYCL" "SYCL0","GET_ROWS_BACK","type=f32,n=256,m=5,r=4,b=1,v=1","support","0","no","SYCL" "SYCL0","GET_ROWS_BACK","type=f16,n=256,m=5,r=4,b=1,v=0","support","0","no","SYCL" @@ -16747,10 +16748,10 @@ zjy 2 "SYCL0","FLASH_ATTN_EXT","hsk=128,hsv=64,nh=4,nr23=[1,1],kv=128,nb=2,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=q1_0,type_V=q4_0,permute=[0,1,2,3]","support","0","no","SYCL" "SYCL0","FLASH_ATTN_EXT","hsk=64,hsv=128,nh=4,nr23=[1,1],kv=128,nb=2,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=q4_0,type_V=q1_0,permute=[0,1,2,3]","support","0","no","SYCL" "SYCL0","FLASH_ATTN_EXT","hsk=128,hsv=64,nh=4,nr23=[1,1],kv=64,nb=2,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=q1_0,type_V=f16,permute=[0,1,2,3]","support","0","no","SYCL" -"SYCL0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","0","no","SYCL" -"SYCL0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","0","no","SYCL" -"SYCL0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","0","no","SYCL" -"SYCL0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","0","no","SYCL" +"SYCL0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","1","yes","SYCL" +"SYCL0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","1","yes","SYCL" +"SYCL0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","1","yes","SYCL" +"SYCL0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","1","yes","SYCL" "SYCL0","OPT_STEP_ADAMW","type=f32,ne=[10,5,4,3]","support","0","no","SYCL" "SYCL0","OPT_STEP_SGD","type=f32,ne=[10,5,4,3]","support","0","no","SYCL" "SYCL0","GATED_DELTA_NET","type=f32,head_count=32,head_size=128,n_seq_tokens=1,n_seqs=1,v_repeat=1,permuted=0,kda=0,K=1","support","1","yes","SYCL" diff --git a/ggml/src/ggml-sycl/cross_entropy_loss.cpp b/ggml/src/ggml-sycl/cross_entropy_loss.cpp new file mode 100644 index 0000000000..c9d12a5590 --- /dev/null +++ b/ggml/src/ggml-sycl/cross_entropy_loss.cpp @@ -0,0 +1,255 @@ +#include "cross_entropy_loss.hpp" + +#include +#include + +template +static __dpct_inline__ void cross_entropy_loss_f32_kernel( + const float * __restrict__ logits, + const float * __restrict__ labels, + float * __restrict__ row_loss, + const int nclasses, + const int nrows, + float * __restrict__ smem, + const sycl::nd_item<3> & item) { + + const int row = item.get_group(2); + const int tid = item.get_local_id(2); + + logits += (int64_t) row * nclasses; + labels += (int64_t) row * nclasses; + + float max_logit = -INFINITY; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float v = logits[i]; + max_logit = sycl::fmax(max_logit, v); + if (has_shared) { + smem[i] = v; + } + } + max_logit = warp_reduce_max(max_logit); + + float sum_exp = 0.0f; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float v = has_shared ? smem[i] : logits[i]; + sum_exp += sycl::exp(v - max_logit); + } + sum_exp = warp_reduce_sum(sum_exp); + const float log_sum = sycl::log(sum_exp); + + float loss = 0.0f; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float v = has_shared ? smem[i] : logits[i]; + loss += (v - max_logit - log_sum) * labels[i]; + } + loss = -warp_reduce_sum(loss) / (float) nrows; + + if (tid == 0) { + row_loss[row] = loss; + } +} + +template +static __dpct_inline__ void cross_entropy_loss_back_f32_kernel( + const float * __restrict__ grad, + const float * __restrict__ logits, + const float * __restrict__ labels, + float * __restrict__ dst, + const int nclasses, + const int nrows, + float * __restrict__ smem, + const sycl::nd_item<3> & item) { + + const int row = item.get_group(2); + const int tid = item.get_local_id(2); + + logits += (int64_t) row * nclasses; + labels += (int64_t) row * nclasses; + dst += (int64_t) row * nclasses; + + float max_logit = -INFINITY; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float v = logits[i]; + max_logit = sycl::fmax(max_logit, v); + if (has_shared) { + smem[i] = v; + } + } + max_logit = warp_reduce_max(max_logit); + + float sum_exp = 0.0f; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float v = sycl::exp((has_shared ? smem[i] : logits[i]) - max_logit); + sum_exp += v; + if (has_shared) { + smem[i] = v; + } else { + dst[i] = v; + } + } + sum_exp = warp_reduce_sum(sum_exp); + const float inv_sum = 1.0f / sum_exp; + + const float d_by_nrows = grad[0] / (float) nrows; + for (int i = tid; i < nclasses; i += WARP_SIZE) { + const float sm_num = has_shared ? smem[i] : dst[i]; + dst[i] = (sm_num * inv_sum - labels[i]) * d_by_nrows; + } +} + +static void cross_entropy_reduce_rows( + ggml_backend_sycl_context & ctx, + const float * row_loss, + float * dst, + const int64_t nrows) { + if (nrows == 1) { + SYCL_CHECK(CHECK_TRY_ERROR( + ctx.stream()->memcpy(dst, row_loss, sizeof(float)))); + return; + } + + ggml_sycl_pool_alloc tmp_alloc(ctx.pool(), nrows); + float * tmp = tmp_alloc.get(); + SYCL_CHECK(CHECK_TRY_ERROR( + ctx.stream()->memcpy(tmp, row_loss, nrows * sizeof(float)))); + + int64_t cur = nrows; + while (cur > 1) { + const int64_t out = (cur + WARP_SIZE - 1) / WARP_SIZE; + const sycl::range<3> block(1, 1, WARP_SIZE); + const sycl::range<3> grid(1, 1, out); + ctx.stream()->parallel_for( + sycl::nd_range<3>(grid * block, block), + [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + const int row = item.get_group(2); + const int tid = item.get_local_id(2); + const int64_t i = (int64_t) row * WARP_SIZE + tid; + float v = i < cur ? tmp[i] : 0.0f; + v = warp_reduce_sum(v); + if (tid == 0) { + tmp[row] = v; + } + }); + cur = out; + } + + SYCL_CHECK(CHECK_TRY_ERROR( + ctx.stream()->memcpy(dst, tmp, sizeof(float)))); +} + +void ggml_sycl_cross_entropy_loss(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_is_scalar(dst)); + + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + const int64_t nclasses = src0->ne[0]; + const int64_t nrows = ggml_nrows(src0); + + const float * logits_d = (const float *) src0->data; + const float * labels_d = (const float *) src1->data; + float * dst_d = (float *) dst->data; + + ggml_sycl_pool_alloc row_loss_alloc(ctx.pool(), nrows); + float * row_loss = row_loss_alloc.get(); + + const sycl::range<3> block(1, 1, WARP_SIZE); + const sycl::range<3> grid(1, 1, nrows); + const size_t nbytes_shared = (size_t) nclasses * sizeof(float); + const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo; + + if (nbytes_shared <= smpbo) { + ctx.stream()->submit([&](sycl::handler & cgh) { + sycl::local_accessor smem(sycl::range<1>(nclasses), cgh); + cgh.parallel_for( + sycl::nd_range<3>(grid * block, block), + [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + cross_entropy_loss_f32_kernel( + logits_d, labels_d, row_loss, + (int) nclasses, (int) nrows, + get_pointer(smem), item); + }); + }); + } else { + ctx.stream()->parallel_for( + sycl::nd_range<3>(grid * block, block), + [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + cross_entropy_loss_f32_kernel( + logits_d, labels_d, row_loss, + (int) nclasses, (int) nrows, + nullptr, item); + }); + } + + cross_entropy_reduce_rows(ctx, row_loss, dst_d, nrows); +} + +void ggml_sycl_cross_entropy_loss_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3); + + const ggml_tensor * grad = dst->src[0]; + const ggml_tensor * src0f = dst->src[1]; + const ggml_tensor * src1f = dst->src[2]; + + GGML_ASSERT(grad->type == GGML_TYPE_F32); + GGML_ASSERT(src0f->type == GGML_TYPE_F32); + GGML_ASSERT(src1f->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_scalar(grad)); + GGML_ASSERT(ggml_is_contiguous(grad)); + GGML_ASSERT(ggml_is_contiguous(src0f)); + GGML_ASSERT(ggml_is_contiguous(src1f)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0f, src1f)); + GGML_ASSERT(ggml_are_same_shape(src0f, dst)); + + SYCL_CHECK(ggml_sycl_set_device(ctx.device)); + + const int64_t nclasses = src0f->ne[0]; + const int64_t nrows = ggml_nrows(src0f); + + const float * grad_d = (const float *) grad->data; + const float * logits_d = (const float *) src0f->data; + const float * labels_d = (const float *) src1f->data; + float * dst_d = (float *) dst->data; + + const sycl::range<3> block(1, 1, WARP_SIZE); + const sycl::range<3> grid(1, 1, nrows); + const size_t nbytes_shared = (size_t) nclasses * sizeof(float); + const size_t smpbo = ggml_sycl_info().devices[ctx.device].smpbo; + + if (nbytes_shared <= smpbo) { + ctx.stream()->submit([&](sycl::handler & cgh) { + sycl::local_accessor smem(sycl::range<1>(nclasses), cgh); + cgh.parallel_for( + sycl::nd_range<3>(grid * block, block), + [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + cross_entropy_loss_back_f32_kernel( + grad_d, logits_d, labels_d, dst_d, + (int) nclasses, (int) nrows, + get_pointer(smem), item); + }); + }); + } else { + ctx.stream()->parallel_for( + sycl::nd_range<3>(grid * block, block), + [=](sycl::nd_item<3> item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] { + cross_entropy_loss_back_f32_kernel( + grad_d, logits_d, labels_d, dst_d, + (int) nclasses, (int) nrows, + nullptr, item); + }); + } +} diff --git a/ggml/src/ggml-sycl/cross_entropy_loss.hpp b/ggml/src/ggml-sycl/cross_entropy_loss.hpp new file mode 100644 index 0000000000..3f1cb817ba --- /dev/null +++ b/ggml/src/ggml-sycl/cross_entropy_loss.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include "common.hpp" + +void ggml_sycl_cross_entropy_loss(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + +void ggml_sycl_cross_entropy_loss_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 48256e2b53..c1d2eb97ec 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -74,6 +74,7 @@ #include "ggml-sycl/solve_tri.hpp" #include "ggml-sycl/gated_delta_net.hpp" #include "ggml-sycl/pool.hpp" +#include "ggml-sycl/cross_entropy_loss.hpp" #define MEM_SIZE_2M 0x00200000 #define MEM_SIZE_1G 0x40000000 @@ -5078,6 +5079,12 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_SOFT_MAX_BACK: ggml_sycl_op_soft_max_back(ctx, dst); break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_sycl_cross_entropy_loss(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + ggml_sycl_cross_entropy_loss_back(ctx, dst); + break; case GGML_OP_ROPE: ggml_sycl_rope(ctx, dst); break; @@ -5892,6 +5899,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_FILL: case GGML_OP_CUMSUM: case GGML_OP_DIAG: + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: return true; case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;