ruPRIM 0.1.6

Parallel primitives, reductions, scans, and indexing for Ruda.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

use ruda_kernel::dsl as kernel_dsl;
use crate::reduce::{ReduceDtypes, ReduceError, VectorizationMode, routines::ReduceBlueprint};
use ruda_kernel::dsl::prelude::*;

#[derive(Debug)]
pub struct ReduceVectorSettings {
    pub vectorization_mode: VectorizationMode,
    pub vector_size_input: VectorSize,
    pub vector_size_output: VectorSize,
}

#[derive(Debug)]
pub struct ReduceLaunchSettings {
    pub ruda_dim: RudaDim,
    pub ruda_count: RudaCount,
    pub address_type: AddressType,
    pub vector: ReduceVectorSettings,
}

#[derive(Debug)]
pub struct ReduceProblem {
    /// Number of elements in reduce axis
    pub reduce_len: usize,
    /// Number of instances of the reduce axis
    pub reduce_count: usize,
    pub axis: usize,
    pub dtypes: ReduceDtypes,
    /// The address type, defined by the max of each handle's `required_address_type`
    pub address_type: AddressType,
}

#[derive(Debug, Clone)]
pub enum BlueprintStrategy<R: Routine> {
    Forced(R::Blueprint, RudaDim),
    Inferred(R::Strategy),
}

pub trait Routine: core::fmt::Debug + Clone + Sized {
    type Strategy: core::fmt::Debug + Clone + Send + 'static;
    type Blueprint: core::fmt::Debug + Clone + Send + 'static;

    fn prepare<R: Runtime>(
        &self,
        client: &ComputeClient<R>,
        problem: ReduceProblem,
        settings: ReduceVectorSettings,
        strategy: BlueprintStrategy<Self>,
    ) -> Result<(ReduceBlueprint, ReduceLaunchSettings), ReduceError>;
}

pub(crate) fn validate_ruda_dim<R: Runtime>(
    client: &ComputeClient<R>,
    ruda_dim: RudaDim,
) -> Result<(), ReduceError> {
    let hardware = &client.properties().hardware;
    let units = ruda_dim.x.checked_mul(ruda_dim.y).and_then(|xy| xy.checked_mul(ruda_dim.z));
    if ruda_dim.x == 0 || ruda_dim.y == 0 || ruda_dim.z == 0 {
        return Err(ReduceError::Validation {
            details: "Ruda dimensions must be nonzero",
        });
    }
    if ruda_dim.x > hardware.max_ruda_dim.0
        || ruda_dim.y > hardware.max_ruda_dim.1
        || ruda_dim.z > hardware.max_ruda_dim.2
        || units.is_none_or(|units| units > hardware.max_units_per_ruda)
    {
        return Err(ReduceError::Validation {
            details: "Ruda dimensions exceed device limits",
        });
    }
    Ok(())
}