add AMD specific includes in cuda_prelude.h (pytorch#3614)

Summary:

X-link: facebookresearch/FBGEMM#691

as title

Reviewed By: q10

Differential Revision: D68638427

Author: Bangsheng Tang

fbgemm_gpu/include/fbgemm_gpu/utils/cuda_prelude.cuh

@@ -9,16 +9,31 @@
 #pragma once
 
 #include <ATen/ATen.h>
+
 #include <cuda.h>
+
+#ifdef __HIP_PLATFORM_AMD__
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#include <ATen/cuda/detail/UnpackRaw.cuh> // For at::cuda::philox::unpack
+
+#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h> // @manual
+#else
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
+#endif
 #include <cassert>
 
 namespace {
 
 inline int get_device_sm_cnt_() {
+#ifdef __HIP_PLATFORM_AMD__
+  hipDeviceProp_t deviceProp;
+  hipGetDeviceProperties(&deviceProp, c10::hip::current_device());
+  return deviceProp.multiProcessorCount;
+#else
   cudaDeviceProp* deviceProp =
       at::cuda::getDeviceProperties(c10::cuda::current_device());
   return deviceProp->multiProcessorCount;
+#endif
 }
 
 } // namespace