1use crate::context::GpuContext;
7use crate::error::{GpuError, GpuResult};
8use std::collections::HashMap;
9use tracing::{debug, info};
10
11#[derive(Debug, Clone)]
13pub struct VulkanOptimizationConfig {
14 pub enable_subgroup_ops: bool,
16 pub enable_push_constants: bool,
18 pub enable_timeline_semaphores: bool,
20 pub descriptor_pool_size: u32,
22 pub enable_async_compute: bool,
24}
25
26impl Default for VulkanOptimizationConfig {
27 fn default() -> Self {
28 Self {
29 enable_subgroup_ops: true,
30 enable_push_constants: true,
31 enable_timeline_semaphores: true,
32 descriptor_pool_size: 1000,
33 enable_async_compute: true,
34 }
35 }
36}
37
38pub struct VulkanFeatureDetector {
40 features: VulkanFeatures,
41}
42
43#[derive(Debug, Clone)]
44pub struct VulkanFeatures {
45 pub subgroup_size: u32,
47 pub subgroup_arithmetic: bool,
49 pub subgroup_ballot: bool,
51 pub subgroup_shuffle: bool,
53 pub timeline_semaphores: bool,
55 pub max_push_constants_size: u32,
57 pub async_compute: bool,
59}
60
61impl Default for VulkanFeatures {
62 fn default() -> Self {
63 Self {
64 subgroup_size: 32,
65 subgroup_arithmetic: true,
66 subgroup_ballot: true,
67 subgroup_shuffle: true,
68 timeline_semaphores: true,
69 max_push_constants_size: 128,
70 async_compute: true,
71 }
72 }
73}
74
75impl VulkanFeatureDetector {
76 pub fn new(context: &GpuContext) -> Self {
78 let features = Self::detect_features(context);
79 info!(
80 "Vulkan features: subgroup_size={}, arithmetic={}, ballot={}, shuffle={}",
81 features.subgroup_size,
82 features.subgroup_arithmetic,
83 features.subgroup_ballot,
84 features.subgroup_shuffle
85 );
86
87 Self { features }
88 }
89
90 pub fn features(&self) -> &VulkanFeatures {
92 &self.features
93 }
94
95 fn detect_features(context: &GpuContext) -> VulkanFeatures {
96 let device_features = context.device().features();
102 let subgroup = device_features.contains(wgpu::Features::SUBGROUP);
103
104 let adapter_info = context.adapter_info();
109 let subgroup_size = adapter_info.subgroup_max_size.max(1);
110
111 VulkanFeatures {
112 subgroup_size,
113 subgroup_arithmetic: subgroup,
114 subgroup_ballot: subgroup,
115 subgroup_shuffle: subgroup,
116 timeline_semaphores: true,
117 max_push_constants_size: 128,
118 async_compute: true,
119 }
120 }
121}
122
123pub struct SubgroupOptimizer {
125 features: VulkanFeatures,
126 config: VulkanOptimizationConfig,
127}
128
129impl SubgroupOptimizer {
130 pub fn new(features: VulkanFeatures, config: VulkanOptimizationConfig) -> Self {
132 Self { features, config }
133 }
134
135 pub fn optimize_shader(&self, shader_code: &str) -> String {
158 if !self.config.enable_subgroup_ops {
159 return shader_code.to_string();
160 }
161
162 let mut prologue = String::new();
163
164 if !shader_code.contains("SUBGROUP_SIZE") {
166 prologue.push_str(&format!(
167 "const SUBGROUP_SIZE: u32 = {}u;\n",
168 self.features.subgroup_size
169 ));
170 }
171
172 let needs_emulation = !self.features.subgroup_arithmetic || !self.features.subgroup_ballot;
175 if needs_emulation {
176 prologue.push_str(Self::emulation_scratch_decl());
177 }
178
179 let mut helpers = String::new();
180 helpers.push_str(&Self::subgroup_arithmetic_helpers(
181 self.features.subgroup_arithmetic,
182 ));
183 helpers.push_str(&Self::subgroup_ballot_helpers(
184 self.features.subgroup_ballot,
185 ));
186
187 format!("{prologue}{helpers}\n{shader_code}")
190 }
191
192 pub const EMU_MAX: u32 = 256;
197
198 fn emulation_scratch_decl() -> &'static str {
200 r#"
201// Workgroup-shared scratch backing the emulated subgroup helpers.
202const SUBGROUP_EMU_MAX: u32 = 256u;
203var<workgroup> sg_emu_scratch: array<f32, SUBGROUP_EMU_MAX>;
204var<workgroup> sg_emu_flags: array<u32, SUBGROUP_EMU_MAX>;
205"#
206 }
207
208 fn subgroup_arithmetic_helpers(native: bool) -> String {
215 if native {
216 r#"
217// Native subgroup arithmetic (device created with Features::SUBGROUP).
218// Reductions/scans run across the hardware subgroup; `lid`/`n` are unused.
219fn subgroup_add(value: f32, lid: u32, n: u32) -> f32 { return subgroupAdd(value); }
220fn subgroup_mul(value: f32, lid: u32, n: u32) -> f32 { return subgroupMul(value); }
221fn subgroup_min(value: f32, lid: u32, n: u32) -> f32 { return subgroupMin(value); }
222fn subgroup_max(value: f32, lid: u32, n: u32) -> f32 { return subgroupMax(value); }
223fn subgroup_inclusive_add(value: f32, lid: u32, n: u32) -> f32 { return subgroupInclusiveAdd(value); }
224fn subgroup_exclusive_add(value: f32, lid: u32, n: u32) -> f32 { return subgroupExclusiveAdd(value); }
225"#
226 .to_string()
227 } else {
228 r#"
229// Emulated subgroup arithmetic — workgroup-wide reduction/scan via shared
230// memory + workgroupBarrier(). Semantics match the native builtins but over
231// the whole 1-D workgroup (n active invocations, n <= SUBGROUP_EMU_MAX).
232fn subgroup_add(value: f32, lid: u32, n: u32) -> f32 {
233 sg_emu_scratch[lid] = value;
234 workgroupBarrier();
235 var stride = 1u;
236 loop {
237 if (stride >= n) { break; }
238 let idx = lid * stride * 2u;
239 if (idx + stride < n) {
240 sg_emu_scratch[idx] = sg_emu_scratch[idx] + sg_emu_scratch[idx + stride];
241 }
242 stride = stride * 2u;
243 workgroupBarrier();
244 }
245 let result = sg_emu_scratch[0];
246 workgroupBarrier();
247 return result;
248}
249fn subgroup_mul(value: f32, lid: u32, n: u32) -> f32 {
250 sg_emu_scratch[lid] = value;
251 workgroupBarrier();
252 var stride = 1u;
253 loop {
254 if (stride >= n) { break; }
255 let idx = lid * stride * 2u;
256 if (idx + stride < n) {
257 sg_emu_scratch[idx] = sg_emu_scratch[idx] * sg_emu_scratch[idx + stride];
258 }
259 stride = stride * 2u;
260 workgroupBarrier();
261 }
262 let result = sg_emu_scratch[0];
263 workgroupBarrier();
264 return result;
265}
266fn subgroup_min(value: f32, lid: u32, n: u32) -> f32 {
267 sg_emu_scratch[lid] = value;
268 workgroupBarrier();
269 var stride = 1u;
270 loop {
271 if (stride >= n) { break; }
272 let idx = lid * stride * 2u;
273 if (idx + stride < n) {
274 sg_emu_scratch[idx] = min(sg_emu_scratch[idx], sg_emu_scratch[idx + stride]);
275 }
276 stride = stride * 2u;
277 workgroupBarrier();
278 }
279 let result = sg_emu_scratch[0];
280 workgroupBarrier();
281 return result;
282}
283fn subgroup_max(value: f32, lid: u32, n: u32) -> f32 {
284 sg_emu_scratch[lid] = value;
285 workgroupBarrier();
286 var stride = 1u;
287 loop {
288 if (stride >= n) { break; }
289 let idx = lid * stride * 2u;
290 if (idx + stride < n) {
291 sg_emu_scratch[idx] = max(sg_emu_scratch[idx], sg_emu_scratch[idx + stride]);
292 }
293 stride = stride * 2u;
294 workgroupBarrier();
295 }
296 let result = sg_emu_scratch[0];
297 workgroupBarrier();
298 return result;
299}
300fn subgroup_inclusive_add(value: f32, lid: u32, n: u32) -> f32 {
301 sg_emu_scratch[lid] = value;
302 workgroupBarrier();
303 var acc = 0.0;
304 for (var i = 0u; i <= lid; i = i + 1u) {
305 acc = acc + sg_emu_scratch[i];
306 }
307 workgroupBarrier();
308 return acc;
309}
310fn subgroup_exclusive_add(value: f32, lid: u32, n: u32) -> f32 {
311 sg_emu_scratch[lid] = value;
312 workgroupBarrier();
313 var acc = 0.0;
314 for (var i = 0u; i < lid; i = i + 1u) {
315 acc = acc + sg_emu_scratch[i];
316 }
317 workgroupBarrier();
318 return acc;
319}
320"#
321 .to_string()
322 }
323 }
324
325 fn subgroup_ballot_helpers(native: bool) -> String {
333 if native {
334 r#"
335// Native subgroup ballot / vote (device created with Features::SUBGROUP).
336fn subgroup_all(predicate: bool, lid: u32, n: u32) -> bool { return subgroupAll(predicate); }
337fn subgroup_any(predicate: bool, lid: u32, n: u32) -> bool { return subgroupAny(predicate); }
338fn subgroup_ballot(predicate: bool, lid: u32, n: u32) -> u32 {
339 let b = subgroupBallot(predicate);
340 return countOneBits(b.x) + countOneBits(b.y) + countOneBits(b.z) + countOneBits(b.w);
341}
342"#
343 .to_string()
344 } else {
345 r#"
346// Emulated subgroup ballot / vote — workgroup-wide via shared flags buffer.
347fn subgroup_all(predicate: bool, lid: u32, n: u32) -> bool {
348 sg_emu_flags[lid] = select(0u, 1u, predicate);
349 workgroupBarrier();
350 var acc = 1u;
351 for (var i = 0u; i < n; i = i + 1u) { acc = acc & sg_emu_flags[i]; }
352 workgroupBarrier();
353 return acc != 0u;
354}
355fn subgroup_any(predicate: bool, lid: u32, n: u32) -> bool {
356 sg_emu_flags[lid] = select(0u, 1u, predicate);
357 workgroupBarrier();
358 var acc = 0u;
359 for (var i = 0u; i < n; i = i + 1u) { acc = acc | sg_emu_flags[i]; }
360 workgroupBarrier();
361 return acc != 0u;
362}
363fn subgroup_ballot(predicate: bool, lid: u32, n: u32) -> u32 {
364 sg_emu_flags[lid] = select(0u, 1u, predicate);
365 workgroupBarrier();
366 var acc = 0u;
367 for (var i = 0u; i < n; i = i + 1u) { acc = acc + sg_emu_flags[i]; }
368 workgroupBarrier();
369 return acc;
370}
371"#
372 .to_string()
373 }
374 }
375}
376
377pub struct PushConstantsManager {
379 max_size: u32,
380 constants: HashMap<String, PushConstant>,
381}
382
383#[derive(Debug, Clone)]
384struct PushConstant {
385 name: String,
386 offset: u32,
387 size: u32,
388 data: Vec<u8>,
389}
390
391impl PushConstantsManager {
392 pub fn new(max_size: u32) -> Self {
394 Self {
395 max_size,
396 constants: HashMap::new(),
397 }
398 }
399
400 pub fn register(&mut self, name: String, size: u32) -> GpuResult<()> {
406 let offset = self.calculate_next_offset();
407
408 if offset + size > self.max_size {
409 return Err(GpuError::invalid_buffer(format!(
410 "Push constant exceeds maximum size: {} + {} > {}",
411 offset, size, self.max_size
412 )));
413 }
414
415 self.constants.insert(
416 name.clone(),
417 PushConstant {
418 name,
419 offset,
420 size,
421 data: vec![0; size as usize],
422 },
423 );
424
425 Ok(())
426 }
427
428 pub fn update(&mut self, name: &str, data: &[u8]) -> GpuResult<()> {
434 let constant = self
435 .constants
436 .get_mut(name)
437 .ok_or_else(|| GpuError::invalid_buffer("Push constant not found"))?;
438
439 if data.len() != constant.size as usize {
440 return Err(GpuError::invalid_buffer("Data size mismatch"));
441 }
442
443 constant.data.copy_from_slice(data);
444
445 debug!("Updated push constant '{}' ({} bytes)", name, data.len());
446
447 Ok(())
448 }
449
450 pub fn total_size(&self) -> u32 {
452 self.constants.values().map(|c| c.size).sum()
453 }
454
455 fn calculate_next_offset(&self) -> u32 {
456 self.constants
457 .values()
458 .map(|c| c.offset + c.size)
459 .max()
460 .unwrap_or(0)
461 }
462}
463
464pub struct DescriptorSetPool {
466 pool_size: u32,
467 allocated: u32,
468 free_sets: Vec<u32>,
469}
470
471impl DescriptorSetPool {
472 pub fn new(pool_size: u32) -> Self {
474 Self {
475 pool_size,
476 allocated: 0,
477 free_sets: Vec::new(),
478 }
479 }
480
481 pub fn allocate(&mut self) -> GpuResult<u32> {
487 if let Some(set_id) = self.free_sets.pop() {
488 debug!("Reused descriptor set {}", set_id);
489 return Ok(set_id);
490 }
491
492 if self.allocated >= self.pool_size {
493 return Err(GpuError::internal(
494 "Descriptor set pool exhausted".to_string(),
495 ));
496 }
497
498 let set_id = self.allocated;
499 self.allocated += 1;
500
501 debug!("Allocated descriptor set {}", set_id);
502
503 Ok(set_id)
504 }
505
506 pub fn free(&mut self, set_id: u32) {
508 if set_id < self.allocated {
509 self.free_sets.push(set_id);
510 debug!("Freed descriptor set {}", set_id);
511 }
512 }
513
514 pub fn reset(&mut self) {
516 self.free_sets.clear();
517 for i in 0..self.allocated {
518 self.free_sets.push(i);
519 }
520 debug!("Reset descriptor set pool");
521 }
522
523 pub fn stats(&self) -> (u32, u32, usize) {
525 (self.pool_size, self.allocated, self.free_sets.len())
526 }
527}
528
529pub struct TimelineSemaphoreManager {
531 semaphores: HashMap<u32, TimelineSemaphore>,
532 next_id: u32,
533}
534
535#[derive(Debug, Clone)]
536struct TimelineSemaphore {
537 id: u32,
538 value: u64,
539 name: String,
540}
541
542impl TimelineSemaphoreManager {
543 pub fn new() -> Self {
545 Self {
546 semaphores: HashMap::new(),
547 next_id: 0,
548 }
549 }
550
551 pub fn create(&mut self, name: String, initial_value: u64) -> u32 {
553 let id = self.next_id;
554 self.next_id += 1;
555
556 self.semaphores.insert(
557 id,
558 TimelineSemaphore {
559 id,
560 value: initial_value,
561 name: name.clone(),
562 },
563 );
564
565 debug!("Created timeline semaphore '{}' (ID: {})", name, id);
566
567 id
568 }
569
570 pub fn signal(&mut self, id: u32, value: u64) -> GpuResult<()> {
576 let sem = self
577 .semaphores
578 .get_mut(&id)
579 .ok_or_else(|| GpuError::internal("Semaphore not found"))?;
580
581 sem.value = value;
582
583 debug!("Signaled semaphore '{}' with value {}", sem.name, value);
584
585 Ok(())
586 }
587
588 pub fn wait(&self, id: u32, value: u64) -> GpuResult<bool> {
594 let sem = self
595 .semaphores
596 .get(&id)
597 .ok_or_else(|| GpuError::internal("Semaphore not found"))?;
598
599 Ok(sem.value >= value)
600 }
601
602 pub fn get_value(&self, id: u32) -> GpuResult<u64> {
608 let sem = self
609 .semaphores
610 .get(&id)
611 .ok_or_else(|| GpuError::internal("Semaphore not found"))?;
612
613 Ok(sem.value)
614 }
615
616 pub fn destroy(&mut self, id: u32) {
618 if let Some(sem) = self.semaphores.remove(&id) {
619 debug!("Destroyed timeline semaphore '{}'", sem.name);
620 }
621 }
622}
623
624impl Default for TimelineSemaphoreManager {
625 fn default() -> Self {
626 Self::new()
627 }
628}
629
630pub struct AsyncComputeQueue {
632 compute_queue: Option<QueueHandle>,
633 graphics_queue: Option<QueueHandle>,
634 transfer_queue: Option<QueueHandle>,
635}
636
637#[derive(Debug, Clone)]
638struct QueueHandle {
639 family_index: u32,
640 queue_index: u32,
641}
642
643impl AsyncComputeQueue {
644 pub fn new() -> Self {
646 Self {
647 compute_queue: Some(QueueHandle {
648 family_index: 0,
649 queue_index: 0,
650 }),
651 graphics_queue: Some(QueueHandle {
652 family_index: 0,
653 queue_index: 0,
654 }),
655 transfer_queue: None,
656 }
657 }
658
659 pub fn is_available(&self) -> bool {
661 self.compute_queue.is_some()
662 }
663
664 pub fn submit_compute(&self, _commands: &[u8]) -> GpuResult<()> {
684 if self.compute_queue.is_none() {
685 return Err(GpuError::unsupported_operation(
686 "Compute queue not available".to_string(),
687 ));
688 }
689
690 Err(GpuError::unsupported_operation(
691 "async Vulkan compute-queue submission is not implemented on the wgpu backend; \
692 use GpuContext/ComputePipeline for real GPU execution"
693 .to_string(),
694 ))
695 }
696
697 pub fn submit_graphics(&self, _commands: &[u8]) -> GpuResult<()> {
708 if self.graphics_queue.is_none() {
709 return Err(GpuError::unsupported_operation(
710 "Graphics queue not available".to_string(),
711 ));
712 }
713
714 Err(GpuError::unsupported_operation(
715 "async Vulkan graphics-queue submission is not implemented on the wgpu backend; \
716 use GpuContext/ComputePipeline for real GPU execution"
717 .to_string(),
718 ))
719 }
720
721 pub fn submit_transfer(&self, _commands: &[u8]) -> GpuResult<()> {
734 if self.transfer_queue.is_none() {
735 return self.submit_graphics(_commands);
737 }
738
739 Err(GpuError::unsupported_operation(
740 "async Vulkan transfer-queue submission is not implemented on the wgpu backend; \
741 use GpuContext/ComputePipeline for real GPU execution"
742 .to_string(),
743 ))
744 }
745}
746
747impl Default for AsyncComputeQueue {
748 fn default() -> Self {
749 Self::new()
750 }
751}
752
753#[cfg(test)]
754mod tests {
755 use super::*;
756
757 #[test]
758 fn test_vulkan_features() {
759 let features = VulkanFeatures::default();
760 assert_eq!(features.subgroup_size, 32);
761 assert!(features.subgroup_arithmetic);
762 assert!(features.subgroup_ballot);
763 }
764
765 #[test]
766 fn test_push_constants_manager() {
767 let mut manager = PushConstantsManager::new(256);
768
769 manager
770 .register("view_matrix".to_string(), 64)
771 .expect("Failed to register");
772 manager
773 .register("light_pos".to_string(), 16)
774 .expect("Failed to register");
775
776 let data = vec![0u8; 64];
777 manager
778 .update("view_matrix", &data)
779 .expect("Failed to update");
780
781 assert!(manager.total_size() <= 256);
782 }
783
784 #[test]
785 fn test_descriptor_set_pool() {
786 let mut pool = DescriptorSetPool::new(10);
787
788 let set1 = pool.allocate().expect("Failed to allocate");
789 let _set2 = pool.allocate().expect("Failed to allocate");
790
791 pool.free(set1);
792
793 let set3 = pool.allocate().expect("Failed to allocate");
794 assert_eq!(set3, set1); let (pool_size, allocated, free) = pool.stats();
797 assert_eq!(pool_size, 10);
798 assert_eq!(allocated, 2);
799 assert_eq!(free, 0);
800 }
801
802 #[test]
803 fn test_timeline_semaphore() {
804 let mut manager = TimelineSemaphoreManager::new();
805
806 let sem = manager.create("test_sem".to_string(), 0);
807
808 manager.signal(sem, 5).expect("Failed to signal");
809
810 assert_eq!(manager.get_value(sem).expect("Failed to get value"), 5);
811 assert!(manager.wait(sem, 3).expect("Failed to wait"));
812 assert!(manager.wait(sem, 5).expect("Failed to wait"));
813 }
814
815 #[test]
816 fn test_async_compute_queue() {
817 let queue = AsyncComputeQueue::new();
818 assert!(queue.is_available());
820
821 let commands = vec![0u8; 64];
825 let compute_err = queue.submit_compute(&commands);
826 assert!(
827 matches!(compute_err, Err(GpuError::UnsupportedOperation { .. })),
828 "submit_compute must return an explicit unsupported-operation error, got {compute_err:?}"
829 );
830
831 let graphics_err = queue.submit_graphics(&commands);
832 assert!(
833 matches!(graphics_err, Err(GpuError::UnsupportedOperation { .. })),
834 "submit_graphics must return an explicit unsupported-operation error, got {graphics_err:?}"
835 );
836
837 let transfer_err = queue.submit_transfer(&commands);
839 assert!(
840 matches!(transfer_err, Err(GpuError::UnsupportedOperation { .. })),
841 "submit_transfer must return an explicit unsupported-operation error, got {transfer_err:?}"
842 );
843 }
844}