Skip to content

Commit

Permalink
Add algo parameter for CAGRA ANN bench (#1687)
Browse files Browse the repository at this point in the history 
To achieve high QPS with small batch size, it is recommended to use the `multi_cta` algorithm variant. CAGRA has an input parameter to specify which algorithm to use, but until now the benchmarks only used the default algorithm.

This PR adds configuration option for changing the algorithm parameter for the CAGRA ANN benchmark wrappers. This will be useful while benchmarking various configurations with CAGRA.

Authors:
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1687
  • Loading branch information
@tfeher
tfeher authored Jul 27, 2023
1 parent d8b9187 commit bcb7404
Show file tree
Hide file tree
Showing 2 changed files with 328 additions and 21 deletions.
335 changes: 314 additions & 21 deletions cpp/bench/ann/conf/deep-100M.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -222,31 +222,324 @@
"name" : "cagra.dim32",
"algo" : "raft_cagra",
"build_param": {
"index_dim" : 32
"index_dim": 32,
"intermediate_graph_degree": 48
},
"file" : "index/deep-100M/cagra/dim32",
"file": "index/deep-100M/cagra/dim32",
"search_params": [
{
"itopk": 32,
"search_width": 1,
"max_iterations": 0,
"algo": "single_cta"
},
{
"itopk": 32,
"search_width": 1,
"max_iterations": 32,
"algo": "single_cta"
},
{
"itopk": 64,
"search_width": 4,
"max_iterations": 16,
"algo": "single_cta"
},
{
"itopk": 64,
"search_width": 1,
"max_iterations": 64,
"algo": "single_cta"
},
{
"itopk": 96,
"search_width": 2,
"max_iterations": 48,
"algo": "single_cta"
},
{
"itopk": 128,
"search_width": 8,
"max_iterations": 16,
"algo": "single_cta"
},
{
"itopk": 128,
"search_width": 2,
"max_iterations": 64,
"algo": "single_cta"
},
{
"itopk": 192,
"search_width": 8,
"max_iterations": 24,
"algo": "single_cta"
},
{
"itopk": 192,
"search_width": 2,
"max_iterations": 96,
"algo": "single_cta"
},
{
"itopk": 256,
"search_width": 8,
"max_iterations": 32,
"algo": "single_cta"
},
{
"itopk": 384,
"search_width": 8,
"max_iterations": 48,
"algo": "single_cta"
},
{
"itopk": 512,
"search_width": 8,
"max_iterations": 64,
"algo": "single_cta"
}
],
"search_result_file": "result/deep-100M/cagra/dim32"
},
{
"name": "cagra.dim32.multi_cta",
"algo": "raft_cagra",
"build_param": {
"index_dim": 32,
"intermediate_graph_degree": 48
},
"file": "index/deep-100M/cagra/dim32",
"search_params": [
{
"itopk": 32,
"search_width": 1,
"max_iterations": 0,
"algo": "multi_cta"
},
{
"itopk": 32,
"search_width": 1,
"max_iterations": 32,
"algo": "multi_cta"
},
{
"itopk": 64,
"search_width": 4,
"max_iterations": 16,
"algo": "multi_cta"
},
{
"itopk": 64,
"search_width": 1,
"max_iterations": 64,
"algo": "multi_cta"
},
{
"itopk": 96,
"search_width": 2,
"max_iterations": 48,
"algo": "multi_cta"
},
{
"itopk": 128,
"search_width": 8,
"max_iterations": 16,
"algo": "multi_cta"
},
{
"itopk": 128,
"search_width": 2,
"max_iterations": 64,
"algo": "multi_cta"
},
{
"itopk": 192,
"search_width": 8,
"max_iterations": 24,
"algo": "multi_cta"
},
{
"itopk": 192,
"search_width": 2,
"max_iterations": 96,
"algo": "multi_cta"
},
{
"itopk": 256,
"search_width": 8,
"max_iterations": 32,
"algo": "multi_cta"
},
{
"itopk": 384,
"search_width": 8,
"max_iterations": 48,
"algo": "multi_cta"
},
{
"itopk": 512,
"search_width": 8,
"max_iterations": 64,
"algo": "multi_cta"
}
],
"search_result_file": "result/deep-100M/cagra/dim32_multi_cta"
},
{
"name": "cagra.dim32.multi_kernel",
"algo": "raft_cagra",
"build_param": {
"index_dim": 32,
"intermediate_graph_degree": 48
},
"file": "index/deep-100M/cagra/dim32",
"search_params": [
{
"itopk": 32,
"search_width": 1,
"max_iterations": 0,
"algo": "multi_kernel"
},
{
"itopk": 32,
"search_width": 1,
"max_iterations": 32,
"algo": "multi_kernel"
},
{
"itopk": 64,
"search_width": 4,
"max_iterations": 16,
"algo": "multi_kernel"
},
{
"itopk": 64,
"search_width": 1,
"max_iterations": 64,
"algo": "multi_kernel"
},
{
"itopk": 96,
"search_width": 2,
"max_iterations": 48,
"algo": "multi_kernel"
},
{
"itopk": 128,
"search_width": 8,
"max_iterations": 16,
"algo": "multi_kernel"
},
{
"itopk": 128,
"search_width": 2,
"max_iterations": 64,
"algo": "multi_kernel"
},
{
"itopk": 192,
"search_width": 8,
"max_iterations": 24,
"algo": "multi_kernel"
},
{
"itopk": 192,
"search_width": 2,
"max_iterations": 96,
"algo": "multi_kernel"
},
{
"itopk": 256,
"search_width": 8,
"max_iterations": 32,
"algo": "multi_kernel"
},
{
"itopk": 384,
"search_width": 8,
"max_iterations": 48,
"algo": "multi_kernel"
},
{
"itopk": 512,
"search_width": 8,
"max_iterations": 64,
"algo": "multi_kernel"
}
],
"search_result_file": "result/deep-100M/cagra/dim32_multi_kernel"
},
{
"name": "cagra.dim64",
"algo": "raft_cagra",
"build_param": {
"index_dim": 64
},
"file": "index/deep-100M/cagra/dim64",
"search_params" : [
{"itopk": 32},
{"itopk": 64},
{"itopk": 128}
{
"itopk": 32,
"search_width": 1,
"max_iterations": 0
},
{
"itopk": 32,
"search_width": 1,
"max_iterations": 32
},
{
"itopk": 64,
"search_width": 4,
"max_iterations": 16
},
{
"itopk": 64,
"search_width": 1,
"max_iterations": 64
},
{
"itopk": 96,
"search_width": 2,
"max_iterations": 48
},
{
"itopk": 128,
"search_width": 8,
"max_iterations": 16
},
{
"itopk": 128,
"search_width": 2,
"max_iterations": 64
},
{
"itopk": 192,
"search_width": 8,
"max_iterations": 24
},
{
"itopk": 192,
"search_width": 2,
"max_iterations": 96
},
{
"itopk": 256,
"search_width": 8,
"max_iterations": 32
},
{
"itopk": 384,
"search_width": 8,
"max_iterations": 48
},
{
"itopk": 512,
"search_width": 8,
"max_iterations": 64
}
],
"search_result_file" : "result/deep-100M/cagra/dim32"
},


