[fix] Supported hkv cpu operator base cuckoo cpu operator

tensorflow · Mar 26, 2024 · 685326f · 685326f
1 parent ca4b266
commit 685326f
Show file tree

Hide file tree

Showing 6 changed files with 1,319 additions and 198 deletions.
diff --git a/tensorflow_recommenders_addons/dynamic_embedding/core/BUILD b/tensorflow_recommenders_addons/dynamic_embedding/core/BUILD
@@ -93,9 +93,11 @@ custom_op_library(
 custom_cuda_op_library(
     name = "_hkv_ops.so",
     srcs = [
+        "kernels/cuckoo_hashtable_op.h",
+        "kernels/hkv_hashtable_op.cc",
         "ops/hkv_hashtable_ops.cc",
         "utils/utils.h",
-    ],
+    ] + glob(["kernels/lookup_impl/lookup_table_op_cpu*"]),
     copts = [
         "-Ofast",
     ],
@@ -110,4 +112,7 @@ custom_cuda_op_library(
         "kernels/cuckoo_hashtable_op_gpu.h",
         "kernels/hkv_hashtable_op_gpu.cu.cc",
     ]) + glob(["kernels/lookup_impl/lookup_table_op_hkv*"]),
+    deps = [
+        "//tensorflow_recommenders_addons/dynamic_embedding/core/lib/cuckoo:cuckoohash",
+    ],
 )