1use alloc::boxed::Box;
2use ruda_core::device::Device;
3use ruda_tensor::{AutodiffContext, DeviceId, DeviceOps, GradientCheckpointingStrategy};
4
5#[derive(Clone, Debug)]
6pub enum DispatchDevice {
7Host(ruda_tensor_host::HostDevice),
8#[cfg(feature = "cuda")]
9Cuda(ruda_tensor_device::cuda::CudaDevice),
10#[cfg(feature = "cpu")]
11Cpu(ruda_driver_cpu::CpuDevice),
12#[cfg(feature = "wgpu")]
13Wgpu(ruda_tensor_wgpu::WgpuDevice),
14#[cfg(feature = "rocm")]
15Rocm(ruda_tensor_rocm::RocmDevice),
16#[cfg(feature = "tch")]
17LibTorch(ruda_tensor_tch::LibTorchDevice),
18#[cfg(feature = "remote")]
19Remote(ruda_tensor_remote::RemoteDevice),
20#[cfg(feature = "capture")]
21Capture(ruda_tensor_capture::CaptureDevice),
22 Autodiff(Box<DispatchDevice>, GradientCheckpointingStrategy),
23}
24
25impl DispatchDevice {
26 #[cfg(feature = "throughput")]
27 pub fn performance_stats(
28 &self,
29 keys: &[crate::ThroughputKey],
30 ) -> alloc::vec::Vec<Result<crate::ThroughputValue, crate::ThroughputError>> {
31 #[allow(unused_imports)]
32 use ruda_kernel::library::throughput::device_throughput;
33 let _ = keys;
34 match self {
35 Self::Host(_) => alloc::vec::Vec::new(),
36 #[cfg(feature = "cuda")]
37 Self::Cuda(device) => device_throughput::<ruda_tensor_device::cuda::CudaRuntime>(device, keys),
38 #[cfg(feature = "cpu")]
39 Self::Cpu(device) => device_throughput::<ruda_driver_cpu::CpuRuntime>(device, keys),
40 #[cfg(feature = "wgpu")]
41 Self::Wgpu(device) => device_throughput::<ruda_tensor_wgpu::WgpuRuntime>(device, keys),
42 #[cfg(feature = "rocm")]
43 Self::Rocm(device) => device_throughput::<ruda_tensor_rocm::HipRuntime>(device, keys),
44 #[cfg(feature = "tch")]
45 Self::LibTorch(_) => alloc::vec::Vec::new(),
46 #[cfg(feature = "remote")]
47 Self::Remote(_) => alloc::vec::Vec::new(),
48 #[cfg(feature = "capture")]
49 Self::Capture(_) => alloc::vec::Vec::new(),
50 Self::Autodiff(device, _) => device.performance_stats(keys),
51 }
52 }
53
54 pub fn identity(&self) -> Option<ruda_core::ir::DeviceIdentity> {
55 match self.inner() {
56 Self::Host(_) => None,
57 #[cfg(feature = "cuda")]
58 Self::Cuda(device) => Some(device.identity()),
59 #[cfg(feature = "cpu")]
60 Self::Cpu(device) => Some(device.identity()),
61 #[cfg(feature = "wgpu")]
62 Self::Wgpu(device) => Some(device.identity()),
63 #[cfg(feature = "rocm")]
64 Self::Rocm(device) => Some(device.identity()),
65 #[cfg(feature = "tch")]
66 Self::LibTorch(_) => None,
67 #[cfg(feature = "remote")]
68 Self::Remote(_) => None,
69 #[cfg(feature = "capture")]
70 Self::Capture(_) => None,
71 Self::Autodiff(..) => unreachable!(),
72 }
73 }
74
75 #[cfg(feature = "capture")]
76 pub fn capture() -> Self { Self::Capture(Default::default()) }
77
78 #[cfg(feature = "capture")]
79 pub fn capture_scope(&self, capture: impl FnOnce(ruda_tensor_capture::CaptureScope) -> ruda_tensor_capture::CompletedCaptureScope)
80 -> Result<ruda_tensor_capture::CapturedGraph, ruda_tensor_capture::CaptureError> {
81 match self {
82 Self::Capture(device) => device.capture_scope(capture),
83 _ => Err(ruda_tensor_capture::CaptureError::InvalidDevice),
84 }
85 }
86
87 pub fn autodiff(device: impl Into<Self>) -> Self {
88 Self::autodiff_with_gradient_checkpointing(device, GradientCheckpointingStrategy::Disabled)
89 }
90
91 pub fn autodiff_with_gradient_checkpointing(device: impl Into<Self>, strategy: GradientCheckpointingStrategy) -> Self {
92 let device = device.into();
93 assert!(!matches!(strategy, GradientCheckpointingStrategy::Custom(_)),
94 "Dispatch accepts Disabled or Balanced checkpointing");
95 #[cfg(feature = "capture")]
96 assert!(!matches!(device.inner(), Self::Capture(_)), "Capture does not support autodiff");
97 Self::Autodiff(Box::new(device.inner().clone()), strategy)
98 }
99
100 pub fn context(&self) -> AutodiffContext {
101 match self {
102 Self::Autodiff(_, strategy) => AutodiffContext::Enabled(*strategy),
103 _ => AutodiffContext::Disabled,
104 }
105 }
106
107 pub fn gradient_checkpointing_strategy(&self) -> GradientCheckpointingStrategy {
108 match self.context() {
109 AutodiffContext::Enabled(strategy) => strategy,
110 AutodiffContext::Disabled => GradientCheckpointingStrategy::Disabled,
111 }
112 }
113
114 pub fn with_context(&self, context: AutodiffContext) -> Self {
115 match context {
116 AutodiffContext::Disabled => self.inner().clone(),
117 AutodiffContext::Enabled(strategy) => Self::autodiff_with_gradient_checkpointing(self.clone(), strategy),
118 }
119 }
120}
121
122impl PartialEq for DispatchDevice {
123 fn eq(&self, other: &Self) -> bool {
124 match (self.inner(), other.inner()) {
125(Self::Host(a), Self::Host(b)) => a == b,
126#[cfg(feature = "cuda")]
127(Self::Cuda(a), Self::Cuda(b)) => a == b,
128#[cfg(feature = "cpu")]
129(Self::Cpu(a), Self::Cpu(b)) => a == b,
130#[cfg(feature = "wgpu")]
131(Self::Wgpu(a), Self::Wgpu(b)) => a == b,
132#[cfg(feature = "rocm")]
133(Self::Rocm(a), Self::Rocm(b)) => a == b,
134#[cfg(feature = "tch")]
135(Self::LibTorch(a), Self::LibTorch(b)) => a == b,
136#[cfg(feature = "remote")]
137(Self::Remote(a), Self::Remote(b)) => a == b,
138#[cfg(feature = "capture")]
139(Self::Capture(a), Self::Capture(b)) => a == b,
140 _ => false,
141 }
142 }
143}
144
145impl Eq for DispatchDevice {}
146
147impl DeviceOps for DispatchDevice {
148 fn defaults(&self) -> Option<ruda_tensor::DeviceSettings> {
149 let settings = match self.inner() {
150 Self::Host(device) => ruda_tensor::get_device_settings::<crate::backends::Host>(device),
151 #[cfg(feature = "cuda")]
152 Self::Cuda(device) => ruda_tensor::get_device_settings::<crate::backends::Cuda>(device),
153 #[cfg(feature = "cpu")]
154 Self::Cpu(device) => ruda_tensor::get_device_settings::<crate::backends::Cpu>(device),
155 #[cfg(feature = "wgpu")]
156 Self::Wgpu(device) => ruda_tensor::get_device_settings::<crate::backends::Wgpu>(device),
157 #[cfg(feature = "rocm")]
158 Self::Rocm(device) => ruda_tensor::get_device_settings::<crate::backends::Rocm>(device),
159 #[cfg(feature = "tch")]
160 Self::LibTorch(device) => ruda_tensor::get_device_settings::<crate::backends::LibTorch>(device),
161 #[cfg(feature = "remote")]
162 Self::Remote(device) => ruda_tensor::get_device_settings::<crate::backends::Remote>(device),
163 #[cfg(feature = "capture")]
164 Self::Capture(device) => ruda_tensor::get_device_settings::<crate::backends::Capture>(device),
165 Self::Autodiff(..) => unreachable!(),
166 };
167 Some(settings)
168 }
169
170 fn inner(&self) -> &Self {
171 match self { Self::Autodiff(device, _) => device.inner(), _ => self }
172 }
173}
174
175impl Default for DispatchDevice {
176 fn default() -> Self {
177 #[cfg(feature = "std")]
178 if let Ok(name) = std::env::var("RUDA_DEVICE") {
179 return match name.to_ascii_lowercase().as_str() {
180"host" => Self::Host(Default::default()),
181#[cfg(feature = "cuda")]
182"cuda" => Self::Cuda(Default::default()),
183#[cfg(feature = "cpu")]
184"cpu" => Self::Cpu(Default::default()),
185#[cfg(feature = "wgpu")]
186"wgpu" => Self::Wgpu(Default::default()),
187#[cfg(feature = "metal")]
188"metal" => Self::Wgpu(ruda_tensor_wgpu::WgpuDevice::default().on(ruda_tensor_wgpu::WgpuBackend::Metal)),
189#[cfg(feature = "vulkan")]
190"vulkan" => Self::Wgpu(ruda_tensor_wgpu::WgpuDevice::default().on(ruda_tensor_wgpu::WgpuBackend::Vulkan)),
191#[cfg(feature = "webgpu")]
192"webgpu" => Self::Wgpu(ruda_tensor_wgpu::WgpuDevice::default().on(ruda_tensor_wgpu::WgpuBackend::WebGpu)),
193#[cfg(feature = "rocm")]
194"rocm" => Self::Rocm(Default::default()),
195#[cfg(feature = "tch")]
196"tch" | "libtorch" => Self::LibTorch(Default::default()),
197#[cfg(feature = "remote")]
198"remote" => Self::Remote(Default::default()),
199 _ => panic!("Unknown or disabled RUDA_DEVICE: {name}"),
200 };
201 }
202#[cfg(feature = "cuda")]
203{ return Self::Cuda(Default::default()); }
204#[cfg(feature = "metal")]
205{ return Self::Wgpu(Default::default()); }
206#[cfg(feature = "rocm")]
207{ return Self::Rocm(Default::default()); }
208#[cfg(feature = "vulkan")]
209{ return Self::Wgpu(Default::default()); }
210#[cfg(feature = "webgpu")]
211{ return Self::Wgpu(Default::default()); }
212#[cfg(feature = "wgpu")]
213{ return Self::Wgpu(Default::default()); }
214#[cfg(feature = "cpu")]
215{ return Self::Cpu(Default::default()); }
216#[cfg(feature = "tch")]
217{ return Self::LibTorch(Default::default()); }
218#[cfg(feature = "remote")]
219{ return Self::Remote(Default::default()); }
220 #[allow(unreachable_code)]
221 Self::Host(Default::default())
222 }
223}
224
225impl Device for DispatchDevice {
226 fn to_id(&self) -> DeviceId {
227 let (tag, id) = match self.inner() {
228Self::Host(device) => (1u16, device.to_id()),
229#[cfg(feature = "cuda")]
230Self::Cuda(device) => (2u16, device.to_id()),
231#[cfg(feature = "cpu")]
232Self::Cpu(device) => (3u16, device.to_id()),
233#[cfg(feature = "wgpu")]
234Self::Wgpu(device) => (4u16, device.to_id()),
235#[cfg(feature = "rocm")]
236Self::Rocm(device) => (5u16, device.to_id()),
237#[cfg(feature = "tch")]
238Self::LibTorch(device) => (6u16, device.to_id()),
239#[cfg(feature = "remote")]
240Self::Remote(device) => (7u16, device.to_id()),
241#[cfg(feature = "capture")]
242Self::Capture(device) => (8u16, device.to_id()),
243 Self::Autodiff(..) => unreachable!(),
244 };
245 assert!(id.type_id < 256, "Native device type exceeds the Dispatch device namespace");
246 DeviceId::new((tag << 8) | id.type_id, id.index_id)
247 }
248
249 fn from_id(id: DeviceId) -> Self {
250 let native = DeviceId::new(id.type_id & 255, id.index_id);
251 let device = match (id.type_id & 0x7fff) >> 8 {
2521 => Self::Host(<ruda_tensor_host::HostDevice as Device>::from_id(native)),
253#[cfg(feature = "cuda")]
2542 => Self::Cuda(<ruda_tensor_device::cuda::CudaDevice as Device>::from_id(native)),
255#[cfg(feature = "cpu")]
2563 => Self::Cpu(<ruda_driver_cpu::CpuDevice as Device>::from_id(native)),
257#[cfg(feature = "wgpu")]
2584 => Self::Wgpu(<ruda_tensor_wgpu::WgpuDevice as Device>::from_id(native)),
259#[cfg(feature = "rocm")]
2605 => Self::Rocm(<ruda_tensor_rocm::RocmDevice as Device>::from_id(native)),
261#[cfg(feature = "tch")]
2626 => Self::LibTorch(<ruda_tensor_tch::LibTorchDevice as Device>::from_id(native)),
263#[cfg(feature = "remote")]
2647 => Self::Remote(<ruda_tensor_remote::RemoteDevice as Device>::from_id(native)),
265#[cfg(feature = "capture")]
2668 => Self::Capture(<ruda_tensor_capture::CaptureDevice as Device>::from_id(native)),
267 _ => panic!("Unknown or disabled Dispatch device: {id}"),
268 };
269 if id.type_id & 0x8000 != 0 { Self::autodiff(device) } else { device }
270 }
271}
272
273impl From<ruda_tensor_host::HostDevice> for DispatchDevice {
274 fn from(device: ruda_tensor_host::HostDevice) -> Self { Self::Host(device) }
275}
276#[cfg(feature = "cuda")]
277impl From<ruda_tensor_device::cuda::CudaDevice> for DispatchDevice {
278 fn from(device: ruda_tensor_device::cuda::CudaDevice) -> Self { Self::Cuda(device) }
279}
280#[cfg(feature = "cpu")]
281impl From<ruda_driver_cpu::CpuDevice> for DispatchDevice {
282 fn from(device: ruda_driver_cpu::CpuDevice) -> Self { Self::Cpu(device) }
283}
284#[cfg(feature = "wgpu")]
285impl From<ruda_tensor_wgpu::WgpuDevice> for DispatchDevice {
286 fn from(device: ruda_tensor_wgpu::WgpuDevice) -> Self { Self::Wgpu(device) }
287}
288#[cfg(feature = "rocm")]
289impl From<ruda_tensor_rocm::RocmDevice> for DispatchDevice {
290 fn from(device: ruda_tensor_rocm::RocmDevice) -> Self { Self::Rocm(device) }
291}
292#[cfg(feature = "tch")]
293impl From<ruda_tensor_tch::LibTorchDevice> for DispatchDevice {
294 fn from(device: ruda_tensor_tch::LibTorchDevice) -> Self { Self::LibTorch(device) }
295}
296#[cfg(feature = "remote")]
297impl From<ruda_tensor_remote::RemoteDevice> for DispatchDevice {
298 fn from(device: ruda_tensor_remote::RemoteDevice) -> Self { Self::Remote(device) }
299}
300
301#[cfg(feature = "capture")]
302impl From<ruda_tensor_capture::CaptureDevice> for DispatchDevice {
303 fn from(device: ruda_tensor_capture::CaptureDevice) -> Self { Self::Capture(device) }
304}