1pub 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 #[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 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 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 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 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 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 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 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 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 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 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 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}