diff --git a/aten/src/ATen/cuda/cub.cuh b/aten/src/ATen/cuda/cub.cuh index 7828c3917fc46..a332e3e63e24f 100644 --- a/aten/src/ATen/cuda/cub.cuh +++ b/aten/src/ATen/cuda/cub.cuh @@ -70,12 +70,17 @@ #define ATEN_CUB_MAXIMUM() NO_ROCM(at_cuda_detail)ROCM_HIPCUB(::cub)::Max() #endif -#if (!defined(USE_ROCM) && !CUB_SUPPORTS_NV_BFLOAT16()) || defined(USE_ROCM) +#if (!defined(USE_ROCM) && !CUB_SUPPORTS_NV_BFLOAT16()) || (defined(USE_ROCM) && CUB_VERSION < 300000) #if !defined(USE_ROCM) namespace at_cuda_detail { #endif +// Pre ROCm 8.0 we require backporting of FpLimits specialization +// for 'c10::BFloat16'. However, ROCm 8.0 lifts the compatible +// CUB version to 3.x which uses libhipcxx and derives numerical +// limits from there. + // backport https://github.com/NVIDIA/cub/pull/306 for c10::BFloat16 template <> diff --git a/aten/src/ATen/cuda/cub_definitions.cuh b/aten/src/ATen/cuda/cub_definitions.cuh index 0d76ae6e8dcfc..b68fee5114e96 100644 --- a/aten/src/ATen/cuda/cub_definitions.cuh +++ b/aten/src/ATen/cuda/cub_definitions.cuh @@ -1,14 +1,20 @@ #pragma once #if !defined(USE_ROCM) -#include // for CUDA_VERSION -#endif - -#if !defined(USE_ROCM) +#include // for CUDA_VERSION #include #else +// Check if we can find HIPCUB_CCCL_VERSION. +#include +// Older versions of hipCUB do not support the CUB V3 API. +#if defined(HIPCUB_CCCL_VERSION) && HIPCUB_CCCL_VERSION >= 300000 +#define CUB_VERSION HIPCUB_CCCL_VERSION +#else +// If we cannot find a compatible CCCL version, fallback to +// a very old version. #define CUB_VERSION 200001 #endif +#endif // cub sort support for __nv_bfloat16 is added to cub 1.13 in: // https://github.com/NVIDIA/cub/pull/306 diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake index 538a2fe1982c5..c4030328940c1 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -189,6 +189,9 @@ if(HIP_FOUND) # Optional components. find_package_and_print_version(hipsparselt) # Will be required when ready. + # ROCm 8.0 and later requires libhipcxx! This should be marked as + # 'REQUIRED' once minimal ROCm version is bumped to 8.0 or later. + find_package_and_print_version(libhipcxx) list(REMOVE_DUPLICATES ROCM_INCLUDE_DIRS)