{
"name" : "cagra.dim64",
"algo" : "raft_cagra",
"build_param": {
"index_dim" : 64
},
"file" : "index/deep-100M/cagra/dim64",
"search_params" : [
{"itopk": 32},
{"itopk": 64},
{"itopk": 128}
],
"search_result_file" : "result/deep-100M/cagra/dim64"
}
]
}
14 changes: 14 additions & 0 deletions cpp/bench/ann/src/raft/raft_benchmark.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,20 @@ void parse_search_param(const nlohmann::json& conf,
if (conf.contains("itopk")) { param.p.itopk_size = conf.at("itopk"); }
if (conf.contains("search_width")) { param.p.search_width = conf.at("search_width"); }
if (conf.contains("max_iterations")) { param.p.max_iterations = conf.at("max_iterations"); }
if (conf.contains("algo")) {
if (conf.at("algo") == "single_cta") {
param.p.algo = raft::neighbors::experimental::cagra::search_algo::SINGLE_CTA;
} else if (conf.at("algo") == "multi_cta") {
param.p.algo = raft::neighbors::experimental::cagra::search_algo::MULTI_CTA;
} else if (conf.at("algo") == "multi_kernel") {
param.p.algo = raft::neighbors::experimental::cagra::search_algo::MULTI_KERNEL;
} else if (conf.at("algo") == "auto") {
param.p.algo = raft::neighbors::experimental::cagra::search_algo::AUTO;
} else {
std::string tmp = conf.at("algo");
THROW("Invalid value for algo: %s", tmp.c_str());
}
}
}
#endif

Expand Down

0 comments on commit bcb7404

Please sign in to comment.