msm_webgpu/cuzk/test/
field.rs

1use std::time::Instant;
2
3use ff::PrimeField;
4use wgpu::CommandEncoderDescriptor;
5
6use crate::cuzk::{
7    gpu::{
8        create_and_write_storage_buffer, create_bind_group, create_bind_group_layout,
9        create_compute_pipeline, create_storage_buffer, execute_pipeline, get_adapter, get_device,
10        read_from_gpu_test,
11    },
12    msm::{PARAMS, WORD_SIZE},
13    shader_manager::ShaderManager,
14    utils::{bytes_to_field, field_to_u8_vec_for_gpu, to_biguint_le},
15};
16
17async fn field_op<F: PrimeField>(op: &str, a: F, b: F) -> F {
18    let a_bytes = field_to_u8_vec_for_gpu(&a, PARAMS.num_words, WORD_SIZE);
19    let b_bytes = field_to_u8_vec_for_gpu(&b, PARAMS.num_words, WORD_SIZE);
20    let input_size = 1;
21    let chunk_size = if input_size >= 65536 { 16 } else { 4 };
22    let num_words = PARAMS.num_words;
23    println!("Input size: {input_size}");
24    println!("Chunk size: {chunk_size}");
25    println!("Num words: {num_words}");
26    println!("Word size: {WORD_SIZE}");
27    println!("Params: {PARAMS:?}");
28
29    let shader_manager = ShaderManager::new(WORD_SIZE, chunk_size, input_size);
30
31    let adapter = get_adapter().await;
32    let (device, queue) = get_device(&adapter).await;
33    let mut encoder = device.create_command_encoder(&CommandEncoderDescriptor {
34        label: Some("Field Encoder"),
35    });
36
37    let shader_code = shader_manager.gen_test_field_shader();
38
39    let a_sb = create_and_write_storage_buffer(Some("A buffer"), &device, &a_bytes);
40    let b_sb = create_and_write_storage_buffer(Some("B buffer"), &device, &b_bytes);
41
42    let result_sb = create_storage_buffer(Some("Result buffer"), &device, (num_words * 4) as u64);
43
44    let bind_group_layout = create_bind_group_layout(
45        Some("Bind group layout"),
46        &device,
47        vec![],
48        vec![&a_sb, &b_sb, &result_sb],
49        vec![],
50    );
51
52    let bind_group = create_bind_group(
53        Some("Bind group"),
54        &device,
55        &bind_group_layout,
56        vec![&a_sb, &b_sb, &result_sb],
57    );
58
59    let compute_pipeline = create_compute_pipeline(
60        Some("Field add shader"),
61        &device,
62        &bind_group_layout,
63        &shader_code,
64        op,
65    )
66    .await;
67
68    execute_pipeline(&mut encoder, compute_pipeline, bind_group, 1, 1, 1).await;
69
70    // Map results back from GPU to CPU.
71    let data = read_from_gpu_test(&device, &queue, encoder, vec![result_sb]).await;
72
73    // Destroy the GPU device object.
74    device.destroy();
75
76    let data_u32 = bytemuck::cast_slice::<u8, u32>(&data[0]);
77
78    let result_biguint = to_biguint_le(data_u32, num_words, WORD_SIZE as u32);
79
80    
81
82    bytes_to_field(&result_biguint.to_bytes_le())
83}
84
85/// Run WebGPU field op sync
86pub fn run_webgpu_field_op<F: PrimeField>(op: &str, a: F, b: F) -> F {
87    pollster::block_on(run_webgpu_field_op_async(op, a, b))
88}
89
90/// Run WebGPU field op async
91pub async fn run_webgpu_field_op_async<F: PrimeField>(op: &str, a: F, b: F) -> F {
92    let now = Instant::now();
93    let result = field_op::<F>(op, a, b).await;
94    println!("Field add time: {:?}", now.elapsed());
95    result
96}
97
98#[cfg(test)]
99mod tests {
100    use crate::{
101        cuzk::{msm::calc_num_words, utils::u8s_to_field_without_assertion},
102        sample_scalars,
103    };
104
105    use super::*;
106    use ff::Field;
107    use halo2curves::bn256::Fq;
108    use rand::thread_rng;
109
110    #[test]
111    fn test_webgpu_field_add() {
112        let scalars = sample_scalars::<Fq>(50);
113        for scalar in scalars.chunks(2) {
114            let a = scalar[0];
115            let b = scalar[1];
116
117            let fast = a + b;
118
119            let result = run_webgpu_field_op::<Fq>("test_field_add", a, b);
120
121            println!("Result: {:?}", result);
122            assert_eq!(fast, result);
123        }
124    }
125
126    #[test]
127    fn test_webgpu_field_sub() {
128        let scalars = sample_scalars::<Fq>(50);
129        for scalar in scalars.chunks(2) {
130            let a = scalar[0];
131            let b = scalar[1];
132
133            let fast = a - b;
134
135            let result = run_webgpu_field_op::<Fq>("test_field_sub", a, b);
136
137            println!("Result: {:?}", result);
138            assert_eq!(fast, result);
139        }
140    }
141
142    #[test]
143    fn test_webgpu_field_mul() {
144        let mut rng = thread_rng();
145        let a = Fq::random(&mut rng);
146        let b = Fq::random(&mut rng);
147
148        let fast = a * b;
149        let result = run_webgpu_field_op::<Fq>("test_field_mul", a, b);
150
151        println!("Result: {:?}", result);
152        assert_eq!(fast, result);
153    }
154
155    #[test]
156    fn test_webgpu_field_barret_mul() {
157        let mut rng = thread_rng();
158        let a = Fq::random(&mut rng);
159        let b = Fq::random(&mut rng);
160
161        let fast = a;
162        let result = run_webgpu_field_op::<Fq>("test_barret_mul", a, b);
163
164        println!("Result: {:?}", result);
165        assert_eq!(fast, result);
166    }
167
168    #[test]
169    fn test_field_to_u8_vec_for_gpu() {
170        // random
171        let mut rng = thread_rng();
172        let a = Fq::random(&mut rng);
173        for word_size in 13..17 {
174            let num_words = calc_num_words(word_size);
175            let bytes = field_to_u8_vec_for_gpu(&a, num_words, word_size);
176            let a_from_bytes = u8s_to_field_without_assertion(&bytes, num_words, word_size);
177            assert_eq!(a, a_from_bytes);
178        }
179    }
180}
msm_webgpu/cuzk/test/field.rs

msm_webgpu/cuzk/test/
field.rs
