fix sparse_permute_1d kernel to support double dtype (#4969)

Summary:
Pull Request resolved: #4969

X-link: https://github.com/facebookresearch/FBGEMM/pull/1986

Feature score collection will generate a weight tensor with fp64 dtype, we need to make sure the distribution kernels support this dtype

Reviewed By: q10

Differential Revision: D83788055

fbshipit-source-id: ca8db60acf638413db839316ab27b3da6dc6785f

Author: emlin

fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp

@@ -943,7 +943,7 @@ std::tuple<Tensor, Tensor, std::optional<Tensor>> permute_1D_sparse_data_cpu(
         FBGEMM_DISPATCH_ALL_TYPES(
             indices.scalar_type(), "permute_1D_indices_weights_kernel_2", [&] {
               using indices_t = scalar_t;
-              FBGEMM_DISPATCH_FLOAT_ONLY(
+              FBGEMM_DISPATCH_FLOAT_AND_DOUBLE(
                   weights.has_value() ? weights.value().scalar_type()
                                       : at::ScalarType::Float,
                   "permute_1D_indices_weights_kernel_3",
@@ -2971,7 +2971,7 @@ std::tuple<Tensor, Tensor, std::optional<Tensor>> permute_sparse_features_cpu(
   permuted_indices = at::empty(permuted_lengths_sum, indices.options());
   AT_DISPATCH_INDEX_TYPES(
       input_offsets.scalar_type(), "permute_data_kernel_1", ([&] {
-        FBGEMM_DISPATCH_FLOAT_ONLY(
+        FBGEMM_DISPATCH_FLOAT_AND_DOUBLE(
             weights.has_value() ? weights.value().scalar_type()
                                 : at::ScalarType::Float,
             "permute_data_kernel_2",

fbgemm_gpu/src/sparse_ops/sparse_permute_1d.cu

@@ -141,7 +141,7 @@ permute_1D_sparse_data_cuda(
                 const auto weights_value_contig = weights_value.contiguous();
                 permuted_weights =
                     at::empty(permuted_indices_size, weights_value.options());
-                FBGEMM_DISPATCH_ALL_TYPES(
+                FBGEMM_DISPATCH_ALL_TYPES_AND_DOUBLE(
                     weights_value.scalar_type(),
                     "permute_1D_data_kernel_3",
                     [&] {