Skip to content

Commit b46ae7b

Browse files
committed
Merge branch 'master' into HEAD
2 parents 048831e + 13268c5 commit b46ae7b

File tree

2 files changed

+32
-20
lines changed

2 files changed

+32
-20
lines changed

ggml-metal.m

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,12 @@ void ggml_metal_graph_compute(
683683
} break;
684684
case GGML_OP_ADD:
685685
{
686+
GGML_ASSERT(ggml_is_contiguous(src0));
687+
688+
// utilize float4
689+
GGML_ASSERT(ne00 % 4 == 0);
690+
const int64_t nb = ne00/4;
691+
686692
if (ggml_nelements(src1) == ne10) {
687693
// src1 is a row
688694
[encoder setComputePipelineState:ctx->pipeline_add_row];
@@ -692,14 +698,20 @@ void ggml_metal_graph_compute(
692698
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
693699
[encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
694700
[encoder setBuffer:id_dst offset:offs_dst atIndex:2];
695-
[encoder setBytes:&ne00 length:sizeof(ne00) atIndex:3];
701+
[encoder setBytes:&nb length:sizeof(nb) atIndex:3];
696702

697-
const int64_t n = ggml_nelements(dst);
703+
const int64_t n = ggml_nelements(dst)/4;
698704

699705
[encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
700706
} break;
701707
case GGML_OP_MUL:
702708
{
709+
GGML_ASSERT(ggml_is_contiguous(src0));
710+
711+
// utilize float4
712+
GGML_ASSERT(ne00 % 4 == 0);
713+
const int64_t nb = ne00/4;
714+
703715
if (ggml_nelements(src1) == ne10) {
704716
// src1 is a row
705717
[encoder setComputePipelineState:ctx->pipeline_mul_row];
@@ -709,9 +721,9 @@ void ggml_metal_graph_compute(
709721
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
710722
[encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
711723
[encoder setBuffer:id_dst offset:offs_dst atIndex:2];
712-
[encoder setBytes:&ne00 length:sizeof(ne00) atIndex:3];
724+
[encoder setBytes:&nb length:sizeof(nb) atIndex:3];
713725

714-
const int64_t n = ggml_nelements(dst);
726+
const int64_t n = ggml_nelements(dst)/4;
715727

716728
[encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
717729
} break;

ggml-metal.metal

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -6,41 +6,41 @@ using namespace metal;
66
#define MIN(x, y) ((x) < (y) ? (x) : (y))
77

88
kernel void kernel_add(
9-
device const float * src0,
10-
device const float * src1,
11-
device float * dst,
9+
device const float4 * src0,
10+
device const float4 * src1,
11+
device float4 * dst,
1212
uint tpig[[thread_position_in_grid]]) {
1313
dst[tpig] = src0[tpig] + src1[tpig];
1414
}
1515

1616
// assumption: src1 is a row
1717
// broadcast src1 into src0
1818
kernel void kernel_add_row(
19-
device const float * src0,
20-
device const float * src1,
21-
device float * dst,
22-
constant int64_t & ne00,
19+
device const float4 * src0,
20+
device const float4 * src1,
21+
device float4 * dst,
22+
constant int64_t & nb,
2323
uint tpig[[thread_position_in_grid]]) {
24-
dst[tpig] = src0[tpig] + src1[tpig % ne00];
24+
dst[tpig] = src0[tpig] + src1[tpig % nb];
2525
}
2626

2727
kernel void kernel_mul(
28-
device const float * src0,
29-
device const float * src1,
30-
device float * dst,
28+
device const float4 * src0,
29+
device const float4 * src1,
30+
device float4 * dst,
3131
uint tpig[[thread_position_in_grid]]) {
3232
dst[tpig] = src0[tpig] * src1[tpig];
3333
}
3434

3535
// assumption: src1 is a row
3636
// broadcast src1 into src0
3737
kernel void kernel_mul_row(
38-
device const float * src0,
39-
device const float * src1,
40-
device float * dst,
41-
constant int64_t & ne00,
38+
device const float4 * src0,
39+
device const float4 * src1,
40+
device float4 * dst,
41+
constant int64_t & nb,
4242
uint tpig[[thread_position_in_grid]]) {
43-
dst[tpig] = src0[tpig] * src1[tpig % ne00];
43+
dst[tpig] = src0[tpig] * src1[tpig % nb];
4444
}
4545

4646
kernel void kernel_scale(

0 commit comments

Comments
 (0)