1pub mod simd;
38
39#[cfg(feature = "wasm")]
40pub mod wasm;
41
42use bytemuck::{Pod, Zeroable};
43use std::borrow::Cow;
44use std::sync::Arc;
45use thiserror::Error;
46use wgpu::util::DeviceExt;
47
48use cliffy_core::GA3;
49pub use simd::{addition_simd, geometric_product_simd, sandwich_simd, SimdBatch};
50
51#[cfg(feature = "wasm")]
52pub use wasm::*;
53
54#[derive(Error, Debug)]
56pub enum GpuError {
57 #[error("Failed to request GPU adapter")]
58 AdapterNotFound,
59
60 #[error("Failed to request GPU device: {0}")]
61 DeviceRequestFailed(#[from] wgpu::RequestDeviceError),
62
63 #[error("Buffer size mismatch: expected {expected}, got {actual}")]
64 BufferSizeMismatch { expected: usize, actual: usize },
65
66 #[error("GPU computation failed: {0}")]
67 ComputeFailed(String),
68
69 #[error("WebGPU not available")]
70 WebGpuNotAvailable,
71}
72
73#[repr(C)]
85#[derive(Clone, Copy, Debug, Pod, Zeroable, Default)]
86pub struct GpuMultivector {
87 pub coeffs: [f32; 8],
88}
89
90impl GpuMultivector {
91 #[must_use]
93 pub fn zero() -> Self {
94 Self { coeffs: [0.0; 8] }
95 }
96
97 #[must_use]
99 pub fn scalar(s: f32) -> Self {
100 let mut mv = Self::zero();
101 mv.coeffs[0] = s;
102 mv
103 }
104
105 #[must_use]
107 pub fn vector(x: f32, y: f32, z: f32) -> Self {
108 let mut mv = Self::zero();
109 mv.coeffs[1] = x;
110 mv.coeffs[2] = y;
111 mv.coeffs[4] = z;
112 mv
113 }
114
115 #[must_use]
117 pub fn get_scalar(&self) -> f32 {
118 self.coeffs[0]
119 }
120
121 #[must_use]
123 pub fn get_vector(&self) -> (f32, f32, f32) {
124 (self.coeffs[1], self.coeffs[2], self.coeffs[4])
125 }
126}
127
128impl From<&GA3> for GpuMultivector {
129 fn from(mv: &GA3) -> Self {
130 let mut coeffs = [0.0f32; 8];
131 let slice = mv.as_slice();
134 for (dst, &c) in coeffs.iter_mut().zip(slice.iter()) {
137 *dst = c as f32;
138 }
139 Self { coeffs }
140 }
141}
142
143impl From<GpuMultivector> for GA3 {
144 fn from(gpu_mv: GpuMultivector) -> Self {
145 let coeffs: Vec<f64> = gpu_mv.coeffs.iter().map(|&c| f64::from(c)).collect();
146 Self::from_slice(&coeffs)
147 }
148}
149
150pub const GPU_DISPATCH_THRESHOLD: usize = 256;
153
154pub struct GpuContext {
159 device: Arc<wgpu::Device>,
160 queue: Arc<wgpu::Queue>,
161 geometric_product_pipeline: wgpu::ComputePipeline,
162 addition_pipeline: wgpu::ComputePipeline,
163 sandwich_pipeline: wgpu::ComputePipeline,
164 exp_pipeline: wgpu::ComputePipeline,
165 rotor_slerp_pipeline: wgpu::ComputePipeline,
166 bind_group_layout: wgpu::BindGroupLayout,
167}
168
169impl GpuContext {
170 #[allow(clippy::too_many_lines)]
182 pub async fn new() -> Result<Self, GpuError> {
183 let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
184 backends: wgpu::Backends::all(),
185 ..wgpu::InstanceDescriptor::new_without_display_handle()
186 });
187
188 let adapter = instance
189 .request_adapter(&wgpu::RequestAdapterOptions {
190 power_preference: wgpu::PowerPreference::HighPerformance,
191 compatible_surface: None,
192 force_fallback_adapter: false,
193 apply_limit_buckets: false,
196 })
197 .await
198 .map_err(|_| GpuError::AdapterNotFound)?;
199
200 let (device, queue) = adapter
201 .request_device(&wgpu::DeviceDescriptor {
202 label: Some("Cliffy GPU Device"),
203 required_features: wgpu::Features::empty(),
204 required_limits: wgpu::Limits::default(),
205 memory_hints: wgpu::MemoryHints::Performance,
206 experimental_features: wgpu::ExperimentalFeatures::default(),
207 trace: wgpu::Trace::default(),
208 })
209 .await?;
210
211 let device = Arc::new(device);
212 let queue = Arc::new(queue);
213
214 let shader_source = include_str!("../shaders/geometric.wgsl");
215 let shader_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
216 label: Some("Geometric Algebra Shader"),
217 source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(shader_source)),
218 });
219
220 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
221 label: Some("Geometric Compute Bind Group Layout"),
222 entries: &[
223 wgpu::BindGroupLayoutEntry {
224 binding: 0,
225 visibility: wgpu::ShaderStages::COMPUTE,
226 ty: wgpu::BindingType::Buffer {
227 ty: wgpu::BufferBindingType::Storage { read_only: true },
228 has_dynamic_offset: false,
229 min_binding_size: None,
230 },
231 count: None,
232 },
233 wgpu::BindGroupLayoutEntry {
234 binding: 1,
235 visibility: wgpu::ShaderStages::COMPUTE,
236 ty: wgpu::BindingType::Buffer {
237 ty: wgpu::BufferBindingType::Storage { read_only: true },
238 has_dynamic_offset: false,
239 min_binding_size: None,
240 },
241 count: None,
242 },
243 wgpu::BindGroupLayoutEntry {
244 binding: 2,
245 visibility: wgpu::ShaderStages::COMPUTE,
246 ty: wgpu::BindingType::Buffer {
247 ty: wgpu::BufferBindingType::Storage { read_only: false },
248 has_dynamic_offset: false,
249 min_binding_size: None,
250 },
251 count: None,
252 },
253 ],
254 });
255
256 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
257 label: Some("Geometric Compute Pipeline Layout"),
258 bind_group_layouts: &[Some(&bind_group_layout)],
259 immediate_size: 0,
260 });
261
262 let geometric_product_pipeline =
263 device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
264 label: Some("Geometric Product Pipeline"),
265 layout: Some(&pipeline_layout),
266 module: &shader_module,
267 entry_point: Some("geometric_product_kernel"),
268 compilation_options: wgpu::PipelineCompilationOptions::default(),
269 cache: None,
270 });
271
272 let addition_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
273 label: Some("Addition Pipeline"),
274 layout: Some(&pipeline_layout),
275 module: &shader_module,
276 entry_point: Some("addition_kernel"),
277 compilation_options: wgpu::PipelineCompilationOptions::default(),
278 cache: None,
279 });
280
281 let sandwich_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
282 label: Some("Sandwich Pipeline"),
283 layout: Some(&pipeline_layout),
284 module: &shader_module,
285 entry_point: Some("sandwich_kernel"),
286 compilation_options: wgpu::PipelineCompilationOptions::default(),
287 cache: None,
288 });
289
290 let exp_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
291 label: Some("Exponential Pipeline"),
292 layout: Some(&pipeline_layout),
293 module: &shader_module,
294 entry_point: Some("exp_kernel"),
295 compilation_options: wgpu::PipelineCompilationOptions::default(),
296 cache: None,
297 });
298
299 let rotor_slerp_pipeline =
300 device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
301 label: Some("Rotor Slerp Pipeline"),
302 layout: Some(&pipeline_layout),
303 module: &shader_module,
304 entry_point: Some("rotor_slerp_kernel"),
305 compilation_options: wgpu::PipelineCompilationOptions::default(),
306 cache: None,
307 });
308
309 Ok(Self {
310 device,
311 queue,
312 geometric_product_pipeline,
313 addition_pipeline,
314 sandwich_pipeline,
315 exp_pipeline,
316 rotor_slerp_pipeline,
317 bind_group_layout,
318 })
319 }
320
321 pub async fn batch_geometric_product(
331 &self,
332 a: &[GA3],
333 b: &[GA3],
334 ) -> Result<Vec<GA3>, GpuError> {
335 if a.len() != b.len() {
336 return Err(GpuError::BufferSizeMismatch {
337 expected: a.len(),
338 actual: b.len(),
339 });
340 }
341
342 let a_gpu: Vec<GpuMultivector> = a.iter().map(std::convert::Into::into).collect();
344 let b_gpu: Vec<GpuMultivector> = b.iter().map(std::convert::Into::into).collect();
345
346 let result = self.run_binary_kernel(&self.geometric_product_pipeline, &a_gpu, &b_gpu)?;
347
348 Ok(result.into_iter().map(Into::into).collect())
349 }
350
351 pub async fn batch_addition(&self, a: &[GA3], b: &[GA3]) -> Result<Vec<GA3>, GpuError> {
358 if a.len() != b.len() {
359 return Err(GpuError::BufferSizeMismatch {
360 expected: a.len(),
361 actual: b.len(),
362 });
363 }
364
365 let a_gpu: Vec<GpuMultivector> = a.iter().map(std::convert::Into::into).collect();
366 let b_gpu: Vec<GpuMultivector> = b.iter().map(std::convert::Into::into).collect();
367
368 let result = self.run_binary_kernel(&self.addition_pipeline, &a_gpu, &b_gpu)?;
369
370 Ok(result.into_iter().map(Into::into).collect())
371 }
372
373 pub async fn batch_sandwich(
382 &self,
383 rotors: &[GA3],
384 vectors: &[GA3],
385 ) -> Result<Vec<GA3>, GpuError> {
386 if rotors.len() != vectors.len() {
387 return Err(GpuError::BufferSizeMismatch {
388 expected: rotors.len(),
389 actual: vectors.len(),
390 });
391 }
392
393 let rotors_gpu: Vec<GpuMultivector> = rotors.iter().map(std::convert::Into::into).collect();
394 let vectors_gpu: Vec<GpuMultivector> =
395 vectors.iter().map(std::convert::Into::into).collect();
396
397 let result = self.run_binary_kernel(&self.sandwich_pipeline, &rotors_gpu, &vectors_gpu)?;
398
399 Ok(result.into_iter().map(Into::into).collect())
400 }
401
402 pub async fn batch_exp(&self, a: &[GA3]) -> Result<Vec<GA3>, GpuError> {
410 let a_gpu: Vec<GpuMultivector> = a.iter().map(std::convert::Into::into).collect();
411
412 let result = self.run_binary_kernel(&self.exp_pipeline, &a_gpu, &a_gpu)?;
414
415 Ok(result.into_iter().map(Into::into).collect())
416 }
417
418 pub async fn batch_rotor_slerp(
427 &self,
428 a: &[GA3],
429 b: &[GA3],
430 t: f32,
431 ) -> Result<Vec<GA3>, GpuError> {
432 if a.len() != b.len() {
433 return Err(GpuError::BufferSizeMismatch {
434 expected: a.len(),
435 actual: b.len(),
436 });
437 }
438
439 let a_gpu: Vec<GpuMultivector> = a.iter().map(std::convert::Into::into).collect();
440 let b_gpu: Vec<GpuMultivector> = b
442 .iter()
443 .map(|mv| {
444 let mut gpu_mv: GpuMultivector = mv.into();
445 gpu_mv.coeffs[0] = t;
446 gpu_mv
447 })
448 .collect();
449
450 let result = self.run_binary_kernel(&self.rotor_slerp_pipeline, &a_gpu, &b_gpu)?;
451
452 Ok(result.into_iter().map(Into::into).collect())
453 }
454
455 fn run_binary_kernel(
457 &self,
458 pipeline: &wgpu::ComputePipeline,
459 a: &[GpuMultivector],
460 b: &[GpuMultivector],
461 ) -> Result<Vec<GpuMultivector>, GpuError> {
462 let count = a.len();
463 if count == 0 {
464 return Ok(Vec::new());
465 }
466
467 let a_buffer = self
469 .device
470 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
471 label: Some("Input A Buffer"),
472 contents: bytemuck::cast_slice(a),
473 usage: wgpu::BufferUsages::STORAGE,
474 });
475
476 let b_buffer = self
477 .device
478 .create_buffer_init(&wgpu::util::BufferInitDescriptor {
479 label: Some("Input B Buffer"),
480 contents: bytemuck::cast_slice(b),
481 usage: wgpu::BufferUsages::STORAGE,
482 });
483
484 let output_size = std::mem::size_of_val(a) as u64;
486 let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
487 label: Some("Output Buffer"),
488 size: output_size,
489 usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
490 mapped_at_creation: false,
491 });
492
493 let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
495 label: Some("Staging Buffer"),
496 size: output_size,
497 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
498 mapped_at_creation: false,
499 });
500
501 let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
503 label: Some("Compute Bind Group"),
504 layout: &self.bind_group_layout,
505 entries: &[
506 wgpu::BindGroupEntry {
507 binding: 0,
508 resource: a_buffer.as_entire_binding(),
509 },
510 wgpu::BindGroupEntry {
511 binding: 1,
512 resource: b_buffer.as_entire_binding(),
513 },
514 wgpu::BindGroupEntry {
515 binding: 2,
516 resource: output_buffer.as_entire_binding(),
517 },
518 ],
519 });
520
521 let mut encoder = self
523 .device
524 .create_command_encoder(&wgpu::CommandEncoderDescriptor {
525 label: Some("Compute Encoder"),
526 });
527
528 {
529 let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
530 label: Some("Compute Pass"),
531 timestamp_writes: None,
532 });
533 compute_pass.set_pipeline(pipeline);
534 compute_pass.set_bind_group(0, &bind_group, &[]);
535
536 let workgroup_count = count.div_ceil(64) as u32;
538 compute_pass.dispatch_workgroups(workgroup_count, 1, 1);
539 }
540
541 encoder.copy_buffer_to_buffer(&output_buffer, 0, &staging_buffer, 0, output_size);
543
544 self.queue.submit(std::iter::once(encoder.finish()));
545
546 let buffer_slice = staging_buffer.slice(..);
548 let (sender, receiver) = std::sync::mpsc::channel();
549 buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
550 let _ = sender.send(result);
551 });
552
553 self.device
554 .poll(wgpu::PollType::Wait {
555 submission_index: None,
556 timeout: None,
557 })
558 .ok();
559
560 receiver
561 .recv()
562 .map_err(|e| GpuError::ComputeFailed(e.to_string()))?
563 .map_err(|e| GpuError::ComputeFailed(format!("{e:?}")))?;
564
565 let data = buffer_slice
567 .get_mapped_range()
568 .map_err(|e| GpuError::ComputeFailed(format!("{e:?}")))?;
569 let result: Vec<GpuMultivector> = bytemuck::cast_slice(&data).to_vec();
570 drop(data);
571 staging_buffer.unmap();
572
573 Ok(result)
574 }
575
576 #[must_use]
578 pub fn should_use_gpu(&self, batch_size: usize) -> bool {
579 batch_size >= GPU_DISPATCH_THRESHOLD
580 }
581
582 #[must_use]
584 pub fn device_info(&self) -> String {
585 "Cliffy GPU Context (wgpu)".to_string()
586 }
587}
588
589pub struct AutoDispatcher {
594 gpu_ctx: Option<GpuContext>,
595 threshold: usize,
596}
597
598impl AutoDispatcher {
599 pub async fn new() -> Self {
601 let gpu_ctx = GpuContext::new().await.ok();
602 Self {
603 gpu_ctx,
604 threshold: GPU_DISPATCH_THRESHOLD,
605 }
606 }
607
608 pub async fn with_threshold(threshold: usize) -> Self {
610 let gpu_ctx = GpuContext::new().await.ok();
611 Self { gpu_ctx, threshold }
612 }
613
614 #[must_use]
616 pub fn cpu_only() -> Self {
617 Self {
618 gpu_ctx: None,
619 threshold: GPU_DISPATCH_THRESHOLD,
620 }
621 }
622
623 #[must_use]
625 pub fn has_gpu(&self) -> bool {
626 self.gpu_ctx.is_some()
627 }
628
629 #[must_use]
631 pub fn threshold(&self) -> usize {
632 self.threshold
633 }
634
635 pub async fn geometric_product(&self, a: &[GA3], b: &[GA3]) -> Result<Vec<GA3>, GpuError> {
645 if let Some(ref ctx) = self.gpu_ctx {
646 if a.len() >= self.threshold {
647 return ctx.batch_geometric_product(a, b).await;
648 }
649 }
650
651 if a.len() != b.len() {
653 return Err(GpuError::BufferSizeMismatch {
654 expected: a.len(),
655 actual: b.len(),
656 });
657 }
658
659 let a_gpu = SimdBatch::from_ga3(a);
660 let b_gpu = SimdBatch::from_ga3(b);
661 let result = SimdBatch::geometric_product(&a_gpu, &b_gpu);
662 Ok(SimdBatch::to_ga3(&result))
663 }
664
665 pub async fn addition(&self, a: &[GA3], b: &[GA3]) -> Result<Vec<GA3>, GpuError> {
675 if let Some(ref ctx) = self.gpu_ctx {
676 if a.len() >= self.threshold {
677 return ctx.batch_addition(a, b).await;
678 }
679 }
680
681 if a.len() != b.len() {
683 return Err(GpuError::BufferSizeMismatch {
684 expected: a.len(),
685 actual: b.len(),
686 });
687 }
688
689 let a_gpu = SimdBatch::from_ga3(a);
690 let b_gpu = SimdBatch::from_ga3(b);
691 let result = SimdBatch::addition(&a_gpu, &b_gpu);
692 Ok(SimdBatch::to_ga3(&result))
693 }
694
695 pub async fn sandwich(&self, rotors: &[GA3], vectors: &[GA3]) -> Result<Vec<GA3>, GpuError> {
705 if let Some(ref ctx) = self.gpu_ctx {
706 if rotors.len() >= self.threshold {
707 return ctx.batch_sandwich(rotors, vectors).await;
708 }
709 }
710
711 if rotors.len() != vectors.len() {
713 return Err(GpuError::BufferSizeMismatch {
714 expected: rotors.len(),
715 actual: vectors.len(),
716 });
717 }
718
719 let rotors_gpu = SimdBatch::from_ga3(rotors);
720 let vectors_gpu = SimdBatch::from_ga3(vectors);
721 let result = SimdBatch::sandwich(&rotors_gpu, &vectors_gpu);
722 Ok(SimdBatch::to_ga3(&result))
723 }
724
725 pub async fn exp(&self, a: &[GA3]) -> Result<Vec<GA3>, GpuError> {
735 if let Some(ref ctx) = self.gpu_ctx {
736 if a.len() >= self.threshold {
737 return ctx.batch_exp(a).await;
738 }
739 }
740
741 let a_gpu = SimdBatch::from_ga3(a);
743 let result = SimdBatch::exp(&a_gpu);
744 Ok(SimdBatch::to_ga3(&result))
745 }
746
747 pub async fn rotor_slerp(&self, a: &[GA3], b: &[GA3], t: f32) -> Result<Vec<GA3>, GpuError> {
757 if let Some(ref ctx) = self.gpu_ctx {
758 if a.len() >= self.threshold {
759 return ctx.batch_rotor_slerp(a, b, t).await;
760 }
761 }
762
763 if a.len() != b.len() {
765 return Err(GpuError::BufferSizeMismatch {
766 expected: a.len(),
767 actual: b.len(),
768 });
769 }
770
771 let a_gpu = SimdBatch::from_ga3(a);
772 let b_gpu = SimdBatch::from_ga3(b);
773 let result = SimdBatch::rotor_slerp(&a_gpu, &b_gpu, t);
774 Ok(SimdBatch::to_ga3(&result))
775 }
776}
777
778#[cfg(test)]
779mod tests {
780 #![allow(clippy::float_cmp)]
781 use super::*;
784
785 #[test]
786 fn test_gpu_multivector_zero() {
787 let mv = GpuMultivector::zero();
788 assert!(mv.coeffs.iter().all(|&c| c == 0.0));
789 }
790
791 #[test]
792 fn test_gpu_multivector_scalar() {
793 let mv = GpuMultivector::scalar(5.0);
794 assert_eq!(mv.get_scalar(), 5.0);
795 }
796
797 #[test]
798 fn test_gpu_multivector_vector() {
799 let mv = GpuMultivector::vector(1.0, 2.0, 3.0);
800 assert_eq!(mv.get_vector(), (1.0, 2.0, 3.0));
801 }
802
803 #[test]
804 fn test_ga3_roundtrip() {
805 use amari_core::Vector;
806 let vec = Vector::<3, 0, 0>::from_components(1.0, 2.0, 3.0);
807 let original = GA3::from_vector(&vec);
808 let gpu: GpuMultivector = (&original).into();
809 let back: GA3 = gpu.into();
810
811 let x = back.get(1);
813 let y = back.get(2);
814 let z = back.get(4);
815 assert!((x - 1.0).abs() < 1e-5);
816 assert!((y - 2.0).abs() < 1e-5);
817 assert!((z - 3.0).abs() < 1e-5);
818 }
819
820 #[test]
821 fn test_dispatch_threshold() {
822 let dispatcher = AutoDispatcher::cpu_only();
824 assert!(!dispatcher.has_gpu());
825 assert_eq!(dispatcher.threshold(), GPU_DISPATCH_THRESHOLD);
826 }
827}