Skip to content

Commit

Permalink
update cmake for dense TBE VBE support (pytorch#2641)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: pytorch#2641

follow up to pytorch#2620

Differential Revision: https://internalfb.com/D57938517
  • Loading branch information
joshuadeng authored and facebook-github-bot committed May 30, 2024
1 parent 2dd8b63 commit 03140fd
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions fbgemm_gpu/FbgemmGpu.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,8 @@ set(GPU_OPTIMIZERS ${COMMON_OPTIMIZERS} ${GPU_ONLY_OPTIMIZERS})
set(VBE_OPTIMIZERS
rowwise_adagrad
rowwise_adagrad_with_counter
sgd)
sgd
dense)

# Optimizers with the GWD support
set(GWD_OPTIMIZERS
Expand Down Expand Up @@ -151,6 +152,8 @@ set(gen_gpu_kernel_source_files
"gen_embedding_backward_dense_indice_weights_codegen_cuda.cu"
"gen_embedding_backward_split_indice_weights_codegen_cuda.cu"
"gen_embedding_backward_ssd_indice_weights_codegen_cuda.cu"
"gen_embedding_forward_dense_weighted_vbe_codegen_cuda.cu"
"gen_embedding_forward_dense_unweighted_vbe_codegen_cuda.cu"
"gen_embedding_forward_split_weighted_vbe_codegen_cuda.cu"
"gen_embedding_forward_split_unweighted_vbe_codegen_cuda.cu"
"gen_batch_index_select_dim0_forward_codegen_cuda.cu"
Expand Down Expand Up @@ -203,7 +206,9 @@ endforeach()
foreach(wdesc weighted unweighted)
list(APPEND gen_gpu_kernel_source_files
"gen_embedding_forward_split_${wdesc}_vbe_kernel.cu"
"gen_embedding_backward_split_${wdesc}_vbe_device_kernel.cuh")
"gen_embedding_backward_split_${wdesc}_vbe_device_kernel.cuh"
"gen_embedding_forward_dense_${wdesc}_vbe_kernel.cu")

endforeach()

# Generate GWD files
Expand Down Expand Up @@ -261,6 +266,9 @@ foreach(optimizer ${ALL_OPTIMIZERS})
"gen_embedding_backward_split_${optimizer}_pt2_cuda_wrapper.cpp")
endforeach()

list(APPEND gen_gpu_host_source_files
"gen_embedding_backward_split_dense.cpp")

foreach(optimizer ${CPU_OPTIMIZERS})
list(APPEND gen_cpu_source_files
"gen_embedding_backward_${optimizer}_split_cpu.cpp")
Expand Down Expand Up @@ -462,7 +470,6 @@ set(fbgemm_gpu_sources_static_cpu
if(NOT FBGEMM_CPU_ONLY)
list(APPEND fbgemm_gpu_sources_static_cpu
codegen/inference/embedding_forward_quantized_host.cpp
codegen/training/backward/embedding_backward_dense_host.cpp
codegen/utils/embedding_bounds_check_host.cpp
src/memory_utils/memory_utils.cpp
src/memory_utils/memory_utils_ops.cpp
Expand Down

0 comments on commit 03140fd

Please sign in to comment.