diff --git a/aten/src/ATen/cuda/cub.cuh b/aten/src/ATen/cuda/cub.cuh index ba097dafef791..5f4af0f42a065 100644 --- a/aten/src/ATen/cuda/cub.cuh +++ b/aten/src/ATen/cuda/cub.cuh @@ -70,7 +70,11 @@ #define ATEN_CUB_MAXIMUM() NO_ROCM(at_cuda_detail)ROCM_HIPCUB(::cub)::Max() #endif -#if defined(USE_ROCM) +// 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. +#if defined(USE_ROCM) && CUB_VERSION < 300000 // backport https://github.com/NVIDIA/cub/pull/306 for c10::BFloat16 diff --git a/aten/src/ATen/cuda/cub_definitions.cuh b/aten/src/ATen/cuda/cub_definitions.cuh index e0d7455d4c22b..7d9eb4096133f 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 support for CUB_WRAPPED_NAMESPACE is added to cub 1.13.1 in: // https://github.com/NVIDIA/cub/pull/326 diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake index a87b16f5ba889..1784a002785b6 100644 --- a/cmake/public/LoadHIP.cmake +++ b/cmake/public/LoadHIP.cmake @@ -228,6 +228,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)