cubecl_std/tensor/
identity.rs

1use cubecl::frontend::TensorBinding;
2use cubecl::prelude::*;
3use cubecl::tensor_vector_size_parallel;
4use cubecl_core as cubecl;
5
6use super::TensorHandle;
7
8#[cube(launch_unchecked, address_type = "dynamic")]
9fn identity_kernel<C: Numeric, N: Size>(
10    output: &mut Tensor<Vector<C, N>>,
11    gap: usize,
12    #[define(C)] _elem: StorageType,
13) {
14    let pos_x = ABSOLUTE_POS_X as usize * output.vector_size();
15    let pos_y = ABSOLUTE_POS_Y as usize;
16    let vector_size = output.vector_size();
17    if pos_y < output.shape(0) && pos_x < output.shape(1) {
18        let mut vector = Vector::new(C::from_int(0));
19        let offs_y = pos_y * output.stride(0);
20
21        let start_pos = offs_y + pos_x;
22        let mut offset = 0;
23        while offset < output.vector_size() {
24            let remainder = (start_pos + offset) % gap;
25            if remainder == 0 {
26                vector.insert(offset, C::from_int(1));
27                offset += gap;
28            } else {
29                offset += gap - remainder;
30            }
31        }
32        output[start_pos / vector_size] = vector;
33    }
34}
35
36/// Launch identity matrix kernel.
37/// Ensure output is a [`TensorHandle`] containing a square matrix.
38/// output will contain the identity matrix.
39pub fn launch<R: Runtime>(client: &ComputeClient<R>, output: &TensorHandle<R>) {
40    let dtype = output.dtype;
41    launch_ref(client, output.clone().binding(), dtype);
42}
43
44/// Launch identity matrix kernel by ref.
45/// Ensure output is a [`TensorHandleRef`] containing a square matrix.
46/// output will contain the identity matrix.
47pub fn launch_ref<R: Runtime>(
48    client: &ComputeClient<R>,
49    output: TensorBinding<R>,
50    dtype: StorageType,
51) {
52    assert_eq!(2, output.shape.len(), "input should be a matrix");
53    assert_eq!(
54        output.shape[0], output.shape[1],
55        "input should be a square matrix"
56    );
57
58    let vectorization_factor = tensor_vector_size_parallel(
59        client.io_optimized_vector_sizes(dtype.size()),
60        &output.shape,
61        &output.strides,
62        1,
63    );
64
65    let cube_dim = CubeDim::new_2d(2, 2);
66    let vectors_x = output.shape[1] as u32 / vectorization_factor as u32;
67    let cube_count_x = vectors_x.div_ceil(cube_dim.x);
68    let cube_count_y = (output.shape[0] as u32).div_ceil(cube_dim.y);
69    let cube_count = CubeCount::new_2d(cube_count_x, cube_count_y);
70
71    let scalar = output.strides[0] + 1;
72    unsafe {
73        identity_kernel::launch_unchecked(
74            client,
75            cube_count,
76            cube_dim,
77            output.required_address_type(dtype.size()),
78            vectorization_factor,
79            output.into_tensor_arg(),
80            scalar,
81            dtype,
82        )
83    }
84}
cubecl_std/tensor/identity.rs

cubecl_std/tensor/
identity.rs
