Remove use of __AMDGCN_WAVEFRONT_SIZE macro with it being deprecated in ROCm 7 (#4619)

Summary:
Usage of `__AMDGCN_WAVEFRONT_SIZE` is a compilation error when using amdclang++ shipped with ROCm 7.x.x. For versions >= 7 we delegate to rocprim to get the right constexpr value that can be used in device code. We continue using __AMDGCN_WAVEFRONT_SIZE for older ROCm versions to maintain compatibility.

Note that the usage of kWarpSize in host code is still non-sensical however the current usage of this variable seems to always be in device functions.

See https://rocm.docs.amd.com/en/latest/about/release-notes.html#amdgpu-wavefront-size-compiler-macro-deprecation for deprecation notice.

Pull Request resolved: #4619

Reviewed By: mnorris11

Differential Revision: D95327171

Pulled By: subhadeepkaran

fbshipit-source-id: f1a7b7d67a7c5a6c2a8068f4608d23873d21ff40

Author: kevinjoseph1995

faiss/gpu/impl/PQCodeDistances-inl.cuh

@@ -20,11 +20,9 @@
 namespace faiss {
 namespace gpu {
 
-#if defined(USE_AMD_ROCM) && __AMDGCN_WAVEFRONT_SIZE == 64u
-#define LAUNCH_BOUND 320
-#else
-#define LAUNCH_BOUND 288
-#endif
+__device__ constexpr inline int getLaunchBound() {
+    return kWarpSize == 32 ? 288 : 320;
+}
 
 // Kernel responsible for calculating distance from residual vector to
 // each product quantizer code centroid
@@ -33,7 +31,7 @@ template <
         typename CentroidT,
         int DimsPerSubQuantizer,
         bool L2Distance>
-__global__ void __launch_bounds__(LAUNCH_BOUND, 3) pqCodeDistances(
+__global__ void __launch_bounds__(getLaunchBound(), 3) pqCodeDistances(
         Tensor<float, 2, true> queries,
         int queriesPerBlock,
         Tensor<CentroidT, 2, true> coarseCentroids,

faiss/gpu/utils/DeviceDefs.cuh

@@ -8,17 +8,28 @@
 #pragma once
 
 #include <cuda.h>
+#ifdef USE_AMD_ROCM
+#include <rocm-core/rocm_version.h> // ROCm version macros
+#if ROCM_VERSION_MAJOR >= 7
+#include <rocprim/intrinsics/arch.hpp> // rocprim::arch::wavefront
+#endif
+#endif
 
 namespace faiss {
 namespace gpu {
 
 #ifdef USE_AMD_ROCM
 
+#if ROCM_VERSION_MAJOR < 7
 #if __AMDGCN_WAVEFRONT_SIZE == 32u
 constexpr int kWarpSize = 32;
 #else
 constexpr int kWarpSize = 64;
 #endif
+#else
+// ROCm 7.0 and above
+constexpr __device__ int kWarpSize = rocprim::arch::wavefront::max_size();
+#endif
 
 // This is a memory barrier for intra-warp writes to shared memory.
 __forceinline__ __device__ void warpFence() {