dspk_ops/core/kernels/transpose.rs
1//! A transpose between buffers with padded rows.
2
3use cubecl::prelude::*;
4use dspk_core::compute::element_position;
5
6/// One unit per element `(r, c)` of a `[rows, cols]` input whose rows are `in_row_stride` apart:
7/// `output[c · out_row_stride + r] = input[r · in_row_stride + c]`. The padding of either buffer
8/// is neither read nor written.
9#[cube(launch)]
10pub fn transpose_strided_kernel<F: Float>(input: &[F], output: &mut [F], rows: u32, cols: u32, in_row_stride: u32, out_row_stride: u32) {
11 let e = element_position() as u32;
12 if e < rows * cols {
13 let r = e / cols;
14 let c = e % cols;
15 output[(c * out_row_stride + r) as usize] = input[(r * in_row_stride + c) as usize];
16 }
17}