Code

Integrate KNN implementation: ivf-pq#789

Code

Merged

rapids-bot[bot] merged 148 commits intorapidsai:branch-22.10rapidsai/raft:branch-22.10from achirkin:fea-knn-ivf-pqachirkin/raft:fea-knn-ivf-pqCopy head branch name to clipboardSep 30, 2022

+4,990-290

Commits

Commits on Jun 24, 2022

• Initial commit: copied everything into a single file

  achirkin

  committed
  4d7b699

  Copy full SHA for 4d7b699

Commits on Aug 4, 2022

• Merge branch 'branch-22.08' into fea-knn-ivf-pq

  achirkin

  committed
  306d4bb

  Copy full SHA for 306d4bb

• Add missing functions and fix small compile errors

  achirkin

  committed
  ebc20be

  Copy full SHA for ebc20be

Commits on Aug 5, 2022

• Added tests (still failing)

  achirkin

  committed
  4f71b09

  Copy full SHA for 4f71b09

Commits on Aug 9, 2022

• WIP replacing chunks of code with raft's helpers

  achirkin

  committed
  1d6402f

  Copy full SHA for 1d6402f

Commits on Aug 10, 2022

• Reset the expected cuda error to cudaSuccess to not crumble following cuda calls

  achirkin

  committed
  f12e5da

  Copy full SHA for f12e5da

• Replace cuann return codes with raft exceptions

  achirkin

  committed
  12fad92

  Copy full SHA for 12fad92

• replace cublas calls with raft wrappers

  achirkin

  committed
  0bb702f

  Copy full SHA for 0bb702f

• Replace _cuann_memset with utils::memzero

  achirkin

  committed
  e1cae88

  Copy full SHA for e1cae88

• Use raft logging in most of the places and wrap all cuda calls into RAFT_CUDA_TRY

  achirkin

  committed
  5b2c799

  Copy full SHA for 5b2c799

• Use helper accumulate_into_selected

  achirkin

  committed
  4f64cd3

  Copy full SHA for 4f64cd3

• Use helper copy_selected

  achirkin

  committed
  0a7324b

  Copy full SHA for 0a7324b

• Wrap remaining cuda calls into RAFT_CUDA_TRY and replace asserts with RAFT_EXPECTS

  achirkin

  committed
  4d11178

  Copy full SHA for 4d11178

Commits on Aug 11, 2022

• WIP: templatizing the api

  achirkin

  committed
  267b35e

  Copy full SHA for 267b35e

• Add more templates and use adjust_centers

  achirkin

  committed
  2541250

  Copy full SHA for 2541250

• Put index as a member of descriptor

  achirkin

  committed
  67bffa4

  Copy full SHA for 67bffa4

Commits on Aug 12, 2022

• Add the index extending operation

  achirkin

  committed
  f951c52

  Copy full SHA for f951c52

Commits on Aug 15, 2022

• Replace cuann similarity with raft distance metric

  achirkin

  committed
  5822844

  Copy full SHA for 5822844

• Use raft logging levels instead of CUANN_DEBUG macro

  achirkin

  committed
  62d32a7

  Copy full SHA for 62d32a7

• Made the API scaffold.

  achirkin

  committed
  48c8702

  Copy full SHA for 48c8702

Commits on Aug 16, 2022

• Cleanup unused functions

  achirkin

  committed
  a74ed5d

  Copy full SHA for a74ed5d

Commits on Aug 17, 2022

• Removed couple helpers and added a few WIP comments

  achirkin

  committed
  f75a5b5

  Copy full SHA for f75a5b5

• Fixing small warnings to make Release config compilable

  achirkin

  committed
  f8ea503

  Copy full SHA for f8ea503

Commits on Aug 18, 2022

• Don't store all k scores in the temporary buffer

  achirkin

  committed
  140f186

  Copy full SHA for 140f186

• Use raft topk code

  achirkin

  committed
  cb70c79

  Copy full SHA for cb70c79

• Replace _cuann_kmeans_predict_MP with raft's kmeans primitives

  achirkin

  committed
  52dbe0d

  Copy full SHA for 52dbe0d

• Remove _cuann_kmeans_predict and use the higher-level kmeans::build_clusters where possible

  achirkin

  committed
  613d551

  Copy full SHA for 613d551

• Cleanup unused functions and variables

  achirkin

  committed
  b939b8b

  Copy full SHA for b939b8b

• Cleanup unused array

  achirkin

  committed
  9a0d371

  Copy full SHA for 9a0d371

Commits on Aug 19, 2022

• WIP swithing to RMM. Cleaned up _cuann_compute_PQ_code

  achirkin

  committed
  7bea6df

  Copy full SHA for 7bea6df

• WIP swithing to RMM

  achirkin

  committed
  d7e7d49

  Copy full SHA for d7e7d49

Commits on Aug 22, 2022

• Replaced all function-local allocations with rmm::uvector/std::vector

  achirkin

  committed
  faaa495

  Copy full SHA for faaa495

• Replacing search workspace with rmm (wip almost done)

  achirkin

  committed
  e83eae0

  Copy full SHA for e83eae0

Commits on Aug 23, 2022

• Switched to device pointers in the api

  achirkin

  committed
  f8bec3e

  Copy full SHA for f8bec3e

• get rid of cuannIvfPqIndexHeader

  achirkin

  committed
  5dc2c4b

  Copy full SHA for 5dc2c4b

• Replaced more runtime branches with template instantiations

  achirkin

  committed
  c42e57c

  Copy full SHA for c42e57c

• Reset cuda error at the end of the tests for better reporting

  achirkin

  committed
  bd22b86

  Copy full SHA for bd22b86

• Reuse larger parts of kmeans clustering

  achirkin

  committed
  9907bb7

  Copy full SHA for 9907bb7

• Fix a legacy workaround to correctly calculate worst-case max number of samples

  achirkin

  committed
  3238591

  Copy full SHA for 3238591

Commits on Aug 24, 2022

• Use the new functions from rapidsai#788

  achirkin

  committed
  1bae646

  Copy full SHA for 1bae646

• Revert kmeans changes

  achirkin

  committed
  066f8c6

  Copy full SHA for 066f8c6

• Reuse the last bits of kmeans

  achirkin

  committed
  e98a214

  Copy full SHA for e98a214

Commits on Aug 25, 2022

• Rename a few variables

  achirkin

  committed
  e4da7f8

  Copy full SHA for e4da7f8

• Initialize topk output indices with zeros do workaround out-of-bounds in ivf-pq in extreme cases

  achirkin

  committed
  d478645

  Copy full SHA for d478645

• Remove element type template T from the index definition

  achirkin

  committed
  d4e645a

  Copy full SHA for d4e645a

• WIP: merging cuannIvfPqDescriptor into ivf_pq::index<IdxT>

  achirkin

  committed
  adf1dca

  Copy full SHA for adf1dca

Commits on Aug 26, 2022

• Merge branch 'branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  fdb5ad0

  Copy full SHA for fdb5ad0

• WIP: merging cuannIvfPqDescriptor into ivf_pq::index<IdxT> - removing unused members

  achirkin

  committed
  e7cb541

  Copy full SHA for e7cb541

• WIP: merging cuannIvfPqDescriptor into ivf_pq::index<IdxT> - more renamings

  achirkin

  committed
  e3b61bd

  Copy full SHA for e3b61bd

• WIP: merging cuannIvfPqDescriptor into ivf_pq::index<IdxT> - moved all data pointers

  achirkin

  committed
  cc00fb7

  Copy full SHA for cc00fb7

• WIP: merging cuannIvfPqDescriptor into ivf_pq::index<IdxT> - continue removing members

  achirkin

  committed
  55a29c1

  Copy full SHA for 55a29c1

• Removed cuannIvfPqDescriptor usage from ivfpq_search

  achirkin

  committed
  961e647

  Copy full SHA for 961e647

• Completely removed cuannIvfPqDescriptor

  achirkin

  committed
  b889394

  Copy full SHA for b889394

Commits on Aug 29, 2022

• Make ivf_pq::index properly movable

  achirkin

  committed
  81abc5d

  Copy full SHA for 81abc5d

• Remove last traces of omp for now

  achirkin

  committed
  adbffd8

  Copy full SHA for adbffd8

• Remove ivf_pq_legacy.cuh

  achirkin

  committed
  ef4e28b

  Copy full SHA for ef4e28b

• fix style

  achirkin

  committed
  adc7cbf

  Copy full SHA for adc7cbf

• Small fixes to docs]

  achirkin

  committed
  1a2f79c

  Copy full SHA for 1a2f79c

• use cudaMemcpy2DAsync in-place of transpose-3d-helper

  achirkin

  committed
  a957ec9

  Copy full SHA for a957ec9

• Add more nvtx annotations and renamings

  achirkin

  committed
  f41ef2c

  Copy full SHA for f41ef2c

• Annotate a performance bottleneck TODO

  achirkin

  committed
  9d292ae

  Copy full SHA for 9d292ae

• Collapsing wrapper code in ivf_pq::search

  achirkin

  committed
  8034aae

  Copy full SHA for 8034aae

• Use ceildiv everywhere

  achirkin

  committed
  413842c

  Copy full SHA for 413842c

Commits on Aug 30, 2022

• Use new_indices in extend

  achirkin

  committed
  f1b2b39

  Copy full SHA for f1b2b39

• Propagated IdxT into the implementation (largely untested)

  achirkin

  committed
  c22d02c

  Copy full SHA for c22d02c

Commits on Sep 1, 2022

• Add support for InnerProduct distance type

  achirkin

  committed
  d5d425b

  Copy full SHA for d5d425b

• Lift the inner-product-float-only restriction

  achirkin

  committed
  20e74f2

  Copy full SHA for 20e74f2

Commits on Sep 2, 2022

• Fix output index type

  achirkin

  committed
  d043c75

  Copy full SHA for d043c75

• Fix dependency on .cuh utils in a .hpp file

  achirkin

  committed
  f4a2852

  Copy full SHA for f4a2852

• A few small renamings in preparation for merging ivfpq_compute_similarity[*] kernels

  achirkin

  committed
  bf90837

  Copy full SHA for bf90837

• Merge branch 'branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  369d49d

  Copy full SHA for 369d49d

Commits on Sep 5, 2022

• Merge ivfpq_compute_similarity variants into a single kernel to reduce code duplication

  achirkin

  committed
  e55d043

  Copy full SHA for e55d043

Commits on Sep 6, 2022

• Merge remote-tracking branch 'rapidsai/branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  2a20b25

  Copy full SHA for 2a20b25

• fix: don't call copy on empty data

  achirkin

  committed
  0e5cea5

  Copy full SHA for 0e5cea5

• Use raft's device-wide topk functions

  achirkin

  committed
  9b6668f

  Copy full SHA for 9b6668f

• clang-tidy refactoring

  achirkin

  committed
  998c9a5

  Copy full SHA for 998c9a5

Commits on Sep 7, 2022

• Fix the rest of clang-tidy warning and document bits

  achirkin

  committed
  66453a1

  Copy full SHA for 66453a1

• Expand comments a little bit

  achirkin

  committed
  68a3664

  Copy full SHA for 68a3664

• rename lut_dtype parameter

  achirkin

  committed
  bcd4100

  Copy full SHA for bcd4100

Commits on Sep 8, 2022

• Rework tests

  achirkin

  committed
  8c88c79

  Copy full SHA for 8c88c79

• Allow reference implementation to have k > 256

  achirkin

  committed
  60d2712

  Copy full SHA for 60d2712

• Allow ivf-pq implementation to have k > 256

  achirkin

  committed
  5a3c3ae

  Copy full SHA for 5a3c3ae

Commits on Sep 9, 2022

• Fix managed memory being destroyed too soon on move assignment

  achirkin

  committed
  fe29a15

  Copy full SHA for fe29a15

• Separate the training logic (build) from adding new data to index (extend)

  achirkin

  committed
  969d31f

  Copy full SHA for 969d31f

Commits on Sep 12, 2022

• .clang-tidy fixes and other small refactorings

  achirkin

  committed
  e9bcd52

  Copy full SHA for e9bcd52

• Ported the ivfpq_encode_core bugfix from upstream

  achirkin

  committed
  0a93be9

  Copy full SHA for 0a93be9

Commits on Sep 13, 2022

• Remove unnecessary type cast

  achirkin

  committed
  add19ef

  Copy full SHA for add19ef

• Fix forgot '#pragma once'

  achirkin

  committed
  b2c3799

  Copy full SHA for b2c3799

• Fix incorrect constexpr bounds for fp16

  achirkin

  committed
  97444d8

  Copy full SHA for 97444d8

