cutile 0.3.1

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

//! Tile kernel compilation, caching, launching, and partitioning for CUDA device operations.

use anyhow::Result;
use cuda_async::error::DeviceError;
use cuda_core::DType;
use cuda_core::{memcpy_dtoh_async, Function};
use cutile_compiler::ast::Module;
use cutile_compiler::compile_api::KernelCompiler;
use cutile_compiler::compiler::{CUDATileFunctionCompiler, CUDATileModules};
use cutile_compiler::cuda_tile_runtime_utils::{
    compile_bytecode_cached, env_flag_enabled, get_compiler_version, get_gpu_name,
    recompile_after_disk_rejection, serialize_tile_ir_bytecode, tileiras_fingerprint,
    toolchain_env_snapshot, Stage2Source, TileirasOptions, ToolchainEnvSnapshot,
};
use cutile_compiler::specialization::{DivHint, SpecializationBits};
use dashmap::DashMap;
use once_cell::sync::OnceCell;
use std::fs;
use std::future::IntoFuture;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};

// JIT diagnostic logging (set CUTILE_JIT_LOG=1, true, yes, or on to enable)

fn jit_log_enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| env_flag_enabled("CUTILE_JIT_LOG"))
}

macro_rules! jit_log {
    ($($arg:tt)*) => {
        if jit_log_enabled() {
            eprintln!("[cutile::jit] {}", format!($($arg)*));
        }
    };
}

static JIT_COMPILE_COUNT: AtomicU64 = AtomicU64::new(0);

/// Process-global JIT compile counter: +1 per successful compile, +0 on cache
/// hits and on failed compiles. Equals the number of distinct kernels cached.
/// Snapshot before a call and check the delta to get exact miss counts.
///
/// A disk-cache hit still counts: the counter tracks in-memory misses, which
/// run the compiler frontend either way. Absent failures,
/// `jit_compile_count == jit_backend_compile_count + jit_disk_hit_count`
/// (both in [`crate::jit_cache`]).
pub fn jit_compile_count() -> u64 {
    JIT_COMPILE_COUNT.load(Ordering::Relaxed)
}

#[inline]
fn record_jit_compile() {
    JIT_COMPILE_COUNT.fetch_add(1, Ordering::Relaxed);
}

use crate::error::*;
use crate::tensor::{
    GridBound, IntoPartition, IntoPartitionArc, KernelInput, KernelOutput, Partition, Tensor,
};

pub use cuda_async::{
    device_buffer::*, device_context::*, device_future::*, device_operation::*, launch::*,
    predicate::*, scheduling_policies::*,
};

pub use cutile_compiler::compiler::utils::CompileOptions;

/// Function-pointer form of the module AST provider generated by
/// `#[cutile::module]`.
pub type ModuleAstFn = fn() -> Module;

/// Cache key for a compiled tile kernel.
///
/// Two kernel invocations that share the same `TileFunctionKey` can reuse the same compiled
/// CUDA module and function, avoiding recompilation. The key captures everything that can
/// change the generated GPU code: module name, function name, generic type/const parameters,
/// tensor stride layouts, (optionally) the launch grid, compile options, source hash,
/// GPU architecture, compiler version, and the `tileiras` binary that assembles the cubin.
///
/// Tensor extents are deliberately absent: `stride_args` records only which
/// dimensions have stride 1, and `spec_args` only power-of-two divisibility. A
/// `[1024, 1024]` and a `[4096, 4096]` matmul share one key, because extents are
/// runtime kernel arguments and do not reach the generated code.
///
/// `source_hash` covers the kernel's own module, not the dependency modules the
/// use-graph links in. Editing a helper module that the kernel
/// calls changes the cubin without changing this field. Within a process this is
/// harmless, since a rebuild restarts it; it is why the on-disk cache keys on the
/// serialized bytecode rather than on this struct.
#[derive(Debug, Eq, PartialEq, Hash, Clone)]
pub struct TileFunctionKey {
    module_name: String,
    function_name: String,
    pub function_generics: Vec<String>,
    pub stride_args: Vec<(String, Vec<i32>)>,
    pub spec_args: Vec<(String, SpecializationBits)>,
    pub scalar_hints: Vec<(String, DivHint)>,
    pub grid: Option<(u32, u32, u32)>,
    pub compile_options: CompileOptions,
    source_hash: String,
    device_id: usize,
    gpu_name: String,
    compiler_version: String,
    /// Output of `tileiras --version`, not `nvcc --version`: the JIT resolves
    /// `tileiras` on its own, so `CUTILE_TILEIRAS_PATH` can point at a binary the
    /// toolkit version knows nothing about.
    tileiras_fingerprint: String,
}

/// Builder for [`TileFunctionKey`].
///
/// With 11 positional arguments it is easy to silently transpose two `String`
/// fields and produce a wrong-but-valid key. The builder makes each field
/// self-documenting and keeps future additions backward-compatible.
///
/// # Example
///
/// ```rust,ignore
/// let key = TileFunctionKey::builder("linalg", "matmul")
///     .generics(vec!["f32".into(), "128".into()])
///     .source_hash(linalg::_SOURCE_HASH)
///     .device_id(device_id)
///     .gpu_name(get_gpu_name(device_id))
///     .compiler_version(get_compiler_version())
///     .tileiras_fingerprint(tileiras_fingerprint())
///     .build();
/// ```
pub struct TileFunctionKeyBuilder {
    module_name: String,
    function_name: String,
    function_generics: Vec<String>,
    stride_args: Vec<(String, Vec<i32>)>,
    spec_args: Vec<(String, SpecializationBits)>,
    scalar_hints: Vec<(String, DivHint)>,
    grid: Option<(u32, u32, u32)>,
    compile_options: CompileOptions,
    source_hash: String,
    device_id: usize,
    gpu_name: String,
    compiler_version: String,
    tileiras_fingerprint: String,
}

impl TileFunctionKeyBuilder {
    pub fn generics(mut self, generics: Vec<String>) -> Self {
        self.function_generics = generics;
        self
    }
    pub fn stride_args(mut self, stride_args: Vec<(String, Vec<i32>)>) -> Self {
        self.stride_args = stride_args;
        self
    }
    pub fn spec_args(mut self, spec_args: Vec<(String, SpecializationBits)>) -> Self {
        self.spec_args = spec_args;
        self
    }
    pub fn scalar_hints(mut self, scalar_hints: Vec<(String, DivHint)>) -> Self {
        self.scalar_hints = scalar_hints;
        self
    }
    pub fn grid(mut self, grid: (u32, u32, u32)) -> Self {
        self.grid = Some(grid);
        self
    }
    pub fn compile_options(mut self, options: CompileOptions) -> Self {
        self.compile_options = options;
        self
    }
    pub fn source_hash(mut self, hash: impl Into<String>) -> Self {
        self.source_hash = hash.into();
        self
    }
    pub fn device_id(mut self, device_id: usize) -> Self {
        self.device_id = device_id;
        self
    }
    pub fn gpu_name(mut self, name: impl Into<String>) -> Self {
        self.gpu_name = name.into();
        self
    }
    pub fn compiler_version(mut self, version: impl Into<String>) -> Self {
        self.compiler_version = version.into();
        self
    }
    /// Output of `tileiras --version`; see [`tileiras_fingerprint`].
    pub fn tileiras_fingerprint(mut self, fingerprint: impl Into<String>) -> Self {
        self.tileiras_fingerprint = fingerprint.into();
        self
    }
    pub fn build(self) -> TileFunctionKey {
        TileFunctionKey {
            module_name: self.module_name,
            function_name: self.function_name,
            function_generics: self.function_generics,
            stride_args: self.stride_args,
            spec_args: self.spec_args,
            scalar_hints: self.scalar_hints,
            grid: self.grid,
            compile_options: self.compile_options,
            source_hash: self.source_hash,
            device_id: self.device_id,
            gpu_name: self.gpu_name,
            compiler_version: self.compiler_version,
            tileiras_fingerprint: self.tileiras_fingerprint,
        }
    }
}

impl TileFunctionKey {
    /// The kernel's module name.
    pub fn module_name(&self) -> &str {
        &self.module_name
    }

    /// The kernel's function name.
    pub fn function_name(&self) -> &str {
        &self.function_name
    }

    /// Start building a key with required `module_name` and `function_name`.
    /// All other fields default to empty / `None` / `default()`.
    pub fn builder(
        module_name: impl Into<String>,
        function_name: impl Into<String>,
    ) -> TileFunctionKeyBuilder {
        TileFunctionKeyBuilder {
            module_name: module_name.into(),
            function_name: function_name.into(),
            function_generics: vec![],
            stride_args: vec![],
            spec_args: vec![],
            scalar_hints: vec![],
            grid: None,
            compile_options: CompileOptions::default(),
            source_hash: String::new(),
            device_id: 0,
            gpu_name: String::new(),
            compiler_version: String::new(),
            tileiras_fingerprint: String::new(),
        }
    }
}

impl FunctionKey for TileFunctionKey {}

/// A resolved launch-site specialization together with the lazy module AST
/// provider needed to derive its persistent L2 cache key.
///
/// Creating this value resolves the specialization identity for a launch, but
/// does not compile or launch the kernel.
///
/// The structured [`TileFunctionKey`] is sufficient for an L1 lookup. Deriving
/// the L2 key additionally runs the compiler frontend because that key hashes
/// serialized Tile IR bytecode. Keeping the AST as a provider preserves the
/// L1-hit fast path: [`Self::l1_cache_key`] never builds the AST or runs the
/// frontend.
pub struct Specialization<F: Fn() -> Module> {
    module_ast_fn: F,
    key: TileFunctionKey,
}

impl<F: Fn() -> Module> Specialization<F> {
    /// Returns the complete structured key used by the in-memory kernel cache.
    pub fn l1_cache_key(&self) -> &TileFunctionKey {
        &self.key
    }

    /// Consumes the specialization and returns its structured L1 key.
    pub fn into_l1_cache_key(self) -> TileFunctionKey {
        self.key
    }

