Add CAGRA support with latest RAFT #175

wphicks · 2023-11-04T00:27:05Z

This PR brings in the latest features from RAFT and significantly refactors the RAFT integration code. The primary goal of this refactor is to more clearly separate Knowhere code from RAFT integration code from RAFT itself. This leads to three layers in the updated integration:

knowhere: Code which directly creates e.g. a new IndexNode type in Knowhere is implemented in such a way as to expose no RAFT symbols or CUDA calls to Knowhere headers or other Knowhere code
raft_knowhere: This namespace is used for code responsible for translating between types, symbols, and concepts in Knowhere to types, symbols and concepts in RAFT
raft_proto: This namespace is used for features that may ultimately be upstreamed to RAFT but which are immediately useful to Knowhere.

CAGRA benchmarks have been substantially simplified in this PR and should run significantly faster. Throughput for batch size 1 is still not as high as CAGRA potentially allows, but it is significantly higher than previous benchmarks. Performance is currently bottlenecked on many small host-to-device transfers, but this can be improved in a follow-up PR. Throughput for larger batch sizes is substantially improved, with a median 17% overhead relative to raw RAFT calls during testing.

Given the significant scope of this PR, I will add some comments in-line, but here is the overall summary of changes:

Update to RAFT 23.12
Update CAGRA integration to improve performance
Avoid post-filtering using RAFT's new filtering feature Use RAFT's new device_resources_manager to simplify and optimize resource initialization
Update build infratructure to build for all supported CUDA architectures Refactor RAFT integration code to more cleanly separate RAFT code from Knowhere code
Avoid exposing RAFT symbols in any Knowhere header
Simplify CAGRA benchmarking
Allow refinement of initial results for all RAFT index types except CAGRA

NOTE: This PR currently points to a fork of RAFT while waiting for rapidsai/raft#1831 to merge. This was impacted by today's GIthub outage. Before merging, we should shift back to the main RAFT repo.

Close #176

Update to RAFT 23.12 Update CAGRA integration to improve performance Avoid post-filtering using RAFT's new filtering feature Use RAFT's new device_resources_manager to simplify and optimize resource initialization Update build infratructure to build for all supported CUDA architectures Refactor RAFT integration code to more cleanly separate RAFT code from Knowhere code Avoid exposing RAFT symbols in any Knowhere header Signed-off-by: William Hicks <[email protected]>

sre-ci-robot · 2023-11-04T00:27:10Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by: wphicks
To complete the pull request process, please assign chasingegg after the PR has been reviewed.
You can assign the PR to them by writing /assign @chasingegg in a comment when ready.

The full list of commands accepted by this bot can be found here.

Needs approval from an approver in each of these files:

OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

sre-ci-robot · 2023-11-04T00:27:15Z

Welcome @wphicks! It looks like this is your first PR to zilliztech/knowhere 🎉

mergify · 2023-11-04T00:27:47Z

@wphicks 🔍 Important: PR Classification Needed!

For efficient project management and a seamless review process, it's essential to classify your PR correctly. Here's how:

If you're fixing a bug, label it as kind/bug.
For small tweaks (less than 20 lines without altering any functionality), please use kind/improvement.
Significant changes that don't modify existing functionalities should be tagged as kind/enhancement.
Adjusting APIs or changing functionality? Go with kind/feature.

For any PR outside the kind/improvement category, ensure you link to the associated issue using the format: “issue: #”.

Thanks for your efforts and contribution to the community!.

wphicks

I've added some explanatory comments inline where I expect there to be questions about specific changes.

wphicks · 2023-11-04T00:31:04Z

CMakeLists.txt

@@ -12,19 +12,28 @@
 # License for the specific language governing permissions and limitations under
 # the License

-cmake_minimum_required(VERSION 3.23.0 FATAL_ERROR)
-project(knowhere CXX C)
+cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)


Required for RAPIDS CMake used in RAFT 23.12.

wphicks · 2023-11-04T00:31:42Z

CMakeLists.txt


 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/")
 include(GNUInstallDirs)
 include(ExternalProject)
 include(cmake/utils/utils.cmake)

+knowhere_option(WITH_RAFT "Build with RAFT indexes" OFF)


Moved this up because CMAKE_CUDA_ARCHITECTURES needs to be filled in before initializing the project.

