## 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/batch.rs

    Line
    Count
    Source
    1
    //! Async GPU command batching for reduced transfer overhead
    2
    //!
    3
    //! This module provides an async API for GPU operations that batches multiple
    4
    //! operations together to minimize CPU↔GPU data transfers.
    5
    //!
    6
    //! # Motivation
    7
    //!
    8
    //! The synchronous GPU API transfers data for each operation:
    9
    //! ```text
    10
    //! vec.relu()      // Upload → GPU compute → Download
    11
    //! vec.scale(2.0)  // Upload → GPU compute → Download
    12
    //! vec.add(&other) // Upload → GPU compute → Download
    13
    //! Total: 6 transfers (3 up, 3 down)
    14
    //! ```
    15
    //!
    16
    //! The async batch API queues operations and executes them together:
    17
    //! ```text
    18
    //! batch.relu(input)
    19
    //! batch.scale(relu_out, 2.0)
    20
    //! batch.add(scaled, other)
    21
    //! batch.execute()  // Upload once → 3 GPU computes → Download once
    22
    //! Total: 2 transfers (1 up, 1 down)  // 3x reduction!
    23
    //! ```
    24
    //!
    25
    //! # Example
    26
    //!
    27
    //! ```rust,no_run
    28
    //! use trueno::backends::gpu::{GpuDevice, GpuCommandBatch};
    29
    //!
    30
    //! # async fn example() -> Result<(), String> {
    31
    //! let device = GpuDevice::new()?;
    32
    //! let mut batch = GpuCommandBatch::new(device);
    33
    //!
    34
    //! // Queue operations (no GPU execution yet)
    35
    //! let input = batch.upload(&[1.0, 2.0, -3.0, 4.0]);
    36
    //! let relu_out = batch.relu(input);
    37
    //! let scaled = batch.scale(relu_out, 2.0);
    38
    //! let other = batch.upload(&[0.5, 0.5, 0.5, 0.5]);
    39
    //! let final_out = batch.add(scaled, other);
    40
    //!
    41
    //! // Execute all operations in single batch
    42
    //! batch.execute().await?;
    43
    //!
    44
    //! // Read final result
    45
    //! let result = batch.read(final_out).await?;
    46
    //! assert_eq!(result, vec![2.5, 4.5, 0.5, 8.5]);
    47
    //! # Ok(())
    48
    //! # }
    49
    //! ```
    50
    51
    use super::GpuDevice;
    52
    use std::collections::HashMap;
    53
    use std::sync::Arc;
    54
    use wgpu;
    55
    56
    /// Unique identifier for a buffer in a batch
    57
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    58
    pub struct BufferId(usize);
    59
    60
    /// GPU operation to be executed in a batch
    61
    #[derive(Debug)]
    62
    enum GpuOp {
    63
        /// ReLU activation: max(0, x)
    64
        Relu { input: BufferId, output: BufferId },
    65
    66
        /// Scalar multiplication: x * scalar
    67
        Scale {
    68
            input: BufferId,
    69
            output: BufferId,
    70
            scalar: f32,
    71
        },
    72
    73
        /// Element-wise addition: a + b
    74
        Add {
    75
            a: BufferId,
    76
            b: BufferId,
    77
            output: BufferId,
    78
        },
    79
    80
        /// Element-wise multiplication: a * b
    81
        Mul {
    82
            a: BufferId,
    83
            b: BufferId,
    84
            output: BufferId,
    85
        },
    86
    87
        /// Dot product: sum(a[i] * b[i])
    88
        Dot {
    89
            a: BufferId,
    90
            b: BufferId,
    91
            output: BufferId, // Single-element buffer for result
    92
        },
    93
    94
        /// Sigmoid activation: 1 / (1 + exp(-x))
    95
        Sigmoid { input: BufferId, output: BufferId },
    96
    97
        /// Hyperbolic tangent: tanh(x)
    98
        Tanh { input: BufferId, output: BufferId },
    99
    100
        /// Swish activation: x * sigmoid(x)
    101
        Swish { input: BufferId, output: BufferId },
    102
    103
        /// GELU activation: x * Φ(x) where Φ is cumulative distribution function
    104
        Gelu { input: BufferId, output: BufferId },
    105
    106
        /// Element-wise subtraction: a - b
    107
        Sub {
    108
            a: BufferId,
    109
            b: BufferId,
    110
            output: BufferId,
    111
        },
    112
    }
    113
    114
    /// Command batch for async GPU execution
    115
    ///
    116
    /// Accumulates GPU operations and executes them together to minimize
    117
    /// CPU↔GPU data transfers.
    118
    pub struct GpuCommandBatch {
    119
        device: Arc<GpuDevice>,
    120
        operations: Vec<GpuOp>,
    121
        buffers: HashMap<BufferId, BufferInfo>,
    122
        next_buffer_id: usize,
    123
    }
    124
    125
    /// Information about a buffer in the batch
    126
    #[derive(Debug)]
    127
    struct BufferInfo {
    128
        /// Size in elements (f32)
    129
        size: usize,
    130
    131
        /// Initial data to upload (if any)
    132
        data: Option<Vec<f32>>,
    133
    134
        /// GPU buffer (created during execute())
    135
        gpu_buffer: Option<wgpu::Buffer>,
    136
    }
    137
    138
    impl GpuCommandBatch {
    139
        /// Create a new command batch
    140
    0
        pub fn new(device: GpuDevice) -> Self {
    141
    0
            Self {
    142
    0
                device: Arc::new(device),
    143
    0
                operations: Vec::new(),
    144
    0
                buffers: HashMap::new(),
    145
    0
                next_buffer_id: 0,
    146
    0
            }
    147
    0
        }
    148
    149
        /// Allocate a new buffer ID
    150
    0
        fn alloc_buffer(&mut self, size: usize, data: Option<Vec<f32>>) -> BufferId {
    151
    0
            let id = BufferId(self.next_buffer_id);
    152
    0
            self.next_buffer_id += 1;
    153
    154
    0
            self.buffers.insert(
    155
    0
                id,
    156
    0
                BufferInfo {
    157
    0
                    size,
    158
    0
                    data,
    159
    0
                    gpu_buffer: None,
    160
    0
                },
    161
            );
    162
    163
    0
            id
    164
    0
        }
    165
    166
        /// Upload data to GPU (queued for batch execution)
    167
        ///
    168
        /// Returns a buffer ID that can be used in subsequent operations.
    169
    0
        pub fn upload(&mut self, data: &[f32]) -> BufferId {
    170
    0
            self.alloc_buffer(data.len(), Some(data.to_vec()))
    171
    0
        }
    172
    173
        /// Allocate an output buffer for an operation
    174
    0
        fn alloc_output(&mut self, size: usize) -> BufferId {
    175
    0
            self.alloc_buffer(size, None)
    176
    0
        }
    177
    178
        /// Queue ReLU operation: max(0, x)
    179
        ///
    180
        /// Returns buffer ID for the output.
    181
    0
        pub fn relu(&mut self, input: BufferId) -> BufferId {
    182
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    183
    184
    0
            let output = self.alloc_output(size);
    185
    186
    0
            self.operations.push(GpuOp::Relu { input, output });
    187
    188
    0
            output
    189
    0
        }
    190
    191
        /// Queue scalar multiplication: x * scalar
    192
        ///
    193
        /// Returns buffer ID for the output.
    194
    0
        pub fn scale(&mut self, input: BufferId, scalar: f32) -> BufferId {
    195
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    196
    197
    0
            let output = self.alloc_output(size);
    198
    199
    0
            self.operations.push(GpuOp::Scale {
    200
    0
                input,
    201
    0
                output,
    202
    0
                scalar,
    203
    0
            });
    204
    205
    0
            output
    206
    0
        }
    207
    208
        /// Queue element-wise addition: a + b
    209
        ///
    210
        /// Returns buffer ID for the output.
    211
        ///
    212
        /// # Panics
    213
        ///
    214
        /// Panics if buffers have different sizes.
    215
    0
        pub fn add(&mut self, a: BufferId, b: BufferId) -> BufferId {
    216
    0
            let size_a = self.buffers.get(&a).expect("Invalid buffer ID").size;
    217
    0
            let size_b = self.buffers.get(&b).expect("Invalid buffer ID").size;
    218
    219
    0
            assert_eq!(
    220
                size_a, size_b,
    221
    0
                "Buffer size mismatch: {} vs {}",
    222
                size_a, size_b
    223
            );
    224
    225
    0
            let output = self.alloc_output(size_a);
    226
    227
    0
            self.operations.push(GpuOp::Add { a, b, output });
    228
    229
    0
            output
    230
    0
        }
    231
    232
        /// Queue element-wise multiplication: a * b
    233
        ///
    234
        /// Returns buffer ID for the output.
    235
        ///
    236
        /// # Panics
    237
        ///
    238
        /// Panics if buffers have different sizes.
    239
    0
        pub fn mul(&mut self, a: BufferId, b: BufferId) -> BufferId {
    240
    0
            let size_a = self.buffers.get(&a).expect("Invalid buffer ID").size;
    241
    0
            let size_b = self.buffers.get(&b).expect("Invalid buffer ID").size;
    242
    243
    0
            assert_eq!(
    244
                size_a, size_b,
    245
    0
                "Buffer size mismatch: {} vs {}",
    246
                size_a, size_b
    247
            );
    248
    249
    0
            let output = self.alloc_output(size_a);
    250
    251
    0
            self.operations.push(GpuOp::Mul { a, b, output });
    252
    253
    0
            output
    254
    0
        }
    255
    256
        /// Queue dot product: sum(a[i] * b[i])
    257
        ///
    258
        /// Returns buffer ID for a single-element output buffer.
    259
        ///
    260
        /// # Panics
    261
        ///
    262
        /// Panics if buffers have different sizes.
    263
    0
        pub fn dot(&mut self, a: BufferId, b: BufferId) -> BufferId {
    264
    0
            let size_a = self.buffers.get(&a).expect("Invalid buffer ID").size;
    265
    0
            let size_b = self.buffers.get(&b).expect("Invalid buffer ID").size;
    266
    267
    0
            assert_eq!(
    268
                size_a, size_b,
    269
    0
                "Buffer size mismatch: {} vs {}",
    270
                size_a, size_b
    271
            );
    272
    273
    0
            let output = self.alloc_output(1); // Dot product returns scalar
    274
    275
    0
            self.operations.push(GpuOp::Dot { a, b, output });
    276
    277
    0
            output
    278
    0
        }
    279
    280
        /// Queue sigmoid activation: 1 / (1 + exp(-x))
    281
        ///
    282
        /// Returns buffer ID for the output.
    283
    0
        pub fn sigmoid(&mut self, input: BufferId) -> BufferId {
    284
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    285
    286
    0
            let output = self.alloc_output(size);
    287
    288
    0
            self.operations.push(GpuOp::Sigmoid { input, output });
    289
    290
    0
            output
    291
    0
        }
    292
    293
        /// Queue hyperbolic tangent: tanh(x)
    294
        ///
    295
        /// Returns buffer ID for the output.
    296
    0
        pub fn tanh(&mut self, input: BufferId) -> BufferId {
    297
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    298
    299
    0
            let output = self.alloc_output(size);
    300
    301
    0
            self.operations.push(GpuOp::Tanh { input, output });
    302
    303
    0
            output
    304
    0
        }
    305
    306
        /// Queue Swish activation: x * sigmoid(x)
    307
        ///
    308
        /// Returns buffer ID for the output.
    309
    0
        pub fn swish(&mut self, input: BufferId) -> BufferId {
    310
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    311
    312
    0
            let output = self.alloc_output(size);
    313
    314
    0
            self.operations.push(GpuOp::Swish { input, output });
    315
    316
    0
            output
    317
    0
        }
    318
    319
        /// Queue GELU activation: x * Φ(x)
    320
        ///
    321
        /// Returns buffer ID for the output.
    322
    0
        pub fn gelu(&mut self, input: BufferId) -> BufferId {
    323
    0
            let size = self.buffers.get(&input).expect("Invalid buffer ID").size;
    324
    325
    0
            let output = self.alloc_output(size);
    326
    327
    0
            self.operations.push(GpuOp::Gelu { input, output });
    328
    329
    0
            output
    330
    0
        }
    331
    332
        /// Queue element-wise subtraction: a - b
    333
        ///
    334
        /// Returns buffer ID for the output.
    335
        ///
    336
        /// # Panics
    337
        ///
    338
        /// Panics if buffers have different sizes.
    339
    0
        pub fn sub(&mut self, a: BufferId, b: BufferId) -> BufferId {
    340
    0
            let size_a = self.buffers.get(&a).expect("Invalid buffer ID").size;
    341
    0
            let size_b = self.buffers.get(&b).expect("Invalid buffer ID").size;
    342
    343
    0
            assert_eq!(
    344
                size_a, size_b,
    345
    0
                "Buffer size mismatch: {} vs {}",
    346
                size_a, size_b
    347
            );
    348
    349
    0
            let output = self.alloc_output(size_a);
    350
    351
    0
            self.operations.push(GpuOp::Sub { a, b, output });
    352
    353
    0
            output
    354
    0
        }
    355
    356
        /// Execute all queued operations on GPU
    357
        ///
    358
        /// This performs all GPU operations in a single batch:
    359
        /// 1. Upload all input buffers once
    360
        /// 2. Execute all operations sequentially on GPU
    361
        /// 3. Results stay on GPU until `read()` is called
    362
    0
        pub async fn execute(&mut self) -> Result<(), String> {
    363
            // Step 1: Create GPU buffers for all BufferIds
    364
    0
            for (buffer_id, buffer_info) in &mut self.buffers {
    365
    0
                let size_bytes = (buffer_info.size * std::mem::size_of::<f32>()) as u64;
    366
    0
    367
    0
                let gpu_buffer = self.device.device.create_buffer(&wgpu::BufferDescriptor {
    368
    0
                    label: Some(&format!("Buffer {:?}", buffer_id)),
    369
    0
                    size: size_bytes,
    370
    0
                    usage: wgpu::BufferUsages::STORAGE
    371
    0
                        | wgpu::BufferUsages::COPY_SRC
    372
    0
                        | wgpu::BufferUsages::COPY_DST,
    373
    0
                    mapped_at_creation: false,
    374
    0
                });
    375
    0
    376
    0
                buffer_info.gpu_buffer = Some(gpu_buffer);
    377
    0
            }
    378
    379
            // Step 2: Upload initial data to buffers that have it
    380
    0
            for buffer_info in self.buffers.values() {
    381
    0
                if let Some(data) = &buffer_info.data {
    382
    0
                    if let Some(gpu_buffer) = &buffer_info.gpu_buffer {
    383
    0
                        self.device
    384
    0
                            .queue
    385
    0
                            .write_buffer(gpu_buffer, 0, bytemuck::cast_slice(data));
    386
    0
                    }
    387
    0
                }
    388
            }
    389
    390
            // Step 3: Execute each operation
    391
    0
            for op in &self.operations {
    392
    0
                self.execute_operation(op).await?;
    393
            }
    394
    395
    0
            Ok(())
    396
    0
        }
    397
    398
        /// Execute a single GPU operation
    399
    0
        async fn execute_operation(&self, op: &GpuOp) -> Result<(), String> {
    400
            use super::shaders;
    401
    402
    0
            match op {
    403
    0
                GpuOp::Relu { input, output } => {
    404
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    405
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    406
    407
    0
                    let input_buffer = input_info
    408
    0
                        .gpu_buffer
    409
    0
                        .as_ref()
    410
    0
                        .ok_or("Input buffer not created")?;
    411
    0
                    let output_buffer = output_info
    412
    0
                        .gpu_buffer
    413
    0
                        .as_ref()
    414
    0
                        .ok_or("Output buffer not created")?;
    415
    416
    0
                    self.execute_unary_op::<()>(
    417
    0
                        shaders::RELU_SHADER,
    418
    0
                        "ReLU",
    419
    0
                        input_buffer,
    420
    0
                        output_buffer,
    421
    0
                        input_info.size,
    422
    0
                        None,
    423
    0
                    )
    424
    0
                    .await?;
    425
                }
    426
    427
                GpuOp::Scale {
    428
    0
                    input,
    429
    0
                    output,
    430
    0
                    scalar,
    431
                } => {
    432
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    433
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    434
    435
    0
                    let input_buffer = input_info
    436
    0
                        .gpu_buffer
    437
    0
                        .as_ref()
    438
    0
                        .ok_or("Input buffer not created")?;
    439
    0
                    let output_buffer = output_info
    440
    0
                        .gpu_buffer
    441
    0
                        .as_ref()
    442
    0
                        .ok_or("Output buffer not created")?;
    443
    444
                    // Create uniform buffer for scalar parameter
    445
                    #[repr(C)]
    446
                    #[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
    447
                    struct ScaleParams {
    448
                        scalar: f32,
    449
                        _padding: [f32; 3], // Uniform buffer alignment
    450
                    }
    451
    452
    0
                    let params = ScaleParams {
    453
    0
                        scalar: *scalar,
    454
    0
                        _padding: [0.0; 3],
    455
    0
                    };
    456
    457
    0
                    self.execute_unary_op(
    458
    0
                        shaders::SCALE_SHADER,
    459
    0
                        "Scale",
    460
    0
                        input_buffer,
    461
    0
                        output_buffer,
    462
    0
                        input_info.size,
    463
    0
                        Some(&params),
    464
    0
                    )
    465
    0
                    .await?;
    466
                }
    467
    468
    0
                GpuOp::Add { a, b, output } => {
    469
    0
                    let a_info = self.buffers.get(a).ok_or("Invalid buffer A ID")?;
    470
    0
                    let b_info = self.buffers.get(b).ok_or("Invalid buffer B ID")?;
    471
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    472
    473
    0
                    let a_buffer = a_info.gpu_buffer.as_ref().ok_or("Buffer A not created")?;
    474
    0
                    let b_buffer = b_info.gpu_buffer.as_ref().ok_or("Buffer B not created")?;
    475
    0
                    let output_buffer = output_info
    476
    0
                        .gpu_buffer
    477
    0
                        .as_ref()
    478
    0
                        .ok_or("Output buffer not created")?;
    479
    480
    0
                    self.execute_binary_op(
    481
    0
                        shaders::VEC_ADD_SHADER,
    482
    0
                        "Add",
    483
    0
                        a_buffer,
    484
    0
                        b_buffer,
    485
    0
                        output_buffer,
    486
    0
                        a_info.size,
    487
    0
                    )
    488
    0
                    .await?;
    489
                }
    490
    491
    0
                GpuOp::Mul { a, b, output } => {
    492
    0
                    let a_info = self.buffers.get(a).ok_or("Invalid buffer A ID")?;
    493
    0
                    let b_info = self.buffers.get(b).ok_or("Invalid buffer B ID")?;
    494
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    495
    496
    0
                    let a_buffer = a_info.gpu_buffer.as_ref().ok_or("Buffer A not created")?;
    497
    0
                    let b_buffer = b_info.gpu_buffer.as_ref().ok_or("Buffer B not created")?;
    498
    0
                    let output_buffer = output_info
    499
    0
                        .gpu_buffer
    500
    0
                        .as_ref()
    501
    0
                        .ok_or("Output buffer not created")?;
    502
    503
    0
                    self.execute_binary_op(
    504
    0
                        shaders::VEC_MUL_SHADER,
    505
    0
                        "Mul",
    506
    0
                        a_buffer,
    507
    0
                        b_buffer,
    508
    0
                        output_buffer,
    509
    0
                        a_info.size,
    510
    0
                    )
    511
    0
                    .await?;
    512
                }
    513
    514
    0
                GpuOp::Dot { a, b, output } => {
    515
    0
                    let a_info = self.buffers.get(a).ok_or("Invalid buffer A ID")?;
    516
    0
                    let b_info = self.buffers.get(b).ok_or("Invalid buffer B ID")?;
    517
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    518
    519
    0
                    let a_buffer = a_info.gpu_buffer.as_ref().ok_or("Buffer A not created")?;
    520
    0
                    let b_buffer = b_info.gpu_buffer.as_ref().ok_or("Buffer B not created")?;
    521
    0
                    let output_buffer = output_info
    522
    0
                        .gpu_buffer
    523
    0
                        .as_ref()
    524
    0
                        .ok_or("Output buffer not created")?;
    525
    526
    0
                    self.execute_binary_op(
    527
    0
                        shaders::DOT_PRODUCT_SHADER,
    528
    0
                        "Dot",
    529
    0
                        a_buffer,
    530
    0
                        b_buffer,
    531
    0
                        output_buffer,
    532
    0
                        a_info.size,
    533
    0
                    )
    534
    0
                    .await?;
    535
                }
    536
    537
    0
                GpuOp::Sigmoid { input, output } => {
    538
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    539
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    540
    541
    0
                    let input_buffer = input_info
    542
    0
                        .gpu_buffer
    543
    0
                        .as_ref()
    544
    0
                        .ok_or("Input buffer not created")?;
    545
    0
                    let output_buffer = output_info
    546
    0
                        .gpu_buffer
    547
    0
                        .as_ref()
    548
    0
                        .ok_or("Output buffer not created")?;
    549
    550
    0
                    self.execute_unary_op::<()>(
    551
    0
                        shaders::SIGMOID_SHADER,
    552
    0
                        "Sigmoid",
    553
    0
                        input_buffer,
    554
    0
                        output_buffer,
    555
    0
                        input_info.size,
    556
    0
                        None,
    557
    0
                    )
    558
    0
                    .await?;
    559
                }
    560
    561
    0
                GpuOp::Tanh { input, output } => {
    562
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    563
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    564
    565
    0
                    let input_buffer = input_info
    566
    0
                        .gpu_buffer
    567
    0
                        .as_ref()
    568
    0
                        .ok_or("Input buffer not created")?;
    569
    0
                    let output_buffer = output_info
    570
    0
                        .gpu_buffer
    571
    0
                        .as_ref()
    572
    0
                        .ok_or("Output buffer not created")?;
    573
    574
    0
                    self.execute_unary_op::<()>(
    575
    0
                        shaders::TANH_SHADER,
    576
    0
                        "Tanh",
    577
    0
                        input_buffer,
    578
    0
                        output_buffer,
    579
    0
                        input_info.size,
    580
    0
                        None,
    581
    0
                    )
    582
    0
                    .await?;
    583
                }
    584
    585
    0
                GpuOp::Swish { input, output } => {
    586
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    587
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    588
    589
    0
                    let input_buffer = input_info
    590
    0
                        .gpu_buffer
    591
    0
                        .as_ref()
    592
    0
                        .ok_or("Input buffer not created")?;
    593
    0
                    let output_buffer = output_info
    594
    0
                        .gpu_buffer
    595
    0
                        .as_ref()
    596
    0
                        .ok_or("Output buffer not created")?;
    597
    598
    0
                    self.execute_unary_op::<()>(
    599
    0
                        shaders::SWISH_SHADER,
    600
    0
                        "Swish",
    601
    0
                        input_buffer,
    602
    0
                        output_buffer,
    603
    0
                        input_info.size,
    604
    0
                        None,
    605
    0
                    )
    606
    0
                    .await?;
    607
                }
    608
    609
    0
                GpuOp::Gelu { input, output } => {
    610
    0
                    let input_info = self.buffers.get(input).ok_or("Invalid input buffer ID")?;
    611
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    612
    613
    0
                    let input_buffer = input_info
    614
    0
                        .gpu_buffer
    615
    0
                        .as_ref()
    616
    0
                        .ok_or("Input buffer not created")?;
    617
    0
                    let output_buffer = output_info
    618
    0
                        .gpu_buffer
    619
    0
                        .as_ref()
    620
    0
                        .ok_or("Output buffer not created")?;
    621
    622
    0
                    self.execute_unary_op::<()>(
    623
    0
                        shaders::GELU_SHADER,
    624
    0
                        "GELU",
    625
    0
                        input_buffer,
    626
    0
                        output_buffer,
    627
    0
                        input_info.size,
    628
    0
                        None,
    629
    0
                    )
    630
    0
                    .await?;
    631
                }
    632
    633
    0
                GpuOp::Sub { a, b, output } => {
    634
    0
                    let a_info = self.buffers.get(a).ok_or("Invalid buffer A ID")?;
    635
    0
                    let b_info = self.buffers.get(b).ok_or("Invalid buffer B ID")?;
    636
    0
                    let output_info = self.buffers.get(output).ok_or("Invalid output buffer ID")?;
    637
    638
    0
                    let a_buffer = a_info.gpu_buffer.as_ref().ok_or("Buffer A not created")?;
    639
    0
                    let b_buffer = b_info.gpu_buffer.as_ref().ok_or("Buffer B not created")?;
    640
    0
                    let output_buffer = output_info
    641
    0
                        .gpu_buffer
    642
    0
                        .as_ref()
    643
    0
                        .ok_or("Output buffer not created")?;
    644
    645
    0
                    self.execute_binary_op(
    646
    0
                        shaders::VEC_SUB_SHADER,
    647
    0
                        "Sub",
    648
    0
                        a_buffer,
    649
    0
                        b_buffer,
    650
    0
                        output_buffer,
    651
    0
                        a_info.size,
    652
    0
                    )
    653
    0
                    .await?;
    654
                }
    655
            }
    656
    657
    0
            Ok(())
    658
    0
        }
    659
    660
        /// Execute a unary operation (one input, one output)
    661
    0
        async fn execute_unary_op<T: bytemuck::Pod>(
    662
    0
            &self,
    663
    0
            shader_source: &str,
    664
    0
            label: &str,
    665
    0
            input_buffer: &wgpu::Buffer,
    666
    0
            output_buffer: &wgpu::Buffer,
    667
    0
            size: usize,
    668
    0
            params: Option<&T>,
    669
    0
        ) -> Result<(), String> {
    670
            // Create shader module
    671
    0
            let shader = self
    672
    0
                .device
    673
    0
                .device
    674
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    675
    0
                    label: Some(&format!("{} Shader", label)),
    676
    0
                    source: wgpu::ShaderSource::Wgsl(shader_source.into()),
    677
    0
                });
    678
    679
            // Create bind group layout entries
    680
    0
            let mut layout_entries = vec![
    681
    0
                wgpu::BindGroupLayoutEntry {
    682
    0
                    binding: 0,
    683
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    684
    0
                    ty: wgpu::BindingType::Buffer {
    685
    0
                        ty: wgpu::BufferBindingType::Storage { read_only: true },
    686
    0
                        has_dynamic_offset: false,
    687
    0
                        min_binding_size: None,
    688
    0
                    },
    689
    0
                    count: None,
    690
    0
                },
    691
    0
                wgpu::BindGroupLayoutEntry {
    692
    0
                    binding: 1,
    693
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    694
    0
                    ty: wgpu::BindingType::Buffer {
    695
    0
                        ty: wgpu::BufferBindingType::Storage { read_only: false },
    696
    0
                        has_dynamic_offset: false,
    697
    0
                        min_binding_size: None,
    698
    0
                    },
    699
    0
                    count: None,
    700
    0
                },
    701
            ];
    702
    703
            // Add uniform binding if params provided
    704
    0
            if params.is_some() {
    705
    0
                layout_entries.push(wgpu::BindGroupLayoutEntry {
    706
    0
                    binding: 2,
    707
    0
                    visibility: wgpu::ShaderStages::COMPUTE,
    708
    0
                    ty: wgpu::BindingType::Buffer {
    709
    0
                        ty: wgpu::BufferBindingType::Uniform,
    710
    0
                        has_dynamic_offset: false,
    711
    0
                        min_binding_size: None,
    712
    0
                    },
    713
    0
                    count: None,
    714
    0
                });
    715
    0
            }
    716
    717
    0
            let bind_group_layout =
    718
    0
                self.device
    719
    0
                    .device
    720
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    721
    0
                        label: Some(&format!("{} Bind Group Layout", label)),
    722
    0
                        entries: &layout_entries,
    723
    0
                    });
    724
    725
            // Create uniform buffer if params provided (needs to live through bind group creation)
    726
    0
            let params_buffer = if let Some(params_data) = params {
    727
    0
                let buffer = self.device.device.create_buffer(&wgpu::BufferDescriptor {
    728
    0
                    label: Some(&format!("{} Params", label)),
    729
    0
                    size: std::mem::size_of::<T>() as u64,
    730
    0
                    usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    731
    0
                    mapped_at_creation: false,
    732
    0
                });
    733
    734
    0
                self.device
    735
    0
                    .queue
    736
    0
                    .write_buffer(&buffer, 0, bytemuck::bytes_of(params_data));
    737
    738
    0
                Some(buffer)
    739
            } else {
    740
    0
                None
    741
            };
    742
    743
            // Create bind group entries
    744
    0
            let mut bind_entries = vec![
    745
    0
                wgpu::BindGroupEntry {
    746
    0
                    binding: 0,
    747
    0
                    resource: input_buffer.as_entire_binding(),
    748
    0
                },
    749
    0
                wgpu::BindGroupEntry {
    750
    0
                    binding: 1,
    751
    0
                    resource: output_buffer.as_entire_binding(),
    752
    0
                },
    753
            ];
    754
    755
            // Add params binding if provided
    756
    0
            if let Some(ref buffer) = params_buffer {
    757
    0
                bind_entries.push(wgpu::BindGroupEntry {
    758
    0
                    binding: 2,
    759
    0
                    resource: buffer.as_entire_binding(),
    760
    0
                });
    761
    0
            }
    762
    763
    0
            let bind_group = self
    764
    0
                .device
    765
    0
                .device
    766
    0
                .create_bind_group(&wgpu::BindGroupDescriptor {
    767
    0
                    label: Some(&format!("{} Bind Group", label)),
    768
    0
                    layout: &bind_group_layout,
    769
    0
                    entries: &bind_entries,
    770
    0
                });
    771
    772
            // Create pipeline
    773
    0
            let pipeline_layout =
    774
    0
                self.device
    775
    0
                    .device
    776
    0
                    .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    777
    0
                        label: Some(&format!("{} Pipeline Layout", label)),
    778
    0
                        bind_group_layouts: &[&bind_group_layout],
    779
    0
                        push_constant_ranges: &[],
    780
    0
                    });
    781
    782
    0
            let pipeline =
    783
    0
                self.device
    784
    0
                    .device
    785
    0
                    .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    786
    0
                        label: Some(&format!("{} Pipeline", label)),
    787
    0
                        layout: Some(&pipeline_layout),
    788
    0
                        module: &shader,
    789
    0
                        entry_point: Some("main"),
    790
    0
                        compilation_options: Default::default(),
    791
    0
                        cache: None,
    792
    0
                    });
    793
    794
            // Execute
    795
    0
            let mut encoder =
    796
    0
                self.device
    797
    0
                    .device
    798
    0
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    799
    0
                        label: Some(&format!("{} Encoder", label)),
    800
    0
                    });
    801
    802
    0
            {
    803
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    804
    0
                    label: Some(&format!("{} Pass", label)),
    805
    0
                    timestamp_writes: None,
    806
    0
                });
    807
    0
    808
    0
                compute_pass.set_pipeline(&pipeline);
    809
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    810
    0
    811
    0
                // Dispatch workgroups (256 threads per workgroup)
    812
    0
                let workgroup_size = 256;
    813
    0
                let num_workgroups = (size as u32).div_ceil(workgroup_size);
    814
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    815
    0
            }
    816
    817
    0
            self.device.queue.submit(Some(encoder.finish()));
    818
    819
    0
            Ok(())
    820
    0
        }
    821
    822
        /// Execute a binary operation (two inputs, one output)
    823
    0
        async fn execute_binary_op(
    824
    0
            &self,
    825
    0
            shader_source: &str,
    826
    0
            label: &str,
    827
    0
            a_buffer: &wgpu::Buffer,
    828
    0
            b_buffer: &wgpu::Buffer,
    829
    0
            output_buffer: &wgpu::Buffer,
    830
    0
            size: usize,
    831
    0
        ) -> Result<(), String> {
    832
            // Create shader module
    833
    0
            let shader = self
    834
    0
                .device
    835
    0
                .device
    836
    0
                .create_shader_module(wgpu::ShaderModuleDescriptor {
    837
    0
                    label: Some(&format!("{} Shader", label)),
    838
    0
                    source: wgpu::ShaderSource::Wgsl(shader_source.into()),
    839
    0
                });
    840
    841
            // Create bind group layout
    842
    0
            let bind_group_layout =
    843
    0
                self.device
    844
    0
                    .device
    845
    0
                    .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
    846
    0
                        label: Some(&format!("{} Bind Group Layout", label)),
    847
    0
                        entries: &[
    848
    0
                            wgpu::BindGroupLayoutEntry {
    849
    0
                                binding: 0,
    850
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    851
    0
                                ty: wgpu::BindingType::Buffer {
    852
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    853
    0
                                    has_dynamic_offset: false,
    854
    0
                                    min_binding_size: None,
    855
    0
                                },
    856
    0
                                count: None,
    857
    0
                            },
    858
    0
                            wgpu::BindGroupLayoutEntry {
    859
    0
                                binding: 1,
    860
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    861
    0
                                ty: wgpu::BindingType::Buffer {
    862
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: true },
    863
    0
                                    has_dynamic_offset: false,
    864
    0
                                    min_binding_size: None,
    865
    0
                                },
    866
    0
                                count: None,
    867
    0
                            },
    868
    0
                            wgpu::BindGroupLayoutEntry {
    869
    0
                                binding: 2,
    870
    0
                                visibility: wgpu::ShaderStages::COMPUTE,
    871
    0
                                ty: wgpu::BindingType::Buffer {
    872
    0
                                    ty: wgpu::BufferBindingType::Storage { read_only: false },
    873
    0
                                    has_dynamic_offset: false,
    874
    0
                                    min_binding_size: None,
    875
    0
                                },
    876
    0
                                count: None,
    877
    0
                            },
    878
    0
                        ],
    879
    0
                    });
    880
    881
    0
            let bind_group = self
    882
    0
                .device
    883
    0
                .device
    884
    0
                .create_bind_group(&wgpu::BindGroupDescriptor {
    885
    0
                    label: Some(&format!("{} Bind Group", label)),
    886
    0
                    layout: &bind_group_layout,
    887
    0
                    entries: &[
    888
    0
                        wgpu::BindGroupEntry {
    889
    0
                            binding: 0,
    890
    0
                            resource: a_buffer.as_entire_binding(),
    891
    0
                        },
    892
    0
                        wgpu::BindGroupEntry {
    893
    0
                            binding: 1,
    894
    0
                            resource: b_buffer.as_entire_binding(),
    895
    0
                        },
    896
    0
                        wgpu::BindGroupEntry {
    897
    0
                            binding: 2,
    898
    0
                            resource: output_buffer.as_entire_binding(),
    899
    0
                        },
    900
    0
                    ],
    901
    0
                });
    902
    903
            // Create pipeline
    904
    0
            let pipeline_layout =
    905
    0
                self.device
    906
    0
                    .device
    907
    0
                    .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
    908
    0
                        label: Some(&format!("{} Pipeline Layout", label)),
    909
    0
                        bind_group_layouts: &[&bind_group_layout],
    910
    0
                        push_constant_ranges: &[],
    911
    0
                    });
    912
    913
    0
            let pipeline =
    914
    0
                self.device
    915
    0
                    .device
    916
    0
                    .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
    917
    0
                        label: Some(&format!("{} Pipeline", label)),
    918
    0
                        layout: Some(&pipeline_layout),
    919
    0
                        module: &shader,
    920
    0
                        entry_point: Some("main"),
    921
    0
                        compilation_options: Default::default(),
    922
    0
                        cache: None,
    923
    0
                    });
    924
    925
            // Execute
    926
    0
            let mut encoder =
    927
    0
                self.device
    928
    0
                    .device
    929
    0
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    930
    0
                        label: Some(&format!("{} Encoder", label)),
    931
    0
                    });
    932
    933
    0
            {
    934
    0
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
    935
    0
                    label: Some(&format!("{} Pass", label)),
    936
    0
                    timestamp_writes: None,
    937
    0
                });
    938
    0
    939
    0
                compute_pass.set_pipeline(&pipeline);
    940
    0
                compute_pass.set_bind_group(0, &bind_group, &[]);
    941
    0
    942
    0
                // Dispatch workgroups (256 threads per workgroup)
    943
    0
                let workgroup_size = 256;
    944
    0
                let num_workgroups = (size as u32).div_ceil(workgroup_size);
    945
    0
                compute_pass.dispatch_workgroups(num_workgroups, 1, 1);
    946
    0
            }
    947
    948
    0
            self.device.queue.submit(Some(encoder.finish()));
    949
    950
    0
            Ok(())
    951
    0
        }
    952
    953
        /// Read buffer data back from GPU
    954
        ///
    955
        /// Must call `execute()` first.
    956
    0
        pub async fn read(&self, buffer_id: BufferId) -> Result<Vec<f32>, String> {
    957
    0
            let buffer_info = self.buffers.get(&buffer_id).ok_or("Invalid buffer ID")?;
    958
    959
    0
            let gpu_buffer = buffer_info
    960
    0
                .gpu_buffer
    961
    0
                .as_ref()
    962
    0
                .ok_or("Buffer not executed yet - call execute() first")?;
    963
    964
    0
            let size_bytes = (buffer_info.size * std::mem::size_of::<f32>()) as u64;
    965
    966
            // Create staging buffer for reading
    967
    0
            let staging_buffer = self.device.device.create_buffer(&wgpu::BufferDescriptor {
    968
    0
                label: Some("Staging Buffer"),
    969
    0
                size: size_bytes,
    970
    0
                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
    971
    0
                mapped_at_creation: false,
    972
    0
            });
    973
    974
            // Copy from GPU buffer to staging buffer
    975
    0
            let mut encoder =
    976
    0
                self.device
    977
    0
                    .device
    978
    0
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
    979
    0
                        label: Some("Read Encoder"),
    980
    0
                    });
    981
    982
    0
            encoder.copy_buffer_to_buffer(gpu_buffer, 0, &staging_buffer, 0, size_bytes);
    983
    984
    0
            self.device.queue.submit(Some(encoder.finish()));
    985
    986
            // Map the staging buffer for reading
    987
    0
            let buffer_slice = staging_buffer.slice(..);
    988
    0
            let (sender, receiver) = futures_intrusive::channel::shared::oneshot_channel();
    989
    990
    0
            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
    991
    0
                sender.send(result).ok();
    992
    0
            });
    993
    994
            // Wait for mapping to complete
    995
    0
            receiver
    996
    0
                .receive()
    997
    0
                .await
    998
    0
                .ok_or("Failed to receive mapping result")?
    999
    0
                .map_err(|e| format!("Buffer mapping failed: {:?}", e))?;
    1000
    1001
            // Read data from mapped buffer
    1002
    0
            let data = {
    1003
    0
                let mapped_range = buffer_slice.get_mapped_range();
    1004
    0
                let float_data: &[f32] = bytemuck::cast_slice(&mapped_range);
    1005
    0
                float_data.to_vec()
    1006
            };
    1007
    1008
    0
            staging_buffer.unmap();
    1009
    1010
    0
            Ok(data)
    1011
    0
        }
    1012
    1013
        /// Get number of queued operations
    1014
    0
        pub fn num_operations(&self) -> usize {
    1015
    0
            self.operations.len()
    1016
    0
        }
    1017
    1018
        /// Get number of buffers
    1019
    0
        pub fn num_buffers(&self) -> usize {
    1020
    0
            self.buffers.len()
    1021
    0
        }
    1022
    }
    1023
    1024
    #[cfg(test)]
    1025
    mod tests {
    1026
        use super::*;
    1027
        use std::sync::OnceLock;
    1028
    1029
        /// Shared GPU device for fast test execution (initialized once)
    1030
        static SHARED_DEVICE: OnceLock<Option<GpuDevice>> = OnceLock::new();
    1031
    1032
        /// Get shared GPU device (fast) or None if unavailable
    1033
        fn get_shared_device() -> Option<GpuDevice> {
    1034
            SHARED_DEVICE
    1035
                .get_or_init(|| {
    1036
                    if GpuDevice::is_available() {
    1037
                        GpuDevice::new().ok()
    1038
                    } else {
    1039
                        None
    1040
                    }
    1041
                })
    1042
                .clone()
    1043
        }
    1044
    1045
        #[test]
    1046
        fn test_buffer_allocation() {
    1047
            let Some(device) = get_shared_device() else {
    1048
                eprintln!("GPU not available, skipping");
    1049
                return;
    1050
            };
    1051
            let mut batch = GpuCommandBatch::new(device);
    1052
    1053
            let buf1 = batch.upload(&[1.0, 2.0, 3.0]);
    1054
            let buf2 = batch.upload(&[4.0, 5.0, 6.0]);
    1055
    1056
            assert_eq!(batch.num_buffers(), 2);
    1057
            assert_ne!(buf1, buf2);
    1058
        }
    1059
    1060
        #[test]
    1061
        fn test_operation_queuing() {
    1062
            let Some(device) = get_shared_device() else {
    1063
                eprintln!("GPU not available, skipping");
    1064
                return;
    1065
            };
    1066
            let mut batch = GpuCommandBatch::new(device);
    1067
    1068
            let input = batch.upload(&[1.0, 2.0, -3.0, 4.0]);
    1069
            let relu_out = batch.relu(input);
    1070
            let scaled = batch.scale(relu_out, 2.0);
    1071
            let other = batch.upload(&[0.5, 0.5, 0.5, 0.5]);
    1072
            let _final_out = batch.add(scaled, other);
    1073
    1074
            assert_eq!(batch.num_operations(), 3); // relu, scale, add
    1075
            assert_eq!(batch.num_buffers(), 5); // input, relu_out, scaled, other, final_out
    1076
        }
    1077
    1078
        #[test]
    1079
        #[should_panic(expected = "Buffer size mismatch")]
    1080
        fn test_size_mismatch_add() {
    1081
            let Some(device) = get_shared_device() else {
    1082
                panic!("Buffer size mismatch"); // Satisfy should_panic when skipping
    1083
            };
    1084
            let mut batch = GpuCommandBatch::new(device);
    1085
    1086
            let a = batch.upload(&[1.0, 2.0]);
    1087
            let b = batch.upload(&[1.0, 2.0, 3.0]);
    1088
            batch.add(a, b); // Should panic
    1089
        }
    1090
    1091
        #[test]
    1092
        #[should_panic(expected = "Buffer size mismatch")]
    1093
        fn test_size_mismatch_mul() {
    1094
            let Some(device) = get_shared_device() else {
    1095
                panic!("Buffer size mismatch"); // Satisfy should_panic when skipping
    1096
            };
    1097
            let mut batch = GpuCommandBatch::new(device);
    1098
    1099
            let a = batch.upload(&[1.0, 2.0]);
    1100
            let b = batch.upload(&[1.0, 2.0, 3.0]);
    1101
            batch.mul(a, b); // Should panic
    1102
        }
    1103
    1104
        #[test]
    1105
        #[should_panic(expected = "Buffer size mismatch")]
    1106
        fn test_size_mismatch_dot() {
    1107
            let Some(device) = get_shared_device() else {
    1108
                panic!("Buffer size mismatch"); // Satisfy should_panic when skipping
    1109
            };
    1110
            let mut batch = GpuCommandBatch::new(device);
    1111
    1112
            let a = batch.upload(&[1.0, 2.0]);
    1113
            let b = batch.upload(&[1.0, 2.0, 3.0]);
    1114
            batch.dot(a, b); // Should panic
    1115
        }
    1116
    1117
        /// Comprehensive async test covering ALL batch operations in a single GPU session.
    1118
        /// This reduces GPU initialization overhead for coverage (1 session vs 10).
    1119
        #[tokio::test]
    1120
        async fn test_all_batch_operations() {
    1121
            let Some(device) = get_shared_device() else {
    1122
                eprintln!("GPU not available, skipping");
    1123
                return;
    1124
            };
    1125
            let mut batch = GpuCommandBatch::new(device);
    1126
    1127
            // Test 1: End-to-end (relu + scale + add)
    1128
            let input1 = batch.upload(&[1.0, 2.0, -3.0, 4.0]);
    1129
            let relu_out = batch.relu(input1);
    1130
            let scaled = batch.scale(relu_out, 2.0);
    1131
            let other = batch.upload(&[0.5, 0.5, 0.5, 0.5]);
    1132
            let add_result = batch.add(scaled, other);
    1133
    1134
            // Test 2: Mul operation
    1135
            let mul_a = batch.upload(&[1.0, 2.0, 3.0, 4.0]);
    1136
            let mul_b = batch.upload(&[2.0, 3.0, 4.0, 5.0]);
    1137
            let mul_result = batch.mul(mul_a, mul_b);
    1138
    1139
            // Test 3: Dot operation
    1140
            let dot_a = batch.upload(&[1.0, 2.0, 3.0, 4.0]);
    1141
            let dot_b = batch.upload(&[2.0, 3.0, 4.0, 5.0]);
    1142
            let dot_result = batch.dot(dot_a, dot_b);
    1143
    1144
            // Test 4: Sigmoid
    1145
            let sig_input = batch.upload(&[-2.0, -1.0, 0.0, 1.0, 2.0]);
    1146
            let sig_result = batch.sigmoid(sig_input);
    1147
    1148
            // Test 5: Tanh
    1149
            let tanh_input = batch.upload(&[-1.0, 0.0, 1.0]);
    1150
            let tanh_result = batch.tanh(tanh_input);
    1151
    1152
            // Test 6: Swish
    1153
            let swish_input = batch.upload(&[0.0, 1.0, 2.0]);
    1154
            let swish_result = batch.swish(swish_input);
    1155
    1156
            // Test 7: GELU
    1157
            let gelu_input = batch.upload(&[-1.0, 0.0, 1.0]);
    1158
            let gelu_result = batch.gelu(gelu_input);
    1159
    1160
            // Test 8: Sub
    1161
            let sub_a = batch.upload(&[5.0, 10.0, 15.0, 20.0]);
    1162
            let sub_b = batch.upload(&[1.0, 2.0, 3.0, 4.0]);
    1163
            let sub_result = batch.sub(sub_a, sub_b);
    1164
    1165
            // Test 9: Chained activations
    1166
            let chain_input = batch.upload(&[-2.0, -1.0, 0.0, 1.0, 2.0]);
    1167
            let chain_relu = batch.relu(chain_input);
    1168
            let chain_sigmoid = batch.sigmoid(chain_relu);
    1169
            let chain_result = batch.tanh(chain_sigmoid);
    1170
    1171
            // Execute all operations in single batch
    1172
            batch.execute().await.unwrap();
    1173
    1174
            // Verify Test 1: relu([1,2,-3,4])=[1,2,0,4] → scale(*2)=[2,4,0,8] → add([0.5])=[2.5,4.5,0.5,8.5]
    1175
            let result1 = batch.read(add_result).await.unwrap();
    1176
            assert_eq!(result1.len(), 4);
    1177
            assert!((result1[0] - 2.5).abs() < 1e-5);
    1178
            assert!((result1[1] - 4.5).abs() < 1e-5);
    1179
            assert!((result1[2] - 0.5).abs() < 1e-5);
    1180
            assert!((result1[3] - 8.5).abs() < 1e-5);
    1181
    1182
            // Verify Test 2: [1*2, 2*3, 3*4, 4*5] = [2, 6, 12, 20]
    1183
            let result2 = batch.read(mul_result).await.unwrap();
    1184
            assert_eq!(result2, vec![2.0, 6.0, 12.0, 20.0]);
    1185
    1186
            // Verify Test 3: Dot product returns a result
    1187
            let result3 = batch.read(dot_result).await.unwrap();
    1188
            assert!(!result3.is_empty());
    1189
    1190
            // Verify Test 4: Sigmoid values
    1191
            let result4 = batch.read(sig_result).await.unwrap();
    1192
            assert_eq!(result4.len(), 5);
    1193
            assert!((result4[0] - 0.119).abs() < 0.01); // sigmoid(-2)
    1194
            assert!((result4[2] - 0.5).abs() < 0.01); // sigmoid(0)
    1195
            assert!((result4[4] - 0.881).abs() < 0.01); // sigmoid(2)
    1196
    1197
            // Verify Test 5: Tanh values
    1198
            let result5 = batch.read(tanh_result).await.unwrap();
    1199
            assert_eq!(result5.len(), 3);
    1200
            assert!((result5[0] - (-0.762)).abs() < 0.01);
    1201
            assert!(result5[1].abs() < 0.01);
    1202
            assert!((result5[2] - 0.762).abs() < 0.01);
    1203
    1204
            // Verify Test 6: Swish values
    1205
            let result6 = batch.read(swish_result).await.unwrap();
    1206
            assert_eq!(result6.len(), 3);
    1207
            assert!(result6[0].abs() < 0.01);
    1208
            assert!((result6[1] - 0.731).abs() < 0.01);
    1209
    1210
            // Verify Test 7: GELU values
    1211
            let result7 = batch.read(gelu_result).await.unwrap();
    1212
            assert_eq!(result7.len(), 3);
    1213
            assert!(result7[1].abs() < 0.01);
    1214
            assert!((result7[2] - 0.841).abs() < 0.05);
    1215
    1216
            // Verify Test 8: Sub [5-1, 10-2, 15-3, 20-4] = [4, 8, 12, 16]
    1217
            let result8 = batch.read(sub_result).await.unwrap();
    1218
            assert_eq!(result8, vec![4.0, 8.0, 12.0, 16.0]);
    1219
    1220
            // Verify Test 9: Chained activations in range
    1221
            let result9 = batch.read(chain_result).await.unwrap();
    1222
            assert_eq!(result9.len(), 5);
    1223
            for &val in &result9 {
    1224
                assert!((-1.0..=1.0).contains(&val), "Value {} out of range", val);
    1225
            }
    1226
        }
    1227
    }