cutile 0.3.1

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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

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

//! `x.load_like(z)` is the method-form spelling of `load_tile_like(x, z)`.
//! Pins that both spellings lower to identical Tile IR (same ops, same
//! bounds-check obligations) and that the method form compiles and runs
//! in a `#[cutile::entry]` kernel.

use cutile::prelude::*;
use cutile_compiler::compiler::utils::CompileOptions;

mod common;

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

    #[cutile::entry()]
    fn add_free_fn<const S: [i32; 1]>(
        z: &mut Tensor<f32, S>,
        x: &Tensor<f32, { [-1] }>,
        y: &Tensor<f32, { [-1] }>,
    ) {
        let tx = load_tile_like(x, z);
        let ty = load_tile_like(y, z);
        z.store(tx + ty);
    }

    #[cutile::entry()]
    fn add_method<const S: [i32; 1]>(
        z: &mut Tensor<f32, S>,
        x: &Tensor<f32, { [-1] }>,
        y: &Tensor<f32, { [-1] }>,
    ) {
        let tx = x.load_like(z);
        let ty = y.load_like(z);
        z.store(tx + ty);
    }
}

use load_like_module::__module_ast_self as load_like_module_ast;
use load_like_module::add_method;

fn compile(name: &str) -> String {
    common::compile_to_ir(
        load_like_module_ast,
        "load_like_module",
        name,
        &["4".to_string()],
        &[("z", &[1]), ("x", &[1]), ("y", &[1])],
        &[],
        &[],
        None,
        &CompileOptions::default(),
    )
    .expect("Failed to compile.")
}

/// Normalize a module dump to its live-op sequence: SSA numbers stripped,
/// dead `constant` lines dropped (the method form's extra inline hop
/// materializes shape-metadata constants the free-fn form doesn't; both
/// are DCE'd by tileiras). What remains — every view construction, load,
/// store, assume, and assert, in order — is exactly the placement-relevant
/// content.
fn live_ops(module: &str) -> Vec<String> {
    module
        .lines()
        .map(str::trim)
        .filter(|line| !line.is_empty())
        .filter(|line| !line.contains("= constant "))
        .map(|line| {
            let mut out = String::new();
            let mut chars = line.chars().peekable();
            while let Some(c) = chars.next() {
                if c == '%' {
                    while chars.peek().is_some_and(|d| d.is_ascii_digit()) {
                        chars.next();
                    }
                    out.push('%');
                } else {
                    out.push(c);
                }
            }
            out
        })
        .collect()
}

#[test]
fn both_spellings_lower_to_the_same_live_ops() {
    common::with_test_stack(|| {
        let free_fn = compile("add_free_fn").replace("add_free_fn", "K");
        let method = compile("add_method").replace("add_method", "K");
        assert_eq!(
            live_ops(&free_fn),
            live_ops(&method),
            "load_like must lower to the same live ops as load_tile_like"
        );
        // Same check-placement outcomes, explicitly: neither spelling may
        // change the number of device asserts.
        assert_eq!(
            free_fn.matches("assert").count(),
            method.matches("assert").count()
        );
    });
}

#[test]
fn method_form_runs_end_to_end() {
    common::with_test_stack(|| {
        let len = 32usize;
        let z_host = add_method(
            api::zeros(&[len]).partition([4]),
            api::arange::<f32>(len),
            api::ones(&[len]),
        )
        .grid(((len / 4) as u32, 1, 1))
        .first()
        .unpartition()
        .to_host_vec()
        .sync()
        .expect("add_method kernel");
        for (i, v) in z_host.iter().enumerate() {
            assert_eq!(*v, i as f32 + 1.0, "index {i}");
        }
    });
}