JIT LTO Cagra Search #1807
There are no files selected for viewing
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
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,74 @@ | ||
| /* | ||
| * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <cuvs/detail/jit_lto/common_fragments.hpp> | ||
|
|
||
| namespace cuvs::neighbors::cagra::detail { | ||
|
|
||
| template <typename DataTag, | ||
| typename IndexTag, | ||
| typename DistanceTag, | ||
| typename QueryTag, | ||
| typename CodebookTag, | ||
| typename TeamTag, | ||
| typename BlockDimTag, | ||
| typename PqBitsTag, | ||
| typename PqLenTag> | ||
| struct fragment_tag_setup_workspace {}; | ||
|
|
||
| template <typename DataTag, | ||
| typename IndexTag, | ||
| typename DistanceTag, | ||
| typename QueryTag, | ||
| typename CodebookTag, | ||
| typename TeamTag, | ||
| typename BlockDimTag, | ||
| typename PqBitsTag, | ||
| typename PqLenTag> | ||
| struct fragment_tag_compute_distance {}; | ||
|
|
||
| template <typename QueryTag, typename DistanceTag, typename MetricTag> | ||
| struct fragment_tag_dist_op {}; | ||
|
|
||
| template <typename DataTag, | ||
| typename IndexTag, | ||
| typename DistanceTag, | ||
| typename QueryTag, | ||
| typename TeamTag, | ||
| typename BlockDimTag, | ||
| typename NormTag> | ||
| struct fragment_tag_apply_normalization_standard {}; | ||
|
|
||
| template <typename DataTag, | ||
| typename SourceIndexTag, | ||
| typename IndexTag, | ||
| typename DistanceTag, | ||
| bool TopkByBitonicSort, | ||
| bool BitonicSortAndMergeMultiWarps> | ||
| struct fragment_tag_search_single_cta {}; | ||
|
|
||
| template <typename DataTag, | ||
| typename SourceIndexTag, | ||
| typename IndexTag, | ||
| typename DistanceTag, | ||
| bool TopkByBitonicSort, | ||
| bool BitonicSortAndMergeMultiWarps> | ||
| struct fragment_tag_search_single_cta_p {}; | ||
|
|
||
| template <typename DataTag, typename SourceIndexTag, typename IndexTag, typename DistanceTag> | ||
| struct fragment_tag_search_multi_cta {}; | ||
|
|
||
| template <typename DataTag, typename IndexTag, typename DistanceTag> | ||
| struct fragment_tag_random_pickup {}; | ||
|
|
||
| template <typename DataTag, typename IndexTag, typename DistanceTag, typename SourceIndexTag> | ||
| struct fragment_tag_compute_distance_to_child_nodes {}; | ||
|
|
||
| template <typename IndexTag, typename DistanceTag, typename SourceIndexTag> | ||
| struct fragment_tag_apply_filter_kernel {}; | ||
|
|
||
| } // namespace cuvs::neighbors::cagra::detail |
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,98 @@ | ||
| /* | ||
|
Member
There was a problem hiding this comment. Please restore |
||
| * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| #pragma once | ||
|
|
||
| #include <cstdint> | ||
|
|
||
| namespace cuvs::detail::jit_lto { | ||
|
|
||
| struct tag_f {}; | ||
| struct tag_h {}; | ||
| struct tag_sc {}; | ||
| struct tag_uc {}; | ||
| struct tag_idx_l {}; | ||
| struct tag_filter_none {}; | ||
| struct tag_filter_bitset {}; | ||
|
|
||
| } // namespace cuvs::detail::jit_lto | ||
|
|
||
| namespace cuvs::neighbors::cagra::detail { | ||
|
|
||
| using cuvs::detail::jit_lto::tag_f; | ||
| using cuvs::detail::jit_lto::tag_filter_bitset; | ||
| using cuvs::detail::jit_lto::tag_filter_none; | ||
| using cuvs::detail::jit_lto::tag_h; | ||
| using cuvs::detail::jit_lto::tag_idx_l; | ||
| using cuvs::detail::jit_lto::tag_sc; | ||
| using cuvs::detail::jit_lto::tag_uc; | ||
|
|
||
| struct tag_idx_ui {}; | ||
| struct tag_dist_f {}; | ||
| struct tag_metric_l2 {}; | ||
| struct tag_metric_inner_product {}; | ||
| struct tag_metric_cosine {}; | ||
| struct tag_metric_hamming {}; | ||
| struct tag_team_8 {}; | ||
| struct tag_team_16 {}; | ||
| struct tag_team_32 {}; | ||
| struct tag_dim_128 {}; | ||
| struct tag_dim_256 {}; | ||
| struct tag_dim_512 {}; | ||
| struct tag_pq_bits_0 {}; | ||
| struct tag_pq_bits_8 {}; | ||
| struct tag_pq_len_0 {}; | ||
| struct tag_pq_len_2 {}; | ||
| struct tag_pq_len_4 {}; | ||
| struct tag_codebook_none {}; | ||
| struct tag_codebook_half {}; | ||
| struct tag_metric_l1 {}; | ||
| struct tag_norm_noop {}; | ||
| struct tag_norm_cosine {}; | ||
|
|
||
| } // namespace cuvs::neighbors::cagra::detail | ||
|
|
||
| namespace cuvs::neighbors::ivf_flat::detail { | ||
|
|
||
| using cuvs::detail::jit_lto::tag_f; | ||
| using cuvs::detail::jit_lto::tag_filter_bitset; | ||
| using cuvs::detail::jit_lto::tag_filter_none; | ||
| using cuvs::detail::jit_lto::tag_h; | ||
| using cuvs::detail::jit_lto::tag_idx_l; | ||
| using cuvs::detail::jit_lto::tag_sc; | ||
| using cuvs::detail::jit_lto::tag_uc; | ||
|
|
||
| struct tag_i8 {}; | ||
| struct tag_u8 {}; | ||
|
|
||
| struct tag_acc_f {}; | ||
| struct tag_acc_h {}; | ||
| struct tag_acc_i32 {}; | ||
| struct tag_acc_u32 {}; | ||
|
|
||
| struct tag_metric_euclidean {}; | ||
| struct tag_metric_inner_product {}; | ||
| struct tag_metric_custom_udf {}; | ||
|
|
||
| struct tag_post_process_identity {}; | ||
| struct tag_post_process_sqrt {}; | ||
| struct tag_post_process_compose {}; | ||
|
|
||
| template <typename DataTag, typename AccTag, typename IdxTag, int Capacity, bool Ascending> | ||
| struct fragment_tag_interleaved_scan {}; | ||
|
|
||
| template <typename DataTag, typename AccTag, bool ComputeNorm, int Veclen> | ||
| struct fragment_tag_load_and_compute_dist {}; | ||
|
|
||
| template <typename DataTag, typename AccTag, typename MetricTag, int Veclen> | ||
| struct fragment_tag_metric {}; | ||
|
|
||
| template <typename IndexTag, typename FilterTag> | ||
| struct fragment_tag_filter {}; | ||
|
|
||
| template <typename PostLambdaTag> | ||
| struct fragment_tag_post_lambda {}; | ||
|
|
||
| } // namespace cuvs::neighbors::ivf_flat::detail | ||
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
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
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
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
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please follow the existing conventions that we use in other matrix files. We use the
*_abbrevsubstitutions for the kernel name and the tag name. Please don't pollute core CMake code with these implementation details.