Skip to main content

versorc_math/gpu/
mod.rs

1// ============================================================================
2// VERSORC-MATH: GPU HARDWARE COMPUTE & VULKAN DISPATCH PRIMITIVES
3// License: Apache-2.0
4// ============================================================================
5
6pub mod shaders;
7
8pub use shaders::{
9    CgaTransformPushConstants, HotStratumE8PushConstants, PgaTransformPushConstants,
10    RmsNormPushConstants, RopePushConstants, SiluMulPushConstants, VecAddPushConstants,
11    CGA_TRANSFORM_SPV, HOT_STRATUM_E8_SPV, PGA_TRANSFORM_SPV, RMSNORM_SPV, ROPE_SPV,
12    SILU_MUL_SPV, VEC_ADD_SPV,
13};
14
15#[cfg(feature = "vulkan")]
16pub mod vulkan {
17    use ash::{vk, Device, Entry, Instance};
18    use std::ffi::CStr;
19    use crate::pga::PgaMotor;
20    use crate::gpu::shaders::*;
21
22    /// Discovered Vulkan Compute Device.
23    #[derive(Debug, Clone)]
24    pub struct DiscoveredGpu {
25        pub name: String,
26        pub is_discrete: bool,
27        pub driver_version: u32,
28        pub api_version: u32,
29    }
30
31    /// Minimal zero-dependency Vulkan Compute Context for math dispatch.
32    pub struct GpuComputeContext {
33        pub entry: Entry,
34        pub instance: Instance,
35        pub physical_device: vk::PhysicalDevice,
36        pub device: Device,
37        pub compute_queue: vk::Queue,
38        pub queue_family_index: u32,
39    }
40
41    impl GpuComputeContext {
42        /// Probes host for available Vulkan devices.
43        pub fn probe_devices() -> Result<Vec<DiscoveredGpu>, String> {
44            let entry = unsafe {
45                Entry::load().map_err(|e| format!("Vulkan loader error: {}", e))?
46            };
47            let app_info = vk::ApplicationInfo::default()
48                .api_version(vk::make_api_version(0, 1, 3, 0));
49            let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
50            let instance = unsafe {
51                entry.create_instance(&create_info, None)
52                    .map_err(|e| format!("vkCreateInstance failed: {}", e))?
53            };
54
55            let pdevices = unsafe {
56                instance.enumerate_physical_devices()
57                    .map_err(|e| format!("Failed to enumerate physical devices: {}", e))?
58            };
59
60            let mut list = Vec::new();
61            for &pdev in &pdevices {
62                let props = unsafe { instance.get_physical_device_properties(pdev) };
63                let name = unsafe { CStr::from_ptr(props.device_name.as_ptr()) }
64                    .to_string_lossy()
65                    .into_owned();
66                let is_discrete = props.device_type == vk::PhysicalDeviceType::DISCRETE_GPU;
67                list.push(DiscoveredGpu {
68                    name,
69                    is_discrete,
70                    driver_version: props.driver_version,
71                    api_version: props.api_version,
72                });
73            }
74
75            unsafe {
76                instance.destroy_instance(None);
77            }
78
79            Ok(list)
80        }
81
82        /// Initializes compute context targeting preferred GPU (preferring discrete if available).
83        pub fn init_compute() -> Result<Self, String> {
84            let entry = unsafe {
85                Entry::load().map_err(|e| format!("Vulkan loader error: {}", e))?
86            };
87            let app_info = vk::ApplicationInfo::default()
88                .api_version(vk::make_api_version(0, 1, 3, 0));
89            let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
90            let instance = unsafe {
91                entry.create_instance(&create_info, None)
92                    .map_err(|e| format!("vkCreateInstance failed: {}", e))?
93            };
94
95            let pdevices = unsafe {
96                instance.enumerate_physical_devices()
97                    .map_err(|e| format!("Failed to enumerate physical devices: {}", e))?
98            };
99
100            if pdevices.is_empty() {
101                return Err("No Vulkan physical devices found".into());
102            }
103
104            // Pick discrete GPU if present, else first
105            let mut selected_pdev = pdevices[0];
106            for &pdev in &pdevices {
107                let props = unsafe { instance.get_physical_device_properties(pdev) };
108                if props.device_type == vk::PhysicalDeviceType::DISCRETE_GPU {
109                    selected_pdev = pdev;
110                    break;
111                }
112            }
113
114            // Find compute queue family
115            let qprops = unsafe { instance.get_physical_device_queue_family_properties(selected_pdev) };
116            let queue_family_index = qprops
117                .iter()
118                .enumerate()
119                .position(|(_, p)| p.queue_flags.contains(vk::QueueFlags::COMPUTE))
120                .ok_or_else(|| "No compute queue family found".to_string())? as u32;
121
122            let priorities = [1.0f32];
123            let queue_create_info = vk::DeviceQueueCreateInfo::default()
124                .queue_family_index(queue_family_index)
125                .queue_priorities(&priorities);
126
127            let dev_create_info = vk::DeviceCreateInfo::default()
128                .queue_create_infos(std::slice::from_ref(&queue_create_info));
129
130            let device = unsafe {
131                instance.create_device(selected_pdev, &dev_create_info, None)
132                    .map_err(|e| format!("vkCreateDevice failed: {}", e))?
133            };
134
135            let compute_queue = unsafe { device.get_device_queue(queue_family_index, 0) };
136
137            Ok(Self {
138                entry,
139                instance,
140                physical_device: selected_pdev,
141                device,
142                compute_queue,
143                queue_family_index,
144            })
145        }
146
147        /// Creates a shader module from SPIR-V byte slice.
148        pub fn create_shader_module(&self, spirv_bytes: &[u8]) -> Result<vk::ShaderModule, String> {
149            assert_eq!(
150                spirv_bytes.len() % 4,
151                0,
152                "SPIR-V byte length must be a multiple of 4"
153            );
154            let mut words = Vec::<u32>::with_capacity(spirv_bytes.len() / 4);
155            for chunk in spirv_bytes.chunks_exact(4) {
156                words.push(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
157            }
158            let info = vk::ShaderModuleCreateInfo::default().code(&words);
159            unsafe {
160                self.device.create_shader_module(&info, None)
161                    .map_err(|e| format!("vkCreateShaderModule failed: {}", e))
162            }
163        }
164
165        /// Finds a host or device memory type satisfying filter flags and properties.
166        pub fn find_memory_type(&self, type_filter: u32, properties: vk::MemoryPropertyFlags) -> Result<u32, String> {
167            let mem_props = unsafe { self.instance.get_physical_device_memory_properties(self.physical_device) };
168            for i in 0..mem_props.memory_type_count {
169                if (type_filter & (1 << i)) != 0 && (mem_props.memory_types[i as usize].property_flags & properties) == properties {
170                    return Ok(i);
171                }
172            }
173            Err("Failed to find suitable memory type".into())
174        }
175
176        /// Allocates a GPU buffer with dedicated memory.
177        pub fn create_buffer(
178            &self,
179            size: vk::DeviceSize,
180            usage: vk::BufferUsageFlags,
181            properties: vk::MemoryPropertyFlags,
182        ) -> Result<(vk::Buffer, vk::DeviceMemory), String> {
183            let buffer_info = vk::BufferCreateInfo::default()
184                .size(size)
185                .usage(usage)
186                .sharing_mode(vk::SharingMode::EXCLUSIVE);
187
188            let buffer = unsafe {
189                self.device.create_buffer(&buffer_info, None)
190                    .map_err(|e| format!("vkCreateBuffer failed: {}", e))?
191            };
192
193            let mem_reqs = unsafe { self.device.get_buffer_memory_requirements(buffer) };
194            let mem_type_index = self.find_memory_type(mem_reqs.memory_type_bits, properties)?;
195
196            let alloc_info = vk::MemoryAllocateInfo::default()
197                .allocation_size(mem_reqs.size)
198                .memory_type_index(mem_type_index);
199
200            let memory = unsafe {
201                self.device.allocate_memory(&alloc_info, None)
202                    .map_err(|e| format!("vkAllocateMemory failed: {}", e))?
203            };
204
205            unsafe {
206                self.device.bind_buffer_memory(buffer, memory, 0)
207                    .map_err(|e| format!("vkBindBufferMemory failed: {}", e))?;
208            }
209
210            Ok((buffer, memory))
211        }
212
213        /// Dispatches 3D PGA motor rigid transformations over an [N, 3] tensor batch directly in GPU VRAM.
214        pub fn dispatch_pga_transform(
215            &self,
216            points: &[[f32; 3]],
217            motor: &PgaMotor,
218        ) -> Result<Vec<[f32; 3]>, String> {
219            let n = points.len();
220            if n == 0 {
221                return Ok(Vec::new());
222            }
223
224            let mut in_data = Vec::<[f32; 4]>::with_capacity(n);
225            for p in points {
226                in_data.push([p[0], p[1], p[2], 1.0]);
227            }
228
229            let size = (n * std::mem::size_of::<[f32; 4]>()) as vk::DeviceSize;
230
231            let (in_buf, in_mem) = self.create_buffer(
232                size,
233                vk::BufferUsageFlags::STORAGE_BUFFER,
234                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
235            )?;
236
237            let (out_buf, out_mem) = self.create_buffer(
238                size,
239                vk::BufferUsageFlags::STORAGE_BUFFER,
240                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
241            )?;
242
243            // Copy input points to GPU VRAM
244            unsafe {
245                let ptr = self.device.map_memory(in_mem, 0, size, vk::MemoryMapFlags::empty())
246                    .map_err(|e| format!("vkMapMemory failed: {}", e))? as *mut [f32; 4];
247                std::ptr::copy_nonoverlapping(in_data.as_ptr(), ptr, n);
248                self.device.unmap_memory(in_mem);
249            }
250
251            let binding_in = vk::DescriptorSetLayoutBinding::default()
252                .binding(0)
253                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
254                .descriptor_count(1)
255                .stage_flags(vk::ShaderStageFlags::COMPUTE);
256
257            let binding_out = vk::DescriptorSetLayoutBinding::default()
258                .binding(1)
259                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
260                .descriptor_count(1)
261                .stage_flags(vk::ShaderStageFlags::COMPUTE);
262
263            let bindings = [binding_in, binding_out];
264            let desc_layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings);
265            let desc_layout = unsafe {
266                self.device.create_descriptor_set_layout(&desc_layout_info, None)
267                    .map_err(|e| format!("vkCreateDescriptorSetLayout failed: {}", e))?
268            };
269
270            let push_range = vk::PushConstantRange::default()
271                .stage_flags(vk::ShaderStageFlags::COMPUTE)
272                .offset(0)
273                .size(std::mem::size_of::<PgaTransformPushConstants>() as u32);
274
275            let push_ranges = [push_range];
276            let set_layouts = [desc_layout];
277            let pipe_layout_info = vk::PipelineLayoutCreateInfo::default()
278                .set_layouts(&set_layouts)
279                .push_constant_ranges(&push_ranges);
280
281            let pipeline_layout = unsafe {
282                self.device.create_pipeline_layout(&pipe_layout_info, None)
283                    .map_err(|e| format!("vkCreatePipelineLayout failed: {}", e))?
284            };
285
286            let shader_module = self.create_shader_module(PGA_TRANSFORM_SPV)?;
287            let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
288            let stage_info = vk::PipelineShaderStageCreateInfo::default()
289                .stage(vk::ShaderStageFlags::COMPUTE)
290                .module(shader_module)
291                .name(entry_name);
292
293            let pipe_create_info = vk::ComputePipelineCreateInfo::default()
294                .stage(stage_info)
295                .layout(pipeline_layout);
296
297            let pipeline = unsafe {
298                self.device.create_compute_pipelines(vk::PipelineCache::null(), &[pipe_create_info], None)
299                    .map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
300            };
301
302            let pool_size = vk::DescriptorPoolSize::default()
303                .ty(vk::DescriptorType::STORAGE_BUFFER)
304                .descriptor_count(2);
305
306            let pool_sizes = [pool_size];
307            let pool_info = vk::DescriptorPoolCreateInfo::default()
308                .max_sets(1)
309                .pool_sizes(&pool_sizes);
310
311            let desc_pool = unsafe {
312                self.device.create_descriptor_pool(&pool_info, None)
313                    .map_err(|e| format!("vkCreateDescriptorPool failed: {}", e))?
314            };
315
316            let alloc_info = vk::DescriptorSetAllocateInfo::default()
317                .descriptor_pool(desc_pool)
318                .set_layouts(&set_layouts);
319
320            let desc_set = unsafe {
321                self.device.allocate_descriptor_sets(&alloc_info)
322                    .map_err(|e| format!("vkAllocateDescriptorSets failed: {}", e))?[0]
323            };
324
325            let buf_in_info = vk::DescriptorBufferInfo::default()
326                .buffer(in_buf)
327                .offset(0)
328                .range(size);
329            let buf_out_info = vk::DescriptorBufferInfo::default()
330                .buffer(out_buf)
331                .offset(0)
332                .range(size);
333
334            let write_in = vk::WriteDescriptorSet::default()
335                .dst_set(desc_set)
336                .dst_binding(0)
337                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
338                .buffer_info(std::slice::from_ref(&buf_in_info));
339
340            let write_out = vk::WriteDescriptorSet::default()
341                .dst_set(desc_set)
342                .dst_binding(1)
343                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
344                .buffer_info(std::slice::from_ref(&buf_out_info));
345
346            unsafe {
347                self.device.update_descriptor_sets(&[write_in, write_out], &[]);
348            }
349
350            let cmd_pool_info = vk::CommandPoolCreateInfo::default()
351                .queue_family_index(self.queue_family_index);
352            let cmd_pool = unsafe {
353                self.device.create_command_pool(&cmd_pool_info, None)
354                    .map_err(|e| format!("vkCreateCommandPool failed: {}", e))?
355            };
356
357            let cmd_alloc_info = vk::CommandBufferAllocateInfo::default()
358                .command_pool(cmd_pool)
359                .level(vk::CommandBufferLevel::PRIMARY)
360                .command_buffer_count(1);
361
362            let cmd_buf = unsafe {
363                self.device.allocate_command_buffers(&cmd_alloc_info)
364                    .map_err(|e| format!("vkAllocateCommandBuffers failed: {}", e))?[0]
365            };
366
367            let pc = PgaTransformPushConstants::new(n as u32, motor);
368
369            unsafe {
370                let begin_info = vk::CommandBufferBeginInfo::default()
371                    .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT);
372                self.device.begin_command_buffer(cmd_buf, &begin_info)
373                    .map_err(|e| format!("vkBeginCommandBuffer failed: {}", e))?;
374
375                self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
376                self.device.cmd_bind_descriptor_sets(
377                    cmd_buf,
378                    vk::PipelineBindPoint::COMPUTE,
379                    pipeline_layout,
380                    0,
381                    &[desc_set],
382                    &[],
383                );
384
385                self.device.cmd_push_constants(
386                    cmd_buf,
387                    pipeline_layout,
388                    vk::ShaderStageFlags::COMPUTE,
389                    0,
390                    bytemuck::bytes_of(&pc),
391                );
392
393                let workgroups = ((n as u32) + 63) / 64;
394                self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
395
396                self.device.end_command_buffer(cmd_buf)
397                    .map_err(|e| format!("vkEndCommandBuffer failed: {}", e))?;
398            }
399
400            let fence = unsafe {
401                self.device.create_fence(&vk::FenceCreateInfo::default(), None)
402                    .map_err(|e| format!("vkCreateFence failed: {}", e))?
403            };
404
405            let cmd_bufs = [cmd_buf];
406            let submit_info = vk::SubmitInfo::default().command_buffers(&cmd_bufs);
407            unsafe {
408                self.device.queue_submit(self.compute_queue, &[submit_info], fence)
409                    .map_err(|e| format!("vkQueueSubmit failed: {}", e))?;
410                self.device.wait_for_fences(&[fence], true, u64::MAX)
411                    .map_err(|e| format!("vkWaitForFences failed: {}", e))?;
412            }
413
414            let mut results = Vec::<[f32; 3]>::with_capacity(n);
415            unsafe {
416                let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
417                    .map_err(|e| format!("vkMapMemory failed: {}", e))? as *const [f32; 4];
418                for i in 0..n {
419                    let p4 = *ptr.add(i);
420                    results.push([p4[0], p4[1], p4[2]]);
421                }
422                self.device.unmap_memory(out_mem);
423            }
424
425            unsafe {
426                self.device.destroy_fence(fence, None);
427                self.device.destroy_command_pool(cmd_pool, None);
428                self.device.destroy_descriptor_pool(desc_pool, None);
429                self.device.destroy_pipeline(pipeline, None);
430                self.device.destroy_pipeline_layout(pipeline_layout, None);
431                self.device.destroy_shader_module(shader_module, None);
432                self.device.destroy_descriptor_set_layout(desc_layout, None);
433                self.device.destroy_buffer(in_buf, None);
434                self.device.free_memory(in_mem, None);
435                self.device.destroy_buffer(out_buf, None);
436                self.device.free_memory(out_mem, None);
437            }
438
439            Ok(results)
440        }
441
442        /// Dispatches dense residual vector addition on GPU.
443        pub fn dispatch_vec_add(&self, a: &[f32], b: &[f32]) -> Result<Vec<f32>, String> {
444            assert_eq!(a.len(), b.len(), "Vectors a and b must have identical length");
445            let n = a.len();
446            if n == 0 {
447                return Ok(Vec::new());
448            }
449
450            let size = (n * std::mem::size_of::<f32>()) as vk::DeviceSize;
451
452            let (in_buf, in_mem) = self.create_buffer(
453                size,
454                vk::BufferUsageFlags::STORAGE_BUFFER,
455                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
456            )?;
457
458            let (out_buf, out_mem) = self.create_buffer(
459                size,
460                vk::BufferUsageFlags::STORAGE_BUFFER,
461                vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
462            )?;
463
464            unsafe {
465                let ptr_in = self.device.map_memory(in_mem, 0, size, vk::MemoryMapFlags::empty())
466                    .map_err(|e| format!("vkMapMemory in failed: {}", e))? as *mut f32;
467                std::ptr::copy_nonoverlapping(a.as_ptr(), ptr_in, n);
468                self.device.unmap_memory(in_mem);
469
470                let ptr_out = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
471                    .map_err(|e| format!("vkMapMemory out failed: {}", e))? as *mut f32;
472                std::ptr::copy_nonoverlapping(b.as_ptr(), ptr_out, n);
473                self.device.unmap_memory(out_mem);
474            }
475
476            let binding0 = vk::DescriptorSetLayoutBinding::default()
477                .binding(0)
478                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
479                .descriptor_count(1)
480                .stage_flags(vk::ShaderStageFlags::COMPUTE);
481
482            let binding1 = vk::DescriptorSetLayoutBinding::default()
483                .binding(1)
484                .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
485                .descriptor_count(1)
486                .stage_flags(vk::ShaderStageFlags::COMPUTE);
487
488            let desc_layout = unsafe {
489                self.device.create_descriptor_set_layout(
490                    &vk::DescriptorSetLayoutCreateInfo::default().bindings(&[binding0, binding1]),
491                    None,
492                ).map_err(|e| format!("create_descriptor_set_layout failed: {}", e))?
493            };
494
495            let push_range = vk::PushConstantRange::default()
496                .stage_flags(vk::ShaderStageFlags::COMPUTE)
497                .offset(0)
498                .size(std::mem::size_of::<VecAddPushConstants>() as u32);
499
500            let pipeline_layout = unsafe {
501                self.device.create_pipeline_layout(
502                    &vk::PipelineLayoutCreateInfo::default()
503                        .set_layouts(&[desc_layout])
504                        .push_constant_ranges(&[push_range]),
505                    None,
506                ).map_err(|e| format!("create_pipeline_layout failed: {}", e))?
507            };
508
509            let shader_module = self.create_shader_module(VEC_ADD_SPV)?;
510            let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
511            let pipeline = unsafe {
512                self.device.create_compute_pipelines(
513                    vk::PipelineCache::null(),
514                    &[vk::ComputePipelineCreateInfo::default()
515                        .stage(vk::PipelineShaderStageCreateInfo::default()
516                            .stage(vk::ShaderStageFlags::COMPUTE)
517                            .module(shader_module)
518                            .name(entry_name))
519                        .layout(pipeline_layout)],
520                    None,
521                ).map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
522            };
523
524            let pool_size = vk::DescriptorPoolSize::default()
525                .ty(vk::DescriptorType::STORAGE_BUFFER)
526                .descriptor_count(2);
527            let desc_pool = unsafe {
528                self.device.create_descriptor_pool(
529                    &vk::DescriptorPoolCreateInfo::default().max_sets(1).pool_sizes(&[pool_size]),
530                    None,
531                ).map_err(|e| format!("create_descriptor_pool failed: {}", e))?
532            };
533
534            let desc_set = unsafe {
535                self.device.allocate_descriptor_sets(
536                    &vk::DescriptorSetAllocateInfo::default().descriptor_pool(desc_pool).set_layouts(&[desc_layout]),
537                ).map_err(|e| format!("allocate_descriptor_sets failed: {}", e))?[0]
538            };
539
540            let buf0_info = vk::DescriptorBufferInfo::default().buffer(in_buf).offset(0).range(size);
541            let buf1_info = vk::DescriptorBufferInfo::default().buffer(out_buf).offset(0).range(size);
542            let write0 = vk::WriteDescriptorSet::default().dst_set(desc_set).dst_binding(0).descriptor_type(vk::DescriptorType::STORAGE_BUFFER).buffer_info(std::slice::from_ref(&buf0_info));
543            let write1 = vk::WriteDescriptorSet::default().dst_set(desc_set).dst_binding(1).descriptor_type(vk::DescriptorType::STORAGE_BUFFER).buffer_info(std::slice::from_ref(&buf1_info));
544
545            unsafe {
546                self.device.update_descriptor_sets(&[write0, write1], &[]);
547            }
548
549            let cmd_pool = unsafe {
550                self.device.create_command_pool(
551                    &vk::CommandPoolCreateInfo::default().queue_family_index(self.queue_family_index),
552                    None,
553                ).map_err(|e| format!("create_command_pool failed: {}", e))?
554            };
555
556            let cmd_buf = unsafe {
557                self.device.allocate_command_buffers(
558                    &vk::CommandBufferAllocateInfo::default().command_pool(cmd_pool).level(vk::CommandBufferLevel::PRIMARY).command_buffer_count(1),
559                ).map_err(|e| format!("allocate_command_buffers failed: {}", e))?[0]
560            };
561
562            let pc = VecAddPushConstants::new(n as u32);
563
564            unsafe {
565                self.device.begin_command_buffer(cmd_buf, &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT))
566                    .map_err(|e| format!("begin_command_buffer failed: {}", e))?;
567                self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
568                self.device.cmd_bind_descriptor_sets(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline_layout, 0, &[desc_set], &[]);
569                self.device.cmd_push_constants(cmd_buf, pipeline_layout, vk::ShaderStageFlags::COMPUTE, 0, bytemuck::bytes_of(&pc));
570                let workgroups = ((n as u32) + 63) / 64;
571                self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
572                self.device.end_command_buffer(cmd_buf)
573                    .map_err(|e| format!("end_command_buffer failed: {}", e))?;
574            }
575
576            let fence = unsafe { self.device.create_fence(&vk::FenceCreateInfo::default(), None).map_err(|e| format!("create_fence failed: {}", e))? };
577            unsafe {
578                self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::default().command_buffers(&[cmd_buf])], fence)
579                    .map_err(|e| format!("queue_submit failed: {}", e))?;
580                self.device.wait_for_fences(&[fence], true, u64::MAX)
581                    .map_err(|e| format!("wait_for_fences failed: {}", e))?;
582            }
583
584            let mut results = vec![0.0f32; n];
585            unsafe {
586                let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
587                    .map_err(|e| format!("map_memory failed: {}", e))? as *const f32;
588                std::ptr::copy_nonoverlapping(ptr, results.as_mut_ptr(), n);
589                self.device.unmap_memory(out_mem);
590            }
591
592            unsafe {
593                self.device.destroy_fence(fence, None);
594                self.device.destroy_command_pool(cmd_pool, None);
595                self.device.destroy_descriptor_pool(desc_pool, None);
596                self.device.destroy_pipeline(pipeline, None);
597                self.device.destroy_pipeline_layout(pipeline_layout, None);
598                self.device.destroy_shader_module(shader_module, None);
599                self.device.destroy_descriptor_set_layout(desc_layout, None);
600                self.device.destroy_buffer(in_buf, None);
601                self.device.free_memory(in_mem, None);
602                self.device.destroy_buffer(out_buf, None);
603                self.device.free_memory(out_mem, None);
604            }
605
606            Ok(results)
607        }
608    }
609
610    impl Drop for GpuComputeContext {
611        fn drop(&mut self) {
612            unsafe {
613                self.device.destroy_device(None);
614                self.instance.destroy_instance(None);
615            }
616        }
617    }
618}