    /// Returns the persistent L2 key this specialization would look up.
    ///
    /// This runs the compiler frontend and canonical bytecode serializer, but
    /// it does not query a JIT store, compile a cubin, or load a CUDA module.
    pub fn l2_cache_key(&self) -> std::result::Result<String, cutile_compiler::error::JITError> {
        let stride_refs: Vec<(&str, &[i32])> = self
            .key
            .stride_args
            .iter()
            .map(|(name, strides)| (name.as_str(), strides.as_slice()))
            .collect();
        let spec_refs: Vec<(&str, SpecializationBits)> = self
            .key
            .spec_args
            .iter()
            .map(|(name, spec)| (name.as_str(), spec.clone()))
            .collect();
        let scalar_hint_refs: Vec<(&str, DivHint)> = self
            .key
            .scalar_hints
            .iter()
            .map(|(name, hint)| (name.as_str(), *hint))
            .collect();

        let mut compiler = KernelCompiler::new(
            &self.module_ast_fn,
            &self.key.module_name,
            &self.key.function_name,
        )
        .target(&self.key.gpu_name)
        .generics(self.key.function_generics.clone())
        .strides(&stride_refs)
        .spec_args(&spec_refs)
        .scalar_hints(&scalar_hint_refs)
        .options(self.key.compile_options.clone());
        if let Some(grid) = self.key.grid {
            compiler = compiler.grid(grid);
        }
        compiler.l2_cache_key()
    }
}

/// Resolves the canonical in-memory key for one macro-generated launch-site
/// specialization.
///
/// The generated launcher already derives the specialization metadata from its
/// materialized arguments. This helper adds the current device and toolchain
/// identity exactly once and keeps the lazy AST provider alongside the key for
/// an optional [`Specialization::l2_cache_key`] call.
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn _specialization_from_context<F: Fn() -> Module>(
    ctx: &ExecutionContext,
    module_ast_fn: F,
    module_name: &str,
    function_name: &str,
    function_generics: Vec<String>,
    stride_args: Vec<(String, Vec<i32>)>,
    spec_args: Vec<(String, SpecializationBits)>,
    scalar_hints: Vec<(String, DivHint)>,
    const_grid: Option<(u32, u32, u32)>,
    compile_options: CompileOptions,
    source_hash: &str,
) -> Specialization<F> {
    let device_id = ctx.get_device_id();
    let gpu_name = get_gpu_name(device_id);
    let mut key_builder = TileFunctionKey::builder(module_name, function_name)
        .generics(function_generics)
        .stride_args(stride_args)
        .spec_args(spec_args)
        .scalar_hints(scalar_hints)
        .compile_options(compile_options)
        .source_hash(source_hash)
        .device_id(device_id)
        .gpu_name(gpu_name)
        .compiler_version(get_compiler_version())
        .tileiras_fingerprint(tileiras_fingerprint());
    if let Some(grid) = const_grid {
        key_builder = key_builder.grid(grid);
    }
    Specialization {
        module_ast_fn,
        key: key_builder.build(),
    }
}

/// One macro-generated launch site's most recent resolution: the volatile
/// parts of the specialization it was resolved for, plus the resolved
/// function and validator.
///
/// The generated launcher probes this before constructing a
/// [`TileFunctionKey`]: on the steady state the probe is a handful of
/// integer/`String` comparisons and two `Arc` clones, with no allocation and
/// no toolchain-fingerprint work. Probe equality implies key equality: the
/// process-constant key fields (names, source hash, compiler version) cannot
/// differ at one site, `gpu_name` is a function of `device_id`, and stride
/// hints are derived from the compared [`SpecializationBits`].
///
/// Toolchain env semantics: the `CUTILE_TILEIRAS_PATH` / toolkit env values
/// are snapshotted when the site fills, not re-read per launch — a
/// mid-process env switch takes effect for cache misses and new sites (the
/// global cache still keys by fingerprint) but does not invalidate an
/// already-hot launch site. Reading the env on every launch costs
/// allocations on the very path this cache exists to strip.
pub struct LaunchSite {
    inner: std::sync::RwLock<Option<std::sync::Arc<SiteResolution>>>,
}

/// The snapshot a [`LaunchSite`] holds. Constructed by the generated
/// launcher on a probe miss, after the normal cache path resolved.
pub struct SiteResolution {
    /// Cache generation at fill time: any eviction from the global cache
    /// invalidates every launch site, so no site outlives a quiesced clear.
    epoch: u64,
    device_id: usize,
    #[allow(dead_code)]
    toolchain: ToolchainEnvSnapshot,
    generics: Vec<String>,
    specs: Vec<SpecializationBits>,
    scalar_hints: Vec<DivHint>,
    const_grid: Option<(u32, u32, u32)>,
    compile_options: CompileOptions,
    function: Arc<Function>,
    validator: Arc<Validator>,
}

impl SiteResolution {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        device_id: usize,
        generics: Vec<String>,
        specs: Vec<SpecializationBits>,
        scalar_hints: Vec<DivHint>,
        const_grid: Option<(u32, u32, u32)>,
        compile_options: CompileOptions,
        function: Arc<Function>,
        validator: Arc<Validator>,
    ) -> Self {
        Self {
            epoch: kernel_cache_epoch(),
            device_id,
            toolchain: toolchain_env_snapshot(),
            generics,
            specs,
            scalar_hints,
            const_grid,
            compile_options,
            function,
            validator,
        }
    }
}

impl LaunchSite {
    pub const fn new() -> Self {
        Self {
            inner: std::sync::RwLock::new(None),
        }
    }

    /// The cached resolution, if the probe matches it exactly.
    pub fn get(
        &self,
        device_id: usize,
        generics: &[String],
        specs: &[&SpecializationBits],
        scalar_hints: &[DivHint],
        const_grid: Option<(u32, u32, u32)>,
        compile_options: &CompileOptions,
    ) -> Option<(Arc<Function>, Arc<Validator>)> {
        let guard = self.inner.read().ok()?;
        let r = guard.as_ref()?;
        if r.epoch == kernel_cache_epoch()
            && r.device_id == device_id
            && r.const_grid == const_grid
            && &r.compile_options == compile_options
            && r.generics.as_slice() == generics
            && r.specs.len() == specs.len()
            && r.specs.iter().zip(specs).all(|(a, b)| a == *b)
            && r.scalar_hints.as_slice() == scalar_hints
        {
            Some((Arc::clone(&r.function), Arc::clone(&r.validator)))
        } else {
            None
        }
    }

    /// Replaces the cached resolution (single entry: last one wins).
    pub fn store(&self, resolution: SiteResolution) {
        if let Ok(mut guard) = self.inner.write() {
            *guard = Some(std::sync::Arc::new(resolution));
        }
    }
}

impl Default for LaunchSite {
    fn default() -> Self {
        Self::new()
    }
}

// ── Global kernel cache (process-wide, cross-thread) ────────────────────────

/// Global kernel cache. `DashMap` for cross-thread sharing; inner `OnceCell` for
/// single-flight compilation dedup (if multiple threads need the same kernel,
/// only one compiles while the rest wait). `once_cell::sync::OnceCell` gives
/// fallible initialization (`get_or_try_init`).
///
/// Keyed on the whole [`TileFunctionKey`], not on a digest of it: 64 bits collide
/// often enough to matter once a process caches many kernels, and a collision
/// here hands back a cubin compiled for a different kernel.
///
/// Intentionally unbounded: no cap or LRU. Capacity management lives in the L2
/// disk cache, not here — the same shape as cutile-python (unbounded in-memory
/// kernel cache, 2 GiB LRU on disk). Bounding L1 is a harder problem than L2:
/// evicting a `CompiledKernel` unloads its `Module`, which may still be
/// executing on the GPU, whereas deleting an L2 file is always safe.
static KERNEL_CACHE: OnceLock<DashMap<TileFunctionKey, Arc<OnceCell<CompiledKernel>>>> =
    OnceLock::new();

/// The process-global L1 cache. Crate-internal ONLY: exposing the raw
/// `DashMap` publicly would hand out safe `.clear()`/`.remove()`/`.retain()`,
/// which bypass the `unsafe` eviction gate below — safe code could unload a
/// `Module` mid-launch (the exact UAF the gate prevents) or re-create the
/// re-entrant `.retain()` deadlock. All external mutation must go through the
/// `unsafe` eviction APIs.
pub(crate) fn get_kernel_cache() -> &'static DashMap<TileFunctionKey, Arc<OnceCell<CompiledKernel>>>
{
    KERNEL_CACHE.get_or_init(DashMap::new)
}

/// Clears L1. Test-support only — `#[doc(hidden)]`, not public API, and
/// `unsafe` for the same reason as the eviction APIs: it can unload a
/// `Module` still executing on the GPU, so the caller must quiesce first.
/// Tests control launch timing, so they satisfy that obligation.
///
/// # Safety
/// See [`clear_kernel_cache`]: quiesce every stream that may run a cached
/// kernel before calling.
#[doc(hidden)]
pub unsafe fn clear_kernel_cache_for_tests() {
    get_kernel_cache().clear();
    bump_kernel_cache_epoch();
}

/// Generation counter for the in-memory cache: bumped by every removal so
/// launch-site caches (which hold their own `Arc<Function>`) re-resolve
/// instead of keeping evicted modules alive past a quiesced clear.
static KERNEL_CACHE_EPOCH: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

fn kernel_cache_epoch() -> u64 {
    KERNEL_CACHE_EPOCH.load(std::sync::atomic::Ordering::Acquire)
}

fn bump_kernel_cache_epoch() {
    KERNEL_CACHE_EPOCH.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
}

