ruda-kernel 0.1.1

Ruda Rust kernel DSL and device tensor operations.
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

use crate::dsl::Runtime;
use super::{RudaTensor, allocation::empty_device_dtype};
use ruda_core::tensor::{Shape, Metadata, ReshapeAction, reshape_action};

/// Reshape a jit tensor to a new shape
pub fn reshape<R: Runtime>(mut tensor: RudaTensor<R>, shape: Shape) -> RudaTensor<R> {
    let analysis = reshape_action(tensor.meta.shape(), tensor.meta.strides(), &shape);

    match analysis {
        ReshapeAction::UpdateStrides { strides } => {
            *tensor.meta = Metadata::new(shape, strides);
            return tensor;
        }
        ReshapeAction::NoChange => return tensor,
        ReshapeAction::Recompute => (),
    }

    let out = empty_device_dtype(
        tensor.client.clone(),
        tensor.device.clone(),
        shape,
        tensor.dtype,
    );

    crate::library::tensor::copy_into(
        &out.client,
        tensor.binding(),
        out.clone().binding(),
        out.dtype.into(),
    );

    out
}

/// Reshape a quantized tensor without recalibrating its values.
pub fn q_reshape<R: Runtime>(tensor: RudaTensor<R>, shape: Shape) -> RudaTensor<R> {
    try_q_reshape(tensor, shape).unwrap_or_else(|_| {
        panic!("No scale-preserving reshape path for this block layout")
    })
}

/// Returns the original tensor and target shape when no scale-preserving path is selected.
pub fn try_q_reshape<R: Runtime>(
    tensor: RudaTensor<R>,
    shape: Shape,
) -> Result<RudaTensor<R>, (RudaTensor<R>, Shape)> {
    super::reshape_quantized::try_reshape(tensor, shape)
}