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

cubecl_wgpu/
lib.rs
