cubecl-server 0.11.1

Toolkit for implementing a CubeCL runtime: memory pools, streams, drivers and the compilation pipeline.
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

use super::{DummyServer, Marker};
use cubecl_common::device::{Device, DeviceService};
use cubecl_ir::MemoryDeviceProperties;
use cubecl_server::server::Server;
use cubecl_server::{
    client::Client,
    logging::ServerLogger,
    memory_management::{MemoryConfiguration, MemoryManagement, MemoryManagementOptions},
    runtime::Runtime,
    storage::BytesStorage,
};
use cubecl_zspace::Shape;
use cubecl_zspace::Strides;
use std::sync::Arc;

/// The dummy device.
#[derive(Clone, Debug, Hash, PartialEq, Eq, Default)]
pub struct DummyDevice;

impl Device for DummyDevice {
    fn from_id(_device_id: cubecl_common::device::DeviceId) -> Self {
        Self
    }

    fn to_id(&self) -> cubecl_common::device::DeviceId {
        cubecl_common::device::DeviceId {
            type_id: 0,
            index_id: 0,
        }
    }
}

pub type DummyClient = Client;

impl<M: Marker> DeviceService for DummyServer<M> {
    fn init(device_id: cubecl_common::device::DeviceId) -> Self {
        init_server(cubecl_common::device::ServiceId::of::<Self>(device_id))
    }

    fn utilities(&self) -> Arc<dyn std::any::Any + Send + Sync> {
        Server::utilities(self) as Arc<dyn std::any::Any + Send + Sync>
    }
}

fn init_server<M: Marker>(service: cubecl_common::device::ServiceId) -> DummyServer<M> {
    let storage = BytesStorage::default();
    let mem_properties = MemoryDeviceProperties::new(1024 * 1024 * 512, 32);

    let memory_management = MemoryManagement::from_configuration(
        storage,
        &mem_properties,
        MemoryConfiguration::default(),
        Arc::new(ServerLogger::default()),
        MemoryManagementOptions::new("Main CPU Memory"),
    );
    DummyServer::new(service, memory_management, mem_properties)
}

pub fn test_client(device: &DummyDevice) -> DummyClient {
    DummyRuntime::client(device)
}

#[derive(Debug, Clone)]
pub struct DummyRuntime;

impl Runtime for DummyRuntime {
    type Server = DummyServer;

    type Device = DummyDevice;

    fn can_read_tensor(_shape: &Shape, _strides: &Strides) -> bool {
        unimplemented!()
    }

    fn target_properties() -> cubecl_ir::TargetProperties {
        unimplemented!()
    }

    fn enumerate_devices(_: u16) -> Vec<cubecl_common::device::DeviceId> {
        vec![cubecl_common::device::DeviceId {
            type_id: 0,
            index_id: 0,
        }]
    }
}