Redefine FBGEMM targets with gpu_cpp_library [8/N] (pytorch#194)

Summary:
X-link: pytorch#3105

Pull Request resolved: facebookresearch/FBGEMM#194

- Rename `embedding_ops_training` and `embedding_ops_inference` CPU targets and move `embedding_ops_common_cpu` into its own target

Reviewed By: sryap

Differential Revision: D62415798

fbshipit-source-id: 6d3a8977fb4622100051f706197f21e510ee0ddc

Author: q10

fbgemm_gpu/codegen/training/python/split_embedding_codegen_lookup_invoker.template

@@ -18,21 +18,18 @@ from .lookup_args{{ sdesc }} import *
 
 {%- if is_fbcode %}
 
+from fbgemm_gpu.utils.loader import load_torch_module, load_torch_module_bc
 # Provide compatibility to downstream packages for eventual migration to the split training / inference packages
 try:
-    try:
-        torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_gpu")
-    except Exception:
-        if torch.version.hip:
-            torch.ops.load_library(
-                "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip_training"
-            )
-        else:
-            torch.ops.load_library(
-                "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cuda_training"
-            )
-    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training")
-
+    load_torch_module(
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_gpu",
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cuda_training",
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip_training",
+    )
+    load_torch_module_bc(
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_cpu",
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training",
+    )
 except Exception:
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops")
     torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu")

fbgemm_gpu/codegen/training/python/split_embedding_optimizer_codegen.template

@@ -22,6 +22,8 @@ import logging
 
 {%- if is_fbcode %}
 
+from fbgemm_gpu.utils.loader import load_torch_module, load_torch_module_bc
+
 def _load_library(unified_path: str, cuda_path: str, hip_path: str) -> None:
     try:
         torch.ops.load_library(unified_path)
@@ -32,17 +34,18 @@ def _load_library(unified_path: str, cuda_path: str, hip_path: str) -> None:
         else:
             torch.ops.load_library(cuda_path)
 
-torch.ops.load_library(
-    "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training"
+load_torch_module_bc(
+    "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_cpu",
+    "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training",
 )
 
-_load_library(
+load_torch_module(
     "//deeplearning/fbgemm/fbgemm_gpu/codegen:optimizer_ops",
     "//deeplearning/fbgemm/fbgemm_gpu/codegen:optimizer_ops_cuda",
     "//deeplearning/fbgemm/fbgemm_gpu/codegen:optimizer_ops_hip",
 )
 
-_load_library(
+load_torch_module(
     "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings",
     "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings",
     "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings_hip",

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py

@@ -31,16 +31,21 @@
     round_up,
     SplitState,
 )
-from fbgemm_gpu.utils.loader import load_torch_module
+from fbgemm_gpu.utils.loader import load_torch_module, load_torch_module_bc
 
 try:
     load_torch_module(
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_inference_gpu",
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cuda_inference",
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip_inference",
     )
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_inference"
+except Exception:
+    pass
+
+try:
+    load_torch_module_bc(
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_inference_cpu",
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_inference",
     )
 except Exception:
     pass

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -50,17 +50,17 @@
     is_torchdynamo_compiling,
 )
 
-from fbgemm_gpu.utils.loader import load_torch_module
+from fbgemm_gpu.utils.loader import load_torch_module, load_torch_module_bc
 
 try:
     load_torch_module(
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_gpu",
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cuda_training",
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip_training",
     )
-
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training"
+    load_torch_module_bc(
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_training_cpu",
+        "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu_training",
     )
 except Exception:
     pass

fbgemm_gpu/fbgemm_gpu/utils/loader.py

@@ -27,3 +27,10 @@ def load_torch_module(
             if not cuda_path:
                 cuda_path = f"{unified_path}_cuda"
             torch.ops.load_library(cuda_path)
+
+
+def load_torch_module_bc(new_path: str, old_path: str) -> None:
+    try:
+        torch.ops.load_library(new_path)
+    except Exception:
+        torch.ops.load_library(old_path)