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

Support more hkv strategy #389

Merged
merged 1 commit into from
May 6, 2024
Merged

Support more hkv strategy #389

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
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 = "79c59b19c03b771cdcb6deb3c6a3213353482f4d07cb1ddb53c4b001a0f58b29",
strip_prefix = "HierarchicalKV-0.1.0-beta.10",
url = "https://github.com/NVIDIA-Merlin/HierarchicalKV/archive/refs/tags/v0.1.0-beta.10.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
13 changes: 7 additions & 6 deletions tensorflow_recommenders_addons/dynamic_embedding/core/kernels/hkv_hashtable_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -664,7 +664,8 @@ class HashTableInsertOp : public HashTableOpKernel {
core::ScopedUnref unref_me(table);

DataTypeVector expected_inputs = {expected_input_0_, table->key_dtype(),
table->value_dtype()};
table->value_dtype(),
DataTypeToEnum<int64>::v()};
OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {}));

const Tensor& keys = ctx->input(1);
Expand Down Expand Up @@ -750,9 +751,9 @@ class HashTableAccumOp : public HashTableOpKernel {
hkv_table::HkvHashTableOfTensors<K, V>* table_cuckoo =
(hkv_table::HkvHashTableOfTensors<K, V>*)table;

DataTypeVector expected_inputs = {expected_input_0_, table->key_dtype(),
table->value_dtype(),
DataTypeToEnum<bool>::v()};
DataTypeVector expected_inputs = {
expected_input_0_, table->key_dtype(), table->value_dtype(),
DataTypeToEnum<bool>::v(), DataTypeToEnum<int64>::v()};
OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {}));

const Tensor& keys = ctx->input(1);
Expand Down Expand Up @@ -975,8 +976,8 @@ REGISTER_KERNEL_BUILDER(
HashTableAccumOp<key_dtype, value_dtype>); \
REGISTER_KERNEL_BUILDER(Name(PREFIX_OP_NAME(HkvHashTableFindWithExists)) \
.Device(DEVICE_CPU) \
.TypeConstraint<key_dtype>("Tin") \
.TypeConstraint<value_dtype>("Tout"), \
.TypeConstraint<key_dtype>("key_dtype") \
.TypeConstraint<value_dtype>("value_dtype"), \
HashTableFindWithExistsOp<key_dtype, value_dtype>); \
REGISTER_KERNEL_BUILDER( \
Name(PREFIX_OP_NAME(HkvHashTableSaveToFileSystem)) \
Expand Down