Redefine FBGEMM targets with gpu_cpp_library [6/N] (pytorch#3092)

Summary:
Pull Request resolved: pytorch#3092

X-link: facebookresearch/FBGEMM#185

- Redefine `split_table_batched_embeddings_gpu` target with `gpu_cpp_library`

Reviewed By: spcyppt

Differential Revision: D62282866

fbshipit-source-id: 217251a2ca2ff2502bb97f2b67ef07b15f5ccedc

fbgemm_gpu/bench/split_embeddings_cache_benchmark.py

@@ -13,7 +13,6 @@
 import click
 import numpy as np
 import torch
-
 from fbgemm_gpu.split_embedding_configs import SparseType
 from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
     CacheAlgorithm,
@@ -23,6 +22,8 @@
     IntNBitTableBatchedEmbeddingBagsCodegen,
 )
 
+from fbgemm_gpu.utils.loader import load_torch_module
+
 from torch import nn, Tensor
 
 logging.basicConfig(level=logging.DEBUG)
@@ -31,16 +32,11 @@
     # pyre-ignore[21]
     from fbgemm_gpu import open_source  # noqa: F401
 except Exception:
-    if torch.version.hip:
-        torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:cumem_utils_hip")
-        torch.ops.load_library(
-            "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings_hip"
-        )
-    else:
-        torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:cumem_utils")
-        torch.ops.load_library(
-            "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings"
-        )
+    for module in [
+        "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings",
+        "//deeplearning/fbgemm/fbgemm_gpu:cumem_utils",
+    ]:
+        load_torch_module(module)
 
 
 # pyre-ignore

fbgemm_gpu/bench/ssd_table_batched_embeddings_benchmark.py

@@ -36,20 +36,15 @@
     SSDTableBatchedEmbeddingBags,
 )
 from fbgemm_gpu.tbe.utils import generate_requests, get_device, round_up, TBERequest
+from fbgemm_gpu.utils.loader import load_torch_module
 from torch.autograd.profiler import record_function
 from torch.profiler import profile
 
 logging.basicConfig(level=logging.DEBUG)
 
-if torch.version.hip:
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu:ssd_split_table_batched_embeddings_hip"
-    )
-else:
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu:ssd_split_table_batched_embeddings"
-    )
-
+load_torch_module(
+    "//deeplearning/fbgemm/fbgemm_gpu:ssd_split_table_batched_embeddings",
+)
 
 logging.basicConfig(level=logging.DEBUG)
 

fbgemm_gpu/codegen/training/python/split_embedding_optimizer_codegen.template

@@ -42,20 +42,21 @@ _load_library(
     "//deeplearning/fbgemm/fbgemm_gpu/codegen:optimizer_ops_hip",
 )
 
+_load_library(
+    "//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",
+)
+
 if torch.version.hip:
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings_hip"
-    )
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_hip_training"
     )
 else:
-    torch.ops.load_library(
-        "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings"
-    )    
     torch.ops.load_library(
         "//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cuda_training"
     )
+
 {%- endif %}