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
40    /// WGSL packs fp8 four lanes to a `u32` and has no type for anything narrower. Rejecting
41    /// that has to reach the caller: a panic on the device thread is caught there, logged as a
42    /// warning, and the caller reads back a zeroed buffer as if the launch had succeeded.
43    mod fp8_lanes {
44        use cubecl_common::e4m3;
45        use cubecl_core::prelude::*;
46        use cubecl_core::{self as cubecl};
47        use cubecl_runtime::server::Handle;
48
49        use super::TestRuntime;
50
51        /// A cast, which the minifloat lowering pass is what rejects.
52        #[cube(launch_unchecked)]
53        fn cast_fp8<N: Size>(input: &[Vector<f32, N>], out: &mut [Vector<f32, N>]) {
54            if ABSOLUTE_POS < input.len() {
55                let codes = Vector::<e4m3, N>::cast_from(input[ABSOLUTE_POS]);
56                out[ABSOLUTE_POS] = Vector::cast_from(codes);
57            }
58        }
59
60        /// No cast at all, so only the WGSL type printer ever sees the fp8.
61        #[cube(launch_unchecked)]
62        fn copy_fp8<N: Size>(input: &[Vector<e4m3, N>], out: &mut [Vector<e4m3, N>]) {
63            if ABSOLUTE_POS < input.len() {
64                out[ABSOLUTE_POS] = input[ABSOLUTE_POS];
65            }
66        }
67
68        fn assert_rejected(client: &ComputeClient<TestRuntime>, out: Handle) {
69            let err = client
70                .read_one(out)
71                .expect_err("two fp8 lanes have no WGSL representation, the launch must fail")
72                .to_string();
73            assert!(
74                err.contains("fp8 on WGSL is packed 4 lanes to a u32"),
75                "the packing rule has to be in the error the caller sees, got: {err}"
76            );
77        }
78
79        #[test]
80        fn cast_at_two_lanes_is_reported() {
81            let client = TestRuntime::client(&Default::default());
82            let input = client.create_from_slice(&[0u8; 64]);
83            let out = client.empty(64);
84            unsafe {
85                cast_fp8::launch_unchecked::<TestRuntime>(
86                    &client,
87                    CubeCount::new_single(),
88                    CubeDim::new_1d(8),
89                    2,
90                    BufferArg::from_raw_parts(input, 16),
91                    BufferArg::from_raw_parts(out.clone(), 16),
92                )
93            };
94            assert_rejected(&client, out);
95        }
96
97        #[test]
98        fn copy_at_two_lanes_is_reported() {
99            let client = TestRuntime::client(&Default::default());
100            let input = client.create_from_slice(&[0u8; 32]);
101            let out = client.empty(32);
102            unsafe {
103                copy_fp8::launch_unchecked::<TestRuntime>(
104                    &client,
105                    CubeCount::new_single(),
106                    CubeDim::new_1d(8),
107                    2,
108                    BufferArg::from_raw_parts(input, 32),
109                    BufferArg::from_raw_parts(out.clone(), 32),
110                )
111            };
112            assert_rejected(&client, out);
113        }
114    }
115}
116
117#[cfg(all(test, feature = "spirv"))]
118#[allow(unexpected_cfgs)]
119mod tests_spirv {
120    pub type TestRuntime = crate::WgpuRuntime;
121    use cubecl_core::flex32;
122    use half::f16;
123
124    cubecl_core::testgen_all!(f32: [f16, flex32, f32], i32: [i8, i16, i32, i64], u32: [u8, u16, u32, u64]);
125    cubecl_std::testgen!();
126    cubecl_std::testgen_tensor_identity!([f16, flex32, f32, u32]);
127    cubecl_std::testgen_quantized_view!(f16);
128}
129
130#[cfg(all(test, feature = "msl"))]
131#[allow(unexpected_cfgs)]
132mod tests_msl {
133    pub type TestRuntime = crate::WgpuRuntime;
134    use half::f16;
135
136    cubecl_core::testgen_all!(f32: [f16, f32], i32: [i16, i32], u32: [u16, u32]);
137    cubecl_std::testgen!();
138    cubecl_std::testgen_tensor_identity!([f16, flex32, f32, u32]);
139    cubecl_std::testgen_quantized_view!(f16);
140}
cubecl_wgpu/lib.rs

cubecl_wgpu/
lib.rs
