tract_gpu/ops/
fused_output.rs

1//! Writing a copy-based op's output in the layout its consumer wants.
2//!
3//! A copy-based op writes its output through `copy_nd`, which takes the
4//! destination's strides, so it can write in any layout those strides describe.
5//! [`FusedOutputLayout`] carries the axis-op chain a consumer needs: the op
6//! declares the chained fact and asks [`make_fused_output_for_node`] for a view
7//! of the arena region to write through, and the chain's nodes disappear --
8//! one copy per axis op becoming none.
9//!
10//! [`fuse_output_axis_op`] moves a chain into its producer. It runs before the
11//! input-side `fuse_axis_op` of each backend, which would otherwise wrap the
12//! producer and hide it.
13
14use tract_core::internal::*;
15
16use crate::fact::DeviceTypedFactExt;
17use crate::rule_ensure;
18use crate::tensor::DeviceTensor;
19use crate::turn_handler::make_tensor_for_node;
20
21use super::change_axes::GpuAxisOp;
22
23/// An op that can write its output through an axis-op chain's strides. The
24/// chain is empty unless [`fuse_output_axis_op`] has moved one in.
25pub trait FusedOutputLayout {
26    fn out_axis_ops(&self) -> &[GpuAxisOp];
27    fn with_out_axis_ops(&self, ops: TVec<GpuAxisOp>) -> Box<dyn TypedOp>;
28}
29
30/// The copy-based ops that honour an output chain.
31pub fn as_fused_output(op: &dyn TypedOp) -> Option<&dyn FusedOutputLayout> {
32    if let Some(op) = op.downcast_ref::<super::pulse::GpuDelay>() {
33        return Some(op);
34    }
35    None
36}
37
38fn eval_axis_op(op: &AxisOp, symbols: &SymbolValues) -> AxisOp {
39    match op {
40        AxisOp::Reshape(skip, from, to) => AxisOp::Reshape(
41            *skip,
42            from.iter().map(|d| d.eval(symbols)).collect(),
43            to.iter().map(|d| d.eval(symbols)).collect(),
44        ),
45        op => op.clone(),
46    }
47}
48
49/// The shape `chain` turns `shape` into.
50pub fn chain_shape(
51    shape: &[usize],
52    chain: &[GpuAxisOp],
53    symbols: &SymbolValues,
54) -> TractResult<TVec<usize>> {
55    let mut shape: TVec<usize> = shape.into();
56    for op in chain {
57        eval_axis_op(&op.inner, symbols).change_shape_array(&mut shape, false)?;
58    }
59    Ok(shape)
60}
61
62fn concrete(dims: &[TDim]) -> TractResult<TVec<usize>> {
63    dims.iter().map(|d| Ok(d.to_usize()?)).collect()
64}
65
66/// The shape and strides addressing a `chain`-shaped tensor in the axis order
67/// its producer writes: `chain` applied backwards, each op's inverse carrying
68/// the shape and the strides together.
69pub fn unchain_strides(
70    shape: &[usize],
71    strides: &[isize],
72    chain: &[GpuAxisOp],
73    symbols: &SymbolValues,
74) -> TractResult<(TVec<usize>, TVec<isize>)> {
75    let mut shape: TVec<usize> = shape.into();
76    let mut strides: TVec<isize> = strides.into();
77    for op in chain.iter().rev() {
78        match eval_axis_op(&op.inner, symbols) {
79            AxisOp::Add(ix) => {
80                ensure!(shape[ix] == 1, "Add({ix}) carries extent {}", shape[ix]);
81                shape.remove(ix);
82                strides.remove(ix);
83            }
84            AxisOp::Rm(ix) => {
85                let stride = strides.get(ix).copied().unwrap_or(1);
86                shape.insert(ix, 1);
87                strides.insert(ix, stride);
88            }
89            AxisOp::Move(from, to) => {
90                let dim = shape.remove(to);
91                shape.insert(from, dim);
92                let stride = strides.remove(to);
93                strides.insert(from, stride);
94            }
95            AxisOp::Reshape(skip, from, to) => {
96                let (from, to) = (concrete(&from)?, concrete(&to)?);
97                let group = &strides[skip..skip + to.len()];
98                let inner = *group.last().context("Reshape to no axis at all")?;
99                let mut nested = inner;
100                for (dim, stride) in to.iter().zip(group).rev() {
101                    ensure!(
102                        *stride == nested,
103                        "Reshaped group is not contiguous, so its split has no strides"
104                    );
105                    nested *= *dim as isize;
106                }
107                let mut split: TVec<isize> = tvec!();
108                let mut stride = inner;
109                for dim in from.iter().rev() {
110                    split.insert(0, stride);
111                    stride *= *dim as isize;
112                }
113                let tail: TVec<usize> = shape[skip + to.len()..].into();
114                shape.truncate(skip);
115                shape.extend(from.iter().copied().chain(tail));
116                let tail: TVec<isize> = strides[skip + to.len()..].into();
117                strides.truncate(skip);
118                strides.extend(split.into_iter().chain(tail));
119            }
120        }
121    }
122    Ok((shape, strides))
123}
124
125/// The view to write the output through, and the tensor to publish. They are
126/// the same tensor when `chain` is empty; otherwise the arena holds the chained
127/// layout and the view addresses it in the op's own axis order.
128pub fn make_fused_output_for_node(
129    ctx: &EvalContext,
130    dt: DatumType,
131    shape: &[usize],
132    chain: &[GpuAxisOp],
133) -> TractResult<(DeviceTensor, DeviceTensor)> {
134    if chain.is_empty() {
135        let output = make_tensor_for_node(ctx, dt, shape)?;
136        return Ok((output.clone(), output));
137    }
138    let chained = chain_shape(shape, chain, ctx.symbols)?;
139    let published = make_tensor_for_node(ctx, dt, &chained)?;
140    let (unchained, strides) =
141        unchain_strides(published.shape(), published.strides(), chain, ctx.symbols)?;
142    ensure!(unchained.as_slice() == shape, "Output chain leads to {unchained:?}, not {shape:?}");
143    let view = published.reshaped(shape.into())?.restrided(strides)?;
144    Ok((view, published))
145}
146
147/// Axis 0 is where a laned turn addresses its seats, so the op has to keep it
148/// as its own leading axis: a chain reordering it or folding it into another axis
149/// cannot be honoured through the strides of one seat.
150fn keeps_lane_axis(op: &GpuAxisOp) -> bool {
151    match op.inner {
152        AxisOp::Add(ix) | AxisOp::Rm(ix) => ix > 0,
153        AxisOp::Move(from, to) => from > 0 && to > 0,
154        AxisOp::Reshape(skip, _, _) => skip > 0,
155    }
156}
157
158/// A copy is only worth fusing while it still writes runs: the op's innermost
159/// axis has to stay the chain's innermost one, or every element lands at its own
160/// stride and the scattered write costs more than the copies it replaces.
161fn keeps_runs(rank: usize, chain: &[GpuAxisOp]) -> bool {
162    let mut axis = rank - 1;
163    let mut rank = rank;
164    for op in chain {
165        let Some(next) = op.inner.transform_axis(axis) else { return false };
166        axis = next;
167        rank = match &op.inner {
168            AxisOp::Add(_) => rank + 1,
169            AxisOp::Rm(_) => rank - 1,
170            AxisOp::Move(_, _) => rank,
171            AxisOp::Reshape(_, from, to) => rank + to.len() - from.len(),
172        };
173    }
174    axis + 1 == rank
175}
176
177/// Move the axis ops between a copy-based producer and its consumer into the
178/// producer, which then writes its output in the consumer's layout.
179pub fn fuse_output_axis_op(
180    _ctx: &(),
181    model: &TypedModel,
182    node: &TypedNode,
183    _node_name: &str,
184    axis_op: &GpuAxisOp,
185) -> TractResult<Option<TypedModelPatch>> {
186    rule_ensure!(keeps_lane_axis(axis_op));
187    rule_if_some!(producer = model.single_prec(node.id)?);
188    rule_if_some!(fused = as_fused_output(producer.op.as_ref()));
189    rule_ensure!(fused.out_axis_ops().is_empty());
190    // An op that may hand its input back writes nothing in that case, so taking a
191    // chain would trade the consumer's copy for one of its own.
192    rule_ensure!(producer.op.forwards_input().is_none());
193    rule_ensure!(model.single_succ(producer.id)?.is_some());
194
195    let mut chain = tvec!(axis_op.clone());
196    let mut cursor = node;
197    while let Some(next) = model.single_succ(cursor.id)? {
198        let Some(next_op) = next.op_as::<GpuAxisOp>().filter(|o| keeps_lane_axis(o)) else {
199            break;
200        };
201        chain.push(next_op.clone());
202        cursor = next;
203    }
204
205    // A Move is the only axis op that changes the layout, so it is the only one
206    // whose copy this saves: the input-side pass folds a bare Add, Rm or Reshape
207    // into its consumer as a view, for nothing.
208    rule_ensure!(chain.iter().any(|op| matches!(op.inner, AxisOp::Move(..))));
209
210    let rank = model.node_output_facts(producer.id)?[0]
211        .as_device_fact()
212        .map(|f| f.shape.rank())
213        .unwrap_or_else(|| model.node_output_facts(producer.id).unwrap()[0].rank());
214    rule_ensure!(keeps_runs(rank, &chain));
215
216    let mut patch = TypedModelPatch::default();
217    let inputs = patch.taps(model, &producer.inputs)?;
218    let out = patch.wire_node(producer.name.clone(), fused.with_out_axis_ops(chain), &inputs)?;
219    patch.shunt_outside(model, cursor.id.into(), out[0])?;
220    Ok(Some(patch))
221}
tract_gpu/ops/fused_output.rs

tract_gpu/ops/
fused_output.rs
