From af1f5480393b8ef8844b3f6e29e5175e96b28ae6 Mon Sep 17 00:00:00 2001 From: Hongyi Jin Date: Mon, 17 Aug 2026 22:22:31 -0400 Subject: [PATCH] [CUDA][BugFix] Guard tensor map enums by toolkit version CUDA tensor-map enum members are identifiers rather than preprocessor macros, so #ifdef silently excluded the CUDA 12.8 packed data types and atomic swizzle modes even when cuda.h declared them. Use CUDA_VERSION consistently for the validation paths that reference those members. --- src/backend/cuda/runtime/cuda_device_api.cc | 30 +++++---------------- 1 file changed, 7 insertions(+), 23 deletions(-) diff --git a/src/backend/cuda/runtime/cuda_device_api.cc b/src/backend/cuda/runtime/cuda_device_api.cc index 01222b408588..c70cf24421b6 100644 --- a/src/backend/cuda/runtime/cuda_device_api.cc +++ b/src/backend/cuda/runtime/cuda_device_api.cc @@ -602,14 +602,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto is_valid_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_NONE || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_32B || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_64B || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B +#if CUDA_VERSION >= 12080 is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B is_valid_swizzle = is_valid_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; #endif TVM_FFI_ICHECK(is_valid_swizzle) @@ -628,15 +624,11 @@ TVM_FFI_STATIC_INIT_BLOCK() { << "Unsupported oobFill enum value: " << static_cast(oob_fill_kind); bool is_packed_16u4_align8 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B - is_packed_16u4_align8 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; -#endif bool is_packed_16u4_align16 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B - is_packed_16u4_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; -#endif bool is_packed_16u6_align16 = false; -#ifdef CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B +#if CUDA_VERSION >= 12080 + is_packed_16u4_align8 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; + is_packed_16u4_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; is_packed_16u6_align16 = cu_dtype == CU_TENSOR_MAP_DATA_TYPE_16U6_ALIGN16B; #endif auto is_packed_align16 = is_packed_16u4_align16 || is_packed_16u6_align16; @@ -645,25 +637,17 @@ TVM_FFI_STATIC_INIT_BLOCK() { cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT32 || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT64 || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; -#ifdef CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ +#if CUDA_VERSION >= 12080 is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_FLOAT32_FTZ; -#endif -#ifdef CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; -#endif -#ifdef CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ is_floating_dtype = is_floating_dtype || cu_dtype == CU_TENSOR_MAP_DATA_TYPE_TFLOAT32_FTZ; #endif auto is_128b_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B +#if CUDA_VERSION >= 12080 is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B_FLIP_8B; -#endif -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B is_128b_swizzle = is_128b_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_64B; #endif @@ -702,7 +686,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { if (is_packed_16u4_align16) { bool supported_swizzle = swizzle_kind == CU_TENSOR_MAP_SWIZZLE_NONE || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; -#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B +#if CUDA_VERSION >= 12080 supported_swizzle = supported_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; #endif TVM_FFI_ICHECK(supported_swizzle)