## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/trueno/src/backends/gpu/device.rs

    Line
    Count
    Source
    1
    //! GPU device initialization and management
    2
    //!
    3
    //! This module provides cross-platform GPU compute via wgpu (WebGPU).
    4
    //!
    5
    //! # Platform differences
    6
    //!
    7
    //! - **Native**: Sync wrappers available using `pollster::block_on`
    8
    //! - **WASM**: Sync wrappers unavailable (can't block main thread); use `*_async` methods
    9
    //!
    10
    //! Use `runtime::sync_available()` to check at runtime.
    11
    12
    #[cfg(any(feature = "gpu", feature = "gpu-wasm"))]
    13
    use super::runtime;
    14
    use super::shaders;
    15
    16
    /// GPU device manager
    17
    #[derive(Clone)]
    18
    pub struct GpuDevice {
    19
        pub device: wgpu::Device,
    20
        pub queue: wgpu::Queue,
    21
    }
    22
    23
    impl GpuDevice {
    24
        /// Initialize GPU device (sync, native only)
    25
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    26
    47
        pub fn new() -> Result<Self, String> {
    27
    47
            runtime::block_on(async { Self::new_async().await })
    28
    47
        }
    29
    30
        /// Initialize GPU device (async, works on all platforms)
    31
    47
        pub async fn new_async() -> Result<Self, String> {
    32
            // Create instance
    33
    47
            let instance = wgpu::Instance::default();
    34
    35
            // Request adapter (GPU)
    36
    47
            let adapter = instance
    37
    47
                .request_adapter(&wgpu::RequestAdapterOptions {
    38
    47
                    power_preference: wgpu::PowerPreference::HighPerformance,
    39
    47
                    compatible_surface: None,
    40
    47
                    force_fallback_adapter: false,
    41
    47
                })
    42
    47
                .await
    43
    47
                .map_err(|e| format!0("Failed to find GPU adapter: {}"0, e))?0;
    44
    45
            // Request device and queue
    46
    47
            let (device, queue) = adapter
    47
    47
                .request_device(&wgpu::DeviceDescriptor {
    48
    47
                    label: Some("Trueno GPU Device"),
    49
    47
                    required_features: wgpu::Features::empty(),
    50
    47
                    required_limits: wgpu::Limits::default(),
    51
    47
                    memory_hints: wgpu::MemoryHints::Performance,
    52
    47
                    experimental_features: Default::default(),
    53
    47
                    trace: Default::default(),
    54
    47
                })
    55
    47
                .await
    56
    47
                .map_err(|e| format!0("Failed to create device: {}"0, e))?0;
    57
    58
    47
            Ok(Self { device, queue })
    59
    47
        }
    60
    61
        /// Check if GPU is available (sync, native only)
    62
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    63
    155
        pub fn is_available() -> bool {
    64
    155
            runtime::block_on(Self::is_available_async())
    65
    155
        }
    66
    67
        /// Check if GPU is available (async, works on all platforms)
    68
    155
        pub async fn is_available_async() -> bool {
    69
    155
            let instance = wgpu::Instance::default();
    70
    155
            instance
    71
    155
                .request_adapter(&wgpu::RequestAdapterOptions {
    72
    155
                    power_preference: wgpu::PowerPreference::HighPerformance,
    73
    155
                    compatible_surface: None,
    74
    155
                    force_fallback_adapter: false,
    75
    155
                })
    76
    155
                .await
    77
    155
                .is_ok()
    78
    155
        }
    79
    80
        /// Execute matrix multiplication on GPU (sync, native only)
    81
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    82
    438
        pub fn matmul(
    83
    438
            &self,
    84
    438
            a: &[f32],
    85
    438
            b: &[f32],
    86
    438
            result: &mut [f32],
    87
    438
            m: usize,
    88
    438
            k: usize,
    89
    438
            n: usize,
    90
    438
        ) -> Result<(), String> {
    91
    438
            runtime::block_on(async { self.matmul_async(a, b, result, m, k, n).await })
    92
    438
        }
    93
    94
        /// Execute matrix multiplication on GPU (async, works on all platforms)
    95
    438
        pub async fn matmul_async(
    96
    438
            &self,
    97
    438
            a: &[f32],
    98
    438
            b: &[f32],
    99
    438
            result: &mut [f32],
    100
    438
            m: usize,
    101
    438
            k: usize,
    102
    438
            n: usize,
    103
    438
        ) -> Result<(), String> {
    104
            // Create shader module
    105
    438
            let shader = self
    106
    438
                .device
    107
    438
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    108
    438
                    label: Some("Matmul Shader"),
    109
    438
                    source: wgpu::ShaderSource::Wgsl(shaders::MATMUL_SHADER.into()),
    110
    438
                });
    111
    112
            // Create buffers
    113
    438
            let a_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    114
    438
                label: Some("Matrix A"),
    115
    438
                size: std::mem::size_of_val(a) as u64,
    116
    438
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    117
    438
                mapped_at_creation: false,
    118
    438
            });
    119
    120
    438
            let b_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    121
    438
                label: Some("Matrix B"),
    122
    438
                size: std::mem::size_of_val(b) as u64,
    123
    438
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    124
    438
                mapped_at_creation: false,
    125
    438
            });
    126
    127
    438
            let c_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    128
    438
                label: Some("Matrix C"),
    129
    438
                size: std::mem::size_of_val(result) as u64,
    130
    438
                usage: wgpu::BufferUsages::STORAGE
    131
    438
                    | wgpu::BufferUsages::COPY_SRC
    132
    438
                    | wgpu::BufferUsages::COPY_DST,
    133
    438
                mapped_at_creation: false,
    134
    438
            });
    135
    136
            // Dimensions uniform buffer
    137
            #[repr(C)]
    138
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    139
            struct Dimensions {
    140
                m: u32,
    141
                k: u32,
    142
                n: u32,
    143
                _padding: u32,
    144
            }
    145
    146
    438
            let dims = Dimensions {
    147
    438
                m: m as u32,
    148
    438
                k: k as u32,
    149
    438
                n: n as u32,
    150
    438
                _padding: 0,
    151
    438
            };
    152
    153
    438
            let dims_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    154
    438
                label: Some("Dimensions"),
    155
    438
                size: std::mem::size_of::<Dimensions>() as u64,
    156
    438
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    157
    438
                mapped_at_creation: false,
    158
    438
            });
    159
    160
            // Write data to buffers
    161
    438
            self.queue
    162
    438
                .write_buffer(&a_buffer, 0, bytemuck::cast_slice(a));
    163
    438
            self.queue
    164
    438
                .write_buffer(&b_buffer, 0, bytemuck::cast_slice(b));
    165
    438
            self.queue
    166
    438
                .write_buffer(&dims_buffer, 0, bytemuck::bytes_of(&dims));
    167
    168
            // Create bind group layout
    169
    438
            let bind_group_layout =
    170
    438
                self.device
    171
    438
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    172
    438
                        label: Some("Matmul Bind Group Layout"),
    173
    438
                        entries: &[
    174
    438
                            wgpu::BindGroupLayoutEntry {
    175
    438
                                binding: 0,
    176
    438
                                visibility: wgpu::ShaderStages::COMPUTE,
    177
    438
                                ty: wgpu::BindingType::Buffer {
    178
    438
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    179
    438
                                    has_dynamic_offset: false,
    180
    438
                                    min_binding_size: None,
    181
    438
                                },
    182
    438
                                count: None,
    183
    438
                            },
    184
    438
                            wgpu::BindGroupLayoutEntry {
    185
    438
                                binding: 1,
    186
    438
                                visibility: wgpu::ShaderStages::COMPUTE,
    187
    438
                                ty: wgpu::BindingType::Buffer {
    188
    438
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    189
    438
                                    has_dynamic_offset: false,
    190
    438
                                    min_binding_size: None,
    191
    438
                                },
    192
    438
                                count: None,
    193
    438
                            },
    194
    438
                            wgpu::BindGroupLayoutEntry {
    195
    438
                                binding: 2,
    196
    438
                                visibility: wgpu::ShaderStages::COMPUTE,
    197
    438
                                ty: wgpu::BindingType::Buffer {
    198
    438
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    199
    438
                                    has_dynamic_offset: false,
    200
    438
                                    min_binding_size: None,
    201
    438
                                },
    202
    438
                                count: None,
    203
    438
                            },
    204
    438
                            wgpu::BindGroupLayoutEntry {
    205
    438
                                binding: 3,
    206
    438
                                visibility: wgpu::ShaderStages::COMPUTE,
    207
    438
                                ty: wgpu::BindingType::Buffer {
    208
    438
                                    ty: wgpu::BufferBindingType::Uniform,
    209
    438
                                    has_dynamic_offset: false,
    210
    438
                                    min_binding_size: None,
    211
    438
                                },
    212
    438
                                count: None,
    213
    438
                            },
    214
    438
                        ],
    215
    438
                    });
    216
    217
            // Create bind group
    218
    438
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    219
    438
                label: Some("Matmul Bind Group"),
    220
    438
                layout: &bind_group_layout,
    221
    438
                entries: &[
    222
    438
                    wgpu::BindGroupEntry {
    223
    438
                        binding: 0,
    224
    438
                        resource: a_buffer.as_entire_binding(),
    225
    438
                    },
    226
    438
                    wgpu::BindGroupEntry {
    227
    438
                        binding: 1,
    228
    438
                        resource: b_buffer.as_entire_binding(),
    229
    438
                    },
    230
    438
                    wgpu::BindGroupEntry {
    231
    438
                        binding: 2,
    232
    438
                        resource: c_buffer.as_entire_binding(),
    233
    438
                    },
    234
    438
                    wgpu::BindGroupEntry {
    235
    438
                        binding: 3,
    236
    438
                        resource: dims_buffer.as_entire_binding(),
    237
    438
                    },
    238
    438
                ],
    239
    438
            });
    240
    241
            // Create pipeline
    242
    438
            let pipeline_layout = self
    243
    438
                .device
    244
    438
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    245
    438
                    label: Some("Matmul Pipeline Layout"),
    246
    438
                    bind_group_layouts: &[&bind_group_layout],
    247
    438
                    push_constant_ranges: &[],
    248
    438
                });
    249
    250
    438
            let pipeline = self
    251
    438
                .device
    252
    438
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    253
    438
                    label: Some("Matmul Pipeline"),
    254
    438
                    layout: Some(&pipeline_layout),
    255
    438
                    module: &shader,
    256
    438
                    entry_point: Some("main"),
    257
    438
                    compilation_options: Default::default(),
    258
    438
                    cache: None,
    259
    438
                });
    260
    261
            // Create staging buffer for reading results
    262
    438
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    263
    438
                label: Some("Staging Buffer"),
    264
    438
                size: std::mem::size_of_val(result) as u64,
    265
    438
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    266
    438
                mapped_at_creation: false,
    267
    438
            });
    268
    269
            // Create command encoder
    270
    438
            let mut encoder = self
    271
    438
                .device
    272
    438
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    273
    438
                    label: Some("Matmul Encoder"),
    274
    438
                });
    275
    276
    438
            {
    277
    438
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    278
    438
                    label: Some("Matmul Pass"),
    279
    438
                    timestamp_writes: None,
    280
    438
                });
    281
    438
                compute_pass.set_pipeline(&pipeline);
    282
    438
                compute_pass.set_bind_group(0, &bind_group, &[]);
    283
    438
    284
    438
                // Dispatch workgroups (16×16 threads per workgroup)
    285
    438
                let workgroup_size_x = 16;
    286
    438
                let workgroup_size_y = 16;
    287
    438
                let num_workgroups_x = (m as u32).div_ceil(workgroup_size_x);
    288
    438
                let num_workgroups_y = (n as u32).div_ceil(workgroup_size_y);
    289
    438
    290
    438
                compute_pass.dispatch_workgroups(num_workgroups_x, num_workgroups_y, 1);
    291
    438
            }
    292
    293
            // Copy result to staging buffer
    294
    438
            encoder.copy_buffer_to_buffer(
    295
    438
                &c_buffer,
    296
                0,
    297
    438
                &staging_buffer,
    298
                0,
    299
    438
                std::mem::size_of_val(result) as u64,
    300
            );
    301
    302
            // Submit commands
    303
    438
            self.queue.submit(Some(encoder.finish()));
    304
    305
            // Read back results
    306
    438
            let buffer_slice = staging_buffer.slice(..);
    307
    438
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    308
    438
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    309
    438
                sender.send(result).ok();
    310
    438
            });
    311
    312
            // Poll device to ensure GPU work completes and callbacks are invoked
    313
    438
            self.device
    314
    438
                .poll(wgpu::PollType::Wait {
    315
    438
                    submission_index: None,
    316
    438
                    timeout: None,
    317
    438
                })
    318
    438
                .ok();
    319
    320
    438
            receiver
    321
    438
                .receive()
    322
    438
                .await
    323
    438
                .ok_or("Failed to receive mapping result")?0
    324
    438
                .map_err(|e| format!0("Buffer mapping failed: {:?}"0, e))?0;
    325
    326
    438
            {
    327
    438
                let data = buffer_slice.get_mapped_range();
    328
    438
                result.copy_from_slice(bytemuck::cast_slice(&data));
    329
    438
            }
    330
    331
    438
            staging_buffer.unmap();
    332
    333
    438
            Ok(())
    334
    438
        }
    335
    336
        /// Execute vector addition on GPU: c = a + b (sync, native only)
    337
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    338
    0
        pub fn vec_add(&self, a: &[f32], b: &[f32], result: &mut [f32]) -> Result<(), String> {
    339
    0
            runtime::block_on(async { self.vec_add_async(a, b, result).await })
    340
    0
        }
    341
    342
        /// Execute vector addition on GPU: c = a + b (async, works on all platforms)
    343
    0
        pub async fn vec_add_async(
    344
    0
            &self,
    345
    0
            a: &[f32],
    346
    0
            b: &[f32],
    347
    0
            result: &mut [f32],
    348
    0
        ) -> Result<(), String> {
    349
    0
            let len = a.len();
    350
    351
            // Create shader module
    352
    0
            let shader = self
    353
    0
                .device
    354
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    355
    0
                    label: Some("Vec Add Shader"),
    356
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::VEC_ADD_SHADER.into()),
    357
    0
                });
    358
    359
            // Create buffers
    360
    0
            let a_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    361
    0
                label: Some("Vector A"),
    362
    0
                size: std::mem::size_of_val(a) as u64,
    363
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    364
    0
                mapped_at_creation: false,
    365
    0
            });
    366
    367
    0
            let b_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    368
    0
                label: Some("Vector B"),
    369
    0
                size: std::mem::size_of_val(b) as u64,
    370
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    371
    0
                mapped_at_creation: false,
    372
    0
            });
    373
    374
    0
            let c_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    375
    0
                label: Some("Vector C"),
    376
    0
                size: std::mem::size_of_val(result) as u64,
    377
    0
                usage: wgpu::BufferUsages::STORAGE
    378
    0
                    | wgpu::BufferUsages::COPY_SRC
    379
    0
                    | wgpu::BufferUsages::COPY_DST,
    380
    0
                mapped_at_creation: false,
    381
    0
            });
    382
    383
            // Write data to buffers
    384
    0
            self.queue
    385
    0
                .write_buffer(&a_buffer, 0, bytemuck::cast_slice(a));
    386
    0
            self.queue
    387
    0
                .write_buffer(&b_buffer, 0, bytemuck::cast_slice(b));
    388
    389
            // Create bind group layout
    390
    0
            let bind_group_layout =
    391
    0
                self.device
    392
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    393
    0
                        label: Some("Vec Add Bind Group Layout"),
    394
    0
                        entries: &[
    395
    0
                            wgpu::BindGroupLayoutEntry {
    396
    0
                                binding: 0,
    397
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    398
    0
                                ty: wgpu::BindingType::Buffer {
    399
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    400
    0
                                    has_dynamic_offset: false,
    401
    0
                                    min_binding_size: None,
    402
    0
                                },
    403
    0
                                count: None,
    404
    0
                            },
    405
    0
                            wgpu::BindGroupLayoutEntry {
    406
    0
                                binding: 1,
    407
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    408
    0
                                ty: wgpu::BindingType::Buffer {
    409
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    410
    0
                                    has_dynamic_offset: false,
    411
    0
                                    min_binding_size: None,
    412
    0
                                },
    413
    0
                                count: None,
    414
    0
                            },
    415
    0
                            wgpu::BindGroupLayoutEntry {
    416
    0
                                binding: 2,
    417
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    418
    0
                                ty: wgpu::BindingType::Buffer {
    419
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    420
    0
                                    has_dynamic_offset: false,
    421
    0
                                    min_binding_size: None,
    422
    0
                                },
    423
    0
                                count: None,
    424
    0
                            },
    425
    0
                        ],
    426
    0
                    });
    427
    428
            // Create bind group
    429
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    430
    0
                label: Some("Vec Add Bind Group"),
    431
    0
                layout: &bind_group_layout,
    432
    0
                entries: &[
    433
    0
                    wgpu::BindGroupEntry {
    434
    0
                        binding: 0,
    435
    0
                        resource: a_buffer.as_entire_binding(),
    436
    0
                    },
    437
    0
                    wgpu::BindGroupEntry {
    438
    0
                        binding: 1,
    439
    0
                        resource: b_buffer.as_entire_binding(),
    440
    0
                    },
    441
    0
                    wgpu::BindGroupEntry {
    442
    0
                        binding: 2,
    443
    0
                        resource: c_buffer.as_entire_binding(),
    444
    0
                    },
    445
    0
                ],
    446
    0
            });
    447
    448
            // Create pipeline
    449
    0
            let pipeline_layout = self
    450
    0
                .device
    451
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    452
    0
                    label: Some("Vec Add Pipeline Layout"),
    453
    0
                    bind_group_layouts: &[&bind_group_layout],
    454
    0
                    push_constant_ranges: &[],
    455
    0
                });
    456
    457
    0
            let pipeline = self
    458
    0
                .device
    459
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    460
    0
                    label: Some("Vec Add Pipeline"),
    461
    0
                    layout: Some(&pipeline_layout),
    462
    0
                    module: &shader,
    463
    0
                    entry_point: Some("main"),
    464
    0
                    compilation_options: Default::default(),
    465
    0
                    cache: None,
    466
    0
                });
    467
    468
            // Create staging buffer for reading results
    469
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    470
    0
                label: Some("Staging Buffer"),
    471
    0
                size: std::mem::size_of_val(result) as u64,
    472
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    473
    0
                mapped_at_creation: false,
    474
    0
            });
    475
    476
            // Create command encoder
    477
    0
            let mut encoder = self
    478
    0
                .device
    479
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    480
    0
                    label: Some("Vec Add Encoder"),
    481
    0
                });
    482
    483
    0
            {
    484
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    485
    0
                    label: Some("Vec Add Pass"),
    486
    0
                    timestamp_writes: None,
    487
    0
                });
    488
    0
                compute_pass.set_pipeline(&pipeline);
    489
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    490
    0
    491
    0
                // Dispatch workgroups (256 threads per workgroup)
    492
    0
                let workgroup_size = 256;
    493
    0
                let num_workgroups = (len as u32).div_ceil(workgroup_size);
    494
    0
    495
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    496
    0
            }
    497
    498
            // Copy result to staging buffer
    499
    0
            encoder.copy_buffer_to_buffer(
    500
    0
                &c_buffer,
    501
                0,
    502
    0
                &staging_buffer,
    503
                0,
    504
    0
                std::mem::size_of_val(result) as u64,
    505
            );
    506
    507
            // Submit commands
    508
    0
            self.queue.submit(Some(encoder.finish()));
    509
    510
            // Read back results
    511
    0
            let buffer_slice = staging_buffer.slice(..);
    512
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    513
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    514
    0
                sender.send(result).ok();
    515
    0
            });
    516
    517
            // Poll device to ensure GPU work completes and callbacks are invoked
    518
    0
            self.device
    519
    0
                .poll(wgpu::PollType::Wait {
    520
    0
                    submission_index: None,
    521
    0
                    timeout: None,
    522
    0
                })
    523
    0
                .ok();
    524
    525
    0
            receiver
    526
    0
                .receive()
    527
    0
                .await
    528
    0
                .ok_or("Failed to receive mapping result")?
    529
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    530
    531
    0
            {
    532
    0
                let data = buffer_slice.get_mapped_range();
    533
    0
                result.copy_from_slice(bytemuck::cast_slice(&data));
    534
    0
            }
    535
    536
    0
            staging_buffer.unmap();
    537
    538
    0
            Ok(())
    539
    0
        }
    540
    541
        /// Generic helper for element-wise GPU operations
    542
        ///
    543
        /// This helper eliminates code duplication between element-wise operations
    544
        /// (relu, clip, sigmoid, tanh, etc.) by abstracting the common GPU compute pattern.
    545
        ///
    546
        /// # Arguments
    547
        ///
    548
        /// * `op_name` - Operation name for labels (e.g., "ReLU", "Clip")
    549
        /// * `shader_source` - WGSL shader source code
    550
        /// * `input` - Input data
    551
        /// * `result` - Output buffer
    552
        /// * `uniform_data` - Optional uniform buffer data (e.g., clip parameters)
    553
    0
        async fn execute_element_wise_op(
    554
    0
            &self,
    555
    0
            op_name: &str,
    556
    0
            shader_source: &str,
    557
    0
            input: &[f32],
    558
    0
            result: &mut [f32],
    559
    0
            uniform_data: Option<&[u8]>,
    560
    0
        ) -> Result<(), String> {
    561
    0
            let len = input.len();
    562
    563
            // Create shader module
    564
    0
            let shader = self
    565
    0
                .device
    566
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    567
    0
                    label: Some(&format!("{} Shader", op_name)),
    568
    0
                    source: wgpu::ShaderSource::Wgsl(shader_source.into()),
    569
    0
                });
    570
    571
            // Create input buffer
    572
    0
            let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    573
    0
                label: Some(&format!("{} Input", op_name)),
    574
    0
                size: std::mem::size_of_val(input) as u64,
    575
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    576
    0
                mapped_at_creation: false,
    577
    0
            });
    578
    579
            // Create output buffer
    580
    0
            let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    581
    0
                label: Some(&format!("{} Output", op_name)),
    582
    0
                size: std::mem::size_of_val(result) as u64,
    583
    0
                usage: wgpu::BufferUsages::STORAGE
    584
    0
                    | wgpu::BufferUsages::COPY_SRC
    585
    0
                    | wgpu::BufferUsages::COPY_DST,
    586
    0
                mapped_at_creation: false,
    587
    0
            });
    588
    589
            // Write input data
    590
    0
            self.queue
    591
    0
                .write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
    592
    593
            // Create optional uniform buffer
    594
    0
            let uniform_buffer = uniform_data.map(|data| {
    595
    0
                let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    596
    0
                    label: Some(&format!("{} Uniform", op_name)),
    597
    0
                    size: data.len() as u64,
    598
    0
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    599
    0
                    mapped_at_creation: false,
    600
    0
                });
    601
    0
                self.queue.write_buffer(&buffer, 0, data);
    602
    0
                buffer
    603
    0
            });
    604
    605
            // Create bind group layout entries (input + output + optional uniform)
    606
    0
            let mut bind_group_entries = vec![
    607
    0
                wgpu::BindGroupLayoutEntry {
    608
    0
                    binding: 0,
    609
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    610
    0
                    ty: wgpu::BindingType::Buffer {
    611
    0
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
    612
    0
                        has_dynamic_offset: false,
    613
    0
                        min_binding_size: None,
    614
    0
                    },
    615
    0
                    count: None,
    616
    0
                },
    617
    0
                wgpu::BindGroupLayoutEntry {
    618
    0
                    binding: 1,
    619
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    620
    0
                    ty: wgpu::BindingType::Buffer {
    621
    0
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
    622
    0
                        has_dynamic_offset: false,
    623
    0
                        min_binding_size: None,
    624
    0
                    },
    625
    0
                    count: None,
    626
    0
                },
    627
            ];
    628
    629
            // Add uniform buffer binding if present
    630
    0
            if uniform_buffer.is_some() {
    631
    0
                bind_group_entries.push(wgpu::BindGroupLayoutEntry {
    632
    0
                    binding: 2,
    633
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    634
    0
                    ty: wgpu::BindingType::Buffer {
    635
    0
                        ty: wgpu::BufferBindingType::Uniform,
    636
    0
                        has_dynamic_offset: false,
    637
    0
                        min_binding_size: None,
    638
    0
                    },
    639
    0
                    count: None,
    640
    0
                });
    641
    0
            }
    642
    643
            // Create bind group layout
    644
    0
            let bind_group_layout =
    645
    0
                self.device
    646
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    647
    0
                        label: Some(&format!("{} Bind Group Layout", op_name)),
    648
    0
                        entries: &bind_group_entries,
    649
    0
                    });
    650
    651
            // Create bind group entries
    652
    0
            let mut bind_entries = vec![
    653
    0
                wgpu::BindGroupEntry {
    654
    0
                    binding: 0,
    655
    0
                    resource: input_buffer.as_entire_binding(),
    656
    0
                },
    657
    0
                wgpu::BindGroupEntry {
    658
    0
                    binding: 1,
    659
    0
                    resource: output_buffer.as_entire_binding(),
    660
    0
                },
    661
            ];
    662
    663
            // Add uniform buffer binding if present
    664
    0
            if let Some(ref uniform_buf) = uniform_buffer {
    665
    0
                bind_entries.push(wgpu::BindGroupEntry {
    666
    0
                    binding: 2,
    667
    0
                    resource: uniform_buf.as_entire_binding(),
    668
    0
                });
    669
    0
            }
    670
    671
            // Create bind group
    672
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    673
    0
                label: Some(&format!("{} Bind Group", op_name)),
    674
    0
                layout: &bind_group_layout,
    675
    0
                entries: &bind_entries,
    676
    0
            });
    677
    678
            // Create pipeline
    679
    0
            let pipeline_layout = self
    680
    0
                .device
    681
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    682
    0
                    label: Some(&format!("{} Pipeline Layout", op_name)),
    683
    0
                    bind_group_layouts: &[&bind_group_layout],
    684
    0
                    push_constant_ranges: &[],
    685
    0
                });
    686
    687
    0
            let pipeline = self
    688
    0
                .device
    689
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    690
    0
                    label: Some(&format!("{} Pipeline", op_name)),
    691
    0
                    layout: Some(&pipeline_layout),
    692
    0
                    module: &shader,
    693
    0
                    entry_point: Some("main"),
    694
    0
                    compilation_options: Default::default(),
    695
    0
                    cache: None,
    696
    0
                });
    697
    698
            // Create staging buffer for reading results
    699
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    700
    0
                label: Some(&format!("{} Staging Buffer", op_name)),
    701
    0
                size: std::mem::size_of_val(result) as u64,
    702
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    703
    0
                mapped_at_creation: false,
    704
    0
            });
    705
    706
            // Create command encoder
    707
    0
            let mut encoder = self
    708
    0
                .device
    709
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    710
    0
                    label: Some(&format!("{} Encoder", op_name)),
    711
    0
                });
    712
    713
    0
            {
    714
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    715
    0
                    label: Some(&format!("{} Pass", op_name)),
    716
    0
                    timestamp_writes: None,
    717
    0
                });
    718
    0
                compute_pass.set_pipeline(&pipeline);
    719
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    720
    0
    721
    0
                // Dispatch workgroups (256 threads per workgroup)
    722
    0
                let workgroup_size = 256;
    723
    0
                let num_workgroups = (len as u32).div_ceil(workgroup_size);
    724
    0
    725
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    726
    0
            }
    727
    728
            // Copy result to staging buffer
    729
    0
            encoder.copy_buffer_to_buffer(
    730
    0
                &output_buffer,
    731
                0,
    732
    0
                &staging_buffer,
    733
                0,
    734
    0
                std::mem::size_of_val(result) as u64,
    735
            );
    736
    737
            // Submit commands
    738
    0
            self.queue.submit(Some(encoder.finish()));
    739
    740
            // Read back results
    741
    0
            let buffer_slice = staging_buffer.slice(..);
    742
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    743
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    744
    0
                sender.send(result).ok();
    745
    0
            });
    746
    747
            // Poll device to ensure GPU work completes and callbacks are invoked
    748
    0
            self.device
    749
    0
                .poll(wgpu::PollType::Wait {
    750
    0
                    submission_index: None,
    751
    0
                    timeout: None,
    752
    0
                })
    753
    0
                .ok();
    754
    755
    0
            receiver
    756
    0
                .receive()
    757
    0
                .await
    758
    0
                .ok_or("Failed to receive mapping result")?
    759
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    760
    761
    0
            {
    762
    0
                let data = buffer_slice.get_mapped_range();
    763
    0
                result.copy_from_slice(bytemuck::cast_slice(&data));
    764
    0
            }
    765
    766
    0
            staging_buffer.unmap();
    767
    768
    0
            Ok(())
    769
    0
        }
    770
    771
        /// Execute ReLU activation on GPU: result[i] = max(0, input[i]) (sync, native only)
    772
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    773
    0
        pub fn relu(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    774
    0
            runtime::block_on(async {
    775
    0
                self.execute_element_wise_op("ReLU", shaders::RELU_SHADER, input, result, None)
    776
    0
                    .await
    777
    0
            })
    778
    0
        }
    779
    780
        /// Execute ReLU activation on GPU (async, works on all platforms)
    781
    0
        pub async fn relu_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    782
    0
            self.execute_element_wise_op("ReLU", shaders::RELU_SHADER, input, result, None)
    783
    0
                .await
    784
    0
        }
    785
    786
        /// Execute leaky ReLU activation on GPU (sync, native only)
    787
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    788
    0
        pub fn leaky_relu(
    789
    0
            &self,
    790
    0
            input: &[f32],
    791
    0
            result: &mut [f32],
    792
    0
            negative_slope: f32,
    793
    0
        ) -> Result<(), String> {
    794
    0
            runtime::block_on(self.leaky_relu_async(input, result, negative_slope))
    795
    0
        }
    796
    797
        /// Execute leaky ReLU activation on GPU (async, works on all platforms)
    798
    0
        pub async fn leaky_relu_async(
    799
    0
            &self,
    800
    0
            input: &[f32],
    801
    0
            result: &mut [f32],
    802
    0
            negative_slope: f32,
    803
    0
        ) -> Result<(), String> {
    804
            #[repr(C)]
    805
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    806
            struct LeakyReluParams {
    807
                negative_slope: f32,
    808
            }
    809
    810
    0
            let params = LeakyReluParams { negative_slope };
    811
    0
            let uniform_data = bytemuck::bytes_of(&params);
    812
    813
    0
            self.execute_element_wise_op(
    814
    0
                "LeakyReLU",
    815
    0
                shaders::LEAKY_RELU_SHADER,
    816
    0
                input,
    817
    0
                result,
    818
    0
                Some(uniform_data),
    819
    0
            )
    820
    0
            .await
    821
    0
        }
    822
    823
        /// Execute ELU activation on GPU (sync, native only)
    824
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    825
    0
        pub fn elu(&self, input: &[f32], result: &mut [f32], alpha: f32) -> Result<(), String> {
    826
    0
            runtime::block_on(self.elu_async(input, result, alpha))
    827
    0
        }
    828
    829
        /// Execute ELU activation on GPU (async, works on all platforms)
    830
    0
        pub async fn elu_async(
    831
    0
            &self,
    832
    0
            input: &[f32],
    833
    0
            result: &mut [f32],
    834
    0
            alpha: f32,
    835
    0
        ) -> Result<(), String> {
    836
            #[repr(C)]
    837
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    838
            struct EluParams {
    839
                alpha: f32,
    840
            }
    841
    842
    0
            let params = EluParams { alpha };
    843
    0
            let uniform_data = bytemuck::bytes_of(&params);
    844
    845
    0
            self.execute_element_wise_op(
    846
    0
                "ELU",
    847
    0
                shaders::ELU_SHADER,
    848
    0
                input,
    849
    0
                result,
    850
    0
                Some(uniform_data),
    851
    0
            )
    852
    0
            .await
    853
    0
        }
    854
    855
        /// Execute sigmoid activation on GPU (sync, native only)
    856
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    857
    0
        pub fn sigmoid(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    858
    0
            runtime::block_on(self.sigmoid_async(input, result))
    859
    0
        }
    860
    861
        /// Execute sigmoid activation on GPU (async, works on all platforms)
    862
    0
        pub async fn sigmoid_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    863
    0
            self.execute_element_wise_op("Sigmoid", shaders::SIGMOID_SHADER, input, result, None)
    864
    0
                .await
    865
    0
        }
    866
    867
        /// Execute tanh activation on GPU (sync, native only)
    868
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    869
    0
        pub fn tanh(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    870
    0
            runtime::block_on(self.tanh_async(input, result))
    871
    0
        }
    872
    873
        /// Execute tanh activation on GPU (async, works on all platforms)
    874
    0
        pub async fn tanh_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    875
    0
            self.execute_element_wise_op("Tanh", shaders::TANH_SHADER, input, result, None)
    876
    0
                .await
    877
    0
        }
    878
    879
        /// Execute swish activation on GPU (sync, native only)
    880
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    881
    0
        pub fn swish(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    882
    0
            runtime::block_on(self.swish_async(input, result))
    883
    0
        }
    884
    885
        /// Execute swish activation on GPU (async, works on all platforms)
    886
    0
        pub async fn swish_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    887
    0
            self.execute_element_wise_op("Swish", shaders::SWISH_SHADER, input, result, None)
    888
    0
                .await
    889
    0
        }
    890
    891
        /// Execute GELU activation on GPU (sync, native only)
    892
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    893
    0
        pub fn gelu(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    894
    0
            runtime::block_on(self.gelu_async(input, result))
    895
    0
        }
    896
    897
        /// Execute GELU activation on GPU (async, works on all platforms)
    898
    0
        pub async fn gelu_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    899
    0
            self.execute_element_wise_op("GELU", shaders::GELU_SHADER, input, result, None)
    900
    0
                .await
    901
    0
        }
    902
    903
        /// Execute clip (clamp) operation on GPU (sync, native only)
    904
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    905
    0
        pub fn clip(
    906
    0
            &self,
    907
    0
            input: &[f32],
    908
    0
            result: &mut [f32],
    909
    0
            min_val: f32,
    910
    0
            max_val: f32,
    911
    0
        ) -> Result<(), String> {
    912
    0
            runtime::block_on(self.clip_async(input, result, min_val, max_val))
    913
    0
        }
    914
    915
        /// Execute clip (clamp) operation on GPU (async, works on all platforms)
    916
    0
        pub async fn clip_async(
    917
    0
            &self,
    918
    0
            input: &[f32],
    919
    0
            result: &mut [f32],
    920
    0
            min_val: f32,
    921
    0
            max_val: f32,
    922
    0
        ) -> Result<(), String> {
    923
            #[repr(C)]
    924
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    925
            struct ClipParams {
    926
                min_val: f32,
    927
                max_val: f32,
    928
            }
    929
    930
    0
            let params = ClipParams { min_val, max_val };
    931
    0
            let uniform_data = bytemuck::bytes_of(&params);
    932
    933
    0
            self.execute_element_wise_op(
    934
    0
                "Clip",
    935
    0
                shaders::CLIP_SHADER,
    936
    0
                input,
    937
    0
                result,
    938
    0
                Some(uniform_data),
    939
    0
            )
    940
    0
            .await
    941
    0
        }
    942
    943
        /// Execute softmax on GPU (sync, native only)
    944
        ///
    945
        /// Multi-pass implementation:
    946
        /// 1. Find max value (parallel reduction)
    947
        /// 2. Compute exp(x - max) (element-wise)
    948
        /// 3. Sum exp values (parallel reduction)
    949
        /// 4. Normalize by sum (element-wise)
    950
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    951
    0
        pub fn softmax(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    952
    0
            runtime::block_on(async { self.softmax_async(input, result).await })
    953
    0
        }
    954
    955
        /// Execute softmax on GPU (async, works on all platforms)
    956
    0
        pub async fn softmax_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    957
            // Pass 1: Find max value
    958
    0
            let max_val = self.reduce_max(input).await?;
    959
    960
            // Pass 2: Compute exp(x - max)
    961
    0
            let exp_vals = self.compute_exp_subtract(input, max_val).await?;
    962
    963
            // Pass 3: Sum exp values
    964
    0
            let sum_exp = self.reduce_sum(&exp_vals).await?;
    965
    966
            // Pass 4: Normalize by sum
    967
    0
            self.normalize_by_sum(&exp_vals, result, sum_exp).await?;
    968
    969
    0
            Ok(())
    970
    0
        }
    971
    972
        /// Execute log_softmax on GPU (sync, native only)
    973
        ///
    974
        /// Multi-pass implementation:
    975
        /// 1. Find max value (parallel reduction)
    976
        /// 2. Compute exp(x - max) (element-wise)
    977
        /// 3. Sum exp values (parallel reduction)
    978
        /// 4. Compute log_softmax (element-wise)
    979
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    980
    0
        pub fn log_softmax(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    981
    0
            runtime::block_on(async { self.log_softmax_async(input, result).await })
    982
    0
        }
    983
    984
        /// Execute log_softmax on GPU (async, works on all platforms)
    985
    0
        pub async fn log_softmax_async(&self, input: &[f32], result: &mut [f32]) -> Result<(), String> {
    986
            // Pass 1: Find max value
    987
    0
            let max_val = self.reduce_max(input).await?;
    988
    989
            // Pass 2: Compute exp(x - max)
    990
    0
            let exp_vals = self.compute_exp_subtract(input, max_val).await?;
    991
    992
            // Pass 3: Sum exp values
    993
    0
            let sum_exp = self.reduce_sum(&exp_vals).await?;
    994
    995
            // Pass 4: Compute log_softmax = x - max - log(sum_exp)
    996
    0
            let log_sum_exp = sum_exp.ln();
    997
    998
            #[repr(C)]
    999
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    1000
            struct LogSoftmaxParams {
    1001
                max_val: f32,
    1002
                log_sum_exp: f32,
    1003
            }
    1004
    1005
    0
            let params = LogSoftmaxParams {
    1006
    0
                max_val,
    1007
    0
                log_sum_exp,
    1008
    0
            };
    1009
    0
            let uniform_data = bytemuck::bytes_of(&params);
    1010
    1011
    0
            self.execute_element_wise_op(
    1012
    0
                "LogSoftmax",
    1013
    0
                shaders::LOG_SOFTMAX_SHADER,
    1014
    0
                input,
    1015
    0
                result,
    1016
    0
                Some(uniform_data),
    1017
    0
            )
    1018
    0
            .await?;
    1019
    1020
    0
            Ok(())
    1021
    0
        }
    1022
    1023
        /// Helper: Parallel max reduction
    1024
    0
        async fn reduce_max(&self, input: &[f32]) -> Result<f32, String> {
    1025
    0
            let len = input.len();
    1026
    0
            let workgroup_size = 256;
    1027
    0
            let num_workgroups = (len as u32).div_ceil(workgroup_size);
    1028
    1029
            // Create shader module
    1030
    0
            let shader = self
    1031
    0
                .device
    1032
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    1033
    0
                    label: Some("Max Reduction Shader"),
    1034
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::MAX_REDUCTION_SHADER.into()),
    1035
    0
                });
    1036
    1037
            // Create input buffer
    1038
    0
            let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1039
    0
                label: Some("Max Reduction Input"),
    1040
    0
                size: std::mem::size_of_val(input) as u64,
    1041
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1042
    0
                mapped_at_creation: false,
    1043
    0
            });
    1044
    1045
            // Result buffer for partial maxes
    1046
    0
            let partial_results = vec![f32::NEG_INFINITY; num_workgroups as usize];
    1047
    0
            let result_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1048
    0
                label: Some("Max Partial Results"),
    1049
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1050
    0
                usage: wgpu::BufferUsages::STORAGE
    1051
    0
                    | wgpu::BufferUsages::COPY_SRC
    1052
    0
                    | wgpu::BufferUsages::COPY_DST,
    1053
    0
                mapped_at_creation: false,
    1054
    0
            });
    1055
    1056
    0
            self.queue
    1057
    0
                .write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
    1058
    1059
            // Create bind group layout
    1060
    0
            let bind_group_layout =
    1061
    0
                self.device
    1062
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    1063
    0
                        label: Some("Max Reduction Bind Group Layout"),
    1064
    0
                        entries: &[
    1065
    0
                            wgpu::BindGroupLayoutEntry {
    1066
    0
                                binding: 0,
    1067
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1068
    0
                                ty: wgpu::BindingType::Buffer {
    1069
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1070
    0
                                    has_dynamic_offset: false,
    1071
    0
                                    min_binding_size: None,
    1072
    0
                                },
    1073
    0
                                count: None,
    1074
    0
                            },
    1075
    0
                            wgpu::BindGroupLayoutEntry {
    1076
    0
                                binding: 1,
    1077
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1078
    0
                                ty: wgpu::BindingType::Buffer {
    1079
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    1080
    0
                                    has_dynamic_offset: false,
    1081
    0
                                    min_binding_size: None,
    1082
    0
                                },
    1083
    0
                                count: None,
    1084
    0
                            },
    1085
    0
                        ],
    1086
    0
                    });
    1087
    1088
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    1089
    0
                label: Some("Max Reduction Bind Group"),
    1090
    0
                layout: &bind_group_layout,
    1091
    0
                entries: &[
    1092
    0
                    wgpu::BindGroupEntry {
    1093
    0
                        binding: 0,
    1094
    0
                        resource: input_buffer.as_entire_binding(),
    1095
    0
                    },
    1096
    0
                    wgpu::BindGroupEntry {
    1097
    0
                        binding: 1,
    1098
    0
                        resource: result_buffer.as_entire_binding(),
    1099
    0
                    },
    1100
    0
                ],
    1101
    0
            });
    1102
    1103
    0
            let pipeline_layout = self
    1104
    0
                .device
    1105
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    1106
    0
                    label: Some("Max Reduction Pipeline Layout"),
    1107
    0
                    bind_group_layouts: &[&bind_group_layout],
    1108
    0
                    push_constant_ranges: &[],
    1109
    0
                });
    1110
    1111
    0
            let pipeline = self
    1112
    0
                .device
    1113
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    1114
    0
                    label: Some("Max Reduction Pipeline"),
    1115
    0
                    layout: Some(&pipeline_layout),
    1116
    0
                    module: &shader,
    1117
    0
                    entry_point: Some("main"),
    1118
    0
                    compilation_options: Default::default(),
    1119
    0
                    cache: None,
    1120
    0
                });
    1121
    1122
    0
            let mut encoder = self
    1123
    0
                .device
    1124
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    1125
    0
                    label: Some("Max Reduction Encoder"),
    1126
    0
                });
    1127
    1128
    0
            {
    1129
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    1130
    0
                    label: Some("Max Reduction Pass"),
    1131
    0
                    timestamp_writes: None,
    1132
    0
                });
    1133
    0
    1134
    0
                compute_pass.set_pipeline(&pipeline);
    1135
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    1136
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    1137
    0
            }
    1138
    1139
            // Create staging buffer
    1140
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1141
    0
                label: Some("Max Staging Buffer"),
    1142
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1143
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    1144
    0
                mapped_at_creation: false,
    1145
    0
            });
    1146
    1147
    0
            encoder.copy_buffer_to_buffer(
    1148
    0
                &result_buffer,
    1149
                0,
    1150
    0
                &staging_buffer,
    1151
                0,
    1152
    0
                std::mem::size_of_val(partial_results.as_slice()) as u64,
    1153
            );
    1154
    1155
    0
            self.queue.submit(Some(encoder.finish()));
    1156
    1157
    0
            let buffer_slice = staging_buffer.slice(..);
    1158
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    1159
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    1160
    0
                sender.send(result).ok();
    1161
    0
            });
    1162
    1163
            // Poll device to ensure GPU work completes and callbacks are invoked
    1164
    0
            self.device
    1165
    0
                .poll(wgpu::PollType::Wait {
    1166
    0
                    submission_index: None,
    1167
    0
                    timeout: None,
    1168
    0
                })
    1169
    0
                .ok();
    1170
    0
            receiver
    1171
    0
                .receive()
    1172
    0
                .await
    1173
    0
                .ok_or("Channel receive failed")?
    1174
    0
                .map_err(|e| format!("Buffer map failed: {:?}", e))?;
    1175
    1176
    0
            let data = buffer_slice.get_mapped_range();
    1177
    0
            let result: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
    1178
    0
            drop(data);
    1179
    0
            staging_buffer.unmap();
    1180
    1181
            // Final reduction on CPU
    1182
    0
            Ok(result.iter().copied().fold(f32::NEG_INFINITY, f32::max))
    1183
    0
        }
    1184
    1185
        /// Helper: Parallel sum reduction
    1186
    0
        async fn reduce_sum(&self, input: &[f32]) -> Result<f32, String> {
    1187
    0
            let len = input.len();
    1188
    0
            let workgroup_size = 256;
    1189
    0
            let num_workgroups = (len as u32).div_ceil(workgroup_size);
    1190
    1191
    0
            let shader = self
    1192
    0
                .device
    1193
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    1194
    0
                    label: Some("Sum Reduction Shader"),
    1195
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::SUM_REDUCTION_SHADER.into()),
    1196
    0
                });
    1197
    1198
    0
            let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1199
    0
                label: Some("Sum Reduction Input"),
    1200
    0
                size: std::mem::size_of_val(input) as u64,
    1201
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1202
    0
                mapped_at_creation: false,
    1203
    0
            });
    1204
    1205
    0
            let partial_results = vec![0.0f32; num_workgroups as usize];
    1206
    0
            let result_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1207
    0
                label: Some("Sum Partial Results"),
    1208
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1209
    0
                usage: wgpu::BufferUsages::STORAGE
    1210
    0
                    | wgpu::BufferUsages::COPY_SRC
    1211
    0
                    | wgpu::BufferUsages::COPY_DST,
    1212
    0
                mapped_at_creation: false,
    1213
    0
            });
    1214
    1215
    0
            self.queue
    1216
    0
                .write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
    1217
    1218
    0
            let bind_group_layout =
    1219
    0
                self.device
    1220
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    1221
    0
                        label: Some("Sum Reduction Bind Group Layout"),
    1222
    0
                        entries: &[
    1223
    0
                            wgpu::BindGroupLayoutEntry {
    1224
    0
                                binding: 0,
    1225
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1226
    0
                                ty: wgpu::BindingType::Buffer {
    1227
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1228
    0
                                    has_dynamic_offset: false,
    1229
    0
                                    min_binding_size: None,
    1230
    0
                                },
    1231
    0
                                count: None,
    1232
    0
                            },
    1233
    0
                            wgpu::BindGroupLayoutEntry {
    1234
    0
                                binding: 1,
    1235
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1236
    0
                                ty: wgpu::BindingType::Buffer {
    1237
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    1238
    0
                                    has_dynamic_offset: false,
    1239
    0
                                    min_binding_size: None,
    1240
    0
                                },
    1241
    0
                                count: None,
    1242
    0
                            },
    1243
    0
                        ],
    1244
    0
                    });
    1245
    1246
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    1247
    0
                label: Some("Sum Reduction Bind Group"),
    1248
    0
                layout: &bind_group_layout,
    1249
    0
                entries: &[
    1250
    0
                    wgpu::BindGroupEntry {
    1251
    0
                        binding: 0,
    1252
    0
                        resource: input_buffer.as_entire_binding(),
    1253
    0
                    },
    1254
    0
                    wgpu::BindGroupEntry {
    1255
    0
                        binding: 1,
    1256
    0
                        resource: result_buffer.as_entire_binding(),
    1257
    0
                    },
    1258
    0
                ],
    1259
    0
            });
    1260
    1261
    0
            let pipeline_layout = self
    1262
    0
                .device
    1263
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    1264
    0
                    label: Some("Sum Reduction Pipeline Layout"),
    1265
    0
                    bind_group_layouts: &[&bind_group_layout],
    1266
    0
                    push_constant_ranges: &[],
    1267
    0
                });
    1268
    1269
    0
            let pipeline = self
    1270
    0
                .device
    1271
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    1272
    0
                    label: Some("Sum Reduction Pipeline"),
    1273
    0
                    layout: Some(&pipeline_layout),
    1274
    0
                    module: &shader,
    1275
    0
                    entry_point: Some("main"),
    1276
    0
                    compilation_options: Default::default(),
    1277
    0
                    cache: None,
    1278
    0
                });
    1279
    1280
    0
            let mut encoder = self
    1281
    0
                .device
    1282
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    1283
    0
                    label: Some("Sum Reduction Encoder"),
    1284
    0
                });
    1285
    1286
    0
            {
    1287
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    1288
    0
                    label: Some("Sum Reduction Pass"),
    1289
    0
                    timestamp_writes: None,
    1290
    0
                });
    1291
    0
    1292
    0
                compute_pass.set_pipeline(&pipeline);
    1293
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    1294
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    1295
    0
            }
    1296
    1297
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1298
    0
                label: Some("Sum Staging Buffer"),
    1299
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1300
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    1301
    0
                mapped_at_creation: false,
    1302
    0
            });
    1303
    1304
    0
            encoder.copy_buffer_to_buffer(
    1305
    0
                &result_buffer,
    1306
                0,
    1307
    0
                &staging_buffer,
    1308
                0,
    1309
    0
                std::mem::size_of_val(partial_results.as_slice()) as u64,
    1310
            );
    1311
    1312
    0
            self.queue.submit(Some(encoder.finish()));
    1313
    1314
    0
            let buffer_slice = staging_buffer.slice(..);
    1315
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    1316
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    1317
    0
                sender.send(result).ok();
    1318
    0
            });
    1319
    1320
            // Poll device to ensure GPU work completes and callbacks are invoked
    1321
    0
            self.device
    1322
    0
                .poll(wgpu::PollType::Wait {
    1323
    0
                    submission_index: None,
    1324
    0
                    timeout: None,
    1325
    0
                })
    1326
    0
                .ok();
    1327
    0
            receiver
    1328
    0
                .receive()
    1329
    0
                .await
    1330
    0
                .ok_or("Channel receive failed")?
    1331
    0
                .map_err(|e| format!("Buffer map failed: {:?}", e))?;
    1332
    1333
    0
            let data = buffer_slice.get_mapped_range();
    1334
    0
            let result: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
    1335
    0
            drop(data);
    1336
    0
            staging_buffer.unmap();
    1337
    1338
            // Final reduction on CPU
    1339
    0
            Ok(result.iter().sum())
    1340
    0
        }
    1341
    1342
        /// Helper: Compute exp(input[i] - max_val)
    1343
    0
        async fn compute_exp_subtract(&self, input: &[f32], max_val: f32) -> Result<Vec<f32>, String> {
    1344
            #[repr(C)]
    1345
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    1346
            struct MaxValue {
    1347
                max_val: f32,
    1348
            }
    1349
    1350
    0
            let params = MaxValue { max_val };
    1351
    0
            let uniform_data = bytemuck::bytes_of(&params);
    1352
    1353
    0
            let mut result = vec![0.0f32; input.len()];
    1354
    0
            self.execute_element_wise_op(
    1355
    0
                "SoftmaxExp",
    1356
    0
                shaders::SOFTMAX_EXP_SHADER,
    1357
    0
                input,
    1358
    0
                &mut result,
    1359
    0
                Some(uniform_data),
    1360
    0
            )
    1361
    0
            .await?;
    1362
    1363
    0
            Ok(result)
    1364
    0
        }
    1365
    1366
        /// Helper: Normalize by sum
    1367
    0
        async fn normalize_by_sum(
    1368
    0
            &self,
    1369
    0
            input: &[f32],
    1370
    0
            result: &mut [f32],
    1371
    0
            sum_val: f32,
    1372
    0
        ) -> Result<(), String> {
    1373
            #[repr(C)]
    1374
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    1375
            struct SumValue {
    1376
                sum_val: f32,
    1377
            }
    1378
    1379
    0
            let params = SumValue { sum_val };
    1380
    0
            let uniform_data = bytemuck::bytes_of(&params);
    1381
    1382
    0
            self.execute_element_wise_op(
    1383
    0
                "SoftmaxNormalize",
    1384
    0
                shaders::SOFTMAX_NORMALIZE_SHADER,
    1385
    0
                input,
    1386
    0
                result,
    1387
    0
                Some(uniform_data),
    1388
    0
            )
    1389
    0
            .await?;
    1390
    1391
    0
            Ok(())
    1392
    0
        }
    1393
    1394
        /// Execute dot product on GPU (sync, native only)
    1395
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    1396
    0
        pub fn dot(&self, a: &[f32], b: &[f32]) -> Result<f32, String> {
    1397
    0
            runtime::block_on(async { self.dot_async(a, b).await })
    1398
    0
        }
    1399
    1400
        /// Execute dot product on GPU (async, works on all platforms)
    1401
    0
        pub async fn dot_async(&self, a: &[f32], b: &[f32]) -> Result<f32, String> {
    1402
    0
            let len = a.len();
    1403
    0
            let workgroup_size = 256;
    1404
    0
            let num_workgroups = (len as u32).div_ceil(workgroup_size);
    1405
    1406
            // Create shader module
    1407
    0
            let shader = self
    1408
    0
                .device
    1409
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    1410
    0
                    label: Some("Dot Product Shader"),
    1411
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::DOT_PRODUCT_SHADER.into()),
    1412
    0
                });
    1413
    1414
            // Create buffers
    1415
    0
            let a_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1416
    0
                label: Some("Vector A"),
    1417
    0
                size: std::mem::size_of_val(a) as u64,
    1418
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1419
    0
                mapped_at_creation: false,
    1420
    0
            });
    1421
    1422
    0
            let b_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1423
    0
                label: Some("Vector B"),
    1424
    0
                size: std::mem::size_of_val(b) as u64,
    1425
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1426
    0
                mapped_at_creation: false,
    1427
    0
            });
    1428
    1429
            // Result buffer for partial sums (one per workgroup)
    1430
    0
            let partial_results = vec![0.0f32; num_workgroups as usize];
    1431
    0
            let result_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1432
    0
                label: Some("Partial Results"),
    1433
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1434
    0
                usage: wgpu::BufferUsages::STORAGE
    1435
    0
                    | wgpu::BufferUsages::COPY_SRC
    1436
    0
                    | wgpu::BufferUsages::COPY_DST,
    1437
    0
                mapped_at_creation: false,
    1438
    0
            });
    1439
    1440
            // Write data to buffers
    1441
    0
            self.queue
    1442
    0
                .write_buffer(&a_buffer, 0, bytemuck::cast_slice(a));
    1443
    0
            self.queue
    1444
    0
                .write_buffer(&b_buffer, 0, bytemuck::cast_slice(b));
    1445
    1446
            // Create bind group layout
    1447
    0
            let bind_group_layout =
    1448
    0
                self.device
    1449
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    1450
    0
                        label: Some("Dot Product Bind Group Layout"),
    1451
    0
                        entries: &[
    1452
    0
                            wgpu::BindGroupLayoutEntry {
    1453
    0
                                binding: 0,
    1454
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1455
    0
                                ty: wgpu::BindingType::Buffer {
    1456
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1457
    0
                                    has_dynamic_offset: false,
    1458
    0
                                    min_binding_size: None,
    1459
    0
                                },
    1460
    0
                                count: None,
    1461
    0
                            },
    1462
    0
                            wgpu::BindGroupLayoutEntry {
    1463
    0
                                binding: 1,
    1464
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1465
    0
                                ty: wgpu::BindingType::Buffer {
    1466
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1467
    0
                                    has_dynamic_offset: false,
    1468
    0
                                    min_binding_size: None,
    1469
    0
                                },
    1470
    0
                                count: None,
    1471
    0
                            },
    1472
    0
                            wgpu::BindGroupLayoutEntry {
    1473
    0
                                binding: 2,
    1474
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1475
    0
                                ty: wgpu::BindingType::Buffer {
    1476
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    1477
    0
                                    has_dynamic_offset: false,
    1478
    0
                                    min_binding_size: None,
    1479
    0
                                },
    1480
    0
                                count: None,
    1481
    0
                            },
    1482
    0
                        ],
    1483
    0
                    });
    1484
    1485
            // Create bind group
    1486
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    1487
    0
                label: Some("Dot Product Bind Group"),
    1488
    0
                layout: &bind_group_layout,
    1489
    0
                entries: &[
    1490
    0
                    wgpu::BindGroupEntry {
    1491
    0
                        binding: 0,
    1492
    0
                        resource: a_buffer.as_entire_binding(),
    1493
    0
                    },
    1494
    0
                    wgpu::BindGroupEntry {
    1495
    0
                        binding: 1,
    1496
    0
                        resource: b_buffer.as_entire_binding(),
    1497
    0
                    },
    1498
    0
                    wgpu::BindGroupEntry {
    1499
    0
                        binding: 2,
    1500
    0
                        resource: result_buffer.as_entire_binding(),
    1501
    0
                    },
    1502
    0
                ],
    1503
    0
            });
    1504
    1505
            // Create pipeline
    1506
    0
            let pipeline_layout = self
    1507
    0
                .device
    1508
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    1509
    0
                    label: Some("Dot Product Pipeline Layout"),
    1510
    0
                    bind_group_layouts: &[&bind_group_layout],
    1511
    0
                    push_constant_ranges: &[],
    1512
    0
                });
    1513
    1514
    0
            let pipeline = self
    1515
    0
                .device
    1516
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    1517
    0
                    label: Some("Dot Product Pipeline"),
    1518
    0
                    layout: Some(&pipeline_layout),
    1519
    0
                    module: &shader,
    1520
    0
                    entry_point: Some("main"),
    1521
    0
                    compilation_options: Default::default(),
    1522
    0
                    cache: None,
    1523
    0
                });
    1524
    1525
            // Create staging buffer for reading results
    1526
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1527
    0
                label: Some("Staging Buffer"),
    1528
    0
                size: std::mem::size_of_val(partial_results.as_slice()) as u64,
    1529
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    1530
    0
                mapped_at_creation: false,
    1531
    0
            });
    1532
    1533
            // Create command encoder
    1534
    0
            let mut encoder = self
    1535
    0
                .device
    1536
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    1537
    0
                    label: Some("Dot Product Encoder"),
    1538
    0
                });
    1539
    1540
    0
            {
    1541
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    1542
    0
                    label: Some("Dot Product Pass"),
    1543
    0
                    timestamp_writes: None,
    1544
    0
                });
    1545
    0
                compute_pass.set_pipeline(&pipeline);
    1546
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    1547
    0
    1548
    0
                // Dispatch workgroups
    1549
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    1550
    0
            }
    1551
    1552
            // Copy result to staging buffer
    1553
    0
            encoder.copy_buffer_to_buffer(
    1554
    0
                &result_buffer,
    1555
                0,
    1556
    0
                &staging_buffer,
    1557
                0,
    1558
    0
                std::mem::size_of_val(partial_results.as_slice()) as u64,
    1559
            );
    1560
    1561
            // Submit commands
    1562
    0
            self.queue.submit(Some(encoder.finish()));
    1563
    1564
            // Read back results
    1565
    0
            let buffer_slice = staging_buffer.slice(..);
    1566
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    1567
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    1568
    0
                sender.send(result).ok();
    1569
    0
            });
    1570
    1571
            // Poll device to ensure GPU work completes and callbacks are invoked
    1572
    0
            self.device
    1573
    0
                .poll(wgpu::PollType::Wait {
    1574
    0
                    submission_index: None,
    1575
    0
                    timeout: None,
    1576
    0
                })
    1577
    0
                .ok();
    1578
    1579
    0
            receiver
    1580
    0
                .receive()
    1581
    0
                .await
    1582
    0
                .ok_or("Failed to receive mapping result")?
    1583
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    1584
    1585
    0
            let final_result = {
    1586
    0
                let data = buffer_slice.get_mapped_range();
    1587
    0
                let partial_sums: &[f32] = bytemuck::cast_slice(&data);
    1588
    1589
                // Sum the partial results from each workgroup on CPU
    1590
    0
                partial_sums.iter().sum()
    1591
            };
    1592
    1593
    0
            staging_buffer.unmap();
    1594
    1595
    0
            Ok(final_result)
    1596
    0
        }
    1597
    1598
        /// Perform 2D convolution on GPU (sync, native only)
    1599
        ///
    1600
        /// # Arguments
    1601
        ///
    1602
        /// * `input` - Input image (row-major)
    1603
        /// * `kernel` - Convolution kernel (row-major)
    1604
        /// * `result` - Output buffer (row-major)
    1605
        /// * `input_rows` - Number of rows in input
    1606
        /// * `input_cols` - Number of columns in input
    1607
        /// * `kernel_rows` - Number of rows in kernel
    1608
        /// * `kernel_cols` - Number of columns in kernel
    1609
        ///
    1610
        /// Output dimensions: (input_rows - kernel_rows + 1) × (input_cols - kernel_cols + 1)
    1611
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    1612
        #[allow(clippy::too_many_arguments)]
    1613
    0
        pub fn convolve2d(
    1614
    0
            &self,
    1615
    0
            input: &[f32],
    1616
    0
            kernel: &[f32],
    1617
    0
            result: &mut [f32],
    1618
    0
            input_rows: usize,
    1619
    0
            input_cols: usize,
    1620
    0
            kernel_rows: usize,
    1621
    0
            kernel_cols: usize,
    1622
    0
        ) -> Result<(), String> {
    1623
    0
            runtime::block_on(async {
    1624
    0
                self.convolve2d_async(
    1625
    0
                    input,
    1626
    0
                    kernel,
    1627
    0
                    result,
    1628
    0
                    input_rows,
    1629
    0
                    input_cols,
    1630
    0
                    kernel_rows,
    1631
    0
                    kernel_cols,
    1632
    0
                )
    1633
    0
                .await
    1634
    0
            })
    1635
    0
        }
    1636
    1637
        /// Perform 2D convolution on GPU (async, works on all platforms)
    1638
        #[allow(clippy::too_many_arguments)]
    1639
    0
        pub async fn convolve2d_async(
    1640
    0
            &self,
    1641
    0
            input: &[f32],
    1642
    0
            kernel: &[f32],
    1643
    0
            result: &mut [f32],
    1644
    0
            input_rows: usize,
    1645
    0
            input_cols: usize,
    1646
    0
            kernel_rows: usize,
    1647
    0
            kernel_cols: usize,
    1648
    0
        ) -> Result<(), String> {
    1649
    0
            let output_rows = input_rows - kernel_rows + 1;
    1650
    0
            let output_cols = input_cols - kernel_cols + 1;
    1651
    1652
            // Create shader module
    1653
    0
            let shader = self
    1654
    0
                .device
    1655
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    1656
    0
                    label: Some("Convolve2D Shader"),
    1657
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::CONVOLVE2D_SHADER.into()),
    1658
    0
                });
    1659
    1660
            // Create buffers
    1661
    0
            let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1662
    0
                label: Some("Input Image"),
    1663
    0
                size: std::mem::size_of_val(input) as u64,
    1664
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1665
    0
                mapped_at_creation: false,
    1666
    0
            });
    1667
    1668
    0
            let kernel_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1669
    0
                label: Some("Kernel"),
    1670
    0
                size: std::mem::size_of_val(kernel) as u64,
    1671
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    1672
    0
                mapped_at_creation: false,
    1673
    0
            });
    1674
    1675
    0
            let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1676
    0
                label: Some("Output"),
    1677
    0
                size: std::mem::size_of_val(result) as u64,
    1678
    0
                usage: wgpu::BufferUsages::STORAGE
    1679
    0
                    | wgpu::BufferUsages::COPY_SRC
    1680
    0
                    | wgpu::BufferUsages::COPY_DST,
    1681
    0
                mapped_at_creation: false,
    1682
    0
            });
    1683
    1684
            // Dimensions uniform buffer
    1685
            #[repr(C)]
    1686
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    1687
            struct ConvDimensions {
    1688
                input_rows: u32,
    1689
                input_cols: u32,
    1690
                kernel_rows: u32,
    1691
                kernel_cols: u32,
    1692
                output_rows: u32,
    1693
                output_cols: u32,
    1694
            }
    1695
    1696
    0
            let dims = ConvDimensions {
    1697
    0
                input_rows: input_rows as u32,
    1698
    0
                input_cols: input_cols as u32,
    1699
    0
                kernel_rows: kernel_rows as u32,
    1700
    0
                kernel_cols: kernel_cols as u32,
    1701
    0
                output_rows: output_rows as u32,
    1702
    0
                output_cols: output_cols as u32,
    1703
    0
            };
    1704
    1705
    0
            let dims_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1706
    0
                label: Some("Conv Dimensions"),
    1707
    0
                size: std::mem::size_of::<ConvDimensions>() as u64,
    1708
    0
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    1709
    0
                mapped_at_creation: false,
    1710
    0
            });
    1711
    1712
            // Write data to buffers
    1713
    0
            self.queue
    1714
    0
                .write_buffer(&input_buffer, 0, bytemuck::cast_slice(input));
    1715
    0
            self.queue
    1716
    0
                .write_buffer(&kernel_buffer, 0, bytemuck::cast_slice(kernel));
    1717
    0
            self.queue
    1718
    0
                .write_buffer(&dims_buffer, 0, bytemuck::bytes_of(&dims));
    1719
    1720
            // Create bind group layout
    1721
    0
            let bind_group_layout =
    1722
    0
                self.device
    1723
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    1724
    0
                        label: Some("Convolve2D Bind Group Layout"),
    1725
    0
                        entries: &[
    1726
    0
                            wgpu::BindGroupLayoutEntry {
    1727
    0
                                binding: 0,
    1728
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1729
    0
                                ty: wgpu::BindingType::Buffer {
    1730
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1731
    0
                                    has_dynamic_offset: false,
    1732
    0
                                    min_binding_size: None,
    1733
    0
                                },
    1734
    0
                                count: None,
    1735
    0
                            },
    1736
    0
                            wgpu::BindGroupLayoutEntry {
    1737
    0
                                binding: 1,
    1738
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1739
    0
                                ty: wgpu::BindingType::Buffer {
    1740
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    1741
    0
                                    has_dynamic_offset: false,
    1742
    0
                                    min_binding_size: None,
    1743
    0
                                },
    1744
    0
                                count: None,
    1745
    0
                            },
    1746
    0
                            wgpu::BindGroupLayoutEntry {
    1747
    0
                                binding: 2,
    1748
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1749
    0
                                ty: wgpu::BindingType::Buffer {
    1750
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    1751
    0
                                    has_dynamic_offset: false,
    1752
    0
                                    min_binding_size: None,
    1753
    0
                                },
    1754
    0
                                count: None,
    1755
    0
                            },
    1756
    0
                            wgpu::BindGroupLayoutEntry {
    1757
    0
                                binding: 3,
    1758
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    1759
    0
                                ty: wgpu::BindingType::Buffer {
    1760
    0
                                    ty: wgpu::BufferBindingType::Uniform,
    1761
    0
                                    has_dynamic_offset: false,
    1762
    0
                                    min_binding_size: None,
    1763
    0
                                },
    1764
    0
                                count: None,
    1765
    0
                            },
    1766
    0
                        ],
    1767
    0
                    });
    1768
    1769
            // Create bind group
    1770
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    1771
    0
                label: Some("Convolve2D Bind Group"),
    1772
    0
                layout: &bind_group_layout,
    1773
    0
                entries: &[
    1774
    0
                    wgpu::BindGroupEntry {
    1775
    0
                        binding: 0,
    1776
    0
                        resource: input_buffer.as_entire_binding(),
    1777
    0
                    },
    1778
    0
                    wgpu::BindGroupEntry {
    1779
    0
                        binding: 1,
    1780
    0
                        resource: kernel_buffer.as_entire_binding(),
    1781
    0
                    },
    1782
    0
                    wgpu::BindGroupEntry {
    1783
    0
                        binding: 2,
    1784
    0
                        resource: output_buffer.as_entire_binding(),
    1785
    0
                    },
    1786
    0
                    wgpu::BindGroupEntry {
    1787
    0
                        binding: 3,
    1788
    0
                        resource: dims_buffer.as_entire_binding(),
    1789
    0
                    },
    1790
    0
                ],
    1791
    0
            });
    1792
    1793
            // Create pipeline layout
    1794
    0
            let pipeline_layout = self
    1795
    0
                .device
    1796
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    1797
    0
                    label: Some("Convolve2D Pipeline Layout"),
    1798
    0
                    bind_group_layouts: &[&bind_group_layout],
    1799
    0
                    push_constant_ranges: &[],
    1800
    0
                });
    1801
    1802
            // Create compute pipeline
    1803
    0
            let pipeline = self
    1804
    0
                .device
    1805
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    1806
    0
                    label: Some("Convolve2D Pipeline"),
    1807
    0
                    layout: Some(&pipeline_layout),
    1808
    0
                    module: &shader,
    1809
    0
                    entry_point: Some("main"),
    1810
    0
                    compilation_options: Default::default(),
    1811
    0
                    cache: None,
    1812
    0
                });
    1813
    1814
            // Create command encoder
    1815
    0
            let mut encoder = self
    1816
    0
                .device
    1817
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    1818
    0
                    label: Some("Convolve2D Encoder"),
    1819
    0
                });
    1820
    1821
            // Compute pass
    1822
    0
            {
    1823
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    1824
    0
                    label: Some("Convolve2D Pass"),
    1825
    0
                    timestamp_writes: None,
    1826
    0
                });
    1827
    0
    1828
    0
                compute_pass.set_pipeline(&pipeline);
    1829
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    1830
    0
    1831
    0
                // Dispatch workgroups: 16×16 threads per workgroup
    1832
    0
                let workgroup_size_x = 16;
    1833
    0
                let workgroup_size_y = 16;
    1834
    0
                let num_workgroups_x = (output_rows as u32).div_ceil(workgroup_size_x);
    1835
    0
                let num_workgroups_y = (output_cols as u32).div_ceil(workgroup_size_y);
    1836
    0
                compute_pass.dispatch_workgroups(num_workgroups_x, num_workgroups_y, 1);
    1837
    0
            }
    1838
    1839
            // Create staging buffer for result readback
    1840
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1841
    0
                label: Some("Staging Buffer"),
    1842
    0
                size: std::mem::size_of_val(result) as u64,
    1843
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    1844
    0
                mapped_at_creation: false,
    1845
    0
            });
    1846
    1847
            // Copy output to staging buffer
    1848
    0
            encoder.copy_buffer_to_buffer(
    1849
    0
                &output_buffer,
    1850
                0,
    1851
    0
                &staging_buffer,
    1852
                0,
    1853
    0
                std::mem::size_of_val(result) as u64,
    1854
            );
    1855
    1856
            // Submit commands
    1857
    0
            self.queue.submit(Some(encoder.finish()));
    1858
    1859
            // Read result back
    1860
    0
            let buffer_slice = staging_buffer.slice(..);
    1861
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    1862
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    1863
    0
                sender.send(result).expect("oneshot channel receiver dropped");
    1864
    0
            });
    1865
    1866
            // Poll device to ensure GPU work completes and callbacks are invoked
    1867
    0
            self.device
    1868
    0
                .poll(wgpu::PollType::Wait {
    1869
    0
                    submission_index: None,
    1870
    0
                    timeout: None,
    1871
    0
                })
    1872
    0
                .ok();
    1873
    1874
    0
            receiver
    1875
    0
                .receive()
    1876
    0
                .await
    1877
    0
                .ok_or("Failed to receive mapping result")?
    1878
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    1879
    1880
    0
            {
    1881
    0
                let data = buffer_slice.get_mapped_range();
    1882
    0
                let output_data: &[f32] = bytemuck::cast_slice(&data);
    1883
    0
                result.copy_from_slice(output_data);
    1884
    0
            }
    1885
    1886
    0
            staging_buffer.unmap();
    1887
    1888
    0
            Ok(())
    1889
    0
        }
    1890
    1891
        /// Execute symmetric eigendecomposition on GPU (sync, native only)
    1892
        ///
    1893
        /// Computes eigenvalues and eigenvectors using Jacobi algorithm with GPU-accelerated
    1894
        /// Givens rotations. Returns (eigenvalues, eigenvector_data) where eigenvector_data
    1895
        /// is in row-major format.
    1896
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    1897
    0
        pub fn symmetric_eigen(
    1898
    0
            &self,
    1899
    0
            matrix: &[f32],
    1900
    0
            n: usize,
    1901
    0
        ) -> Result<(Vec<f32>, Vec<f32>), String> {
    1902
    0
            runtime::block_on(async { self.symmetric_eigen_async(matrix, n).await })
    1903
    0
        }
    1904
    1905
        /// Execute symmetric eigendecomposition on GPU (async, works on all platforms)
    1906
        ///
    1907
        /// Computes eigenvalues and eigenvectors using Jacobi algorithm with GPU-accelerated
    1908
        /// Givens rotations.
    1909
    0
        pub async fn symmetric_eigen_async(
    1910
    0
            &self,
    1911
    0
            matrix: &[f32],
    1912
    0
            n: usize,
    1913
    0
        ) -> Result<(Vec<f32>, Vec<f32>), String> {
    1914
    0
            if matrix.len() != n * n {
    1915
    0
                return Err(format!(
    1916
    0
                    "Matrix size mismatch: expected {} elements for {}x{} matrix, got {}",
    1917
    0
                    n * n,
    1918
    0
                    n,
    1919
    0
                    n,
    1920
    0
                    matrix.len()
    1921
    0
                ));
    1922
    0
            }
    1923
    1924
    0
            if n == 0 {
    1925
    0
                return Ok((Vec::new(), Vec::new()));
    1926
    0
            }
    1927
    1928
            // For small matrices, use CPU (GPU overhead not worth it)
    1929
    0
            if n < 64 {
    1930
    0
                return self.symmetric_eigen_cpu(matrix, n);
    1931
    0
            }
    1932
    1933
            // Create shader module for Jacobi rotation
    1934
    0
            let rotation_shader = self
    1935
    0
                .device
    1936
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    1937
    0
                    label: Some("Jacobi Rotation Shader"),
    1938
    0
                    source: wgpu::ShaderSource::Wgsl(shaders::JACOBI_ROTATION_SHADER.into()),
    1939
    0
                });
    1940
    1941
            // Create buffers
    1942
    0
            let matrix_size = (n * n * std::mem::size_of::<f32>()) as u64;
    1943
    1944
    0
            let matrix_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1945
    0
                label: Some("Matrix Buffer"),
    1946
    0
                size: matrix_size,
    1947
    0
                usage: wgpu::BufferUsages::STORAGE
    1948
    0
                    | wgpu::BufferUsages::COPY_DST
    1949
    0
                    | wgpu::BufferUsages::COPY_SRC,
    1950
    0
                mapped_at_creation: false,
    1951
    0
            });
    1952
    1953
    0
            let eigenvectors_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1954
    0
                label: Some("Eigenvectors Buffer"),
    1955
    0
                size: matrix_size,
    1956
    0
                usage: wgpu::BufferUsages::STORAGE
    1957
    0
                    | wgpu::BufferUsages::COPY_DST
    1958
    0
                    | wgpu::BufferUsages::COPY_SRC,
    1959
    0
                mapped_at_creation: false,
    1960
    0
            });
    1961
    1962
            // JacobiParams uniform buffer
    1963
            #[repr(C)]
    1964
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    1965
            struct JacobiParams {
    1966
                n: u32,
    1967
                p: u32,
    1968
                q: u32,
    1969
                c: f32,
    1970
                s: f32,
    1971
                _padding: [u32; 3],
    1972
            }
    1973
    1974
    0
            let params_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    1975
    0
                label: Some("Jacobi Params"),
    1976
    0
                size: std::mem::size_of::<JacobiParams>() as u64,
    1977
    0
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    1978
    0
                mapped_at_creation: false,
    1979
    0
            });
    1980
    1981
            // Initialize eigenvectors to identity matrix
    1982
    0
            let mut eigenvectors = vec![0.0f32; n * n];
    1983
    0
            for i in 0..n {
    1984
    0
                eigenvectors[i * n + i] = 1.0;
    1985
    0
            }
    1986
    1987
            // Write initial data
    1988
    0
            self.queue
    1989
    0
                .write_buffer(&matrix_buffer, 0, bytemuck::cast_slice(matrix));
    1990
    0
            self.queue
    1991
    0
                .write_buffer(&eigenvectors_buffer, 0, bytemuck::cast_slice(&eigenvectors));
    1992
    1993
            // Create bind group layout
    1994
    0
            let bind_group_layout =
    1995
    0
                self.device
    1996
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    1997
    0
                        label: Some("Jacobi Bind Group Layout"),
    1998
    0
                        entries: &[
    1999
    0
                            wgpu::BindGroupLayoutEntry {
    2000
    0
                                binding: 0,
    2001
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2002
    0
                                ty: wgpu::BindingType::Buffer {
    2003
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    2004
    0
                                    has_dynamic_offset: false,
    2005
    0
                                    min_binding_size: None,
    2006
    0
                                },
    2007
    0
                                count: None,
    2008
    0
                            },
    2009
    0
                            wgpu::BindGroupLayoutEntry {
    2010
    0
                                binding: 1,
    2011
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2012
    0
                                ty: wgpu::BindingType::Buffer {
    2013
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    2014
    0
                                    has_dynamic_offset: false,
    2015
    0
                                    min_binding_size: None,
    2016
    0
                                },
    2017
    0
                                count: None,
    2018
    0
                            },
    2019
    0
                            wgpu::BindGroupLayoutEntry {
    2020
    0
                                binding: 2,
    2021
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2022
    0
                                ty: wgpu::BindingType::Buffer {
    2023
    0
                                    ty: wgpu::BufferBindingType::Uniform,
    2024
    0
                                    has_dynamic_offset: false,
    2025
    0
                                    min_binding_size: None,
    2026
    0
                                },
    2027
    0
                                count: None,
    2028
    0
                            },
    2029
    0
                        ],
    2030
    0
                    });
    2031
    2032
            // Create bind group
    2033
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    2034
    0
                label: Some("Jacobi Bind Group"),
    2035
    0
                layout: &bind_group_layout,
    2036
    0
                entries: &[
    2037
    0
                    wgpu::BindGroupEntry {
    2038
    0
                        binding: 0,
    2039
    0
                        resource: matrix_buffer.as_entire_binding(),
    2040
    0
                    },
    2041
    0
                    wgpu::BindGroupEntry {
    2042
    0
                        binding: 1,
    2043
    0
                        resource: eigenvectors_buffer.as_entire_binding(),
    2044
    0
                    },
    2045
    0
                    wgpu::BindGroupEntry {
    2046
    0
                        binding: 2,
    2047
    0
                        resource: params_buffer.as_entire_binding(),
    2048
    0
                    },
    2049
    0
                ],
    2050
    0
            });
    2051
    2052
            // Create pipeline
    2053
    0
            let pipeline_layout = self
    2054
    0
                .device
    2055
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    2056
    0
                    label: Some("Jacobi Pipeline Layout"),
    2057
    0
                    bind_group_layouts: &[&bind_group_layout],
    2058
    0
                    push_constant_ranges: &[],
    2059
    0
                });
    2060
    2061
    0
            let rotation_pipeline =
    2062
    0
                self.device
    2063
    0
                    .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    2064
    0
                        label: Some("Jacobi Rotation Pipeline"),
    2065
    0
                        layout: Some(&pipeline_layout),
    2066
    0
                        module: &rotation_shader,
    2067
    0
                        entry_point: Some("main"),
    2068
    0
                        compilation_options: wgpu::PipelineCompilationOptions::default(),
    2069
    0
                        cache: None,
    2070
    0
                    });
    2071
    2072
            // Jacobi iteration
    2073
    0
            let max_sweeps = 50;
    2074
    0
            let tolerance = 1e-7 * (matrix.iter().map(|x| x * x).sum::<f32>().sqrt()).max(1.0);
    2075
    2076
            // Working copy of matrix for CPU-side pivot selection
    2077
    0
            let mut a = matrix.to_vec();
    2078
    2079
    0
            for _sweep in 0..max_sweeps {
    2080
    0
                let mut converged = true;
    2081
    2082
                // Cyclic Jacobi: process all pairs (i, j) where i < j
    2083
    0
                for i in 0..n {
    2084
    0
                    for j in (i + 1)..n {
    2085
    0
                        let aij = a[i * n + j];
    2086
    2087
    0
                        if aij.abs() < tolerance {
    2088
    0
                            continue;
    2089
    0
                        }
    2090
    2091
    0
                        converged = false;
    2092
    2093
                        // Compute rotation parameters
    2094
    0
                        let aii = a[i * n + i];
    2095
    0
                        let ajj = a[j * n + j];
    2096
    2097
    0
                        let tau = (ajj - aii) / (2.0 * aij);
    2098
    0
                        let t = if tau >= 0.0 {
    2099
    0
                            1.0 / (tau + (1.0 + tau * tau).sqrt())
    2100
                        } else {
    2101
    0
                            -1.0 / (-tau + (1.0 + tau * tau).sqrt())
    2102
                        };
    2103
    2104
    0
                        let c = 1.0 / (1.0 + t * t).sqrt();
    2105
    0
                        let s = t * c;
    2106
    2107
                        // Update params and dispatch GPU
    2108
    0
                        let params = JacobiParams {
    2109
    0
                            n: n as u32,
    2110
    0
                            p: i as u32,
    2111
    0
                            q: j as u32,
    2112
    0
                            c,
    2113
    0
                            s,
    2114
    0
                            _padding: [0; 3],
    2115
    0
                        };
    2116
    2117
    0
                        self.queue
    2118
    0
                            .write_buffer(&params_buffer, 0, bytemuck::bytes_of(&params));
    2119
    2120
                        // Create command encoder and dispatch
    2121
    0
                        let mut encoder =
    2122
    0
                            self.device
    2123
    0
                                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    2124
    0
                                    label: Some("Jacobi Rotation Encoder"),
    2125
    0
                                });
    2126
    2127
    0
                        {
    2128
    0
                            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    2129
    0
                                label: Some("Jacobi Rotation Pass"),
    2130
    0
                                timestamp_writes: None,
    2131
    0
                            });
    2132
    0
                            pass.set_pipeline(&rotation_pipeline);
    2133
    0
                            pass.set_bind_group(0, &bind_group, &[]);
    2134
    0
                            pass.dispatch_workgroups((n as u32).div_ceil(256), 1, 1);
    2135
    0
                        }
    2136
    2137
    0
                        self.queue.submit(Some(encoder.finish()));
    2138
    2139
                        // Update local copy of diagonal and off-diagonal
    2140
    0
                        a[i * n + i] = aii - t * aij;
    2141
    0
                        a[j * n + j] = ajj + t * aij;
    2142
    0
                        a[i * n + j] = 0.0;
    2143
    0
                        a[j * n + i] = 0.0;
    2144
    2145
                        // Update off-diagonal elements in rows/columns i and j
    2146
    0
                        for k in 0..n {
    2147
    0
                            if k != i && k != j {
    2148
    0
                                let aki = a[k * n + i];
    2149
    0
                                let akj = a[k * n + j];
    2150
    0
                                a[k * n + i] = c * aki - s * akj;
    2151
    0
                                a[i * n + k] = a[k * n + i];
    2152
    0
                                a[k * n + j] = s * aki + c * akj;
    2153
    0
                                a[j * n + k] = a[k * n + j];
    2154
    0
                            }
    2155
                        }
    2156
                    }
    2157
                }
    2158
    2159
    0
                if converged {
    2160
    0
                    break;
    2161
    0
                }
    2162
            }
    2163
    2164
            // Read back results
    2165
    0
            let staging_matrix = self.device.create_buffer(&wgpu::BufferDescriptor {
    2166
    0
                label: Some("Staging Matrix"),
    2167
    0
                size: matrix_size,
    2168
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    2169
    0
                mapped_at_creation: false,
    2170
    0
            });
    2171
    2172
    0
            let staging_eigenvectors = self.device.create_buffer(&wgpu::BufferDescriptor {
    2173
    0
                label: Some("Staging Eigenvectors"),
    2174
    0
                size: matrix_size,
    2175
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    2176
    0
                mapped_at_creation: false,
    2177
    0
            });
    2178
    2179
    0
            let mut encoder = self
    2180
    0
                .device
    2181
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    2182
    0
                    label: Some("Copy Encoder"),
    2183
    0
                });
    2184
    2185
    0
            encoder.copy_buffer_to_buffer(&matrix_buffer, 0, &staging_matrix, 0, matrix_size);
    2186
    0
            encoder.copy_buffer_to_buffer(
    2187
    0
                &eigenvectors_buffer,
    2188
                0,
    2189
    0
                &staging_eigenvectors,
    2190
                0,
    2191
    0
                matrix_size,
    2192
            );
    2193
    2194
    0
            self.queue.submit(Some(encoder.finish()));
    2195
    2196
            // Map and read eigenvectors
    2197
    0
            let eigenvector_slice = staging_eigenvectors.slice(..);
    2198
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    2199
    0
            eigenvector_slice.map_async(wgpu::MapMode::Read, move |result| {
    2200
    0
                sender.send(result).expect("oneshot channel receiver dropped");
    2201
    0
            });
    2202
    2203
    0
            self.device
    2204
    0
                .poll(wgpu::PollType::Wait {
    2205
    0
                    submission_index: None,
    2206
    0
                    timeout: None,
    2207
    0
                })
    2208
    0
                .ok();
    2209
    2210
    0
            receiver
    2211
    0
                .receive()
    2212
    0
                .await
    2213
    0
                .ok_or("Failed to receive mapping result")?
    2214
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    2215
    2216
    0
            let mut result_eigenvectors = vec![0.0f32; n * n];
    2217
    0
            {
    2218
    0
                let data = eigenvector_slice.get_mapped_range();
    2219
    0
                let output_data: &[f32] = bytemuck::cast_slice(&data);
    2220
    0
                result_eigenvectors.copy_from_slice(output_data);
    2221
    0
            }
    2222
    0
            staging_eigenvectors.unmap();
    2223
    2224
            // Extract eigenvalues from diagonal of working matrix
    2225
    0
            let eigenvalues: Vec<f32> = (0..n).map(|i| a[i * n + i]).collect();
    2226
    2227
    0
            Ok((eigenvalues, result_eigenvectors))
    2228
    0
        }
    2229
    2230
        /// CPU fallback for small matrices (GPU overhead not worthwhile)
    2231
    0
        fn symmetric_eigen_cpu(
    2232
    0
            &self,
    2233
    0
            matrix: &[f32],
    2234
    0
            n: usize,
    2235
    0
        ) -> Result<(Vec<f32>, Vec<f32>), String> {
    2236
    0
            let max_sweeps = 50;
    2237
    0
            let tolerance = 1e-7 * (matrix.iter().map(|x| x * x).sum::<f32>().sqrt()).max(1.0);
    2238
    2239
    0
            let mut a = matrix.to_vec();
    2240
    0
            let mut v = vec![0.0f32; n * n];
    2241
    0
            for i in 0..n {
    2242
    0
                v[i * n + i] = 1.0;
    2243
    0
            }
    2244
    2245
    0
            for _sweep in 0..max_sweeps {
    2246
    0
                let mut converged = true;
    2247
    2248
    0
                for i in 0..n {
    2249
    0
                    for j in (i + 1)..n {
    2250
    0
                        let aij = a[i * n + j];
    2251
    2252
    0
                        if aij.abs() < tolerance {
    2253
    0
                            continue;
    2254
    0
                        }
    2255
    2256
    0
                        converged = false;
    2257
    2258
    0
                        let aii = a[i * n + i];
    2259
    0
                        let ajj = a[j * n + j];
    2260
    2261
    0
                        let tau = (ajj - aii) / (2.0 * aij);
    2262
    0
                        let t = if tau >= 0.0 {
    2263
    0
                            1.0 / (tau + (1.0 + tau * tau).sqrt())
    2264
                        } else {
    2265
    0
                            -1.0 / (-tau + (1.0 + tau * tau).sqrt())
    2266
                        };
    2267
    2268
    0
                        let c = 1.0 / (1.0 + t * t).sqrt();
    2269
    0
                        let s = t * c;
    2270
    2271
                        // Update diagonal
    2272
    0
                        a[i * n + i] = aii - t * aij;
    2273
    0
                        a[j * n + j] = ajj + t * aij;
    2274
    0
                        a[i * n + j] = 0.0;
    2275
    0
                        a[j * n + i] = 0.0;
    2276
    2277
                        // Update off-diagonal
    2278
    0
                        for k in 0..n {
    2279
    0
                            if k != i && k != j {
    2280
    0
                                let aki = a[k * n + i];
    2281
    0
                                let akj = a[k * n + j];
    2282
    0
                                a[k * n + i] = c * aki - s * akj;
    2283
    0
                                a[i * n + k] = a[k * n + i];
    2284
    0
                                a[k * n + j] = s * aki + c * akj;
    2285
    0
                                a[j * n + k] = a[k * n + j];
    2286
    0
                            }
    2287
                        }
    2288
    2289
                        // Update eigenvectors
    2290
    0
                        for k in 0..n {
    2291
    0
                            let vki = v[k * n + i];
    2292
    0
                            let vkj = v[k * n + j];
    2293
    0
                            v[k * n + i] = c * vki - s * vkj;
    2294
    0
                            v[k * n + j] = s * vki + c * vkj;
    2295
    0
                        }
    2296
                    }
    2297
                }
    2298
    2299
    0
                if converged {
    2300
    0
                    break;
    2301
    0
                }
    2302
            }
    2303
    2304
    0
            let eigenvalues: Vec<f32> = (0..n).map(|i| a[i * n + i]).collect();
    2305
    0
            Ok((eigenvalues, v))
    2306
    0
        }
    2307
    2308
        /// 2D Tiled Sum Reduction on GPU (sync, native only)
    2309
        ///
    2310
        /// Uses 16×16 workgroups for efficient parallel reduction with
    2311
        /// optimal memory coalescing. GPU version of `tiled_sum_2d`.
    2312
        ///
    2313
        /// # Arguments
    2314
        ///
    2315
        /// * `data` - Input 2D data in row-major order
    2316
        /// * `width` - Number of columns
    2317
        /// * `height` - Number of rows
    2318
        ///
    2319
        /// # Returns
    2320
        ///
    2321
        /// Sum of all elements
    2322
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2323
    0
        pub fn tiled_sum_2d(&self, data: &[f32], width: usize, height: usize) -> Result<f32, String> {
    2324
    0
            runtime::block_on(self.tiled_sum_2d_async(data, width, height))
    2325
    0
        }
    2326
    2327
        /// 2D Tiled Sum Reduction on GPU (async, works on all platforms)
    2328
    0
        pub async fn tiled_sum_2d_async(
    2329
    0
            &self,
    2330
    0
            data: &[f32],
    2331
    0
            width: usize,
    2332
    0
            height: usize,
    2333
    0
        ) -> Result<f32, String> {
    2334
    0
            self.tiled_reduce_2d_async(
    2335
    0
                data,
    2336
    0
                width,
    2337
    0
                height,
    2338
    0
                shaders::TILED_SUM_REDUCTION_SHADER,
    2339
    0
                "TiledSum",
    2340
                0.0, // identity for sum
    2341
    0
                |partials| partials.iter().sum(),
    2342
            )
    2343
    0
            .await
    2344
    0
        }
    2345
    2346
        /// 2D Tiled Max Reduction on GPU (sync, native only)
    2347
        ///
    2348
        /// Uses 16×16 workgroups for efficient parallel max reduction.
    2349
        /// GPU version of `tiled_max_2d`.
    2350
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2351
    0
        pub fn tiled_max_2d(&self, data: &[f32], width: usize, height: usize) -> Result<f32, String> {
    2352
    0
            runtime::block_on(self.tiled_max_2d_async(data, width, height))
    2353
    0
        }
    2354
    2355
        /// 2D Tiled Max Reduction on GPU (async, works on all platforms)
    2356
    0
        pub async fn tiled_max_2d_async(
    2357
    0
            &self,
    2358
    0
            data: &[f32],
    2359
    0
            width: usize,
    2360
    0
            height: usize,
    2361
    0
        ) -> Result<f32, String> {
    2362
    0
            self.tiled_reduce_2d_async(
    2363
    0
                data,
    2364
    0
                width,
    2365
    0
                height,
    2366
    0
                shaders::TILED_MAX_REDUCTION_SHADER,
    2367
    0
                "TiledMax",
    2368
                f32::NEG_INFINITY, // identity for max
    2369
    0
                |partials| partials.iter().copied().fold(f32::NEG_INFINITY, f32::max),
    2370
            )
    2371
    0
            .await
    2372
    0
        }
    2373
    2374
        /// 2D Tiled Min Reduction on GPU (sync, native only)
    2375
        ///
    2376
        /// Uses 16×16 workgroups for efficient parallel min reduction.
    2377
        /// GPU version of `tiled_min_2d`.
    2378
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2379
    0
        pub fn tiled_min_2d(&self, data: &[f32], width: usize, height: usize) -> Result<f32, String> {
    2380
    0
            runtime::block_on(self.tiled_min_2d_async(data, width, height))
    2381
    0
        }
    2382
    2383
        /// 2D Tiled Min Reduction on GPU (async, works on all platforms)
    2384
    0
        pub async fn tiled_min_2d_async(
    2385
    0
            &self,
    2386
    0
            data: &[f32],
    2387
    0
            width: usize,
    2388
    0
            height: usize,
    2389
    0
        ) -> Result<f32, String> {
    2390
    0
            self.tiled_reduce_2d_async(
    2391
    0
                data,
    2392
    0
                width,
    2393
    0
                height,
    2394
    0
                shaders::TILED_MIN_REDUCTION_SHADER,
    2395
    0
                "TiledMin",
    2396
                f32::INFINITY, // identity for min
    2397
    0
                |partials| partials.iter().copied().fold(f32::INFINITY, f32::min),
    2398
            )
    2399
    0
            .await
    2400
    0
        }
    2401
    2402
        /// Generic 2D tiled reduction helper
    2403
        #[allow(clippy::too_many_arguments)]
    2404
    0
        async fn tiled_reduce_2d_async<F>(
    2405
    0
            &self,
    2406
    0
            data: &[f32],
    2407
    0
            width: usize,
    2408
    0
            height: usize,
    2409
    0
            shader_source: &str,
    2410
    0
            op_name: &str,
    2411
    0
            identity: f32,
    2412
    0
            combine: F,
    2413
    0
        ) -> Result<f32, String>
    2414
    0
        where
    2415
    0
            F: Fn(&[f32]) -> f32,
    2416
    0
        {
    2417
    0
            if data.is_empty() || width == 0 || height == 0 {
    2418
    0
                return Ok(identity);
    2419
    0
            }
    2420
    2421
            // Calculate workgroup dimensions (16×16 tiles)
    2422
    0
            let workgroup_size_x: u32 = 16;
    2423
    0
            let workgroup_size_y: u32 = 16;
    2424
    0
            let num_workgroups_x = (width as u32).div_ceil(workgroup_size_x);
    2425
    0
            let num_workgroups_y = (height as u32).div_ceil(workgroup_size_y);
    2426
    0
            let total_workgroups = (num_workgroups_x * num_workgroups_y) as usize;
    2427
    2428
            // Create shader module
    2429
    0
            let shader = self
    2430
    0
                .device
    2431
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    2432
    0
                    label: Some(&format!("{} Shader", op_name)),
    2433
    0
                    source: wgpu::ShaderSource::Wgsl(shader_source.into()),
    2434
    0
                });
    2435
    2436
            // Create input buffer
    2437
    0
            let input_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    2438
    0
                label: Some(&format!("{} Input", op_name)),
    2439
    0
                size: std::mem::size_of_val(data) as u64,
    2440
    0
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
    2441
    0
                mapped_at_creation: false,
    2442
    0
            });
    2443
    2444
            // Create partial results buffer
    2445
    0
            let partial_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    2446
    0
                label: Some(&format!("{} Partial Results", op_name)),
    2447
    0
                size: (total_workgroups * std::mem::size_of::<f32>()) as u64,
    2448
    0
                usage: wgpu::BufferUsages::STORAGE
    2449
    0
                    | wgpu::BufferUsages::COPY_SRC
    2450
    0
                    | wgpu::BufferUsages::COPY_DST,
    2451
    0
                mapped_at_creation: false,
    2452
    0
            });
    2453
    2454
            // Dimensions uniform buffer
    2455
            #[repr(C)]
    2456
            #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    2457
            struct Dimensions {
    2458
                width: u32,
    2459
                height: u32,
    2460
            }
    2461
    2462
    0
            let dims = Dimensions {
    2463
    0
                width: width as u32,
    2464
    0
                height: height as u32,
    2465
    0
            };
    2466
    2467
    0
            let dims_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    2468
    0
                label: Some(&format!("{} Dimensions", op_name)),
    2469
    0
                size: std::mem::size_of::<Dimensions>() as u64,
    2470
    0
                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    2471
    0
                mapped_at_creation: false,
    2472
    0
            });
    2473
    2474
            // Write data
    2475
    0
            self.queue
    2476
    0
                .write_buffer(&input_buffer, 0, bytemuck::cast_slice(data));
    2477
    0
            self.queue
    2478
    0
                .write_buffer(&dims_buffer, 0, bytemuck::bytes_of(&dims));
    2479
    2480
            // Create bind group layout
    2481
    0
            let bind_group_layout =
    2482
    0
                self.device
    2483
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    2484
    0
                        label: Some(&format!("{} Bind Group Layout", op_name)),
    2485
    0
                        entries: &[
    2486
    0
                            wgpu::BindGroupLayoutEntry {
    2487
    0
                                binding: 0,
    2488
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2489
    0
                                ty: wgpu::BindingType::Buffer {
    2490
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    2491
    0
                                    has_dynamic_offset: false,
    2492
    0
                                    min_binding_size: None,
    2493
    0
                                },
    2494
    0
                                count: None,
    2495
    0
                            },
    2496
    0
                            wgpu::BindGroupLayoutEntry {
    2497
    0
                                binding: 1,
    2498
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2499
    0
                                ty: wgpu::BindingType::Buffer {
    2500
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    2501
    0
                                    has_dynamic_offset: false,
    2502
    0
                                    min_binding_size: None,
    2503
    0
                                },
    2504
    0
                                count: None,
    2505
    0
                            },
    2506
    0
                            wgpu::BindGroupLayoutEntry {
    2507
    0
                                binding: 2,
    2508
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    2509
    0
                                ty: wgpu::BindingType::Buffer {
    2510
    0
                                    ty: wgpu::BufferBindingType::Uniform,
    2511
    0
                                    has_dynamic_offset: false,
    2512
    0
                                    min_binding_size: None,
    2513
    0
                                },
    2514
    0
                                count: None,
    2515
    0
                            },
    2516
    0
                        ],
    2517
    0
                    });
    2518
    2519
            // Create bind group
    2520
    0
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
    2521
    0
                label: Some(&format!("{} Bind Group", op_name)),
    2522
    0
                layout: &bind_group_layout,
    2523
    0
                entries: &[
    2524
    0
                    wgpu::BindGroupEntry {
    2525
    0
                        binding: 0,
    2526
    0
                        resource: input_buffer.as_entire_binding(),
    2527
    0
                    },
    2528
    0
                    wgpu::BindGroupEntry {
    2529
    0
                        binding: 1,
    2530
    0
                        resource: partial_buffer.as_entire_binding(),
    2531
    0
                    },
    2532
    0
                    wgpu::BindGroupEntry {
    2533
    0
                        binding: 2,
    2534
    0
                        resource: dims_buffer.as_entire_binding(),
    2535
    0
                    },
    2536
    0
                ],
    2537
    0
            });
    2538
    2539
            // Create pipeline
    2540
    0
            let pipeline_layout = self
    2541
    0
                .device
    2542
    0
                .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    2543
    0
                    label: Some(&format!("{} Pipeline Layout", op_name)),
    2544
    0
                    bind_group_layouts: &[&bind_group_layout],
    2545
    0
                    push_constant_ranges: &[],
    2546
    0
                });
    2547
    2548
    0
            let pipeline = self
    2549
    0
                .device
    2550
    0
                .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    2551
    0
                    label: Some(&format!("{} Pipeline", op_name)),
    2552
    0
                    layout: Some(&pipeline_layout),
    2553
    0
                    module: &shader,
    2554
    0
                    entry_point: Some("main"),
    2555
    0
                    compilation_options: Default::default(),
    2556
    0
                    cache: None,
    2557
    0
                });
    2558
    2559
            // Create staging buffer
    2560
    0
            let staging_buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
    2561
    0
                label: Some(&format!("{} Staging", op_name)),
    2562
    0
                size: (total_workgroups * std::mem::size_of::<f32>()) as u64,
    2563
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    2564
    0
                mapped_at_creation: false,
    2565
    0
            });
    2566
    2567
            // Create command encoder
    2568
    0
            let mut encoder = self
    2569
    0
                .device
    2570
    0
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    2571
    0
                    label: Some(&format!("{} Encoder", op_name)),
    2572
    0
                });
    2573
    2574
    0
            {
    2575
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    2576
    0
                    label: Some(&format!("{} Pass", op_name)),
    2577
    0
                    timestamp_writes: None,
    2578
    0
                });
    2579
    0
                compute_pass.set_pipeline(&pipeline);
    2580
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    2581
    0
                compute_pass.dispatch_workgroups(num_workgroups_x, num_workgroups_y, 1);
    2582
    0
            }
    2583
    2584
            // Copy result to staging buffer
    2585
    0
            encoder.copy_buffer_to_buffer(
    2586
    0
                &partial_buffer,
    2587
                0,
    2588
    0
                &staging_buffer,
    2589
                0,
    2590
    0
                (total_workgroups * std::mem::size_of::<f32>()) as u64,
    2591
            );
    2592
    2593
            // Submit commands
    2594
    0
            self.queue.submit(Some(encoder.finish()));
    2595
    2596
            // Read back results
    2597
    0
            let buffer_slice = staging_buffer.slice(..);
    2598
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    2599
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    2600
    0
                sender.send(result).ok();
    2601
    0
            });
    2602
    2603
            // Poll device
    2604
    0
            self.device
    2605
    0
                .poll(wgpu::PollType::Wait {
    2606
    0
                    submission_index: None,
    2607
    0
                    timeout: None,
    2608
    0
                })
    2609
    0
                .ok();
    2610
    2611
    0
            receiver
    2612
    0
                .receive()
    2613
    0
                .await
    2614
    0
                .ok_or("Failed to receive mapping result")?
    2615
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    2616
    2617
    0
            let final_result = {
    2618
    0
                let data = buffer_slice.get_mapped_range();
    2619
    0
                let partials: &[f32] = bytemuck::cast_slice(&data);
    2620
    0
                combine(partials)
    2621
            };
    2622
    2623
    0
            staging_buffer.unmap();
    2624
    2625
    0
            Ok(final_result)
    2626
    0
        }
    2627
    }
    2628
    2629
    #[cfg(all(test, feature = "gpu", not(target_arch = "wasm32")))]
    2630
    mod tests {
    2631
        use super::*;
    2632
    2633
        #[test]
    2634
        fn test_is_available_consistency() {
    2635
            // EXTREME TDD: Kill mutant that replaces is_available() with hardcoded false
    2636
            // Test that is_available() is consistent with GpuDevice::new()
    2637
            let available = GpuDevice::is_available();
    2638
            let device_result = GpuDevice::new();
    2639
    2640
            if available {
    2641
                // If is_available() returns true, device creation should succeed
    2642
                assert!(
    2643
                    device_result.is_ok(),
    2644
                    "is_available() returned true, but GpuDevice::new() failed"
    2645
                );
    2646
            } else {
    2647
                // If is_available() returns false, we can't make assertions about new()
    2648
                // (it might still succeed in some edge cases, but typically should fail)
    2649
                // The key test is: mutant always returns false, so on GPU systems this fails
    2650
                eprintln!(
    2651
                    "GPU not available (is_available=false), device creation result: {:?}",
    2652
                    device_result.is_err()
    2653
                );
    2654
            }
    2655
        }
    2656
    2657
        #[test]
    2658
        fn test_reduce_sum_not_hardcoded() {
    2659
            // EXTREME TDD: Kill mutant that replaces reduce_sum with Ok(-1.0)
    2660
            if !GpuDevice::is_available() {
    2661
                eprintln!("GPU not available, skipping test");
    2662
                return;
    2663
            }
    2664
    2665
            let device = GpuDevice::new().expect("Failed to create GPU device");
    2666
            let input = vec![1.0, 2.0, 3.0, 4.0, 5.0]; // sum = 15.0
    2667
    2668
            // reduce_sum is async, so we use runtime::block_on
    2669
            let result = runtime::block_on(device.reduce_sum(&input)).expect("reduce_sum failed");
    2670
    2671
            // Kill mutant: verify result is NOT -1.0
    2672
            assert_ne!(
    2673
                result, -1.0,
    2674
                "reduce_sum returned hardcoded -1.0 (mutant not killed)"
    2675
            );
    2676
    2677
            // Verify correct computation
    2678
            let expected: f32 = input.iter().sum();
    2679
            assert!(
    2680
                (result - expected).abs() < 1e-4,
    2681
                "reduce_sum({:?}) = {} (expected {})",
    2682
                input,
    2683
                result,
    2684
                expected
    2685
            );
    2686
        }
    2687
    }