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
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
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use ruda_core::tensor::{QTensorPrimitive, TensorMetadata};
use crate::dsl::{Runtime, server::MemoryLayoutStrategy};
use super::{RudaTensor, allocation::empty_qtensor};

/// Make a jit tensor contiguous.
pub fn into_contiguous<R: Runtime>(tensor: RudaTensor<R>) -> RudaTensor<R> {
    if tensor.qparams.is_some() {
        let (values, scales) = tensor.quantized_handles().unwrap();
        if values.is_contiguous() && scales.is_contiguous() {
            return tensor;
        }
        return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Contiguous);
    }

    if tensor.is_contiguous() {
        return tensor;
    }

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

    let output = crate::library::tensor::into_contiguous(&client, tensor.binding(), dtype.into());

    RudaTensor::new(
        client.clone(),
        output.handle,
        *output.metadata,
        device,
        dtype,
    )
}

/// Make a jit tensor contiguous with an aligned last stride. Tensor is considered already contiguous
/// if runtime can read it as is. This is equivalent in practice.
#[cfg_attr(
    feature = "device-tensor-tracing",
    tracing::instrument(level = "trace", skip(tensor))
)]
pub fn into_contiguous_aligned<R: Runtime>(tensor: RudaTensor<R>) -> RudaTensor<R> {
    if tensor.qparams.is_some() {
        let (values, scales) = tensor.quantized_handles().unwrap();
        if R::can_read_tensor(values.meta.shape(), values.meta.strides())
            && R::can_read_tensor(scales.meta.shape(), scales.meta.strides())
        {
            return tensor;
        }
        return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Optimized);
    }

    if R::can_read_tensor(tensor.meta.shape(), tensor.meta.strides()) {
        return tensor;
    }

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

    let output =
        crate::library::tensor::into_contiguous_pitched(&client, tensor.binding(), dtype.into());

    RudaTensor::new(
        client.clone(),
        output.handle,
        *output.metadata,
        device,
        dtype,
    )
}

#[cfg_attr(
    feature = "device-tensor-tracing",
    tracing::instrument(level = "trace", skip(tensor))
)]
fn into_contiguous_quantized<R: Runtime>(
    tensor: RudaTensor<R>,
    strategy: MemoryLayoutStrategy,
) -> RudaTensor<R> {
    let output = empty_qtensor(tensor.shape(), *tensor.scheme(), &tensor.device, strategy);
    let (values, scales) = tensor.quantized_handles().unwrap();
    let (out_values, out_scales) = output.quantized_handles().unwrap();

    let (client, dtype_scales, dtype_value) = (scales.client.clone(), scales.dtype, values.dtype);

    crate::library::tensor::copy_into(
        &client,
        values.binding(),
        out_values.binding(),
        dtype_value.into(),
    );

    crate::library::tensor::copy_into(
        &client,
        scales.binding(),
        out_scales.binding(),
        dtype_scales.into(),
    );

    output
}