char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

YiChen Lvmetal: per-op source split + parallel compile (#26561)b615f5b4b

char/aliceai-t5-moe
3.8 KiB148 linesraw
1#include "common.h"
2
3kernel void kernel_pool_2d_max_f32(
4        constant    ggml_metal_kargs_pool_2d & args,
5        device  const float * src0,
6        device        float * dst,
7        uint        gid[[thread_position_in_grid]]) {
8
9    if (gid >= args.np) {
10        return;
11    }
12
13    const int idx = gid;
14    const int I_HW = args.IH * args.IW;
15    const int O_HW = args.OH * args.OW;
16    const int nc = idx / O_HW;
17    const int cur_oh = idx % O_HW / args.OW;
18    const int cur_ow = idx % O_HW % args.OW;
19
20    device const float * i_ptr = src0 + nc * I_HW;
21    device       float * o_ptr = dst  + nc * O_HW;
22
23    const int start_h = cur_oh * args.s1 - args.p1;
24    const int bh = MAX(0,  start_h);
25    const int eh = MIN(args.IH, start_h + args.k1);
26    const int start_w = cur_ow * args.s0 - args.p0;
27    const int bw = MAX(0,  start_w);
28    const int ew = MIN(args.IW, start_w + args.k0);
29
30    float res = -INFINITY;
31
32    for (int i = bh; i < eh; i += 1) {
33        for (int j = bw; j < ew; j += 1) {
34            res = MAX(res, i_ptr[i * args.IW + j]);
35        }
36    }
37
38    o_ptr[cur_oh * args.OW + cur_ow] = res;
39}
40
41kernel void kernel_pool_2d_avg_f32(
42        constant    ggml_metal_kargs_pool_2d & args,
43        device  const float * src0,
44        device        float * dst,
45        uint        gid[[thread_position_in_grid]]) {
46
47    if (gid >= args.np) {
48        return;
49    }
50
51    const int idx = gid;
52    const int I_HW = args.IH * args.IW;
53    const int O_HW = args.OH * args.OW;
54    const int nc = idx / O_HW;
55    const int cur_oh = idx % O_HW / args.OW;
56    const int cur_ow = idx % O_HW % args.OW;
57
58    device const float * i_ptr = src0 + nc * I_HW;
59    device       float * o_ptr = dst  + nc * O_HW;
60
61    const int start_h = cur_oh * args.s1 - args.p1;
62    const int bh = MAX(0,  start_h);
63    const int eh = MIN(args.IH, start_h + args.k1);
64    const int start_w = cur_ow * args.s0 - args.p0;
65    const int bw = MAX(0,  start_w);
66    const int ew = MIN(args.IW, start_w + args.k0);
67    // const float scale = 1. / ((eh - bh) * (ew - bw));
68    const float scale = 1. / (args.k0 * args.k1);
69
70    float res = 0;
71
72    for (int i = bh; i < eh; i += 1) {
73        for (int j = bw; j < ew; j += 1) {
74            float cur = i_ptr[i * args.IW + j];
75            res += cur * scale;
76        }
77    }
78
79    o_ptr[cur_oh * args.OW + cur_ow] = res;
80}
81
82
83kernel void kernel_pool_1d_max_f32(
84        constant        ggml_metal_kargs_pool_1d & args,
85        device  const   float * src,
86        device          float * dst,
87        uint            gid [[thread_position_in_grid]]
88) {
89
90    if (gid >= args.np) {
91        return;
92    }
93
94    const int ow  = (int)gid % args.OW;
95    const int row = (int)gid / args.OW;
96
97    const int base = ow * args.s0 - args.p0;
98
99    float acc = -INFINITY;
100
101    const int src_off = row * args.IW;
102    const int dst_off = row * args.OW;
103
104    for (int ki = 0; ki < args.k0; ++ki) {
105        int j = base + ki;
106        if (j < 0 || j >= args.IW){
107            continue;
108        }
109        float v = src[src_off + j];
110        acc = max(acc, v);
111    }
112
113    dst[dst_off + ow] = acc;
114}
115
116kernel void kernel_pool_1d_avg_f32(
117        constant        ggml_metal_kargs_pool_1d & args,
118        device  const   float * src,
119        device          float * dst,
120        uint            gid [[thread_position_in_grid]]
121) {
122
123    if (gid >= args.np) {
124        return;
125    }
126
127    const int ow  = (int)gid % args.OW;
128    const int row = (int)gid / args.OW;
129
130    const int base = ow * args.s0 - args.p0;
131
132    float acc = 0.0f;
133    int   cnt = 0;
134
135    const int src_off = row * args.IW;
136    const int dst_off = row * args.OW;
137
138    for (int ki = 0; ki < args.k0; ++ki) {
139        const int j = base + ki;
140        if (j < 0 || j >= args.IW) {
141            continue;
142        }
143        acc += src[src_off + j];
144        cnt += 1;
145    }
146
147    dst[dst_off + ow] = (cnt > 0) ? (acc / (float)cnt) : 0.0f;
148}