cutile 0.2.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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

/*
 * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
/// Common test utilities and constants shared across all test modules.
use cutile::compile_api::KernelCompiler;
use cutile_compiler::ast::Module;
use cutile_compiler::compiler::utils::CompileOptions;
use cutile_compiler::error::JITError;
use cutile_compiler::specialization::{DivHint, SpecializationBits};

/// Stack size for test threads.
///
/// Tests require larger stack sizes due to:
/// - Deep MLIR AST structures during compilation
/// - Multiple unary operations in single test kernels
/// - Nested function calls in the compiler
///
/// Binary search determined minimum requirements:
/// - Basic tests: ~2.121 MB
/// - With assume variants: ~2.612 MB
/// - With reduce/scan operations: ~2.7 MB
/// - With all unary math operations: ~5 MB (after adding absf, negf, negi, floor)
/// - tensor_views module tests require a bit more headroom.
/// Using 8 MB provides an adequate safety margin for all tests.
pub const TEST_STACK_SIZE: usize = 8_000_000; // 8 MB

/// Helper to run a test with the required stack size.
///
/// # Example
///
/// ```rust,ignore
/// #[test]
/// fn my_test() {
///     common::with_test_stack(|| {
///         // Your test code here
///     });
/// }
/// ```
pub fn with_test_stack<F, R>(f: F) -> R
where
    F: FnOnce() -> R + Send + 'static,
    R: Send + 'static,
{
    std::thread::Builder::new()
        .stack_size(TEST_STACK_SIZE)
        .spawn(f)
        .expect("Failed to spawn test thread")
        .join()
        .expect("Test thread panicked")
}

#[allow(clippy::too_many_arguments, dead_code)]
pub fn compile_to_ir<F>(
    module_ast_fn: F,
    module_name: &str,
    function_name: &str,
    generics: &[String],
    strides: &[(&str, &[i32])],
    spec_args: &[(&str, &SpecializationBits)],
    scalar_hints: &[(&str, &DivHint)],
    const_grid: Option<(u32, u32, u32)>,
    options: &CompileOptions,
) -> Result<String, JITError>
where
    F: Fn() -> Module,
{
    let spec_args = spec_args
        .iter()
        .map(|(name, spec)| (*name, (*spec).clone()))
        .collect::<Vec<_>>();
    let scalar_hints = scalar_hints
        .iter()
        .map(|(name, hint)| (*name, **hint))
        .collect::<Vec<_>>();

    let mut compiler = KernelCompiler::new(module_ast_fn, module_name, function_name)
        .target("sm_120")
        .generics(generics.to_vec())
        .strides(strides)
        .spec_args(&spec_args)
        .scalar_hints(&scalar_hints)
        .options(options.clone());

    if let Some(grid) = const_grid {
        compiler = compiler.grid(grid);
    }

    compiler.compile().map(|artifacts| artifacts.ir_text())
}