teeny_cuda/compiler/
aot.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
17use teeny_compiler::compiler::backend::llvm::compiler::LlvmCompiler;
18use teeny_compiler::compiler::target::cuda::Target;
19use teeny_core::graph::Graph;
20use teeny_core::model::{Lowering, LoweringMode};
21
22use crate::compiler::graph::CudaGraphCompiler;
23use crate::compiler::options::Options;
24use crate::errors::Result;
25use crate::model::CudaModel;
26
27/// Ahead-of-time compile an already-traced `graph` to PTX files on disk,
28/// driven by an already-parsed [`Options`] (see [`Options::parse`]).
29///
30/// This is the same `LlvmCompiler` + `CudaGraphCompiler` sequence used for
31/// JIT compilation at runtime (e.g. in `models/teeny-vision/examples/mnist.rs`),
32/// just parameterized by CLI-driven `options`/`cache_dir` instead of hardcoded
33/// values. No live CUDA device/context is required — only `.load()` on the
34/// returned `CudaModel` needs one.
35///
36/// `cache_dir` is where compiled PTX is written/read. Passing the same
37/// directory a later run resolves via `TEENYC_CACHE_DIR` pre-warms that
38/// runtime JIT cache. The `teenyc` binary is resolved via
39/// [`teeny_compiler::compiler::find_teenyc`], matching the existing JIT compile path.
40pub fn compile_graph<'a, L: Lowering<'a>>(
41    graph: &Graph,
42    lowering: &L,
43    mode: LoweringMode,
44    options: &Options,
45    cache_dir: &str,
46    force: bool,
47) -> Result<CudaModel<'a>> {
48    let teenyc_path = teeny_compiler::compiler::find_teenyc()?;
49
50    let mut compiler = LlvmCompiler::new(teenyc_path, cache_dir.to_string())?;
51    if let Some(ptx_version) = options.ptx_version {
52        compiler = compiler.with_ptx_version(ptx_version);
53    }
54    if let Some(log_level) = options.log_level {
55        compiler = compiler.with_log_level(log_level);
56    }
57    let graph_compiler = CudaGraphCompiler::new(compiler);
58    let target = Target::new(options.gpu_name);
59
60    graph_compiler.compile_model(graph, lowering, &target, mode, force)
61}
teeny_cuda/compiler/aot.rs

teeny_cuda/compiler/
aot.rs
