singe-cublas 0.1.0-alpha.5

Safe Rust wrappers for the NVIDIA cuBLAS dense linear algebra library (with cuBLASLt).
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

use std::error::Error;

use singe_cublas::{
    lt::{
        context::Context as LtContext,
        descriptor::MatrixLayout,
        matmul::{MatmulDescriptor, matmul},
        types::BatchMode,
    },
    types::{ComputeType, Operation},
};
use singe_cuda::{
    context::Context as CudaContext, data_type::DataType, device::Device, memory::DeviceMemory,
    types::f16,
};

const M: u64 = 2;
const N: u64 = 2;
const K: u64 = 2;
const BATCH_COUNT: i32 = 2;
const WORKSPACE_BYTES: usize = 1 << 20;

fn main() -> Result<(), Box<dyn Error>> {
    let cuda = CudaContext::create_for_device(Device::new(0))?;
    let lt = LtContext::create(&cuda)?;

    let one = f16::from_f32(1.0);
    let two = f16::from_f32(2.0);
    let three = f16::from_f32(3.0);
    let four = f16::from_f32(4.0);

    let a_batches = [
        DeviceMemory::from_slice(&[one, three, two, four])?,
        DeviceMemory::from_slice(&[two, four, one, three])?,
    ];
    let b_batches = [
        DeviceMemory::from_slice(&[one, one, one, one])?,
        DeviceMemory::from_slice(&[one, two, three, four])?,
    ];
    let c_batches = [
        DeviceMemory::<f16>::zeroes((M * N) as usize)?,
        DeviceMemory::<f16>::zeroes((M * N) as usize)?,
    ];
    let mut d_batches = [
        DeviceMemory::<f16>::zeroes((M * N) as usize)?,
        DeviceMemory::<f16>::zeroes((M * N) as usize)?,
    ];

    let a_ptrs = DeviceMemory::from_slice(
        &a_batches
            .iter()
            .map(DeviceMemory::as_ptr)
            .collect::<Vec<*const f16>>(),
    )?;
    let b_ptrs = DeviceMemory::from_slice(
        &b_batches
            .iter()
            .map(DeviceMemory::as_ptr)
            .collect::<Vec<*const f16>>(),
    )?;
    let c_ptrs = DeviceMemory::from_slice(
        &c_batches
            .iter()
            .map(DeviceMemory::as_ptr)
            .collect::<Vec<*const f16>>(),
    )?;
    let mut d_ptrs = DeviceMemory::from_slice(
        &d_batches
            .iter_mut()
            .map(|batch| batch.as_mut_ptr())
            .collect::<Vec<*mut f16>>(),
    )?;
    let mut workspace = DeviceMemory::<u8>::create(WORKSPACE_BYTES)?;

    let mut a_layout = MatrixLayout::create(DataType::F16, M, K, M as i64)?;
    a_layout.set_batch_count(BATCH_COUNT)?;
    a_layout.set_batch_mode(BatchMode::PointerArray)?;

    let mut b_layout = MatrixLayout::create(DataType::F16, K, N, K as i64)?;
    b_layout.set_batch_count(BATCH_COUNT)?;
    b_layout.set_batch_mode(BatchMode::PointerArray)?;

    let mut c_layout = MatrixLayout::create(DataType::F16, M, N, M as i64)?;
    c_layout.set_batch_count(BATCH_COUNT)?;
    c_layout.set_batch_mode(BatchMode::PointerArray)?;

    let mut d_layout = MatrixLayout::create(DataType::F16, M, N, M as i64)?;
    d_layout.set_batch_count(BATCH_COUNT)?;
    d_layout.set_batch_mode(BatchMode::PointerArray)?;

    let mut desc = MatmulDescriptor::create(ComputeType::F32, DataType::F32)?;
    desc.set_transpose_a(Operation::NonTranspose)?;
    desc.set_transpose_b(Operation::NonTranspose)?;

    let alpha = 1.0_f32;
    let beta = 0.0_f32;
    matmul(
        &lt,
        &desc,
        &alpha,
        &a_ptrs,
        &a_layout,
        &b_ptrs,
        &b_layout,
        &beta,
        &c_ptrs,
        &c_layout,
        &mut d_ptrs,
        &d_layout,
        None,
        Some(&mut workspace),
        None,
    )?;

    for (index, batch) in d_batches.iter().enumerate() {
        println!(
            "lt_hshgemm_pointer_array_batch_simple batch {index}: {:?}",
            batch.copy_to_host_vec()?
        );
    }
    Ok(())
}