ruda-kernel 0.1.2

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

use crate::dsl::Runtime;
use super::{allocation::empty_qtensor_optimized, RudaTensor};
use ruda_core::tensor::{TensorMetadata, QuantScheme};

/// Convert the tensor to a lower precision data type based on the quantization scheme and parameters.
pub fn quantize<R>(
    tensor: RudaTensor<R>,
    scheme: &QuantScheme,
    scale: RudaTensor<R>,
) -> RudaTensor<R>
where
    R: Runtime,
{
    let output = empty_qtensor_optimized(tensor.shape(), *scheme, &tensor.device);
    let (out_values, out_params) = output.clone().quantized_handles().unwrap();
    let dtype = tensor.dtype;
    let scale_dtype = scale.dtype;

    crate::quantization::quantize::launch_ref_with_scale_dtype(
        &output.client,
        tensor.binding(),
        out_values.binding(),
        scale.binding(),
        out_params.binding(),
        scheme,
        dtype.into(),
        scale_dtype.into(),
    )
    .expect("Kernel to never fail");

    output
}