[release/2.9] [ROCm] derive the hipcub compatible cccl version from 'HIPCUB_CCCL_VERSION'

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 <>

aten/src/ATen/cuda/cub_definitions.cuh

@@ -1,14 +1,20 @@
 #pragma once
 
 #if !defined(USE_ROCM)
-#include <cuda.h>  // for CUDA_VERSION
-#endif
-
-#if !defined(USE_ROCM)
+#include <cuda.h> // for CUDA_VERSION
 #include <cub/version.cuh>
 #else
+// Check if we can find HIPCUB_CCCL_VERSION.
+#include <hipcub/hipcub_version.hpp>
+// 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

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)