rublas 0.2.5

Linear algebra and matrix multiplication 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
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

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

use crate::kernel_ir::components::{
    batch::partitioned_matmul::config::PartitionedBatchConfig, stage::NumStages,
};
use crate::kernel_ir::{
    components::batch::partitioned_matmul::matmul::PartitionedBatchMatmul,
    components::batch::partitioned_matmul::matmul::matmul_entry,
    components::batch::partitioned_matmul::partition::GlobalPartitionMatmul,
    components::global::GlobalMatmulFamily,
};
use crate::kernel_ir::{
    definition::RudaMappingLaunch,
    definition::MatmulProblem,
    definition::MatmulVectorSizes,
    definition::TilingBlueprint,
    definition::{MatmulElems, MatmulSetupError, MatmulTypes},
    launch::*,
    {components::RudaDimResource, launch::RuntimeConfig},
    {components::batch::BatchMatmulFamily, launch::ConfigRuntimeArg},
};
use ruda_kernel::dsl::ir::DeviceProperties;
use ruda_kernel::dsl::prelude::*;

/// Simple partitioned batch matmul family for any precision
pub struct PartitionedBatchMatmulFamily<
    RC: RuntimeConfig,
    GMM: GlobalMatmulFamily<RC>,
    S: GlobalPartitionMatmul,
> {
    _rc: PhantomData<RC>,
    _gmm: PhantomData<GMM>,
    _s: PhantomData<S>,
}

impl<RC: RuntimeConfig, GMM: GlobalMatmulFamily<RC>, S: GlobalPartitionMatmul> BatchMatmulFamily<RC>
    for PartitionedBatchMatmulFamily<RC, GMM, S>
{
    type Matmul<MP: MatmulTypes> = PartitionedBatchMatmul<RC, MP, GMM::Matmul<MP>, S>;
    type Config = PartitionedBatchConfig<GMM::Config>;
    type Blueprint = TilingBlueprint;

    fn expand_config(
        device_props: &DeviceProperties,
        blueprint: &Self::Blueprint,
        dtypes: &MatmulElems,
        vector_sizes: &MatmulVectorSizes,
    ) -> Result<Self::Config, MatmulSetupError> {
        let global_config = GMM::expand_config(device_props, blueprint, dtypes, vector_sizes)?;

        Ok(PartitionedBatchConfig::new(
            global_config,
            blueprint.tiling_scheme.global_partition_size,
        ))
    }

    fn num_stages() -> NumStages {
        GMM::num_stages()
    }

    unsafe fn launch_unchecked<MA: MatmulArgs<Config = RC>, R: Runtime>(
        client: &ComputeClient<R>,
        ruda_dim: RudaDim,
        ruda_count: RudaCount,
        address_type: AddressType,
        input: InputRuntimeArg<MA, R>,
        output: OutputRuntimeArg<MA, R>,
        config: ConfigRuntimeArg<MA, R>,
        ruda_count_input: RudaMappingLaunch<R>,
        blueprint: Self::Blueprint,
        dtypes: &MatmulElems,
        vector_sizes: &MatmulVectorSizes,
    ) -> Result<(), LaunchError> {
        unsafe {
            matmul_entry::launch_unchecked::<MA, Lhs, LhsSize, Rhs, RhsSize, Acc, AccSize, GMM, S, R>(
                client,
                ruda_count,
                ruda_dim,
                address_type,
                input,
                output,
                config,
                ruda_count_input,
                blueprint,
                dtypes.clone(),
                concat!(
                    include_str!("../../global/memory/iterator.rs"),
                    include_str!("../../global/read/reader/full_reader.rs"),
                    include_str!("../../global/single_stage/simple/matmul.rs"),
                ).to_owned(),
                [dtypes.lhs_global, dtypes.rhs_global, dtypes.acc_global],
                [vector_sizes.lhs, vector_sizes.rhs, vector_sizes.out],
            )
        };

        Ok(())
    }

    fn rudadim_resource(
        blueprint: &Self::Blueprint,
        dtypes: &MatmulElems,
        vector_sizes: &MatmulVectorSizes,
    ) -> Result<RudaDimResource, MatmulSetupError> {
        GMM::rudadim_resource(blueprint, dtypes, vector_sizes)
    }

    fn prepare_kernel<MA: MatmulArgs<Config = RC>, R: Runtime>(
        client: &ComputeClient<R>, ruda_dim: RudaDim, ruda_count: RudaCount,
        address_type: AddressType, input: InputRuntimeArg<MA, R>,
        output: OutputRuntimeArg<MA, R>, config: ConfigRuntimeArg<MA, R>,
        mapping: RudaMappingLaunch<R>, blueprint: Self::Blueprint,
        dtypes: &MatmulElems, vectors: &MatmulVectorSizes,
    ) -> Option<PreparedKernel<R>> {
        Some(matmul_entry::prepare::<MA, Lhs, LhsSize, Rhs, RhsSize, Acc, AccSize, GMM, S, R>(
            client, ruda_count, ruda_dim, address_type, input, output, config, mapping,
            blueprint, dtypes.clone(), concat!(
                include_str!("../../global/memory/iterator.rs"),
                include_str!("../../global/read/reader/full_reader.rs"),
                include_str!("../../global/single_stage/simple/matmul.rs"),
            ).to_owned(), [dtypes.lhs_global, dtypes.rhs_global, dtypes.acc_global],
            [vectors.lhs, vectors.rhs, vectors.out]))
    }

    fn validate_blueprint<R: Runtime>(
        client: &ComputeClient<R>,
        blueprint: &Self::Blueprint,
        problem: &MatmulProblem,
        dtypes: &MatmulElems,
        vector_sizes: &MatmulVectorSizes,
    ) -> Result<(), MatmulSetupError> {
        GMM::validate_blueprint(client, blueprint, problem, dtypes, vector_sizes)
    }
}