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, 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    /// Discovered Vulkan Compute Device.
24    #[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    /// RAII guard that safely deallocates and destroys all intermediate Vulkan resources
33    /// in case of early return, error, or normal function exit.
34    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    /// Minimal zero-dependency Vulkan Compute Context for math dispatch.
99    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        /// Probes host for available Vulkan devices.
110        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        /// Initializes compute context targeting preferred GPU (preferring discrete if available).
150        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            // Pick discrete GPU if present, else first
172            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            // Find compute queue family
182            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        /// Creates a shader module from SPIR-V byte slice.
215        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        /// Finds a host or device memory type satisfying filter flags and properties.
233        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        /// Allocates a GPU buffer with dedicated memory.
244        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        /// Dispatches 3D PGA motor rigid transformations over an [N, 3] tensor batch directly in GPU VRAM.
281        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            // Copy input points to GPU VRAM
317            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        /// Dispatches dense residual vector addition on GPU.
510        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        /// Dispatches a batch Cl(3, 0) multivector operation (op_code 0..=12) on GPU VRAM.
678        /// Unary ops (8..=12: Hodge Dual, Reversion, Involution, Conjugation, Undual) do not require b.
679        /// Binary ops (0..=7: GP, Wedge, LeftContract, Hestenes, Meet, Commutator, Anticommutator, Sandwich) require b.len() == a.len().
680        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        /// Dispatches Poincare/Hyperbolic batch operations (0=Mobius Add, 1=Geodesic Distance, 2=Spherical Inversion, 3=Poincare to Klein, 4=Klein to Poincare, 5=Einstein Midpoint) on GPU VRAM.
897        /// Uniform stride: Always returns Vec<[f32; 3]>. For distance (op_code 1), distance is at [0] with [1] and [2] set to 0.0.
898        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        /// Dispatches Poincare geodesic distance batch on GPU, returning a flat Vec<f32>.
1117        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        /// Dispatches 3D Conformal Geometric Algebra (CGA) transformations (Spherical Inversion, Dilation, Translation) on GPU VRAM.
1127        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}