Skip to content

Commit 17e98d4

Browse files
NeoZhangJianyuNeoZhangJianyu
authored
fix mul_mat_id() for new input, make the ut pass (#6682)
1 parent 1958f7e commit 17e98d4

File tree

1 file changed

+50
-46
lines changed

1 file changed

+50
-46
lines changed

‎ggml-sycl.cpp

Lines changed: 50 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -15996,73 +15996,76 @@ static void ggml_sycl_mul_mat_id_sycl(ggml_tensor * dst) {
1599615996
static void ggml_sycl_mul_mat_id(const ggml_tensor *src0,
1599715997
const ggml_tensor *src1,
1599815998
ggml_tensor *dst) try {
15999-
#if 0
16000-
ggml_sycl_mul_mat_id_sycl(dst);
16001-
// TODO: mmq/mmv support
16002-
#endif
15999+
GGML_ASSERT(src0->backend != GGML_BACKEND_TYPE_GPU_SPLIT &&
16000+
"mul_mat_id does not support split buffers");
16001+
const ggml_tensor *ids = dst->src[2];
16002+
const dpct::queue_ptr stream = g_syclStreams[g_main_device][0];
1600316003

16004-
const int64_t nb11 = src1->nb[1];
16005-
const int64_t nb1 = dst->nb[1];
16004+
const size_t nb11 = src1->nb[1];
16005+
const size_t nb1 = dst->nb[1];
1600616006

16007-
const struct ggml_tensor * ids = src0;
16008-
const int32_t id = ((int32_t *) dst->op_params)[0];
16009-
const int32_t n_as = ((int32_t *) dst->op_params)[1];
16007+
const int32_t id = ((int32_t *)dst->op_params)[0];
16008+
const int32_t n_as = src0->ne[2];
1601016009

1601116010
std::vector<char> ids_host(ggml_nbytes(ids));
16011+
const char *ids_dev = (const char *)ids->data;
1601216012

16013-
const dpct::queue_ptr stream = g_syclStreams[g_main_device][0];
16014-
16015-
if (ids->backend == GGML_BACKEND_TYPE_GPU) {
16016-
const char * ids_dev = (const char *)((const ggml_tensor_extra_gpu *)ids->extra)->data_device[g_main_device];
16017-
SYCL_CHECK(CHECK_TRY_ERROR(
16018-
stream->memcpy(ids_host.data(), ids_dev, ggml_nbytes(ids)).wait()));
16019-
// SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
16020-
} else {
16021-
memcpy(ids_host.data(), ids->data, ggml_nbytes(ids));
16022-
}
16013+
SYCL_CHECK(CHECK_TRY_ERROR(
16014+
stream->memcpy(ids_host.data(), ids_dev, ggml_nbytes(ids))));
16015+
SYCL_CHECK(CHECK_TRY_ERROR(stream->wait()));
1602316016

16024-
const ggml_tensor_extra_gpu * src1_extra = (const ggml_tensor_extra_gpu *) src1->extra;
16025-
const ggml_tensor_extra_gpu * dst_extra = (const ggml_tensor_extra_gpu *) dst->extra;
16017+
const ggml_tensor_extra_gpu *src0_extra =
16018+
(const ggml_tensor_extra_gpu *)src0->extra;
16019+
const ggml_tensor_extra_gpu *src1_extra =
16020+
(const ggml_tensor_extra_gpu *)src1->extra;
16021+
const ggml_tensor_extra_gpu *dst_extra =
16022+
(const ggml_tensor_extra_gpu *)dst->extra;
1602616023

16024+
ggml_tensor_extra_gpu src0_row_extra;
1602716025
ggml_tensor_extra_gpu src1_row_extra;
1602816026
ggml_tensor_extra_gpu dst_row_extra;
1602916027

16028+
ggml_tensor src0_row = *src0;
1603016029
ggml_tensor src1_row = *src1;
1603116030
ggml_tensor dst_row = *dst;
1603216031

1603316032
src1_row.backend = GGML_BACKEND_TYPE_GPU;
1603416033
dst_row.backend = GGML_BACKEND_TYPE_GPU;
1603516034

16035+
src0_row.extra = &src0_row_extra;
1603616036
src1_row.extra = &src1_row_extra;
1603716037
dst_row.extra = &dst_row_extra;
1603816038

16039-
char * src1_original = src1->backend == GGML_BACKEND_TYPE_CPU ?
16040-
(char *) src1->data : (char *) src1_extra->data_device[g_main_device];
16041-
char * dst_original = dst->backend == GGML_BACKEND_TYPE_CPU ?
16042-
(char *) dst->data : (char *) dst_extra->data_device[g_main_device];
16039+
char *src0_original = src1->backend == GGML_BACKEND_TYPE_CPU
16040+
? (char *)src0->data
16041+
: (char *)src0_extra->data_device[g_main_device];
16042+
char *src1_original = src1->backend == GGML_BACKEND_TYPE_CPU
16043+
? (char *)src1->data
16044+
: (char *)src1_extra->data_device[g_main_device];
16045+
char *dst_original = dst->backend == GGML_BACKEND_TYPE_CPU
16046+
? (char *)dst->data
16047+
: (char *)dst_extra->data_device[g_main_device];
1604316048

16044-
if (src1->ne[1] == 1) {
16045-
GGML_ASSERT(src1->backend == GGML_BACKEND_TYPE_GPU);
16046-
GGML_ASSERT(dst->backend == GGML_BACKEND_TYPE_GPU);
16049+
src0_row.ne[2] = 1;
16050+
src0_row.ne[3] = 1;
16051+
src0_row.nb[3] = src0->nb[2];
1604716052

16053+
if (src1->ne[1] == 1) {
1604816054
for (int64_t i01 = 0; i01 < ids->ne[1]; i01++) {
16049-
//int32_t row_id;
16050-
//SYCL_CHECK(syclMemcpyAsync(&row_id, ids_dev + i01*ids->nb[1] + id*ids->nb[0], sizeof(int32_t), syclMemcpyDeviceToHost, g_syclStreams[g_main_device][0]));
16051-
//SYCL_CHECK(syclStreamSynchronize(g_syclStreams[g_main_device][0]));
16052-
16053-
const int32_t row_id = *(const int32_t *) (ids_host.data() + i01*ids->nb[1] + id*ids->nb[0]);
16055+
const int32_t row_id =
16056+
*(const int32_t *)(ids_host.data() + i01 * ids->nb[1] +
16057+
id * ids->nb[0]);
1605416058

1605516059
GGML_ASSERT(row_id >= 0 && row_id < n_as);
1605616060

16057-
const struct ggml_tensor * src0_row = dst->src[row_id + 2];
16061+
src0_row_extra.data_device[g_main_device] =
16062+
src0_original + row_id * src0->nb[2];
16063+
src1_row_extra.data_device[g_main_device] =
16064+
src1_original + i01 * src1->nb[1];
16065+
dst_row_extra.data_device[g_main_device] =
16066+
dst_original + i01 * dst->nb[1];
1605816067

16059-
src1_row_extra.data_device[g_main_device] = src1_original + i01*src1->nb[1];
16060-
src1_row.data = (char *) src1->data + i01*src1->nb[1]; // TODO why is this set?
16061-
16062-
dst_row_extra.data_device[g_main_device] = dst_original + i01*dst->nb[1];
16063-
dst_row.data = (char *) dst->data + i01*dst->nb[1]; // TODO why is this set?
16064-
16065-
ggml_sycl_mul_mat(src0_row, &src1_row, &dst_row);
16068+
ggml_sycl_mul_mat(&src0_row, &src1_row, &dst_row);
1606616069
}
1606716070
} else {
1606816071
sycl_pool_alloc<char> src1_contiguous(sizeof(float)*ggml_nelements(src1));
@@ -16072,8 +16075,6 @@ static void ggml_sycl_mul_mat_id(const ggml_tensor *src0,
1607216075
dst_row_extra.data_device[g_main_device] = dst_contiguous.get();
1607316076

1607416077
for (int32_t row_id = 0; row_id < n_as; ++row_id) {
16075-
const struct ggml_tensor * src0_row = dst->src[row_id + 2];
16076-
1607716078
int64_t num_src1_rows = 0;
1607816079
for (int64_t i01 = 0; i01 < ids->ne[1]; i01++) {
1607916080
const int32_t row_id_i = *(const int32_t *) (ids_host.data() + i01*ids->nb[1] + id*ids->nb[0]);
@@ -16086,14 +16087,17 @@ static void ggml_sycl_mul_mat_id(const ggml_tensor *src0,
1608616087

1608716088
SYCL_CHECK(CHECK_TRY_ERROR(
1608816089
stream->memcpy(src1_contiguous.get() + num_src1_rows * nb11,
16089-
src1_original + i01 * nb11, nb11).wait()));
16090+
src1_original + i01 * nb11, nb11)));
1609016091
num_src1_rows++;
1609116092
}
1609216093

1609316094
if (num_src1_rows == 0) {
1609416095
continue;
1609516096
}
1609616097

16098+
src0_row_extra.data_device[g_main_device] =
16099+
src0_original + row_id * src0->nb[2];
16100+
1609716101
src1_row.ne[1] = num_src1_rows;
1609816102
dst_row.ne[1] = num_src1_rows;
1609916103

@@ -16105,7 +16109,7 @@ static void ggml_sycl_mul_mat_id(const ggml_tensor *src0,
1610516109
dst_row.nb[2] = num_src1_rows*nb1;
1610616110
dst_row.nb[3] = num_src1_rows*nb1;
1610716111

16108-
ggml_sycl_mul_mat(src0_row, &src1_row, &dst_row);
16112+
ggml_sycl_mul_mat(&src0_row, &src1_row, &dst_row);
1610916113

1611016114
num_src1_rows = 0;
1611116115
for (int64_t i01 = 0; i01 < ids->ne[1]; i01++) {
@@ -16119,7 +16123,7 @@ static void ggml_sycl_mul_mat_id(const ggml_tensor *src0,
1611916123

1612016124
SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(
1612116125
dst_original + i01 * nb1,
16122-
dst_contiguous.get() + num_src1_rows * nb1, nb1).wait()));
16126+
dst_contiguous.get() + num_src1_rows * nb1, nb1)));
1612316127
num_src1_rows++;
1612416128
}
1612516129
}

0 commit comments

Comments
 (0)