Revert D51688407: Refactor embedding_bounds_check

Differential Revision: D51688407 Original commit changeset: e021212f7c71 Original Phabricator Diff: D51688407 fbshipit-source-id: fa2c01bbd0f48609561dc6a353f9088bb91f0177
pytorch · Dec 1, 2023 · c58679a · c58679a
1 parent bf980a7
commit c58679a
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 6 deletions.
diff --git a/fbgemm_gpu/CMakeLists.txt b/fbgemm_gpu/CMakeLists.txt
@@ -572,6 +572,7 @@ set(fbgemm_gpu_sources_static_cpu
     codegen/embedding_forward_split_cpu.cpp
     codegen/embedding_forward_quantized_host_cpu.cpp
     codegen/embedding_backward_dense_host_cpu.cpp
+    codegen/embedding_bounds_check_host_cpu.cpp
     src/merge_pooled_embedding_ops/merge_pooled_embedding_ops_cpu.cpp
     src/permute_pooled_embedding_ops/permute_pooled_embedding_function.cpp
     src/permute_pooled_embedding_ops/permute_pooled_embedding_ops_cpu.cpp
@@ -591,13 +592,13 @@ set(fbgemm_gpu_sources_static_cpu
     src/split_embeddings_cache/lru_cache_populate_byte.cpp
     src/split_embeddings_cache/lxu_cache.cpp
     src/split_embeddings_cache/split_embeddings_cache_ops.cpp
-    src/split_embeddings_utils/embedding_bounds_check_cpu.cpp
     codegen/batch_index_select_dim0_cpu_host.cpp)
 
 if(NOT FBGEMM_CPU_ONLY)
   list(APPEND fbgemm_gpu_sources_static_cpu
     codegen/embedding_forward_quantized_host.cpp
     codegen/embedding_backward_dense_host.cpp
+    codegen/embedding_bounds_check_host.cpp
     src/memory_utils/memory_utils.cpp
     src/memory_utils/memory_utils_ops.cpp
     src/layout_transform_ops/layout_transform_ops_gpu.cpp
@@ -609,7 +610,6 @@ if(NOT FBGEMM_CPU_ONLY)
     src/split_embeddings_cache/split_embeddings_cache_ops.cu
     src/metric_ops/metric_ops_host.cpp
     src/embedding_inplace_ops/embedding_inplace_update_gpu.cpp
-    src/split_embeddings_utils/embedding_bounds_check_gpu.cpp
     src/input_combine_ops/input_combine_gpu.cpp
     codegen/batch_index_select_dim0_host.cpp)
 
@@ -639,6 +639,7 @@ endif()
 
 if(NOT FBGEMM_CPU_ONLY)
   set(fbgemm_gpu_sources_static_gpu
+      codegen/embedding_bounds_check.cu
       codegen/embedding_forward_quantized_split_lookup.cu
       src/memory_utils/memory_utils.cu
       src/memory_utils/memory_utils_ops.cu
@@ -702,7 +703,6 @@ if(NOT FBGEMM_CPU_ONLY)
       src/split_embeddings_cache/lxu_cache.cu
       src/split_embeddings_cache/linearize_cache_indices.cu
       src/split_embeddings_cache/reset_weight_momentum.cu
-      src/split_embeddings_utils/embedding_bounds_check.cu
       src/split_embeddings_utils/generate_vbe_metadata.cu
       src/split_embeddings_utils/get_infos_metadata.cu
       src/split_embeddings_utils/radix_sort_pairs.cu

diff --git a/...mbeddings_utils/embedding_bounds_check.cu → fbgemm_gpu/codegen/embedding_bounds_check.cu b/...mbeddings_utils/embedding_bounds_check.cu → fbgemm_gpu/codegen/embedding_bounds_check.cu
@@ -6,9 +6,10 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include "fbgemm_gpu/embedding_backward_template_helpers.cuh"
+
 #include <c10/cuda/CUDADeviceAssertion.h>
 #include <c10/cuda/CUDAException.h>
-#include "fbgemm_gpu/embedding_backward_template_helpers.cuh"
 
 using Tensor = at::Tensor;
 using namespace fbgemm_gpu;

diff --git a/...ings_utils/embedding_bounds_check_gpu.cpp → ...u/codegen/embedding_bounds_check_host.cpp b/...ings_utils/embedding_bounds_check_gpu.cpp → ...u/codegen/embedding_bounds_check_host.cpp
@@ -9,7 +9,9 @@
 #include <ATen/ATen.h>
 #include <ATen/TypeDefault.h>
 #include <ATen/core/op_registration/op_registration.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <torch/library.h>
+
 #include "fbgemm_gpu/sparse_ops_utils.h"
 
 using Tensor = at::Tensor;

diff --git a/...ings_utils/embedding_bounds_check_cpu.cpp → ...degen/embedding_bounds_check_host_cpu.cpp b/...ings_utils/embedding_bounds_check_cpu.cpp → ...degen/embedding_bounds_check_host_cpu.cpp
diff --git a/fbgemm_gpu/docs/Doxyfile.in b/fbgemm_gpu/docs/Doxyfile.in
@@ -932,10 +932,10 @@ INPUT                  = "../include/fbgemm_gpu" \
                          "../codegen/embedding_backward_split_host_cpu_template.cpp" \
                          "../codegen/embedding_forward_quantized_host.cpp" \
                          "../codegen/embedding_forward_quantized_host_cpu.cpp" \
+                         "../codegen/embedding_bounds_check_host.cpp" \
+                         "../codegen/embedding_bounds_check_host_cpu.cpp" \
                          "../codegen/embedding_backward_dense_host.cpp" \
                          "../codegen/embedding_forward_template_helpers.cuh " \
-                         "../src/split_embeddings_utils/embedding_bounds_check_cpu.cpp" \
-                         "../src/split_embeddings_utils/embedding_bounds_check_gpu.cpp" \
                          "../src/permute_pooled_embedding_ops_gpu.cpp" \
 
 # This tag can be used to specify the character encoding of the source files