Create separate targets for training and inference (pytorch#1757)

Summary:
Pull Request resolved: pytorch#1757

- Create separate targets for training and inference

- Redefine the old `embedding_ops` and `embedding_ops` as an empty
target with `exported_defs` pointing to the new split targets

Differential Revision: D45687293

fbshipit-source-id: 3bb7d0dd95211a2ae464528a5714e2bbe041cffe

Author: q10

fbgemm_gpu/codegen/embedding_ops_placeholder.cpp

@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+/*
+    This is placeholder code to force compilation and generation of an
+    `libdeeplearning_fbgemm_fbgemm_gpu_codegen_embedding_ops.so` file, which
+    allows downstream PyTorch code to contlinue loading the `embedding_ops“
+    and `embedding_ops_cpu` (now-)shim targets correctly.
+*/
+namespace fbgemm_gpu {}

fbgemm_gpu/codegen/split_embedding_codegen_lookup_invoker.template

@@ -10,8 +10,8 @@ from .lookup_args import *
 
 
 {% if is_fbcode %}
-torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops")
-torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu")
+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")
 torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:cumem_utils")
 torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops")
 torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py

@@ -14,11 +14,19 @@
 from math import log2
 from typing import Dict, List, NamedTuple, Optional, Tuple, Type, Union
 
+import fbgemm_gpu
 import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers
 import torch  # usort:skip
 from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType, SparseType
 from torch import nn, Tensor  # usort:skip
 
+# pyre-fixme[16]: Module `fbgemm_gpu` has no attribute `open_source`.
+open_source: bool = getattr(fbgemm_gpu, "open_source", False)
+
+if not open_source:
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops")
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/codegen:embedding_ops_cpu")
+
 DEFAULT_ASSOC = 32 if torch.version.hip is None else 64
 # Maximum number of times prefetch() can be called without
 # a corresponding forward() call