ruprim/reduce/
mod.rs

1//! This provides different implementations of the reduce algorithm which
2//! can run on multiple GPU backends using Ruda.
3//!
4//! A reduction is a tensor operation mapping a rank `R` tensor to a rank `R - 1`
5//! by agglomerating all elements along a given axis with some binary operator.
6//! This is often also called folding.
7//!
8//! This crate provides a main entrypoint as the [`reduce`] function which allows to automatically
9//! perform a reduction for a given instruction implementing the [`ReduceInstruction`] trait and a given [`ReduceStrategy`].
10//! It also provides implementation of the [`ReduceInstruction`] trait for common operations in the [`instructions`] module.
11//! Finally, it provides many reusable primitives to perform different general reduction algorithms in the [`primitives`] module.
12
13#![allow(
14    clippy::type_complexity,
15    reason = "Too sensitive, triggers on tuple of vector."
16)]
17
18use ruda_kernel::dsl as kernel_dsl;
19pub mod components;
20pub mod launch;
21pub mod routines;
22
23#[cfg(feature = "tensor-reduce")]
24pub mod tensor;
25
26mod error;
27
28#[cfg(feature = "cpu-reference")]
29pub mod cpu_reference;
30
31pub use crate::reduce::launch::ReduceStrategy;
32use crate::reduce::{components::instructions::ReduceOperationConfig, launch::launch_reduce};
33pub use components::{
34    args::init_tensors,
35    config::*,
36    instructions::{ReduceFamily, ReduceInstruction},
37    precision::ReducePrecision,
38};
39use ruda_kernel::dsl::prelude::*;
40pub use error::*;
41pub use launch::{ReduceDtypes, reduce_kernel};
42pub use routines::shared_sum::shared_sum;
43
44/// Reduce the given `axis` of the `input` tensor using the instruction `Inst` and write the result into `output`.
45///
46/// An optional [`ReduceStrategy`] can be provided to force the reduction to use a specific algorithm. If omitted, a best effort
47/// is done to try and pick the best strategy supported for the provided `client`.
48///
49/// Return an error if `strategy` is `Some(strategy)` and the specified strategy is not supported by the `client`.
50/// Also returns an error if the `axis` is larger than the `input` rank or if the shape of `output` is invalid.
51/// The shape of `output` must be the same as input except with a value of 1 for the given `axis`.
52///
53///
54/// # Example
55///
56/// This examples show how to sum the rows of a small `2 x 2` matrix into a `1 x 2` vector.
57/// For more details, see the Ruda documentation.
58///
59/// ```ignore
60/// use ruprim::instructions::Sum;
61///
62/// let client = /* ... */;
63/// let size_f32 = std::mem::size_of::<f32>();
64/// let axis = 0; // 0 for rows, 1 for columns in the case of a matrix.
65///
66/// // Create input and output handles.
67/// let input_handle = client.create(f32::as_bytes(&[0, 1, 2, 3]));
68/// let input = unsafe {
69///     TensorBinding::from_raw_parts(
70///         &input_handle,
71///         &[2, 1],
72///         &[2, 2],
73///         size_f32,
74///     )
75/// };
76///
77/// let output_handle = client.empty(2 * size_f32);
78/// let output = unsafe {
79///     TensorBinding::from_raw_parts(
80///         &output_handle,
81///         &output_stride,
82///         &output_shape,
83///         size_f32,
84///     )
85/// };
86///
87/// // Here `R` is a `ruda_kernel::dsl::Runtime`.
88/// let result = reduce::<R, f32, f32, Sum>(&client, input, output, axis, None);
89///
90/// if result.is_ok() {
91///        let binding = output_handle.binding();
92///        let bytes = client.read_one(binding);
93///        let output_values = f32::from_bytes(&bytes);
94///        println!("Output = {:?}", output_values); // Should print [1, 5].
95/// }
96/// ```
97pub fn reduce<R: Runtime>(
98    client: &ComputeClient<R>,
99    input: TensorBinding<R>,
100    output: TensorBinding<R>,
101    axis: usize,
102    strategy: ReduceStrategy,
103    operation: ReduceOperationConfig,
104    dtypes: ReduceDtypes,
105) -> Result<(), ReduceError> {
106    validate_axis(input.shape.len(), axis)?;
107    validate_shapes(
108        &input.shape,
109        &output.shape,
110        axis,
111        match operation {
112            ReduceOperationConfig::ArgTopK(k) => Some(k),
113            ReduceOperationConfig::TopK(k) => Some(k),
114            _ => None,
115        },
116    )?;
117
118    launch_reduce::<R>(client, input, output, axis, strategy, dtypes, operation)
119}
120
121// Check that the given axis is less than the rank of the input.
122fn validate_axis(rank: usize, axis: usize) -> Result<(), ReduceError> {
123    if axis >= rank {
124        return Err(ReduceError::InvalidAxis { axis, rank });
125    }
126    Ok(())
127}
128
129// Check that the output shape match the input shape with the given axis set to 1.
130fn validate_shapes(
131    input_shape: &[usize],
132    output_shape: &[usize],
133    axis: usize,
134    k: Option<usize>,
135) -> Result<(), ReduceError> {
136    let mut expected_shape = input_shape.to_vec();
137    let k = k.unwrap_or(1);
138
139    if expected_shape[axis] < k {
140        return Err(ReduceError::ReduceAxisTooSmall {
141            axis_length: expected_shape[axis],
142            k,
143        });
144    }
145
146    expected_shape[axis] = k;
147    if output_shape != expected_shape {
148        return Err(ReduceError::MismatchOutputShape {
149            expected_shape,
150            output_shape: output_shape.to_vec(),
151        });
152    }
153    Ok(())
154}
ruprim/reduce/mod.rs

ruprim/reduce/
mod.rs
