ruprim/reduce/routines/
base.rs

1use ruda_kernel::dsl as kernel_dsl;
2use crate::reduce::{ReduceDtypes, ReduceError, VectorizationMode, routines::ReduceBlueprint};
3use ruda_kernel::dsl::prelude::*;
4
5#[derive(Debug)]
6pub struct ReduceVectorSettings {
7    pub vectorization_mode: VectorizationMode,
8    pub vector_size_input: VectorSize,
9    pub vector_size_output: VectorSize,
10}
11
12#[derive(Debug)]
13pub struct ReduceLaunchSettings {
14    pub ruda_dim: RudaDim,
15    pub ruda_count: RudaCount,
16    pub address_type: AddressType,
17    pub vector: ReduceVectorSettings,
18}
19
20#[derive(Debug)]
21pub struct ReduceProblem {
22    /// Number of elements in reduce axis
23    pub reduce_len: usize,
24    /// Number of instances of the reduce axis
25    pub reduce_count: usize,
26    pub axis: usize,
27    pub dtypes: ReduceDtypes,
28    /// The address type, defined by the max of each handle's `required_address_type`
29    pub address_type: AddressType,
30}
31
32#[derive(Debug, Clone)]
33pub enum BlueprintStrategy<R: Routine> {
34    Forced(R::Blueprint, RudaDim),
35    Inferred(R::Strategy),
36}
37
38pub trait Routine: core::fmt::Debug + Clone + Sized {
39    type Strategy: core::fmt::Debug + Clone + Send + 'static;
40    type Blueprint: core::fmt::Debug + Clone + Send + 'static;
41
42    fn prepare<R: Runtime>(
43        &self,
44        client: &ComputeClient<R>,
45        problem: ReduceProblem,
46        settings: ReduceVectorSettings,
47        strategy: BlueprintStrategy<Self>,
48    ) -> Result<(ReduceBlueprint, ReduceLaunchSettings), ReduceError>;
49}
50
51pub(crate) fn validate_ruda_dim<R: Runtime>(
52    client: &ComputeClient<R>,
53    ruda_dim: RudaDim,
54) -> Result<(), ReduceError> {
55    let hardware = &client.properties().hardware;
56    let units = ruda_dim.x.checked_mul(ruda_dim.y).and_then(|xy| xy.checked_mul(ruda_dim.z));
57    if ruda_dim.x == 0 || ruda_dim.y == 0 || ruda_dim.z == 0 {
58        return Err(ReduceError::Validation {
59            details: "Ruda dimensions must be nonzero",
60        });
61    }
62    if ruda_dim.x > hardware.max_ruda_dim.0
63        || ruda_dim.y > hardware.max_ruda_dim.1
64        || ruda_dim.z > hardware.max_ruda_dim.2
65        || units.is_none_or(|units| units > hardware.max_units_per_ruda)
66    {
67        return Err(ReduceError::Validation {
68            details: "Ruda dimensions exceed device limits",
69        });
70    }
71    Ok(())
72}
ruprim/reduce/routines/base.rs

ruprim/reduce/routines/
base.rs
