Skip to content

Commit 2dcc02d

Browse files
am17anam17an
committed
Non-broadcast fuse works
1 parent fbbd94c commit 2dcc02d

File tree

4 files changed

+123
-1
lines changed

4 files changed

+123
-1
lines changed

‎ggml/src/ggml-cuda/binbcast.cu

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include "binbcast.cuh"
2+
#include <driver_types.h>
23
#include <cstdint>
34

45
static __device__ __forceinline__ float op_repeat(const float a, const float b) {
@@ -56,6 +57,37 @@ static __global__ void k_bin_bcast(const src0_t * src0, const src1_t * src1, dst
5657
}
5758
}
5859

60+
template<typename src_t, typename dst_t>
61+
static __global__ void k_fused_add(const src_t ** src, int n_srcs, dst_t * dst,
62+
const int ne0, const int ne1, const int ne2, const int ne3,
63+
const int s1, const int s2, const int s3) {
64+
65+
const int i0s = blockDim.x*blockIdx.x + threadIdx.x;
66+
const int i1 = (blockDim.y*blockIdx.y + threadIdx.y);
67+
const int i2 = (blockDim.z*blockIdx.z + threadIdx.z) / ne3;
68+
const int i3 = (blockDim.z*blockIdx.z + threadIdx.z) % ne3;
69+
70+
if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3) {
71+
return;
72+
}
73+
74+
const size_t i_src0 = i3*s3 + i2*s2 + i1*s1;
75+
const size_t i_src1 = i3*s3 + i2*s2 + i1*s1;
76+
const size_t i_src2 = i3*s3 + i2*s2 + i1*s1;
77+
const size_t i_dst = i3*s3 + i2*s2 + i1*s1;
78+
79+
dst_t * dst_row = dst + i_dst;
80+
81+
for (int i0 = i0s; i0 < ne0; i0 += blockDim.x*gridDim.x) {
82+
float sum = 0.;
83+
for (int i = 0 ; i < n_srcs; ++i) {
84+
const src_t * src_row = src[i] + i_dst; // use same offset as dst
85+
sum += (float)src_row[i0];
86+
}
87+
dst_row[i0] = (dst_t)sum;
88+
}
89+
}
90+
5991
template<float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t>
6092
static __global__ void k_bin_bcast_unravel(const src0_t * src0, const src1_t * src1, dst_t * dst,
6193
int ne0, int ne1, int ne2, int ne3,
@@ -331,6 +363,51 @@ void ggml_cuda_op_div(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
331363
ggml_cuda_op_bin_bcast<bin_bcast_cuda<op_div>>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream());
332364
}
333365

366+
367+
void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse) {
368+
369+
printf("ggml_cuda_op_fused_add: %d\n", n_fuse);
370+
371+
cudaStream_t stream = ctx.stream();
372+
373+
GGML_ASSERT(1 <= n_fuse && n_fuse <= 8);
374+
375+
// Collect device pointers on host
376+
const float * h_src[n_fuse + 1];
377+
for(int i = 0 ; i < n_fuse + 1; ++i) {
378+
h_src[i] = (const float*)dst->src[i]->data;
379+
}
380+
381+
// Allocate device array for pointers and copy
382+
const float ** d_src;
383+
cudaMalloc((void **) &d_src, (n_fuse + 1) * sizeof(float *));
384+
cudaMemcpy(d_src, h_src, (n_fuse + 1) * sizeof(float *), cudaMemcpyHostToDevice);
385+
386+
//All layouts are same in the fused ops
387+
const int ne0 = dst->ne[0];
388+
const int ne1 = dst->ne[1];
389+
const int ne2 = dst->ne[2];
390+
const int ne3 = dst->ne[3];
391+
const int s1 = dst->nb[1] / sizeof(float);
392+
const int s2 = dst->nb[2] / sizeof(float);
393+
const int s3 = dst->nb[3] / sizeof(float);
394+
395+
const int block_size = 128;
396+
dim3 block_dims;
397+
block_dims.x = std::min<unsigned int>(ne0, block_size);
398+
block_dims.y = std::min<unsigned int>(ne1, block_size / block_dims.x);
399+
block_dims.z = std::min(std::min<unsigned int>(ne2 * ne3, block_size / block_dims.x / block_dims.y), 64U);
400+
401+
dim3 block_nums((ne0 + block_dims.x - 1) / block_dims.x,
402+
(ne1 + block_dims.y - 1) / block_dims.y,
403+
(ne2 * ne3 + block_dims.z - 1) / block_dims.z);
404+
405+
k_fused_add<float, float>
406+
<<<block_nums, block_dims, 0, stream>>>(d_src, n_fuse + 1, (float *) dst->data, ne0, ne1, ne2, ne3, s1, s2, s3);
407+
408+
cudaFree(d_src);
409+
}
410+
334411
void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
335412
const ggml_tensor * src0 = dst->src[0];
336413

‎ggml/src/ggml-cuda/binbcast.cuh

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,6 @@ void ggml_cuda_op_sub(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
66
void ggml_cuda_op_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
77
void ggml_cuda_op_div(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
88

9+
void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse);
10+
911
void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

‎ggml/src/ggml-cuda/ggml-cuda.cu

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2902,6 +2902,49 @@ static void evaluate_and_capture_cuda_graph(ggml_backend_cuda_context * cuda_ctx
29022902

29032903
static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr);
29042904
if (!disable_fusion) {
2905+
2906+
if (node->op == GGML_OP_ADD) {
2907+
2908+
int n_fuse = 0;
2909+
2910+
ggml_op ops[8] = {GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD};
2911+
2912+
for (; n_fuse <= 6; ++n_fuse){
2913+
2914+
if (!ggml_can_fuse(cgraph, i + n_fuse, ops + n_fuse, 2)) {
2915+
break;
2916+
}
2917+
2918+
if (cgraph->nodes[i + n_fuse] != cgraph->nodes[i + n_fuse + 1]->src[0]) {
2919+
break;
2920+
}
2921+
2922+
if (!ggml_are_same_layout(cgraph->nodes[i + n_fuse]->src[1], cgraph->nodes[i + n_fuse + 1]->src[1])) {
2923+
break;
2924+
}
2925+
2926+
}
2927+
2928+
n_fuse++;
2929+
2930+
if (n_fuse > 1) {
2931+
printf("fused ADD: %d\n", n_fuse);
2932+
2933+
for (int j = 0; j < n_fuse - 1; ++j) {
2934+
node->src[j + 2] = cgraph->nodes[i + j + 1]->src[1];
2935+
}
2936+
2937+
cgraph->nodes[i + n_fuse - 1]->data = node->data;
2938+
2939+
ggml_cuda_op_fused_add(*cuda_ctx, node, n_fuse);
2940+
2941+
i += n_fuse -1;
2942+
2943+
continue;
2944+
}
2945+
}
2946+
2947+
29052948
if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) {
29062949
ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]);
29072950
i += 2;

‎tests/test-backend-ops.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5920,7 +5920,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
59205920
}
59215921

59225922
// fusion
5923-
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {2, 1, 1, 1}, 2));
5923+
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {1, 1, 1, 1}, 2));
59245924
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 2, 1, 1}, 3));
59255925
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {1, 1, 2, 1}, 4));
59265926
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 2}, 5));

0 commit comments

Comments
 (0)