Follow up on BC issue for open sourcing TBE inplace update op

Reviewed By: jspark1105 Differential Revision: D41717190 fbshipit-source-id: 5a32a5c183db1364082638ee945d19c33a0de950
pytorch · Dec 4, 2022 · 79eb9f9 · 79eb9f9
1 parent 4cd267c
commit 79eb9f9
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 9 deletions.
diff --git a/fbgemm_gpu/src/embedding_inplace_update.h → ...ude/fbgemm_gpu/embedding_inplace_update.h b/fbgemm_gpu/src/embedding_inplace_update.h → ...ude/fbgemm_gpu/embedding_inplace_update.h
@@ -71,4 +71,22 @@ void embedding_inplace_update_cuda(
     c10::optional<Tensor> lxu_cache_weights = c10::nullopt,
     c10::optional<Tensor> lxu_cache_locations = c10::nullopt);
 
+void embedding_inplace_update_cpu(
+    Tensor dev_weights,
+    Tensor uvm_weights,
+    Tensor weights_placements,
+    Tensor weights_offsets,
+    Tensor weights_tys,
+    Tensor D_offsets,
+    Tensor update_weights,
+    Tensor update_table_idx,
+    Tensor update_row_idx,
+    Tensor update_offsets,
+    const int64_t row_alignment,
+    c10::optional<Tensor> lxu_cache_weights =
+        c10::nullopt, // Not used, to match cache interface for CUDA op
+    c10::optional<Tensor> lxu_cache_locations =
+        c10::nullopt // Not used, to match cache interface for CUDA op
+);
+
 } // namespace fbgemm_gpu
diff --git a/fbgemm_gpu/src/embedding_inplace_update.cu b/fbgemm_gpu/src/embedding_inplace_update.cu
@@ -10,7 +10,7 @@
 
 #include <c10/cuda/CUDAGuard.h>
 
-#include "embedding_inplace_update.h"
+#include "fbgemm_gpu/embedding_inplace_update.h"
 #include "fbgemm_gpu/fbgemm_cuda_utils.cuh"
 
 using Tensor = at::Tensor;

diff --git a/fbgemm_gpu/src/embedding_inplace_update_cpu.cpp b/fbgemm_gpu/src/embedding_inplace_update_cpu.cpp
@@ -12,7 +12,7 @@
 #include <ATen/ATen.h>
 #include <torch/library.h>
 
-#include "embedding_inplace_update.h"
+#include "fbgemm_gpu/embedding_inplace_update.h"
 
 using Tensor = at::Tensor;
 
@@ -72,11 +72,8 @@ void embedding_inplace_update_cpu(
     Tensor update_row_idx,
     Tensor update_offsets,
     const int64_t row_alignment,
-    c10::optional<Tensor> lxu_cache_weights =
-        c10::nullopt, // Not used, to match cache interface for CUDA op
-    c10::optional<Tensor> lxu_cache_locations =
-        c10::nullopt // Not used, to match cache interface for CUDA op
-) {
+    c10::optional<Tensor> lxu_cache_weights,
+    c10::optional<Tensor> lxu_cache_locations) {
   TENSOR_ON_CPU(dev_weights);
   TENSOR_ON_CPU(uvm_weights);
   TENSOR_ON_CPU(weights_placements);

diff --git a/fbgemm_gpu/src/embedding_inplace_update_gpu.cpp b/fbgemm_gpu/src/embedding_inplace_update_gpu.cpp
@@ -8,7 +8,7 @@
 #include <ATen/core/op_registration/op_registration.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <torch/library.h>
-#include "embedding_inplace_update.h"
+#include "fbgemm_gpu/embedding_inplace_update.h"
 
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   DISPATCH_TO_CUDA(

diff --git a/fbgemm_gpu/test/embedding_inplace_update_test.cpp b/fbgemm_gpu/test/embedding_inplace_update_test.cpp
@@ -6,7 +6,7 @@
  */
 #include <folly/Random.h>
 #include <gtest/gtest.h>
-#include "deeplearning/fbgemm/fbgemm_gpu/src/embedding_inplace_update.h"
+#include "fbgemm_gpu/embedding_inplace_update.h"
 
 using namespace ::testing;
 using namespace fbgemm_gpu;