teeny_core/compiler/
mod.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 alloc::format;
18use alloc::string::String;
19
20use crate::device::program::Kernel;
21use crate::errors::Result;
22
23/// A compilation target (e.g. a GPU architecture).
24pub trait Target: Sized {
25    /// The `-Ctarget-cpu`-style string identifying this target, if applicable.
26    fn target_cpu(&self) -> Option<String> {
27        None
28    }
29}
30
31/// Something capable of compiling a [`Kernel`] for a [`Target`].
32pub trait Compiler {
33    /// Compiles `kernel` for `target`, returning the path to the compiled artifact.
34    /// `force` recompiles even if a cached artifact exists.
35    fn compile(&self, kernel: &impl Kernel, target: &impl Target, force: bool) -> Result<String>;
36}
37
38/// GPU compute capability (CUDA SM version).
39///
40/// Minimum supported: sm_75 (Turing). Triton's MMA acceleration requires sm_75+;
41/// sm_70/sm_72 only have a deprecated FMA fallback path.
42#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
43pub enum Capability {
44    /// Turing: RTX 20xx, GTX 16xx, T4.
45    Sm75,
46    /// Ampere datacenter: A100, A30.
47    Sm80,
48    /// Ampere: RTX 30xx, A40, A10, A16.
49    Sm86,
50    /// Ampere embedded: Jetson Orin (AGX/NX/Nano).
51    Sm87,
52    /// Ada Lovelace: RTX 40xx, L4, L40S.
53    Sm89,
54    /// Hopper: H100, H200.
55    Sm90,
56    /// Blackwell datacenter: B100, B200, GB200.
57    Sm100,
58    /// Blackwell: RTX 50xx (GB10x).
59    Sm120,
60}
61
62impl Capability {
63    /// Looks up the `Capability` matching a CUDA `(major, minor)` compute-capability version.
64    pub fn from_major_minor(major: i32, minor: i32) -> Option<Self> {
65        match (major, minor) {
66            (7, 5) => Some(Self::Sm75),
67            (8, 0) => Some(Self::Sm80),
68            (8, 6) => Some(Self::Sm86),
69            (8, 7) => Some(Self::Sm87),
70            (8, 9) => Some(Self::Sm89),
71            (9, 0) => Some(Self::Sm90),
72            (10, 0) => Some(Self::Sm100),
73            (12, 0) => Some(Self::Sm120),
74            _ => None,
75        }
76    }
77}
78
79impl core::fmt::Display for Capability {
80    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
81        let s = match self {
82            Self::Sm75 => "sm_75",
83            Self::Sm80 => "sm_80",
84            Self::Sm86 => "sm_86",
85            Self::Sm87 => "sm_87",
86            Self::Sm89 => "sm_89",
87            Self::Sm90 => "sm_90",
88            Self::Sm100 => "sm_100",
89            Self::Sm120 => "sm_120",
90        };
91        f.write_str(s)
92    }
93}
94
95impl core::str::FromStr for Capability {
96    type Err = String;
97
98    /// Accepts the canonical `sm_90` form as well as `sm-90`, `sm90` and bare
99    /// `90` (case-insensitive) — the latter forms are convenient for CLI args
100    /// like `--options capability=sm_90`.
101    fn from_str(s: &str) -> core::result::Result<Self, Self::Err> {
102        let normalized = s.trim().to_ascii_lowercase();
103        let digits = normalized
104            .strip_prefix("sm_")
105            .or_else(|| normalized.strip_prefix("sm-"))
106            .or_else(|| normalized.strip_prefix("sm"))
107            .unwrap_or(normalized.as_str());
108
109        match digits {
110            "75" => Ok(Self::Sm75),
111            "80" => Ok(Self::Sm80),
112            "86" => Ok(Self::Sm86),
113            "87" => Ok(Self::Sm87),
114            "89" => Ok(Self::Sm89),
115            "90" => Ok(Self::Sm90),
116            "100" => Ok(Self::Sm100),
117            "120" => Ok(Self::Sm120),
118            _ => Err(format!(
119                "unknown capability '{s}'; expected one of sm_75, sm_80, sm_86, sm_87, sm_89, sm_90, sm_100, sm_120"
120            )),
121        }
122    }
123}
124
125#[cfg(test)]
126mod capability_from_str_tests {
127    use super::Capability;
128
129    #[test]
130    fn accepts_canonical_and_tolerant_forms() {
131        assert_eq!("sm_90".parse::<Capability>().unwrap(), Capability::Sm90);
132        assert_eq!("sm-90".parse::<Capability>().unwrap(), Capability::Sm90);
133        assert_eq!("SM90".parse::<Capability>().unwrap(), Capability::Sm90);
134        assert_eq!("90".parse::<Capability>().unwrap(), Capability::Sm90);
135        assert_eq!("sm_120".parse::<Capability>().unwrap(), Capability::Sm120);
136    }
137
138    #[test]
139    fn rejects_unknown_capability() {
140        assert!("sm_61".parse::<Capability>().is_err());
141    }
142}
teeny_core/compiler/mod.rs

teeny_core/compiler/
mod.rs
