Skip to content

Commit

Permalink
Reduce the generated code volume
Browse files Browse the repository at this point in the history
  • Loading branch information
@achirkin
achirkin committed Sep 3, 2024
1 parent 7599331 commit 4d9241e
Show file tree
Hide file tree
Showing 79 changed files with 794 additions and 4,347 deletions.
1,894 changes: 220 additions & 1,674 deletions cpp/src/neighbors/detail/cagra/compute_distance-ext.cuh
View file

Large diffs are not rendered by default.

604 changes: 74 additions & 530 deletions cpp/src/neighbors/detail/cagra/compute_distance.cu
View file

Large diffs are not rendered by default.

30 changes: 8 additions & 22 deletions cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
namespace cuvs::neighbors::cagra::detail {{
using namespace cuvs::distance;
{content}
}} // namespace cuvs::neighbors::cagra::detail
Expand All @@ -69,7 +70,7 @@
half_uint64=("half", "uint64_t", "float"),
)

metric_prefix = 'cuvs::distance::DistanceType::'
metric_prefix = 'DistanceType::'

specs = []
descs = []
Expand All @@ -90,17 +91,10 @@
# CAGRA
for metric in ['L2Expanded', 'InnerProduct']:
path = f"compute_distance_standard_{metric}_{type_path}_dim{mxdim}_t{team}.cu"
includes = '#include "compute_distance_standard.cuh"'
includes = '#include "compute_distance_standard-impl.cuh"'
params = f"{metric_prefix}{metric}, {team}, {mxdim}, {data_t}, {idx_t}, {distance_t}"
spec = f"standard_descriptor_spec<{params}>"
desc = f"standard_dataset_descriptor_t<{params}>"
content = f"""
template struct {desc};
template <>
const void* {spec}::init_kernel = reinterpret_cast<const void*>(&standard_dataset_descriptor_init_kernel<{params}>);
template struct {spec};
"""
descs.append(desc)
content = f"""template struct {spec};"""
specs.append(spec)
with open(path, "w") as f:
f.write(template.format(includes=includes, content=content))
Expand All @@ -112,17 +106,10 @@
for pq_bit in pq_bits:
for metric in ['L2Expanded']:
path = f"compute_distance_vpq_{metric}_{type_path}_dim{mxdim}_t{team}_{pq_bit}pq_{pq_len}subd_{code_book_t}.cu"
includes = '#include "compute_distance_vpq.cuh"'
includes = '#include "compute_distance_vpq-impl.cuh"'
params = f"{metric_prefix}{metric}, {team}, {mxdim}, {pq_bit}, {pq_len}, {code_book_t}, {data_t}, {idx_t}, {distance_t}"
spec = f"vpq_descriptor_spec<{params}>"
desc = f"cagra_q_dataset_descriptor_t<{params}>"
content = f"""
template struct {desc};
template <>
const void* {spec}::init_kernel = reinterpret_cast<const void*>(&vpq_dataset_descriptor_init_kernel<{params}>);
template struct {spec};
"""
descs.append(desc)
content = f"""template struct {spec};"""
specs.append(spec)
with open(path, "w") as f:
f.write(template.format(includes=includes, content=content))
Expand All @@ -132,12 +119,11 @@
includes = '''
#pragma once
#include "compute_distance_standard.cuh"
#include "compute_distance_vpq.cuh"
#include "compute_distance_standard.hpp"
#include "compute_distance_vpq.hpp"
'''
newline = "\n"
contents = f'''
{newline.join(map(lambda s: "extern template struct " + s + ";", descs))}
{newline.join(map(lambda s: "extern template struct " + s + ";", specs))}
extern template struct
Expand Down
75 changes: 24 additions & 51 deletions ...etail/cagra/compute_distance_standard.cuh → .../cagra/compute_distance_standard-impl.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
*/
#pragma once

#include "compute_distance.hpp"
#include "compute_distance_standard.hpp"

#include <cuvs/distance/distance.hpp>
#include <cuvs/neighbors/common.hpp>
Expand Down Expand Up @@ -262,57 +262,30 @@ template <cuvs::distance::DistanceType Metric,
typename DataT,
typename IndexT,
typename DistanceT>
struct standard_descriptor_spec : public instance_spec<DataT, IndexT, DistanceT> {
using base_type = instance_spec<DataT, IndexT, DistanceT>;
using typename base_type::data_type;
using typename base_type::distance_type;
using typename base_type::host_type;
using typename base_type::index_type;

template <typename DatasetT>
constexpr static inline bool accepts_dataset()
{
return is_strided_dataset_v<DatasetT>;
}

using descriptor_type =
dataset_descriptor_host<DataT, IndexT, DistanceT>
standard_descriptor_spec<Metric, TeamSize, DatasetBlockDim, DataT, IndexT, DistanceT>::init_(
const cagra::search_params& params,
const DataT* ptr,
IndexT size,
uint32_t dim,
uint32_t ld,
rmm::cuda_stream_view stream)
{
using desc_type =
standard_dataset_descriptor_t<Metric, TeamSize, DatasetBlockDim, DataT, IndexT, DistanceT>;
static const void* init_kernel;

template <typename DatasetT>
static auto init(const cagra::search_params& params,
const DatasetT& dataset,
cuvs::distance::DistanceType metric,
rmm::cuda_stream_view stream) -> host_type
{
descriptor_type dd_host{nullptr,
nullptr,
dataset.view().data_handle(),
IndexT(dataset.n_rows()),
dataset.dim(),
dataset.stride()};
host_type result{dd_host, stream, DatasetBlockDim};
void* args[] = // NOLINT
{&result.dev_ptr,
&descriptor_type::ptr(dd_host.args),
&dd_host.size,
&dd_host.args.dim,
&descriptor_type::ld(dd_host.args)};
RAFT_CUDA_TRY(cudaLaunchKernel(init_kernel, 1, 1, args, 0, stream));
return result;
}
using base_type = typename desc_type::base_type;
desc_type dd_host{nullptr, nullptr, ptr, size, dim, ld};
host_type result{dd_host, stream, DatasetBlockDim};

template <typename DatasetT>
static auto priority(const cagra::search_params& params,
const DatasetT& dataset,
cuvs::distance::DistanceType metric) -> double
{
// If explicit team_size is specified and doesn't match the instance, discard it
if (params.team_size != 0 && TeamSize != params.team_size) { return -1.0; }
if (Metric != metric) { return -1.0; }
// Otherwise, favor the closest dataset dimensionality.
return 1.0 / (0.1 + std::abs(double(dataset.dim()) - double(DatasetBlockDim)));
}
};
standard_dataset_descriptor_init_kernel<Metric,
TeamSize,
DatasetBlockDim,
DataT,
IndexT,
DistanceT>
<<<1, 1, 0, stream>>>(result.dev_ptr, ptr, size, dim, desc_type::ld(dd_host.args));
RAFT_CUDA_TRY(cudaPeekAtLastError());
return result;
}

} // namespace cuvs::neighbors::cagra::detail
80 changes: 80 additions & 0 deletions cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/*
* Copyright (c) 2024, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once

#include "compute_distance.hpp"

#include <cuvs/distance/distance.hpp>

#include <type_traits>

namespace cuvs::neighbors::cagra::detail {

template <cuvs::distance::DistanceType Metric,
uint32_t TeamSize,
uint32_t DatasetBlockDim,
typename DataT,
typename IndexT,
typename DistanceT>
struct standard_descriptor_spec : public instance_spec<DataT, IndexT, DistanceT> {
using base_type = instance_spec<DataT, IndexT, DistanceT>;
using typename base_type::data_type;
using typename base_type::distance_type;
using typename base_type::host_type;
using typename base_type::index_type;

template <typename DatasetT>
constexpr static inline bool accepts_dataset()
{
return is_strided_dataset_v<DatasetT>;
}

template <typename DatasetT>
static auto init(const cagra::search_params& params,
const DatasetT& dataset,
cuvs::distance::DistanceType metric,
rmm::cuda_stream_view stream) -> host_type
{
return init_(params,
dataset.view().data_handle(),
IndexT(dataset.n_rows()),
dataset.dim(),
dataset.stride(),
stream);
}

template <typename DatasetT>
static auto priority(const cagra::search_params& params,
const DatasetT& dataset,
cuvs::distance::DistanceType metric) -> double
{
// If explicit team_size is specified and doesn't match the instance, discard it
if (params.team_size != 0 && TeamSize != params.team_size) { return -1.0; }
if (Metric != metric) { return -1.0; }
// Otherwise, favor the closest dataset dimensionality.
return 1.0 / (0.1 + std::abs(double(dataset.dim()) - double(DatasetBlockDim)));
}

private:
static dataset_descriptor_host<DataT, IndexT, DistanceT> init_(const cagra::search_params& params,
const DataT* ptr,
IndexT size,
uint32_t dim,
uint32_t ld,
rmm::cuda_stream_view stream);
};

} // namespace cuvs::neighbors::cagra::detail
25 changes: 3 additions & 22 deletions ...c/neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim128_t8.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,31 +23,12 @@
*
*/

#include "compute_distance_standard.cuh"
#include "compute_distance_standard-impl.cuh"

namespace cuvs::neighbors::cagra::detail {

template struct standard_dataset_descriptor_t<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint32_t,
float>;
template <>
const void* standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint32_t,
float>::init_kernel =
reinterpret_cast<const void*>(
&standard_dataset_descriptor_init_kernel<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint32_t,
float>);
template struct standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
using namespace cuvs::distance;
template struct standard_descriptor_spec<DistanceType::InnerProduct,
8,
128,
float,
Expand Down
25 changes: 3 additions & 22 deletions .../neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim256_t16.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,31 +23,12 @@
*
*/

#include "compute_distance_standard.cuh"
#include "compute_distance_standard-impl.cuh"

namespace cuvs::neighbors::cagra::detail {

template struct standard_dataset_descriptor_t<cuvs::distance::DistanceType::InnerProduct,
16,
256,
float,
uint32_t,
float>;
template <>
const void* standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
16,
256,
float,
uint32_t,
float>::init_kernel =
reinterpret_cast<const void*>(
&standard_dataset_descriptor_init_kernel<cuvs::distance::DistanceType::InnerProduct,
16,
256,
float,
uint32_t,
float>);
template struct standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
using namespace cuvs::distance;
template struct standard_descriptor_spec<DistanceType::InnerProduct,
16,
256,
float,
Expand Down
25 changes: 3 additions & 22 deletions .../neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim512_t32.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,31 +23,12 @@
*
*/

#include "compute_distance_standard.cuh"
#include "compute_distance_standard-impl.cuh"

namespace cuvs::neighbors::cagra::detail {

template struct standard_dataset_descriptor_t<cuvs::distance::DistanceType::InnerProduct,
32,
512,
float,
uint32_t,
float>;
template <>
const void* standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
32,
512,
float,
uint32_t,
float>::init_kernel =
reinterpret_cast<const void*>(
&standard_dataset_descriptor_init_kernel<cuvs::distance::DistanceType::InnerProduct,
32,
512,
float,
uint32_t,
float>);
template struct standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
using namespace cuvs::distance;
template struct standard_descriptor_spec<DistanceType::InnerProduct,
32,
512,
float,
Expand Down
25 changes: 3 additions & 22 deletions ...c/neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint64_dim128_t8.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,31 +23,12 @@
*
*/

#include "compute_distance_standard.cuh"
#include "compute_distance_standard-impl.cuh"

namespace cuvs::neighbors::cagra::detail {

template struct standard_dataset_descriptor_t<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint64_t,
float>;
template <>
const void* standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint64_t,
float>::init_kernel =
reinterpret_cast<const void*>(
&standard_dataset_descriptor_init_kernel<cuvs::distance::DistanceType::InnerProduct,
8,
128,
float,
uint64_t,
float>);
template struct standard_descriptor_spec<cuvs::distance::DistanceType::InnerProduct,
using namespace cuvs::distance;
template struct standard_descriptor_spec<DistanceType::InnerProduct,
8,
128,
float,
Expand Down
Loading

0 comments on commit 4d9241e

Please sign in to comment.