ruDNN 0.1.55

Neural network operators for Ruda.
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

use ruda_kernel::dsl as kernel_dsl;
use std::marker::PhantomData;

use ruda_kernel::dsl::ir::AddressType;
use ruda_kernel::dsl::ir::DeviceProperties;
use ruda_kernel::dsl::server::LaunchError;

use crate::attention::kernel_ir::{
    components::{
        batch::{
            BatchAttentionFamily,
            entry_point::attention,
            simple::{SimpleBatchAttention, config::SimpleBatchConfig},
        },
        global::GlobalAttentionFamily,
    },
    definition::{
        AttentionBlueprint, AttentionElems, AttentionPrecision, AttentionSetupError,
        AttentionVectorSizes, RudaMappingLaunch, InputRuntimeArg, OutputRuntimeArg,
        launch_types::*,
    },
    launch::AttentionArgs,
};

pub struct SimpleBatchAttentionFamily<GA: GlobalAttentionFamily> {
    _phantom: PhantomData<GA>,
}

impl<GA: GlobalAttentionFamily> BatchAttentionFamily for SimpleBatchAttentionFamily<GA> {
    type Attention<AP: AttentionPrecision> = SimpleBatchAttention<AP, GA::Attention<AP>>;
    type Config = SimpleBatchConfig<GA::Config>;
    type Blueprint = AttentionBlueprint;

    unsafe fn launch_unchecked<'a, AA: AttentionArgs, R: ruda_kernel::dsl::Runtime>(
        client: &ruda_kernel::dsl::prelude::ComputeClient<R>,
        ruda_dim: ruda_kernel::dsl::RudaDim,
        ruda_count: ruda_kernel::dsl::RudaCount,
        address_type: AddressType,
        input: InputRuntimeArg<AA, R>,
        output: OutputRuntimeArg<AA, R>,
        ruda_mapping: RudaMappingLaunch<R>,
        dtypes: &AttentionElems,
        vector_sizes: &AttentionVectorSizes,
        blueprint: Self::Blueprint,
    ) -> Result<(), LaunchError> {
        unsafe {
            attention::launch_unchecked::<AA, QG, QGS, KG, KGS, VG, VGS, MSK, MSKS, OG, OGS, Self, R>(
                client,
                ruda_count,
                ruda_dim,
                address_type,
                input,
                output,
                ruda_mapping,
                blueprint,
                dtypes.clone(),
                dtypes.into(),
                vector_sizes.into(),
            )
        };

        Ok(())
    }

    fn expand_config(
        device_props: &DeviceProperties,
        blueprint: Self::Blueprint,
        dtypes: &AttentionElems,
    ) -> Result<Self::Config, AttentionSetupError> {
        let global_config = GA::expand_config(device_props, &blueprint, dtypes)?;

        Ok(SimpleBatchConfig::new(global_config))
    }
}