cubecl_wgpu/
lib.rs

1#[macro_use]
2extern crate derive_new;
3
4extern crate alloc;
5
6mod backend;
7mod compiler;
8mod compute;
9mod device;
10mod graphics;
11mod runtime;
12
13pub use compiler::base::*;
14pub use compute::*;
15pub use device::*;
16pub use graphics::*;
17pub use runtime::*;
18/// The exact wgpu version used by this runtime, for sharing resources with applications.
19pub use wgpu;
20
21#[cfg(feature = "spirv")]
22pub use backend::vulkan;
23
24#[cfg(all(feature = "msl", target_os = "macos"))]
25pub use backend::metal;
26
27#[cfg(all(test, not(feature = "spirv"), not(feature = "msl")))]
28#[allow(unexpected_cfgs)]
29mod tests {
30    pub type TestRuntime = crate::WgpuRuntime;
31    use half::f16;
32
33    // Include 64-bit types (i64, u64) for WGSL as wgpu supports them. These don't exist on
34    // native WebGPU however.
35    //
36    // Also include f16, this is an extension but supported by wgpu and WebGPU.
37    cubecl_core::testgen_all!(f32: [f16, f32], i32: [i32, i64], u32: [u32, u64]);
38    cubecl_std::testgen!();
39    cubecl_std::testgen_tensor_identity!([flex32, f32, u32]);
40    cubecl_std::testgen_quantized_view!(f32);
41    cubecl_core::testgen_profiling!();
42
43    /// A kernel that brings its own WGSL, the way a template kernel downstream
44    /// does: no representation for the runtime to read, only text for naga.
45    mod precompiled {
46        use cubecl_core::prelude::*;
47        use cubecl_ir::{UIntKind, metadata::Info, settings::Dim3};
48        use cubecl_server::kernel::{
49            CubeKernel, KernelDefinition, KernelMetadata, PrecompiledSource,
50        };
51        use cubecl_server::runtime::Runtime;
52        use cubecl_server::server::KernelArguments;
53
54        use super::TestRuntime;
55
56        const SOURCE: &str = r#"
57@group(0) @binding(0) var<storage, read_write> data: array<f32>;
58
59@compute @workgroup_size(1)
60fn double(@builtin(global_invocation_id) id: vec3<u32>) {
61    data[id.x] = data[id.x] * 2.0;
62}
63"#;
64
65        struct Double;
66
67        impl KernelMetadata for Double {
68            fn id(&self) -> KernelId {
69                KernelId::new::<Self>()
70            }
71
72            fn address_type(&self) -> ElemType {
73                ElemType::UInt(UIntKind::U32)
74            }
75        }
76
77        impl CubeKernel for Double {
78            fn define(&self) -> KernelDefinition {
79                let settings = KernelSettings::new(
80                    Dim3::new_single(),
81                    ExecutionMode::Checked,
82                    AddressType::U32,
83                );
84                KernelDefinition {
85                    body: Scope::root(settings.clone()),
86                    settings,
87                    info: Info::default(),
88                }
89            }
90
91            fn source(&self) -> Option<PrecompiledSource> {
92                Some(PrecompiledSource {
93                    source: SOURCE.to_string(),
94                    entrypoint_name: "double".to_string(),
95                    lang: "wgsl",
96                })
97            }
98        }
99
100        #[test]
101        fn a_hand_written_wgsl_kernel_launches() {
102            let client = TestRuntime::client(&Default::default());
103            let input = [1.0f32, 2.0, 3.0, 4.0];
104            let handle = client.create_from_slice(bytemuck::cast_slice(&input));
105
106            client.launch(
107                Box::new(Double),
108                CubeCount::Static(input.len() as u32, 1, 1),
109                KernelArguments::new().with_buffer(handle.clone().binding()),
110            );
111
112            let bytes = client.read_one(handle).expect("the launch ran");
113            let output: &[f32] = bytemuck::cast_slice(&bytes);
114            assert_eq!(output, [2.0, 4.0, 6.0, 8.0]);
115        }
116    }
117
118    /// WGSL packs fp8 four lanes to a `u32` and has no type for anything narrower. Rejecting
119    /// that has to reach the caller: a panic on the device thread is caught there, logged as a
120    /// warning, and the caller reads back a zeroed buffer as if the launch had succeeded.
121    mod fp8_lanes {
122        use cubecl_common::e4m3;
123        use cubecl_core::prelude::*;
124        use cubecl_core::{self as cubecl};
125        use cubecl_server::runtime::Runtime;
126        use cubecl_server::server::Handle;
127
128        use super::TestRuntime;
129
130        /// A cast, which the minifloat lowering pass is what rejects.
131        #[cube(launch_unchecked)]
132        fn cast_fp8<N: Size>(input: &[Vector<f32, N>], out: &mut [Vector<f32, N>]) {
133            if ABSOLUTE_POS < input.len() {
134                let codes = Vector::<e4m3, N>::cast_from(input[ABSOLUTE_POS]);
135                out[ABSOLUTE_POS] = Vector::cast_from(codes);
136            }
137        }
138
139        /// No cast at all, so only the WGSL type printer ever sees the fp8.
140        #[cube(launch_unchecked)]
141        fn copy_fp8<N: Size>(input: &[Vector<e4m3, N>], out: &mut [Vector<e4m3, N>]) {
142            if ABSOLUTE_POS < input.len() {
143                out[ABSOLUTE_POS] = input[ABSOLUTE_POS];
144            }
145        }
146
147        fn assert_rejected(client: &Client, out: Handle) {
148            let err = client
149                .read_one(out)
150                .expect_err("two fp8 lanes have no WGSL representation, the launch must fail")
151                .to_string();
152            assert!(
153                err.contains("fp8 on WGSL is packed 4 lanes to a u32"),
154                "the packing rule has to be in the error the caller sees, got: {err}"
155            );
156        }
157
158        #[test]
159        fn cast_at_two_lanes_is_reported() {
160            let client = TestRuntime::client(&Default::default());
161            let input = client.create_from_slice(&[0u8; 64]);
162            let out = client.empty(64);
163            unsafe {
164                cast_fp8::launch_unchecked(
165                    &client,
166                    CubeCount::new_single(),
167                    CubeDim::new_1d(8),
168                    2,
169                    BufferArg::from_raw_parts(input, 16),
170                    BufferArg::from_raw_parts(out.clone(), 16),
171                )
172            };
173            assert_rejected(&client, out);
174        }
175
176        #[test]
177        fn copy_at_two_lanes_is_reported() {
178            let client = TestRuntime::client(&Default::default());
179            let input = client.create_from_slice(&[0u8; 32]);
180            let out = client.empty(32);
181            unsafe {
182                copy_fp8::launch_unchecked(
183                    &client,
184                    CubeCount::new_single(),
185                    CubeDim::new_1d(8),
186                    2,
187                    BufferArg::from_raw_parts(input, 32),
188                    BufferArg::from_raw_parts(out.clone(), 32),
189                )
190            };
191            assert_rejected(&client, out);
192        }
193    }
194
195    cubecl_core::testgen_complex_validation!();
196}
197
198#[cfg(all(test, feature = "spirv"))]
199#[allow(unexpected_cfgs)]
200mod tests_spirv {
201    pub type TestRuntime = crate::WgpuRuntime;
202    use cubecl_core::flex32;
203    use half::f16;
204
205    cubecl_core::testgen_all!(f32: [f16, flex32, f32], i32: [i8, i16, i32, i64], u32: [u8, u16, u32, u64]);
206    cubecl_std::testgen!();
207    cubecl_std::testgen_tensor_identity!([f16, flex32, f32, u32]);
208    cubecl_std::testgen_quantized_view!(f16);
209    cubecl_core::testgen_profiling!();
210}
211
212#[cfg(all(test, feature = "msl"))]
213#[allow(unexpected_cfgs)]
214mod tests_msl {
215    pub type TestRuntime = crate::WgpuRuntime;
216    use half::{bf16, f16};
217
218    cubecl_core::testgen_all!(f32: [f16, bf16, f32], i32: [i8, i16, i32, i64], u32: [u8, u16, u32, u64]);
219    cubecl_std::testgen!();
220    cubecl_std::testgen_tensor_identity!([f16, flex32, f32, u32]);
221    cubecl_std::testgen_quantized_view!(f16);
222    cubecl_core::testgen_profiling!();
223    // No `testgen_profiling_nested!`: Metal writes a pass's timestamps only once its command
224    // buffer completes, so a window's end waits for the device, and an enclosing window measures
225    // that idle gap. Add it once the readback is deferred instead of waited for.
226    cubecl_core::testgen_complex_validation!();
227}
cubecl_wgpu/lib.rs

cubecl_wgpu/
lib.rs
