cutile 0.0.2

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

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

//! Compile-only coverage for example-style `load_tile_like` calls.
//!
//! The examples usually rely on inference (`let tile = load_tile_like(...)`)
//! instead of ascribing the returned tile type. These tests exercise the same
//! JIT path without requiring a CUDA device.

use cutile;
use cutile_compiler::compiler::utils::CompileOptions;
use cutile_compiler::compiler::{CUDATileFunctionCompiler, CUDATileModules};

mod common;

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

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

    #[cutile::entry()]
    fn saxpy_like<const S: [i32; 2]>(
        y: &mut Tensor<f32, S>,
        a: f32,
        x: &Tensor<f32, { [-1, -1] }>,
    ) {
        let tile_a = a.broadcast(y.shape());
        let tile_x = load_tile_like(x, y);
        let tile_y = y.load();
        y.store(tile_a * tile_x + tile_y);
    }

    #[cutile::entry()]
    fn generic_saxpy_like<T: ElementType, const S: [i32; 2]>(
        y: &mut Tensor<T, S>,
        a: T,
        x: &Tensor<T, { [-1, -1] }>,
    ) {
        let tile_a = a.broadcast(y.shape());
        let tile_x = load_tile_like(x, y);
        let tile_y = y.load();
        y.store(tile_a * tile_x + tile_y);
    }
}

use load_tile_like_examples_module::__module_ast_self;

fn compile(kernel: &str, generics: &[String], strides: &[(&str, &[i32])]) -> String {
    let modules = CUDATileModules::from_kernel(__module_ast_self())
        .expect("Failed to create CUDATileModules");
    let compiler = CUDATileFunctionCompiler::new(
        &modules,
        "load_tile_like_examples_module",
        kernel,
        generics,
        strides,
        &[],
        &[],
        None,
        "sm_120".to_string(),
        &CompileOptions::default(),
    )
    .expect("Failed to create compiler");
    let mlir = compiler.compile().expect("Failed to compile").to_string();
    println!("=== MLIR for {kernel} ===\n{mlir}");
    mlir
}

#[test]
fn compiles_add_refs_style_1d_inference() {
    common::with_test_stack(|| {
        let mlir = compile(
            "add_refs_like",
            &[4.to_string()],
            &[("z", &[1]), ("x", &[1]), ("y", &[1])],
        );
        assert!(mlir.contains("load_view_tko"));
        assert_eq!(mlir.matches("load_view_tko").count(), 2);
    });
}

#[test]
fn compiles_saxpy_style_2d_inference() {
    common::with_test_stack(|| {
        let mlir = compile(
            "saxpy_like",
            &[2.to_string(), 4.to_string()],
            &[("y", &[4, 1]), ("x", &[4, 1])],
        );
        assert!(mlir.contains("load_view_tko"));
        assert!(mlir.contains("tile<2x4xf32>"));
    });
}

#[test]
fn compiles_generic_saxpy_style_2d_inference() {
    common::with_test_stack(|| {
        let mlir = compile(
            "generic_saxpy_like",
            &["f32".to_string(), 2.to_string(), 4.to_string()],
            &[("y", &[4, 1]), ("x", &[4, 1])],
        );
        assert!(mlir.contains("load_view_tko"));
        assert!(mlir.contains("tile<2x4xf32>"));
    });
}