enki-gpu 0.2.3

Pure Rust heterogeneous GPU compute platform with JIT compilation
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

use super::core::GpuVec;
use crate::enki_api::type_safety::GpuTypeMatch;
use anu::nam_args_api::{
    AccessIntent, ArgDescriptor, ArgValue, GpuBufferTarget, GpuType, IngressContext,
    InputResourceRecord, ResourceAccessKind,
};

impl<T: Copy + Send + Sync + 'static> GpuType for &GpuVec<T> {
    fn describe() -> ArgDescriptor {
        ArgDescriptor::spmd_cell_const::<T>(false)
    }

    fn collect<'a>(&self, ctx: &mut IngressContext<'a>) {
        let desc = <Self as GpuType>::describe();
        ctx.push_arg(ArgValue::BufferBDA(self.device_address), desc);
        ctx.push_resource_binding(
            format!("GpuVec_Slot_{}", self.slot_index),
            self.slot_index,
            self.state.clone(),
            AccessIntent::Read,
        );
    }
}

impl<T: Copy + Send + Sync + 'static> GpuType for &mut GpuVec<T> {
    fn describe() -> ArgDescriptor {
        ArgDescriptor::spmd_cell_mut::<T>(false)
    }

    fn collect<'a>(&self, ctx: &mut IngressContext<'a>) {
        let desc = <Self as GpuType>::describe();
        ctx.push_arg(ArgValue::BufferBDA(self.device_address), desc);
        ctx.push_resource_binding(
            format!("GpuVec_Slot_{}", self.slot_index),
            self.slot_index,
            self.state.clone(),
            AccessIntent::ReadWrite,
        );
    }
}

impl<T: Copy + Send + Sync + 'static> GpuBufferTarget<T> for GpuVec<T> {}

impl<'target, T: Copy + Send + Sync + 'static> GpuTypeMatch<'target> for &'target GpuVec<T> {
    type Target = &'target T;

    fn describe() -> ArgDescriptor {
        ArgDescriptor::spmd_cell_const::<T>(false)
    }

    fn collect<'a>(&self, ctx: &mut IngressContext<'a>) {
        let desc = <Self as GpuTypeMatch>::describe();
        ctx.push_arg(ArgValue::BufferBDA(self.device_address()), desc);
        ctx.push_resource_binding(
            format!("GpuVec_Slot_{}", self.slot_index()),
            self.slot_index(),
            self.state.clone(),
            AccessIntent::Read,
        );
    }

    fn as_input_record(&self, arg_index: usize) -> InputResourceRecord {
        InputResourceRecord {
            arg_index,
            root_id: Some(self.slot_index() as usize),
            access_kind: ResourceAccessKind::PerCell {
                element_count: self.len,
            },
            is_mutable: false,
            element_type_name: std::any::type_name::<T>(),
        }
    }
}

impl<'target, T: Copy + Send + Sync + 'static> GpuTypeMatch<'target> for &'target mut GpuVec<T> {
    type Target = &'target mut T;

    fn describe() -> ArgDescriptor {
        ArgDescriptor::spmd_cell_mut::<T>(false)
    }

    fn collect<'a>(&self, ctx: &mut IngressContext<'a>) {
        let desc = <Self as GpuTypeMatch>::describe();
        ctx.push_arg(ArgValue::BufferBDA(self.device_address()), desc);
        ctx.push_resource_binding(
            format!("GpuVec_Slot_{}", self.slot_index()),
            self.slot_index(),
            self.state.clone(),
            AccessIntent::ReadWrite,
        );
    }

    fn as_input_record(&self, arg_index: usize) -> InputResourceRecord {
        InputResourceRecord {
            arg_index,
            root_id: Some(self.slot_index as usize),
            access_kind: ResourceAccessKind::PerCell {
                element_count: self.len,
            },
            is_mutable: true,
            element_type_name: std::any::type_name::<T>(),
        }
    }
}