cubecl_cpp/cuda/
arch.rs

1use std::fmt::Display;
2
3use crate::shared::Architecture;
4
5#[derive(Debug)]
6pub struct CudaArchitecture {
7    pub version: u32,
8    /// Compute capability cannot say this: TU116 and TU117 report 7.5 like every other
9    /// Turing, and `mma.sync` still runs on them, on the FP16 pipeline at a fiftieth the rate.
10    pub tensor_cores: bool,
11}
12
13impl CudaArchitecture {
14    /// Decided by the marketing name because CUDA exposes no attribute for it. No GTX-branded
15    /// Turing has tensor cores and every Turing that does shipped under another brand, so the
16    /// two coincide exactly at 7.5. The professional TU117 parts (T400 to T1000) are not caught.
17    pub fn has_tensor_cores(version: u32, name: &str) -> bool {
18        match version {
19            // Tensor cores arrive with Volta.
20            ..70 => false,
21            // The GTX-branded Turings are the only later dies without them.
22            75 => !name.to_uppercase().contains("GTX"),
23            _ => true,
24        }
25    }
26}
27
28impl Display for CudaArchitecture {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        write!(f, "{}", self.version)
31    }
32}
33
34impl Architecture for CudaArchitecture {
35    fn warp_size(&self) -> u32 {
36        32
37    }
38
39    fn is_wmma_capable(&self) -> bool {
40        self.tensor_cores
41    }
42
43    fn is_mfma_capable(&self) -> bool {
44        false
45    }
46
47    fn get_version(&self) -> u32 {
48        self.version
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::CudaArchitecture;
55
56    #[test]
57    fn turing_without_tensor_cores_is_told_apart_from_turing_with_them() {
58        assert!(!CudaArchitecture::has_tensor_cores(
59            75,
60            "NVIDIA GeForce GTX 1660 SUPER"
61        ));
62        assert!(CudaArchitecture::has_tensor_cores(
63            75,
64            "NVIDIA GeForce RTX 2060"
65        ));
66    }
67
68    #[test]
69    fn the_gtx_exception_applies_only_to_turing() {
70        assert!(CudaArchitecture::has_tensor_cores(
71            80,
72            "NVIDIA A100-GTX-ish"
73        ));
74        // Tesla T4 is TU104: 7.5 with tensor cores, no GTX in the name.
75        assert!(CudaArchitecture::has_tensor_cores(75, "Tesla T4"));
76    }
77
78    #[test]
79    fn nothing_before_volta_has_them() {
80        assert!(!CudaArchitecture::has_tensor_cores(
81            61,
82            "NVIDIA GeForce GTX 1080"
83        ));
84        assert!(!CudaArchitecture::has_tensor_cores(
85            52,
86            "NVIDIA GeForce GTX 980"
87        ));
88        // Volta itself has them: the boundary is inclusive.
89        assert!(CudaArchitecture::has_tensor_cores(70, "Tesla V100-SXM2"));
90    }
91
92    #[test]
93    fn an_unnamed_device_keeps_what_its_version_says() {
94        assert!(CudaArchitecture::has_tensor_cores(
95            75,
96            "unknown CUDA device"
97        ));
98    }
99}
cubecl_cpp/cuda/arch.rs

cubecl_cpp/cuda/
arch.rs
