parallelize pruned_array_lookup_cpu

Summary: as titled, parallelize pruned_array_lookup_cpu across multiple threads. Differential Revision: D47923801 fbshipit-source-id: c94f0ee66ca7d0550dbaa67eaf5bebdaa6c3c948
pytorch · Aug 1, 2023 · da77edf · da77edf
1 parent f5d37c3
commit da77edf
Showing 1 changed file with 21 additions and 18 deletions.
diff --git a/fbgemm_gpu/codegen/embedding_forward_quantized_cpu_template.cpp b/fbgemm_gpu/codegen/embedding_forward_quantized_cpu_template.cpp
@@ -11,6 +11,7 @@
 
 #include <ATen/ATen.h>
 #include <ATen/Context.h>
+#include <ATen/Parallel.h>
 
 #include "fbgemm_gpu/cpu_utils.h"
 #include "fbgemm_gpu/dispatch_macros.h"
@@ -524,24 +525,26 @@ Tensor pruned_array_lookup_cpu(
 
     const auto index_remappings_acc = index_remappings.data_ptr<int32_t>();
     const auto index_remappings_offsets_acc = index_remappings_offsets.data_ptr<int64_t>();
-for (const auto t : c10::irange(T)) {
-        int64_t index_remappings_start = index_remappings_offsets_acc[t];
-        int64_t index_remappings_end = index_remappings_offsets_acc[t + 1];
-        int64_t capacity = index_remappings_end - index_remappings_start;
-        int32_t indices_start = offsets_acc[t * B];
-        int32_t indices_end = offsets_acc[(t + 1) * B];
-        if (capacity > 0) {
-for (const auto i : c10::irange(indices_start,indices_end)) {
-                int32_t idx = indices_acc[i];
-                dense_indices_acc[i] = index_remappings_acc[index_remappings_start + idx];
-            }
-        } else {
-            std::memcpy(
-                dense_indices_acc + indices_start,
-                indices_acc + indices_start,
-                (indices_end - indices_start) * sizeof(int32_t));
-        }
-    }
+    at::parallel_for(0, T, 1, [&](int64_t begin, int64_t end) {
+      for (const auto t : c10::irange(begin, end)) {
+          int64_t index_remappings_start = index_remappings_offsets_acc[t];
+          int64_t index_remappings_end = index_remappings_offsets_acc[t + 1];
+          int64_t capacity = index_remappings_end - index_remappings_start;
+          int32_t indices_start = offsets_acc[t * B];
+          int32_t indices_end = offsets_acc[(t + 1) * B];
+          if (capacity > 0) {
+            for (const auto i : c10::irange(indices_start,indices_end)) {
+                  int32_t idx = indices_acc[i];
+                  dense_indices_acc[i] = index_remappings_acc[index_remappings_start + idx];
+              }
+          } else {
+              std::memcpy(
+                  dense_indices_acc + indices_start,
+                  indices_acc + indices_start,
+                  (indices_end - indices_start) * sizeof(int32_t));
+          }
+      }
+    });
     return dense_indices;
 }