cubek-fft 0.2.0

CubeK: Fast Fourier Transform
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

use cubecl::prelude::*;
use cubecl::{
    self,
    std::tensor::layout::{Coords1d, Layout, LayoutExpand},
};

#[derive(CubeType, Clone, Copy)]
/// Allows to work on the last dimension of the signal/spectrum (one window),
/// abstracting batches
pub struct BatchSignalLayout {
    num_samples: usize,
    stride_samples: usize,
    batch_offset: usize,
    #[cube(comptime)]
    vector_size: usize,
}

#[cube]
impl BatchSignalLayout {
    pub fn new<F: Numeric>(tensor: &Tensor<F>, batch_index: usize, #[comptime] dim: usize) -> Self {
        let rank = tensor.rank();
        let mut batch_offset = 0;
        let mut temp_idx = batch_index;

        for i in 0..rank {
            if i != dim {
                let size = tensor.shape(i);
                let stride = tensor.stride(i);

                let coord = temp_idx % size;
                batch_offset += coord * stride;
                temp_idx /= size;
            }
        }

        BatchSignalLayout {
            num_samples: tensor.shape(dim),
            stride_samples: tensor.stride(dim),
            batch_offset,
            vector_size: 1usize,
        }
    }
}

#[cube]
impl Layout for BatchSignalLayout {
    type Coordinates = Coords1d;
    type SourceCoordinates = Coords1d;

    fn to_source_pos(&self, coords: Self::Coordinates) -> usize {
        (self.batch_offset + coords * self.stride_samples) / self.vector_size
    }

    fn to_source_pos_checked(&self, coords: Self::Coordinates) -> (usize, bool) {
        (self.to_source_pos(coords), self.is_in_bounds(coords))
    }

    fn shape(&self) -> Self::Coordinates {
        self.num_samples
    }

    fn is_in_bounds(&self, pos: Self::Coordinates) -> bool {
        pos < self.num_samples
    }
}