GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 100.0% 6 / 0 / 6
Functions: 100.0% 2 / 0 / 2
Branches: -% 0 / 0 / 0
include/brotensor/ops/activation.h
Line Branch Exec Source
1 #pragma once
2
3 // brotensor ops/activation.h — Pointwise activations + GLU gates (relu/tanh/sigmoid/silu/gelu/snake/elu/geglu/swiglu).
4
5 #include "../tensor.h"
6 #include <cstdint>
7
8 namespace brotensor {
9
10
11 // y = max(x, 0). Shapes match; y resized if mis-shaped. x and y may alias.
12 void relu_forward(const Tensor& x, Tensor& y);
13
14
15 // dX = dY * (x > 0). dX resized to match x; may alias dY.
16 void relu_backward(const Tensor& x, const Tensor& dY, Tensor& dX);
17
18
19 // y = tanh(x). y resized to match x.
20 void tanh_forward(const Tensor& x, Tensor& y);
21
22
23 // dX = dY * (1 - y*y). `y` is the cached forward output (not raw x).
24 void tanh_backward(const Tensor& y, const Tensor& dY, Tensor& dX);
25
26
27 // y = 1 / (1 + exp(-x)).
28 void sigmoid_forward(const Tensor& x, Tensor& y);
29
30
31 // dX = dY * y * (1 - y). `y` is the cached forward output.
32 void sigmoid_backward(const Tensor& y, const Tensor& dY, Tensor& dX);
33
34
35 // Elementwise ReLU / Tanh over (B,D). Y resized to match X; X and Y may alias.
36 void relu_forward_batched(const Tensor& X_BD, Tensor& Y_BD);
37
38 void tanh_forward_batched(const Tensor& X_BD, Tensor& Y_BD);
39
40
41 // Elementwise activation backward over (B,D), same shapes throughout.
42 // relu: dX = dY*(X>0), reads X_BD (forward input).
43 // tanh: dX = dY*(1-Y*Y), reads Y_BD (forward output).
44 void relu_backward_batched(const Tensor& X_BD, const Tensor& dY_BD,
45 Tensor& dX_BD);
46
47 void tanh_backward_batched(const Tensor& Y_BD, const Tensor& dY_BD,
48 Tensor& dX_BD);
49
50
51 // SiLU / Swish: y = x*sigmoid(x). Dispatched FP32/FP16 on x.dtype; y resized +
52 // dtype-set to match x. x and y may alias.
53 void silu_forward(const Tensor& x, Tensor& y);
54
55
56 // SiLU backward, reads the raw forward input x:
57 // dX = dY * sigmoid(x) * (1 + x*(1-sigmoid(x))).
58 // Dispatched FP32/FP16 on x.dtype; dX resized + dtype-set to match x; may alias dY.
59 void silu_backward(const Tensor& x, const Tensor& dY, Tensor& dX);
60
61
62 // GELU, tanh approximation (PyTorch approximate="tanh"):
63 // y = 0.5*x*(1 + tanh(sqrt(2/pi)*(x + 0.044715*x^3))).
64 // Dispatched FP32/FP16 on x.dtype.
65 void gelu_forward(const Tensor& x, Tensor& y);
66
67
68 // GELU (tanh-approx) backward, reads x. With k=sqrt(2/pi),
69 // u=k*(x+0.044715*x^3), t=tanh(u):
70 // dX = dY * [0.5*(1+t) + 0.5*x*(1-t^2)*k*(1+3*0.044715*x^2)].
71 // Dispatched FP32/FP16 on x.dtype; dX resized + dtype-set to match x; may alias dY.
72 void gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX);
73
74
75 // Exact GELU (erf form, PyTorch approximate="none", diffusers default):
76 // y = 0.5*x*(1 + erf(x/sqrt(2))).
77 // Distinct from the tanh-approx gelu_forward. Dispatched FP32/FP16 on x.dtype;
78 // y resized + dtype-set to match x. x and y may alias.
79 void gelu_exact_forward(const Tensor& x, Tensor& y);
80
81
82 // Exact-GELU backward, reads x:
83 // dX = dY * [0.5*(1+erf(x/sqrt(2))) + x*phi(x)], phi = standard normal pdf.
84 // Dispatched FP32/FP16 on x.dtype; dX resized + dtype-set to match x; may alias dY.
85 void gelu_exact_backward(const Tensor& x, const Tensor& dY,
86 Tensor& dX);
87
88
89 // QuickGELU: y = x*sigmoid(1.702*x). OpenAI CLIP's activation (SD1.5 text
90 // encoder). Dispatched FP32/FP16 on x.dtype; y resized + dtype-set to match x.
91 // x and y may alias.
92 void quick_gelu_forward(const Tensor& x, Tensor& y);
93
94
95 // QuickGELU backward, reads x. With s = sigmoid(1.702*x):
96 // dX = dY * (s + x*1.702*s*(1-s)).
97 // Dispatched FP32/FP16 on x.dtype; dX resized + dtype-set to match x; may alias dY.
98 void quick_gelu_backward(const Tensor& x, const Tensor& dY,
99 Tensor& dX);
100
101
102 // GEGLU: input (B,2*D) split along the last dim into A=(B,D) and B_half=(B,D);
103 // output (B,D) = A * gelu(B_half) (tanh-approx). Dispatched FP32/FP16 on
104 // X.dtype; Y resized + dtype-set to match X.
105 void geglu_forward(const Tensor& X, Tensor& Y);
106
107
108 // GEGLU backward. With g = gelu(B_half) (tanh-approx):
109 // dA = dY*g; dB_half = dY*A*gelu'(B_half).
110 // dX = concat(dA, dB_half) along the last dim (A then B_half). Dispatched
111 // FP32/FP16 on X.dtype; dX resized + dtype-set to match X.
112 void geglu_backward(const Tensor& X, const Tensor& dY,
113 Tensor& dX);
114
115
116 // Exact-GELU GEGLU: same split as geglu_forward but output = A*gelu_exact(B_half),
117 // using the exact erf-based GELU. Matches diffusers' default GEGLU. Dispatched
118 // FP32/FP16 on X.dtype; Y resized + dtype-set to match X.
119 void geglu_exact_forward(const Tensor& X, Tensor& Y);
120
121
122 // Exact-GELU GEGLU backward. With g = gelu_exact(B_half):
123 // dA = dY*g; dB_half = dY*A*gelu_exact'(B_half).
124 // dX = concat(dA, dB_half) along the last dim (A then B_half). Dispatched
125 // FP32/FP16 on X.dtype; dX resized + dtype-set to match X.
126 void geglu_exact_backward(const Tensor& X, const Tensor& dY,
127 Tensor& dX);
128
129
130 // SwiGLU (Llama FFN gate): input (B,2*D) split along the last dim into A=(B,D)
131 // and B_half=(B,D); output (B,D) = silu(A) * B_half. Dispatched FP32/FP16 on
132 // X.dtype; Y resized + dtype-set to match X.
133 void swiglu_forward(const Tensor& X, Tensor& Y);
134
135
136 // SwiGLU backward. With s = silu(A):
137 // dA = dY*B_half*silu'(A); dB_half = dY*s.
138 // dX = concat(dA, dB_half) along the last dim (A then B_half). Dispatched on
139 // X.dtype; dX resized + dtype-set to match X.
140 void swiglu_backward(const Tensor& X, const Tensor& dY,
141 Tensor& dX);
142
143
144 // ─── Vocoder / codec activations (audio) ───────────────────────────────────
145 //
146 // FP32-only, implemented on all three backends (CPU / CUDA / Metal). NCL
147 // layout — the (N,C,L) dims are passed as int args; element (n,c,l) is at flat
148 // index (n*C+c)*L + l.
149
150 // Snake activation (BigVGAN / DAC vocoder), per-channel learnable alpha (and
151 // optional beta):
152 // plain snake (beta == null): y = x + (1/alpha_c)*sin^2(alpha_c*x)
153 // snakebeta (beta != null): y = x + (1/beta_c) *sin^2(alpha_c*x)
154 // alpha/beta are per-channel (broadcast across the (n,l) plane). The reciprocal
155 // denominator is sign-preserved-floored at magnitude 1e-9 to avoid NaN/Inf.
156 // X, Y: (N,C*L). alpha: (C,1) or (1,C). beta: (C,1)/(1,C) or null.
157 // Y resized + dtype-set to match X; X and Y may alias. FP32, CPU-resident.
158 void snake_forward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
159 int N, int C, int L, Tensor& Y);
160
161
162 // Snake backward, reads the raw forward input X. With s=sin(a*x), c=cos(a*x),
163 // a=alpha_c, denom=(beta?beta_c:a), r=1/denom (sign-guarded as in the forward):
164 // dy/dx = 1 + 2*a*r*s*c
165 // dy/dalpha = 2*r*x*s*c (plain snake also adds the -r^2*s^2 term,
166 // since denom==alpha there)
167 // dy/dbeta = -r^2*s^2 (snakebeta only)
168 // dX: (N,C*L) overwritten (resized + dtype-set to X).
169 // dAlpha: (C,1) accumulated — caller zeros.
170 // dBeta: (C,1) accumulated — caller zeros; non-null exactly when beta is.
171 void snake_backward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
172 const Tensor& dY, int N, int C, int L,
173 Tensor& dX, Tensor& dAlpha, Tensor* dBeta);
174
175
176 // ELU (EnCodec activation), elementwise:
177 // y = x if x > 0
178 // y = alpha*(exp(x) - 1) otherwise
179 // y resized to match x; x and y may alias. CPU FP32-only.
180 void elu_forward(const Tensor& x, float alpha, Tensor& y);
181
182 1 inline void elu_forward(const Tensor& x, Tensor& y) {
183 1 elu_forward(x, /*alpha=*/1.0f, y);
184 1 }
185
186
187 // ELU backward, reads the raw forward input x:
188 // dX = dY * (x > 0 ? 1 : alpha*exp(x)).
189 // dX overwritten (resized to match x); may alias dY. CPU FP32-only.
190 void elu_backward(const Tensor& x, const Tensor& dY, float alpha, Tensor& dX);
191
192 1 inline void elu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
193 1 elu_backward(x, dY, /*alpha=*/1.0f, dX);
194 1 }
195
196
197 // Leaky ReLU (HiFi-GAN activation), elementwise:
198 // y = x > 0 ? x : negative_slope*x.
199 // y resized to match x; x and y may alias. CPU FP32-only; the CUDA forward
200 // is dtype-dispatched on x (FP32/FP16/BF16, FP32 math per element).
201 void leaky_relu_forward(const Tensor& x, float negative_slope, Tensor& y);
202
203
204 // Leaky ReLU backward, reads the raw forward input x:
205 // dX = dY * (x > 0 ? 1 : negative_slope).
206 // dX overwritten (resized to match x); may alias dY. CPU FP32-only.
207 void leaky_relu_backward(const Tensor& x, const Tensor& dY,
208 float negative_slope, Tensor& dX);
209
210 } // namespace brotensor
211