Mirrors/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggerganov/llama.cpp.git synced 2025-01-26 03:12:23 +01:00
Code Issues Packages Projects Releases Wiki Activity

ggml : multi-thread mul and diag_mask ops (#1428)

Browse Source
This commit is contained in:
Georgi Gerganov 2023-05-13 16:48:03 +03:00 committed by GitHub
parent 905d87b70a
commit 66841fdb0e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
ggml.c
View File

@ -7765,12 +7765,13 @@ static void ggml_compute_forward_mul_f32(
const struct ggml_tensor * src0,
const struct ggml_tensor * src1,
struct ggml_tensor * dst) {
assert(params->ith == 0);
assert(ggml_are_same_shape(src0, src1) && ggml_are_same_shape(src0, dst));
if (params->type == GGML_TASK_INIT || params->type == GGML_TASK_FINALIZE) {
return;
}
const int ith = params->ith;
const int nth = params->nth;
const int nr = ggml_nrows(src0);
const int64_t ne0 = src0->ne[0];
@ -7796,7 +7797,7 @@ static void ggml_compute_forward_mul_f32(
GGML_ASSERT(nb00 == sizeof(float));
if (nb10 == sizeof(float)) {
for (int ir = 0; ir < nr; ++ir) {
for (int ir = ith; ir < nr; ir += nth) {
// src0, src1 and dst are same shape => same indices
const int i3 = ir/(ne2*ne1);
const int i2 = (ir - i3*ne2*ne1)/ne1;
@ -7822,7 +7823,7 @@ static void ggml_compute_forward_mul_f32(
}
} else {
// src1 is not contiguous
for (int ir = 0; ir < nr; ++ir) {
for (int ir = ith; ir < nr; ir += nth) {
// src0, src1 and dst are same shape => same indices
const int i3 = ir/(ne2*ne1);
const int i2 = (ir - i3*ne2*ne1)/ne1;
@ -10317,7 +10318,6 @@ static void ggml_compute_forward_diag_mask_f32(
const struct ggml_tensor * src1,
struct ggml_tensor * dst,
const float value) {
assert(params->ith == 0);
assert(src1->type == GGML_TYPE_I32);
assert(ggml_nelements(src1) == 2);
@ -10325,6 +10325,9 @@ static void ggml_compute_forward_diag_mask_f32(
return;
}
const int ith = params->ith;
const int nth = params->nth;
const int n_past = ((int32_t *) src1->data)[0];
const bool inplace = (bool)((int32_t *) src1->data)[1];
@ -10343,7 +10346,7 @@ static void ggml_compute_forward_diag_mask_f32(
assert(src0->nb[0] == sizeof(float));
for (int k = 0; k < nz; k++) {
for (int j = 0; j < nr; j++) {
for (int j = ith; j < nr; j += nth) {
for (int i = n_past; i < nc; i++) {
if (i > n_past + j) {
*(float *)((char *) dst->data + k*dst->nb[2] + j*dst->nb[1] + i*dst->nb[0]) = value;
@ -13609,7 +13612,6 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
work_size = MAX(work_size, cur);
} break;
case GGML_OP_SUB:
case GGML_OP_MUL:
case GGML_OP_DIV:
case GGML_OP_SQR:
case GGML_OP_SQRT:
@ -13626,18 +13628,10 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
{
node->n_tasks = 1;
} break;
case GGML_OP_MUL:
case GGML_OP_GELU:
{
node->n_tasks = n_threads;
} break;
case GGML_OP_SILU:
{
node->n_tasks = n_threads;
} break;
case GGML_OP_SILU_BACK:
{
node->n_tasks = n_threads;
} break;
case GGML_OP_NORM:
case GGML_OP_RMS_NORM:
case GGML_OP_RMS_NORM_BACK:
@ -13715,11 +13709,11 @@ void ggml_graph_compute(struct ggml_context * ctx, struct ggml_cgraph * cgraph)
case GGML_OP_GET_ROWS:
case GGML_OP_GET_ROWS_BACK:
case GGML_OP_DIAG:
case GGML_OP_DIAG_MASK_INF:
case GGML_OP_DIAG_MASK_ZERO:
{
node->n_tasks = 1;
} break;
case GGML_OP_DIAG_MASK_INF:
case GGML_OP_SOFT_MAX:
case GGML_OP_ROPE:
case GGML_OP_ROPE_BACK: