diff --git a/ggml/src/ggml-sycl/getrows.cpp b/ggml/src/ggml-sycl/getrows.cpp index 2113f3563..4e84bd222 100644 --- a/ggml/src/ggml-sycl/getrows.cpp +++ b/ggml/src/ggml-sycl/getrows.cpp @@ -245,6 +245,84 @@ static void get_rows_sycl_float(ggml_backend_sycl_context & ctx, const ggml_tens GGML_UNUSED(ctx); } +template +static void k_get_rows_back_float(const src0_t * src0, const int32_t * src1, float * dst, + const int64_t ncols, const int64_t nrows_grad_10, const int64_t nrows_grad_11, const int64_t nrows_dst, + const size_t s01, const size_t s02, + const size_t s10, const size_t s11, + const size_t s1, + const int64_t block_num_y, + const sycl::nd_item<3> & item_ct1) { + const int64_t col = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2); + if (col >= ncols) { + return; + } + + // block_num_y is clamped, so stride over destination rows like CUDA k_get_rows_back_float + for (int64_t dst_row = item_ct1.get_group(1); dst_row < nrows_dst; dst_row += block_num_y) { + float sum = 0.0f; + + const int64_t nrows_grad_total = nrows_grad_10 * nrows_grad_11; + for (int64_t i = 0; i < nrows_grad_total; ++i) { + const int64_t i10 = i % nrows_grad_10; + const int64_t i11 = i / nrows_grad_10; + if (src1[i10*s10 + i11*s11] != dst_row) { + continue; + } + sum += (float) src0[col + i10*s01 + i11*s02]; + } + + dst[col + dst_row*s1] = sum; + } +} + +template +static void get_rows_back_sycl_float(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, + const ggml_tensor * src1, ggml_tensor * dst, + const src0_t * src0_dd, const int32_t * src1_dd, + float * dst_dd, queue_ptr stream) { + + GGML_TENSOR_BINARY_OP_LOCALS + + GGML_ASSERT(ne02*ne03 == 1); + GGML_ASSERT(ne12*ne13 == 1); + GGML_ASSERT(ne2*ne3 == 1); + GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); + GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); + GGML_ASSERT(dst->nb[0] == ggml_type_size(dst->type)); + + const int64_t ncols = ne00; + const int64_t nrows_grad_10 = ne10; + const int64_t nrows_grad_11 = ne11; + const int64_t nrows_dst = ne1; + + const size_t s01 = nb01 / sizeof(src0_t); + const size_t s02 = nb02 / sizeof(src0_t); + + const size_t s10 = nb10 / sizeof(int32_t); + const size_t s11 = nb11 / sizeof(int32_t); + + const size_t s1 = nb1 / sizeof(float); + + const sycl::range<3> block_dims(1, 1, SYCL_GET_ROWS_BLOCK_SIZE); + const int64_t block_num_x = (ncols + SYCL_GET_ROWS_BLOCK_SIZE - 1) / SYCL_GET_ROWS_BLOCK_SIZE; + const int64_t block_num_y = std::min(nrows_dst, (int64_t) UINT16_MAX); + const sycl::range<3> block_nums(1, block_num_y, block_num_x); + + stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims), + [=](sycl::nd_item<3> item_ct1) { + k_get_rows_back_float(src0_dd, src1_dd, dst_dd, + ncols, nrows_grad_10, nrows_grad_11, nrows_dst, + s01, s02, s10, s11, s1, + block_num_y, item_ct1); + }); + + GGML_UNUSED(src0); + GGML_UNUSED(src1); + GGML_UNUSED(dst); + GGML_UNUSED(ctx); +} + void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { GGML_ASSERT(dst->src[1]->type == GGML_TYPE_I32); GGML_ASSERT(dst->type == GGML_TYPE_F32 || dst->type == GGML_TYPE_I32 ); @@ -366,3 +444,30 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { GGML_ABORT("fatal error"); } } + +void ggml_sycl_op_get_rows_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(dst)); + + switch (src0->type) { + case GGML_TYPE_F16: + get_rows_back_sycl_float(ctx, src0, src1, dst, (const sycl::half *) src0->data, + (const int32_t *) src1->data, (float *) dst->data, + ctx.stream()); + break; + case GGML_TYPE_F32: + get_rows_back_sycl_float(ctx, src0, src1, dst, (const float *) src0->data, + (const int32_t *) src1->data, (float *) dst->data, + ctx.stream()); + break; + default: + GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0->type)); + break; + } +} diff --git a/ggml/src/ggml-sycl/getrows.hpp b/ggml/src/ggml-sycl/getrows.hpp index 1c560cd9f..0388e6c7f 100644 --- a/ggml/src/ggml-sycl/getrows.hpp +++ b/ggml/src/ggml-sycl/getrows.hpp @@ -16,5 +16,6 @@ #include "common.hpp" void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor *dst); +void ggml_sycl_op_get_rows_back(ggml_backend_sycl_context & ctx, ggml_tensor *dst); #endif // GGML_SYCL_GETROWS_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index e599d2d84..f1fe7eea3 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -3638,6 +3638,11 @@ static void ggml_sycl_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * ds ggml_sycl_op_get_rows(ctx, dst); } +static void ggml_sycl_get_rows_back(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { + scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); + ggml_sycl_op_get_rows_back(ctx, dst); +} + static void ggml_sycl_norm(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1); ggml_sycl_op_norm(ctx, dst); @@ -5437,6 +5442,9 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg case GGML_OP_GET_ROWS: ggml_sycl_get_rows(ctx, dst); break; + case GGML_OP_GET_ROWS_BACK: + ggml_sycl_get_rows_back(ctx, dst); + break; case GGML_OP_SET: ggml_sycl_op_set(ctx, dst); break; @@ -6463,6 +6471,12 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons return false; } } + case GGML_OP_GET_ROWS_BACK: + // return true; + return op->type == GGML_TYPE_F32 && + (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && + op->src[1]->type == GGML_TYPE_I32 && + op->ne[2] == 1 && op->ne[3] == 1; case GGML_OP_SET: return (op->type == GGML_TYPE_F32) && (op->src[0] && op->src[1]) &&