• Fixed the last failing tests (negative scores with unsigned float type)

  achirkin

  committed
  e01ccdf

  Copy full SHA for e01ccdf

Commits on Sep 14, 2022

• Simplified ivfpq_encode

  achirkin

  committed
  8d15e75

  Copy full SHA for 8d15e75

Commits on Sep 19, 2022

• Merge remote-tracking branch 'rapidsai/branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  38c244c

  Copy full SHA for 38c244c

• Rename reorganized includes

  achirkin

  committed
  db23eb8

  Copy full SHA for db23eb8

• Move a preprocessing code to GPU (up to 10x speedup in PER_SUBSPACE mode)

  achirkin

  committed
  c05953f

  Copy full SHA for c05953f

Commits on Sep 20, 2022

• Fix a bug causing recall drops and add support for signed fp8 internal rep

  achirkin

  committed
  c8e1431

  Copy full SHA for c8e1431

Commits on Sep 21, 2022

• Restore the old kernel setup code for now, as the last commit introduced a slowdown

  achirkin

  committed
  528076a

  Copy full SHA for 528076a

• Remove a template parameter PqT that can be computed using other parameters

  achirkin

  committed
  22b3631

  Copy full SHA for 22b3631

Commits on Sep 22, 2022

• Small fixes to headers

  achirkin

  committed
  af5da32

  Copy full SHA for af5da32

• Refactor kernel selecting pragmas into templates

  achirkin

  committed
  68dab8d

  Copy full SHA for 68dab8d

• Add specializations to ivfpq_compute_similarity kernel

  achirkin

  committed
  5ef0067

  Copy full SHA for 5ef0067

• Add a missing inline specifier

  achirkin

  committed
  90bac7c

  Copy full SHA for 90bac7c

• Add more specializations to ivf_pq::search and split test instantiations

  achirkin

  committed
  f2fbd65

  Copy full SHA for f2fbd65

• Add benchmarks and split them by types to enable faster parallel compilation

  achirkin

  committed
  a77d453

  Copy full SHA for a77d453

• Merge remote-tracking branch 'rapidsai/branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  1f6d353

  Copy full SHA for 1f6d353

• remove instantiations from the header

  achirkin

  committed
  213d66c

  Copy full SHA for 213d66c

• Adapt to changes in mdarray headers

  achirkin

  committed
  10bb02a

  Copy full SHA for 10bb02a

Commits on Sep 23, 2022

• Move the make_rotation_matrix to GPU

  achirkin

  committed
  5529872

  Copy full SHA for 5529872

• Move the specialization includes between headers

  achirkin

  committed
  2f8f77c

  Copy full SHA for 2f8f77c

• Merge two template parameters to reduce the raft compile times by ~4-5x at the cost of up to 15% search performance penalty

  achirkin

  committed
  443916f

  Copy full SHA for 443916f

• Moving more bits to GPU and removing managed_memory

  achirkin

  committed
  0acfdd7

  Copy full SHA for 0acfdd7

Commits on Sep 26, 2022

• Merge remote-tracking branch 'rapidsai/branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  a5a0577

  Copy full SHA for a5a0577

• Factor out calculate_offsets_and_indices

  achirkin

  committed
  ee4aeee

  Copy full SHA for ee4aeee

• Sort the cluster by their size in the index; this allows to quickly estimate the maximum number of samples during search without keeping extra arrays in the index.

  achirkin

  committed
  ad84f01

  Copy full SHA for ad84f01

• Proceeding further with removing the managed arrays

  achirkin

  committed
  f5049a8

  Copy full SHA for f5049a8

Commits on Sep 27, 2022

• Remove managed memory from the index

  achirkin

  committed
  6a4d444

  Copy full SHA for 6a4d444

• Slightly improve performance by reducing the number of arguments in device kernel/functions

  achirkin

  committed
  a02767b

  Copy full SHA for a02767b

• Allow to disable random rotation with force_random_rotation index parameter

  achirkin

  committed
  ad10141

  Copy full SHA for ad10141

Commits on Sep 28, 2022

• Use a lambda in place of the copy_fill helper.

  achirkin

  committed
  90c6b9f

  Copy full SHA for 90c6b9f

