ruda_tensor_wgpu/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3extern crate alloc;
4
5#[cfg(feature = "template")]
6pub use ruda_kernel::{
7    dsl::prelude::KernelMetadata,
8    tensor::contiguous::into_contiguous,
9    kernel_source,
10    template::{KernelSource, SourceKernel, SourceTemplate},
11};
12#[cfg(feature = "template")]
13mod template;
14#[cfg(feature = "template")]
15pub use template::build_info;
16
17pub use ruda_tensor_device::{BoolElement, FloatElement, IntElement};
18pub use ruda_tensor_device::{DeviceBackend, RudaTensor};
19pub use ruda_kernel::dsl::RudaDim;
20pub use ruda_core::flex32;
21
22pub use ruda_driver_wgpu::{
23    AutoCompiler, MemoryConfiguration, RuntimeOptions, WgpuDevice, WgpuResource, WgpuRuntime,
24    WgpuSetup, WgpuStorage, init_device, init_setup, init_setup_async,
25};
26// Vulkan and WebGpu would have conflicting type names
27pub mod graphics {
28    pub use ruda_driver_wgpu::{AutoGraphicsApi, Dx12, GraphicsApi, Metal, OpenGl, Vulkan, WebGpu};
29}
30
31#[cfg(feature = "wgsl")]
32pub use ruda_driver_wgpu::WgslCompiler;
33#[cfg(feature = "spirv")]
34pub use ruda_driver_wgpu::vulkan::VkSpirvCompiler;
35
36#[cfg(feature = "fusion")]
37/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
38///
39/// This backend can target multiple graphics APIs, including:
40///   - [Vulkan][crate::graphics::Vulkan] on Linux, Windows, and Android.
41///   - [OpenGL](crate::graphics::OpenGl) on Linux, Windows, and Android.
42///   - [DirectX 12](crate::graphics::Dx12) on Windows.
43///   - [Metal][crate::graphics::Metal] on Apple hardware.
44///   - [WebGPU](crate::graphics::WebGpu) on supported browsers and `wasm` runtimes.
45///
46/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
47/// you have to manually initialize the runtime. For example:
48///
49/// ```rust, ignore
50/// fn custom_init() {
51///     let device = Default::default();
52///     ruda::backend::wgpu::init_setup::<ruda::backend::wgpu::graphics::Vulkan>(
53///         &device,
54///         Default::default(),
55///     );
56/// }
57/// ```
58/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
59/// It's also possible to use an existing wgpu device, by using `init_device`.
60///
61/// # Notes
62///
63/// This version of the wgpu backend uses [ruda_fusion] to compile and optimize streams of tensor
64/// operations for improved performance.
65///
66/// You can disable the `fusion` feature flag to remove that functionality, which might be
67/// necessary on `wasm` for now.
68pub type Wgpu<F = f32, I = i32, B = u32> =
69    ruda_fusion::Fusion<DeviceBackend<ruda_driver_wgpu::WgpuRuntime, F, I, B>>;
70
71#[cfg(not(feature = "fusion"))]
72/// Tensor backend that uses the wgpu crate for executing GPU compute shaders.
73///
74/// This backend can target multiple graphics APIs, including:
75///   - [Vulkan] on Linux, Windows, and Android.
76///   - [OpenGL](crate::OpenGl) on Linux, Windows, and Android.
77///   - [DirectX 12](crate::Dx12) on Windows.
78///   - [Metal] on Apple hardware.
79///   - [WebGPU](crate::WebGpu) on supported browsers and `wasm` runtimes.
80///
81/// To configure the wgpu backend, eg. to select what graphics API to use or what memory strategy to use,
82/// you have to manually initialize the runtime. For example:
83///
84/// ```rust, ignore
85/// fn custom_init() {
86///     let device = Default::default();
87///     ruda::backend::wgpu::init_setup::<ruda::backend::wgpu::graphics::Vulkan>(
88///         &device,
89///         Default::default(),
90///     );
91/// }
92/// ```
93/// will mean the given device (in this case the default) will be initialized to use Vulkan as the graphics API.
94/// It's also possible to use an existing wgpu device, by using `init_device`.
95///
96/// # Notes
97///
98/// This version of the wgpu backend doesn't use [ruda_fusion] to compile and optimize streams of tensor
99/// operations.
100///
101/// You can enable the `fusion` feature flag to add that functionality, which might improve
102/// performance.
103pub type Wgpu<F = f32, I = i32, B = u32> = DeviceBackend<ruda_driver_wgpu::WgpuRuntime, F, I, B>;
104
105#[cfg(feature = "vulkan")]
106/// Tensor backend that leverages the Vulkan graphics API to execute GPU compute shaders compiled to SPIR-V.
107pub type Vulkan<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
108
109#[cfg(feature = "webgpu")]
110/// Tensor backend that uses the wgpu crate to execute GPU compute shaders written in WGSL.
111pub type WebGpu<F = f32, I = i32, B = u32> = Wgpu<F, I, B>;
112
113#[cfg(feature = "metal")]
114/// Tensor backend that leverages the Metal graphics API to execute GPU compute shaders compiled to MSL.
115pub type Metal<F = f32, I = i32, B = u8> = Wgpu<F, I, B>;
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use ruda_tensor::{Backend, BoolStore, DType, QTensorPrimitive};
121
122    #[test]
123    fn should_support_dtypes() {
124        type B = Wgpu;
125        let device = Default::default();
126
127        assert!(B::supports_dtype(&device, DType::F32));
128        assert!(B::supports_dtype(&device, DType::I64));
129        assert!(B::supports_dtype(&device, DType::I32));
130        assert!(B::supports_dtype(&device, DType::U64));
131        assert!(B::supports_dtype(&device, DType::U32));
132        assert!(B::supports_dtype(
133            &device,
134            DType::QFloat(RudaTensor::<WgpuRuntime>::default_scheme())
135        ));
136        // Registered as supported type but we don't actually use it?
137        assert!(B::supports_dtype(&device, DType::Bool(BoolStore::Native)));
138
139        #[cfg(feature = "vulkan")]
140        {
141            assert!(B::supports_dtype(&device, DType::F16));
142            assert!(B::supports_dtype(&device, DType::I16));
143            assert!(B::supports_dtype(&device, DType::I8));
144            assert!(B::supports_dtype(&device, DType::U16));
145            assert!(B::supports_dtype(&device, DType::U8));
146
147            assert!(!B::supports_dtype(&device, DType::F64));
148            assert!(!B::supports_dtype(&device, DType::Flex32));
149            // Not supported for any arithmetics, but buffer, conversion and possibly matmul (hw dependent)
150            assert!(!B::supports_dtype(&device, DType::BF16));
151        }
152
153        #[cfg(feature = "metal")]
154        {
155            assert!(B::supports_dtype(&device, DType::F16));
156            assert!(B::supports_dtype(&device, DType::I16));
157            assert!(B::supports_dtype(&device, DType::I8));
158            assert!(B::supports_dtype(&device, DType::U16));
159            assert!(B::supports_dtype(&device, DType::U8));
160
161            assert!(!B::supports_dtype(&device, DType::F64));
162            assert!(!B::supports_dtype(&device, DType::BF16));
163            assert!(!B::supports_dtype(&device, DType::Flex32));
164        }
165
166        // On macOS without the `metal` feature, wgpu still uses Metal at runtime,
167        // which doesn't support F64 or BF16.
168        #[cfg(all(not(any(feature = "vulkan", feature = "metal")), target_os = "macos"))]
169        {
170            assert!(B::supports_dtype(&device, DType::Flex32));
171            assert!(B::supports_dtype(&device, DType::F16));
172
173            assert!(!B::supports_dtype(&device, DType::F64));
174            assert!(!B::supports_dtype(&device, DType::BF16));
175            assert!(!B::supports_dtype(&device, DType::I16));
176            assert!(!B::supports_dtype(&device, DType::I8));
177            assert!(!B::supports_dtype(&device, DType::U16));
178            assert!(!B::supports_dtype(&device, DType::U8));
179        }
180
181        #[cfg(not(any(feature = "vulkan", feature = "metal", target_os = "macos")))]
182        {
183            assert!(B::supports_dtype(&device, DType::F64));
184            assert!(B::supports_dtype(&device, DType::Flex32));
185            assert!(B::supports_dtype(&device, DType::F16));
186
187            assert!(!B::supports_dtype(&device, DType::BF16));
188            assert!(!B::supports_dtype(&device, DType::I16));
189            assert!(!B::supports_dtype(&device, DType::I8));
190            assert!(!B::supports_dtype(&device, DType::U16));
191            assert!(!B::supports_dtype(&device, DType::U8));
192        }
193    }
194}
ruda_tensor_wgpu/lib.rs

ruda_tensor_wgpu/
lib.rs