/// Removes every kernel from the process-global in-memory cache.
///
/// Entries removed here drop the cache's reference; the underlying CUDA
/// module unloads (releasing its device memory) when the LAST holder
/// drops, so host-side users are protected by refcount. What refcounts
/// cannot see is the GPU: a launched kernel executes after the launch
/// call returns. A tuning objective between trials is exactly the
/// situation this API exists for: sweeps churn specializations by design,
/// each holding device memory, while the cache is intentionally unbounded
/// for steady-state engines (capacity policy lives in the L2 disk cache).
///
/// In-flight compiles are unaffected: a thread mid-compile holds its own
/// `Arc` to its single-flight slot and completes into it; the next
/// request for that key recompiles (or is served by the disk cache).
///
/// Returns the number of entries removed. Freed device bytes are not
/// tracked host-side; per-module sizes are not observable through the
/// driver's module API.
///
/// # Safety
/// The caller must quiesce first: synchronize every stream that may still
/// be running any cached kernel before calling. Unloading a `Module` whose
/// grid is still executing on the device is undefined behavior, and host
/// refcounts cannot observe in-flight GPU work — only the caller knows
/// which streams are idle. This is why the eviction APIs are `unsafe`.
#[cfg(feature = "experimental-tune")]
pub unsafe fn clear_kernel_cache() -> usize {
    unsafe { retain_kernels(|_| false) }
}

/// Removes one specialization from the in-memory cache; returns whether
/// it was present.
///
/// # Safety
/// Same quiesce obligation as [`clear_kernel_cache`]: the caller must
/// ensure no stream is still running this kernel before evicting it.
#[cfg(feature = "experimental-tune")]
pub unsafe fn evict_kernel(key: &TileFunctionKey) -> bool {
    let removed = get_kernel_cache().remove(key).is_some();
    if removed {
        bump_kernel_cache_epoch();
    }
    removed
}

/// Keeps only specializations whose key satisfies `pred`; returns the
/// number of entries removed.
///
/// `pred` may freely query the cache (`contains_cuda_function`,
/// `evict_kernel`, or even trigger a compile): it runs with no cache lock
/// held. Use the [`TileFunctionKey::module_name`]/
/// [`TileFunctionKey::function_name`] accessors to scope a predicate to
/// your own kernel.
///
/// # Safety
/// Same quiesce obligation as [`clear_kernel_cache`].
#[cfg(feature = "experimental-tune")]
pub unsafe fn retain_kernels(mut pred: impl FnMut(&TileFunctionKey) -> bool) -> usize {
    let cache = get_kernel_cache();
    // Snapshot every key first, fully draining the iterator so no shard lock
    // is held, THEN evaluate `pred` and remove. Evaluating `pred` inside
    // `DashMap::retain` (or during iteration) holds a shard lock, so the
    // instant the predicate re-enters the cache — a lookup, an evict, a
    // compile — it self-deadlocks that shard, wedging every subsequent JIT
    // lookup in the process. Useful predicates ("evict only my kernel's
    // specializations") want exactly that re-entry.
    let keys: Vec<TileFunctionKey> = cache.iter().map(|entry| entry.key().clone()).collect();
    let mut removed = 0;
    for key in keys {
        if !pred(&key) && cache.remove(&key).is_some() {
            removed += 1;
        }
    }
    if removed > 0 {
        bump_kernel_cache_epoch();
    }
    removed
}

/// Get (or create) the single-flight compilation slot for `key`.
///
/// The returned `OnceCell` lets the caller `get_or_try_init` the compile
/// exactly once across threads. The DashMap shard lock is released before
/// this returns, so the slow compile never holds it.
///
/// Hits take the read path (shard read lock, no allocation); only a miss falls
/// back to `entry()` (write lock + owned key).
pub fn kernel_cache_slot(key: &TileFunctionKey) -> Arc<OnceCell<CompiledKernel>> {
    let cache = get_kernel_cache();
    if let Some(existing) = cache.get(key) {
        return Arc::clone(existing.value());
    }
    // `get` returned None holding no lock, so the write path is deadlock-free;
    // `or_insert_with` still resolves a concurrent insert into one slot per key.
    Arc::clone(
        cache
            .entry(key.clone())
            .or_insert_with(|| Arc::new(OnceCell::new()))
            .value(),
    )
}

/// Check whether a kernel with the given key has already been compiled and cached.
pub fn contains_cuda_function(key: &TileFunctionKey) -> bool {
    get_kernel_cache()
        .get(key)
        .is_some_and(|slot| slot.value().get().is_some())
}

/// Reads Tile IR text from a file.
///
/// This helper function reads intermediate representation files from disk, the
/// counterpart of the `dump_mlir_dir` entry attribute's [`write_ir`].
///
/// ## Parameters
///
/// - `path`: Path to the IR file to read
///
/// ## Returns
///
/// The file contents as a UTF-8 string, or an I/O error if reading fails.
#[expect(unused)]
fn read_ir(path: String) -> Result<String, std::io::Error> {
    let s = String::from_utf8(fs::read(path)?).expect("Unable to convert from utf8 to string.");
    Ok(s)
}

/// Writes Tile IR text to a file for debugging.
///
/// This helper function writes intermediate representation to disk when kernel functions
/// are marked with `dump_mlir_dir` entry attributes. The filename
/// includes the module name, function name, and cache hash for uniqueness.
///
/// ## Parameters
///
/// - `module_name`: Name of the module containing the kernel
/// - `function_name`: Name of the kernel function
/// - `cache_hash_str`: Unique hash identifying this compilation
/// - `extension`: File extension (usually "mlir" for the MLIR-like Tile IR text)
/// - `dir`: Directory to write the file to
/// - `contents`: IR contents to write
///
/// ## Errors
///
/// Returns an error if the file cannot be written (e.g. the directory does not
/// exist or is not writable). This runs inside the single-flight compile, so a
/// panic here would poison the launch instead of failing it.
fn write_ir(
    module_name: &str,
    function_name: &str,
    cache_hash_str: &str,
    extension: &str,
    dir: &str,
    contents: &str,
) -> Result<(), Error> {
    let filename = format!("{module_name}_{function_name}_{cache_hash_str}.{extension}");
    let path = PathBuf::from(dir).join(filename);
    fs::write(&path, contents).map_err(|e| {
        Error::Anyhow(anyhow::anyhow!(
            "failed to write the IR dump for {module_name}::{function_name} to {path:?} \
             (dump_mlir_dir = {dir:?}): {e}"
        ))
    })?;
    println!("IR written to {path:?}");
    Ok(())
}

// ── Single-flight compilation dedup is handled by once_cell::sync::OnceCell ──

/// Compiles one tile-function specialization to a CUBIN and loads it into a
/// [`CompiledKernel`].
///
/// This is the single compile-and-load core behind [`compile_from_context`],
/// which serves both real `.sync()` / `.await` launches and the `.compile()`
/// warmup terminal. It runs the compiler, honors the `print_ir` /
/// `dump_mlir_dir` entry attributes,
/// lowers to a CUBIN, loads the module, and resolves `function_entry`, emitting
/// per-stage `CUTILE_JIT_TIMING` along the way.
///
/// Callers own the cache concerns: they build the [`TileFunctionKey`], dedup via
/// the cache slot, and call [`record_jit_compile`]. This function assumes it
/// runs exactly once per cache miss and does no caching itself.
#[allow(clippy::too_many_arguments)]
fn compile_and_load_kernel(
    modules: &CUDATileModules,
    module_name: &str,
    function_name: &str,
    function_entry: &str,
    generics: &[String],
    stride_args: &[(String, Vec<i32>)],
    spec_args: &[(String, SpecializationBits)],
    scalar_hints: &[(String, DivHint)],
    const_grid: Option<(u32, u32, u32)>,
    gpu_name: &str,
    compile_options: &CompileOptions,
    device_id: usize,
    key_str: &str,
) -> Result<CompiledKernel, Error> {
    let t0 = std::time::Instant::now();

    let stride_args_refs: Vec<(&str, &[i32])> = stride_args
        .iter()
        .map(|x| (x.0.as_str(), x.1.as_slice()))
        .collect();
    let spec_args_refs: Vec<(&str, &SpecializationBits)> =
        spec_args.iter().map(|x| (x.0.as_str(), &x.1)).collect();
    let scalar_hints_refs: Vec<(&str, &DivHint)> =
        scalar_hints.iter().map(|x| (x.0.as_str(), &x.1)).collect();

    let stage1_start = std::time::Instant::now();
    let (tile_module, validator, check_stats) = {
        let compiler = CUDATileFunctionCompiler::new(
            modules,
            module_name,
            function_name,
            generics,
            &stride_args_refs,
            &spec_args_refs,
            &scalar_hints_refs,
            const_grid,
            gpu_name.to_string(),
            compile_options,
        )?;
        let tile_module = compiler.compile()?;
        // AFTER compile, not before: the launch-check accumulator fills
        // DURING compilation, and this snapshot is the one the generated
        // launcher enforces. Taken early, every hoisted check is silently
        // dropped at launch while the compiler has already discharged the
        // in-kernel assert on its promise — out-of-bounds accesses then run
        // unchecked (caught by the differential placement harness; pinned by
        // `launch_checks_are_enforced_at_launch`).
        let validator = Arc::new(compiler.get_validator());
        let check_stats = (
            compiler.check_stats.discharged.get(),
            compiler.check_stats.hoisted.get(),
            compiler.check_stats.in_place.get(),
        );
        (tile_module, validator, check_stats)
    };
    let stage1_ms = stage1_start.elapsed().as_secs_f64() * 1000.0;

    let stage2_start = std::time::Instant::now();
    {
        let print_ir =
            modules.get_entry_arg_bool_by_function_name(module_name, function_name, "print_ir")?;
        let dump_mlir_dir = modules.get_entry_arg_string_by_function_name(
            module_name,
            function_name,
            "dump_mlir_dir",
        )?;
        // `to_mlir_text` renders the whole module; only pay for it when asked.
        if print_ir || dump_mlir_dir.is_some() {
            let ir_text = tile_module.to_mlir_text();
            if print_ir {
                println!("COMPILED IR: {module_name}::{function_name}\n{ir_text}");
            }
            if let Some(path) = dump_mlir_dir {
                write_ir(
                    module_name,
                    function_name,
                    key_str,
                    "mlir",
                    path.as_str(),
                    ir_text.as_str(),
                )?;
            }
        }
    }
    let (bytecode, bc_version) = serialize_tile_ir_bytecode(&tile_module)?;
    let tileiras_opts = TileirasOptions::from_compile_options(compile_options);
    let (cubin, mut stage2_source) =
        compile_bytecode_cached(&bytecode, bc_version, gpu_name, &tileiras_opts)?;
    let mut stage2_ms = stage2_start.elapsed().as_secs_f64() * 1000.0;

    // A retry recompile (below) runs inside the stage-3 window but is really
    // stage-2 work; track it so the timing line attributes it to stage2 (which
    // then reports source=tileiras) instead of inflating stage3.
    let mut recompile_ms = 0.0;
    let stage3_start = std::time::Instant::now();
    // SAFETY: `cubin` is a complete image. It either came straight out of
    // tileiras in this process, or from the disk cache, whose read path
    // verifies the entry's SHA-256 payload checksum before handing the bytes
    // back (a torn or corrupt entry is a miss, never a load).
    let module = match unsafe { load_module_from_bytes(&cubin, device_id) } {
        Ok(module) => module,
        // A disk-served cubin the driver rejects (partial write the checksum
        // missed, driver/toolkit skew, …) must not fail the launch: evict that
        // exact entry and recompile with tileiras, bypassing the cache read so a
        // still-present bad entry can't be re-served. `mem::replace` moves the
        // store/key out and leaves `Tileiras`, which is now the true source of
        // the loaded cubin. Only a second failure is a real error.
        Err(e) => match std::mem::replace(&mut stage2_source, Stage2Source::Tileiras) {
            Stage2Source::DiskCache { store, key } => {
                jit_log!(
                    "{module_name}::{function_name} → cached cubin rejected by the driver ({e}); \
                     evicting and recompiling"
                );
                let recompile_start = std::time::Instant::now();
                let cubin = recompile_after_disk_rejection(
                    store.as_ref(),
                    &key,
                    &bytecode,
                    gpu_name,
                    &tileiras_opts,
                )?;
                recompile_ms = recompile_start.elapsed().as_secs_f64() * 1000.0;
                stage2_ms += recompile_ms;
                // SAFETY: a fresh tileiras compile, complete by construction.
                unsafe { load_module_from_bytes(&cubin, device_id) }?
            }
            Stage2Source::Tileiras => return Err(e.into()),
        },
    };
    let function = Arc::new(module.load_function(function_entry).map_err(|e| {
        Error::KernelLaunch(KernelLaunchError(format!(
            "failed to load '{function_entry}' from compiled cubin: {e}"
        )))
    })?);
    // Exclude the retry recompile: it was moved into stage2_ms above, so the
    // stage-3 figure stays "module load only". `max(0.0)` guards float noise.
    let stage3_ms = (stage3_start.elapsed().as_secs_f64() * 1000.0 - recompile_ms).max(0.0);

    jit_log!(
        "{module_name}::{function_name} → JIT compiled in {:.1?}",
        t0.elapsed()
    );
    if std::env::var_os("CUTILE_JIT_TIMING").is_some() {
        let stage2_source = match stage2_source {
            Stage2Source::Tileiras => "tileiras",
            Stage2Source::DiskCache { .. } => "disk",
        };
        eprintln!(
            "CUTILE_JIT_TIMING module={module_name} function={function_name} key={key_str} stage1_ms={stage1_ms:.3} stage2_ms={stage2_ms:.3} stage2_source={stage2_source} stage3_ms={stage3_ms:.3} checks_discharged={} checks_hoisted={} checks_in_place={} generics={}",
            check_stats.0,
            check_stats.1,
            check_stats.2,
            generics.join(","),
        );
    }

    Ok(CompiledKernel {
        module,
        function,
        validator,
    })
}

/// Compiles a tile function to CUDA and caches it for reuse.
///
/// Handles the complete compilation pipeline from Rust to CUDA:
/// 1. Checks the global kernel cache (process-wide, cross-thread)
/// 2. If not cached, compiles the module AST to Tile IR bytecode, then to a cubin
/// 3. Stores the result in the global kernel cache
///
/// **Compilation dedup**: When multiple threads need the same kernel, `OnceCell::get_or_try_init`
/// ensures only one thread performs compilation while others block. Once initialization completes,
/// all threads see the same cached result.
///
/// The caching key is a [`TileFunctionKey`]: module and function name, generics, stride and
/// specialization arguments, scalar hints, optional constant grid, compile options, source
/// hash, and the device/toolchain identity, ensuring correct reuse across specializations.
///
/// ## Arguments
///
/// * `ctx` - Execution context containing device information
/// * `kernel_ast` - Closure producing the kernel's module AST (the `__module_ast_self`
///   generated by `#[cutile::module]`); called only on a cache miss
/// * `module_name` - Name of the module containing the function
/// * `function_name` - Name of the function to compile
/// * `function_entry` - Entry point name in the compiled CUDA code
/// * `function_generics` - Type and const generic arguments (e.g., `["f32", "256"]`)
/// * `stride_args` - Per tensor argument, which strides are 1 (`1`) or unknown (`-1`)
/// * `spec_args` - Per tensor argument, its [`SpecializationBits`]
/// * `scalar_hints` - Divisibility hints for integer scalar and pointer arguments
/// * `const_grid` - Optional compile-time constant grid dimensions
/// * `compile_options` - Compiler and tileiras options for this specialization
/// * `source_hash` - The kernel module's `_SOURCE_HASH`
///
/// Returns the loaded [`Function`] together with the [`Validator`] the launcher checks the
/// runtime arguments against.
///
/// ## Examples
///
/// ```rust,ignore
/// use cutile::tile_kernel::{compile_from_context, CompileOptions};
///
/// let ctx = ExecutionContext::new(stream);
/// let (function, validator) = compile_from_context(
///     &ctx,
///     my_module::__module_ast_self,
///     "my_module",
///     "my_function",
///     "my_function_entry",
///     vec!["f32".to_string(), "128".to_string()],
///     vec![("x".to_string(), vec![1])],
///     vec![("x".to_string(), x.spec().clone())],
///     vec![],
///     None,
///     CompileOptions::default(),
///     my_module::_SOURCE_HASH,
/// )?;
/// ```
#[allow(clippy::too_many_arguments)]
pub fn compile_from_context<F: Fn() -> Module>(
    ctx: &ExecutionContext,
    kernel_ast: F,
    module_name: &str,
    function_name: &str,
    function_entry: &str,
    function_generics: Vec<String>,
    stride_args: Vec<(String, Vec<i32>)>,
    spec_args: Vec<(String, SpecializationBits)>,
    scalar_hints: Vec<(String, DivHint)>,
    const_grid: Option<(u32, u32, u32)>,
    compile_options: CompileOptions,
    source_hash: &str,
) -> Result<(Arc<Function>, Arc<Validator>), Error> {
    let specialization = _specialization_from_context(
        ctx,
        kernel_ast,
        module_name,
        function_name,
        function_generics,
        stride_args,
        spec_args,
        scalar_hints,
        const_grid,
        compile_options,
        source_hash,
    );
    let key = specialization.l1_cache_key().clone();
    let device_id = key.device_id;
    let gpu_name = key.gpu_name.clone();
    let slot = kernel_cache_slot(&key);

    // Use OnceCell::get_or_try_init for single-flight compilation dedup.
    // Only one thread executes the closure; others block and see the result.
    let compiled = match slot.get_or_try_init(|| -> Result<CompiledKernel, Error> {
        jit_log!("{module_name}::{function_name} → JIT compiling...");
        // Build the module ASTs lazily — only on a real cache miss.
        let modules = CUDATileModules::from_kernel((specialization.module_ast_fn)())?;
        let kernel = compile_and_load_kernel(
            &modules,
            module_name,
            function_name,
            function_entry,
            &key.function_generics,
            &key.stride_args,
            &key.spec_args,
            &key.scalar_hints,
            const_grid,
            &gpu_name,
            &key.compile_options,
            device_id,
            &key.display_hash(),
        )?;
        // Count only a successful compile: a failed attempt leaves the slot empty
        // and retries, so counting at the top would double-count on retry and
        // break the "+1 per cached kernel" contract.
        record_jit_compile();
        Ok(kernel)
    }) {
        Ok(compiled) => compiled,
        Err(e) => {
            // A failed compile leaves an empty slot; evict it so repeated failing
            // specializations don't grow the cache unbounded.
            //
            // On failure, once_cell gives the cell to a blocked waiter to retry.
            // To avoid removing the slot while that waiter is still compiling
            // (which would orphan its success and break single-flight), drop our
            // own `slot` first, then (under the shard write lock) remove only
            // when the cell is still empty and `strong_count == 1`.
            drop(slot);
            get_kernel_cache().remove_if(&key, |_, cell| {
                cell.get().is_none() && Arc::strong_count(cell) == 1
            });
            return Err(e);
        }
    };

    Ok((
        Arc::clone(&compiled.function),
        Arc::clone(&compiled.validator),
    ))
}

/// Validates that all partition grids match the expected launch grid.
pub fn validate_grids(
    grid: (u32, u32, u32),
    partition_grids: &[(u32, u32, u32)],
) -> Result<(), Error> {
    // Make sure we're not trying to map mutable references to incorrect launch grid.
    if let Some(partition_grid) = partition_grids.iter().find(|&&i| i != grid) {
        Err(Error::KernelLaunch(KernelLaunchError(format!(
            "{:?} != {:?}",
            grid, partition_grid
        ))))
    } else {
        Ok(())
    }
}

/// Validates the launch grid against every binding's [`GridBound`]:
/// exact-coverage bindings must equal the grid; partial-coverage
/// (`partition_prefix`) bindings must bound it per axis. `launch <= bound`
/// per axis is the sound direction — a per-axis prefix embeds identically
/// into the block grid, uncovered blocks are simply never visited — while
/// `launch > bound` on ANY axis is genuine out-of-bounds and always an
/// error. Per-axis, never total-count: delinearizing against a different
/// grid shape would remap CTAs to the wrong blocks.
pub fn validate_grid_bounds(grid: (u32, u32, u32), bounds: &[GridBound]) -> Result<(), Error> {
    for bound in bounds {
        let GridBound::Exact(expected) = bound else {
            continue;
        };
        if *expected != grid {
            return Err(Error::KernelLaunch(KernelLaunchError(format!(
                "launch grid {:?} does not match the inferred partition grid {:?}",
                grid, expected
            ))));
        }
    }
    for bound in bounds {
        let GridBound::AtMost(max) = bound else {
            continue;
        };
        let launch = [grid.0, grid.1, grid.2];
        let max_axes = [max.0, max.1, max.2];
        if let Some(axis) = (0..3).find(|&k| launch[k] > max_axes[k]) {
            return Err(Error::KernelLaunch(KernelLaunchError(format!(
                "launch grid {:?} exceeds the partial-coverage partition grid {:?} on axis {axis}",
                grid, max
            ))));
        }
    }
    Ok(())
}

/// Runs the full set of launch-time checks before `cuLaunchKernel`: the
/// built-in grid family (all partition grids match the launch grid) plus any
/// compiler-emitted checks hoisted out of the device kernel.
///
/// `param_shapes[i]` is the runtime extent vector of the i-th kernel parameter
/// (empty for non-tensor params). This is the host end of launch-time check
/// hoisting: the compiler evacuated these checks from the kernel, so they run
/// here once per launch instead of per-thread on the device. `validate_grids`
/// is folded in as the first, always-present family; compiler-emitted checks
/// are canonical [`Predicate`]s evaluated against the parameter extents.
pub fn validate_launch(
    launch_checks: &[LaunchCheck],
    grid: (u32, u32, u32),
    partition_bounds: &[GridBound],
    param_shapes: &[Vec<i32>],
    view_shapes: &[Vec<i32>],
) -> Result<(), Error> {
    // Built-in family: launch grid vs. partition grid bounds.
    validate_grid_bounds(grid, partition_bounds)?;
    // Compiler-emitted families (empty unless a kernel hoisted a check).
    for check in launch_checks {
        evaluate_launch_check(check, param_shapes, view_shapes, grid)?;
    }
    Ok(())
}

/// Runs only the compiler-emitted launch checks (the grid family is already
/// validated by `infer_launch_grid`). Called from the generated launcher after
/// grid inference, both arrays indexed in signature order (empty for
/// non-tensor params):
/// - `param_shapes[i]` — the i-th parameter's *root* extents (the whole
///   tensor). Resolves [`Atom::Dim`], the frame declared `preconditions` are
///   stated in.
/// - `view_shapes[i]` — the i-th parameter's *kernel-visible view* extents:
///   the partition slab for a `&mut Tensor` output, the whole tensor
///   otherwise. Resolves [`Atom::ViewExtent`].
/// - `launch_grid` — the grid the kernel will actually be launched with.
///   Resolves [`Atom::NumTileBlocks`]: the block-id axiom rung discharges
///   `tile_block_id(k)` accesses in the kernel against a launch check over
///   this exact grid, so validating any other grid would unsound the rung.
pub fn validate_launch_checks(
    launch_checks: &[LaunchCheck],
    param_shapes: &[Vec<i32>],
    view_shapes: &[Vec<i32>],
    launch_grid: (u32, u32, u32),
) -> Result<(), Error> {
    for check in launch_checks {
        evaluate_launch_check(check, param_shapes, view_shapes, launch_grid)?;
    }
    Ok(())
}

/// Evaluates one hoisted [`LaunchCheck`] by interpreting its canonical
/// [`Predicate`] against the runtime parameter extents, each atom against the
/// array holding its frame. Fails closed: a predicate whose atoms cannot be
/// resolved (a missing parameter/axis, or a non-launch-known `Iv` atom that
/// should never appear here) is an error, not a silent skip.
fn evaluate_launch_check(
    check: &LaunchCheck,
    param_shapes: &[Vec<i32>],
    view_shapes: &[Vec<i32>],
    launch_grid: (u32, u32, u32),
) -> Result<(), Error> {
    // Resolve each atom to its runtime value, in the atom's own frame.
    let resolve_atom = |atom: &Atom| -> Option<i64> {
        match atom {
            Atom::Dim { param, axis } => param_shapes
                .get(*param)
                .and_then(|shape| shape.get(*axis))
                .map(|&extent| extent as i64),
            Atom::ViewExtent { param, axis } => view_shapes
                .get(*param)
                .and_then(|shape| shape.get(*axis))
                .map(|&extent| extent as i64),
            // ceil(root extent / tile). The mint site guarantees tile >= 1;
            // fail closed on a malformed atom rather than dividing by zero.
            Atom::TileCount { param, axis, tile } => {
                if *tile < 1 {
                    return None;
                }
                param_shapes
                    .get(*param)
                    .and_then(|shape| shape.get(*axis))
                    .map(|&extent| (extent as i64 + *tile as i64 - 1) / *tile as i64)
            }
            // The grid axis extents: the host fixes the grid before launch,
            // and the block-id axiom rung's checks are stated over it. Only
            // three grid axes exist; anything else fails closed.
            Atom::NumTileBlocks(k) => match k {
                0 => Some(launch_grid.0 as i64),
                1 => Some(launch_grid.1 as i64),
                2 => Some(launch_grid.2 as i64),
                _ => None,
            },
            // A device-runtime induction variable and the block-id register
            // are not launch-known; they never appear in a launch check (the
            // axiom rung replaces the block id with its grid bound), so fail
            // closed if one somehow does.
            Atom::Iv(_) | Atom::TileBlockId(_) => None,
        }
    };
    match check.predicate.eval(&resolve_atom) {
        Some(true) => Ok(()),
        Some(false) => Err(Error::KernelLaunch(KernelLaunchError(format!(
            "launch check failed: {}",
            check.cause
        )))),
        None => Err(Error::KernelLaunch(KernelLaunchError(format!(
            "launch check has unresolved operands (extent unavailable at launch): {}",
            check.cause
        )))),
    }
}

/// Infers the launch grid for a kernel from partitioned tensor inputs.
///
/// If a grid is explicitly specified (non-zero), it is used directly. Otherwise, the grid
/// is inferred from partitioned tensor inputs. All inferred grids must match, or the
/// function will return an error.
///
/// ## Errors
///
/// Returns an error if no grid is specified and no inferred grids are available, or if inferred
/// grids from different inputs don't match.
pub fn infer_launch_grid(
    grid: (u32, u32, u32),
    bounds: &[GridBound],
) -> Result<(u32, u32, u32), Error> {
    let exact: Vec<(u32, u32, u32)> = bounds
        .iter()
        .filter_map(|b| match b {
            GridBound::Exact(g) => Some(*g),
            GridBound::AtMost(_) => None,
        })
        .collect();
    if grid != (0, 0, 0) {
        // A launch grid was specified.
        validate_grid_bounds(grid, bounds)?;
        return Ok(grid);
    }
    // Try to infer the launch grid. Only an EXACT binding can define it: a
    // partial-coverage binding is an upper bound, and inferring the bound
    // itself would silently reconstruct full coverage — the thing the
    // caller opted out of.
    if exact.is_empty() {
        if bounds.is_empty() {
            return kernel_launch_error_result("Launch grid required.");
        }
        return kernel_launch_error_result(
            "Launch grid required: a partial-coverage (partition_prefix) binding \
             only bounds the grid; specify the grid explicitly or bind with \
             partition().",
        );
    }
    let grid = exact[0];
    validate_grid_bounds(grid, bounds)?;
    Ok(grid)
}

/// A compiled CUDA kernel generated from Rust code that can be launched on the GPU.
///
/// `TileKernel` extends [`DeviceOp`] with kernel-specific functionality. Kernels are
/// automatically generated from Rust functions marked with `#[cutile::entry]` and compiled
/// to Tile IR bytecode, then to a CUDA cubin at runtime.
///
/// The trait provides methods for configuring kernel launch parameters such as grid dimensions,
/// type generics, and shared memory. Grid dimensions can be set explicitly or inferred from
/// partitioned tensor inputs.
///
/// ## Examples
///
/// ### Basic kernel launch
///
/// ```rust,ignore
/// #[cutile::module]
/// mod my_module {
///     use cutile::core::*;
///
///     #[cutile::entry]
///     fn hello_world() {
///         let pid = get_tile_block_id();
///         cuda_tile_print!("Hello from block {}\n", pid.0);
///     }
/// }
///
/// // Launch with explicit grid
/// my_module::hello_world()
///     .grid((4, 1, 1))
///     .sync_on(&stream)?;
/// ```
///
/// ### Kernel with arguments and grid inference
///
/// ```rust,ignore
/// // Output-first convention: &mut param is the first argument.
/// // Grid is inferred from partitioned tensors.
/// // The unified launcher accepts both plain values and DeviceOps.
/// let result = add(
///     api::zeros(&[256]).partition([64]),
///     api::ones(&[256]),
///     api::ones(&[256]),
/// )
/// .first()        // extract the &mut output
/// .unpartition()  // recover Tensor from Partition
/// .to_host_vec()
/// .sync()?;
/// ```
///
/// ### Using with async composition
///
/// ```rust,ignore
/// async fn pipeline() -> impl DeviceOp<Output=Tensor<f32>> {
///     let x = api::randn(0.0, 1.0, [128, 128]).await;
///
///     // Chain kernel operations
///     let y = my_kernel_1(x.clone())
///         .grid((8, 8, 1))
///         .await;
///
///     let z = my_kernel_2(y)
///         .grid((4, 4, 1))
///         .await;
///
///     z
/// }
/// ```
pub trait TileKernel<ARGS: Send, DI, STORED: Send = ARGS>: DeviceOp<Output = ARGS>
where
    DI: DeviceOp<Output = STORED>,
{
    /// Compiles the kernel from its module AST, returning the CUDA function
    /// and validator.
    ///
    /// This is the internal compile-and-cache entry point used by the generated
    /// launcher (both the `.sync()`/`.await` launch path and the `.compile()`
    /// warmup terminal). The user-facing `.compile()` terminal is a separate,
    /// no-argument method generated per kernel; this one keeps the descriptive
    /// name `jit_compile` so it does not collide with it.
    ///
    /// `kernel_ast` is invoked once on cache miss to obtain the kernel's own
    /// [`Module`] (typically the macro-generated `__module_ast_self` fn).
    /// Dep modules are discovered by walking the kernel's `use` statements
    /// against the linker registry.
    #[allow(clippy::too_many_arguments)]
    fn jit_compile<F: Fn() -> Module>(
        &mut self,
        ctx: &ExecutionContext,
        kernel_ast: F,
        module_name: &str,
        function_name: &str,
        function_entry: &str,
        function_generics: Vec<String>,
        stride_args: Vec<(String, Vec<i32>)>,
        spec_args: Vec<(String, SpecializationBits)>,
        scalar_hints: Vec<(String, DivHint)>,
        grid: Option<(u32, u32, u32)>,
        compile_options: CompileOptions,
        source_hash: &str,
    ) -> Result<(Arc<Function>, Arc<Validator>), Error> {
        compile_from_context(
            ctx,
            kernel_ast,
            module_name,
            function_name,
            function_entry,
            function_generics,
            stride_args,
            spec_args,
            scalar_hints,
            grid,
            compile_options,
            source_hash,
        )
    }
    /// Sets the type and const generic arguments for this kernel.
    fn generics(self, generics: Vec<String>) -> Self;
    /// Sets a compile-time constant grid, enabling grid-dependent optimizations.
    fn const_grid(self, grid: (u32, u32, u32)) -> Self;
    /// Sets the runtime launch grid dimensions.
    fn grid(self, grid: (u32, u32, u32)) -> Self;
    /// Sets the runtime compile options (occupancy, num_cta_in_cga).
    fn compile_options(self, options: CompileOptions) -> Self;
    /// Infers the launch grid from partitioned tensor inputs, or uses the explicit grid.
    fn infer_launch_grid(&self, bounds: &[GridBound]) -> Result<(u32, u32, u32), Error> {
        let grid = self.get_launch_grid();
        infer_launch_grid(grid, bounds)
    }
    /// Returns the currently configured launch grid dimensions.
    fn get_launch_grid(&self) -> (u32, u32, u32);
    /// Returns the dynamic shared memory size in bytes. Defaults to 0.
    fn get_launch_smem(&self) -> u32 {
        0
    }
    /// Returns the thread block dimensions. Defaults to `(1, 1, 1)`.
    fn get_launch_block(&self) -> (u32, u32, u32) {
        (1, 1, 1)
    }
    // fn validate(validator: &Validator) -> Result<(), Error> {

    // }
    // fn validate_arc<T: DType>(
    //     &self,
    //     func_name: String,
    //     var_name: String,
    //     arc: &Arc<Tensor<T>>,
    //     shape: &[i32],
    // ) -> Result<(), KernelLauncherError> {
    //     let input_shape = &arc.shape;
    //     if input_shape != shape {
    //         return Err(KernelLauncherError::InvalidTensorShape(format!(
    //             "Unexpected shape {:?} for argument {} for function {}.",
    //             input_shape, var_name, func_name
    //         )));
    //     }
    //     Ok(())

    //     // if input_shape.len() != shape.len() {
    //     //     return Err(KernelLauncherError::InvalidTensorShape(format!("Unexpected rank {} for argument {} for function {}.",
    //     //         input_shape.len(),
    //     //         var_name,
    //     //         func_name
    //     //     )));
    //     // }
    //     // for i in 0..input_shape.len() {
    //     //     let input_dim = input_shape[i];
    //     //     let param_dim = shape[i];
    //     //     if param_dim == -1 {
    //     //         continue;
    //     //     }
    //     //     if input_dim != param_dim {
    //     //         return Err(KernelLauncherError::InvalidTensorShape(format!("Unexpected rank {} for argument {} for function {}.",
    //     //             input_shape.len(),
    //     //             var_name,
    //     //             func_name
    //     //         )));
    //     //     }
    //     // }
    // }
}

/// Implements kernel argument passing for `Tensor` when wrapped in `Arc`.
///
/// Pushes the device pointer, shape, and stride information to the kernel launcher
/// in the order expected by compiled tile functions.
impl<T: DType> ArcKernelArgument for Tensor<T> {
    fn push_arg_arc(self: &Arc<Self>, launcher: &mut AsyncKernelLaunch) {
        // TODO (hme): document safety
        unsafe {
            launcher.push_device_ptr(self.cu_deviceptr());
        }
        for dim in self.shape.iter() {
            launcher.push_arg(*dim);
        }
        for stride in self.strides.iter() {
            launcher.push_arg(*stride);
        }
    }
}

/// Implements kernel argument passing for partitioned tensors.
///
/// Pushes the device pointer, tensor shape and strides, followed by partition shape
/// and strides. This allows kernels to access both the full tensor and the partition
/// information for block-level indexing.
impl<T: DType> KernelArgument for &Partition<Tensor<T>> {
    fn push_arg(self, launcher: &mut AsyncKernelLaunch) {
        // TODO (hme): document safety
        unsafe {
            launcher.push_device_ptr(self.object.cu_deviceptr());
        }
        for dim in self.object.shape.iter() {
            launcher.push_arg(*dim);
        }
        for stride in self.object.strides.iter() {
            launcher.push_arg(*stride);
        }
        for dim in self.partition_shape.iter() {
            launcher.push_arg(*dim as i32);
        }
        for stride in self.partition_strides.iter() {
            launcher.push_arg(*stride as i32);
        }
    }
}

/// Same as above but for borrowed mutable tensor partitions.
impl<T: DType> KernelArgument for &Partition<&mut Tensor<T>> {
    fn push_arg(self, launcher: &mut AsyncKernelLaunch) {
        unsafe {
            launcher.push_device_ptr(self.object.cu_deviceptr());
        }
        for dim in self.object.shape.iter() {
            launcher.push_arg(*dim);
        }
        for stride in self.object.strides.iter() {
            launcher.push_arg(*stride);
        }
        for dim in self.partition_shape.iter() {
            launcher.push_arg(*dim as i32);
        }
        for stride in self.partition_strides.iter() {
            launcher.push_arg(*stride as i32);
        }
    }
}

// Partition

/// Extension trait that enables partitioning device operations into tiles.
///
/// This trait allows async operations that produce tensors to be partitioned before
/// execution, enabling automatic grid inference for tile kernels. The partition divides
/// the tensor into blocks that map to CUDA thread blocks.
///
/// ## Examples
///
/// ```rust,ignore
/// use cutile::tile_kernel::PartitionOp;
///
/// // Partition a tensor operation before it executes
/// let x = api::ones(&[1024]).partition([128]);  // Creates 8 partitions
///
/// // Use partitioned tensors with kernels for automatic grid inference
/// let y = api::randn(0.0, 1.0, [256, 256]).partition([64, 64]);  // 4x4 grid
/// let result = my_kernel(y).await;  // Grid (4, 4, 1) inferred automatically
/// ```
pub trait PartitionOp<I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
    /// Partitions the output of this device operation into tiles of the given shape.
    ///
    /// The partition shape determines how the tensor is divided across CUDA thread blocks.
    fn partition<const RANK: usize>(
        self,
        partition_shape: [usize; RANK],
    ) -> DeviceOperationPartition<RANK, I, DI>;
}

impl<I, DI> PartitionOp<I, DI> for DI
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
    fn partition<const RANK: usize>(
        self,
        partition_shape: [usize; RANK],
    ) -> DeviceOperationPartition<RANK, I, DI>
    where
        Self: Sized,
    {
        DeviceOperationPartition::<RANK, I, DI> {
            partition_shape,
            op: self,
        }
    }
}

/// A device operation that partitions its output into tiles.
///
/// This wrapper executes the underlying device operation and then partitions its result
/// according to the specified partition shape. The resulting partitioned tensor can be
/// used with tile kernels to automatically infer launch grid dimensions.
///
/// Created by calling `.partition()` on any device operation that produces a partitionable output.
///
/// ## Examples
///
/// ```rust,ignore
/// // Create a partitioned tensor operation
/// let z = api::zeros(&[1024]).partition([64]);
///
/// // Pass directly to kernel — grid inferred from partition
/// let result = my_kernel(z, x, y).first().unpartition().sync()?;
/// ```
pub struct DeviceOperationPartition<const RANK: usize, I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
    partition_shape: [usize; RANK],
    op: DI,
}

unsafe impl<const RANK: usize, I, DI> Send for DeviceOperationPartition<RANK, I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
}

impl<const RANK: usize, I, DI> DeviceOp for DeviceOperationPartition<RANK, I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
    type Output = Partition<I>;

    unsafe fn execute(
        self,
        context: &ExecutionContext,
    ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
        let val = self.op.execute(context)?;
        Ok(val.partition(self.partition_shape))
    }
}

impl<const RANK: usize, I, DI> IntoFuture for DeviceOperationPartition<RANK, I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I>,
{
    type Output = Result<Partition<I>, DeviceError>;
    type IntoFuture = DeviceFuture<Partition<I>, DeviceOperationPartition<RANK, I, DI>>;
    fn into_future(self) -> Self::IntoFuture {
        match with_default_device_policy(|policy| {
            let stream = policy.next_stream()?;
            Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
        }) {
            Ok(Ok(future)) => future,
            Ok(Err(e)) => DeviceFuture::failed(e),
            Err(e) => DeviceFuture::failed(e),
        }
    }
}

// Unwrap Partition

