diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 491316f74912..5a29dbfe9f48 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1161,9 +1161,72 @@ void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, // ggml_compute_forward_mul_mat +// Optional element-wise epilogue applied to the matmul result before it is +// written out, i.e. dst = op(mul_mat(src0, src1), other). `dst` and `other` +// have the same shape as the matmul destination and float rows. See +// ggml_compute_forward_mul_mat_fused(). +enum ggml_mul_mat_epilogue_op { + GGML_MUL_MAT_EPILOGUE_ADD, + GGML_MUL_MAT_EPILOGUE_MUL, + GGML_MUL_MAT_EPILOGUE_SILU, +}; + +struct ggml_mul_mat_epilogue { + enum ggml_mul_mat_epilogue_op op; + struct ggml_tensor * dst; + const struct ggml_tensor * other; // NULL for unary epilogues +}; + +// Apply the epilogue as a separate pass over an already materialized matmul +// result. This is the fallback for kernels that write `dst` themselves (the +// extra-buffer "accelerator" kernels and the tinyBLAS/GGML_LLAMAFILE path), +// where the epilogue cannot be folded into the accumulator tile. The result is +// identical to running the element-wise op as its own graph node, and the node +// itself is still saved. +static void ggml_compute_forward_mul_mat_epilogue_pass( + const struct ggml_compute_params * params, + const struct ggml_tensor * dst, + const struct ggml_mul_mat_epilogue * epilogue) { + + const int64_t ne0 = dst->ne[0]; + const int64_t ne1 = dst->ne[1]; + const int64_t ne2 = dst->ne[2]; + + const int64_t nr = ggml_nrows(dst); + + // rows per thread + const int64_t dr = (nr + params->nth - 1) / params->nth; + + const int64_t ir_start = dr * params->ith; + const int64_t ir_end = MIN(ir_start + dr, nr); + + for (int64_t ir = ir_start; ir < ir_end; ++ir) { + const int64_t i3 = ir / (ne2 * ne1); + const int64_t i2 = (ir - i3 * ne2 * ne1) / ne1; + const int64_t i1 = (ir - i3 * ne2 * ne1 - i2 * ne1); + + const float * mm_src = (const float *) ((const char *) dst->data + (i1 * dst->nb[1] + i2 * dst->nb[2] + i3 * dst->nb[3])); + + float * ep_dst = (float *) ((char *) epilogue->dst->data + (i1 * epilogue->dst->nb[1] + i2 * epilogue->dst->nb[2] + i3 * epilogue->dst->nb[3])); + + if (epilogue->op == GGML_MUL_MAT_EPILOGUE_SILU) { + ggml_vec_silu_f32(ne0, ep_dst, mm_src); + } else { + const float * ep_src = (const float *) ((const char *) epilogue->other->data + (i1 * epilogue->other->nb[1] + i2 * epilogue->other->nb[2] + i3 * epilogue->other->nb[3])); + + if (epilogue->op == GGML_MUL_MAT_EPILOGUE_ADD) { + ggml_vec_add_f32(ne0, ep_dst, mm_src, ep_src); + } else { + ggml_vec_mul_f32(ne0, ep_dst, mm_src, ep_src); + } + } + } +} + static void ggml_compute_forward_mul_mat_one_chunk( const struct ggml_compute_params * params, struct ggml_tensor * dst, + const struct ggml_mul_mat_epilogue * add_epilogue, const enum ggml_type type, const int64_t num_rows_per_vec_dot, const int64_t ir0_start, @@ -1243,23 +1306,51 @@ static void ggml_compute_forward_mul_mat_one_chunk( vec_dot(ne00, &tmp[ir0 - iir0], (num_rows_per_vec_dot > 1 ? 16 : 0), src0_row + ir0 * nb01, (num_rows_per_vec_dot > 1 ? nb01 : 0), src1_col, (num_rows_per_vec_dot > 1 ? src1_col_stride : 0), num_rows_per_vec_dot); } - for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) { - memcpy(&dst_col[iir0 + cn * nb1 / nb0], tmp + (cn * 16), (MIN(iir0 + blck_0, ir0_end) - iir0) * sizeof(float)); + const int64_t nc = MIN(iir0 + blck_0, ir0_end) - iir0; + + if (add_epilogue) { + // fused element-wise epilogue: apply the following op straight + // out of the accumulator tile instead of materializing the + // matmul result and reading it back in a second pass + for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) { + const int64_t ia1 = i1 + cn; + + float * ep_dst = (float *) ((char *) add_epilogue->dst->data + (ia1 * add_epilogue->dst->nb[1] + i2 * add_epilogue->dst->nb[2] + i3 * add_epilogue->dst->nb[3])); + + if (add_epilogue->op == GGML_MUL_MAT_EPILOGUE_SILU) { + ggml_vec_silu_f32(nc, ep_dst + iir0, tmp + (cn * 16)); + } else { + const float * ep_src = (const float *) ((char *) add_epilogue->other->data + (ia1 * add_epilogue->other->nb[1] + i2 * add_epilogue->other->nb[2] + i3 * add_epilogue->other->nb[3])); + + if (add_epilogue->op == GGML_MUL_MAT_EPILOGUE_ADD) { + ggml_vec_add_f32(nc, ep_dst + iir0, tmp + (cn * 16), ep_src + iir0); + } else { + ggml_vec_mul_f32(nc, ep_dst + iir0, tmp + (cn * 16), ep_src + iir0); + } + } + } + } else { + for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) { + memcpy(&dst_col[iir0 + cn * nb1 / nb0], tmp + (cn * 16), nc * sizeof(float)); + } } } } } } -void ggml_compute_forward_mul_mat( +static void ggml_compute_forward_mul_mat_impl( const struct ggml_compute_params * params, - struct ggml_tensor * dst) { + struct ggml_tensor * dst, + const struct ggml_mul_mat_epilogue * add_epilogue) { const struct ggml_tensor * src0 = dst->src[0]; const struct ggml_tensor * src1 = dst->src[1]; const int32_t hint = ggml_get_op_params_i32(dst, 1); if (hint == GGML_HINT_SRC0_IS_HADAMARD && !params->use_ref) { + // the fused epilogue is only wired into the generic path below + GGML_ASSERT(add_epilogue == NULL); ggml_compute_forward_fwht(params, dst); return; } @@ -1314,6 +1405,12 @@ void ggml_compute_forward_mul_mat( src1->type, dst->type)) goto UseGgmlGemm1; + if (add_epilogue) { + // tinyBLAS wrote dst itself, so the epilogue could not be folded + // into the accumulator tile - run it as a second pass instead + ggml_barrier(params->threadpool); + ggml_compute_forward_mul_mat_epilogue_pass(params, dst, add_epilogue); + } return; } UseGgmlGemm1:; @@ -1382,6 +1479,12 @@ UseGgmlGemm1:; vec_dot_type, dst->type)) goto UseGgmlGemm2; + if (add_epilogue) { + // tinyBLAS wrote dst itself, so the epilogue could not be folded + // into the accumulator tile - run it as a second pass instead + ggml_barrier(params->threadpool); + ggml_compute_forward_mul_mat_epilogue_pass(params, dst, add_epilogue); + } return; } UseGgmlGemm2:; @@ -1441,7 +1544,7 @@ UseGgmlGemm2:; if ((nr0 % 2 != 0) || (ne11 % 2 != 0) || ((ir0_end - ir0_start) % 2 != 0) || ((ir1_end - ir1_start) % 2 != 0)) { num_rows_per_vec_dot = 1; } - ggml_compute_forward_mul_mat_one_chunk(params, dst, src0->type, num_rows_per_vec_dot, ir0_start, ir0_end, ir1_start, ir1_end); + ggml_compute_forward_mul_mat_one_chunk(params, dst, add_epilogue, src0->type, num_rows_per_vec_dot, ir0_start, ir0_end, ir1_start, ir1_end); if (nth >= nchunk0 * nchunk1) { break; @@ -1451,6 +1554,51 @@ UseGgmlGemm2:; } } +void ggml_compute_forward_mul_mat( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + + ggml_compute_forward_mul_mat_impl(params, dst, /*add_epilogue =*/ NULL); +} + +// Fused MUL_MAT + element-wise op: computes dst_ep = op(mul_mat(src0, src1), ...) +// in a single pass. The epilogue is applied straight out of the accumulator +// tile, so the intermediate matmul output is never written to (and read back +// from) memory, and the graph loses one node - and with it one global thread +// barrier. +void ggml_compute_forward_mul_mat_fused( + struct ggml_compute_params * params, + struct ggml_tensor * dst_mul_mat, + struct ggml_tensor * dst_ep) { + + GGML_ASSERT(dst_ep != NULL); + GGML_ASSERT(dst_ep->src[0] == dst_mul_mat || dst_ep->src[1] == dst_mul_mat); + + struct ggml_mul_mat_epilogue epilogue = { + /*.op =*/ GGML_MUL_MAT_EPILOGUE_SILU, + /*.dst =*/ dst_ep, + /*.other =*/ NULL, + }; + + if (dst_ep->op == GGML_OP_ADD || dst_ep->op == GGML_OP_MUL) { + epilogue.op = dst_ep->op == GGML_OP_ADD ? GGML_MUL_MAT_EPILOGUE_ADD : GGML_MUL_MAT_EPILOGUE_MUL; + epilogue.other = (dst_ep->src[0] == dst_mul_mat) ? dst_ep->src[1] : dst_ep->src[0]; + } else { + GGML_ASSERT(dst_ep->op == GGML_OP_UNARY && ggml_get_unary_op(dst_ep) == GGML_UNARY_OP_SILU); + } + + // extra_buffer op? those kernels (repack, AMX, KleidiAI, ...) own the whole + // matmul - including its work buffer layout - so they must keep handling it. + // They write dst themselves, so the epilogue runs as a second pass. + if (ggml_cpu_extra_compute_forward(params, dst_mul_mat)) { + ggml_barrier(params->threadpool); + ggml_compute_forward_mul_mat_epilogue_pass(params, dst_mul_mat, &epilogue); + return; + } + + ggml_compute_forward_mul_mat_impl(params, dst_mul_mat, &epilogue); +} + // ggml_compute_forward_mul_mat_id #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id)*ids->ne[0]*ids->ne[1] + (i1)] @@ -3026,7 +3174,7 @@ static bool ggml_cpu_disable_fusion = false; // initialized once in ggml_cpu_in static int ggml_cpu_try_fuse_ops( const struct ggml_cgraph * cgraph, const int node_n, - const struct ggml_compute_params * params, + struct ggml_compute_params * params, const struct ggml_cplan * cplan) { if (ggml_cpu_disable_fusion || cplan->use_ref) { @@ -3035,6 +3183,13 @@ static int ggml_cpu_try_fuse_ops( struct ggml_tensor * node = cgraph->nodes[node_n]; + // empty tensors are skipped by ggml_compute_forward(), so the fused kernels + // (which are called instead of it) must not see them - their strides do not + // satisfy the usual layout preconditions + if (ggml_is_empty(node)) { + return 0; + } + if (node->op == GGML_OP_RMS_NORM) { // RMS_NORM + MUL fusion const enum ggml_op fuse_ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL }; @@ -3054,6 +3209,40 @@ static int ggml_cpu_try_fuse_ops( } } + if (node->op == GGML_OP_MUL_MAT && + ggml_get_op_params_i32(node, 1) != GGML_HINT_SRC0_IS_HADAMARD && + node_n + 1 < cgraph->n_nodes) { + // MUL_MAT + element-wise epilogue fusion. In a transformer layer this + // covers the two residual adds (MUL_MAT + ADD), the SwiGLU gate + // (MUL_MAT + SILU) and the gate/up product (MUL_MAT + MUL). + struct ggml_tensor * ep_node = cgraph->nodes[node_n + 1]; + + const bool is_binary = ep_node->op == GGML_OP_ADD || ep_node->op == GGML_OP_MUL; + const bool is_silu = ep_node->op == GGML_OP_UNARY && ggml_get_unary_op(ep_node) == GGML_UNARY_OP_SILU; + + const enum ggml_op fuse_ops[] = { GGML_OP_MUL_MAT, ep_node->op }; + + if ((is_binary || is_silu) && ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) { + // the epilogue writes whole float rows of the destination, so it needs + // identically shaped, non-broadcast, row-contiguous operands + const struct ggml_tensor * other = is_binary + ? ((ep_node->src[0] == node) ? ep_node->src[1] : ep_node->src[0]) + : NULL; + + if (ep_node->type == GGML_TYPE_F32 && + ep_node->nb[0] == sizeof(float) && + !ggml_is_transposed(ep_node) && + (!is_binary || (other->type == GGML_TYPE_F32 && + other->nb[0] == sizeof(float) && + ggml_are_same_shape(other, node) && + !ggml_is_transposed(other)))) { + + ggml_compute_forward_mul_mat_fused(params, node, ep_node); + return 1; + } + } + } + return 0; } diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4c1642a67603..32e0e5df5f4e 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -117,6 +117,7 @@ void ggml_compute_forward_cross_entropy_loss(const struct ggml_compute_params * void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat_fused(struct ggml_compute_params * params, struct ggml_tensor * dst_mul_mat, struct ggml_tensor * dst_ep); void ggml_compute_forward_fwht(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus