improvement: speed up kmeans on large dataset #1219
+43
−67
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
LHT129
requested review from
ShawnShawnYou,
inabao,
jiaweizone and
wxyucs
as code owners
LHT129
added
kind/improvement
|
Important Installation incomplete: to start using Gemini Code Assist, please ask the organization owner(s) to visit the Gemini Code Assist Admin Console and sign the Terms of Services. |
Reviewer's GuideThis PR optimizes object pooling by introducing thread-local caches to minimize global locking and enhances HGraph index setup in KMeans by parameterizing the degree and marking the index as immutable. Class diagram for updated ResourceObjectPool with thread-local poolclassDiagram
class ResourceObjectPool {
+TakeOne() std::shared_ptr<T>
+ReturnOne(std::shared_ptr<T>&)
+GetSize() uint64_t
-resize(uint64_t)
-get_local_pool() std::deque<std::shared_ptr<T>>&
-pool_ std::unique_ptr<Deque<std::shared_ptr<T>>>
-pool_size_ std::atomic<uint64_t>
-owned_allocator_ std::shared_ptr<Allocator>
-kLocalPoolCapacity static const uint64_t
}
class Deque
class Allocator
ResourceObjectPool --> Deque : uses
ResourceObjectPool --> Allocator : owns
Class diagram for updated KMeansCluster HGraph index setupclassDiagram
class KMeansCluster {
+find_nearest_one_with_hgraph(const float* query, ...)
-dim_ int
-allocator_ Allocator
-thread_pool_ ThreadPool
}
class InnerIndexInterface {
+FastCreateIndex(string, param)
}
class HGraphIndex {
+Build(base)
+SetImmutable()
}
class Dataset {
+Make()
}
KMeansCluster --> InnerIndexInterface : creates index
KMeansCluster --> HGraphIndex : builds and sets immutable
KMeansCluster --> Dataset : uses
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey there - I've reviewed your changes and they look great!
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location> `src/utils/resource_object_pool.h:71-75` </location>
<code_context>
- std::unique_lock<std::mutex> lock(mutex_);
- if (pool_->empty()) {
+ auto& local_pool = this->get_local_pool();
+ if (not local_pool.empty()) {
+ std::shared_ptr<T> obj = local_pool.front();
+ local_pool.pop_front();
+ return obj;
+ }
</code_context>
<issue_to_address>
**suggestion:** Consider resetting objects from the local pool before returning.
Objects from the local pool are returned without Reset(), unlike those from the shared pool. This may cause inconsistent state if Reset() is required before reuse. Please review whether Reset() should be called for local pool objects as well.
```suggestion
if (not local_pool.empty()) {
std::shared_ptr<T> obj = local_pool.front();
local_pool.pop_front();
if (obj) {
obj->Reset();
}
return obj;
}
```
</issue_to_address>
### Comment 2
<location> `src/impl/cluster/kmeans_cluster.cpp:255` </location>
<code_context>
param.allocator_ = std::make_shared<SafeAllocator>(this->allocator_);
param.thread_pool_ = this->thread_pool_;
param.metric_ = MetricType::METRIC_TYPE_L2SQR;
+ auto max_degree = std::max(32, dim_ / 8);
- auto hgraph = InnerIndexInterface::FastCreateIndex("hgraph|32|fp32", param);
</code_context>
<issue_to_address>
**question:** Check for potential integer division issues with dim_ / 8.
If dim_ is less than 8, max_degree will default to 32. Please verify this is correct for small values of dim_.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| param.allocator_ = std::make_shared<SafeAllocator>(this->allocator_); | ||
| param.thread_pool_ = this->thread_pool_; | ||
| param.metric_ = MetricType::METRIC_TYPE_L2SQR; | ||
| auto max_degree = std::max(32, dim_ / 8); |
There was a problem hiding this comment.
question: Check for potential integer division issues with dim_ / 8.
If dim_ is less than 8, max_degree will default to 32. Please verify this is correct for small values of dim_.
wxyucs
reviewed
Oct 10, 2025
src/utils/resource_object_pool.h
Outdated
|
|
||
| private: | ||
| std::shared_ptr<Allocator> owned_allocator_{nullptr}; | ||
| static const uint64_t kLocalPoolCapacity = 4; // Tunable parameter |
There was a problem hiding this comment.
style: rename kLocalPoolCapacity to local_pool_capacity_ or LOCAL_POOL_CAPACITY
src/utils/resource_object_pool.h
Outdated
| private: | ||
| inline void | ||
| resize(uint64_t size) { | ||
|
|
src/utils/resource_object_pool.h
Outdated
| return; | ||
| } | ||
| { | ||
| std::lock_guard<std::mutex> lock(mutex_); |
There was a problem hiding this comment.
change std::lock_guard to std::scoped_lock
LHT129
force-pushed
the
kmeans
branch
2 times, most recently
from
428773e to
e3992fc
Compare
Codecov Report❌ Patch coverage is @@ Coverage Diff @@
## main #1219 +/- ##
==========================================
+ Coverage 91.42% 91.52% +0.10%
==========================================
Files 318 318
Lines 17622 17621 -1
==========================================
+ Hits 16111 16128 +17
+ Misses 1511 1493 -18
Flags with carried forward coverage won't be shown. Click here to find out more.
Continue to review full report in Codecov by Sentry.
🚀 New features to boost your workflow:
|
- vt pool waste too much time on fast search in HGraph Index Signed-off-by: LHT129 <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary by Sourcery
Optimize k-means clustering performance by reducing contention in the resource pool and tuning the HGraph index parameters dynamically.
Enhancements: