teeny_cli/
lib.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Shared `--device`/`--options` CLI configuration for ahead-of-time kernel
18//! compilation, meant to be `#[command(flatten)]`-ed into a downstream
19//! project's own binary (e.g. a `vision-rs` demo) and driven end-to-end by
20//! `cargo teeny aot` in `cargo-teeny`.
21//!
22//! `--device` selects a backend; everything after that (including
23//! `--options`) is opaque here and handed to that backend's own parser —
24//! today that's `teeny_cuda::compiler::options::Options::parse` for
25//! `--device cuda`. Adding a new backend means adding a match arm in
26//! [`aot_compile`], not touching [`AotArgs`].
27
28#![warn(missing_docs)]
29
30use std::path::PathBuf;
31
32use anyhow::{Result, anyhow};
33use teeny_core::graph::{DtypeRepr, Shape, SymTensor};
34use teeny_core::model::{Lowering, LoweringMode};
35use teeny_core::nn::Layer;
36
37/// Default fallback cache directory when neither `--cache-dir` nor
38/// `$TEENYC_CACHE_DIR` is set. Matches the default used by the runtime JIT
39/// compile path (`teeny_compiler::compiler::driver::cuda::compile_kernel`).
40const DEFAULT_CACHE_DIR: &str = "/tmp/teenyc_cache";
41
42/// Shared AOT-compile CLI arguments. Flatten this into your own `clap`
43/// struct: `#[command(flatten)] aot: teeny_cli::AotArgs`.
44#[derive(clap::Args, Clone, Debug)]
45pub struct AotArgs {
46    /// Target backend to compile for (e.g. `cuda`).
47    #[arg(long)]
48    pub device: String,
49
50    /// Backend-specific compiler options as comma-separated `key=value`
51    /// pairs (e.g. `"capability=sm_90,maxnreg=16"`). Parsed by whichever
52    /// backend `--device` selects.
53    #[arg(long)]
54    pub options: Option<String>,
55
56    /// Directory to write/read compiled kernels. Falls back to
57    /// `$TEENYC_CACHE_DIR`, then `/tmp/teenyc_cache` if neither is set.
58    #[arg(long)]
59    pub cache_dir: Option<PathBuf>,
60
61    /// Recompile even if a cached artifact already exists.
62    #[arg(long)]
63    pub force: bool,
64}
65
66impl AotArgs {
67    /// Resolve the effective cache directory: `--cache-dir` > `$TEENYC_CACHE_DIR` > default.
68    pub fn resolve_cache_dir(&self) -> PathBuf {
69        self.cache_dir
70            .clone()
71            .or_else(|| std::env::var_os("TEENYC_CACHE_DIR").map(PathBuf::from))
72            .unwrap_or_else(|| PathBuf::from(DEFAULT_CACHE_DIR))
73    }
74}
75
76/// Trace `model` with a symbolic sample input of `input_dtype`/`input_shape`,
77/// then ahead-of-time compile every kernel the traced graph references for
78/// the backend/config selected by `args`.
79///
80/// The compiled artifacts are written to `args.resolve_cache_dir()` as a
81/// side effect; this returns `()` rather than a backend-specific model type
82/// so that callers (and `--device` dispatch here) don't need to unify return
83/// types across backends. Callers that need the compiled model in-process
84/// (e.g. to immediately `.load()` it onto a live device) should call the
85/// backend's own compile entry point directly instead.
86#[cfg_attr(not(feature = "cuda"), allow(unused_variables))]
87pub fn aot_compile<'a, M, L>(
88    model: &M,
89    input_dtype: DtypeRepr,
90    input_shape: Shape,
91    lowering: &L,
92    mode: LoweringMode,
93    args: &AotArgs,
94) -> Result<()>
95where
96    M: Layer<SymTensor>,
97    L: Lowering<'a>,
98{
99    match args.device.as_str() {
100        #[cfg(feature = "cuda")]
101        "cuda" => {
102            let (input, graph) = SymTensor::input(input_dtype, input_shape);
103            let _ = model.call(input);
104            let graph = graph.borrow();
105
106            let options = teeny_cuda::compiler::options::Options::parse(
107                args.options.as_deref().unwrap_or(""),
108            )?;
109
110            let cache_dir = args.resolve_cache_dir();
111            std::fs::create_dir_all(&cache_dir)?;
112
113            teeny_cuda::compiler::aot::compile_graph(
114                &graph,
115                lowering,
116                mode,
117                &options,
118                cache_dir.to_string_lossy().as_ref(),
119                args.force,
120            )?;
121
122            Ok(())
123        }
124        other => Err(anyhow!(
125            "unsupported --device '{other}'; supported backends: {}",
126            supported_backends()
127        )),
128    }
129}
130
131fn supported_backends() -> &'static str {
132    #[cfg(feature = "cuda")]
133    {
134        "cuda"
135    }
136    #[cfg(not(feature = "cuda"))]
137    {
138        "(none enabled; rebuild teeny-cli with --features cuda)"
139    }
140}
teeny_cli/lib.rs

teeny_cli/
lib.rs