• Update cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh

  Show description for 54bbaac
  

  achirkin

  and

  tfeher

  authored
  54bbaac

  Copy full SHA for 54bbaac

• Update cpp/include/raft/spatial/knn/detail/ivf_pq_build.cuh

  Show description for 4ecc6a8
  

  achirkin

  and

  tfeher

  authored
  4ecc6a8

  Copy full SHA for 4ecc6a8

• Update cpp/include/raft/util/cudart_utils.hpp

  Show description for 9c5c463
  

  achirkin

  and

  tfeher

  authored
  9c5c463

  Copy full SHA for 9c5c463

• Update cpp/include/raft/spatial/knn/ivf_pq_types.hpp

  Show description for 4cdf0ce
  

  achirkin

  and

  tfeher

  authored
  4cdf0ce

  Copy full SHA for 4cdf0ce

• Update cpp/include/raft/spatial/knn/detail/ivf_pq_search.cuh

  Show description for 653331d
  

  achirkin

  and

  tfeher

  authored
  653331d

  Copy full SHA for 653331d

• clang-format

  achirkin

  committed
  41a8730

  Copy full SHA for 41a8730

• Add a missing knn.cuh header for the brute-force knn

  achirkin

  committed
  02fde4b

  Copy full SHA for 02fde4b

• Remove unused cpu code and an extra level of indirection

  achirkin

  committed
  4704684

  Copy full SHA for 4704684

• Update cpp/include/raft/spatial/knn/ivf_pq.cuh

  Show description for 887cd7d
  

  achirkin

  and

  tfeher

  authored
  887cd7d

  Copy full SHA for 887cd7d

• Various small fixes

  achirkin

  committed
  5fe21da

  Copy full SHA for 5fe21da

• More documentation

  achirkin

  committed
  2df832f

  Copy full SHA for 2df832f

• Various small fixes

  achirkin

  committed
  2323431

  Copy full SHA for 2323431

• More documentation

  achirkin

  committed
  e661e5c

  Copy full SHA for e661e5c

• Remove unrelated changes

  achirkin

  committed
  de6db27

  Copy full SHA for de6db27

• Rename pq_width->pq_book_size and refactor PER_CLUSTER training

  achirkin

  committed
  1ae1471

  Copy full SHA for 1ae1471

• Factor out select_residuals

  achirkin

  committed
  854c804

  Copy full SHA for 854c804

Commits on Sep 29, 2022

• Factor out select_clusters function

  achirkin

  committed
  d1b006b

  Copy full SHA for d1b006b

• Add a helper to select a worker function instance

  achirkin

  committed
  1784f8b

  Copy full SHA for 1784f8b

• Update docs

  achirkin

  committed
  8f514b6

  Copy full SHA for 8f514b6

• Use cub in preprocessing kernels and refactor the variable names

  achirkin

  committed
  119619b

  Copy full SHA for 119619b

• refactor ivfpq_make_outputs into postprocess_distances and postprocess_neighbors

  achirkin

  committed
  e8aeb71

  Copy full SHA for e8aeb71

Commits on Sep 30, 2022

• Added some documentation to the tests and more debug-logging to ensure the test coverage

  achirkin

  committed
  96b48ad

  Copy full SHA for 96b48ad

• Format the logging strings a bit

  achirkin

  committed
  89e8327

  Copy full SHA for 89e8327

• Add __launch_bounds__ everywhere to avoid occasional 'too many resources requested for launch'

  achirkin

  committed
  729b0e4

  Copy full SHA for 729b0e4

• Factor out train_per_subset and train_per_cluster for readability

  achirkin

  committed
  ea94db3

  Copy full SHA for ea94db3

• Factor out get_max_batch_size

  achirkin

  committed
  0f839fe

  Copy full SHA for 0f839fe

• Factor out kernel selection code

  achirkin

  committed
  9b64e55

  Copy full SHA for 9b64e55

• Merge remote-tracking branch 'rapidsai/branch-22.10' into fea-knn-ivf-pq

  achirkin

  committed
  bf73cf5

  Copy full SHA for bf73cf5

• Expand test coverage a little bit

  achirkin

  committed
  6389eb7

  Copy full SHA for 6389eb7

• Fix a cudaErrorIllegalAddress that can happen if the probed clusters have not enough data to return

  achirkin

  committed
  cf1a196

  Copy full SHA for cf1a196