ferrotensor_backend/engine/
elementwise.rs

1//! Element-wise kernels: arithmetic, activations and their backward passes.
2
3use super::{groups64, ComputeEngine, Kernel};
4use crate::{UnifiedTensor, WgpuContext};
5
6impl ComputeEngine {
7    /// `out = a + b`, element-wise; `a`, `b` and `out` have the same length.
8    pub fn add(
9        &self,
10        ctx: &WgpuContext,
11        a: &UnifiedTensor,
12        b: &UnifiedTensor,
13        out: &UnifiedTensor,
14    ) {
15        self.binary(ctx, &self.add, a, b, out);
16    }
17
18    /// `out = a * b`, element-wise.
19    pub fn mul(
20        &self,
21        ctx: &WgpuContext,
22        a: &UnifiedTensor,
23        b: &UnifiedTensor,
24        out: &UnifiedTensor,
25    ) {
26        self.binary(ctx, &self.mul, a, b, out);
27    }
28
29    /// `out = a - b`, element-wise.
30    pub fn sub(
31        &self,
32        ctx: &WgpuContext,
33        a: &UnifiedTensor,
34        b: &UnifiedTensor,
35        out: &UnifiedTensor,
36    ) {
37        self.binary(ctx, &self.sub, a, b, out);
38    }
39
40    pub(super) fn binary(
41        &self,
42        ctx: &WgpuContext,
43        kernel: &Kernel,
44        a: &UnifiedTensor,
45        b: &UnifiedTensor,
46        out: &UnifiedTensor,
47    ) {
48        self.run(
49            ctx,
50            kernel,
51            &[(0, a.buffer()), (1, b.buffer()), (2, out.buffer())],
52            groups64(a.shape.iter().product()),
53        );
54    }
55
56    /// Gradients of `a + b`: writes `grad_out` to both `grad_a` and `grad_b`.
57    pub fn add_backward(
58        &self,
59        ctx: &WgpuContext,
60        grad_out: &UnifiedTensor,
61        grad_a: &UnifiedTensor,
62        grad_b: &UnifiedTensor,
63    ) {
64        self.run(
65            ctx,
66            &self.add_bw,
67            &[
68                (3, grad_out.buffer()),
69                (4, grad_a.buffer()),
70                (5, grad_b.buffer()),
71            ],
72            groups64(grad_out.shape.iter().product()),
73        );
74    }
75
76    /// Gradients of `a - b`: `grad_a = grad_out`, `grad_b = -grad_out`.
77    pub fn sub_backward(
78        &self,
79        ctx: &WgpuContext,
80        grad_out: &UnifiedTensor,
81        grad_a: &UnifiedTensor,
82        grad_b: &UnifiedTensor,
83    ) {
84        self.run(
85            ctx,
86            &self.sub_bw,
87            &[
88                (3, grad_out.buffer()),
89                (4, grad_a.buffer()),
90                (5, grad_b.buffer()),
91            ],
92            groups64(grad_out.shape.iter().product()),
93        );
94    }
95
96    /// Gradients of `a * b`: `grad_a = grad_out * b`, `grad_b = grad_out * a`.
97    pub fn mul_backward(
98        &self,
99        ctx: &WgpuContext,
100        grad_out: &UnifiedTensor,
101        a: &UnifiedTensor,
102        b: &UnifiedTensor,
103        grad_a: &UnifiedTensor,
104        grad_b: &UnifiedTensor,
105    ) {
106        self.run(
107            ctx,
108            &self.mul_bw,
109            &[
110                (3, grad_out.buffer()),
111                (4, a.buffer()),
112                (5, b.buffer()),
113                (6, grad_a.buffer()),
114                (7, grad_b.buffer()),
115            ],
116            groups64(grad_out.shape.iter().product()),
117        );
118    }
119
120    /// `out = x * sigmoid(x)`, element-wise.
121    pub fn silu(&self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor) {
122        self.run(
123            ctx,
124            &self.silu,
125            &[(0, input.buffer()), (1, out.buffer())],
126            groups64(input.shape.iter().product()),
127        );
128    }
129
130    /// Gradient of SiLU: `grad_in = grad_out * s * (1 + x * (1 - s))` with
131    /// `s = sigmoid(input)`.
132    pub fn silu_backward(
133        &self,
134        ctx: &WgpuContext,
135        grad_out: &UnifiedTensor,
136        input: &UnifiedTensor,
137        grad_in: &UnifiedTensor,
138    ) {
139        self.run(
140            ctx,
141            &self.silu_bw,
142            &[
143                (3, grad_out.buffer()),
144                (4, input.buffer()),
145                (5, grad_in.buffer()),
146            ],
147            groups64(grad_out.shape.iter().product()),
148        );
149    }
150
151    /// `out = max(x, 0)`, element-wise.
152    pub fn relu(&self, ctx: &WgpuContext, input: &UnifiedTensor, out: &UnifiedTensor) {
153        self.run(
154            ctx,
155            &self.relu,
156            &[(0, input.buffer()), (1, out.buffer())],
157            groups64(input.shape.iter().product()),
158        );
159    }
160
161    /// Gradient of ReLU: `grad_out` where `input > 0`, else 0.
162    pub fn relu_backward(
163        &self,
164        ctx: &WgpuContext,
165        grad_out: &UnifiedTensor,
166        input: &UnifiedTensor,
167        grad_in: &UnifiedTensor,
168    ) {
169        self.run(
170            ctx,
171            &self.relu_bw,
172            &[
173                (3, grad_out.buffer()),
174                (4, input.buffer()),
175                (5, grad_in.buffer()),
176            ],
177            groups64(grad_out.shape.iter().product()),
178        );
179    }
180
181    /// Fused SiLU-Gate: out = SiLU(up) * gate
182    /// Reduces memory bandwidth by 2x compared to separate ops
183    pub fn silu_gate_fused(
184        &self,
185        ctx: &WgpuContext,
186        up: &UnifiedTensor,
187        gate: &UnifiedTensor,
188        out: &UnifiedTensor,
189    ) {
190        self.binary(ctx, &self.silu_gate, up, gate, out);
191    }
192
193    /// Fused SiLU-Gate backward: computes gradients for up and gate
194    pub fn silu_gate_backward(
195        &self,
196        ctx: &WgpuContext,
197        up: &UnifiedTensor,
198        gate: &UnifiedTensor,
199        grad_out: &UnifiedTensor,
200        grad_up: &UnifiedTensor,
201        grad_gate: &UnifiedTensor,
202    ) {
203        self.run(
204            ctx,
205            &self.silu_gate_bw,
206            &[
207                (0, up.buffer()),
208                (1, gate.buffer()),
209                (3, grad_out.buffer()),
210                (4, grad_up.buffer()),
211                (5, grad_gate.buffer()),
212            ],
213            groups64(up.shape.iter().product()),
214        );
215    }
216}
ferrotensor_backend/engine/elementwise.rs

ferrotensor_backend/engine/
elementwise.rs
