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

    Line
    Count
    Source
    1
    //! PartitionView - Tiling Strategy for GPU Compute
    2
    //!
    3
    //! Divides a TensorView into tiles for efficient GPU processing.
    4
    //! Enables automatic work distribution across thread blocks.
    5
    //!
    6
    //! # cuda-tile-behavior.md References
    7
    //!
    8
    //! - Section 3.2: Two-Level Memory Hierarchy
    9
    //! - Falsification tests #36-45: PartitionView correctness
    10
    //!
    11
    //! # Academic Foundation
    12
    //!
    13
    //! Based on Volkov & Demmel (2008): Power-of-two tiles achieve 95%+ peak throughput.
    14
    15
    use super::tensor_view::TensorView;
    16
    use std::marker::PhantomData;
    17
    18
    /// A tiling strategy over a TensorView.
    19
    ///
    20
    /// PartitionView divides a tensor into tiles of a specified shape,
    21
    /// enabling efficient GPU processing with shared memory optimization.
    22
    ///
    23
    /// # Type Parameters
    24
    ///
    25
    /// * `T` - Element type of the underlying tensor
    26
    ///
    27
    /// # cuda-tile-behavior.md References
    28
    ///
    29
    /// - Falsification test #36: Tile count calculation is correct
    30
    /// - Falsification test #37: Tile iteration covers all elements
    31
    /// - Falsification test #38: Edge tiles are handled correctly
    32
    #[derive(Debug)]
    33
    pub struct PartitionView<T> {
    34
        /// The underlying tensor being partitioned
    35
        tensor: TensorView<T>,
    36
        /// Shape of each tile
    37
        tile_shape: [usize; 4],
    38
        /// Phantom data for type safety
    39
        _marker: PhantomData<T>,
    40
    }
    41
    42
    /// Information about a single tile within a partition.
    43
    #[derive(Debug, Clone, PartialEq, Eq)]
    44
    pub struct TileInfo {
    45
        /// Tile index in each dimension
    46
        pub tile_idx: [usize; 4],
    47
        /// Starting element index in each dimension
    48
        pub start: [usize; 4],
    49
        /// Size of this tile in each dimension (may be smaller at edges)
    50
        pub size: [usize; 4],
    51
        /// Whether this is an edge tile (smaller than full tile size)
    52
        pub is_edge: bool,
    53
    }
    54
    55
    impl<T> PartitionView<T> {
    56
        /// Create a new PartitionView with the given tile shape.
    57
        ///
    58
        /// # Arguments
    59
        ///
    60
        /// * `tensor` - The tensor to partition
    61
        /// * `tile_shape` - Shape of each tile
    62
        ///
    63
        /// # Panics
    64
        ///
    65
        /// Panics if any tile dimension is zero.
    66
        ///
    67
        /// # cuda-tile-behavior.md References
    68
        ///
    69
        /// - Falsification test #36: Tile count calculation is correct
    70
    0
        pub fn new(tensor: TensorView<T>, tile_shape: [usize; 4]) -> Self {
    71
    0
            assert!(
    72
    0
                tile_shape.iter().all(|&d| d > 0),
    73
    0
                "Tile dimensions must be non-zero"
    74
            );
    75
    0
            Self {
    76
    0
                tensor,
    77
    0
                tile_shape,
    78
    0
                _marker: PhantomData,
    79
    0
            }
    80
    0
        }
    81
    82
        /// Create a PartitionView with power-of-two tile sizes.
    83
        ///
    84
        /// This is recommended for GPU compute as it enables efficient
    85
        /// memory coalescing and avoids bank conflicts.
    86
        ///
    87
        /// # Arguments
    88
        ///
    89
        /// * `tensor` - The tensor to partition
    90
        /// * `tile_log2` - Log2 of tile size for each dimension
    91
        ///
    92
        /// # cuda-tile-behavior.md References
    93
        ///
    94
        /// - Falsification test #1: Power-of-two tiles improve GPU occupancy
    95
    0
        pub fn new_power_of_two(tensor: TensorView<T>, tile_log2: [usize; 4]) -> Self {
    96
    0
            let tile_shape = [
    97
    0
                1 << tile_log2[0],
    98
    0
                1 << tile_log2[1],
    99
    0
                1 << tile_log2[2],
    100
    0
                1 << tile_log2[3],
    101
    0
            ];
    102
    0
            Self::new(tensor, tile_shape)
    103
    0
        }
    104
    105
        /// Create a PartitionView with 2D tiles (for matrix operations).
    106
        ///
    107
        /// # Arguments
    108
        ///
    109
        /// * `tensor` - The tensor to partition
    110
        /// * `tile_rows` - Number of rows per tile
    111
        /// * `tile_cols` - Number of columns per tile
    112
    0
        pub fn new_2d(tensor: TensorView<T>, tile_rows: usize, tile_cols: usize) -> Self {
    113
    0
            Self::new(tensor, [tile_rows, tile_cols, 1, 1])
    114
    0
        }
    115
    116
        /// Get the underlying tensor.
    117
    0
        pub fn tensor(&self) -> &TensorView<T> {
    118
    0
            &self.tensor
    119
    0
        }
    120
    121
        /// Get the tile shape.
    122
    0
        pub fn tile_shape(&self) -> &[usize; 4] {
    123
    0
            &self.tile_shape
    124
    0
        }
    125
    126
        /// Get the number of tiles in each dimension.
    127
        ///
    128
        /// # cuda-tile-behavior.md References
    129
        ///
    130
        /// - Falsification test #36: Tile count calculation is correct
    131
    0
        pub fn tile_count(&self) -> [usize; 4] {
    132
    0
            let tensor_shape = self.tensor.shape();
    133
    0
            [
    134
    0
                tensor_shape[0].div_ceil(self.tile_shape[0]),
    135
    0
                tensor_shape[1].div_ceil(self.tile_shape[1]),
    136
    0
                tensor_shape[2].div_ceil(self.tile_shape[2]),
    137
    0
                tensor_shape[3].div_ceil(self.tile_shape[3]),
    138
    0
            ]
    139
    0
        }
    140
    141
        /// Get the total number of tiles.
    142
    0
        pub fn total_tiles(&self) -> usize {
    143
    0
            let count = self.tile_count();
    144
    0
            count.iter().product()
    145
    0
        }
    146
    147
        /// Get information about a specific tile.
    148
        ///
    149
        /// # Arguments
    150
        ///
    151
        /// * `tile_idx` - Index of the tile in each dimension
    152
        ///
    153
        /// # Returns
    154
        ///
    155
        /// TileInfo containing the tile's position and size.
    156
        ///
    157
        /// # cuda-tile-behavior.md References
    158
        ///
    159
        /// - Falsification test #38: Edge tiles are handled correctly
    160
    0
        pub fn get_tile(&self, tile_idx: [usize; 4]) -> Option<TileInfo> {
    161
    0
            let tile_count = self.tile_count();
    162
    163
            // Validate tile index
    164
    0
            for i in 0..4 {
    165
    0
                if tile_idx[i] >= tile_count[i] {
    166
    0
                    return None;
    167
    0
                }
    168
            }
    169
    170
    0
            let tensor_shape = self.tensor.shape();
    171
    0
            let mut start = [0usize; 4];
    172
    0
            let mut size = [0usize; 4];
    173
    0
            let mut is_edge = false;
    174
    175
    0
            for i in 0..4 {
    176
    0
                start[i] = tile_idx[i] * self.tile_shape[i];
    177
    0
                let remaining = tensor_shape[i] - start[i];
    178
    0
                size[i] = remaining.min(self.tile_shape[i]);
    179
    180
                // Check if this is an edge tile
    181
    0
                if size[i] < self.tile_shape[i] {
    182
    0
                    is_edge = true;
    183
    0
                }
    184
            }
    185
    186
    0
            Some(TileInfo {
    187
    0
                tile_idx,
    188
    0
                start,
    189
    0
                size,
    190
    0
                is_edge,
    191
    0
            })
    192
    0
        }
    193
    194
        /// Get a TensorView for a specific tile.
    195
        ///
    196
        /// # Arguments
    197
        ///
    198
        /// * `tile_idx` - Index of the tile in each dimension
    199
        ///
    200
        /// # Returns
    201
        ///
    202
        /// A TensorView representing the tile, or None if index is invalid.
    203
    0
        pub fn get_tile_view(&self, tile_idx: [usize; 4]) -> Option<TensorView<T>> {
    204
    0
            let info = self.get_tile(tile_idx)?;
    205
    206
            // Create a sliced view for this tile
    207
    0
            let mut view = self.tensor.clone();
    208
    209
    0
            for i in 0..4 {
    210
    0
                if self.tensor.shape()[i] > 1 {
    211
    0
                    view = view.slice_dim(i, info.start[i]..info.start[i] + info.size[i]);
    212
    0
                }
    213
            }
    214
    215
    0
            Some(view)
    216
    0
        }
    217
    218
        /// Iterate over all tiles.
    219
        ///
    220
        /// # cuda-tile-behavior.md References
    221
        ///
    222
        /// - Falsification test #37: Tile iteration covers all elements
    223
    0
        pub fn iter_tiles(&self) -> TileIterator<'_, T> {
    224
    0
            TileIterator {
    225
    0
                partition: self,
    226
    0
                current: [0, 0, 0, 0],
    227
    0
                done: false,
    228
    0
            }
    229
    0
        }
    230
    231
        /// Check if tiles are power-of-two sized.
    232
        ///
    233
        /// Power-of-two tiles are preferred for GPU compute.
    234
    0
        pub fn is_power_of_two_tiles(&self) -> bool {
    235
    0
            self.tile_shape.iter().all(|&d| d.is_power_of_two())
    236
    0
        }
    237
    238
        /// Get the number of elements per tile (maximum).
    239
    0
        pub fn elements_per_tile(&self) -> usize {
    240
    0
            self.tile_shape.iter().product()
    241
    0
        }
    242
    243
        /// Get recommended workgroup size for GPU dispatch.
    244
        ///
    245
        /// Returns (x, y, z) workgroup dimensions based on tile shape.
    246
    0
        pub fn recommended_workgroup_size(&self) -> (u32, u32, u32) {
    247
            // Common GPU workgroup limits
    248
            const MAX_WORKGROUP_SIZE: usize = 256;
    249
            const MAX_DIM: usize = 16;
    250
    251
    0
            let tile_2d = [self.tile_shape[0], self.tile_shape[1]];
    252
    253
            // For 2D tiles, use 2D workgroups
    254
    0
            if tile_2d[0] > 1 && tile_2d[1] > 1 {
    255
    0
                let x = tile_2d[1].min(MAX_DIM) as u32;
    256
    0
                let y = tile_2d[0].min(MAX_DIM) as u32;
    257
    0
                let z = 1;
    258
    0
                (x, y, z)
    259
            } else {
    260
                // 1D workgroup
    261
    0
                let size = self.elements_per_tile().min(MAX_WORKGROUP_SIZE);
    262
    0
                (size as u32, 1, 1)
    263
            }
    264
    0
        }
    265
    }
    266
    267
    impl<T> Clone for PartitionView<T> {
    268
    0
        fn clone(&self) -> Self {
    269
    0
            Self {
    270
    0
                tensor: self.tensor.clone(),
    271
    0
                tile_shape: self.tile_shape,
    272
    0
                _marker: PhantomData,
    273
    0
            }
    274
    0
        }
    275
    }
    276
    277
    /// Iterator over tiles in a PartitionView.
    278
    pub struct TileIterator<'a, T> {
    279
        partition: &'a PartitionView<T>,
    280
        current: [usize; 4],
    281
        done: bool,
    282
    }
    283
    284
    impl<T> Iterator for TileIterator<'_, T> {
    285
        type Item = TileInfo;
    286
    287
    0
        fn next(&mut self) -> Option<Self::Item> {
    288
    0
            if self.done {
    289
    0
                return None;
    290
    0
            }
    291
    292
    0
            let tile = self.partition.get_tile(self.current)?;
    293
    0
            let tile_count = self.partition.tile_count();
    294
    295
            // Advance to next tile (row-major order)
    296
    0
            self.current[3] += 1;
    297
    0
            for i in (0..4).rev() {
    298
    0
                if self.current[i] >= tile_count[i] {
    299
    0
                    self.current[i] = 0;
    300
    0
                    if i > 0 {
    301
    0
                        self.current[i - 1] += 1;
    302
    0
                    } else {
    303
    0
                        self.done = true;
    304
    0
                    }
    305
                } else {
    306
    0
                    break;
    307
                }
    308
            }
    309
    310
    0
            Some(tile)
    311
    0
        }
    312
    313
    0
        fn size_hint(&self) -> (usize, Option<usize>) {
    314
    0
            let total = self.partition.total_tiles();
    315
    0
            (total, Some(total))
    316
    0
        }
    317
    }
    318
    319
    impl<T> ExactSizeIterator for TileIterator<'_, T> {}
    320
    321
    #[cfg(test)]
    322
    mod tests {
    323
        use super::*;
    324
    325
        // cuda-tile-behavior.md: Falsification test #36
    326
        #[test]
    327
        fn test_tile_count_exact_fit() {
    328
            let tensor = TensorView::<f32>::new([16, 32, 1, 1]);
    329
            let partition = PartitionView::new(tensor, [4, 8, 1, 1]);
    330
    331
            assert_eq!(partition.tile_count(), [4, 4, 1, 1]);
    332
            assert_eq!(partition.total_tiles(), 16);
    333
        }
    334
    335
        #[test]
    336
        fn test_tile_count_with_remainder() {
    337
            let tensor = TensorView::<f32>::new([17, 33, 1, 1]);
    338
            let partition = PartitionView::new(tensor, [4, 8, 1, 1]);
    339
    340
            // 17/4 = 5 (rounded up), 33/8 = 5 (rounded up)
    341
            assert_eq!(partition.tile_count(), [5, 5, 1, 1]);
    342
            assert_eq!(partition.total_tiles(), 25);
    343
        }
    344
    345
        // cuda-tile-behavior.md: Falsification test #37
    346
        #[test]
    347
        fn test_tile_iteration_covers_all() {
    348
            let tensor = TensorView::<f32>::new([8, 8, 1, 1]);
    349
            let partition = PartitionView::new(tensor, [4, 4, 1, 1]);
    350
    351
            let tiles: Vec<_> = partition.iter_tiles().collect();
    352
            assert_eq!(tiles.len(), 4);
    353
    354
            // Verify all tiles
    355
            assert_eq!(tiles[0].tile_idx, [0, 0, 0, 0]);
    356
            assert_eq!(tiles[1].tile_idx, [0, 1, 0, 0]);
    357
            assert_eq!(tiles[2].tile_idx, [1, 0, 0, 0]);
    358
            assert_eq!(tiles[3].tile_idx, [1, 1, 0, 0]);
    359
        }
    360
    361
        // cuda-tile-behavior.md: Falsification test #38
    362
        #[test]
    363
        fn test_edge_tiles() {
    364
            let tensor = TensorView::<f32>::new([10, 10, 1, 1]);
    365
            let partition = PartitionView::new(tensor, [8, 8, 1, 1]);
    366
    367
            // First tile: full size
    368
            let tile_0 = partition.get_tile([0, 0, 0, 0]).unwrap();
    369
            assert_eq!(tile_0.size, [8, 8, 1, 1]);
    370
            assert!(!tile_0.is_edge);
    371
    372
            // Edge tile: partial size
    373
            let tile_1 = partition.get_tile([1, 1, 0, 0]).unwrap();
    374
            assert_eq!(tile_1.size, [2, 2, 1, 1]); // 10 - 8 = 2
    375
            assert!(tile_1.is_edge);
    376
        }
    377
    378
        #[test]
    379
        fn test_get_tile_view() {
    380
            let tensor = TensorView::<f32>::new([16, 16, 1, 1]);
    381
            let partition = PartitionView::new(tensor, [8, 8, 1, 1]);
    382
    383
            let tile_view = partition.get_tile_view([1, 1, 0, 0]).unwrap();
    384
            assert_eq!(tile_view.shape()[0], 8);
    385
            assert_eq!(tile_view.shape()[1], 8);
    386
            assert_eq!(tile_view.offset(), 8 * 16 + 8); // Row 8, Col 8
    387
        }
    388
    389
        #[test]
    390
        fn test_power_of_two_tiles() {
    391
            let tensor = TensorView::<f32>::new([256, 256, 1, 1]);
    392
            let partition = PartitionView::new_power_of_two(tensor, [4, 4, 0, 0]);
    393
    394
            assert_eq!(partition.tile_shape(), &[16, 16, 1, 1]);
    395
            assert!(partition.is_power_of_two_tiles());
    396
        }
    397
    398
        #[test]
    399
        fn test_non_power_of_two_detection() {
    400
            let tensor = TensorView::<f32>::new([100, 100, 1, 1]);
    401
            let partition = PartitionView::new(tensor, [12, 12, 1, 1]);
    402
    403
            assert!(!partition.is_power_of_two_tiles());
    404
        }
    405
    406
        #[test]
    407
        fn test_2d_partition() {
    408
            let tensor = TensorView::<f32>::new_2d(100, 200);
    409
            let partition = PartitionView::new_2d(tensor, 16, 32);
    410
    411
            assert_eq!(partition.tile_shape(), &[16, 32, 1, 1]);
    412
            assert_eq!(partition.tile_count(), [7, 7, 1, 1]); // ceil(100/16), ceil(200/32)
    413
        }
    414
    415
        #[test]
    416
        fn test_elements_per_tile() {
    417
            let tensor = TensorView::<f32>::new([64, 64, 1, 1]);
    418
            let partition = PartitionView::new(tensor, [8, 8, 1, 1]);
    419
    420
            assert_eq!(partition.elements_per_tile(), 64);
    421
        }
    422
    423
        #[test]
    424
        fn test_workgroup_size_2d() {
    425
            let tensor = TensorView::<f32>::new([64, 64, 1, 1]);
    426
            let partition = PartitionView::new(tensor, [16, 16, 1, 1]);
    427
    428
            let (x, y, z) = partition.recommended_workgroup_size();
    429
            assert_eq!((x, y, z), (16, 16, 1));
    430
        }
    431
    432
        #[test]
    433
        fn test_workgroup_size_1d() {
    434
            let tensor = TensorView::<f32>::new_1d(1024);
    435
            let partition = PartitionView::new(tensor, [256, 1, 1, 1]);
    436
    437
            let (x, y, z) = partition.recommended_workgroup_size();
    438
            assert_eq!((x, y, z), (256, 1, 1));
    439
        }
    440
    441
        #[test]
    442
        fn test_invalid_tile_index() {
    443
            let tensor = TensorView::<f32>::new([8, 8, 1, 1]);
    444
            let partition = PartitionView::new(tensor, [4, 4, 1, 1]);
    445
    446
            assert!(partition.get_tile([5, 0, 0, 0]).is_none());
    447
            assert!(partition.get_tile([0, 5, 0, 0]).is_none());
    448
        }
    449
    450
        #[test]
    451
        fn test_iterator_size_hint() {
    452
            let tensor = TensorView::<f32>::new([16, 16, 1, 1]);
    453
            let partition = PartitionView::new(tensor, [4, 4, 1, 1]);
    454
    455
            let iter = partition.iter_tiles();
    456
            assert_eq!(iter.size_hint(), (16, Some(16)));
    457
            assert_eq!(iter.len(), 16);
    458
        }
    459
    460
        #[test]
    461
        fn test_tile_info_start_positions() {
    462
            let tensor = TensorView::<f32>::new([20, 20, 1, 1]);
    463
            let partition = PartitionView::new(tensor, [8, 8, 1, 1]);
    464
    465
            let tile_00 = partition.get_tile([0, 0, 0, 0]).unwrap();
    466
            assert_eq!(tile_00.start, [0, 0, 0, 0]);
    467
    468
            let tile_11 = partition.get_tile([1, 1, 0, 0]).unwrap();
    469
            assert_eq!(tile_11.start, [8, 8, 0, 0]);
    470
    471
            let tile_22 = partition.get_tile([2, 2, 0, 0]).unwrap();
    472
            assert_eq!(tile_22.start, [16, 16, 0, 0]);
    473
        }
    474
    475
        // cuda-tile-behavior.md: Falsification test #39 - Tile coverage completeness
    476
        #[test]
    477
        fn test_complete_coverage() {
    478
            let tensor = TensorView::<f32>::new([15, 17, 1, 1]);
    479
            let partition = PartitionView::new(tensor, [4, 4, 1, 1]);
    480
    481
            // Count all elements covered by tiles
    482
            let mut total_elements = 0;
    483
            for tile in partition.iter_tiles() {
    484
                total_elements += tile.size[0] * tile.size[1];
    485
            }
    486
    487
            assert_eq!(total_elements, 15 * 17);
    488
        }
    489
    490
        // cuda-tile-behavior.md: Falsification test #40 - Clone behavior
    491
        #[test]
    492
        fn test_partition_clone() {
    493
            let tensor = TensorView::<f32>::new([32, 32, 1, 1]);
    494
            let partition = PartitionView::new(tensor, [8, 8, 1, 1]);
    495
            let cloned = partition.clone();
    496
    497
            assert_eq!(partition.tile_shape(), cloned.tile_shape());
    498
            assert_eq!(partition.tile_count(), cloned.tile_count());
    499
        }
    500
    501
        #[test]
    502
        #[should_panic(expected = "Tile dimensions must be non-zero")]
    503
        fn test_zero_tile_dimension_panics() {
    504
            let tensor = TensorView::<f32>::new([16, 16, 1, 1]);
    505
            let _partition = PartitionView::new(tensor, [0, 8, 1, 1]);
    506
        }
    507
    508
        #[test]
    509
        fn test_single_tile() {
    510
            let tensor = TensorView::<f32>::new([8, 8, 1, 1]);
    511
            let partition = PartitionView::new(tensor, [16, 16, 1, 1]); // Tile larger than tensor
    512
    513
            assert_eq!(partition.tile_count(), [1, 1, 1, 1]);
    514
            assert_eq!(partition.total_tiles(), 1);
    515
    516
            let tile = partition.get_tile([0, 0, 0, 0]).unwrap();
    517
            assert_eq!(tile.size, [8, 8, 1, 1]); // Clamped to tensor size
    518
            assert!(tile.is_edge); // Smaller than full tile
    519
        }
    520
    }