teeny_compiler/compiler/driver/
cuda.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_core::compiler::{Capability, Compiler};
18use teeny_core::device::program::Kernel;
19
20use crate::compiler::backend::llvm::compiler::LlvmCompiler;
21use crate::compiler::target::cuda::Target;
22use crate::errors::Result;
23
24/// Highest SM version our Triton MLIR codegen has validated support for.
25///
26/// LLVM's NVPTX backend emits `.target sm_NNNa` for each architecture:
27///   sm_90  → sm_90a  (Hopper:   TMA, wgmma)
28///   sm_100 → sm_100a (Blackwell DC: B100/B200/GB200)
29///   sm_120 → sm_120a (Blackwell consumer: RTX 50xx)
30///
31/// Architecture-specific PTX (the `a` suffix) runs only on its own
32/// architecture and forward via native execution — NOT via driver JIT
33/// cross-architecture.  Each GPU therefore needs code compiled for its own SM
34/// version to avoid `ptxas fatal: ... cannot be compiled to future architecture`.
35///
36/// The teenyc backend (LLVM 20+) validates all SM versions up to sm_120.
37/// If a future architecture is released before the backend adds support,
38/// extend this constant and the match arm below.
39#[allow(dead_code)]
40const MAX_CODEGEN_CAPABILITY: Capability = Capability::Sm120;
41
42/// Compiles `kernel` for `target` via the LLVM backend, using `TEENYC_PATH` (or `teenyc` on
43/// `$PATH`) and the default cache directory. Set `force` to recompile even if a cached artifact
44/// exists.
45pub fn compile_kernel(kernel: &impl Kernel, target: &Target, force: bool) -> Result<String> {
46    let teenyc_path = std::env::var("TEENYC_PATH").unwrap_or_else(|_| "teenyc".to_string());
47    let cache_dir = crate::compiler::default_cache_dir();
48
49    let effective_cpu = clamp_capability(target.capability).to_string();
50    let compiler = LlvmCompiler::new(teenyc_path, cache_dir)?.with_target_cpu(effective_cpu);
51    compiler.compile(kernel, target, force)
52}
53
54/// Clamp `cap` to `MAX_CODEGEN_CAPABILITY` for any architecture newer than
55/// what the backend supports.  All SM versions up to sm_120 are natively
56/// supported, so no clamping is needed for current hardware.  If a new SM
57/// version is added to `Capability` before teenyc validates it, add a match
58/// arm here that maps it to `MAX_CODEGEN_CAPABILITY`.
59fn clamp_capability(cap: Capability) -> Capability {
60    cap
61}
teeny_compiler/compiler/driver/cuda.rs

teeny_compiler/compiler/driver/
cuda.rs
