cutile 0.3.0

cuTile Rust lets programmers safely author and execute tile kernels directly in Rust.
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
66
67
68

/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

//! Value test: `atomic_red_view_tko` accumulates concurrently into one tile.
//!
//! Every block atomically adds a ones-tile into tile [0, 0] of the output;
//! that tile must equal the block count in every element (a non-atomic
//! store would lose updates), and every other tile must stay zero.

use cutile::prelude::*;

use crate::common;

#[cutile::module]
mod atomic_red_kernels {
    use cutile::core::*;
    use cutile::tileir::*;

    #[cutile::entry()]
    unsafe fn atomic_red_accumulate(out: &Tensor<f32, { [-1, -1] }>) {
        let mut out_part: PartitionMut<f32, { [16, 16] }> =
            unsafe { out.partition_full_mut(const_shape![16, 16]) };
        let tile: Tile<f32, { [16, 16] }> = constant(1.0f32, const_shape![16, 16]);
        let _tok: Token = unsafe {
            atomic_red_view_tko(
                &mut out_part,
                tile,
                [0i32, 0i32],
                atomic::AddF,
                ordering::Relaxed,
                scope::Device,
            )
        };
    }
}

use atomic_red_kernels::atomic_red_accumulate;

#[test]
fn atomic_red_view_accumulates_across_blocks() {
    common::with_test_stack(|| {
        let device = cuda_core::Device::new(0).expect("device");
        let stream = device.new_stream().expect("stream");

        const BLOCKS: usize = 32;
        let out = api::zeros::<f32>(&[512, 16])
            .sync_on(&stream)
            .expect("zeros");
        unsafe { atomic_red_accumulate(&out) }
            .grid((BLOCKS as u32, 1, 1))
            .sync_on(&stream)
            .expect("launch");
        let host: Vec<f32> = out.dup().to_host_vec().sync_on(&stream).expect("to_host");
        assert_eq!(host.len(), 512 * 16);
        let (first_tile, rest) = host.split_at(16 * 16);
        assert!(
            first_tile.iter().all(|&v| v == BLOCKS as f32),
            "expected tile [0, 0] to accumulate to {BLOCKS}, got {:?}",
            &first_tile[..8]
        );
        assert!(
            rest.iter().all(|&v| v == 0.0),
            "expected all other tiles untouched"
        );
    });
}