svod-model 0.1.0-alpha.3

Pretrained models inference abstraction.
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
97
98
99
100
101
102
103
104
105
106
107

extern crate self as svod_model;

use svod_macros::jit_wrapper;
use svod_tensor::{PrepareConfig, Tensor};

use crate::jit::InputSpec;

struct AddModel;

impl AddModel {
    fn forward(&self, x: &Tensor, y: &Tensor) -> crate::jit::Result<Tensor> {
        x.try_add(y).map_err(|e| crate::jit::JitError::Tensor { source: Box::new(e) })
    }
}

jit_wrapper! {
    AddJit(AddModel) {
        x: Tensor,
        y: Tensor,

        build(x, y) {
            model.forward(x, y)
        }
    }
}

#[test]
fn test_jit_single_input_prepare_and_execute() {
    let mut jit = AddJit::new(AddModel);

    let cfg = PrepareConfig::default();
    jit.prepare_with_config(InputSpec::f32(&[3]), InputSpec::f32(&[3]), &cfg).unwrap();

    copy_tensor_to_buffer(&Tensor::from_slice([1.0f32, 2.0, 3.0]), jit.x_mut().unwrap());
    copy_tensor_to_buffer(&Tensor::from_slice([10.0f32, 20.0, 30.0]), jit.y_mut().unwrap());

    jit.execute().unwrap();

    let output = jit.output().unwrap();
    let mut result = vec![0.0f32; 3];
    output.copyout(unsafe { std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut u8, 12) }).unwrap();
    assert_eq!(result, vec![11.0, 22.0, 33.0]);
}

#[test]
fn test_jit_replay() {
    let mut jit = AddJit::new(AddModel);
    jit.prepare(InputSpec::f32(&[3]), InputSpec::f32(&[3])).unwrap();

    copy_tensor_to_buffer(&Tensor::from_slice([1.0f32, 2.0, 3.0]), jit.x_mut().unwrap());
    copy_tensor_to_buffer(&Tensor::from_slice([10.0f32, 20.0, 30.0]), jit.y_mut().unwrap());
    jit.execute().unwrap();

    copy_tensor_to_buffer(&Tensor::from_slice([100.0f32, 200.0, 300.0]), jit.x_mut().unwrap());
    copy_tensor_to_buffer(&Tensor::from_slice([1.0f32, 2.0, 3.0]), jit.y_mut().unwrap());
    jit.execute().unwrap();

    let output = jit.output().unwrap();
    let mut result = vec![0.0f32; 3];
    output.copyout(unsafe { std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut u8, 12) }).unwrap();
    assert_eq!(result, vec![101.0, 202.0, 303.0]);
}

#[test]
fn test_jit_multiple_replays() {
    let mut jit = AddJit::new(AddModel);
    jit.prepare(InputSpec::f32(&[3]), InputSpec::f32(&[3])).unwrap();

    for i in 0..5 {
        copy_tensor_to_buffer(&Tensor::from_slice([i as f32; 3]), jit.x_mut().unwrap());
        copy_tensor_to_buffer(&Tensor::from_slice([(i + 1) as f32; 3]), jit.y_mut().unwrap());

        jit.execute().unwrap();

        let output = jit.output().unwrap();
        let mut result = vec![0.0f32; 3];
        output.copyout(unsafe { std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut u8, 12) }).unwrap();
        assert_eq!(result, vec![(i + i + 1) as f32; 3]);
    }
}

#[test]
fn test_jit_profiled_execution_apis() {
    let mut jit = AddJit::new(AddModel);
    jit.prepare(InputSpec::f32(&[3]), InputSpec::f32(&[3])).unwrap();

    copy_tensor_to_buffer(&Tensor::from_slice([1.0f32, 2.0, 3.0]), jit.x_mut().unwrap());
    copy_tensor_to_buffer(&Tensor::from_slice([10.0f32, 20.0, 30.0]), jit.y_mut().unwrap());

    let profiles = jit.execute_profiled().unwrap();
    assert!(!profiles.is_empty(), "expected at least one kernel profile");

    let profiles_with_vars = jit.execute_with_vars_profiled(&[]).unwrap();
    assert!(!profiles_with_vars.is_empty(), "expected at least one kernel profile with vars");

    let output = jit.output().unwrap();
    let mut result = vec![0.0f32; 3];
    output.copyout(unsafe { std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut u8, 12) }).unwrap();
    assert_eq!(result, vec![11.0, 22.0, 33.0]);
}

fn copy_tensor_to_buffer(tensor: &Tensor, dst: &mut svod_device::Buffer) {
    let src_buf = tensor.buffer().unwrap();
    let mut data = vec![0u8; src_buf.size()];
    src_buf.copyout(&mut data).unwrap();
    dst.copyin(&data).unwrap();
}