logicaffeine_compile/codegen_cuda/
chip.rs

1//! The GPU **chip / family** knob — which NVIDIA silicon the CUDA backend targets.
2//!
3//! This is the CUDA analog of the Tenstorrent backend's per-part [`crate::spatial::target::SpatialTarget`]
4//! (wormhole / blackhole): a described target part whose SM architecture, shared-memory budget, and realized
5//! numeric formats make the fastest kernel *chip-specific*. The SAME recognized LOGOS tensor op lowers to a
6//! kernel tuned for the local **RTX 3060** (Ampere, sm_86) — where we prove correctness — or for a cloud
7//! **B200** (Blackwell, sm_100), the target we submit to NVIDIA's SOL-ExecBench against.
8//!
9//! Choosing the chip is a genuine *trade-off* (hence a knob, not a constant): it trades **portability** against
10//! **hardware exploitation**. A kernel compiled for `sm_100` won't run on an `sm_86` host; a kernel compiled
11//! for `sm_86` leaves the B200's 5th-gen tensor cores, fp8/nvfp4, and larger shared memory unused.
12
13use crate::codegen_sva::fabric_target::Precision;
14
15/// An NVIDIA GPU architecture family — fixes the tensor-core generation and which numeric formats exist.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum GpuFamily {
18    /// sm_80 / sm_86 — 3rd-gen tensor cores: bf16 / fp16 / tf32, int8 IMMA. A100, RTX 30xx.
19    Ampere,
20    /// sm_89 — 4th-gen tensor cores + fp8 (e4m3/e5m2). Ada: RTX 40xx / L40.
21    Ada,
22    /// sm_90 — 4th-gen tensor cores, wgmma / TMA, fp8. Hopper: H100 / H200.
23    Hopper,
24    /// sm_100 — 5th-gen tensor cores, tcgen05, fp8 + NVFP4 16-element block scaling. Blackwell: B200 / GB200.
25    Blackwell,
26}
27
28/// A described CUDA target part — the hardware limits and realized numeric formats that make the fastest
29/// kernel chip-specific, plus the SM arch its kernels are compiled for. Construct with a named part
30/// ([`CudaChip::rtx3060`], [`CudaChip::b200`], …); the emitter reads it to pick the arch flag, the shared-memory
31/// budget, and which precisions are legal.
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub struct CudaChip {
34    /// Human-readable part name — carried into the emitted kernel's provenance header.
35    pub name: &'static str,
36    /// The architecture family (tensor-core generation).
37    pub family: GpuFamily,
38    /// SM architecture number (80 / 86 / 89 / 90 / 100) — the `sm_XX` the kernel compiles for.
39    pub sm_arch: u32,
40    /// Opt-in dynamic shared-memory ceiling per block, in KiB (bounds the resident row/tile a kernel stages).
41    pub smem_kib_per_block: usize,
42    /// Threads per block ceiling.
43    pub max_threads_per_block: usize,
44    /// Streaming multiprocessors: the blocks a launch must offer to fill the part.
45    pub sms: usize,
46    /// The part's device memory, in bytes: what a lane plans its program's weights, caches and arrays on, less what the
47    /// driver and the program's context take (`lane::residence::CONTEXT_RESERVE`).
48    pub mem_bytes: usize,
49    /// The numeric formats this part's tensor / CUDA cores realize — the precision legality gate (a lowering
50    /// may pick a precision only if the target chip provides it).
51    pub precisions: &'static [Precision],
52}
53
54/// Ampere/Ada realize bf16/fp16/fp32 (+ the int formats, int32 the CUDA cores' own word) but no fp8; the tensor path
55/// picks among these.
56const AMPERE_PRECISIONS: &[Precision] = &[
57    Precision::Int8,
58    Precision::Int16,
59    Precision::Int32,
60    Precision::Int64,
61    Precision::Bf16,
62    Precision::Fp16,
63    Precision::Fp32,
64];
65/// Hopper/Blackwell add fp8. (NVFP4 joins here in Phase 5, once `Precision::Nvfp4` and its witness ship.)
66const HOPPER_PLUS_PRECISIONS: &[Precision] = &[
67    Precision::Int8,
68    Precision::Int16,
69    Precision::Int32,
70    Precision::Int64,
71    Precision::Bf16,
72    Precision::Fp16,
73    Precision::Fp32,
74    Precision::Fp8,
75];
76
77impl CudaChip {
78    /// **RTX 3060** — the local dev GPU: Ampere, sm_86, 28 SMs, ~99 KiB opt-in shared memory, 12 GiB. Where we prove
79    /// correctness and relative speed before spending B200 time.
80    pub fn rtx3060() -> Self {
81        CudaChip {
82            name: "rtx3060",
83            family: GpuFamily::Ampere,
84            sm_arch: 86,
85            smem_kib_per_block: 99,
86            max_threads_per_block: 1024,
87            sms: 28,
88            mem_bytes: 12 << 30,
89            precisions: AMPERE_PRECISIONS,
90        }
91    }
92
93    /// **A100** — Ampere data-center part, sm_80, 108 SMs, ~164 KiB opt-in shared memory, the 80 GB part.
94    pub fn a100() -> Self {
95        CudaChip {
96            name: "a100",
97            family: GpuFamily::Ampere,
98            sm_arch: 80,
99            smem_kib_per_block: 164,
100            max_threads_per_block: 1024,
101            sms: 108,
102            mem_bytes: 80 << 30,
103            precisions: AMPERE_PRECISIONS,
104        }
105    }
106
107    /// **H100** (SXM / NVL) — Hopper, sm_90, 132 SMs, ~228 KiB opt-in shared memory, 80 GB (the SXM part), fp8.
108    pub fn h100() -> Self {
109        CudaChip {
110            name: "h100",
111            family: GpuFamily::Hopper,
112            sm_arch: 90,
113            smem_kib_per_block: 228,
114            max_threads_per_block: 1024,
115            sms: 132,
116            mem_bytes: 80 << 30,
117            precisions: HOPPER_PLUS_PRECISIONS,
118        }
119    }
120
121    /// **B200** — Blackwell, sm_100, 148 SMs, ~228 KiB opt-in shared memory, 192 GB, fp8 (+ NVFP4 in Phase 5). The
122    /// SOL-ExecBench target.
123    pub fn b200() -> Self {
124        CudaChip {
125            name: "b200",
126            family: GpuFamily::Blackwell,
127            sm_arch: 100,
128            smem_kib_per_block: 228,
129            max_threads_per_block: 1024,
130            sms: 148,
131            mem_bytes: 192 << 30,
132            precisions: HOPPER_PLUS_PRECISIONS,
133        }
134    }
135
136    /// The `nvcc` arch flag that targets this part. Hopper/Blackwell get the architecture-specific `a` suffix
137    /// (required to unlock wgmma/TMA on Hopper and tcgen05 on Blackwell); Ampere/Ada use the plain form.
138    pub fn nvcc_arch_flag(&self) -> String {
139        let suffix = match self.family {
140            GpuFamily::Hopper | GpuFamily::Blackwell => "a",
141            GpuFamily::Ampere | GpuFamily::Ada => "",
142        };
143        format!("-arch=sm_{}{}", self.sm_arch, suffix)
144    }
145
146    /// Whether this part's tensor / CUDA path realizes a numeric format.
147    pub fn supports(&self, p: Precision) -> bool {
148        self.precisions.contains(&p)
149    }
150
151    /// Can a kernel built for THIS chip be *numerically* executed on a host GPU of SM arch `host_sm`? Kernels
152    /// are forward-compatible: a lower-arch kernel runs on a higher-arch host, never the reverse. This routes
153    /// the run-vs-compile-check split — everything sm_86-legal is proven on the local 3060; a B200 kernel can
154    /// only be compile-checked locally and is numerically validated on the cloud B200.
155    pub fn runs_numerically_on(&self, host_sm: u32) -> bool {
156        self.sm_arch <= host_sm
157    }
158}
logicaffeine_compile/codegen_cuda/chip.rs

logicaffeine_compile/codegen_cuda/
chip.rs