benchmark/hdf5/benchmark_float_qps.cpp

wphicks · 2023-11-04T00:33:05Z

benchmark/hdf5/benchmark_float_qps.cpp

+    test_cagra(const knowhere::Json& cfg) {
+        auto conf = cfg;
+
+        auto find_smallest_max_iters = [&](float expected_recall) -> int32_t {


Finding the best max_iterations has higher impact than searching over itopk

benchmark/hdf5/benchmark_float_qps.cpp

wphicks · 2023-11-04T00:40:41Z

src/common/raft/integration/raft_knowhere_index.hpp

@@ -0,0 +1,125 @@
+/**


This file is the header actually included elsewhere in Knowhere. It exposes no RAFT symbols and does not require CUDA compilation.

src/common/raft/proto/ivf_to_sample_filter.cuh

wphicks · 2023-11-04T00:41:51Z

src/index/gpu_raft/gpu_raft.h

@@ -0,0 +1,292 @@
+/**


This file provides a generic template for Knowhere indexes based on RAFT.

wphicks · 2023-11-04T00:42:49Z

src/index/gpu_raft/gpu_raft_ivf_pq_config.h

-class RaftIvfFlatConfig : public IvfFlatConfig {
- public:
+struct GpuRaftIvfPqConfig : public IvfPqConfig {
+    CFG_FLOAT refine_ratio;


This newly-introduced parameter allows additional refinement after an initial selection of candidates from an index search.

wphicks · 2023-11-04T00:43:40Z

tests/ut/test_gpu_search.cc

+        return json;
+    };
+
+    auto refined_gen = [](auto&& upstream_gen) {


Helper for generating identical configurations with refinement.

Signed-off-by: William Hicks <[email protected]>

sre-ci-robot added the do-not-merge/work-in-progress label Nov 4, 2023

sre-ci-robot requested review from chasingegg and zhengbuqian November 4, 2023 00:27

sre-ci-robot added the size/XXL label Nov 4, 2023

mergify bot added the dco-passed label Nov 4, 2023

mergify bot added the do-not-merge/missing-related-issue label Nov 4, 2023

wphicks mentioned this pull request Nov 4, 2023

Provide support for RAFT CAGRA indexes #176

Closed

wphicks commented Nov 4, 2023

View reviewed changes

wphicks added 2 commits November 3, 2023 22:44

Fix todo items discovered during self-review

b9a1b34

Signed-off-by: William Hicks <[email protected]>

Update to mainline RAFT and run linters

c869f27

Signed-off-by: William Hicks <[email protected]>

wphicks marked this pull request as ready for review November 13, 2023 20:09

sre-ci-robot removed the do-not-merge/work-in-progress label Nov 13, 2023

Add workaround for refinement issue

075c4e2

Signed-off-by: William Hicks <[email protected]>

mergify bot added needs-dco and removed dco-passed labels Nov 16, 2023

wphicks force-pushed the fea-raft_refactor branch from e988f47 to 075c4e2 Compare November 16, 2023 16:04

mergify bot added dco-passed and removed needs-dco labels Nov 16, 2023

Merge branch 'main' into fea-raft_refactor

e7a9e58

Signed-off-by: William Hicks <[email protected]>

Presburger closed this Nov 24, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add CAGRA support with latest RAFT #175

Add CAGRA support with latest RAFT #175

wphicks commented Nov 4, 2023 •

edited

Loading

sre-ci-robot commented Nov 4, 2023

sre-ci-robot commented Nov 4, 2023

mergify bot commented Nov 4, 2023

wphicks left a comment •

edited

Loading

wphicks Nov 4, 2023

wphicks Nov 4, 2023

wphicks Nov 4, 2023

wphicks Nov 4, 2023

wphicks Nov 4, 2023

wphicks Nov 4, 2023

wphicks Nov 4, 2023

Add CAGRA support with latest RAFT #175

Add CAGRA support with latest RAFT #175

Conversation

wphicks commented Nov 4, 2023 • edited Loading

sre-ci-robot commented Nov 4, 2023

sre-ci-robot commented Nov 4, 2023

mergify bot commented Nov 4, 2023

wphicks left a comment • edited Loading

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks Nov 4, 2023

Choose a reason for hiding this comment

wphicks commented Nov 4, 2023 •

edited

Loading

wphicks left a comment •

edited

Loading