cubecl_std/tensor/contiguous/
launch.rs

1use crate::tensor::{TensorHandle, copy_gpu_ref, try_launch_copy_perpendicular};
2use cubecl_core::{client::Client, ir::ElemType, prelude::TensorBinding};
3
4/// Make a jit tensor contiguous.
5pub fn into_contiguous(client: &Client, input: TensorBinding, dtype: ElemType) -> TensorHandle {
6    let num_elems: usize = input.shape.iter().product();
7
8    let handle = client.empty(num_elems * dtype.size());
9    let output = TensorHandle::new_contiguous(input.shape.to_vec(), handle, dtype);
10
11    copy_into(client, input, output.clone().binding(), dtype);
12
13    output
14}
15
16/// Make a jit tensor contiguous, using the pitched allocator if available.
17/// See [`create_tensor`](cubecl_runtime::client::Client::create_tensor).
18pub fn into_contiguous_pitched(
19    client: &Client,
20    input: TensorBinding,
21    dtype: ElemType,
22) -> TensorHandle {
23    if input.shape.len() <= 1 {
24        return into_contiguous(client, input, dtype);
25    }
26
27    let output = TensorHandle::empty(client, input.shape.clone(), dtype);
28
29    copy_into(client, input, output.clone().binding(), dtype);
30
31    output
32}
33
34/// Copies the input tensor into the output tensor following the strides.
35pub fn copy_into(client: &Client, input: TensorBinding, output: TensorBinding, dtype: ElemType) {
36    let rank = input.strides.len();
37
38    // The perpendicular kernel reads the input at the output's coordinates, so it only
39    // expresses a change of layout. A reshape keeping the rank would otherwise reach it and
40    // index the input out of bounds.
41    let same_shape = input.shape == output.shape;
42    if input.strides[rank - 1] == 1 || !same_shape {
43        return copy_gpu_ref(client, input, output, dtype);
44    }
45
46    if !try_launch_copy_perpendicular(client, &input, &output, dtype) {
47        copy_gpu_ref(client, input, output, dtype);
48    }
49}
cubecl_std/tensor/contiguous/launch.rs

cubecl_std/tensor/contiguous/
launch.rs
