tract_gpu/
utils.rs

1use tract_core::internal::*;
2use tract_core::tract_linalg::block_quant::BlockQuant;
3use tract_linalg::block_quant::{BlockQuantFact, BlockQuantStorage, Q4_0};
4
5use crate::fact::*;
6use crate::tensor::DeviceTensor;
7
8pub fn facts_to_device_facts(
9    facts: &[&TypedFact],
10    resolve_facts: impl Fn(&[&TypedFact]) -> TractResult<TVec<TypedFact>>,
11) -> TractResult<TVec<TypedFact>> {
12    if facts.iter().all(|it| it.as_device_fact().is_some()) {
13        let device_facts = facts
14            .iter()
15            .map(|it| it.to_device_fact().map(|it| it.as_ref()))
16            .collect::<TractResult<TVec<_>>>()?;
17        let output_facts = (resolve_facts)(device_facts.as_slice())?;
18        Ok(output_facts
19            .into_iter()
20            .map(|it| Ok(DeviceFact::new(DeviceTensorOrigin::FromDevice, it)?.into_exotic_fact()))
21            .collect::<TractResult<_>>()?)
22    } else if facts.iter().all(|it| it.as_device_fact().is_none()) {
23        (resolve_facts)(facts)
24    } else {
25        bail!("Inconsistent facts: mix of device and host facts");
26    }
27}
28
29pub fn get_device_facts<'a, 'b: 'a, T>(
30    facts: &'a [&'b TypedFact],
31    map_facts: impl Fn(&[&'b TypedFact]) -> TractResult<T>,
32) -> TractResult<T> {
33    if facts.iter().all(|it| it.as_device_fact().is_some()) {
34        let device_facts = facts
35            .iter()
36            .map(|it| it.to_device_fact().map(|it| it.as_ref()))
37            .collect::<TractResult<TVec<_>>>()?;
38        (map_facts)(device_facts.as_slice())
39    } else if facts.iter().all(|it| it.as_device_fact().is_none()) {
40        (map_facts)(facts)
41    } else {
42        bail!("Inconsistent facts: mix of device and host facts");
43    }
44}
45
46pub fn get_device_fact<'a, T: 'a>(
47    fact: &'a TypedFact,
48    map_fact: impl Fn(&'a TypedFact) -> TractResult<T>,
49) -> TractResult<T> {
50    if fact.as_device_fact().is_some() {
51        (map_fact)(fact.to_device_fact()?)
52    } else {
53        (map_fact)(fact)
54    }
55}
56
57pub fn as_quant_fact<'a>(
58    fact: &'a TypedFact,
59    format: &dyn BlockQuant,
60) -> Option<&'a BlockQuantFact> {
61    fact.exotic_fact
62        .as_ref()
63        .and_then(|of| of.downcast_ref::<BlockQuantFact>())
64        .filter(|&bqf| bqf.format.dyn_eq(format))
65}
66
67pub fn as_q40_tensor(a: &Tensor) -> Option<&BlockQuantStorage> {
68    a.storage_as::<BlockQuantStorage>().filter(|bqs| bqs.format().dyn_eq(&Q4_0))
69}
70
71pub fn get_quant_fact(t: &DeviceTensor, format: &dyn BlockQuant) -> Option<BlockQuantFact> {
72    if let DeviceTensor::Owned(t) = t {
73        t.exotic_fact()
74            .and_then(|of| of.downcast_ref::<BlockQuantFact>())
75            .cloned()
76            .filter(|bqf| bqf.format.dyn_eq(format))
77    } else {
78        None
79    }
80}
81
82// --- Shared array/copy utilities ---
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
85pub enum BroadcastKind {
86    Unicast,
87    ByScalarLeft,
88    ByScalarRight,
89    Nd1,
90    Nd2,
91    Nd3,
92    Nd4,
93    Nd5,
94    Nd6,
95}
96
97impl BroadcastKind {
98    pub const ALL: [BroadcastKind; 9] = [
99        Self::Unicast,
100        Self::ByScalarLeft,
101        Self::ByScalarRight,
102        Self::Nd1,
103        Self::Nd2,
104        Self::Nd3,
105        Self::Nd4,
106        Self::Nd5,
107        Self::Nd6,
108    ];
109
110    pub fn from_rank(rank: usize) -> TractResult<Self> {
111        match rank {
112            1 => Ok(Self::Nd1),
113            2 => Ok(Self::Nd2),
114            3 => Ok(Self::Nd3),
115            4 => Ok(Self::Nd4),
116            5 => Ok(Self::Nd5),
117            6 => Ok(Self::Nd6),
118            _ => bail!("Unsupported rank {rank} for broadcasting"),
119        }
120    }
121
122    pub fn name(&self) -> &'static str {
123        match self {
124            Self::Unicast => "unicast",
125            Self::ByScalarLeft => "by_scalar_lhs",
126            Self::ByScalarRight => "by_scalar_rhs",
127            Self::Nd1 => "nd1",
128            Self::Nd2 => "nd2",
129            Self::Nd3 => "nd3",
130            Self::Nd4 => "nd4",
131            Self::Nd5 => "nd5",
132            Self::Nd6 => "nd6",
133        }
134    }
135
136    /// Map datum type to the copy kernel type name based on element size.
137    /// Copy kernels only care about element size, not the actual type.
138    pub fn copy_tname(dt: DatumType) -> &'static str {
139        match dt.size_of() {
140            1 => "u8",
141            2 => "u16",
142            4 => "u32",
143            8 => "u64",
144            _ => panic!("Unsupported element size {} for copy kernel", dt.size_of()),
145        }
146    }
147
148    pub fn copy_kernel_name(&self, dt: DatumType, prefix: &str) -> TractResult<String> {
149        Ok(format!("{prefix}copy_{}_{}", self.name(), Self::copy_tname(dt)))
150    }
151
152    pub fn all_copy_kernel_names(prefix: &str) -> Vec<String> {
153        let copy_types = ["u8", "u16", "u32", "u64"];
154        Self::ALL
155            .into_iter()
156            .flat_map(|bk| {
157                copy_types
158                    .into_iter()
159                    .map(move |tname| format!("{prefix}copy_{}_{tname}", bk.name()))
160            })
161            .collect()
162    }
163}
164
165/// One tensor's layout as a kernel addresses it: an extent per axis and a stride
166/// per axis, a broadcast axis holding a stride of zero.
167pub struct Layout<'a> {
168    pub shape: &'a mut TVec<usize>,
169    pub strides: &'a mut TVec<isize>,
170}
171
172/// Merges the adjacent axes every layout walks as one run, until the rank fits
173/// `max_rank`, and answers the rank reached.
174///
175/// A pair merges when, for every layout, either the outer stride is the inner
176/// one times its extent -- so the two axes spell out one contiguous run -- or the
177/// axis is broadcast on both halves and stays broadcast. Merging is exact: it
178/// renames a walk the kernel was already going to make, and the extents the
179/// kernel is given are the output's.
180pub fn merge_axes_to_fit(layouts: &mut [Layout], max_rank: usize) -> usize {
181    let mut rank = layouts.first().map(|l| l.shape.len()).unwrap_or(0);
182    let mut axis = 0;
183    while rank > max_rank && axis + 1 < rank {
184        let inner_extent = layouts.last().map(|l| l.shape[axis + 1]).unwrap_or(1);
185        let runs = |l: &Layout| {
186            l.strides[axis] == l.strides[axis + 1] * inner_extent as isize
187                || (l.strides[axis] == 0 && l.strides[axis + 1] == 0)
188        };
189        if !layouts.iter().all(runs) {
190            axis += 1;
191            continue;
192        }
193        for l in layouts.iter_mut() {
194            // A layout broadcast along both halves is broadcast along the merged
195            // axis, and says so with an extent of one.
196            l.shape[axis] = if l.strides[axis] == 0 && l.strides[axis + 1] == 0 {
197                1
198            } else {
199                l.shape[axis] * l.shape[axis + 1]
200            };
201            l.shape.remove(axis + 1);
202            l.strides[axis] = l.strides[axis + 1];
203            l.strides.remove(axis + 1);
204        }
205        rank -= 1;
206    }
207    rank
208}
209
210/// Highest rank the backends' `copy_nd` kernels address.
211pub const MAX_COPY_RANK: usize = 6;
212
213/// A `copy_nd` zone: the output's extents, and the strides each side walks it by.
214pub struct CopyGeometry {
215    pub shape: TVec<usize>,
216    pub input_strides: TVec<isize>,
217    pub output_strides: TVec<isize>,
218}
219
220/// Merges a copy's adjacent axes with [`merge_axes_to_fit`] until it fits
221/// [`MAX_COPY_RANK`]. The rank answered still exceeds it when the two layouts
222/// leave no pair to merge, which a caller must check before dispatching.
223pub fn merge_copy_axes(
224    input_strides: &[isize],
225    output_shape: &[usize],
226    output_strides: &[isize],
227) -> CopyGeometry {
228    let mut in_shape: TVec<usize> = output_shape.into();
229    let mut geo = CopyGeometry {
230        shape: output_shape.into(),
231        input_strides: input_strides.into(),
232        output_strides: output_strides.into(),
233    };
234    merge_axes_to_fit(
235        &mut [
236            Layout { shape: &mut in_shape, strides: &mut geo.input_strides },
237            Layout { shape: &mut geo.shape, strides: &mut geo.output_strides },
238        ],
239        MAX_COPY_RANK,
240    );
241    geo
242}
243
244pub fn compute_broadcast_strides<T: num_traits::Zero + Copy + 'static>(
245    shape: &[usize],
246    strides: &[isize],
247) -> TractResult<TVec<T>>
248where
249    isize: num_traits::AsPrimitive<T>,
250{
251    use num_traits::AsPrimitive;
252    ensure!(
253        shape.len() == strides.len(),
254        "Mismatch between shape and strides length while computing broadcast strides"
255    );
256    Ok(strides
257        .iter()
258        .zip(shape)
259        .map(|(s, dim)| if *dim == 1 { T::zero() } else { s.as_() })
260        .collect::<TVec<T>>())
261}
262
263pub fn reshape_to_rank_2(shape: &[usize], axis: usize) -> TVec<usize> {
264    let dim_axis_0 = shape[0..axis].iter().product::<usize>();
265    let dim_axis_2 = shape[axis..].iter().product::<usize>();
266    tvec![dim_axis_0, dim_axis_2]
267}
268
269pub fn reshape_to_rank_3(shape: &[usize], axis: usize) -> TVec<usize> {
270    let dim_axis_0 = shape[0..axis].iter().product::<usize>();
271    let dim_axis_1 = shape[axis];
272    let dim_axis_2 = shape[axis + 1..].iter().product::<usize>();
273    tvec![dim_axis_0, dim_axis_1, dim_axis_2]
274}
275
276pub fn check_strides_validity(shape: TVec<usize>, strides: TVec<isize>) -> TractResult<()> {
277    let mut zipped_shape_strides: Vec<_> = shape.into_iter().zip(strides).collect();
278    zipped_shape_strides.sort_by_key(|&(_, stride)| stride);
279
280    let mut prev_stride = 1;
281    for (dim, stride) in zipped_shape_strides {
282        ensure!((stride == prev_stride) || (dim == 1), "Invalid strides");
283        prev_stride *= dim as isize;
284    }
285    Ok(())
286}
tract_gpu/utils.rs

tract_gpu/
utils.rs
