Fix missing operator registration (#4042)

Summary:
X-link: facebookresearch/FBGEMM#1129

- Fix missing operator registration

Pull Request resolved: #4042

Reviewed By: sryap, spcyppt

Differential Revision: D73866259

Pulled By: q10

fbshipit-source-id: 3c87a4090fa908a99fe1b38460c1973b19a213f3

Author: q10

fbgemm_gpu/FbgemmGpu.cmake

@@ -39,6 +39,7 @@ set(fbgemm_gpu_sources_cpu_static
     src/layout_transform_ops/layout_transform_ops_cpu.cpp
     src/quantize_ops/quantize_ops_cpu.cpp
     src/quantize_ops/quantize_ops_meta.cpp
+    src/sparse_ops/sparse_async_batched_cumsum.cpp
     src/sparse_ops/sparse_ops_cpu.cpp
     src/sparse_ops/sparse_ops_meta.cpp
     ${tbe_eeg_cpu_sources})
@@ -101,6 +102,7 @@ if(NOT FBGEMM_BUILD_VARIANT STREQUAL BUILD_VARIANT_CPU)
       src/quantize_ops/quantize_msfp.cu
       src/quantize_ops/quantize_padded_fp8_rowwise.cu
       src/quantize_ops/quantize_mx.cu
+      src/sparse_ops/sparse_async_batched_cumsum.cu
       src/sparse_ops/sparse_block_bucketize_features.cu
       src/sparse_ops/sparse_bucketize_features.cu
       src/sparse_ops/sparse_batched_unary_embeddings.cu

fbgemm_gpu/src/sparse_ops/sparse_async_batched_cumsum.cu

@@ -6,11 +6,19 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#ifdef __HIP_PLATFORM_AMD__
+#include <hipcub/block/block_scan.hpp>
+#else
 #include <cub/block/block_scan.cuh>
+#endif
 #include "common.cuh"
 
 static constexpr uint32_t kMaxThreads = 1024;
 
+#ifdef __HIP_PLATFORM_AMD__
+namespace cub = hipcub;
+#endif
+
 namespace fbgemm_gpu {
 
 C10_ALWAYS_INLINE uint32_t next_power_of_2(uint32_t n) {