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: ElemType,
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_dynamic(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(client: &Client, output: &TensorHandle) {
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(client: &Client, output: TensorBinding, dtype: ElemType) {
48    assert_eq!(2, output.shape.len(), "input should be a matrix");
49    assert_eq!(
50        output.shape[0], output.shape[1],
51        "input should be a square matrix"
52    );
53
54    let vectorization_factor = tensor_vector_size_parallel(
55        client.io_optimized_vector_sizes(dtype.size()),
56        &output.shape,
57        &output.strides,
58        1,
59    );
60
61    let cube_dim = CubeDim::new_2d(2, 2);
62    let vectors_x = output.shape[1] as u32 / vectorization_factor as u32;
63    let cube_count_x = vectors_x.div_ceil(cube_dim.x);
64    let cube_count_y = (output.shape[0] as u32).div_ceil(cube_dim.y);
65    let cube_count = CubeCount::new_2d(cube_count_x, cube_count_y);
66
67    let scalar = output.strides[0] + 1;
68    unsafe {
69        identity_kernel::launch_unchecked(
70            client,
71            cube_count,
72            cube_dim,
73            output.required_address_type(dtype.size()),
74            vectorization_factor,
75            output.into_tensor_arg(),
76            scalar,
77            dtype,
78        )
79    }
80}
cubecl_std/tensor/identity.rs

cubecl_std/tensor/
identity.rs
