Open source inplace update operators (#1482)

Summary: Pull Request resolved: #1482 Reviewed By: jspark1105 Differential Revision: D41562775 fbshipit-source-id: 9bc3397c7ba8ad53d3aefa52609f84f297318a6b
pytorch · Dec 4, 2022 · 4cd267c · 4cd267c
1 parent 6120f21
commit 4cd267c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/fbgemm_gpu/codegen/split_embedding_codegen_lookup_invoker.template b/fbgemm_gpu/codegen/split_embedding_codegen_lookup_invoker.template
@@ -18,6 +18,14 @@ torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_cpu")
 torch.ops.load_library(
     "//deeplearning/fbgemm/fbgemm_gpu:split_table_batched_embeddings"
 )
+try:
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:embedding_inplace_update")
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:embedding_inplace_update_cpu")
+except OSError:
+    # Keep for BC: will be deprecated soon.
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/fb:embedding_inplace_update")
+    torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu/fb:embedding_inplace_update_cpu")
+
 {% else %}
 #import os
 #torch.ops.load_library(os.path.join(os.path.join(os.path.dirname(os.path.dirname(__file__)), "fbgemm_gpu_py.so")))

diff --git a/fbgemm_gpu/test/split_table_batched_embeddings_test.py b/fbgemm_gpu/test/split_table_batched_embeddings_test.py
@@ -4685,9 +4685,8 @@ def test_embedding_inplace_update(
         )
 
         weights_ty_list = [weights_ty] * T
-        # if open_source:
-        #     test_internal = False
-        test_internal = False
+        if open_source:
+            test_internal = False
 
         # create two embedding bag op with random weights
         locations = [location] * T