cubek-reduce 0.3.0-pre.2

CubeK: Reduce Kernels
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
73
74
75

use cubecl::{ir::StorageType, server::LaunchError};
use thiserror::Error;

#[derive(Error, Debug, Clone)]
/// This error should be caught and properly handled.
pub enum ReduceError {
    /// Indicate that the hardware / API doesn't support SIMT plane instructions.
    #[error(
        "Trying to launch a kernel using plane instructions, but there are not supported by the hardware."
    )]
    PlanesUnavailable,
    /// When the cube count is bigger than the max supported.
    #[error("The cube count is larger than the max supported.")]
    CubeCountTooLarge,

    /// A generic validation error
    #[error("A generic validation error: {details}")]
    Validation { details: &'static str },

    /// Indicate that min_plane_dim != max_plane_dim, thus the exact plane_dim is not fixed.
    #[error(
        "Trying to launch a kernel using plane instructions, but the min and max plane dimensions are different."
    )]
    ImprecisePlaneDim,
    /// Indicate the axis is too large.
    #[error("The provided axis ({axis}) must be smaller than the input tensor rank ({rank}).")]
    InvalidAxis { axis: usize, rank: usize },
    /// Indicate that the shape of the input tensor is too small for the given input and axis.
    #[error(
        "The input reduce axis length (currently {axis_length:?}) should be at least k ({k:?})."
    )]
    ReduceAxisTooSmall { axis_length: usize, k: usize },
    /// Indicate that the shape of the output tensor is invalid for the given input and axis.
    #[error("The output shape (currently {output_shape:?}) should be {expected_shape:?}.")]
    MismatchOutputShape {
        expected_shape: Vec<usize>,
        output_shape: Vec<usize>,
    },
    /// Indicate that the indices output shape does not match the values output shape.
    #[error(
        "The indices shape (currently {indices_shape:?}) should match the values shape ({values_shape:?})."
    )]
    MismatchIndicesShape {
        values_shape: Vec<usize>,
        indices_shape: Vec<usize>,
    },
    /// Indicate that the indices output strides do not match the values output strides.
    #[error(
        "The indices strides (currently {indices_strides:?}) should match the values strides ({values_strides:?})."
    )]
    MismatchIndicesStrides {
        values_strides: Vec<usize>,
        indices_strides: Vec<usize>,
    },
    /// Indicate the operation cannot report indices alongside its values.
    #[error(
        "The operation {operation} has no index to report; reduce_with_indices only supports TopK, ArgTopK, Max, ArgMax, Min and ArgMin."
    )]
    IndicesUnsupported { operation: &'static str },

    /// Indicate that we can't launch a shared sum because the atomic addition is not supported.
    #[error("Atomic add not supported by the client for {0}")]
    MissingAtomicAdd(StorageType),

    /// The selected blueprint stages more accumulators in shared memory than the
    /// device allows.
    #[error(
        "Reduce blueprint requires {requested} bytes of shared memory, but only {available} bytes are available on the device."
    )]
    SharedMemoryOverflow { requested: usize, available: usize },

    /// An error happened during launch.
    #[error("An error happened during launch\nCaused by:\n  {0}")]
    Launch(LaunchError),
}