diff --git a/cutile-rs/cutile-book/guide/debugging-and-profiling.md b/cutile-rs/cutile-book/guide/debugging-and-profiling.md index 684a608cd1..53da472232 100644 --- a/cutile-rs/cutile-book/guide/debugging-and-profiling.md +++ b/cutile-rs/cutile-book/guide/debugging-and-profiling.md @@ -148,8 +148,8 @@ debug = false Cargo provides only `DEBUG=true|false` to build scripts, not the exact debug level. Any enabled level currently selects device-debug mode, including `debug = "line-tables-only"` and `"limited"`. -It also selects device optimization level 0 unless explicitly overridden, -even in a release profile with `debug = true`. Automatic line-only mapping +Device-debug mode compiles at optimization level 0, even in a release +profile with `debug = true`. Automatic line-only mapping is not implemented. Use the explicit `line` override below for optimized profiling. See Cargo's [profile settings][cargo-profiles] and [build-script environment][cargo-build-env]. @@ -175,8 +175,8 @@ only for a kernel's crate does not. Setting the environment variable only when running an already-built binary does not change device compilation. `CompileOptions::new()` and `CompileOptions::default()` inherit this build -default. Use `debug_info()` to replace it for one launch. It sets both debug -flags, so it can also turn off a build-time `full` default: +default. Use `debug_info()` to replace it for one launch. It replaces the +inherited mode, so it can also turn off a build-time `full` default: ```rust use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; @@ -187,8 +187,7 @@ my_kernel(args) ``` Each option is part of the JIT cache key, so the modes do not share a cached -compiled kernel. The individual flags remain available, as does sanitizer -instrumentation: +compiled kernel. Sanitizer instrumentation is a separate option: ```rust use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; @@ -207,18 +206,37 @@ my_kernel(args) my_kernel(args).compile_options(CompileOptions::new().sanitize_memcheck(true)).sync()?; ``` -- `debug_info(level)` replaces both debug flags. The individual - `device_debug(bool)` and `lineinfo(bool)` setters change only their own - flag, leaving other inherited settings in place. -- `device_debug(true)` passes `--device-debug` to the device compiler and - implies optimization level 0. An explicit `opt_level` takes precedence, - but `tileiras` currently rejects full debug information at optimized - levels. The frontend also stops hoisting bounds checks out of loops. - Checks it proved unnecessary or moved to launch time are unaffected; - they never reach device code in any mode. -- `lineinfo(true)` passes `--lineinfo`: source-line correlation for Nsight Compute and Nsight Systems without changing code generation. This is the option for profiling optimized kernels. +- `debug_info(DebugInfoLevel::Full)` selects device-debug mode, which passes + `--opt-level 0 --device-debug`. The frontend also stops hoisting bounds + checks out of loops. Checks it proved unnecessary or moved to launch time + are unaffected; they never reach device code in any mode. +- `debug_info(DebugInfoLevel::Line)` passes `--lineinfo`: source-line correlation for Nsight Compute and Nsight Systems without changing code generation. This is the option for profiling optimized kernels. +- `debug_info(DebugInfoLevel::None)` and `Line` keep the current + optimization level, or use level 3 when leaving device-debug mode. +- `opt_level(n)` passes `--opt-level n` and keeps line information. + `tileiras` supports device debugging only at level 0, so this also leaves + device-debug mode. Whichever of `debug_info` and `opt_level` is called + last decides. - `sanitize_memcheck(true)` passes `--sanitize=memcheck` for `compute-sanitizer --tool memcheck`. -- `opt_level(n)` selects `--opt-level` directly; the default is 3. + +To switch an already-built application without changing its code, set these +variables when running it. Each applies after the code's options, as if the +matching setter were called last. They are read once per process: + +| Variable | Values | Acts like | +| --- | --- | --- | +| `CUTILE_JIT_DEBUG` | `none`, `line`, `full` | `debug_info(...)` | +| `CUTILE_JIT_OPT_LEVEL` | `0`, `1`, `2`, `3` | `opt_level(n)` | +| `CUTILE_JIT_SANITIZE` | `memcheck`, `none` | `sanitize_memcheck(...)` | + +```bash +CUTILE_JIT_DEBUG=full cuda-gdb --args ./target/release/my_program +CUTILE_JIT_SANITIZE=memcheck compute-sanitizer --tool memcheck ./target/release/my_program +``` + +An unset variable leaves the code's choice in place. An invalid value, or +`CUTILE_JIT_DEBUG=full` together with `CUTILE_JIT_OPT_LEVEL`, fails JIT +compilation with an error that names the variable. Full debug mode describes source locations and inline frames, not Rust variables or values. Tile IR currently supports control-flow debugging of @@ -275,6 +293,7 @@ without a GPU: ```bash cargo test -p cutile --test debug_info +cargo test -p cutile --test jit_env_overrides bash scripts/test_debug_defaults.sh ``` diff --git a/cutile-rs/cutile-compiler/src/check_optimizations.rs b/cutile-rs/cutile-compiler/src/check_optimizations.rs index f128397e76..371ddd72af 100644 --- a/cutile-rs/cutile-compiler/src/check_optimizations.rs +++ b/cutile-rs/cutile-compiler/src/check_optimizations.rs @@ -103,41 +103,43 @@ impl CheckOptimizations { mod tests { use super::*; use crate::cuda_tile_runtime_utils::{TileirasOptions, DEFAULT_OPT_LEVEL}; - use crate::hints::CompileOptions; + use crate::hints::{CodegenMode, CompileOptions, DebugInfoLevel}; #[test] - fn device_debug_implies_opt_zero_unless_explicit() { - let debug = - TileirasOptions::from_compile_options(&CompileOptions::new().device_debug(true)); - assert_eq!(debug.opt_level, 0); - assert!(debug.device_debug); - - let explicit = TileirasOptions::from_compile_options( - &CompileOptions::new().device_debug(true).opt_level(2), + fn full_debug_compiles_at_level_zero() { + let debug = TileirasOptions::from_compile_options( + &CompileOptions::new().debug_info(DebugInfoLevel::Full), ); - assert_eq!(explicit.opt_level, 2, "an explicit level wins"); + assert_eq!(debug.codegen, CodegenMode::Debug); + assert_eq!(debug.codegen.opt_level(), 0); let release = TileirasOptions::from_compile_options( - &CompileOptions::new().debug_info(crate::hints::DebugInfoLevel::None), + &CompileOptions::new().debug_info(DebugInfoLevel::None), ); - assert_eq!(release.opt_level, DEFAULT_OPT_LEVEL); + assert_eq!(release.codegen.opt_level(), DEFAULT_OPT_LEVEL); assert_eq!(release, TileirasOptions::default()); } #[test] fn flags_byte_is_injective_over_the_flag_combinations() { let mut seen = std::collections::BTreeSet::new(); - for dd in [false, true] { - for li in [false, true] { - for sm in [false, true] { - let o = TileirasOptions { - opt_level: 3, - device_debug: dd, - lineinfo: li, - sanitize_memcheck: sm, - }; - assert!(seen.insert(o.flags_byte()), "flags_byte collision"); - } + for codegen in [ + CodegenMode::Debug, + CodegenMode::Optimized { + level: 3, + lineinfo: false, + }, + CodegenMode::Optimized { + level: 3, + lineinfo: true, + }, + ] { + for sanitize_memcheck in [false, true] { + let o = TileirasOptions { + codegen, + sanitize_memcheck, + }; + assert!(seen.insert(o.flags_byte()), "flags_byte collision"); } } assert_eq!( diff --git a/cutile-rs/cutile-compiler/src/compile_api.rs b/cutile-rs/cutile-compiler/src/compile_api.rs index 31eba9b182..00fac81a48 100644 --- a/cutile-rs/cutile-compiler/src/compile_api.rs +++ b/cutile-rs/cutile-compiler/src/compile_api.rs @@ -232,11 +232,12 @@ impl crate::ast::Module> KernelCompiler { /// The returned string is the same 64-character lowercase SHA-256 key that /// the runtime's L2 cache lookup would use for this specialization. pub fn l2_cache_key(self) -> Result { - let gpu_name = self.gpu_name.clone(); + let resolved = self.with_jit_env_overrides()?; + let gpu_name = resolved.gpu_name.clone(); let tileiras_opts = crate::cuda_tile_runtime_utils::TileirasOptions::from_compile_options( - &self.compile_options, + &resolved.compile_options, ); - let artifacts = self.compile()?; + let artifacts = resolved.compile_resolved()?; if let Some(version) = artifacts.bytecode_version { let bytes = artifacts.bytecode()?; Ok(crate::cuda_tile_runtime_utils::current_l2_key_for_bytecode( @@ -255,6 +256,15 @@ impl crate::ast::Module> KernelCompiler { /// /// This is a pure compilation step — no GPU or CUDA driver is needed. pub fn compile(self) -> Result { + self.with_jit_env_overrides()?.compile_resolved() + } + + fn with_jit_env_overrides(mut self) -> Result { + self.compile_options = self.compile_options.with_jit_env_overrides()?; + Ok(self) + } + + fn compile_resolved(self) -> Result { let module_ast = (self.module_ast_fn)(); let modules = CUDATileModules::from_kernel(module_ast)?; diff --git a/cutile-rs/cutile-compiler/src/compiler/_function.rs b/cutile-rs/cutile-compiler/src/compiler/_function.rs index 8dc61241f8..ab77725f66 100644 --- a/cutile-rs/cutile-compiler/src/compiler/_function.rs +++ b/cutile-rs/cutile-compiler/src/compiler/_function.rs @@ -647,7 +647,7 @@ impl<'m> CUDATileFunctionCompiler<'m> { // in-kernel code motion (see CheckOptimizations::device_debug for // why discharge and launch relocation are untouched). let mut check_opts = crate::check_optimizations::CheckOptimizations::from_env(); - if compile_options.device_debug { + if compile_options.codegen == crate::hints::CodegenMode::Debug { check_opts.hoist_to_preheaders = false; } Ok(CUDATileFunctionCompiler { diff --git a/cutile-rs/cutile-compiler/src/cuda_tile_runtime_utils.rs b/cutile-rs/cutile-compiler/src/cuda_tile_runtime_utils.rs index 203dbdd5c1..963c2fa523 100644 --- a/cutile-rs/cutile-compiler/src/cuda_tile_runtime_utils.rs +++ b/cutile-rs/cutile-compiler/src/cuda_tile_runtime_utils.rs @@ -7,6 +7,7 @@ //! Provides GPU detection and bytecode compilation helpers. use crate::error::JITError; +use crate::hints::CodegenMode; use cuda_core::{get_device_sm_name, Device}; use cutile_ir::bytecode::{write_bytecode_version, BytecodeVersion}; use std::collections::HashMap; @@ -519,7 +520,7 @@ fn negotiate_bytecode_version( .map_err(|error| JITError::Generic(error.to_string())) } -/// `--opt-level` passed to `tileiras`. Not configurable yet. +/// Default `--opt-level` for optimized code. /// Numeric, not a string: the disk-cache key and entry header store it as one /// byte. pub const DEFAULT_OPT_LEVEL: u8 = 3; @@ -646,12 +647,8 @@ pub(crate) fn current_l2_key_for_module( /// field can never share a cubin. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct TileirasOptions { - /// `--opt-level `. - pub opt_level: u8, - /// `--device-debug`: generate debug information. - pub device_debug: bool, - /// `--lineinfo`: generate line-number information only. - pub lineinfo: bool, + /// `--opt-level `, plus `--device-debug` or `--lineinfo`. + pub codegen: CodegenMode, /// `--sanitize=memcheck`: instrument memory accesses for the sanitizer. pub sanitize_memcheck: bool, } @@ -659,28 +656,17 @@ pub struct TileirasOptions { impl Default for TileirasOptions { fn default() -> Self { Self { - opt_level: DEFAULT_OPT_LEVEL, - device_debug: false, - lineinfo: false, + codegen: CodegenMode::RELEASE, sanitize_memcheck: false, } } } impl TileirasOptions { - /// Resolves the launch-facing [`crate::hints::CompileOptions`] into the - /// stage-2 flags. `device_debug` implies `--opt-level 0` unless an - /// explicit level was requested. + /// The stage-2 subset of the launch-facing [`crate::hints::CompileOptions`]. pub fn from_compile_options(options: &crate::hints::CompileOptions) -> Self { - let opt_level = options.opt_level.unwrap_or(if options.device_debug { - 0 - } else { - DEFAULT_OPT_LEVEL - }); Self { - opt_level, - device_debug: options.device_debug, - lineinfo: options.lineinfo, + codegen: options.codegen, sanitize_memcheck: options.sanitize_memcheck, } } @@ -688,9 +674,11 @@ impl TileirasOptions { /// The boolean flags packed into one byte, for the cache-entry header /// and the L2 key material. pub fn flags_byte(&self) -> u8 { - (self.device_debug as u8) - | ((self.lineinfo as u8) << 1) - | ((self.sanitize_memcheck as u8) << 2) + let codegen = match self.codegen { + CodegenMode::Debug => 0b001, + CodegenMode::Optimized { lineinfo, .. } => (lineinfo as u8) << 1, + }; + codegen | ((self.sanitize_memcheck as u8) << 2) } } @@ -726,13 +714,14 @@ fn run_tileiras_at( JITError::Generic(format!("Failed to write bytecode for {bc_filename}: {e}")) })?; - let opt_level_arg = opts.opt_level.to_string(); + let opt_level_arg = opts.codegen.opt_level().to_string(); let mut args = vec!["--gpu-name", gpu_name, "--opt-level", &opt_level_arg]; - if opts.device_debug { - args.push("--device-debug"); - } - if opts.lineinfo { - args.push("--lineinfo"); + match opts.codegen { + CodegenMode::Debug => args.push("--device-debug"), + CodegenMode::Optimized { lineinfo: true, .. } => args.push("--lineinfo"), + CodegenMode::Optimized { + lineinfo: false, .. + } => {} } if opts.sanitize_memcheck { args.push("--sanitize=memcheck"); @@ -858,7 +847,7 @@ fn compile_bytecode_cached_at( let params = EntryParams { bc_sha256: Sha256::digest(bytecode).into(), gpu_name, - opt_level: opts.opt_level, + opt_level: opts.codegen.opt_level(), flags: opts.flags_byte(), tileiras_fp, }; diff --git a/cutile-rs/cutile-compiler/src/hints.rs b/cutile-rs/cutile-compiler/src/hints.rs index 0158ecfd3d..7a55d7537b 100644 --- a/cutile-rs/cutile-compiler/src/hints.rs +++ b/cutile-rs/cutile-compiler/src/hints.rs @@ -8,6 +8,7 @@ //! Pure Rust — no melior or tile-ir dependency. use crate::ast::SourceLocation; +use crate::cuda_tile_runtime_utils::DEFAULT_OPT_LEVEL; use crate::error::{JITError, SpannedJITError}; use quote::ToTokens; use std::collections::BTreeMap; @@ -99,6 +100,64 @@ pub enum DebugInfoLevel { Full, } +const BUILD_DEBUG_INFO: DebugInfoLevel = match env!("CUTILE_BUILD_DEBUG_INFO").as_bytes() { + b"full" => DebugInfoLevel::Full, + b"line" => DebugInfoLevel::Line, + _ => DebugInfoLevel::None, +}; + +/// How both compile stages generate device code. +/// +/// Debugging owns level 0 because `tileiras` accepts `--device-debug` at no +/// other level, so every value maps to flags `tileiras` accepts. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum CodegenMode { + /// `--opt-level 0 --device-debug`. The frontend also stops hoisting + /// bounds checks out of loops, so every check that runs on the device + /// sits at the source line that wrote it. Checks the compiler discharged + /// by proof or moved to launch time never reach device code in any mode. + Debug, + /// `--opt-level `, plus `--lineinfo` for profiler correlation. + Optimized { level: u8, lineinfo: bool }, +} + +impl CodegenMode { + pub const RELEASE: Self = Self::Optimized { + level: DEFAULT_OPT_LEVEL, + lineinfo: false, + }; + + pub fn opt_level(self) -> u8 { + match self { + Self::Debug => 0, + Self::Optimized { level, .. } => level, + } + } + + fn with_debug_info(self, debug_info: DebugInfoLevel) -> Self { + let level = match self { + Self::Debug => DEFAULT_OPT_LEVEL, + Self::Optimized { level, .. } => level, + }; + match debug_info { + DebugInfoLevel::Full => Self::Debug, + DebugInfoLevel::Line => Self::Optimized { + level, + lineinfo: true, + }, + DebugInfoLevel::None => Self::Optimized { + level, + lineinfo: false, + }, + } + } + + fn with_opt_level(self, level: u8) -> Self { + let lineinfo = matches!(self, Self::Optimized { lineinfo: true, .. }); + Self::Optimized { level, lineinfo } + } +} + /// Runtime compile options for kernel JIT compilation. /// /// These options control kernel-level compilation hints that can vary between @@ -111,26 +170,16 @@ pub enum DebugInfoLevel { /// `Full` for `debug = "line-tables-only"` or `"limited"`, not `Line`. /// `CUDA_RUST_DEBUG=none|line|full` overrides this default at build time; /// changing it when running an already-built app has no effect. Use -/// [`Self::debug_info`] to override the default for one compilation. +/// [`Self::debug_info`] to override the default for one compilation. At run +/// time, the `CUTILE_JIT_*` variables in [`crate::jit_env`] override what the +/// code selected. #[derive(Debug, Eq, PartialEq, Hash, Clone)] pub struct CompileOptions { pub occupancy: Option, pub num_cta_in_cga: Option, pub max_divisibility: Option, pub num_worker_warps_per_cta: Option, - /// `tileiras` optimization level (`--opt-level`). `None` means the - /// default: 3, or 0 when `device_debug` is set. - pub opt_level: Option, - /// Compile for debugging (`tileiras --device-debug`): the frontend stops - /// hoisting bounds checks out of loops, so every check that runs on the - /// device sits at the source line that wrote it, and the backend - /// generates debug information. Checks the compiler discharged by proof - /// or moved to launch time never reach device code in any mode. Implies - /// `--opt-level 0` unless `opt_level` is set explicitly. - pub device_debug: bool, - /// Emit line-number information (`tileiras --lineinfo`) for profiler - /// correlation, without the rest of the debug contract. - pub lineinfo: bool, + pub codegen: CodegenMode, /// Instrument memory accesses for Compute Sanitizer's memcheck tool /// (`tileiras --sanitize=memcheck`). pub sanitize_memcheck: bool, @@ -143,9 +192,7 @@ impl Default for CompileOptions { num_cta_in_cga: None, max_divisibility: None, num_worker_warps_per_cta: None, - opt_level: None, - device_debug: env!("CUTILE_BUILD_DEBUG_INFO") == "full", - lineinfo: env!("CUTILE_BUILD_DEBUG_INFO") == "line", + codegen: CodegenMode::RELEASE.with_debug_info(BUILD_DEBUG_INFO), sanitize_memcheck: false, } } @@ -156,14 +203,11 @@ impl CompileOptions { Self::default() } - /// Selects exactly one debug-info mode, replacing both debug flags. - /// - /// Other options, including an explicitly set optimization level, are - /// preserved. With no explicit optimization level, `Full` uses level 0 - /// and `None` / `Line` use level 3. + /// Selects the debug-info mode. `Full` selects [`CodegenMode::Debug`]; + /// `None` and `Line` select optimized code at the current level, or at + /// level 3 when leaving debug mode. pub fn debug_info(mut self, level: DebugInfoLevel) -> Self { - self.device_debug = level == DebugInfoLevel::Full; - self.lineinfo = level == DebugInfoLevel::Line; + self.codegen = self.codegen.with_debug_info(level); self } @@ -187,18 +231,10 @@ impl CompileOptions { self } + /// Selects optimized code at `opt_level`, keeping line info. This leaves + /// debug mode, which `tileiras` supports only at level 0. pub fn opt_level(mut self, opt_level: u8) -> Self { - self.opt_level = Some(opt_level); - self - } - - pub fn device_debug(mut self, device_debug: bool) -> Self { - self.device_debug = device_debug; - self - } - - pub fn lineinfo(mut self, lineinfo: bool) -> Self { - self.lineinfo = lineinfo; + self.codegen = self.codegen.with_opt_level(opt_level); self } @@ -336,3 +372,116 @@ impl OptimizationHints { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::cuda_tile_runtime_utils::TileirasOptions; + + /// `(opt_level, device_debug, lineinfo)` as `tileiras` gets them. + fn tileiras(options: &CompileOptions) -> (u8, bool, bool) { + let flags = TileirasOptions::from_compile_options(options); + let device_debug = flags.codegen == CodegenMode::Debug; + let lineinfo = matches!(flags.codegen, CodegenMode::Optimized { lineinfo: true, .. }); + (flags.codegen.opt_level(), device_debug, lineinfo) + } + + #[test] + fn every_setter_sequence_yields_flags_tileiras_accepts() { + let setters: Vec CompileOptions> = vec![ + |o| o.debug_info(DebugInfoLevel::None), + |o| o.debug_info(DebugInfoLevel::Line), + |o| o.debug_info(DebugInfoLevel::Full), + |o| o.opt_level(0), + |o| o.opt_level(3), + ]; + for first in &setters { + for second in &setters { + for third in &setters { + let options = third(second(first(CompileOptions::new()))); + let (opt_level, device_debug, lineinfo) = tileiras(&options); + assert!( + !device_debug || (opt_level == 0 && !lineinfo), + "{options:?} resolves to an invalid tileiras combination" + ); + } + } + } + } + + #[test] + fn the_last_setter_decides_between_debug_and_optimized_code() { + let full = DebugInfoLevel::Full; + assert_eq!( + tileiras(&CompileOptions::new().debug_info(full).opt_level(3)), + (3, false, false) + ); + assert_eq!( + tileiras(&CompileOptions::new().opt_level(3).debug_info(full)), + (0, true, false) + ); + } + + #[test] + fn optimized_code_keeps_its_level_and_line_info_across_setters() { + let line = DebugInfoLevel::Line; + assert_eq!( + tileiras(&CompileOptions::new().opt_level(2).debug_info(line)), + (2, false, true) + ); + assert_eq!( + tileiras(&CompileOptions::new().debug_info(line).opt_level(1)), + (1, false, true) + ); + assert_eq!( + tileiras( + &CompileOptions::new() + .debug_info(DebugInfoLevel::Full) + .debug_info(DebugInfoLevel::None) + ), + (3, false, false) + ); + } + + #[test] + fn cache_bytes_match_the_previous_encoding() { + let full = DebugInfoLevel::Full; + let line = DebugInfoLevel::Line; + for (options, opt_level, flags) in [ + ( + CompileOptions::new().debug_info(DebugInfoLevel::None), + 3, + 0b000, + ), + (CompileOptions::new().debug_info(line), 3, 0b010), + (CompileOptions::new().debug_info(full), 0, 0b001), + ( + CompileOptions::new().opt_level(1).debug_info(line), + 1, + 0b010, + ), + ( + CompileOptions::new() + .debug_info(full) + .sanitize_memcheck(true), + 0, + 0b101, + ), + ( + CompileOptions::new() + .debug_info(DebugInfoLevel::None) + .opt_level(0) + .sanitize_memcheck(true), + 0, + 0b100, + ), + ] { + let tileiras = TileirasOptions::from_compile_options(&options); + assert_eq!( + (tileiras.codegen.opt_level(), tileiras.flags_byte()), + (opt_level, flags), + "{options:?}" + ); + } + } +} diff --git a/cutile-rs/cutile-compiler/src/jit_cache.rs b/cutile-rs/cutile-compiler/src/jit_cache.rs index 9fc75fdb24..55f48a6635 100644 --- a/cutile-rs/cutile-compiler/src/jit_cache.rs +++ b/cutile-rs/cutile-compiler/src/jit_cache.rs @@ -752,7 +752,7 @@ pub fn l2_key( put_field(&mut h, gpu_name.as_bytes()); // The stage-2 flags are key material: a debug or sanitizer cubin must // never be served for a release request (or vice versa). - put_field(&mut h, &[opts.opt_level, opts.flags_byte()]); + put_field(&mut h, &[opts.codegen.opt_level(), opts.flags_byte()]); put_field(&mut h, tileiras_fp.as_bytes()); hex(&h.finalize()) } @@ -889,6 +889,7 @@ pub fn decode_entry(bytes: &[u8], params: &EntryParams<'_>) -> Option> { #[cfg(test)] mod tests { use super::*; + use crate::hints::CodegenMode; use cutile_ir::bytecode::BytecodeVersion; #[cfg(unix)] @@ -1016,7 +1017,10 @@ mod tests { fn opts(opt_level: u8) -> crate::cuda_tile_runtime_utils::TileirasOptions { crate::cuda_tile_runtime_utils::TileirasOptions { - opt_level, + codegen: CodegenMode::Optimized { + level: opt_level, + lineinfo: false, + }, ..Default::default() } } @@ -1042,10 +1046,13 @@ mod tests { // Every stage-2 flag is key material: a debug, lineinfo, or // sanitizer cubin must never share a slot with a release one. let mut flagged = opts(3); - flagged.device_debug = true; + flagged.codegen = CodegenMode::Debug; assert_ne!(base, l2_key(b"bc", V, "sm_90", &flagged, "fp")); let mut flagged = opts(3); - flagged.lineinfo = true; + flagged.codegen = CodegenMode::Optimized { + level: 3, + lineinfo: true, + }; assert_ne!(base, l2_key(b"bc", V, "sm_90", &flagged, "fp")); let mut flagged = opts(3); flagged.sanitize_memcheck = true; diff --git a/cutile-rs/cutile-compiler/src/jit_env.rs b/cutile-rs/cutile-compiler/src/jit_env.rs new file mode 100644 index 0000000000..185525a26b --- /dev/null +++ b/cutile-rs/cutile-compiler/src/jit_env.rs @@ -0,0 +1,260 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! Process-wide `CUTILE_JIT_*` overrides of the code's [`CompileOptions`]. +//! +//! They let a wrapper switch an already-built application into debug or +//! sanitizer compilation. Both compile stages read the overridden options, so +//! the frontend's check placement and the `tileiras` flags always agree. + +use crate::error::JITError; +use crate::hints::{CompileOptions, DebugInfoLevel}; +use std::env::{self, VarError}; +use std::sync::OnceLock; + +pub const JIT_DEBUG_ENV: &str = "CUTILE_JIT_DEBUG"; +pub const JIT_OPT_LEVEL_ENV: &str = "CUTILE_JIT_OPT_LEVEL"; +pub const JIT_SANITIZE_ENV: &str = "CUTILE_JIT_SANITIZE"; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +struct JitEnvOverrides { + debug_info: Option, + opt_level: Option, + sanitize_memcheck: Option, +} + +impl JitEnvOverrides { + fn from_env() -> Result { + Self::parse( + env::var(JIT_DEBUG_ENV), + env::var(JIT_OPT_LEVEL_ENV), + env::var(JIT_SANITIZE_ENV), + ) + } + + fn parse( + debug_info: Result, + opt_level: Result, + sanitize: Result, + ) -> Result { + let overrides = Self { + debug_info: choice( + JIT_DEBUG_ENV, + debug_info, + &[ + ("none", DebugInfoLevel::None), + ("line", DebugInfoLevel::Line), + ("full", DebugInfoLevel::Full), + ], + )?, + opt_level: choice( + JIT_OPT_LEVEL_ENV, + opt_level, + &[("0", 0), ("1", 1), ("2", 2), ("3", 3)], + )?, + sanitize_memcheck: choice( + JIT_SANITIZE_ENV, + sanitize, + &[("memcheck", true), ("none", false)], + )?, + }; + if overrides.debug_info == Some(DebugInfoLevel::Full) && overrides.opt_level.is_some() { + return Err(format!( + "{JIT_DEBUG_ENV}=full compiles at level 0; unset {JIT_OPT_LEVEL_ENV}" + )); + } + Ok(overrides) + } + + fn apply(self, options: &CompileOptions) -> CompileOptions { + let mut options = options.clone(); + if let Some(level) = self.debug_info { + options = options.debug_info(level); + } + if let Some(level) = self.opt_level { + options = options.opt_level(level); + } + if let Some(memcheck) = self.sanitize_memcheck { + options = options.sanitize_memcheck(memcheck); + } + options + } +} + +fn choice( + name: &str, + value: Result, + choices: &[(&str, T)], +) -> Result, String> { + let expected = || { + let texts: Vec<_> = choices.iter().map(|&(text, _)| text).collect(); + texts.join(", ") + }; + match value { + Err(VarError::NotPresent) => Ok(None), + Err(VarError::NotUnicode(_)) => Err(format!( + "{name} is not valid Unicode; expected one of {}", + expected() + )), + Ok(value) => choices + .iter() + .find(|&&(text, _)| text == value) + .map(|&(_, choice)| Some(choice)) + .ok_or_else(|| format!("invalid {name}={value:?}; expected one of {}", expected())), + } +} + +impl CompileOptions { + /// These options with the process's `CUTILE_JIT_*` overrides applied, + /// which is what both compile stages see. The environment is read once + /// per process. + pub fn with_jit_env_overrides(&self) -> Result { + static OVERRIDES: OnceLock> = OnceLock::new(); + let overrides = OVERRIDES + .get_or_init(JitEnvOverrides::from_env) + .as_ref() + .map_err(|message| JITError::Generic(message.clone()))?; + Ok(overrides.apply(self)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hints::CodegenMode; + use std::ffi::OsString; + + const UNSET: Result = Err(VarError::NotPresent); + + fn set(value: &str) -> Result { + Ok(value.to_string()) + } + + fn debug_only(level: DebugInfoLevel) -> JitEnvOverrides { + JitEnvOverrides { + debug_info: Some(level), + ..Default::default() + } + } + + #[test] + fn unset_variables_keep_code_options() { + let overrides = JitEnvOverrides::parse(UNSET, UNSET, UNSET).unwrap(); + let code = CompileOptions::new() + .opt_level(2) + .debug_info(DebugInfoLevel::Line) + .sanitize_memcheck(true); + assert_eq!(overrides.apply(&code), code); + } + + #[test] + fn parses_every_documented_value() { + for (text, level) in [ + ("none", DebugInfoLevel::None), + ("line", DebugInfoLevel::Line), + ("full", DebugInfoLevel::Full), + ] { + let overrides = JitEnvOverrides::parse(set(text), UNSET, UNSET).unwrap(); + assert_eq!(overrides.debug_info, Some(level)); + } + for level in 0..=3u8 { + let overrides = JitEnvOverrides::parse(UNSET, set(&level.to_string()), UNSET).unwrap(); + assert_eq!(overrides.opt_level, Some(level)); + } + for (text, memcheck) in [("memcheck", true), ("none", false)] { + let overrides = JitEnvOverrides::parse(UNSET, UNSET, set(text)).unwrap(); + assert_eq!(overrides.sanitize_memcheck, Some(memcheck)); + } + } + + #[test] + fn invalid_values_name_the_variable() { + let not_unicode = Err(VarError::NotUnicode(OsString::from("full"))); + for (parsed, name) in [ + ( + JitEnvOverrides::parse(set("FULL"), UNSET, UNSET), + JIT_DEBUG_ENV, + ), + ( + JitEnvOverrides::parse(not_unicode, UNSET, UNSET), + JIT_DEBUG_ENV, + ), + ( + JitEnvOverrides::parse(UNSET, set("4"), UNSET), + JIT_OPT_LEVEL_ENV, + ), + ( + JitEnvOverrides::parse(UNSET, UNSET, set("address")), + JIT_SANITIZE_ENV, + ), + ] { + let message = parsed.unwrap_err(); + assert!(message.contains(name), "{message}"); + } + } + + #[test] + fn env_full_debug_outranks_an_optimized_level_in_code() { + let code = CompileOptions::new() + .debug_info(DebugInfoLevel::None) + .opt_level(3); + let effective = debug_only(DebugInfoLevel::Full).apply(&code); + assert_eq!(effective.codegen, CodegenMode::Debug); + } + + #[test] + fn full_debug_with_an_opt_level_is_rejected() { + let message = JitEnvOverrides::parse(set("full"), set("2"), UNSET).unwrap_err(); + assert!( + message.contains(JIT_DEBUG_ENV) && message.contains(JIT_OPT_LEVEL_ENV), + "{message}" + ); + } + + #[test] + fn env_line_info_and_opt_level_combine() { + let overrides = JitEnvOverrides::parse(set("line"), set("1"), UNSET).unwrap(); + let effective = overrides.apply(&CompileOptions::new().debug_info(DebugInfoLevel::Full)); + assert_eq!( + effective.codegen, + CodegenMode::Optimized { + level: 1, + lineinfo: true + } + ); + } + + #[test] + fn env_debug_matches_the_code_setter() { + let debug = CompileOptions::new() + .debug_info(DebugInfoLevel::Full) + .occupancy(2); + let optimized = CompileOptions::new().opt_level(2); + for code in [debug, optimized] { + for level in [ + DebugInfoLevel::None, + DebugInfoLevel::Line, + DebugInfoLevel::Full, + ] { + assert_eq!( + debug_only(level).apply(&code), + code.clone().debug_info(level) + ); + } + } + } + + #[test] + fn env_sanitize_overrides_in_both_directions() { + for (code, env, expected) in [(false, true, true), (true, false, false)] { + let overrides = JitEnvOverrides { + sanitize_memcheck: Some(env), + ..Default::default() + }; + let effective = overrides.apply(&CompileOptions::new().sanitize_memcheck(code)); + assert_eq!(effective.sanitize_memcheck, expected); + } + } +} diff --git a/cutile-rs/cutile-compiler/src/lib.rs b/cutile-rs/cutile-compiler/src/lib.rs index 222fdc34c7..b2a57fdfd2 100644 --- a/cutile-rs/cutile-compiler/src/lib.rs +++ b/cutile-rs/cutile-compiler/src/lib.rs @@ -17,6 +17,7 @@ pub mod error; pub mod generics; pub mod hints; pub mod jit_cache; +pub mod jit_env; mod kernel_entry_generator; pub mod kernel_naming; pub mod registry; diff --git a/cutile-rs/cutile/src/tile_kernel.rs b/cutile-rs/cutile/src/tile_kernel.rs index 503c826e95..a388d0a612 100644 --- a/cutile-rs/cutile/src/tile_kernel.rs +++ b/cutile-rs/cutile/src/tile_kernel.rs @@ -83,7 +83,7 @@ pub use cuda_async::{ }; pub use cutile_compiler::compiler::utils::CompileOptions; -pub use cutile_compiler::hints::DebugInfoLevel; +pub use cutile_compiler::hints::{CodegenMode, DebugInfoLevel}; /// Function-pointer form of the module AST provider generated by /// `#[cutile::module]`. @@ -746,6 +746,7 @@ fn compile_and_load_kernel( key_str: &str, ) -> Result { let t0 = std::time::Instant::now(); + let compile_options = &compile_options.with_jit_env_overrides()?; let stride_args_refs: Vec<(&str, &[i32])> = stride_args .iter() diff --git a/cutile-rs/cutile/tests/check_placement_ratchet.rs b/cutile-rs/cutile/tests/check_placement_ratchet.rs index cd6514d682..518c456e95 100644 --- a/cutile-rs/cutile/tests/check_placement_ratchet.rs +++ b/cutile-rs/cutile/tests/check_placement_ratchet.rs @@ -23,6 +23,7 @@ //! test. New behavior gets new tests elsewhere. use cutile_compiler::compiler::utils::CompileOptions; +use cutile_compiler::hints::DebugInfoLevel; mod common; @@ -157,7 +158,7 @@ mod access_counts { .generics(generics) .strides(&strides) // Optimized-mode pins; debug-mode expectations are tested below. - .options(CompileOptions::default().device_debug(false)); + .options(CompileOptions::default().debug_info(DebugInfoLevel::None)); if let Some(grid) = p.grid { compiler = compiler.grid(grid); } @@ -287,7 +288,7 @@ fn debug_counts( .target("sm_120") .generics(generics) .strides(&strides) - .options(CompileOptions::new().device_debug(true)); + .options(CompileOptions::new().debug_info(DebugInfoLevel::Full)); if let Some(grid) = grid { compiler = compiler.grid(grid); } diff --git a/cutile-rs/cutile/tests/checked_store_corpus.rs b/cutile-rs/cutile/tests/checked_store_corpus.rs index 88c7ec17e8..797829f67d 100644 --- a/cutile-rs/cutile/tests/checked_store_corpus.rs +++ b/cutile-rs/cutile/tests/checked_store_corpus.rs @@ -20,6 +20,7 @@ //! permuted case is only reachable through loads. use cutile_compiler::compiler::utils::CompileOptions; +use cutile_compiler::hints::DebugInfoLevel; mod common; @@ -120,7 +121,7 @@ fn compile_mlir( &[], &[], None, - &CompileOptions::default().device_debug(false), + &CompileOptions::default().debug_info(DebugInfoLevel::None), ) .map_err(|err| err.to_string()) } @@ -139,7 +140,7 @@ fn artifacts( .target("sm_120") .generics(generics) .strides(strides) - .options(CompileOptions::default().device_debug(false)) + .options(CompileOptions::default().debug_info(DebugInfoLevel::None)) .compile() .unwrap_or_else(|e| panic!("compile {name}: {e}")) } diff --git a/cutile-rs/cutile/tests/common/mod.rs b/cutile-rs/cutile/tests/common/mod.rs index 61b88faf18..4f2399a0b5 100644 --- a/cutile-rs/cutile/tests/common/mod.rs +++ b/cutile-rs/cutile/tests/common/mod.rs @@ -10,6 +10,7 @@ use cutile::compile_api::{CheckPlacementCounts, KernelCompiler}; use cutile_compiler::ast::Module; use cutile_compiler::compiler::utils::CompileOptions; use cutile_compiler::error::JITError; +use cutile_compiler::jit_env::{JIT_DEBUG_ENV, JIT_OPT_LEVEL_ENV, JIT_SANITIZE_ENV}; use cutile_compiler::specialization::{DivHint, SpecializationBits}; /// Runtime prerequisites only. Failed discovery is a test failure; an @@ -162,3 +163,35 @@ where .compile() .map(|artifacts| (artifacts.ir_text(), artifacts.check_counts())) } + +/// Names the case a child process was spawned to run. +const JIT_ENV_CASE: &str = "CUTILE_TEST_JIT_ENV_CASE"; + +/// The case this process runs, when [`run_jit_env_case`] spawned it. +#[allow(dead_code)] +pub fn jit_env_case() -> Option { + std::env::var(JIT_ENV_CASE).ok() +} + +/// Runs `test` in a fresh process whose only `CUTILE_JIT_*` settings are the +/// ones in `vars`, and returns its stdout and stderr. The overrides are read +/// once per process, so each setting needs a process of its own. +#[allow(dead_code)] +pub fn run_jit_env_case(test: &str, case: &str, vars: &[(&str, &str)]) -> String { + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args([test, "--exact", "--nocapture", "--test-threads=1"]) + .env_remove(JIT_DEBUG_ENV) + .env_remove(JIT_OPT_LEVEL_ENV) + .env_remove(JIT_SANITIZE_ENV) + .envs(vars.iter().copied()) + .env(JIT_ENV_CASE, case) + .output() + .expect("spawn child test process"); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + assert!( + output.status.success() && stdout.contains("1 passed"), + "case {case} with {vars:?} failed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + format!("{stdout}{stderr}") +} diff --git a/cutile-rs/cutile/tests/debug_info.rs b/cutile-rs/cutile/tests/debug_info.rs index 4b1734f85c..34d2e9d78c 100644 --- a/cutile-rs/cutile/tests/debug_info.rs +++ b/cutile-rs/cutile/tests/debug_info.rs @@ -11,7 +11,7 @@ mod common; mod debug_helpers_file; use cutile::compile_api::{CompileArtifacts, KernelCompiler}; -use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; +use cutile::tile_kernel::{CodegenMode, CompileOptions, DebugInfoLevel}; use cutile_ir::ir::{DebugScope, Location, Module, OpId, Operation}; use std::path::PathBuf; use std::process::Command; @@ -129,10 +129,16 @@ fn explicit_debug_modes_override_build_default_and_change_cache_keys() { use std::collections::HashSet; let mut l1 = HashSet::new(); let mut l2 = HashSet::new(); - for (level, opt, debug, line) in [ - (DebugInfoLevel::None, 3, false, false), - (DebugInfoLevel::Line, 3, false, true), - (DebugInfoLevel::Full, 0, true, false), + for (level, codegen) in [ + (DebugInfoLevel::None, CodegenMode::RELEASE), + ( + DebugInfoLevel::Line, + CodegenMode::Optimized { + level: 3, + lineinfo: true, + }, + ), + (DebugInfoLevel::Full, CodegenMode::Debug), ] { let options = CompileOptions::new() .debug_info(DebugInfoLevel::Full) @@ -141,10 +147,7 @@ fn explicit_debug_modes_override_build_default_and_change_cache_keys() { .sanitize_memcheck(true); assert_eq!(options.occupancy, Some(2)); let flags = TileirasOptions::from_compile_options(&options); - assert_eq!( - (flags.opt_level, flags.device_debug, flags.lineinfo), - (opt, debug, line) - ); + assert_eq!(flags.codegen, codegen); assert!(flags.sanitize_memcheck); assert!(l1.insert( cutile::tile_kernel::TileFunctionKey::builder("m", "k") @@ -171,7 +174,6 @@ fn build_time_default() { .ok() .or_else(|| option_env!("CUDA_RUST_DEBUG").map(str::to_owned)); let options = CompileOptions::new(); - assert!(!(options.device_debug && options.lineinfo)); if let Some(expected) = expected { let level = match expected.as_str() { "none" => DebugInfoLevel::None, diff --git a/cutile-rs/cutile/tests/differential_placement.rs b/cutile-rs/cutile/tests/differential_placement.rs index c12c41fb47..ecf6d3a7cf 100644 --- a/cutile-rs/cutile/tests/differential_placement.rs +++ b/cutile-rs/cutile/tests/differential_placement.rs @@ -41,6 +41,7 @@ use cutile::compile_api::KernelCompiler; use cutile::prelude::*; use cutile::tensor::{IntoPartition, ToHostVec}; use cutile_compiler::compiler::utils::CompileOptions; +use cutile_compiler::hints::DebugInfoLevel; mod common; @@ -313,7 +314,7 @@ fn hash_outputs(v: &[f32]) -> u64 { /// surfaces as an `Err` from a `sync()`. fn execute_case(name: &str) -> Result, String> { // Exercise placement optimizations even when Cargo enables device debug. - let options = CompileOptions::new().device_debug(false); + let options = CompileOptions::new().debug_info(DebugInfoLevel::None); fn e(err: E) -> String { err.to_string() } @@ -495,7 +496,7 @@ fn differential_compile_runner() { .target("sm_120") .generics(vec![B.to_string()]) .strides(&[("z", &[1]), ("x", &[1])]) - .options(CompileOptions::default().device_debug(false)) + .options(CompileOptions::default().debug_info(DebugInfoLevel::None)) .compile() .unwrap_or_else(|err| panic!("compile {case}: {err}")); let counts = artifacts.check_counts(); diff --git a/cutile-rs/cutile/tests/gpu.rs b/cutile-rs/cutile/tests/gpu.rs index 10e48e0459..8169f93edf 100644 --- a/cutile-rs/cutile/tests/gpu.rs +++ b/cutile-rs/cutile/tests/gpu.rs @@ -83,6 +83,9 @@ mod borrow_raw_parts; #[path = "gpu/device_debug.rs"] mod device_debug; +#[path = "gpu/jit_env_overrides.rs"] +mod jit_env_overrides; + // --------------------------------------------------------------------------- // Soundness guards: shape/storage invariants, launcher-side validation, and // graph-capture input discipline. diff --git a/cutile-rs/cutile/tests/gpu/audit_check_hoisting.rs b/cutile-rs/cutile/tests/gpu/audit_check_hoisting.rs index 98768c0c79..80e5a62e17 100644 --- a/cutile-rs/cutile/tests/gpu/audit_check_hoisting.rs +++ b/cutile-rs/cutile/tests/gpu/audit_check_hoisting.rs @@ -12,7 +12,7 @@ use std::sync::Arc; use cutile::prelude::*; -use cutile::tile_kernel::CompileOptions; +use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; use crate::audit_common::{self, host, report_outcome, run_in_subprocess, upload, Outcome}; use crate::common; @@ -132,7 +132,7 @@ fn a_loop_body_with_continue_keeps_its_check_in_place() { let (z, _x, _limit) = check_hoisting_module::continue_before(api::zeros::(&[B]).partition([B]), x, 4) .generics(vec![B.to_string()]) - .compile_options(CompileOptions::new().device_debug(false)) + .compile_options(CompileOptions::new().debug_info(DebugInfoLevel::None)) .sync() .expect("continue_before with every attained access in bounds"); let expected: Vec = (0..B) @@ -167,7 +167,7 @@ fn stepped_loops_are_checked_at_the_last_attained_index() { let (z, _x) = check_hoisting_module::stepped_dynamic(api::zeros::(&[B]).partition([B]), x) .generics(vec![B.to_string()]) - .compile_options(CompileOptions::new().device_debug(false)) + .compile_options(CompileOptions::new().debug_info(DebugInfoLevel::None)) .sync() .expect("nine tiles cover the attained indices {0, 4, 8}"); let expected: Vec = (0..B) @@ -222,7 +222,7 @@ fn execute_trap_case(case: &str) -> Result<(), String> { let x = upload((0..(8 * B) as i32).map(|v| v as f32).collect()); check_hoisting_module::stepped_dynamic(api::zeros::(&[B]).partition([B]), x) .generics(vec![B.to_string()]) - .compile_options(CompileOptions::new().device_debug(false)) + .compile_options(CompileOptions::new().debug_info(DebugInfoLevel::None)) .sync() .map_err(|err| err.to_string())?; Ok(()) diff --git a/cutile-rs/cutile/tests/gpu/audit_common.rs b/cutile-rs/cutile/tests/gpu/audit_common.rs index 5c1973a07c..837f20a561 100644 --- a/cutile-rs/cutile/tests/gpu/audit_common.rs +++ b/cutile-rs/cutile/tests/gpu/audit_common.rs @@ -17,6 +17,7 @@ use cutile::prelude::*; use cutile_compiler::ast::Module; use cutile_compiler::compiler::utils::CompileOptions; use cutile_compiler::error::JITError; +use cutile_compiler::hints::DebugInfoLevel; /// Compiles one kernel of `module_ast_fn`'s module for `sm_120` and returns /// its IR text and placement counts. @@ -33,7 +34,7 @@ pub fn compile Module>( .strides(strides) // These audits pin optimized placement, independent of Cargo's // development-profile device-debug default. - .options(CompileOptions::default().device_debug(false)) + .options(CompileOptions::default().debug_info(DebugInfoLevel::None)) .compile() .map(|artifacts| (artifacts.ir_text(), artifacts.check_counts())) } diff --git a/cutile-rs/cutile/tests/gpu/jit_env_overrides.rs b/cutile-rs/cutile/tests/gpu/jit_env_overrides.rs new file mode 100644 index 0000000000..2b81f5eb06 --- /dev/null +++ b/cutile-rs/cutile/tests/gpu/jit_env_overrides.rs @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! The launcher applies `CUTILE_JIT_*` before both compile stages and keys +//! its disk cache by the effective mode. Each case launches in a fresh +//! process against one shared store, so a cubin served from disk can only +//! have come from an earlier process with an equivalent mode. + +use crate::common; +use cutile::api; +use cutile::jit_cache::{self, FileSystemJitStore}; +use cutile::prelude::*; +use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; +use cutile_compiler::jit_env::{JIT_DEBUG_ENV, JIT_OPT_LEVEL_ENV, JIT_SANITIZE_ENV}; +use std::collections::BTreeMap; +use std::sync::Arc; + +#[cutile::module] +mod launch_env_module { + use cutile::core::*; + + /// Release hoists this kernel's one check to the loop preheader. + #[cutile::entry] + fn dynamic_loop(z: &mut Tensor, x: &Tensor, n: i32) { + let p = x.partition(shape![B]); + for j in 0i32..n { + let t = p.load([j]); + z.store(t); + } + } +} + +const B: usize = 64; +const TILES: usize = 4; +const STORE_ENV: &str = "CUTILE_TEST_JIT_ENV_STORE"; + +fn code(name: &str) -> CompileOptions { + let release = CompileOptions::new() + .debug_info(DebugInfoLevel::None) + .opt_level(3); + match name { + "release" => release, + "line" => release.debug_info(DebugInfoLevel::Line), + "full" => release.debug_info(DebugInfoLevel::Full), + _ => panic!("unknown code options {name:?}"), + } +} + +/// The loop stores every tile of `x` in turn, so `z` ends as the last one. +fn launch_and_check(options: CompileOptions) { + let x: Vec = (0..B * TILES).map(|v| v as f32).collect(); + let x: Arc> = api::copy_host_vec_to_device(&Arc::new(x)) + .sync() + .expect("upload x") + .into(); + let (z, _x, _n) = + launch_env_module::dynamic_loop(api::zeros::(&[B]).partition([B]), x, TILES as i32) + .generics(vec![B.to_string()]) + .compile_options(options) + .sync() + .expect("dynamic_loop"); + let z = z.unpartition().to_host_vec().sync().expect("download z"); + let last: Vec = ((TILES - 1) * B..TILES * B).map(|v| v as f32).collect(); + assert_eq!(z, last); +} + +#[test] +fn disk_cache_follows_the_effective_mode() { + const TEST: &str = "jit_env_overrides::disk_cache_follows_the_effective_mode"; + if let Some(name) = common::jit_env_case() { + let store = std::env::var_os(STORE_ENV).expect("store directory"); + common::with_test_stack(move || { + jit_cache::enable(Arc::new( + FileSystemJitStore::new(&store).expect("open store"), + )); + launch_and_check(code(&name)); + }); + return; + } + + let store = std::env::temp_dir().join(format!("cutile_jit_env_store_{}", std::process::id())); + std::fs::create_dir_all(&store).unwrap(); + let store_dir = store.to_str().unwrap(); + let launch = |name: &str, vars: &[(&str, &str)]| { + let vars = [vars, &[(STORE_ENV, store_dir), ("CUTILE_JIT_TIMING", "1")]].concat(); + let output = common::run_jit_env_case(TEST, name, &vars); + let timing = output + .lines() + .find(|line| { + line.starts_with("CUTILE_JIT_TIMING") && line.contains(" function=dynamic_loop ") + }) + .unwrap_or_else(|| panic!("no timing line for dynamic_loop:\n{output}")); + let fields: BTreeMap<_, _> = timing + .split(' ') + .filter_map(|field| field.split_once('=')) + .collect(); + ( + fields["stage2_source"].to_string(), + format!("{},{}", fields["checks_hoisted"], fields["checks_in_place"]), + ) + }; + + for (name, vars, source, placement) in [ + ("release", &[(JIT_DEBUG_ENV, "full")][..], "tileiras", "0,1"), + ("full", &[], "disk", "0,1"), + ("release", &[], "tileiras", "1,0"), + ("full", &[(JIT_DEBUG_ENV, "none")], "disk", "1,0"), + ("release", &[(JIT_DEBUG_ENV, "line")], "tileiras", "1,0"), + ("line", &[], "disk", "1,0"), + ("line", &[(JIT_OPT_LEVEL_ENV, "1")], "tileiras", "1,0"), + ( + "release", + &[(JIT_SANITIZE_ENV, "memcheck")], + "tileiras", + "1,0", + ), + ] { + assert_eq!( + launch(name, vars), + (source.to_string(), placement.to_string()), + "{vars:?} over {name:?}" + ); + } + let _ = std::fs::remove_dir_all(&store); +} diff --git a/cutile-rs/cutile/tests/jit_env_overrides.rs b/cutile-rs/cutile/tests/jit_env_overrides.rs new file mode 100644 index 0000000000..b8709b78f0 --- /dev/null +++ b/cutile-rs/cutile/tests/jit_env_overrides.rs @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +//! `CUTILE_JIT_*` overrides reach both compile stages: the frontend's check +//! placement and the cache key that names the `tileiras` flags. Every case +//! runs in a child process that inherits no `CUTILE_JIT_*` setting. + +use cutile::compile_api::KernelCompiler; +use cutile::tile_kernel::{CompileOptions, DebugInfoLevel}; +use cutile_compiler::jit_env::{JIT_DEBUG_ENV, JIT_OPT_LEVEL_ENV, JIT_SANITIZE_ENV}; +use std::collections::BTreeMap; + +mod common; + +#[cutile::module] +mod env_module { + use cutile::core::*; + + /// Release hoists this kernel's one check to the loop preheader. + #[cutile::entry] + fn dynamic_loop(z: &mut Tensor, x: &Tensor, n: i32) { + let p = x.partition(shape![B]); + for j in 0i32..n { + let t = p.load([j]); + z.store(t); + } + } +} + +const HOISTED: (u32, u32) = (1, 0); +const IN_PLACE: (u32, u32) = (0, 1); + +fn code(name: &str) -> CompileOptions { + let release = CompileOptions::new().debug_info(DebugInfoLevel::None); + match name { + "release" => release.opt_level(3), + "line" => release.debug_info(DebugInfoLevel::Line).opt_level(3), + "line_o1" => release.debug_info(DebugInfoLevel::Line).opt_level(1), + "full" => release.debug_info(DebugInfoLevel::Full), + "memcheck" => release.opt_level(3).sanitize_memcheck(true), + _ => panic!("unknown code options {name:?}"), + } +} + +fn compiler( + options: CompileOptions, +) -> KernelCompiler cutile::cutile_compiler::ast::Module> { + KernelCompiler::new( + env_module::__module_ast_self as fn() -> cutile::cutile_compiler::ast::Module, + "env_module", + "dynamic_loop", + ) + .target("sm_120") + .generics(vec!["64".into()]) + .strides(&[("z", &[1]), ("x", &[1])]) + .options(options) +} + +fn placement(options: CompileOptions) -> (u32, u32) { + let counts = compiler(options).compile().unwrap().check_counts(); + (counts.hoisted, counts.in_place) +} + +fn compile_error(options: CompileOptions) -> String { + match compiler(options).compile() { + Ok(_) => panic!("compilation must fail"), + Err(error) => error.to_string(), + } +} + +/// Runs `body` when this process is the child for `test`; otherwise spawns +/// that child with `vars`. +fn isolated(test: &str, vars: &[(&str, &str)], body: impl FnOnce() + Send + 'static) { + if common::jit_env_case().is_some() { + common::with_test_stack(body); + } else { + common::run_jit_env_case(test, "run", vars); + } +} + +#[test] +fn env_none_turns_off_debug_in_code() { + isolated( + "env_none_turns_off_debug_in_code", + &[(JIT_DEBUG_ENV, "none")], + || assert_eq!(placement(code("full")), HOISTED), + ); +} + +#[test] +fn env_full_keeps_checks_at_their_access_site() { + isolated( + "env_full_keeps_checks_at_their_access_site", + &[(JIT_DEBUG_ENV, "full")], + || assert_eq!(placement(code("release")), IN_PLACE), + ); +} + +#[test] +fn env_line_with_an_opt_level_stays_optimized() { + isolated( + "env_line_with_an_opt_level_stays_optimized", + &[(JIT_DEBUG_ENV, "line"), (JIT_OPT_LEVEL_ENV, "1")], + || assert_eq!(placement(code("full")), HOISTED), + ); +} + +#[test] +fn invalid_value_fails_compilation() { + isolated( + "invalid_value_fails_compilation", + &[(JIT_DEBUG_ENV, "FULL")], + || { + let error = compile_error(code("release")); + assert!(error.contains(JIT_DEBUG_ENV), "{error}"); + }, + ); +} + +#[test] +fn full_debug_with_an_opt_level_fails_compilation() { + isolated( + "full_debug_with_an_opt_level_fails_compilation", + &[(JIT_DEBUG_ENV, "full"), (JIT_OPT_LEVEL_ENV, "2")], + || { + let error = compile_error(code("release")); + assert!( + error.contains(JIT_DEBUG_ENV) && error.contains(JIT_OPT_LEVEL_ENV), + "{error}" + ); + }, + ); +} + +#[test] +fn overrides_are_read_once_per_process() { + isolated("overrides_are_read_once_per_process", &[], || { + assert_eq!(placement(code("release")), HOISTED); + std::env::set_var(JIT_DEBUG_ENV, "full"); + assert_eq!(placement(code("release")), HOISTED); + }); +} + +/// One compile per process, so the key a process reports is computed under +/// that process's overrides alone. Equal keys in two processes mean one +/// would serve the other's disk-cache entry. +#[test] +fn cache_key_follows_the_effective_mode() { + const TEST: &str = "cache_key_follows_the_effective_mode"; + if let Some(name) = common::jit_env_case() { + let options = code(&name); + let (hoisted, in_place) = placement(options.clone()); + let key = compiler(options).l2_cache_key().unwrap(); + println!("PROBE key={key} placement={hoisted},{in_place}"); + return; + } + let probe = |name: &str, vars: &[(&str, &str)]| { + let output = common::run_jit_env_case(TEST, name, vars); + let line = output + .lines() + .find_map(|line| line.split_once("PROBE ").map(|(_, probe)| probe)) + .unwrap_or_else(|| panic!("no probe line:\n{output}")); + let fields: BTreeMap<_, _> = line + .split(' ') + .filter_map(|field| field.split_once('=')) + .collect(); + (fields["key"].to_string(), fields["placement"].to_string()) + }; + + let release = probe("release", &[]); + let line = probe("line", &[]); + let line_o1 = probe("line_o1", &[]); + let full = probe("full", &[]); + let memcheck = probe("memcheck", &[]); + let keys = [&release.0, &line.0, &line_o1.0, &full.0, &memcheck.0]; + for (i, a) in keys.iter().enumerate() { + for b in &keys[i + 1..] { + assert_ne!(a, b, "distinct modes must not share a cache entry"); + } + } + assert_eq!(release.1, "1,0"); + assert_eq!(full.1, "0,1"); + + for (expected, name, vars) in [ + (&release, "full", &[(JIT_DEBUG_ENV, "none")][..]), + (&line, "release", &[(JIT_DEBUG_ENV, "line")]), + ( + &line_o1, + "full", + &[(JIT_DEBUG_ENV, "line"), (JIT_OPT_LEVEL_ENV, "1")], + ), + (&full, "release", &[(JIT_DEBUG_ENV, "full")]), + (&memcheck, "release", &[(JIT_SANITIZE_ENV, "memcheck")]), + ] { + assert_eq!( + &probe(name, vars), + expected, + "{vars:?} over {name:?} must match the same mode set in code" + ); + } +} diff --git a/cutile-rs/cutile/tests/partition_index_schedules.rs b/cutile-rs/cutile/tests/partition_index_schedules.rs index ba7d871301..767d4fb842 100644 --- a/cutile-rs/cutile/tests/partition_index_schedules.rs +++ b/cutile-rs/cutile/tests/partition_index_schedules.rs @@ -8,6 +8,7 @@ #![allow(deprecated)] use cutile_compiler::compiler::utils::CompileOptions; +use cutile_compiler::hints::DebugInfoLevel; mod common; @@ -1420,7 +1421,7 @@ fn kv_loop_bounds_checks_hoist_to_the_preheader() { &[], &[], Some((4, 1, 1)), - &CompileOptions::default().device_debug(false), + &CompileOptions::default().debug_info(DebugInfoLevel::None), ) .map_err(|err| err.to_string()) .expect("Failed to compile hoisted-checks kernel"); @@ -1474,7 +1475,7 @@ fn compile_hoist_kernel(function_name: &str, generics: &[&str]) -> String { &[], &[], None, - &CompileOptions::default().device_debug(false), + &CompileOptions::default().debug_info(DebugInfoLevel::None), ) .map_err(|err| err.to_string()) .expect("Failed to compile hoist kernel") diff --git a/cutile-rs/scripts/run_cpu_tests.sh b/cutile-rs/scripts/run_cpu_tests.sh index 52f457cd3c..746f12bd39 100755 --- a/cutile-rs/scripts/run_cpu_tests.sh +++ b/cutile-rs/scripts/run_cpu_tests.sh @@ -58,6 +58,7 @@ for test_target in \ flash_attention_compile \ global_memory \ integer_ops \ + jit_env_overrides \ kernel_compiler \ load_tile_like_examples \ macro_submodule \