Skip to main content

arr_o_gpu/arr_o_gpu/array/method/
get_heap.rs

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