1pub mod shaders;
7
8pub use shaders::{
9 CgaTransformPushConstants, GaCl3BatchPushConstants, HotStratumE8PushConstants,
10 PgaTransformPushConstants, PoincareBatchPushConstants, RmsNormPushConstants,
11 RopePushConstants, SiluMulPushConstants, VecAddPushConstants, CGA_TRANSFORM_SPV,
12 GA_CL3_BATCH_SPV, HOT_STRATUM_E8_SPV, PGA_TRANSFORM_SPV, POINCARE_MOBIUS_BATCH_SPV,
13 RMSNORM_SPV, ROPE_SPV, SILU_MUL_SPV, VEC_ADD_SPV,
14};
15
16#[cfg(feature = "vulkan")]
17pub mod vulkan {
18 use ash::{vk, Device, Entry, Instance};
19 use std::ffi::CStr;
20 use crate::pga::PgaMotor;
21 use crate::gpu::shaders::*;
22
23 #[derive(Debug, Clone)]
25 pub struct DiscoveredGpu {
26 pub name: String,
27 pub is_discrete: bool,
28 pub driver_version: u32,
29 pub api_version: u32,
30 }
31
32 struct VulkanResourceGuard<'a> {
35 device: &'a Device,
36 buffers: Vec<vk::Buffer>,
37 memories: Vec<vk::DeviceMemory>,
38 descriptor_set_layouts: Vec<vk::DescriptorSetLayout>,
39 pipeline_layouts: Vec<vk::PipelineLayout>,
40 pipelines: Vec<vk::Pipeline>,
41 descriptor_pools: Vec<vk::DescriptorPool>,
42 command_pools: Vec<vk::CommandPool>,
43 fences: Vec<vk::Fence>,
44 shader_modules: Vec<vk::ShaderModule>,
45 }
46
47 impl<'a> VulkanResourceGuard<'a> {
48 fn new(device: &'a Device) -> Self {
49 Self {
50 device,
51 buffers: Vec::with_capacity(4),
52 memories: Vec::with_capacity(4),
53 descriptor_set_layouts: Vec::with_capacity(2),
54 pipeline_layouts: Vec::with_capacity(2),
55 pipelines: Vec::with_capacity(2),
56 descriptor_pools: Vec::with_capacity(2),
57 command_pools: Vec::with_capacity(2),
58 fences: Vec::with_capacity(2),
59 shader_modules: Vec::with_capacity(2),
60 }
61 }
62 }
63
64 impl<'a> Drop for VulkanResourceGuard<'a> {
65 fn drop(&mut self) {
66 unsafe {
67 for fence in self.fences.drain(..) {
68 self.device.destroy_fence(fence, None);
69 }
70 for pool in self.command_pools.drain(..) {
71 self.device.destroy_command_pool(pool, None);
72 }
73 for pool in self.descriptor_pools.drain(..) {
74 self.device.destroy_descriptor_pool(pool, None);
75 }
76 for pipeline in self.pipelines.drain(..) {
77 self.device.destroy_pipeline(pipeline, None);
78 }
79 for layout in self.pipeline_layouts.drain(..) {
80 self.device.destroy_pipeline_layout(layout, None);
81 }
82 for module in self.shader_modules.drain(..) {
83 self.device.destroy_shader_module(module, None);
84 }
85 for layout in self.descriptor_set_layouts.drain(..) {
86 self.device.destroy_descriptor_set_layout(layout, None);
87 }
88 for buf in self.buffers.drain(..) {
89 self.device.destroy_buffer(buf, None);
90 }
91 for mem in self.memories.drain(..) {
92 self.device.free_memory(mem, None);
93 }
94 }
95 }
96 }
97
98 pub struct GpuComputeContext {
100 pub entry: Entry,
101 pub instance: Instance,
102 pub physical_device: vk::PhysicalDevice,
103 pub device: Device,
104 pub compute_queue: vk::Queue,
105 pub queue_family_index: u32,
106 }
107
108 impl GpuComputeContext {
109 pub fn probe_devices() -> Result<Vec<DiscoveredGpu>, String> {
111 let entry = unsafe {
112 Entry::load().map_err(|e| format!("Vulkan loader error: {}", e))?
113 };
114 let app_info = vk::ApplicationInfo::default()
115 .api_version(vk::make_api_version(0, 1, 3, 0));
116 let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
117 let instance = unsafe {
118 entry.create_instance(&create_info, None)
119 .map_err(|e| format!("vkCreateInstance failed: {}", e))?
120 };
121
122 let pdevices = unsafe {
123 instance.enumerate_physical_devices()
124 .map_err(|e| format!("Failed to enumerate physical devices: {}", e))?
125 };
126
127 let mut list = Vec::new();
128 for &pdev in &pdevices {
129 let props = unsafe { instance.get_physical_device_properties(pdev) };
130 let name = unsafe { CStr::from_ptr(props.device_name.as_ptr()) }
131 .to_string_lossy()
132 .into_owned();
133 let is_discrete = props.device_type == vk::PhysicalDeviceType::DISCRETE_GPU;
134 list.push(DiscoveredGpu {
135 name,
136 is_discrete,
137 driver_version: props.driver_version,
138 api_version: props.api_version,
139 });
140 }
141
142 unsafe {
143 instance.destroy_instance(None);
144 }
145
146 Ok(list)
147 }
148
149 pub fn init_compute() -> Result<Self, String> {
151 let entry = unsafe {
152 Entry::load().map_err(|e| format!("Vulkan loader error: {}", e))?
153 };
154 let app_info = vk::ApplicationInfo::default()
155 .api_version(vk::make_api_version(0, 1, 3, 0));
156 let create_info = vk::InstanceCreateInfo::default().application_info(&app_info);
157 let instance = unsafe {
158 entry.create_instance(&create_info, None)
159 .map_err(|e| format!("vkCreateInstance failed: {}", e))?
160 };
161
162 let pdevices = unsafe {
163 instance.enumerate_physical_devices()
164 .map_err(|e| format!("Failed to enumerate physical devices: {}", e))?
165 };
166
167 if pdevices.is_empty() {
168 return Err("No Vulkan physical devices found".into());
169 }
170
171 let mut selected_pdev = pdevices[0];
173 for &pdev in &pdevices {
174 let props = unsafe { instance.get_physical_device_properties(pdev) };
175 if props.device_type == vk::PhysicalDeviceType::DISCRETE_GPU {
176 selected_pdev = pdev;
177 break;
178 }
179 }
180
181 let qprops = unsafe { instance.get_physical_device_queue_family_properties(selected_pdev) };
183 let queue_family_index = qprops
184 .iter()
185 .enumerate()
186 .position(|(_, p)| p.queue_flags.contains(vk::QueueFlags::COMPUTE))
187 .ok_or_else(|| "No compute queue family found".to_string())? as u32;
188
189 let priorities = [1.0f32];
190 let queue_create_info = vk::DeviceQueueCreateInfo::default()
191 .queue_family_index(queue_family_index)
192 .queue_priorities(&priorities);
193
194 let dev_create_info = vk::DeviceCreateInfo::default()
195 .queue_create_infos(std::slice::from_ref(&queue_create_info));
196
197 let device = unsafe {
198 instance.create_device(selected_pdev, &dev_create_info, None)
199 .map_err(|e| format!("vkCreateDevice failed: {}", e))?
200 };
201
202 let compute_queue = unsafe { device.get_device_queue(queue_family_index, 0) };
203
204 Ok(Self {
205 entry,
206 instance,
207 physical_device: selected_pdev,
208 device,
209 compute_queue,
210 queue_family_index,
211 })
212 }
213
214 pub fn create_shader_module(&self, spirv_bytes: &[u8]) -> Result<vk::ShaderModule, String> {
216 assert_eq!(
217 spirv_bytes.len() % 4,
218 0,
219 "SPIR-V byte length must be a multiple of 4"
220 );
221 let mut words = Vec::<u32>::with_capacity(spirv_bytes.len() / 4);
222 for chunk in spirv_bytes.chunks_exact(4) {
223 words.push(u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
224 }
225 let info = vk::ShaderModuleCreateInfo::default().code(&words);
226 unsafe {
227 self.device.create_shader_module(&info, None)
228 .map_err(|e| format!("vkCreateShaderModule failed: {}", e))
229 }
230 }
231
232 pub fn find_memory_type(&self, type_filter: u32, properties: vk::MemoryPropertyFlags) -> Result<u32, String> {
234 let mem_props = unsafe { self.instance.get_physical_device_memory_properties(self.physical_device) };
235 for i in 0..mem_props.memory_type_count {
236 if (type_filter & (1 << i)) != 0 && (mem_props.memory_types[i as usize].property_flags & properties) == properties {
237 return Ok(i);
238 }
239 }
240 Err("Failed to find suitable memory type".into())
241 }
242
243 pub fn create_buffer(
245 &self,
246 size: vk::DeviceSize,
247 usage: vk::BufferUsageFlags,
248 properties: vk::MemoryPropertyFlags,
249 ) -> Result<(vk::Buffer, vk::DeviceMemory), String> {
250 let buffer_info = vk::BufferCreateInfo::default()
251 .size(size)
252 .usage(usage)
253 .sharing_mode(vk::SharingMode::EXCLUSIVE);
254
255 let buffer = unsafe {
256 self.device.create_buffer(&buffer_info, None)
257 .map_err(|e| format!("vkCreateBuffer failed: {}", e))?
258 };
259
260 let mem_reqs = unsafe { self.device.get_buffer_memory_requirements(buffer) };
261 let mem_type_index = self.find_memory_type(mem_reqs.memory_type_bits, properties)?;
262
263 let alloc_info = vk::MemoryAllocateInfo::default()
264 .allocation_size(mem_reqs.size)
265 .memory_type_index(mem_type_index);
266
267 let memory = unsafe {
268 self.device.allocate_memory(&alloc_info, None)
269 .map_err(|e| format!("vkAllocateMemory failed: {}", e))?
270 };
271
272 unsafe {
273 self.device.bind_buffer_memory(buffer, memory, 0)
274 .map_err(|e| format!("vkBindBufferMemory failed: {}", e))?;
275 }
276
277 Ok((buffer, memory))
278 }
279
280 pub fn dispatch_pga_transform(
282 &self,
283 points: &[[f32; 3]],
284 motor: &PgaMotor,
285 ) -> Result<Vec<[f32; 3]>, String> {
286 let n = points.len();
287 if n == 0 {
288 return Ok(Vec::new());
289 }
290
291 let mut guard = VulkanResourceGuard::new(&self.device);
292
293 let mut in_data = Vec::<[f32; 4]>::with_capacity(n);
294 for p in points {
295 in_data.push([p[0], p[1], p[2], 1.0]);
296 }
297
298 let size = (n * std::mem::size_of::<[f32; 4]>()) as vk::DeviceSize;
299
300 let (in_buf, in_mem) = self.create_buffer(
301 size,
302 vk::BufferUsageFlags::STORAGE_BUFFER,
303 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
304 )?;
305 guard.buffers.push(in_buf);
306 guard.memories.push(in_mem);
307
308 let (out_buf, out_mem) = self.create_buffer(
309 size,
310 vk::BufferUsageFlags::STORAGE_BUFFER,
311 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
312 )?;
313 guard.buffers.push(out_buf);
314 guard.memories.push(out_mem);
315
316 unsafe {
318 let ptr = self.device.map_memory(in_mem, 0, size, vk::MemoryMapFlags::empty())
319 .map_err(|e| format!("vkMapMemory failed: {}", e))? as *mut [f32; 4];
320 std::ptr::copy_nonoverlapping(in_data.as_ptr(), ptr, n);
321 self.device.unmap_memory(in_mem);
322 }
323
324 let binding_in = vk::DescriptorSetLayoutBinding::default()
325 .binding(0)
326 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
327 .descriptor_count(1)
328 .stage_flags(vk::ShaderStageFlags::COMPUTE);
329
330 let binding_out = vk::DescriptorSetLayoutBinding::default()
331 .binding(1)
332 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
333 .descriptor_count(1)
334 .stage_flags(vk::ShaderStageFlags::COMPUTE);
335
336 let bindings = [binding_in, binding_out];
337 let desc_layout_info = vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings);
338 let desc_layout = unsafe {
339 self.device.create_descriptor_set_layout(&desc_layout_info, None)
340 .map_err(|e| format!("vkCreateDescriptorSetLayout failed: {}", e))?
341 };
342 guard.descriptor_set_layouts.push(desc_layout);
343
344 let push_range = vk::PushConstantRange::default()
345 .stage_flags(vk::ShaderStageFlags::COMPUTE)
346 .offset(0)
347 .size(std::mem::size_of::<PgaTransformPushConstants>() as u32);
348
349 let push_ranges = [push_range];
350 let set_layouts = [desc_layout];
351 let pipe_layout_info = vk::PipelineLayoutCreateInfo::default()
352 .set_layouts(&set_layouts)
353 .push_constant_ranges(&push_ranges);
354
355 let pipeline_layout = unsafe {
356 self.device.create_pipeline_layout(&pipe_layout_info, None)
357 .map_err(|e| format!("vkCreatePipelineLayout failed: {}", e))?
358 };
359 guard.pipeline_layouts.push(pipeline_layout);
360
361 let shader_module = self.create_shader_module(PGA_TRANSFORM_SPV)?;
362 guard.shader_modules.push(shader_module);
363
364 let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
365 let stage_info = vk::PipelineShaderStageCreateInfo::default()
366 .stage(vk::ShaderStageFlags::COMPUTE)
367 .module(shader_module)
368 .name(entry_name);
369
370 let pipe_create_info = vk::ComputePipelineCreateInfo::default()
371 .stage(stage_info)
372 .layout(pipeline_layout);
373
374 let pipeline = unsafe {
375 self.device.create_compute_pipelines(vk::PipelineCache::null(), &[pipe_create_info], None)
376 .map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
377 };
378 guard.pipelines.push(pipeline);
379
380 let pool_size = vk::DescriptorPoolSize::default()
381 .ty(vk::DescriptorType::STORAGE_BUFFER)
382 .descriptor_count(2);
383
384 let pool_sizes = [pool_size];
385 let pool_info = vk::DescriptorPoolCreateInfo::default()
386 .max_sets(1)
387 .pool_sizes(&pool_sizes);
388
389 let desc_pool = unsafe {
390 self.device.create_descriptor_pool(&pool_info, None)
391 .map_err(|e| format!("vkCreateDescriptorPool failed: {}", e))?
392 };
393 guard.descriptor_pools.push(desc_pool);
394
395 let alloc_info = vk::DescriptorSetAllocateInfo::default()
396 .descriptor_pool(desc_pool)
397 .set_layouts(&set_layouts);
398
399 let desc_set = unsafe {
400 self.device.allocate_descriptor_sets(&alloc_info)
401 .map_err(|e| format!("vkAllocateDescriptorSets failed: {}", e))?[0]
402 };
403
404 let buf_in_info = vk::DescriptorBufferInfo::default()
405 .buffer(in_buf)
406 .offset(0)
407 .range(size);
408 let buf_out_info = vk::DescriptorBufferInfo::default()
409 .buffer(out_buf)
410 .offset(0)
411 .range(size);
412
413 let write_in = vk::WriteDescriptorSet::default()
414 .dst_set(desc_set)
415 .dst_binding(0)
416 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
417 .buffer_info(std::slice::from_ref(&buf_in_info));
418
419 let write_out = vk::WriteDescriptorSet::default()
420 .dst_set(desc_set)
421 .dst_binding(1)
422 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
423 .buffer_info(std::slice::from_ref(&buf_out_info));
424
425 unsafe {
426 self.device.update_descriptor_sets(&[write_in, write_out], &[]);
427 }
428
429 let cmd_pool_info = vk::CommandPoolCreateInfo::default()
430 .queue_family_index(self.queue_family_index);
431 let cmd_pool = unsafe {
432 self.device.create_command_pool(&cmd_pool_info, None)
433 .map_err(|e| format!("vkCreateCommandPool failed: {}", e))?
434 };
435 guard.command_pools.push(cmd_pool);
436
437 let cmd_alloc_info = vk::CommandBufferAllocateInfo::default()
438 .command_pool(cmd_pool)
439 .level(vk::CommandBufferLevel::PRIMARY)
440 .command_buffer_count(1);
441
442 let cmd_buf = unsafe {
443 self.device.allocate_command_buffers(&cmd_alloc_info)
444 .map_err(|e| format!("vkAllocateCommandBuffers failed: {}", e))?[0]
445 };
446
447 let pc = PgaTransformPushConstants::new(n as u32, motor);
448
449 unsafe {
450 let begin_info = vk::CommandBufferBeginInfo::default()
451 .flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT);
452 self.device.begin_command_buffer(cmd_buf, &begin_info)
453 .map_err(|e| format!("vkBeginCommandBuffer failed: {}", e))?;
454
455 self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
456 self.device.cmd_bind_descriptor_sets(
457 cmd_buf,
458 vk::PipelineBindPoint::COMPUTE,
459 pipeline_layout,
460 0,
461 &[desc_set],
462 &[],
463 );
464
465 self.device.cmd_push_constants(
466 cmd_buf,
467 pipeline_layout,
468 vk::ShaderStageFlags::COMPUTE,
469 0,
470 bytemuck::bytes_of(&pc),
471 );
472
473 let workgroups = ((n as u32) + 63) / 64;
474 self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
475
476 self.device.end_command_buffer(cmd_buf)
477 .map_err(|e| format!("vkEndCommandBuffer failed: {}", e))?;
478 }
479
480 let fence = unsafe {
481 self.device.create_fence(&vk::FenceCreateInfo::default(), None)
482 .map_err(|e| format!("vkCreateFence failed: {}", e))?
483 };
484 guard.fences.push(fence);
485
486 let cmd_bufs = [cmd_buf];
487 let submit_info = vk::SubmitInfo::default().command_buffers(&cmd_bufs);
488 unsafe {
489 self.device.queue_submit(self.compute_queue, &[submit_info], fence)
490 .map_err(|e| format!("vkQueueSubmit failed: {}", e))?;
491 self.device.wait_for_fences(&[fence], true, u64::MAX)
492 .map_err(|e| format!("vkWaitForFences failed: {}", e))?;
493 }
494
495 let mut results = Vec::<[f32; 3]>::with_capacity(n);
496 unsafe {
497 let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
498 .map_err(|e| format!("vkMapMemory failed: {}", e))? as *const [f32; 4];
499 for i in 0..n {
500 let p4 = *ptr.add(i);
501 results.push([p4[0], p4[1], p4[2]]);
502 }
503 self.device.unmap_memory(out_mem);
504 }
505
506 Ok(results)
507 }
508
509 pub fn dispatch_vec_add(&self, a: &[f32], b: &[f32]) -> Result<Vec<f32>, String> {
511 assert_eq!(a.len(), b.len(), "Vectors a and b must have identical length");
512 let n = a.len();
513 if n == 0 {
514 return Ok(Vec::new());
515 }
516
517 let mut guard = VulkanResourceGuard::new(&self.device);
518
519 let size = (n * std::mem::size_of::<f32>()) as vk::DeviceSize;
520
521 let (in_buf, in_mem) = self.create_buffer(
522 size,
523 vk::BufferUsageFlags::STORAGE_BUFFER,
524 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
525 )?;
526 guard.buffers.push(in_buf);
527 guard.memories.push(in_mem);
528
529 let (out_buf, out_mem) = self.create_buffer(
530 size,
531 vk::BufferUsageFlags::STORAGE_BUFFER,
532 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
533 )?;
534 guard.buffers.push(out_buf);
535 guard.memories.push(out_mem);
536
537 unsafe {
538 let ptr_in = self.device.map_memory(in_mem, 0, size, vk::MemoryMapFlags::empty())
539 .map_err(|e| format!("vkMapMemory in failed: {}", e))? as *mut f32;
540 std::ptr::copy_nonoverlapping(a.as_ptr(), ptr_in, n);
541 self.device.unmap_memory(in_mem);
542
543 let ptr_out = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
544 .map_err(|e| format!("vkMapMemory out failed: {}", e))? as *mut f32;
545 std::ptr::copy_nonoverlapping(b.as_ptr(), ptr_out, n);
546 self.device.unmap_memory(out_mem);
547 }
548
549 let binding0 = vk::DescriptorSetLayoutBinding::default()
550 .binding(0)
551 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
552 .descriptor_count(1)
553 .stage_flags(vk::ShaderStageFlags::COMPUTE);
554
555 let binding1 = vk::DescriptorSetLayoutBinding::default()
556 .binding(1)
557 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
558 .descriptor_count(1)
559 .stage_flags(vk::ShaderStageFlags::COMPUTE);
560
561 let desc_layout = unsafe {
562 self.device.create_descriptor_set_layout(
563 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&[binding0, binding1]),
564 None,
565 ).map_err(|e| format!("create_descriptor_set_layout failed: {}", e))?
566 };
567 guard.descriptor_set_layouts.push(desc_layout);
568
569 let push_range = vk::PushConstantRange::default()
570 .stage_flags(vk::ShaderStageFlags::COMPUTE)
571 .offset(0)
572 .size(std::mem::size_of::<VecAddPushConstants>() as u32);
573
574 let pipeline_layout = unsafe {
575 self.device.create_pipeline_layout(
576 &vk::PipelineLayoutCreateInfo::default()
577 .set_layouts(&[desc_layout])
578 .push_constant_ranges(&[push_range]),
579 None,
580 ).map_err(|e| format!("create_pipeline_layout failed: {}", e))?
581 };
582 guard.pipeline_layouts.push(pipeline_layout);
583
584 let shader_module = self.create_shader_module(VEC_ADD_SPV)?;
585 guard.shader_modules.push(shader_module);
586
587 let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
588 let pipeline = unsafe {
589 self.device.create_compute_pipelines(
590 vk::PipelineCache::null(),
591 &[vk::ComputePipelineCreateInfo::default()
592 .stage(vk::PipelineShaderStageCreateInfo::default()
593 .stage(vk::ShaderStageFlags::COMPUTE)
594 .module(shader_module)
595 .name(entry_name))
596 .layout(pipeline_layout)],
597 None,
598 ).map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
599 };
600 guard.pipelines.push(pipeline);
601
602 let pool_size = vk::DescriptorPoolSize::default()
603 .ty(vk::DescriptorType::STORAGE_BUFFER)
604 .descriptor_count(2);
605 let desc_pool = unsafe {
606 self.device.create_descriptor_pool(
607 &vk::DescriptorPoolCreateInfo::default().max_sets(1).pool_sizes(&[pool_size]),
608 None,
609 ).map_err(|e| format!("create_descriptor_pool failed: {}", e))?
610 };
611 guard.descriptor_pools.push(desc_pool);
612
613 let desc_set = unsafe {
614 self.device.allocate_descriptor_sets(
615 &vk::DescriptorSetAllocateInfo::default().descriptor_pool(desc_pool).set_layouts(&[desc_layout]),
616 ).map_err(|e| format!("allocate_descriptor_sets failed: {}", e))?[0]
617 };
618
619 let buf0_info = vk::DescriptorBufferInfo::default().buffer(in_buf).offset(0).range(size);
620 let buf1_info = vk::DescriptorBufferInfo::default().buffer(out_buf).offset(0).range(size);
621 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));
622 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));
623
624 unsafe {
625 self.device.update_descriptor_sets(&[write0, write1], &[]);
626 }
627
628 let cmd_pool = unsafe {
629 self.device.create_command_pool(
630 &vk::CommandPoolCreateInfo::default().queue_family_index(self.queue_family_index),
631 None,
632 ).map_err(|e| format!("create_command_pool failed: {}", e))?
633 };
634 guard.command_pools.push(cmd_pool);
635
636 let cmd_buf = unsafe {
637 self.device.allocate_command_buffers(
638 &vk::CommandBufferAllocateInfo::default().command_pool(cmd_pool).level(vk::CommandBufferLevel::PRIMARY).command_buffer_count(1),
639 ).map_err(|e| format!("allocate_command_buffers failed: {}", e))?[0]
640 };
641
642 let pc = VecAddPushConstants::new(n as u32);
643
644 unsafe {
645 self.device.begin_command_buffer(cmd_buf, &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT))
646 .map_err(|e| format!("begin_command_buffer failed: {}", e))?;
647 self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
648 self.device.cmd_bind_descriptor_sets(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline_layout, 0, &[desc_set], &[]);
649 self.device.cmd_push_constants(cmd_buf, pipeline_layout, vk::ShaderStageFlags::COMPUTE, 0, bytemuck::bytes_of(&pc));
650 let workgroups = ((n as u32) + 63) / 64;
651 self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
652 self.device.end_command_buffer(cmd_buf)
653 .map_err(|e| format!("end_command_buffer failed: {}", e))?;
654 }
655
656 let fence = unsafe { self.device.create_fence(&vk::FenceCreateInfo::default(), None).map_err(|e| format!("create_fence failed: {}", e))? };
657 guard.fences.push(fence);
658
659 unsafe {
660 self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::default().command_buffers(&[cmd_buf])], fence)
661 .map_err(|e| format!("queue_submit failed: {}", e))?;
662 self.device.wait_for_fences(&[fence], true, u64::MAX)
663 .map_err(|e| format!("wait_for_fences failed: {}", e))?;
664 }
665
666 let mut results = vec![0.0f32; n];
667 unsafe {
668 let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
669 .map_err(|e| format!("map_memory failed: {}", e))? as *const f32;
670 std::ptr::copy_nonoverlapping(ptr, results.as_mut_ptr(), n);
671 self.device.unmap_memory(out_mem);
672 }
673
674 Ok(results)
675 }
676
677 pub fn dispatch_cl3_batch(
681 &self,
682 a: &[[f32; 8]],
683 b: &[[f32; 8]],
684 op_code: u32,
685 ) -> Result<Vec<[f32; 8]>, String> {
686 let n = a.len();
687 if n == 0 {
688 return Ok(Vec::new());
689 }
690
691 let is_unary = op_code >= 8 && op_code <= 12;
692 if !is_unary && b.len() != n {
693 return Err(format!(
694 "Binary Cl(3, 0) operator op_code {} requires b.len() == a.len() (got {} vs {})",
695 op_code, b.len(), n
696 ));
697 }
698
699 let mut guard = VulkanResourceGuard::new(&self.device);
700 let size = (n * std::mem::size_of::<[f32; 8]>()) as vk::DeviceSize;
701
702 let (a_buf, a_mem) = self.create_buffer(
703 size,
704 vk::BufferUsageFlags::STORAGE_BUFFER,
705 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
706 )?;
707 guard.buffers.push(a_buf);
708 guard.memories.push(a_mem);
709
710 let (b_buf, b_mem) = self.create_buffer(
711 size,
712 vk::BufferUsageFlags::STORAGE_BUFFER,
713 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
714 )?;
715 guard.buffers.push(b_buf);
716 guard.memories.push(b_mem);
717
718 let (c_buf, c_mem) = self.create_buffer(
719 size,
720 vk::BufferUsageFlags::STORAGE_BUFFER,
721 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
722 )?;
723 guard.buffers.push(c_buf);
724 guard.memories.push(c_mem);
725
726 unsafe {
727 let ptr_a = self.device.map_memory(a_mem, 0, size, vk::MemoryMapFlags::empty())
728 .map_err(|e| format!("map a_mem failed: {}", e))? as *mut [f32; 8];
729 std::ptr::copy_nonoverlapping(a.as_ptr(), ptr_a, n);
730 self.device.unmap_memory(a_mem);
731
732 let ptr_b = self.device.map_memory(b_mem, 0, size, vk::MemoryMapFlags::empty())
733 .map_err(|e| format!("map b_mem failed: {}", e))? as *mut [f32; 8];
734 if !b.is_empty() {
735 let copy_count = b.len().min(n);
736 std::ptr::copy_nonoverlapping(b.as_ptr(), ptr_b, copy_count);
737 if copy_count < n {
738 std::ptr::write_bytes(ptr_b.add(copy_count), 0, n - copy_count);
739 }
740 } else {
741 std::ptr::write_bytes(ptr_b, 0, n);
742 }
743 self.device.unmap_memory(b_mem);
744 }
745
746 let bindings = [
747 vk::DescriptorSetLayoutBinding::default()
748 .binding(0)
749 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
750 .descriptor_count(1)
751 .stage_flags(vk::ShaderStageFlags::COMPUTE),
752 vk::DescriptorSetLayoutBinding::default()
753 .binding(1)
754 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
755 .descriptor_count(1)
756 .stage_flags(vk::ShaderStageFlags::COMPUTE),
757 vk::DescriptorSetLayoutBinding::default()
758 .binding(2)
759 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
760 .descriptor_count(1)
761 .stage_flags(vk::ShaderStageFlags::COMPUTE),
762 ];
763
764 let desc_layout = unsafe {
765 self.device.create_descriptor_set_layout(
766 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
767 None,
768 ).map_err(|e| format!("create_descriptor_set_layout failed: {}", e))?
769 };
770 guard.descriptor_set_layouts.push(desc_layout);
771
772 let push_range = vk::PushConstantRange::default()
773 .stage_flags(vk::ShaderStageFlags::COMPUTE)
774 .offset(0)
775 .size(std::mem::size_of::<GaCl3BatchPushConstants>() as u32);
776
777 let pipeline_layout = unsafe {
778 self.device.create_pipeline_layout(
779 &vk::PipelineLayoutCreateInfo::default()
780 .set_layouts(&[desc_layout])
781 .push_constant_ranges(&[push_range]),
782 None,
783 ).map_err(|e| format!("create_pipeline_layout failed: {}", e))?
784 };
785 guard.pipeline_layouts.push(pipeline_layout);
786
787 let shader_module = self.create_shader_module(GA_CL3_BATCH_SPV)?;
788 guard.shader_modules.push(shader_module);
789
790 let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
791 let stage_info = vk::PipelineShaderStageCreateInfo::default()
792 .stage(vk::ShaderStageFlags::COMPUTE)
793 .module(shader_module)
794 .name(entry_name);
795
796 let pipeline = unsafe {
797 self.device.create_compute_pipelines(
798 vk::PipelineCache::null(),
799 &[vk::ComputePipelineCreateInfo::default().stage(stage_info).layout(pipeline_layout)],
800 None,
801 ).map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
802 };
803 guard.pipelines.push(pipeline);
804
805 let pool_size = vk::DescriptorPoolSize::default()
806 .ty(vk::DescriptorType::STORAGE_BUFFER)
807 .descriptor_count(3);
808
809 let desc_pool = unsafe {
810 self.device.create_descriptor_pool(
811 &vk::DescriptorPoolCreateInfo::default().max_sets(1).pool_sizes(&[pool_size]),
812 None,
813 ).map_err(|e| format!("create_descriptor_pool failed: {}", e))?
814 };
815 guard.descriptor_pools.push(desc_pool);
816
817 let desc_set = unsafe {
818 self.device.allocate_descriptor_sets(
819 &vk::DescriptorSetAllocateInfo::default().descriptor_pool(desc_pool).set_layouts(&[desc_layout]),
820 ).map_err(|e| format!("allocate_descriptor_sets failed: {}", e))?[0]
821 };
822
823 let buf_a_info = vk::DescriptorBufferInfo::default().buffer(a_buf).offset(0).range(size);
824 let buf_b_info = vk::DescriptorBufferInfo::default().buffer(b_buf).offset(0).range(size);
825 let buf_c_info = vk::DescriptorBufferInfo::default().buffer(c_buf).offset(0).range(size);
826
827 let writes = [
828 vk::WriteDescriptorSet::default()
829 .dst_set(desc_set)
830 .dst_binding(0)
831 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
832 .buffer_info(std::slice::from_ref(&buf_a_info)),
833 vk::WriteDescriptorSet::default()
834 .dst_set(desc_set)
835 .dst_binding(1)
836 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
837 .buffer_info(std::slice::from_ref(&buf_b_info)),
838 vk::WriteDescriptorSet::default()
839 .dst_set(desc_set)
840 .dst_binding(2)
841 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
842 .buffer_info(std::slice::from_ref(&buf_c_info)),
843 ];
844
845 unsafe { self.device.update_descriptor_sets(&writes, &[]); }
846
847 let cmd_pool = unsafe {
848 self.device.create_command_pool(
849 &vk::CommandPoolCreateInfo::default().queue_family_index(self.queue_family_index),
850 None,
851 ).map_err(|e| format!("create_command_pool failed: {}", e))?
852 };
853 guard.command_pools.push(cmd_pool);
854
855 let cmd_buf = unsafe {
856 self.device.allocate_command_buffers(
857 &vk::CommandBufferAllocateInfo::default().command_pool(cmd_pool).level(vk::CommandBufferLevel::PRIMARY).command_buffer_count(1),
858 ).map_err(|e| format!("allocate_command_buffers failed: {}", e))?[0]
859 };
860
861 let pc = GaCl3BatchPushConstants::new(n as u32, op_code);
862
863 unsafe {
864 self.device.begin_command_buffer(cmd_buf, &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT))
865 .map_err(|e| format!("begin_command_buffer failed: {}", e))?;
866 self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
867 self.device.cmd_bind_descriptor_sets(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline_layout, 0, &[desc_set], &[]);
868 self.device.cmd_push_constants(cmd_buf, pipeline_layout, vk::ShaderStageFlags::COMPUTE, 0, bytemuck::bytes_of(&pc));
869 let workgroups = ((n as u32) + 63) / 64;
870 self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
871 self.device.end_command_buffer(cmd_buf)
872 .map_err(|e| format!("end_command_buffer failed: {}", e))?;
873 }
874
875 let fence = unsafe { self.device.create_fence(&vk::FenceCreateInfo::default(), None).map_err(|e| format!("create_fence failed: {}", e))? };
876 guard.fences.push(fence);
877
878 unsafe {
879 self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::default().command_buffers(&[cmd_buf])], fence)
880 .map_err(|e| format!("queue_submit failed: {}", e))?;
881 self.device.wait_for_fences(&[fence], true, u64::MAX)
882 .map_err(|e| format!("wait_for_fences failed: {}", e))?;
883 }
884
885 let mut results = vec![[0.0f32; 8]; n];
886 unsafe {
887 let ptr = self.device.map_memory(c_mem, 0, size, vk::MemoryMapFlags::empty())
888 .map_err(|e| format!("map_memory failed: {}", e))? as *const [f32; 8];
889 std::ptr::copy_nonoverlapping(ptr, results.as_mut_ptr(), n);
890 self.device.unmap_memory(c_mem);
891 }
892
893 Ok(results)
894 }
895
896 pub fn dispatch_poincare_batch(
899 &self,
900 u: &[[f32; 3]],
901 v: &[[f32; 3]],
902 op_code: u32,
903 sphere_center: [f32; 3],
904 sphere_radius: f32,
905 ) -> Result<Vec<[f32; 3]>, String> {
906 let n = u.len();
907 if n == 0 {
908 return Ok(Vec::new());
909 }
910
911 let is_binary = op_code == 0 || op_code == 1 || op_code == 5;
912 if is_binary && v.len() != n {
913 return Err(format!(
914 "Binary Poincare operator op_code {} requires v.len() == u.len() (got {} vs {})",
915 op_code, v.len(), n
916 ));
917 }
918
919 let mut guard = VulkanResourceGuard::new(&self.device);
920 let size = (n * std::mem::size_of::<[f32; 3]>()) as vk::DeviceSize;
921
922 let (u_buf, u_mem) = self.create_buffer(
923 size,
924 vk::BufferUsageFlags::STORAGE_BUFFER,
925 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
926 )?;
927 guard.buffers.push(u_buf);
928 guard.memories.push(u_mem);
929
930 let (v_buf, v_mem) = self.create_buffer(
931 size,
932 vk::BufferUsageFlags::STORAGE_BUFFER,
933 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
934 )?;
935 guard.buffers.push(v_buf);
936 guard.memories.push(v_mem);
937
938 let (out_buf, out_mem) = self.create_buffer(
939 size,
940 vk::BufferUsageFlags::STORAGE_BUFFER,
941 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
942 )?;
943 guard.buffers.push(out_buf);
944 guard.memories.push(out_mem);
945
946 unsafe {
947 let ptr_u = self.device.map_memory(u_mem, 0, size, vk::MemoryMapFlags::empty())
948 .map_err(|e| format!("map u_mem failed: {}", e))? as *mut [f32; 3];
949 std::ptr::copy_nonoverlapping(u.as_ptr(), ptr_u, n);
950 self.device.unmap_memory(u_mem);
951
952 let ptr_v = self.device.map_memory(v_mem, 0, size, vk::MemoryMapFlags::empty())
953 .map_err(|e| format!("map v_mem failed: {}", e))? as *mut [f32; 3];
954 if !v.is_empty() {
955 let copy_count = v.len().min(n);
956 std::ptr::copy_nonoverlapping(v.as_ptr(), ptr_v, copy_count);
957 if copy_count < n {
958 std::ptr::write_bytes(ptr_v.add(copy_count), 0, n - copy_count);
959 }
960 } else {
961 std::ptr::write_bytes(ptr_v, 0, n);
962 }
963 self.device.unmap_memory(v_mem);
964 }
965
966 let bindings = [
967 vk::DescriptorSetLayoutBinding::default()
968 .binding(0)
969 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
970 .descriptor_count(1)
971 .stage_flags(vk::ShaderStageFlags::COMPUTE),
972 vk::DescriptorSetLayoutBinding::default()
973 .binding(1)
974 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
975 .descriptor_count(1)
976 .stage_flags(vk::ShaderStageFlags::COMPUTE),
977 vk::DescriptorSetLayoutBinding::default()
978 .binding(2)
979 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
980 .descriptor_count(1)
981 .stage_flags(vk::ShaderStageFlags::COMPUTE),
982 ];
983
984 let desc_layout = unsafe {
985 self.device.create_descriptor_set_layout(
986 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
987 None,
988 ).map_err(|e| format!("create_descriptor_set_layout failed: {}", e))?
989 };
990 guard.descriptor_set_layouts.push(desc_layout);
991
992 let push_range = vk::PushConstantRange::default()
993 .stage_flags(vk::ShaderStageFlags::COMPUTE)
994 .offset(0)
995 .size(std::mem::size_of::<PoincareBatchPushConstants>() as u32);
996
997 let pipeline_layout = unsafe {
998 self.device.create_pipeline_layout(
999 &vk::PipelineLayoutCreateInfo::default()
1000 .set_layouts(&[desc_layout])
1001 .push_constant_ranges(&[push_range]),
1002 None,
1003 ).map_err(|e| format!("create_pipeline_layout failed: {}", e))?
1004 };
1005 guard.pipeline_layouts.push(pipeline_layout);
1006
1007 let shader_module = self.create_shader_module(POINCARE_MOBIUS_BATCH_SPV)?;
1008 guard.shader_modules.push(shader_module);
1009
1010 let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
1011 let stage_info = vk::PipelineShaderStageCreateInfo::default()
1012 .stage(vk::ShaderStageFlags::COMPUTE)
1013 .module(shader_module)
1014 .name(entry_name);
1015
1016 let pipeline = unsafe {
1017 self.device.create_compute_pipelines(
1018 vk::PipelineCache::null(),
1019 &[vk::ComputePipelineCreateInfo::default().stage(stage_info).layout(pipeline_layout)],
1020 None,
1021 ).map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
1022 };
1023 guard.pipelines.push(pipeline);
1024
1025 let pool_size = vk::DescriptorPoolSize::default()
1026 .ty(vk::DescriptorType::STORAGE_BUFFER)
1027 .descriptor_count(3);
1028
1029 let desc_pool = unsafe {
1030 self.device.create_descriptor_pool(
1031 &vk::DescriptorPoolCreateInfo::default().max_sets(1).pool_sizes(&[pool_size]),
1032 None,
1033 ).map_err(|e| format!("create_descriptor_pool failed: {}", e))?
1034 };
1035 guard.descriptor_pools.push(desc_pool);
1036
1037 let desc_set = unsafe {
1038 self.device.allocate_descriptor_sets(
1039 &vk::DescriptorSetAllocateInfo::default().descriptor_pool(desc_pool).set_layouts(&[desc_layout]),
1040 ).map_err(|e| format!("allocate_descriptor_sets failed: {}", e))?[0]
1041 };
1042
1043 let buf_u_info = vk::DescriptorBufferInfo::default().buffer(u_buf).offset(0).range(size);
1044 let buf_v_info = vk::DescriptorBufferInfo::default().buffer(v_buf).offset(0).range(size);
1045 let buf_out_info = vk::DescriptorBufferInfo::default().buffer(out_buf).offset(0).range(size);
1046
1047 let writes = [
1048 vk::WriteDescriptorSet::default()
1049 .dst_set(desc_set)
1050 .dst_binding(0)
1051 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1052 .buffer_info(std::slice::from_ref(&buf_u_info)),
1053 vk::WriteDescriptorSet::default()
1054 .dst_set(desc_set)
1055 .dst_binding(1)
1056 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1057 .buffer_info(std::slice::from_ref(&buf_v_info)),
1058 vk::WriteDescriptorSet::default()
1059 .dst_set(desc_set)
1060 .dst_binding(2)
1061 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1062 .buffer_info(std::slice::from_ref(&buf_out_info)),
1063 ];
1064
1065 unsafe { self.device.update_descriptor_sets(&writes, &[]); }
1066
1067 let cmd_pool = unsafe {
1068 self.device.create_command_pool(
1069 &vk::CommandPoolCreateInfo::default().queue_family_index(self.queue_family_index),
1070 None,
1071 ).map_err(|e| format!("create_command_pool failed: {}", e))?
1072 };
1073 guard.command_pools.push(cmd_pool);
1074
1075 let cmd_buf = unsafe {
1076 self.device.allocate_command_buffers(
1077 &vk::CommandBufferAllocateInfo::default().command_pool(cmd_pool).level(vk::CommandBufferLevel::PRIMARY).command_buffer_count(1),
1078 ).map_err(|e| format!("allocate_command_buffers failed: {}", e))?[0]
1079 };
1080
1081 let pc = PoincareBatchPushConstants::new(n as u32, op_code, sphere_center, sphere_radius);
1082
1083 unsafe {
1084 self.device.begin_command_buffer(cmd_buf, &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT))
1085 .map_err(|e| format!("begin_command_buffer failed: {}", e))?;
1086 self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
1087 self.device.cmd_bind_descriptor_sets(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline_layout, 0, &[desc_set], &[]);
1088 self.device.cmd_push_constants(cmd_buf, pipeline_layout, vk::ShaderStageFlags::COMPUTE, 0, bytemuck::bytes_of(&pc));
1089 let workgroups = ((n as u32) + 63) / 64;
1090 self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
1091 self.device.end_command_buffer(cmd_buf)
1092 .map_err(|e| format!("end_command_buffer failed: {}", e))?;
1093 }
1094
1095 let fence = unsafe { self.device.create_fence(&vk::FenceCreateInfo::default(), None).map_err(|e| format!("create_fence failed: {}", e))? };
1096 guard.fences.push(fence);
1097
1098 unsafe {
1099 self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::default().command_buffers(&[cmd_buf])], fence)
1100 .map_err(|e| format!("queue_submit failed: {}", e))?;
1101 self.device.wait_for_fences(&[fence], true, u64::MAX)
1102 .map_err(|e| format!("wait_for_fences failed: {}", e))?;
1103 }
1104
1105 let mut results = vec![[0.0f32; 3]; n];
1106 unsafe {
1107 let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
1108 .map_err(|e| format!("map_memory failed: {}", e))? as *const [f32; 3];
1109 std::ptr::copy_nonoverlapping(ptr, results.as_mut_ptr(), n);
1110 self.device.unmap_memory(out_mem);
1111 }
1112
1113 Ok(results)
1114 }
1115
1116 pub fn dispatch_poincare_distance(
1118 &self,
1119 u: &[[f32; 3]],
1120 v: &[[f32; 3]],
1121 ) -> Result<Vec<f32>, String> {
1122 let pts = self.dispatch_poincare_batch(u, v, 1, [0.0; 3], 1.0)?;
1123 Ok(pts.into_iter().map(|p| p[0]).collect())
1124 }
1125
1126 pub fn dispatch_cga_transform(
1128 &self,
1129 points: &[[f32; 3]],
1130 pc: &CgaTransformPushConstants,
1131 ) -> Result<Vec<[f32; 3]>, String> {
1132 let n = points.len();
1133 if n == 0 {
1134 return Ok(Vec::new());
1135 }
1136
1137 let mut guard = VulkanResourceGuard::new(&self.device);
1138
1139 let mut in_data = Vec::<[f32; 4]>::with_capacity(n);
1140 for p in points {
1141 in_data.push([p[0], p[1], p[2], 1.0]);
1142 }
1143
1144 let size = (n * std::mem::size_of::<[f32; 4]>()) as vk::DeviceSize;
1145
1146 let (in_buf, in_mem) = self.create_buffer(
1147 size,
1148 vk::BufferUsageFlags::STORAGE_BUFFER,
1149 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
1150 )?;
1151 guard.buffers.push(in_buf);
1152 guard.memories.push(in_mem);
1153
1154 let (out_buf, out_mem) = self.create_buffer(
1155 size,
1156 vk::BufferUsageFlags::STORAGE_BUFFER,
1157 vk::MemoryPropertyFlags::HOST_VISIBLE | vk::MemoryPropertyFlags::HOST_COHERENT,
1158 )?;
1159 guard.buffers.push(out_buf);
1160 guard.memories.push(out_mem);
1161
1162 unsafe {
1163 let ptr = self.device.map_memory(in_mem, 0, size, vk::MemoryMapFlags::empty())
1164 .map_err(|e| format!("map in_mem failed: {}", e))? as *mut [f32; 4];
1165 std::ptr::copy_nonoverlapping(in_data.as_ptr(), ptr, n);
1166 self.device.unmap_memory(in_mem);
1167 }
1168
1169 let bindings = [
1170 vk::DescriptorSetLayoutBinding::default()
1171 .binding(0)
1172 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1173 .descriptor_count(1)
1174 .stage_flags(vk::ShaderStageFlags::COMPUTE),
1175 vk::DescriptorSetLayoutBinding::default()
1176 .binding(1)
1177 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1178 .descriptor_count(1)
1179 .stage_flags(vk::ShaderStageFlags::COMPUTE),
1180 ];
1181
1182 let desc_layout = unsafe {
1183 self.device.create_descriptor_set_layout(
1184 &vk::DescriptorSetLayoutCreateInfo::default().bindings(&bindings),
1185 None,
1186 ).map_err(|e| format!("create_descriptor_set_layout failed: {}", e))?
1187 };
1188 guard.descriptor_set_layouts.push(desc_layout);
1189
1190 let push_range = vk::PushConstantRange::default()
1191 .stage_flags(vk::ShaderStageFlags::COMPUTE)
1192 .offset(0)
1193 .size(std::mem::size_of::<CgaTransformPushConstants>() as u32);
1194
1195 let pipeline_layout = unsafe {
1196 self.device.create_pipeline_layout(
1197 &vk::PipelineLayoutCreateInfo::default()
1198 .set_layouts(&[desc_layout])
1199 .push_constant_ranges(&[push_range]),
1200 None,
1201 ).map_err(|e| format!("create_pipeline_layout failed: {}", e))?
1202 };
1203 guard.pipeline_layouts.push(pipeline_layout);
1204
1205 let shader_module = self.create_shader_module(CGA_TRANSFORM_SPV)?;
1206 guard.shader_modules.push(shader_module);
1207
1208 let entry_name = CStr::from_bytes_with_nul(b"main\0").unwrap();
1209 let stage_info = vk::PipelineShaderStageCreateInfo::default()
1210 .stage(vk::ShaderStageFlags::COMPUTE)
1211 .module(shader_module)
1212 .name(entry_name);
1213
1214 let pipeline = unsafe {
1215 self.device.create_compute_pipelines(
1216 vk::PipelineCache::null(),
1217 &[vk::ComputePipelineCreateInfo::default().stage(stage_info).layout(pipeline_layout)],
1218 None,
1219 ).map_err(|(_, e)| format!("create_compute_pipelines failed: {:?}", e))?[0]
1220 };
1221 guard.pipelines.push(pipeline);
1222
1223 let pool_size = vk::DescriptorPoolSize::default()
1224 .ty(vk::DescriptorType::STORAGE_BUFFER)
1225 .descriptor_count(2);
1226
1227 let desc_pool = unsafe {
1228 self.device.create_descriptor_pool(
1229 &vk::DescriptorPoolCreateInfo::default().max_sets(1).pool_sizes(&[pool_size]),
1230 None,
1231 ).map_err(|e| format!("create_descriptor_pool failed: {}", e))?
1232 };
1233 guard.descriptor_pools.push(desc_pool);
1234
1235 let desc_set = unsafe {
1236 self.device.allocate_descriptor_sets(
1237 &vk::DescriptorSetAllocateInfo::default().descriptor_pool(desc_pool).set_layouts(&[desc_layout]),
1238 ).map_err(|e| format!("allocate_descriptor_sets failed: {}", e))?[0]
1239 };
1240
1241 let buf_in_info = vk::DescriptorBufferInfo::default().buffer(in_buf).offset(0).range(size);
1242 let buf_out_info = vk::DescriptorBufferInfo::default().buffer(out_buf).offset(0).range(size);
1243
1244 let writes = [
1245 vk::WriteDescriptorSet::default()
1246 .dst_set(desc_set)
1247 .dst_binding(0)
1248 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1249 .buffer_info(std::slice::from_ref(&buf_in_info)),
1250 vk::WriteDescriptorSet::default()
1251 .dst_set(desc_set)
1252 .dst_binding(1)
1253 .descriptor_type(vk::DescriptorType::STORAGE_BUFFER)
1254 .buffer_info(std::slice::from_ref(&buf_out_info)),
1255 ];
1256
1257 unsafe { self.device.update_descriptor_sets(&writes, &[]); }
1258
1259 let cmd_pool = unsafe {
1260 self.device.create_command_pool(
1261 &vk::CommandPoolCreateInfo::default().queue_family_index(self.queue_family_index),
1262 None,
1263 ).map_err(|e| format!("create_command_pool failed: {}", e))?
1264 };
1265 guard.command_pools.push(cmd_pool);
1266
1267 let cmd_buf = unsafe {
1268 self.device.allocate_command_buffers(
1269 &vk::CommandBufferAllocateInfo::default().command_pool(cmd_pool).level(vk::CommandBufferLevel::PRIMARY).command_buffer_count(1),
1270 ).map_err(|e| format!("allocate_command_buffers failed: {}", e))?[0]
1271 };
1272
1273 unsafe {
1274 self.device.begin_command_buffer(cmd_buf, &vk::CommandBufferBeginInfo::default().flags(vk::CommandBufferUsageFlags::ONE_TIME_SUBMIT))
1275 .map_err(|e| format!("begin_command_buffer failed: {}", e))?;
1276 self.device.cmd_bind_pipeline(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline);
1277 self.device.cmd_bind_descriptor_sets(cmd_buf, vk::PipelineBindPoint::COMPUTE, pipeline_layout, 0, &[desc_set], &[]);
1278 self.device.cmd_push_constants(cmd_buf, pipeline_layout, vk::ShaderStageFlags::COMPUTE, 0, bytemuck::bytes_of(pc));
1279 let workgroups = ((n as u32) + 63) / 64;
1280 self.device.cmd_dispatch(cmd_buf, workgroups, 1, 1);
1281 self.device.end_command_buffer(cmd_buf)
1282 .map_err(|e| format!("end_command_buffer failed: {}", e))?;
1283 }
1284
1285 let fence = unsafe { self.device.create_fence(&vk::FenceCreateInfo::default(), None).map_err(|e| format!("create_fence failed: {}", e))? };
1286 guard.fences.push(fence);
1287
1288 unsafe {
1289 self.device.queue_submit(self.compute_queue, &[vk::SubmitInfo::default().command_buffers(&[cmd_buf])], fence)
1290 .map_err(|e| format!("queue_submit failed: {}", e))?;
1291 self.device.wait_for_fences(&[fence], true, u64::MAX)
1292 .map_err(|e| format!("wait_for_fences failed: {}", e))?;
1293 }
1294
1295 let mut results = Vec::<[f32; 3]>::with_capacity(n);
1296 unsafe {
1297 let ptr = self.device.map_memory(out_mem, 0, size, vk::MemoryMapFlags::empty())
1298 .map_err(|e| format!("map_memory failed: {}", e))? as *const [f32; 4];
1299 for i in 0..n {
1300 let p4 = *ptr.add(i);
1301 results.push([p4[0], p4[1], p4[2]]);
1302 }
1303 self.device.unmap_memory(out_mem);
1304 }
1305
1306 Ok(results)
1307 }
1308 }
1309
1310 impl Drop for GpuComputeContext {
1311 fn drop(&mut self) {
1312 unsafe {
1313 self.device.destroy_device(None);
1314 self.instance.destroy_instance(None);
1315 }
1316 }
1317 }
1318}