Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[feat] Add HKV which is based on merlin HierarchicalHV #356

Closed
wants to merge 6 commits into from
Closed

[feat] Add HKV which is based on merlin HierarchicalHV #356

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
8925972
[feat] Hkv code draft
Lifann Jun 16, 2023
1187213
[fix] Replace sparse_fill_empty_rows_op with the tf native interface
LinGeLin Nov 2, 2023
68bf529
[feat] Hkv code improve, fix bugs in the draft version, and enhance d…
LinGeLin Aug 11, 2023
b50e931
[feat] Update hkv beta.8 to support evict strategy
LinGeLin Dec 14, 2023
5d67021
[feat] support insert with scores by setting gen_scoures_fn
LinGeLin Dec 27, 2023
d8608cb
CFR
LinGeLin Dec 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/make_wheel_Linux_x86.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ fi

# if tensorflow version >= 2.6.0 and <= 2.11.9
if [[ "$TF_VERSION" =~ ^2\.([6-9]|10|11)\.[0-9]$ ]] ; then
export BUILD_IMAGE="tfra/nosla-cuda11.2.1-cudnn8-ubuntu20.04-manylinux2014-python$PY_VERSION"
export BUILD_IMAGE="tfra/nosla-cuda11.2.2-cudnn8-ubuntu20.04-manylinux2014-python$PY_VERSION"
export TF_CUDA_VERSION="11.2"
export TF_CUDNN_VERSION="8.1"
elif [ $TF_VERSION == "2.4.1" ] ; then
Expand Down
52 changes: 3 additions & 49 deletions WORKSPACE
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,55 +57,9 @@ http_archive(
http_archive(
name = "hkv",
build_file = "//build_deps/toolchains/hkv:hkv.BUILD",
# TODO(LinGeLin) remove this when update hkv
patch_cmds = [
"""sed -i.bak '1772i\\'$'\\n ThrustAllocator<uint8_t> thrust_allocator_;\\n' include/merlin_hashtable.cuh""",
"""sed -i.bak '225i\\'$'\\n thrust_allocator_.set_allocator(allocator_);\\n' include/merlin_hashtable.cuh""",
"sed -i.bak 's/thrust::sort_by_key(thrust_par.on(stream)/thrust::sort_by_key(thrust_par(thrust_allocator_).on(stream)/' include/merlin_hashtable.cuh",
"sed -i.bak 's/reduce(thrust_par.on(stream)/reduce(thrust_par(thrust_allocator_).on(stream)/' include/merlin_hashtable.cuh",
"""sed -i.bak '125i\\'$'\\n template <typename T>\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '126i\\'$'\\n struct ThrustAllocator : thrust::device_malloc_allocator<T> {\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '127i\\'$'\\n public:\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '128i\\'$'\\n typedef thrust::device_malloc_allocator<T> super_t;\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '129i\\'$'\\n typedef typename super_t::pointer pointer;\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '130i\\'$'\\n typedef typename super_t::size_type size_type;\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '131i\\'$'\\n public:\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '132i\\'$'\\n pointer allocate(size_type n) {\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '133i\\'$'\\n void* ptr = nullptr;\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '134i\\'$'\\n MERLIN_CHECK(\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '135i\\'$'\\n allocator_ != nullptr,\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '136i\\'$'\\n "[ThrustAllocator] set_allocator should be called in advance!");\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '137i\\'$'\\n allocator_->alloc(MemoryType::Device, &ptr, sizeof(T) * n);\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '138i\\'$'\\n return pointer(reinterpret_cast<T*>(ptr));\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '139i\\'$'\\n }\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '140i\\'$'\\n void deallocate(pointer p, size_type n) {\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '141i\\'$'\\n MERLIN_CHECK(\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '142i\\'$'\\n allocator_ != nullptr,\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '143i\\'$'\\n "[ThrustAllocator] set_allocator should be called in advance!");\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '144i\\'$'\\n allocator_->free(MemoryType::Device, reinterpret_cast<void*>(p.get()));\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '145i\\'$'\\n }\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '146i\\'$'\\n void set_allocator(BaseAllocator* allocator) { allocator_ = allocator; }\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '147i\\'$'\\n public:\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '148i\\'$'\\n BaseAllocator* allocator_ = nullptr;\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '149i\\'$'\\n };\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '20i\\'$'\\n #include <thrust/device_malloc_allocator.h>\\n' include/merlin/allocator.cuh""",
"""sed -i.bak '367i\\'$'\\n for (auto addr : (*table)->buckets_address) {\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '368i\\'$'\\n allocator->free(MemoryType::Device, addr);\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '369i\\'$'\\n }\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '370i\\'$'\\n /*\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '382i\\'$'\\n */\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '224i\\'$'\\n uint8_t* address = nullptr;\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '225i\\'$'\\n allocator->alloc(MemoryType::Device, (void**)&(address), bucket_memory_size * (end - start));\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '226i\\'$'\\n (*table)->buckets_address.push_back(address);\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '228i\\'$'\\n allocate_bucket_others<K, V, S><<<1, 1>>>((*table)->buckets, i, address + (bucket_memory_size * (i-start)), reserve_size, bucket_max_size);\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '229i\\'$'\\n /*\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '235i\\'$'\\n */\\n' include/merlin/core_kernels.cuh""",
"""sed -i.bak '22i\\'$'\\n#include <vector>\\n' include/merlin/types.cuh""",
"""sed -i.bak '143i\\'$'\\n std::vector<uint8_t*> buckets_address;\\n' include/merlin/types.cuh""",
],
sha256 = "f8179c445a06a558262946cda4d8ae7252d313e73f792586be9b1bc0c993b1cf",
strip_prefix = "HierarchicalKV-0.1.0-beta.6",
url = "https://github.com/NVIDIA-Merlin/HierarchicalKV/archive/refs/tags/v0.1.0-beta.6.tar.gz",
sha256 = "3839f91b703b401fd6d2449c034662b6f8d6563e5b9b71b4c25b217cf1cd63fd",
strip_prefix = "HierarchicalKV-0.1.0-beta.8",
url = "https://github.com/NVIDIA-Merlin/HierarchicalKV/archive/refs/tags/v0.1.0-beta.8.tar.gz",
)

tf_configure(
Expand Down
3 changes: 2 additions & 1 deletion tensorflow_recommenders_addons/dynamic_embedding/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
'CuckooHashTable',
'CuckooHashTableConfig',
'CuckooHashTableCreator',
'HkvEvictStrategy',
'HkvHashTable',
'HkvHashTableConfig',
'HkvHashTableCreator',
Expand Down Expand Up @@ -55,7 +56,7 @@
from tensorflow_recommenders_addons.dynamic_embedding.python.ops import data_flow_ops as data_flow
from tensorflow_recommenders_addons.dynamic_embedding.python.ops.dynamic_embedding_creator import (
KVCreator, CuckooHashTableConfig, CuckooHashTableCreator,
HkvHashTableConfig, HkvHashTableCreator, RedisTableConfig,
HkvHashTableConfig, HkvHashTableCreator, HkvEvictStrategy, RedisTableConfig,
RedisTableCreator, FileSystemSaver)
from tensorflow_recommenders_addons.dynamic_embedding.python.ops.cuckoo_hashtable_ops import (
CuckooHashTable,)
Expand Down
Loading
Loading