combs_core/lib.rs
1//! # combs-core
2//!
3//! Backend type aliases, device helpers, and the memory-pool facade for the
4//! Combs Engine L0 Rust core.
5//!
6//! Phase 1 runs entirely on the wgpu backend (Metal on Apple Silicon) with
7//! f32 compute. f16 compute and custom CubeCL kernels are later-phase
8//! optimizations; the aliases here are the single place to change.
9
10pub use burn::backend::wgpu;
11
12pub mod quant;
13
14use burn::backend::wgpu::{RuntimeOptions, WgpuDevice, WgpuSetup, graphics::AutoGraphicsApi};
15
16/// The default inference backend: autotuned, fusing wgpu/CubeCL backend.
17///
18/// With burn 0.21 + the `fusion` feature this expands to
19/// `Fusion<CubeBackend<WgpuRuntime, f32, i32, u32>>`.
20pub type CombsBackend = burn::backend::Wgpu<f32, i32, u32>;
21
22/// The default device handle type.
23pub type CombsDevice = WgpuDevice;
24
25/// Returns the default wgpu device (best available GPU; on macOS this is the
26/// Metal device). Honors cubecl's `CUBECL_WGPU_DEFAULT_DEVICE` override.
27pub fn init_device() -> CombsDevice {
28 WgpuDevice::default()
29}
30
31/// Basic information about a wgpu adapter.
32#[derive(Debug, Clone)]
33pub struct DeviceInfo {
34 /// Human-readable adapter name (e.g. "Apple M3 Pro").
35 pub name: String,
36 /// Graphics backend in use (e.g. "Metal").
37 pub backend: String,
38 /// Device type (e.g. "IntegratedGpu").
39 pub device_type: String,
40 /// Driver name + info string.
41 pub driver: String,
42}
43
44/// Hardware capabilities consumed by the application-layer device planner
45/// (sharding, KV budget, prefill chunk sizing). Serialized to JSON across
46/// the FFI boundary.
47#[derive(Debug, Clone, serde::Serialize)]
48pub struct DeviceCaps {
49 /// Human-readable adapter name (e.g. "Apple M3 Pro").
50 pub name: String,
51 /// Graphics backend in use ("Metal", "Vulkan", "Dx12", "Gl", ...).
52 pub backend: String,
53 /// Device type ("IntegratedGpu", "DiscreteGpu", ...).
54 pub device_type: String,
55 /// `max_storage_buffer_binding_size`: the hard cap on a single GPU
56 /// buffer — the sharding limit on mobile devices.
57 pub max_storage_buffer_binding_size: u64,
58 /// `max_buffer_size`: the largest single allocation the driver allows.
59 pub max_buffer_size: u64,
60 /// Largest compute workgroup dimension.
61 pub max_compute_workgroup_size_x: u32,
62 /// `max_compute_invocations_per_workgroup`.
63 pub max_compute_invocations_per_workgroup: u32,
64 /// Debug dump of the adapter's enabled feature set (wgpu 29 no longer
65 /// exposes WebGPU extension features like `SHADER_F16` through the
66 /// public adapter API, so we surface the raw list for the planner).
67 pub features: String,
68}
69
70/// Queries the adapter for its limits/features and returns [`DeviceCaps`].
71///
72/// Like [`device_info`], this performs the cubecl runtime setup for the
73/// device, so it is safe (and cheap) to use the device afterwards.
74pub fn device_caps(device: &CombsDevice) -> DeviceCaps {
75 let setup: WgpuSetup =
76 burn::backend::wgpu::init_setup::<AutoGraphicsApi>(device, RuntimeOptions::default());
77 let info = setup.adapter.get_info();
78 let limits = setup.adapter.limits();
79 let features = setup.adapter.features();
80 DeviceCaps {
81 name: info.name,
82 backend: format!("{:?}", info.backend),
83 device_type: format!("{:?}", info.device_type),
84 max_storage_buffer_binding_size: limits.max_storage_buffer_binding_size as u64,
85 max_buffer_size: limits.max_buffer_size,
86 max_compute_workgroup_size_x: limits.max_compute_workgroup_size_x,
87 max_compute_invocations_per_workgroup: limits.max_compute_invocations_per_workgroup,
88 features: format!("{:?}", features),
89 }
90}
91
92/// Initializes the wgpu runtime for `device` and returns adapter information.
93///
94/// Note: this performs the cubecl runtime setup for the device (the same setup
95/// burn performs lazily on first tensor use), so it is safe to use the device
96/// for compute afterwards.
97pub fn device_info(device: &CombsDevice) -> DeviceInfo {
98 let setup: WgpuSetup =
99 burn::backend::wgpu::init_setup::<AutoGraphicsApi>(device, RuntimeOptions::default());
100 let info = setup.adapter.get_info();
101 DeviceInfo {
102 name: info.name,
103 backend: format!("{:?}", info.backend),
104 device_type: format!("{:?}", info.device_type),
105 driver: format!("{} ({})", info.driver, info.driver_info),
106 }
107}
108
109/// Facade over the GPU buffer pool.
110///
111/// # Phase 1 status: documented no-op
112///
113/// The plan's hand-rolled slab/coalescing pool was deliberately replaced by
114/// cubecl's built-in allocator, which already does pooled slab allocation and
115/// reuse (configured through [`burn::backend::wgpu::MemoryConfiguration`]).
116/// burn 0.21 does not publicly expose cubecl 0.10's
117/// `ComputeClient::memory_cleanup`, and `memory_persistent_allocations` does
118/// not exist in cubecl 0.10's public API at all, so this facade currently does
119/// nothing. It exists so that the runtime can call `pool.cleanup()` /
120/// `pool.pin_persistent()` today and Phase 2 can back those calls with real
121/// handles (persistent KV/weight arenas) without changing call sites.
122#[derive(Debug, Default, Clone, Copy)]
123pub struct BufferPool;
124
125impl BufferPool {
126 /// Creates a new pool facade.
127 pub fn new() -> Self {
128 BufferPool
129 }
130
131 /// Pin long-lived allocations (weights, KV arena) so the pool never
132 /// releases them. No-op in Phase 1 — cubecl's pooled allocator keeps
133 /// freed blocks for reuse anyway.
134 pub fn pin_persistent(&self) {
135 // no-op: see type-level docs.
136 }
137
138 /// Release cached free blocks back to the driver. No-op in Phase 1 —
139 /// cubecl 0.10's `memory_cleanup` is not reachable through burn's public
140 /// API.
141 pub fn cleanup(&self) {
142 // no-op: see type-level docs.
143 }
144}