1#include "binary-ops.h"
2
3#if defined(GGML_USE_ACCELERATE)
4#include <Accelerate/Accelerate.h>
5
6using vDSP_fn_t = void (*)(const float *, vDSP_Stride, const float *, vDSP_Stride, float *, vDSP_Stride, vDSP_Length);
7#endif
8
9static inline float op_add(float a, float b) {
10 return a + b;
11}
12
13static inline float op_sub(float a, float b) {
14 return a - b;
15}
16
17static inline float op_mul(float a, float b) {
18 return a * b;
19}
20
21static inline float op_div(float a, float b) {
22 return a / b;
23}
24
25template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
26static inline void vec_binary_op_contiguous(const int64_t n, dst_t * z, const src0_t * x, const src1_t * y) {
27 constexpr auto src0_to_f32 = type_conversion_table<src0_t>::to_f32;
28 constexpr auto src1_to_f32 = type_conversion_table<src1_t>::to_f32;
29 constexpr auto f32_to_dst = type_conversion_table<dst_t >::from_f32;
30
31 for (int i = 0; i < n; i++) {
32 z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(y[i])));
33 }
34}
35
36template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
37static inline void vec_binary_op_non_contiguous(const int64_t n, const int64_t ne10, const int64_t nb10, dst_t * z, const src0_t * x, const src1_t * y) {
38 constexpr auto src0_to_f32 = type_conversion_table<src0_t>::to_f32;
39 constexpr auto src1_to_f32 = type_conversion_table<src1_t>::to_f32;
40 constexpr auto f32_to_dst = type_conversion_table<dst_t >::from_f32;
41
42 for (int i = 0; i < n; i++) {
43 int i10 = i % ne10;
44 const src1_t * y_ptr = (const src1_t *)((const char *)y + i10*nb10);
45 z[i] = f32_to_dst(op(src0_to_f32(x[i]), src1_to_f32(*y_ptr)));
46 }
47}
48
49template <float (*op)(float, float), typename src0_t, typename src1_t, typename dst_t>
50static void apply_binary_op(const ggml_compute_params * params, ggml_tensor * dst) {
51 const ggml_tensor * src0 = dst->src[0];
52 const ggml_tensor * src1 = dst->src[1];
53
54 GGML_ASSERT(ggml_can_repeat(src1, src0) && ggml_are_same_shape(src0, dst));
55
56 GGML_TENSOR_BINARY_OP_LOCALS
57
58 GGML_ASSERT( nb0 == sizeof(dst_t));
59 GGML_ASSERT(nb00 == sizeof(src0_t));
60
61 const auto [ir0, ir1] = get_thread_range(params, src0);
62 const bool is_src1_contiguous_rows = ggml_is_contiguous_rows(src1);
63
64#ifdef GGML_USE_ACCELERATE
65 vDSP_fn_t vDSP_op = nullptr;
66 // TODO - avoid the f32-only check using type 'trait' lookup tables and row-based src-to-float conversion functions
67 if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
68 if (op == op_add) {
69 vDSP_op = vDSP_vadd;
70 } else if (op == op_sub) {
71 vDSP_op = vDSP_vsub;
72 } else if (op == op_mul) {
73 vDSP_op = vDSP_vmul;
74 } else if (op == op_div) {
75 vDSP_op = vDSP_vdiv;
76 }
77 }
78#endif
79
80 for (int64_t ir = ir0; ir < ir1; ++ir) {
81 const int64_t i03 = ir/(ne02*ne01);
82 const int64_t i02 = (ir - i03*ne02*ne01)/ne01;
83 const int64_t i01 = (ir - i03*ne02*ne01 - i02*ne01);
84
85 const int64_t i13 = i03 % ne13;
86 const int64_t i12 = i02 % ne12;
87 const int64_t i11 = i01 % ne11;
88
89 dst_t * dst_ptr = (dst_t *) ((char *) dst->data + i03*nb3 + i02*nb2 + i01*nb1 );
90 const src0_t * src0_ptr = (const src0_t *) ((const char *) src0->data + i03*nb03 + i02*nb02 + i01*nb01);
91 const src1_t * src1_ptr = (const src1_t *) ((const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11);
92
93 if (is_src1_contiguous_rows) {
94 // src1 is broadcastable across src0 and dst in i1, i2, i3
95 const int64_t nr0 = ne00 / ne10;
96
97 for (int64_t r = 0; r < nr0; ++r) {
98#ifdef GGML_USE_ACCELERATE
99 if constexpr (std::is_same_v<src0_t, float> && std::is_same_v<src1_t, float> && std::is_same_v<dst_t, float>) {
100 if (vDSP_op != nullptr) {
101 vDSP_op(src1_ptr, 1, src0_ptr + r*ne10, 1, dst_ptr + r*ne10, 1, ne10);
102 continue;
103 }
104 }
105#endif
106 vec_binary_op_contiguous<op>(ne10, dst_ptr + r*ne10, src0_ptr + r*ne10, src1_ptr);
107 }
108 } else {
109 vec_binary_op_non_contiguous<op>(ne0, ne10, nb10, dst_ptr, src0_ptr, src1_ptr);
110 }
111 }
112}
113
114// TODO: Use the 'traits' lookup table (for type conversion fns), instead of a mass of 'if' conditions with long templates
115template <float (*op)(float, float)>
116static void binary_op(const ggml_compute_params * params, ggml_tensor * dst) {
117 const ggml_tensor * src0 = dst->src[0];
118 const ggml_tensor * src1 = dst->src[1];
119
120 /* */ if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { // all f32
121 apply_binary_op<op, float, float, float>(params, dst);
122 } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) { // all f16
123 apply_binary_op<op, ggml_fp16_t, ggml_fp16_t, ggml_fp16_t>(params, dst);
124 } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) { // all bf16
125 apply_binary_op<op, ggml_bf16_t, ggml_bf16_t, ggml_bf16_t>(params, dst);
126 } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) {
127 apply_binary_op<op, ggml_bf16_t, float, ggml_bf16_t>(params, dst);
128 } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
129 apply_binary_op<op, ggml_bf16_t, float, float>(params, dst);
130 } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
131 apply_binary_op<op, ggml_fp16_t, float, ggml_fp16_t>(params, dst);
132 } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
133 apply_binary_op<op, ggml_fp16_t, float, float>(params, dst);
134 } else {
135 GGML_ABORT("%s: unsupported types: dst: %s, src0: %s, src1: %s\n", __func__,
136 ggml_type_name(dst->type), ggml_type_name(src0->type), ggml_type_name(src1->type));
137 }
138}
139
140void ggml_compute_forward_add_non_quantized(const ggml_compute_params * params, ggml_tensor * dst) {
141 binary_op<op_add>(params, dst);
142}
143
144void ggml_compute_forward_sub(const ggml_compute_params * params, ggml_tensor * dst) {
145 binary_op<op_sub>(params, dst);
146}
147
148void ggml_compute_forward_mul(const ggml_compute_params * params, ggml_tensor * dst) {
149 binary_op<op_mul>(params, dst);
150}
151
152void ggml_compute_forward_div(const ggml_compute_params * params, ggml_tensor * dst) {
153 binary_op<op_div>(params, dst);
154}