Skip to main content

arr_o_gpu/arr_o_gpu/module/method/
get_heap.rs

1use wgpu::{
2    BufferUsages, MapMode,
3    wgt::{BufferDescriptor, CommandEncoderDescriptor, PollType},
4};
5
6use crate::ArrOgpuModule;
7
8impl ArrOgpuModule {
9    pub fn get_heap(&self) -> Vec<f32> {
10        let wgpu_init = &self.wgpu_module.read().unwrap();
11        let size = (std::mem::size_of::<f32>() * (*self.maximum as usize)) as u64;
12
13        let copy_buffer = wgpu_init.device.create_buffer(
14            &(BufferDescriptor {
15                label: Some("create copy buffer for get_heap"),
16                mapped_at_creation: false,
17                size,
18                usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
19            }),
20        );
21
22        let mut encoder = wgpu_init.device.create_command_encoder(
23            &(CommandEncoderDescriptor {
24                label: Some("create encoder for get_heap"),
25            }),
26        );
27
28        encoder.copy_buffer_to_buffer(&self.heap_buffer, 0, &copy_buffer, 0, size);
29        let index = wgpu_init.queue.submit(Some(encoder.finish()));
30
31        let buffer_slice = copy_buffer.slice(..);
32        buffer_slice.map_async(MapMode::Read, |e| e.unwrap());
33        wgpu_init
34            .device
35            .poll(PollType::Wait {
36                submission_index: Some(index),
37                timeout: None,
38            })
39            .unwrap();
40
41        let data = buffer_slice.get_mapped_range();
42        let heap: Vec<f32> = bytemuck::cast_slice(&data).into();
43        drop(data);
44        copy_buffer.unmap();
45
46        heap
47    }
48}