Vectorize index remapping for pruned embedding tables (pytorch#607)

Summary: Pull Request resolved: pytorch#607 Vectorize index remapping for pruned tables ~~TODO: Benchmarking~~ Benchmarking numbers added in the summary of D28361076 Reviewed By: jspark1105 Differential Revision: D28346429 fbshipit-source-id: d0dee350b132d171d6b683b99775ee292eb103b2
q10 · Jun 7, 2021 · 76cf190 · 76cf190
1 parent b3f1af5
commit 76cf190
Show file tree

Hide file tree

Showing 7 changed files with 672 additions and 0 deletions.
diff --git a/defs.bzl b/defs.bzl
@@ -103,6 +103,7 @@ def get_fbgemm_avx512_srcs(msvc = False):
     return [
         #All the source files that use avx512 instructions statically
         "src/FbgemmBfloat16ConvertAvx512.cc",
+        "src/EmbeddingSpMDMAvx512.cc",
         "src/FbgemmFloat16ConvertAvx512.cc",
         "src/FbgemmSparseDenseAvx512.cc",
         "src/FbgemmSparseDenseInt8Avx512.cc",

diff --git a/include/fbgemm/FbgemmEmbedding.h b/include/fbgemm/FbgemmEmbedding.h
@@ -258,6 +258,28 @@ FBGEMM_API bool EmbeddingSpMDMBlockSize1_(
     bool is_weight_positional = false,
     bool use_offsets = true);
 
+template <typename IndexType, bool HAS_WEIGHTS>
+void compressed_indices_remap_avx512(
+    std::int32_t offsets_numel,
+    const IndexType* indices,
+    const int32_t* compressed_indices_mapping,
+    const IndexType* offsets,
+    const float* weights, // optional, can be null,
+    IndexType* out_indices,
+    IndexType* out_offsets,
+    float* out_weights);
+
 } // namespace internal
 
+template <typename IndexType>
+FBGEMM_API void compressed_indices_remap(
+    std::int32_t offsets_numel,
+    const IndexType* indices,
+    const int32_t* compressed_indices_mapping,
+    const IndexType* offsets,
+    const float* weights, // optional, can be null,
+    IndexType* out_indices,
+    IndexType* out_offsets,
+    float* out_weights);
+
 } // namespace fbgemm
diff --git a/src/EmbeddingSpMDM.cc b/src/EmbeddingSpMDM.cc
@@ -1192,4 +1192,70 @@ INSTANTIATE_SPMDM_INDEX_T(uint8_t)
 #undef INSTANTIATE_SPMDM_NAME
 #undef INSTANTIATE_SPMDM_BASE
 
+template <typename IndexType>
+void compressed_indices_remap(
+    std::int32_t offsets_len,
+    const IndexType* indices,
+    const int32_t* compressed_indices_mapping,
+    const IndexType* offsets,
+    const float* weights, // optional, can be null,
+    IndexType* out_indices,
+    IndexType* out_offsets,
+    float* out_weights) {
+  if (!cpuinfo_initialize()) {
+    throw std::runtime_error("Failed to initialize cpuinfo!");
+  }
+
+  const inst_set_t isa = fbgemmInstructionSet();
+  if (isZmm(isa)) {
+    if (weights == nullptr) {
+      internal::compressed_indices_remap_avx512<IndexType, false>(
+          offsets_len,
+          indices,
+          compressed_indices_mapping,
+          offsets,
+          weights,
+          out_indices,
+          out_offsets,
+          out_weights);
+    } else {
+      internal::compressed_indices_remap_avx512<IndexType, true>(
+          offsets_len,
+          indices,
+          compressed_indices_mapping,
+          offsets,
+          weights,
+          out_indices,
+          out_offsets,
+          out_weights);
+    }
+  } else {
+    compressed_indices_remap_ref<IndexType>(
+        offsets_len,
+        indices,
+        compressed_indices_mapping,
+        offsets,
+        weights,
+        out_indices,
+        out_offsets,
+        out_weights);
+  }
+}
+
+#define INSTANTIATE_REMAP_BASE(INDEX_TYPE)           \
+  template FBGEMM_API void compressed_indices_remap( \
+      std::int32_t offsets_numel,                    \
+      const INDEX_TYPE* indices,                     \
+      const int32_t* compressed_indices_mapping,     \
+      const INDEX_TYPE* offsets,                     \
+      const float* weights,                          \
+      INDEX_TYPE* out_indices,                       \
+      INDEX_TYPE* out_offsets,                       \
+      float* out_weights);
+
+INSTANTIATE_REMAP_BASE(int32_t);
+INSTANTIATE_REMAP_BASE(int64_t);
+
+#undef INSTANTIATE_REMAP_BASE
+
 } // namespace fbgemm