Skip to content

Commit

Permalink
Refactor embedding_bounds_check (pytorch#2178)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: pytorch#2178

- Migrate embedding_bounds_check out of codegen and into split_embeddings_utils

Reviewed By: sryap

Differential Revision: D51688407

fbshipit-source-id: e021212f7c7123e770998cd6115f4a99461b1562
  • Loading branch information
q10 authored and facebook-github-bot committed Nov 30, 2023
1 parent 327dcf9 commit bf980a7
Show file tree
Hide file tree
Showing 5 changed files with 6 additions and 9 deletions.
6 changes: 3 additions & 3 deletions fbgemm_gpu/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -572,7 +572,6 @@ 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
Expand All @@ -592,13 +591,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
Expand All @@ -610,6 +609,7 @@ 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)

Expand Down Expand Up @@ -639,7 +639,6 @@ 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
Expand Down Expand Up @@ -703,6 +702,7 @@ 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
Expand Down
4 changes: 2 additions & 2 deletions fbgemm_gpu/docs/Doxyfile.in
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,9 @@
* 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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,7 @@
#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;
Expand Down

0 comments on commit bf980a7

Please sign in to comment.