Gate CUDA 12.8 features using runtime versions

crates/cubecl-cuda/src/compute/server.rs

@@ -8,6 +8,7 @@ use crate::{
         stream::CudaStreamBackend,
         sync::Fence,
     },
+    runtime::CudaApiVersions,
 };
 use cubecl_common::{
     backtrace::BackTrace, bytes::Bytes, profile::ProfileDuration, stream_id::StreamId,
@@ -50,6 +51,7 @@ pub(crate) const MB: usize = 1024 * 1024;
 #[derive(Debug)]
 pub struct CudaServer {
     ctx: CudaContext,
+    cuda_api_versions: CudaApiVersions,
     device_id: DeviceId,
     streams: MultiStream<CudaStreamBackend>,
     utilities: Arc<ServerUtilities<Self>>,
@@ -555,6 +557,7 @@ impl CudaServer {
     /// Create a new cuda server.
     pub(crate) fn new(
         ctx: CudaContext,
+        cuda_api_versions: CudaApiVersions,
         mem_props: MemoryDeviceProperties,
         mem_config: MemoryConfiguration,
         mem_alignment: usize,
@@ -573,6 +576,7 @@ impl CudaServer {
 
         Self {
             ctx,
+            cuda_api_versions,
             device_id,
             streams: MultiStream::new(
                 utilities.logger.clone(),
@@ -663,6 +667,7 @@ impl CudaServer {
             .compilation_options
             .supports_features
             .grid_constants;
+        let supports_cuda_12_8 = self.cuda_api_versions.supports_cuda_12_8();
         let mut command = self.command(
             stream_id,
             bindings.buffers.iter(),
@@ -746,6 +751,25 @@ impl CudaServer {
                 .collect();
             let elem_stride: Vec<_> = map.elem_stride.iter().rev().map(|s| *s as u32).collect();
 
+            #[cfg(cuda_12080)]
+            if !supports_cuda_12_8
+                && (matches!(&map.format, TensorMapFormat::Im2colWide(_))
+                    || matches!(
+                        map.swizzle,
+                        TensorMapSwizzle::B128Atom32B
+                            | TensorMapSwizzle::B128Atom32BFlip8B
+                            | TensorMapSwizzle::B128Atom64B
+                    )
+                    || matches!(map.storage_ty, StorageType::Packed(ty, 2) if ty.size_bits() == 4))
+            {
+                return Err(LaunchError::Unknown {
+                    reason: "CUDA driver and NVRTC 12.8 or newer are required for this tensor map"
+                        .into(),
+                    backtrace: BackTrace::capture(),
+                }
+                .into());
+            }
+
             match &map.format {
                 // SAFETY: `map_ptr` is a zeroed `MaybeUninit<CUtensorMap>`. `device_ptr` is a
                 // valid device pointer. Shape, strides, tile_size, and elem_stride vectors

crates/cubecl-cuda/src/runtime.rs

@@ -31,9 +31,55 @@ use cubecl_cpp::{
 use cubecl_runtime::{
     allocator::PitchedMemoryLayoutPolicy, client::ComputeClient, logging::ServerLogger,
 };
-use cudarc::driver::sys::{CUDA_VERSION, cuDeviceTotalMem_v2};
+use cudarc::driver::sys::{cuDeviceTotalMem_v2, cuDriverGetVersion};
 use std::{mem::MaybeUninit, sync::Arc};
 
+const CUDA_12_8: i32 = 12080;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct CudaApiVersions {
+    driver: i32,
+    nvrtc: i32,
+}
+
+impl CudaApiVersions {
+    const fn new(driver: i32, nvrtc: i32) -> Self {
+        Self { driver, nvrtc }
+    }
+
+    const fn effective(self) -> i32 {
+        if self.driver < self.nvrtc {
+            self.driver
+        } else {
+            self.nvrtc
+        }
+    }
+
+    pub(crate) const fn supports_cuda_12_8(self) -> bool {
+        self.effective() >= CUDA_12_8
+    }
+
+    fn detect() -> Self {
+        let mut driver = 0;
+        let mut nvrtc_major = 0;
+        let mut nvrtc_minor = 0;
+
+        // SAFETY: Both functions only write integer version components to the
+        // valid output pointers supplied here. The CUDA driver is initialized
+        // immediately before this query in `CudaServer::init`.
+        unsafe {
+            cuDriverGetVersion(&mut driver)
+                .result()
+                .expect("failed to query CUDA driver version");
+            cudarc::nvrtc::sys::nvrtcVersion(&mut nvrtc_major, &mut nvrtc_minor)
+                .result()
+                .expect("failed to query NVRTC version");
+        }
+
+        Self::new(driver, nvrtc_major * 1000 + nvrtc_minor * 10)
+    }
+}
+
 /// Options configuring the CUDA runtime.
 #[derive(Default)]
 pub struct RuntimeOptions {
@@ -51,6 +97,7 @@ impl DeviceService for CudaServer {
 
         // To get the supported WMMA features, and memory properties, we have to initialize the server immediately.
         cudarc::driver::result::init().unwrap();
+        let cuda_api_versions = CudaApiVersions::detect();
         let device_index = device.index as i32;
         let device_ptr = cudarc::driver::result::device::get(device_index).unwrap();
         let arch_major;
@@ -217,7 +264,7 @@ impl DeviceService for CudaServer {
                 .matmul
                 .ldmatrix
                 .insert(ElemType::Float(FloatKind::BF16).into());
-            comp_opts.supports_features.fast_tanh = CUDA_VERSION >= 12080;
+            comp_opts.supports_features.fast_tanh = cuda_api_versions.supports_cuda_12_8();
         }
 
         if arch_version >= 80 {
@@ -265,7 +312,7 @@ impl DeviceService for CudaServer {
             );
         }
 
-        if arch_version >= 100 {
+        if arch_version >= 100 && cuda_api_versions.supports_cuda_12_8() {
             device_props.features.tma.insert(Tma::Im2colWide);
             // Breaks swizzle so disable for now and fix in a PR specifically for this
             // if CUDA_VERSION >= 12090 {
@@ -296,7 +343,7 @@ impl DeviceService for CudaServer {
                 TypeUsage::Conversion | TypeUsage::Buffer,
             );
 
-            if CUDA_VERSION >= 12080 {
+            if cuda_api_versions.supports_cuda_12_8() {
                 device_props.features.tma.insert(Tma::SwizzleAtomicity);
             }
         }
@@ -320,6 +367,7 @@ impl DeviceService for CudaServer {
 
         CudaServer::new(
             cuda_ctx,
+            cuda_api_versions,
             mem_properties,
             options.memory_config,
             mem_alignment,
@@ -398,3 +446,26 @@ impl Runtime for CudaRuntime {
             .collect()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::CudaApiVersions;
+
+    #[test]
+    fn cuda_12_8_capabilities_require_both_driver_and_nvrtc() {
+        let cases = [
+            (12040, 12040, 12040, false),
+            (12060, 12060, 12060, false),
+            (12080, 12080, 12080, true),
+            (13000, 12080, 12080, true),
+            (12080, 12040, 12040, false),
+            (12040, 12080, 12040, false),
+        ];
+
+        for (driver, nvrtc, effective, supports_cuda_12_8) in cases {
+            let versions = CudaApiVersions::new(driver, nvrtc);
+            assert_eq!(versions.effective(), effective);
+            assert_eq!(versions.supports_cuda_12_8(), supports_cuda_12_8);
+        }
+    }
+}

crates/cubecl-cuda/tests/cudarc_feature_contract.rs

@@ -20,3 +20,21 @@ fn cudarc_build_dependency_only_enables_build_script_requirements() {
         r#"[build-dependencies] cudarc = { version = "0.19.0", default-features = false, features = ["std", "driver", "dynamic-loading", "cuda-12080"] }"#
     ));
 }
+
+#[test]
+fn cuda_12_8_tensor_map_symbol_is_runtime_guarded() {
+    let runtime = include_str!("../src/runtime.rs");
+    let server = include_str!("../src/compute/server.rs");
+
+    assert!(runtime.contains("cuDriverGetVersion"));
+    assert!(runtime.contains("nvrtcVersion"));
+
+    let tensor_map_loop = server
+        .find("for TensorMapBinding")
+        .expect("tensor map encoding loop must exist");
+    let symbol_call = server[tensor_map_loop..]
+        .find("cuTensorMapEncodeIm2colWide(")
+        .expect("Im2colWide symbol call must exist");
+    let guarded_prefix = &server[tensor_map_loop..tensor_map_loop + symbol_call];
+    assert!(guarded_prefix.contains("supports_cuda_12_8"));
+}