Fix PT2 wrapper registrations (pytorch#805)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#805

- Fix PT2 wrapper registrations

X-link: pytorch#3721

Reviewed By: spcyppt

Differential Revision: D69971546

Pulled By: q10

fbshipit-source-id: 9e964f14b25c24ad3240e4e65d87fa2ae08d5b21

Author: q10

.github/scripts/fbgemm_gpu_install.bash

@@ -98,6 +98,7 @@ __install_check_operator_registrations () {
   else
     local test_operators=(
       "torch.ops.fbgemm.asynchronous_inclusive_cumsum"
+      "torch.ops.fbgemm.split_embedding_codegen_lookup_sgd_function_pt2"
     )
   fi
 

fbgemm_gpu/cmake/TbeTraining.cmake

@@ -26,13 +26,15 @@ handle_genfiles_rocm(gen_gpu_files_forward_split)
 get_tbe_sources_list(static_cpu_files_training)
 get_tbe_sources_list(gen_cpu_files_training)
 get_tbe_sources_list(gen_gpu_files_training)
+get_tbe_sources_list(gen_cpu_files_training_pt2)
 get_tbe_sources_list(gen_gpu_files_training_pt2)
 get_tbe_sources_list(gen_gpu_files_training_dense)
 get_tbe_sources_list(gen_gpu_files_training_split_host)
 get_tbe_sources_list(gen_gpu_files_training_gwd)
 get_tbe_sources_list(gen_gpu_files_training_vbe)
 handle_genfiles_rocm(gen_cpu_files_training)
 handle_genfiles_rocm(gen_gpu_files_training)
+handle_genfiles_rocm(gen_cpu_files_training_pt2)
 handle_genfiles_rocm(gen_gpu_files_training_pt2)
 handle_genfiles_rocm(gen_gpu_files_training_dense)
 handle_genfiles_rocm(gen_gpu_files_training_split_host)
@@ -166,6 +168,8 @@ gpu_cpp_library(
     SHARED
   INCLUDE_DIRS
     ${fbgemm_sources_include_directories}
+  CPU_SRCS
+    ${gen_cpu_files_training_pt2}
   GPU_SRCS
     ${gen_gpu_files_training_pt2}
   NVCC_FLAGS

fbgemm_gpu/cmake/tbe_sources.py

@@ -318,13 +318,13 @@
 static_cpu_files_common = [
     "codegen/utils/embedding_bounds_check_host_cpu.cpp",
     "codegen/training/forward/embedding_forward_split_cpu.cpp",
+    "codegen/training/pt2/pt2_autograd_utils.cpp",
 ]
 
 static_gpu_files_common = [
     "codegen/utils/embedding_bounds_check_v1.cu",
     "codegen/utils/embedding_bounds_check_v2.cu",
     "codegen/utils/embedding_bounds_check_host.cpp",
-    "codegen/training/pt2/pt2_autograd_utils.cpp",
 ]
 
 gen_cpu_files_training = (
@@ -335,17 +335,13 @@
         "gen_embedding_backward_split_{}_cpu.cpp".format(optimizer)
         for optimizer in ALL_OPTIMIZERS
     ]
-    + [
-        "gen_embedding_backward_split_{}_pt2_cpu_wrapper.cpp".format(optimizer)
-        for optimizer in ALL_OPTIMIZERS
-    ]
     + [
         "gen_embedding_backward_{}_split_cpu.cpp".format(optimizer)
         for optimizer in CPU_OPTIMIZERS
     ]
 )
 
-gen_gpu_files_training_pt2 = (
+gen_cpu_files_training_pt2 = (
     [
         "gen_embedding_split_{}_pt2_autograd.cpp".format(optimizer)
         for optimizer in ALL_OPTIMIZERS
@@ -355,15 +351,19 @@
         for optimizer in SSD_OPTIMIZERS
     ]
     + [
-        "gen_embedding_backward_split_{}_pt2_cuda_wrapper.cpp".format(optimizer)
+        "gen_embedding_backward_split_{}_pt2_cpu_wrapper.cpp".format(optimizer)
         for optimizer in ALL_OPTIMIZERS
     ]
-    + [
-        "gen_embedding_backward_ssd_{}_pt2_cuda_wrapper.cpp".format(optimizer)
-        for optimizer in SSD_OPTIMIZERS
-    ]
 )
 
+gen_gpu_files_training_pt2 = [
+    "gen_embedding_backward_split_{}_pt2_cuda_wrapper.cpp".format(optimizer)
+    for optimizer in ALL_OPTIMIZERS
+] + [
+    "gen_embedding_backward_ssd_{}_pt2_cuda_wrapper.cpp".format(optimizer)
+    for optimizer in SSD_OPTIMIZERS
+]
+
 gen_gpu_files_training_dense = [
     # Dense host and kernel, and forward-quantized host src files
     fstring.format(wdesc)