/// A device operation that unwraps a partitioned tensor back to a regular tensor.
///
/// This operation removes the partition structure from a tensor, converting a
/// `Partition<Tensor<T>>` back to `Tensor<T>`. This is useful after kernel operations
/// that work on partitioned inputs but need to return regular tensors for further
/// processing.
///
/// Created by calling `unwrap_partition()` on a device operation that produces a partition.
///
/// ## Examples
///
/// ```rust,ignore
/// use cutile::tile_kernel::unwrap_partition;
///
/// // After a kernel operation on partitioned tensors
/// let x = api::ones(&[256]).partition([64]);
/// let y = my_kernel(x).await;  // Returns Partition<Tensor<f32>>
///
/// // Unwrap back to a regular tensor
/// let z = unwrap_partition(y).await;  // Now Tensor<f32>
/// ```
pub struct UnwrapPartition<I: Send, DI>
where
    DI: DeviceOp<Output = Partition<I>>,
{
    pub(crate) op: DI,
}

unsafe impl<I: Send, DI> Send for UnwrapPartition<I, DI> where DI: DeviceOp<Output = Partition<I>> {}

impl<I: Send, DI> DeviceOp for UnwrapPartition<I, DI>
where
    DI: DeviceOp<Output = Partition<I>>,
{
    type Output = I;

    unsafe fn execute(
        self,
        context: &ExecutionContext,
    ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
        let val = self.op.execute(context)?;
        Ok(val.unpartition())
    }
}

impl<I: Send, DI> IntoFuture for UnwrapPartition<I, DI>
where
    DI: DeviceOp<Output = Partition<I>>,
{
    type Output = Result<I, DeviceError>;
    type IntoFuture = DeviceFuture<I, UnwrapPartition<I, DI>>;
    fn into_future(self) -> Self::IntoFuture {
        match with_default_device_policy(|policy| {
            let stream = policy.next_stream()?;
            Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
        }) {
            Ok(Ok(future)) => future,
            Ok(Err(e)) => DeviceFuture::failed(e),
            Err(e) => DeviceFuture::failed(e),
        }
    }
}

/// Unwraps a partitioned device operation back to a regular tensor operation.
///
/// Converts a device operation that produces a `Partition<T>` into one
/// that produces `T` directly. Useful for converting partitioned kernel outputs
/// back to regular tensors for further processing.
///
/// ## Examples
///
/// ```rust,ignore
/// use cutile::tile_kernel::unwrap_partition;
///
/// async fn process_data() -> Tensor<f32> {
///     let x = api::randn(0.0, 1.0, [1024]).partition([128]);
///     let processed = my_tiled_kernel(x);  // Returns Partition<Tensor<f32>>
///
///     // Unwrap to get a regular tensor
///     unwrap_partition(processed).await
/// }
/// ```
pub fn unwrap_partition<I: Send, DI>(op: DI) -> UnwrapPartition<I, DI>
where
    DI: DeviceOp<Output = Partition<I>>,
{
    UnwrapPartition { op }
}

// Partitioning and unpartitioning allocate nothing themselves, so they are
// graph-recordable exactly when the op they wrap is.
impl<const RANK: usize, I, DI> GraphNode for DeviceOperationPartition<RANK, I, DI>
where
    I: Send + IntoPartition + IntoPartitionArc,
    DI: DeviceOp<Output = I> + GraphNode,
{
}

impl<I: Send, DI> GraphNode for UnwrapPartition<I, DI> where
    DI: DeviceOp<Output = Partition<I>> + GraphNode
{
}

// ── Launcher input combinators ──────────────────────────────────────────────
//
// The generated launcher `my_kernel(arg0, arg1, ..)` turns each argument into
// a DeviceOp (`IntoDeviceOp`), applies KernelOutput::prepare / KernelInput::
// prepare to tensor params, and hands the launcher struct one op producing the
// whole argument tuple. These three types are that op, spelled with nameable
// types so the launcher's return type can carry it: `Launcher<.., KernelArgs<
// (PrepareOutput<Op0, T>, PrepareInput<Op1, T>, Op2)>>`. Each implements
// `GraphNode` only when the ops it wraps do, which is what lets
// `impl GraphNode for Launcher<.., DI> where DI: GraphNode` hold for
// pre-allocated inputs (`&Tensor`, `Arc<Tensor>`, `&TensorView`, partitions of
// existing tensors, plain values — all `Value<T>`) and fail to hold for an
// allocating op such as `api::zeros(..).partition(..)`, whose allocation node
// would return a different address on graph replay.

/// Applies [`KernelInput::prepare`] to the output of `DI`.
pub struct PrepareInput<DI, T> {
    op: DI,
    _elem: std::marker::PhantomData<fn() -> T>,
}

impl<DI, T> PrepareInput<DI, T> {
    pub fn new(op: DI) -> Self {
        Self {
            op,
            _elem: std::marker::PhantomData,
        }
    }
}

impl<T: DType, K: KernelInput<T>, DI: DeviceOp<Output = K>> DeviceOp for PrepareInput<DI, T> {
    type Output = K::Stored;

    unsafe fn execute(self, context: &ExecutionContext) -> Result<K::Stored, DeviceError> {
        Ok(K::prepare(self.op.execute(context)?))
    }
}

impl<T: DType, K: KernelInput<T>, DI: DeviceOp<Output = K> + GraphNode> GraphNode
    for PrepareInput<DI, T>
{
}

impl<T: DType, K: KernelInput<T>, DI: DeviceOp<Output = K>> IntoFuture for PrepareInput<DI, T> {
    type Output = Result<K::Stored, DeviceError>;
    type IntoFuture = DeviceFuture<K::Stored, PrepareInput<DI, T>>;
    fn into_future(self) -> Self::IntoFuture {
        match with_default_device_policy(|policy| {
            let stream = policy.next_stream()?;
            Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
        }) {
            Ok(Ok(future)) => future,
            Ok(Err(e)) => DeviceFuture::failed(e),
            Err(e) => DeviceFuture::failed(e),
        }
    }
}

/// Applies [`KernelOutput::prepare`] to the output of `DI`.
pub struct PrepareOutput<DI, T> {
    op: DI,
    _elem: std::marker::PhantomData<fn() -> T>,
}

impl<DI, T> PrepareOutput<DI, T> {
    pub fn new(op: DI) -> Self {
        Self {
            op,
            _elem: std::marker::PhantomData,
        }
    }
}

impl<T: DType, K: KernelOutput<T>, DI: DeviceOp<Output = K>> DeviceOp for PrepareOutput<DI, T> {
    type Output = K::Stored;

    unsafe fn execute(self, context: &ExecutionContext) -> Result<K::Stored, DeviceError> {
        Ok(K::prepare(self.op.execute(context)?))
    }
}

impl<T: DType, K: KernelOutput<T>, DI: DeviceOp<Output = K> + GraphNode> GraphNode
    for PrepareOutput<DI, T>
{
}

impl<T: DType, K: KernelOutput<T>, DI: DeviceOp<Output = K>> IntoFuture for PrepareOutput<DI, T> {
    type Output = Result<K::Stored, DeviceError>;
    type IntoFuture = DeviceFuture<K::Stored, PrepareOutput<DI, T>>;
    fn into_future(self) -> Self::IntoFuture {
        match with_default_device_policy(|policy| {
            let stream = policy.next_stream()?;
            Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
        }) {
            Ok(Ok(future)) => future,
            Ok(Err(e)) => DeviceFuture::failed(e),
            Err(e) => DeviceFuture::failed(e),
        }
    }
}

/// Executes a tuple of device operations in order and yields the tuple of
/// their outputs — the argument op of a generated kernel launcher.
pub struct KernelArgs<Ops>(pub Ops);

macro_rules! impl_kernel_args {
    ($(($op:ident, $out:ident)),*) => {
        impl<$($op: DeviceOp),*> DeviceOp for KernelArgs<($($op,)*)> {
            type Output = ($(<$op as DeviceOp>::Output,)*);

            #[allow(unused_variables, clippy::unused_unit)]
            unsafe fn execute(
                self,
                context: &ExecutionContext,
            ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
                let ($($out,)*) = self.0;
                Ok(($($out.execute(context)?,)*))
            }
        }

        impl<$($op: GraphNode),*> GraphNode for KernelArgs<($($op,)*)> {}

        impl<$($op: DeviceOp),*> IntoFuture for KernelArgs<($($op,)*)> {
            type Output = Result<<Self as DeviceOp>::Output, DeviceError>;
            type IntoFuture = DeviceFuture<<Self as DeviceOp>::Output, Self>;
            fn into_future(self) -> Self::IntoFuture {
                match with_default_device_policy(|policy| {
                    let stream = policy.next_stream()?;
                    Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
                }) {
                    Ok(Ok(future)) => future,
                    Ok(Err(e)) => DeviceFuture::failed(e),
                    Err(e) => DeviceFuture::failed(e),
                }
            }
        }
    };
}

// Emit the `KernelArgs` impls for every arity from 0 up to the number of
// pairs listed below. Entry kernels in the workspace currently go up to 44
// parameters (`attention_decode_kernel_grouped` in cutile-kernels); 64 leaves
// headroom, and a kernel beyond it fails to compile at the launcher with
// "the trait bound `KernelArgs<...>: DeviceOp` is not satisfied".
macro_rules! impl_kernel_args_cascade {
    (@acc [$($acc:tt),*]) => {
        impl_kernel_args!($($acc),*);
    };
    (@acc [$($acc:tt),*] $next:tt $(, $rest:tt)*) => {
        impl_kernel_args!($($acc),*);
        impl_kernel_args_cascade!(@acc [$($acc,)* $next] $($rest),*);
    };
    ($($pairs:tt),* $(,)?) => {
        impl_kernel_args_cascade!(@acc [] $($pairs),*);
    };
}

