tract_gpu/ops/
pulse.rs

1#![allow(unpredictable_function_pointer_comparisons)]
2use crate::device::{DeviceContext, get_context};
3use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice};
4use crate::turn_handler::make_tensor_for_node;
5use crate::utils::compute_broadcast_strides;
6use std::ops::Range;
7use tract_core::internal::*;
8use tract_core::ops::array::PadMode;
9use tract_pulse_opl::ops::{AffineChunkTrim, Delay, PulsePad};
10
11// ─── GpuDelay ────────────────────────────────────────────────────────────────
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub struct GpuDelay {
15    pub inner: Delay,
16}
17
18impl GpuDelay {
19    pub fn new(inner: &Delay) -> Self {
20        Self { inner: inner.clone() }
21    }
22}
23
24impl Op for GpuDelay {
25    fn name(&self) -> StaticName {
26        "GpuDelay".into()
27    }
28
29    fn info(&self) -> TractResult<Vec<String>> {
30        self.inner.info()
31    }
32
33    op_as_typed_op!();
34}
35
36impl EvalOp for GpuDelay {
37    not_out_of_plan!();
38
39    fn state(&self, ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
40        Ok(Some(Box::new(GpuDelayState {
41            node_id: ctx.node_id,
42            buffer: None,
43            shift_scratch: None,
44            lanes: 0,
45        })))
46    }
47}
48
49impl TypedOp for GpuDelay {
50    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
51        crate::utils::facts_to_device_facts(inputs, |facts| self.inner.output_facts(facts))
52            .with_context(|| format!("Error while computing output facts for {}", self.name()))
53    }
54
55    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
56        crate::utils::get_device_facts(inputs, |facts| self.inner.cost(facts))
57    }
58
59    as_op!();
60}
61
62/// Copy `len` steps along `axis` from `src` into `dst`, addressing one lane of
63/// each: axis 0 is the lane axis, or the copy spans it when the lanes are absent
64/// and the tensors carry a single stream's state.
65#[allow(clippy::too_many_arguments)]
66fn copy_lane(
67    ctx: &dyn DeviceContext,
68    dst: &DeviceTensor,
69    dst_lane: Option<usize>,
70    dst_start: usize,
71    src: &DeviceTensor,
72    src_lane: Option<usize>,
73    src_start: usize,
74    axis: usize,
75    len: usize,
76) -> TractResult<()> {
77    let mut zone: TVec<usize> = src.shape().into();
78    zone[axis] = len;
79    let mut dst_origin = tvec!(0; dst.rank());
80    let mut src_origin = tvec!(0; src.rank());
81    if let (Some(dst_lane), Some(src_lane)) = (dst_lane, src_lane) {
82        zone[0] = 1;
83        dst_origin[0] = dst_lane;
84        src_origin[0] = src_lane;
85    }
86    dst_origin[axis] = dst_start;
87    src_origin[axis] = src_start;
88    ctx.copy_with_origins(&zone, dst, &dst_origin, dst.strides(), src, &src_origin, src.strides())
89}
90
91/// Zero each of `lanes` whole, or the whole tensor when it holds a single
92/// stream's state and has no lane axis.
93fn zero_lanes(
94    ctx: &dyn DeviceContext,
95    dst: &DeviceTensor,
96    lanes: &[LaneId],
97    laned: bool,
98) -> TractResult<()> {
99    let zero = Tensor::zero_dt(dst.datum_type(), &[])?.into_device()?;
100    let flat: TVec<usize> = tvec!(0; dst.rank());
101    let broadcast: TVec<isize> = tvec!(0; dst.rank());
102    let mut zone: TVec<usize> = dst.shape().into();
103    if laned {
104        zone[0] = 1;
105    }
106    for lane in lanes {
107        let mut origin = tvec!(0; dst.rank());
108        if laned {
109            origin[0] = lane.0;
110        }
111        ctx.copy_with_origins(&zone, dst, &origin, dst.strides(), &zero, &flat, &broadcast)?;
112    }
113    Ok(())
114}
115
116/// Broadcast `value` over `range` along `axis` of one lane of `dst`.
117fn fill_lane(
118    ctx: &dyn DeviceContext,
119    dst: &DeviceTensor,
120    lane: Option<usize>,
121    axis: usize,
122    range: Range<usize>,
123    value: &DeviceTensor,
124) -> TractResult<()> {
125    let mut zone: TVec<usize> = dst.shape().into();
126    zone[axis] = range.len();
127    let mut origin = tvec!(0; dst.rank());
128    if let Some(lane) = lane {
129        zone[0] = 1;
130        origin[0] = lane;
131    }
132    origin[axis] = range.start;
133    ctx.copy_with_origins(
134        &zone,
135        dst,
136        &origin,
137        dst.strides(),
138        value,
139        &tvec!(0; dst.rank()),
140        &tvec!(0; dst.rank()),
141    )
142}
143
144/// The streaming context preceding the current pulse. `lanes` is the extent of
145/// the buffer's lane axis, 1 when the state serves a single stream and the
146/// buffer has no lane axis at all.
147#[derive(Debug, Clone)]
148pub struct GpuDelayState {
149    pub node_id: usize,
150    pub buffer: Option<DeviceTensor>,
151    pub shift_scratch: Option<DeviceTensor>,
152    lanes: usize,
153}
154
155impl GpuDelayState {
156    /// One seat's delay: `seat` indexes the batch axis of `input` and `output`,
157    /// `lane` the lane axis of the buffer. Both are absent when the state serves
158    /// a single stream.
159    fn delay_seat(
160        &mut self,
161        ctx: &dyn DeviceContext,
162        op: &Delay,
163        input: &DeviceTensor,
164        output: &mut DeviceTensor,
165        seat: Option<usize>,
166        lane: Option<usize>,
167    ) -> TractResult<()> {
168        let axis = op.axis;
169        let buffered = op.delay + op.overlap;
170        let input_pulse = input.shape()[axis];
171        let output_pulse = input_pulse + op.overlap;
172        let from_input = input_pulse.saturating_sub(op.delay);
173        let from_buffer = output_pulse.saturating_sub(from_input);
174        let buffer = self.buffer.as_ref().unwrap();
175
176        copy_lane(ctx, output, seat, 0, buffer, lane, 0, axis, from_buffer)?;
177        copy_lane(ctx, output, seat, from_buffer, input, seat, 0, axis, from_input)?;
178
179        if buffered < input_pulse {
180            copy_lane(ctx, buffer, lane, 0, input, seat, input_pulse - buffered, axis, buffered)?;
181        } else {
182            // CUDA memcpy is undefined for overlapping regions in the same
183            // buffer (parallel threads), so shift the lane left by input_pulse
184            // through a scratch buffer.
185            let keep = buffered - input_pulse;
186            let scratch = match self.shift_scratch.as_ref() {
187                Some(scratch) => scratch,
188                None => {
189                    let mut shape: TVec<usize> = buffer.shape().into();
190                    if lane.is_some() {
191                        shape[0] = 1;
192                    }
193                    self.shift_scratch
194                        .insert(DeviceTensor::uninitialized_dt(input.datum_type(), &shape)?)
195                }
196            };
197            let scratch_lane = lane.map(|_| 0);
198            copy_lane(ctx, scratch, scratch_lane, 0, buffer, lane, input_pulse, axis, keep)?;
199            copy_lane(ctx, buffer, lane, 0, scratch, scratch_lane, 0, axis, keep)?;
200            copy_lane(ctx, buffer, lane, keep, input, seat, 0, axis, input_pulse)?;
201        }
202        Ok(())
203    }
204}
205
206impl OpState for GpuDelayState {
207    fn eval(
208        &mut self,
209        ctx: &EvalContext,
210        op: &dyn Op,
211        inputs: TVec<TValue>,
212    ) -> TractResult<TVec<TValue>> {
213        let input = args_1!(inputs);
214        let op = &op.downcast_ref::<GpuDelay>().ok_or_else(|| format_err!("Wrong Op type"))?.inner;
215        let device_input = input.as_device_tensor().context("Expected a GPU tensor")?;
216        let mut output_shape: TVec<usize> = device_input.shape().into();
217        output_shape[op.axis] = device_input.shape()[op.axis] + op.overlap;
218        let dt = device_input.datum_type();
219        let device = get_context()?;
220        let max_lanes = ctx.seating.max_lanes();
221        if self.buffer.is_none() {
222            let mut shape: TVec<usize> = device_input.shape().into();
223            shape[op.axis] = op.delay + op.overlap;
224            if max_lanes > 1 {
225                ensure!(op.axis > 0, "GpuDelay on axis 0 leaves no axis 0 for the lanes");
226                shape[0] = max_lanes;
227            }
228            self.buffer = Some(Tensor::zero_dt(dt, &shape)?.into_device()?);
229            self.lanes = max_lanes;
230        }
231        ensure!(
232            self.lanes == max_lanes,
233            "GpuDelay buffer holds {} lanes, this turn seats {max_lanes} of them",
234            self.lanes
235        );
236        let mut output = make_tensor_for_node(ctx, dt, &output_shape)?;
237        if max_lanes > 1 {
238            ensure!(
239                device_input.shape()[0] == ctx.seating.occupancy(),
240                "GpuDelay input carries {} streams, this turn seats {}",
241                device_input.shape()[0],
242                ctx.seating.occupancy()
243            );
244        }
245        for ix in 0..ctx.seating.occupancy() {
246            let (seat, lane) = ctx.seating.address(ix);
247            self.delay_seat(&*device, op, device_input, &mut output, seat, lane)?;
248        }
249        Ok(tvec!(output.into_tensor().into()))
250    }
251
252    fn reset_lanes(&mut self, lanes: &[LaneId]) -> TractResult<()> {
253        let Some(buffer) = self.buffer.as_ref() else { return Ok(()) };
254        ensure!(
255            lanes.iter().all(|l| l.0 < self.lanes),
256            "GpuDelay buffer holds {} lanes, asked to reset {lanes:?}",
257            self.lanes
258        );
259        zero_lanes(&*get_context()?, buffer, lanes, self.lanes > 1)
260    }
261}
262
263// ─── GpuPulsePad ─────────────────────────────────────────────────────────────
264
265#[derive(Debug, Clone, PartialEq, Eq)]
266pub struct GpuPulsePad {
267    pub op: PulsePad,
268    pub device_cst: Option<DeviceTensor>,
269}
270
271impl GpuPulsePad {
272    pub fn new(op: &PulsePad, dt: DatumType) -> TractResult<Self> {
273        let device_cst = if let PadMode::Constant(c) = &op.mode {
274            Some(c.cast_to_dt(dt)?.into_owned().into_device()?)
275        } else {
276            None
277        };
278        Ok(Self { op: op.clone(), device_cst })
279    }
280}
281
282impl std::hash::Hash for GpuPulsePad {
283    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
284        self.op.hash(state);
285    }
286}
287
288impl Op for GpuPulsePad {
289    fn name(&self) -> StaticName {
290        "GpuPulsePad".into()
291    }
292
293    fn info(&self) -> TractResult<Vec<String>> {
294        self.op.info()
295    }
296
297    op_as_typed_op!();
298}
299
300impl EvalOp for GpuPulsePad {
301    not_out_of_plan!();
302
303    fn state(&self, ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
304        Ok(Some(Box::new(GpuPulsePadState {
305            node_id: ctx.node_id,
306            current_pos: tvec!(),
307            last_valid_frame: None,
308            lanes: 0,
309        })))
310    }
311}
312
313impl TypedOp for GpuPulsePad {
314    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
315        crate::utils::facts_to_device_facts(inputs, |facts| self.op.output_facts(facts))
316            .with_context(|| format!("Error while computing output facts for {}", self.name()))
317    }
318
319    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
320        crate::utils::get_device_facts(inputs, |facts| self.op.cost(facts))
321    }
322
323    as_op!();
324}
325
326/// One padding state per lane: `current_pos` is each lane's position in its own
327/// stream and `last_valid_frame` each lane's last frame of valid input, for edge
328/// padding. `lanes` is the extent of their lane axis, 1 when the state serves a
329/// single stream and they have no lane axis at all.
330#[derive(Debug, Clone, Hash, PartialEq, Eq)]
331struct GpuPulsePadState {
332    node_id: usize,
333    current_pos: TVec<usize>,
334    last_valid_frame: Option<DeviceTensor>,
335    lanes: usize,
336}
337
338/// Repeat one frame of `src` over `range` along `axis` of one lane of `dst`.
339/// The frame index has to go into the source offset rather than an origin: the
340/// axis stride is zeroed to broadcast the frame, and an origin term on that axis
341/// is scaled by that same stride.
342#[allow(clippy::too_many_arguments)]
343fn repeat_frame(
344    ctx: &dyn DeviceContext,
345    dst: &DeviceTensor,
346    dst_lane: Option<usize>,
347    axis: usize,
348    range: Range<usize>,
349    src: &DeviceTensor,
350    src_lane: Option<usize>,
351    frame: usize,
352) -> TractResult<()> {
353    let mut zone: TVec<usize> = dst.shape().into();
354    zone[axis] = range.len();
355    let mut dst_offset = range.start * dst.strides()[axis] as usize;
356    let mut src_offset = frame * src.strides()[axis] as usize;
357    if let (Some(dst_lane), Some(src_lane)) = (dst_lane, src_lane) {
358        zone[0] = 1;
359        dst_offset += dst_lane * dst.strides()[0] as usize;
360        src_offset += src_lane * src.strides()[0] as usize;
361    }
362    if zone.iter().product::<usize>() == 0 {
363        return Ok(());
364    }
365    let size = dst.datum_type().size_of();
366    let mut src_strides: TVec<isize> = src.strides().into();
367    src_strides[axis] = 0;
368    ctx.copy_nd(src, src_offset * size, &src_strides, dst, dst_offset * size, &zone, dst.strides())
369}
370
371impl GpuPulsePadState {
372    fn save_frame(
373        &mut self,
374        ctx: &dyn DeviceContext,
375        op: &PulsePad,
376        input: &DeviceTensor,
377        frame: usize,
378        seat: Option<usize>,
379        lane: Option<usize>,
380    ) -> TractResult<()> {
381        let frames = match self.last_valid_frame.as_ref() {
382            Some(frames) => frames,
383            None => {
384                let mut shape: TVec<usize> = input.shape().into();
385                shape[op.axis] = 1;
386                if lane.is_some() {
387                    shape[0] = self.lanes;
388                }
389                let frames = Tensor::zero_dt(input.datum_type(), &shape)?.into_device()?;
390                self.last_valid_frame.insert(frames)
391            }
392        };
393        copy_lane(ctx, frames, lane, 0, input, seat, frame, op.axis, 1)
394    }
395
396    fn pad(
397        &mut self,
398        ctx: &EvalContext,
399        gpu_op: &GpuPulsePad,
400        input: &DeviceTensor,
401    ) -> TractResult<DeviceTensor> {
402        let device = get_context()?;
403        let op = &gpu_op.op;
404        let pulse = input.shape()[op.axis];
405        let end_input = op.end_input.eval(ctx.symbols).to_usize().unwrap_or(usize::MAX);
406        let after = op.after.eval(ctx.symbols).to_usize().unwrap_or(usize::MAX);
407        let max_lanes = ctx.seating.max_lanes();
408        if self.lanes == 0 {
409            self.lanes = max_lanes;
410            self.current_pos = tvec!(0; max_lanes);
411        }
412        ensure!(
413            self.lanes == max_lanes,
414            "GpuPulsePad holds {} lanes, this turn seats {max_lanes} of them",
415            self.lanes
416        );
417        let occupancy = ctx.seating.occupancy();
418        if max_lanes > 1 {
419            ensure!(op.axis > 0, "GpuPulsePad on axis 0 leaves no axis 0 for the lanes");
420            ensure!(
421                input.shape()[0] == occupancy,
422                "GpuPulsePad input carries {} streams, this turn seats {occupancy}",
423                input.shape()[0]
424            );
425        }
426        // Seats whose pulse is neither entirely valid input nor entirely outside
427        // it, with the stream position they start at. Every other seat is copied
428        // over unchanged.
429        let mut to_pad: TVec<(Option<usize>, Option<usize>, usize)> = tvec!();
430        for ix in 0..occupancy {
431            let (seat, lane) = ctx.seating.address(ix);
432            let pulse_begin = self.current_pos[lane.unwrap_or(0)];
433            let pulse_end = pulse_begin + pulse;
434            self.current_pos[lane.unwrap_or(0)] += pulse - op.overlap;
435            if let PadMode::Edge = op.mode
436                && after != 0
437                && pulse_begin < end_input
438            {
439                let latest_valid_frame = (end_input - pulse_begin).min(pulse) - 1;
440                self.save_frame(&*device, op, input, latest_valid_frame, seat, lane)?;
441            }
442            let valid = pulse_begin >= op.begin_input && pulse_end <= end_input;
443            let outside = pulse_end <= op.begin_input - op.before
444                || pulse_begin >= end_input.saturating_add(after);
445            if !valid && !outside {
446                to_pad.push((seat, lane, pulse_begin));
447            }
448        }
449
450        // Start with a copy of input.  The fused-axis-op chain may have
451        // installed a non-contiguous view (Move only permutes strides,
452        // never materialises), so a flat memcpy would read the buffer in
453        // pre-Move order; copy_nd honours `input.strides()` instead.
454        let output = make_tensor_for_node(ctx, input.datum_type(), input.shape())?;
455        device.copy_nd(input, 0, input.strides(), &output, 0, input.shape(), output.strides())?;
456
457        for (seat, lane, pulse_begin) in to_pad {
458            if pulse_begin < op.begin_input {
459                let fill_up_to = (op.begin_input - pulse_begin).min(pulse);
460                match &op.mode {
461                    PadMode::Constant(_) => fill_lane(
462                        &*device,
463                        &output,
464                        seat,
465                        op.axis,
466                        0..fill_up_to,
467                        gpu_op.device_cst.as_ref().unwrap(),
468                    )?,
469                    PadMode::Edge => repeat_frame(
470                        &*device,
471                        &output,
472                        seat,
473                        op.axis,
474                        0..fill_up_to,
475                        input,
476                        seat,
477                        fill_up_to,
478                    )?,
479                    _ => unimplemented!(),
480                }
481            }
482
483            if pulse_begin + pulse > end_input && after > 0 {
484                let fill_from = pulse - (pulse_begin + pulse - end_input).min(pulse);
485                match &op.mode {
486                    PadMode::Constant(_) => fill_lane(
487                        &*device,
488                        &output,
489                        seat,
490                        op.axis,
491                        fill_from..pulse,
492                        gpu_op.device_cst.as_ref().unwrap(),
493                    )?,
494                    PadMode::Edge => repeat_frame(
495                        &*device,
496                        &output,
497                        seat,
498                        op.axis,
499                        fill_from..pulse,
500                        self.last_valid_frame.as_ref().unwrap(),
501                        lane,
502                        0,
503                    )?,
504                    _ => unimplemented!(),
505                }
506            }
507        }
508        Ok(output)
509    }
510}
511
512impl OpState for GpuPulsePadState {
513    fn eval(
514        &mut self,
515        ctx: &EvalContext,
516        op: &dyn Op,
517        inputs: TVec<TValue>,
518    ) -> TractResult<TVec<TValue>> {
519        let input = args_1!(inputs);
520        let gpu_op =
521            op.downcast_ref::<GpuPulsePad>().ok_or_else(|| format_err!("Wrong Op type"))?;
522        let device_input = input.as_device_tensor().context("Expected a GPU tensor")?;
523        let output = self.pad(ctx, gpu_op, device_input)?;
524        Ok(tvec!(output.into_tensor().into_tvalue()))
525    }
526
527    fn reset_lanes(&mut self, lanes: &[LaneId]) -> TractResult<()> {
528        if self.lanes == 0 {
529            return Ok(());
530        }
531        ensure!(
532            lanes.iter().all(|l| l.0 < self.lanes),
533            "GpuPulsePad holds {} lanes, asked to reset {lanes:?}",
534            self.lanes
535        );
536        for lane in lanes {
537            self.current_pos[lane.0] = 0;
538        }
539        if let Some(frames) = self.last_valid_frame.as_ref() {
540            zero_lanes(&*get_context()?, frames, lanes, self.lanes > 1)?;
541        }
542        Ok(())
543    }
544}
545
546// ─── GpuAffineChunkTrim ──────────────────────────────────────────────────────
547
548#[derive(Debug, Clone, PartialEq, Eq, Hash)]
549pub struct GpuAffineChunkTrim {
550    pub inner: AffineChunkTrim,
551}
552
553impl GpuAffineChunkTrim {
554    pub fn new(inner: &AffineChunkTrim) -> Self {
555        Self { inner: inner.clone() }
556    }
557}
558
559impl Op for GpuAffineChunkTrim {
560    fn name(&self) -> StaticName {
561        "GpuAffineChunkTrim".into()
562    }
563
564    fn info(&self) -> TractResult<Vec<String>> {
565        self.inner.info()
566    }
567
568    op_as_typed_op!();
569}
570
571impl EvalOp for GpuAffineChunkTrim {
572    op_out_of_plan!();
573
574    fn eval(&self, ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
575        let input_value = args_1!(inputs);
576        let input = input_value.to_device_tensor()?;
577        let axis = self.inner.axis;
578        let n = input.shape()[axis];
579        let take = if n.saturating_sub(self.inner.typed_trim) >= self.inner.target_per_pulse {
580            n - self.inner.typed_trim
581        } else {
582            n
583        };
584        if take == n {
585            return Ok(tvec!(input_value));
586        }
587        let mut o_shape: TVec<usize> = input.shape().into();
588        o_shape[axis] = take;
589        let output = make_tensor_for_node(ctx, input.datum_type(), &o_shape)?;
590        let broadcast_strides = compute_broadcast_strides(&o_shape, input.strides())?;
591        let device = get_context()?;
592        device.copy_nd(
593            input,
594            0,
595            &broadcast_strides,
596            &output,
597            0,
598            output.shape(),
599            output.strides(),
600        )?;
601        Ok(tvec![output.into_tensor().into_tvalue()])
602    }
603}
604
605impl TypedOp for GpuAffineChunkTrim {
606    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
607        crate::utils::facts_to_device_facts(inputs, |facts| self.inner.output_facts(facts))
608            .with_context(|| format!("Error while computing output facts for {}", self.name()))
609    }
610
611    as_op!();
612}
tract_gpu/ops/pulse.rs

tract_gpu/ops/
pulse.rs
