1use std::num::NonZeroUsize;
2
3use crate::{
4 dag::NodeId,
5 device::{function, operation::CoreDeviceOps, tensor::Shape},
6 graph::{
7 DeviceFunction, Graph, GraphNodeIdTy,
8 ir::{
9 BackendMarker, GraphIR, GraphIRError,
10 node::AnnotatedNode,
11 operation::{GraphIROperationBase, GraphIROperationCompilable, GraphIROperationError, util},
12 },
13 },
14};
15
16#[derive(Clone, Copy, Debug, PartialEq)]
17pub struct AbsPowerError {
18 pub a: AnnotatedNode,
19 pub b: AnnotatedNode,
20 pub power: f32,
21}
22
23impl<B: BackendMarker> GraphIROperationBase<B> for AbsPowerError {
24 fn nodes(&self) -> Vec<AnnotatedNode> {
25 vec![self.a, self.b]
26 }
27
28 fn output_shape(&self, ir: &GraphIR<B>) -> Result<Shape, GraphIRError> {
29 util::check_dense_eq(ir, &self.a, true)?;
30 util::check_dense_eq(ir, &self.b, true)?;
31 util::check_same_batching(ir, &[&self.a, &self.b])?;
32
33 if self.a.shape == self.b.shape {
34 Ok(self.a.shape)
35 } else {
36 Err(GraphIRError::Op(GraphIROperationError::MismatchedInputShapes(vec![self.a.shape, self.b.shape])))
37 }
38 }
39}
40
41impl<B: BackendMarker> GraphIROperationCompilable<B> for AbsPowerError {
42 fn forward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
43 let output = graph.get_ref(output_node, GraphNodeIdTy::Values);
44 let bsn = util::batch_size_node::<B>(graph, &[self.a, self.b]);
45
46 let mut func = DeviceFunction::default();
47
48 func.push(function::MaybeUpdateBatchSize {
49 input: graph.get_ref(bsn, GraphNodeIdTy::Values),
50 output: output.clone(),
51 });
52
53 func.push(function::AbsPowerError {
54 a: graph.get_ref(self.a.idx, GraphNodeIdTy::Values),
55 b: graph.get_ref(self.b.idx, GraphNodeIdTy::Values),
56 power: self.power,
57 output,
58 });
59
60 func
61 }
62
63 fn backward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
64 let a = graph.get_ref(self.a.idx, GraphNodeIdTy::Values);
65 let b = graph.get_ref(self.b.idx, GraphNodeIdTy::Values);
66 let output_grad = graph.get_ref(output_node, GraphNodeIdTy::Gradients);
67
68 let mut func = DeviceFunction::default();
69
70 if let Some(grd) = graph.maybe_get_ref(self.a.idx, GraphNodeIdTy::Gradients) {
71 func.push(function::MaybeUpdateBatchSize { input: a.clone(), output: grd.clone() });
72 func.push(function::AbsPowerErrorBackward {
73 a: a.clone(),
74 b: b.clone(),
75 c: output_grad.clone(),
76 output: grd,
77 power: self.power,
78 });
79 }
80
81 if let Some(grd) = graph.maybe_get_ref(self.b.idx, GraphNodeIdTy::Gradients) {
82 func.push(function::MaybeUpdateBatchSize { input: b.clone(), output: grd.clone() });
83 func.push(function::AbsPowerErrorBackward {
84 a: b.clone(),
85 b: a.clone(),
86 c: output_grad.clone(),
87 output: grd,
88 power: self.power,
89 });
90 }
91
92 func
93 }
94}
95
96#[derive(Clone, Debug)]
97pub struct Concat {
98 pub a: AnnotatedNode,
99 pub b: AnnotatedNode,
100}
101
102impl<B: BackendMarker> GraphIROperationBase<B> for Concat {
103 fn nodes(&self) -> Vec<AnnotatedNode> {
104 vec![self.a, self.b]
105 }
106
107 fn output_shape(&self, ir: &GraphIR<B>) -> Result<Shape, GraphIRError> {
108 util::check_dense_eq(ir, &self.a, true)?;
109 util::check_dense_eq(ir, &self.b, true)?;
110 util::check_same_batching(ir, &[&self.a, &self.b])?;
111
112 let ash = self.a.shape;
113
114 if ash.cols() != 1 {
115 return Err(GraphIRError::Op(GraphIROperationError::InvalidInputShape(ash)));
116 }
117
118 if ash.cols() == self.b.shape.cols() {
119 Ok(Shape::new(ash.rows() + self.b.shape.rows(), ash.cols()))
120 } else {
121 Err(GraphIRError::Op(GraphIROperationError::MismatchedInputShapes(vec![ash, self.b.shape])))
122 }
123 }
124}
125
126impl<B: BackendMarker> GraphIROperationCompilable<B> for Concat {
127 fn forward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
128 let output = graph.get_ref(output_node, GraphNodeIdTy::Values);
129 let a = graph.get_ref(self.a.idx, GraphNodeIdTy::Values);
130 let b = graph.get_ref(self.b.idx, GraphNodeIdTy::Values);
131
132 let mut func = DeviceFunction::default();
133
134 func.push(function::MaybeUpdateBatchSize { input: a.clone(), output: output.clone() });
135
136 func.push(function::CopyOrAddStrided {
137 input: a.clone(),
138 output: output.clone(),
139 input_offset: 0,
140 output_offset: 0,
141 add: false,
142 len_is_out: false,
143 });
144
145 func.push(function::CopyOrAddStrided {
146 input: b,
147 output,
148 input_offset: 0,
149 output_offset: a.borrow().shape().size(),
150 add: false,
151 len_is_out: false,
152 });
153
154 func
155 }
156
157 fn backward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
158 let input = graph.get_ref(output_node, GraphNodeIdTy::Gradients);
159
160 let mut func = DeviceFunction::default();
161
162 if let Some(output) = graph.maybe_get_ref(self.a.idx, GraphNodeIdTy::Gradients) {
163 func.push(function::MaybeUpdateBatchSize { input: input.clone(), output: output.clone() });
164
165 func.push(function::CopyOrAddStrided {
166 input: input.clone(),
167 output,
168 input_offset: 0,
169 output_offset: 0,
170 add: true,
171 len_is_out: true,
172 });
173 }
174
175 if let Some(output) = graph.maybe_get_ref(self.b.idx, GraphNodeIdTy::Gradients) {
176 func.push(function::MaybeUpdateBatchSize { input: input.clone(), output: output.clone() });
177
178 let input_offset = graph.get_ref(self.a.idx, GraphNodeIdTy::Values).borrow().shape().size();
179 func.push(function::CopyOrAddStrided {
180 input,
181 output,
182 input_offset,
183 output_offset: 0,
184 add: true,
185 len_is_out: true,
186 });
187 }
188
189 func
190 }
191}
192
193#[derive(Clone, Debug)]
194pub struct FusedPairwiseMulConcat {
195 pub a: AnnotatedNode,
196 pub b: AnnotatedNode,
197}
198
199impl<B: BackendMarker> GraphIROperationBase<B> for FusedPairwiseMulConcat {
200 fn nodes(&self) -> Vec<AnnotatedNode> {
201 vec![self.a, self.b]
202 }
203
204 fn output_shape(&self, ir: &GraphIR<B>) -> Result<Shape, GraphIRError> {
205 util::check_dense_eq(ir, &self.a, true)?;
206 util::check_dense_eq(ir, &self.b, true)?;
207 util::check_same_batching(ir, &[&self.a, &self.b])?;
208
209 let ash = self.a.shape;
210
211 if ash.cols() != 1 {
212 return Err(GraphIRError::Op(GraphIROperationError::InvalidInputShape(ash)));
213 }
214
215 if ash.cols() == self.b.shape.cols() {
216 Ok(Shape::new(ash.rows() / 2 + self.b.shape.rows() / 2, ash.cols()))
217 } else {
218 Err(GraphIRError::Op(GraphIROperationError::MismatchedInputShapes(vec![ash, self.b.shape])))
219 }
220 }
221}
222
223impl<B: BackendMarker> GraphIROperationCompilable<B> for FusedPairwiseMulConcat {
224 fn forward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
225 let output = graph.get_ref(output_node, GraphNodeIdTy::Values);
226 let a = graph.get_ref(self.a.idx, GraphNodeIdTy::Values);
227 let b = graph.get_ref(self.b.idx, GraphNodeIdTy::Values);
228
229 let mut func = DeviceFunction::default();
230
231 func.push(function::MaybeUpdateBatchSize { input: a.clone(), output: output.clone() });
232
233 func.push(function::PairwiseMul { offset: 0, input: a.clone(), output: output.clone() });
234
235 func.push(function::PairwiseMul { offset: a.borrow().shape().size() / 2, input: b, output });
236
237 func
238 }
239
240 fn backward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
241 let input = graph.get_ref(output_node, GraphNodeIdTy::Gradients);
242 let a = graph.get_ref(self.a.idx, GraphNodeIdTy::Values);
243 let offset = a.borrow().shape().size() / 2;
244
245 let mut func = DeviceFunction::default();
246
247 if let Some(output) = graph.maybe_get_ref(self.a.idx, GraphNodeIdTy::Gradients) {
248 func.push(function::MaybeUpdateBatchSize { input: input.clone(), output: output.clone() });
249
250 func.push(function::PairwiseMulBackward { offset: 0, input: input.clone(), values: a, output });
251 }
252
253 if let Some(output) = graph.maybe_get_ref(self.b.idx, GraphNodeIdTy::Gradients) {
254 let values = graph.get_ref(self.b.idx, GraphNodeIdTy::Values);
255
256 func.push(function::MaybeUpdateBatchSize { input: input.clone(), output: output.clone() });
257
258 func.push(function::PairwiseMulBackward { offset, input, values, output });
259 }
260
261 func
262 }
263}
264
265#[derive(Clone, Debug)]
266pub struct Select {
267 pub input: AnnotatedNode,
268 pub buckets: AnnotatedNode,
269}
270
271impl<B: BackendMarker> GraphIROperationBase<B> for Select {
272 fn nodes(&self) -> Vec<AnnotatedNode> {
273 vec![self.input, self.buckets]
274 }
275
276 fn output_shape(&self, ir: &GraphIR<B>) -> Result<Shape, GraphIRError> {
277 util::check_dense_eq(ir, &self.input, true)?;
278 util::check_dense_eq(ir, &self.buckets, false)?;
279
280 if util::check_not_batched(ir, &self.buckets).is_ok() {
281 util::check_not_batched(ir, &self.input)?;
282 }
283
284 let is = self.input.shape;
285 let bs = self.buckets.shape;
286
287 if is.cols() == bs.cols() && is.rows() % bs.rows() == 0 {
288 Ok(Shape::new(is.rows() / bs.rows(), is.cols()))
289 } else {
290 Err(GraphIRError::Op(GraphIROperationError::MismatchedInputShapes(vec![is, bs])))
291 }
292 }
293}
294
295impl<B: BackendMarker> GraphIROperationCompilable<B> for Select
296where
297 B::Backend: CoreDeviceOps,
298{
299 fn forward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
300 let input = graph.get_ref(self.input.idx, GraphNodeIdTy::Values);
301 let output = graph.get_ref(output_node, GraphNodeIdTy::Values);
302 let buckets = graph.get_ref(self.buckets.idx, GraphNodeIdTy::Values);
303
304 let mut func = DeviceFunction::default();
305
306 func.push(function::MaybeUpdateBatchSize { input: buckets.clone(), output: output.clone() });
307 func.push(function::Select { input, output, buckets });
308
309 func
310 }
311
312 fn backward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
313 let mut func = DeviceFunction::default();
314
315 if let Some(output) = graph.maybe_get_ref(self.input.idx, GraphNodeIdTy::Gradients) {
316 let input = graph.get_ref(output_node, GraphNodeIdTy::Gradients);
317 let buckets = graph.get_ref(self.buckets.idx, GraphNodeIdTy::Values);
318 let values = graph.get_ref(self.input.idx, GraphNodeIdTy::Values);
319
320 func.push(function::MaybeUpdateBatchSize { input: values, output: output.clone() });
321 func.push(function::SelectBackprop { input, output, buckets });
322 }
323
324 func
325 }
326}
327
328#[derive(Clone, Debug)]
329pub struct SoftmaxCrossEntropy {
330 pub logits: AnnotatedNode,
331 pub targets: AnnotatedNode,
332}
333
334impl<B: BackendMarker> GraphIROperationBase<B> for SoftmaxCrossEntropy {
335 fn nodes(&self) -> Vec<AnnotatedNode> {
336 vec![self.logits, self.targets]
337 }
338
339 fn output_shape(&self, ir: &GraphIR<B>) -> Result<Shape, GraphIRError> {
340 util::check_dense_eq(ir, &self.logits, true)?;
341 util::check_dense_eq(ir, &self.targets, true)?;
342 util::check_same_batching(ir, &[&self.logits, &self.targets])?;
343 util::check_no_grad(ir, &[&self.targets])?;
344
345 let shape = self.logits.shape;
346
347 if shape != self.targets.shape {
348 Err(GraphIRError::Op(GraphIROperationError::MismatchedInputShapes(vec![shape, self.targets.shape])))
349 } else {
350 Ok(shape)
351 }
352 }
353
354 fn ancillary_buffers(&self, ir: &GraphIR<B>) -> Result<Vec<(Shape, Option<NonZeroUsize>, bool)>, GraphIRError> {
355 let batched = ir.get(self.logits.idx)?.ty().batched;
356 Ok(vec![(self.logits.shape, None, batched)])
357 }
358}
359
360impl<B: BackendMarker> GraphIROperationCompilable<B> for SoftmaxCrossEntropy
361where
362 B::Backend: CoreDeviceOps,
363{
364 fn forward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
365 let logits = graph.get_ref(self.logits.idx, GraphNodeIdTy::Values);
366 let targets = graph.get_ref(self.targets.idx, GraphNodeIdTy::Values);
367 let smax = graph.get_ref(output_node, GraphNodeIdTy::Ancillary(0));
368 let output = graph.get_ref(output_node, GraphNodeIdTy::Values);
369
370 let mut func = DeviceFunction::default();
371
372 func.push(function::MaybeUpdateBatchSize { input: logits.clone(), output: smax.clone() });
373 func.push(function::MaybeUpdateBatchSize { input: logits.clone(), output: output.clone() });
374
375 func.push(function::Softmax { input: logits, output: smax.clone() });
376 func.push(function::CrossEntropy { a: smax, b: targets, output });
377
378 func
379 }
380
381 fn backward_pass(&self, graph: &Graph<B::Backend>, output_node: NodeId) -> DeviceFunction<B::Backend> {
382 let mut func = DeviceFunction::default();
383
384 if let Some(output) = graph.maybe_get_ref(self.logits.idx, GraphNodeIdTy::Gradients) {
385 let softmax = graph.get_ref(output_node, GraphNodeIdTy::Ancillary(0));
386 let output_grads = graph.get_ref(output_node, GraphNodeIdTy::Gradients);
387 let targets = graph.get_ref(self.targets.idx, GraphNodeIdTy::Values);
388
389 func.push(function::MaybeUpdateBatchSize { input: softmax.clone(), output: output.clone() });
390 func.push(function::SoftmaxCrossEntropyBackward { softmax, output_grads, targets, output });
391 }
392
393 func
394 }
395}