impl_kernel_args_cascade!(
    (A0, a0),
    (A1, a1),
    (A2, a2),
    (A3, a3),
    (A4, a4),
    (A5, a5),
    (A6, a6),
    (A7, a7),
    (A8, a8),
    (A9, a9),
    (A10, a10),
    (A11, a11),
    (A12, a12),
    (A13, a13),
    (A14, a14),
    (A15, a15),
    (A16, a16),
    (A17, a17),
    (A18, a18),
    (A19, a19),
    (A20, a20),
    (A21, a21),
    (A22, a22),
    (A23, a23),
    (A24, a24),
    (A25, a25),
    (A26, a26),
    (A27, a27),
    (A28, a28),
    (A29, a29),
    (A30, a30),
    (A31, a31),
    (A32, a32),
    (A33, a33),
    (A34, a34),
    (A35, a35),
    (A36, a36),
    (A37, a37),
    (A38, a38),
    (A39, a39),
    (A40, a40),
    (A41, a41),
    (A42, a42),
    (A43, a43),
    (A44, a44),
    (A45, a45),
    (A46, a46),
    (A47, a47),
    (A48, a48),
    (A49, a49),
    (A50, a50),
    (A51, a51),
    (A52, a52),
    (A53, a53),
    (A54, a54),
    (A55, a55),
    (A56, a56),
    (A57, a57),
    (A58, a58),
    (A59, a59),
    (A60, a60),
    (A61, a61),
    (A62, a62),
    (A63, a63)
);

// ToHostVec

/// A device operation that copies a tensor from device memory to a host `Vec<T>`.
pub struct TensorToHostVec<T: DType, DI>
where
    DI: DeviceOp<Output = Tensor<T>>,
{
    pub(crate) op: DI,
}

unsafe impl<T: DType, DI> Send for TensorToHostVec<T, DI> where DI: DeviceOp<Output = Tensor<T>> {}

impl<T: DType, DI> DeviceOp for TensorToHostVec<T, DI>
where
    DI: DeviceOp<Output = Tensor<T>>,
{
    type Output = Vec<T>;

    unsafe fn execute(
        self,
        context: &ExecutionContext,
    ) -> Result<<Self as DeviceOp>::Output, DeviceError> {
        let tensor = self.op.execute(context)?;
        let cu_deviceptr = tensor.cu_deviceptr();
        let size = tensor.size();
        // The `Vec` owns the host buffer from the start, so an early return
        // frees it, and unlike a bare `alloc` it is well-defined for a
        // zero-size request and never yields null.
        let mut host = Vec::<T>::with_capacity(size);
        if size > 0 {
            unsafe {
                memcpy_dtoh_async(
                    host.as_mut_ptr(),
                    cu_deviceptr,
                    size,
                    context.get_cuda_stream(),
                )
            }?;
        }
        // SAFETY: `cuMemcpyDtoHAsync` into pageable host memory (a `Vec`'s
        // heap buffer is pageable) returns only once the copy has completed,
        // so all `size` elements are initialized here, and `size` is exactly
        // the capacity reserved above.
        unsafe { host.set_len(size) };
        Ok(host)
    }
}

impl<T: DType, DI> IntoFuture for TensorToHostVec<T, DI>
where
    DI: DeviceOp<Output = Tensor<T>>,
{
    type Output = Result<Vec<T>, DeviceError>;
    type IntoFuture = DeviceFuture<Vec<T>, TensorToHostVec<T, DI>>;
    fn into_future(self) -> Self::IntoFuture {
        match with_default_device_policy(|policy| {
            let stream = policy.next_stream()?;
            Ok(DeviceFuture::scheduled(self, ExecutionContext::new(stream)))
        }) {
            Ok(Ok(future)) => future,
            Ok(Err(e)) => DeviceFuture::failed(e),
            Err(e) => DeviceFuture::failed(e),
        }
    }
}

/// Extension trait for converting a tensor device operation into a host `Vec<T>` operation.
pub trait ToHostVecOp<T: DType> {
    /// Wraps this operation to copy the resulting tensor to a host `Vec<T>`.
    fn to_host_vec(self) -> impl DeviceOp<Output = Vec<T>>
    where
        Self: DeviceOp<Output = Tensor<T>>,
    {
        TensorToHostVec { op: self }
    }
}

impl<T: DType, DI> ToHostVecOp<T> for DI where DI: DeviceOp<Output = Tensor<T>> {}

#[cfg(test)]
mod launch_check_tests {
    use super::*;

    fn nonzero(param: usize, axis: usize) -> LaunchCheck {
        LaunchCheck {
            predicate: Predicate::nonzero(Term::atom(Atom::Dim { param, axis })),
            cause: "extent > 0".to_string(),
        }
    }

    fn view_nonzero(param: usize, axis: usize) -> LaunchCheck {
        LaunchCheck {
            predicate: Predicate::nonzero(Term::atom(Atom::ViewExtent { param, axis })),
            cause: "view extent > 0".to_string(),
        }
    }

    #[test]
    fn empty_checks_run_only_the_grid_family() {
        // Matching grids pass; no compiler checks means no extent evaluation.
        assert!(validate_launch(&[], (4, 1, 1), &[GridBound::Exact((4, 1, 1))], &[], &[]).is_ok());
    }

    #[test]
    fn grid_family_still_rejects_mismatched_partition_grid() {
        assert!(validate_launch(&[], (4, 1, 1), &[GridBound::Exact((2, 1, 1))], &[], &[]).is_err());
    }

    #[test]
    fn dim_nonzero_passes_for_positive_extent() {
        let shapes = vec![vec![128, 256]];
        assert!(validate_launch(&[nonzero(0, 0)], (1, 1, 1), &[], &shapes, &[]).is_ok());
    }

    #[test]
    fn dim_nonzero_rejects_zero_extent() {
        let shapes = vec![vec![0, 256]];
        assert!(validate_launch(&[nonzero(0, 0)], (1, 1, 1), &[], &shapes, &[]).is_err());
    }

    #[test]
    fn dim_nonzero_fails_closed_on_missing_parameter() {
        // Check references param 1 axis 0, but only one param was supplied.
        let shapes = vec![vec![128]];
        assert!(validate_launch(&[nonzero(1, 0)], (1, 1, 1), &[], &shapes, &[]).is_err());
    }

    #[test]
    fn each_atom_resolves_against_its_own_frame() {
        // Root says 256 rows; the kernel-visible view (the per-CTA slab) says
        // zero. A root-frame check passes while the view-frame check rejects:
        // the frames are not interchangeable, and the atom picks the array.
        let roots = vec![vec![256, 256]];
        let views = vec![vec![0, 256]];
        assert!(validate_launch(&[nonzero(0, 0)], (1, 1, 1), &[], &roots, &views).is_ok());
        assert!(validate_launch(&[view_nonzero(0, 0)], (1, 1, 1), &[], &roots, &views).is_err());
    }

    #[test]
    fn view_atoms_fail_closed_without_view_shapes() {
        let roots = vec![vec![256, 256]];
        assert!(validate_launch(&[view_nonzero(0, 0)], (1, 1, 1), &[], &roots, &[]).is_err());
    }

    #[test]
    fn prefix_bound_admits_a_per_axis_prefix_and_nothing_more() {
        use GridBound::{AtMost, Exact};
        // Equal and per-axis-smaller launches pass; exceeding ANY axis fails.
        assert!(validate_grid_bounds((3, 2, 1), &[AtMost((3, 2, 1))]).is_ok());
        assert!(validate_grid_bounds((2, 2, 1), &[AtMost((3, 2, 1))]).is_ok());
        assert!(validate_grid_bounds((2, 1, 1), &[AtMost((3, 2, 1))]).is_ok());
        assert!(validate_grid_bounds((4, 1, 1), &[AtMost((3, 2, 1))]).is_err());
        assert!(validate_grid_bounds((1, 3, 1), &[AtMost((3, 2, 1))]).is_err());
        // Per-axis, never total-count: 6 = 3*2 total blocks but the wrong
        // shape must be rejected (delinearization would remap CTAs).
        assert!(validate_grid_bounds((6, 1, 1), &[AtMost((3, 2, 1))]).is_err());
        // Exact bindings keep strict equality even alongside a prefix one.
        assert!(validate_grid_bounds((2, 1, 1), &[Exact((3, 1, 1)), AtMost((3, 1, 1))]).is_err());
        assert!(validate_grid_bounds((3, 1, 1), &[Exact((3, 1, 1)), AtMost((4, 1, 1))]).is_ok());
    }

    #[test]
    fn prefix_bound_cannot_define_the_launch_grid() {
        use GridBound::{AtMost, Exact};
        // Inference needs an exact binding; a bound alone is not a grid.
        let err = infer_launch_grid((0, 0, 0), &[AtMost((3, 1, 1))]).unwrap_err();
        assert!(
            err.to_string().contains("partial-coverage"),
            "the error should say why inference refused: {err}"
        );
        // With an exact sibling, inference works and the bound still gates.
        assert_eq!(
            infer_launch_grid((0, 0, 0), &[Exact((3, 1, 1)), AtMost((4, 1, 1))]).unwrap(),
            (3, 1, 1)
        );
        assert!(infer_launch_grid((0, 0, 0), &[Exact((3, 1, 1)), AtMost((2, 1, 1))]).is_err());
        // An explicit grid validates against both kinds.
        assert_eq!(
            infer_launch_grid((2, 1, 1), &[AtMost((3, 1, 1))]).unwrap(),
            (2, 1, 1)
        );
        assert!(infer_launch_grid((4, 1, 1), &[AtMost((3, 1, 1))]).is_err());
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn unexpected_ast_provider() -> Module {
        panic!("L1 cache-key access must not invoke the AST provider")
    }

    #[test]
    fn l1_cache_key_does_not_invoke_ast_provider() {
        let expected = TileFunctionKey::builder("module", "kernel").build();
        let specialization = Specialization {
            module_ast_fn: unexpected_ast_provider as ModuleAstFn,
            key: expected.clone(),
        };

        assert_eq!(specialization.l1_cache_key(), &expected);
    }
}