Skip to content

Commit

Permalink
Splitting fused l2 knn specializations (#461)
Browse files Browse the repository at this point in the history
Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #461
  • Loading branch information
cjnolet authored Jan 25, 2022
1 parent 98a6ccc commit c39ecd8
Show file tree
Hide file tree
Showing 8 changed files with 185 additions and 97 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ The folder structure mirrors other RAPIDS repos (cuDF, cuML, cuGraph...), with t
- `ci`: Scripts for running CI in PRs
- `conda`: conda recipes and development conda environments
- `cpp`: Source code for all C++ code. The code is currently header-only, therefore it is in the `include` folder (with no `src`).
- `cpp`: Source code for all C++ code. Headers are in the `include` folder and compiled template specializations for the shared libraries are in `src`.
- `docs`: Source code and scripts for building library documentation
- `python`: Source code for all Python source code.
Expand Down
34 changes: 19 additions & 15 deletions build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ ARGS=$*
# script, and that this script resides in the repo dir!
REPODIR=$(cd $(dirname $0); pwd)

VALIDARGS="clean libraft pyraft docs -v -g --compilelibs --allgpuarch --nvtx --show_depr_warn -h --buildgtest --buildfaiss"
VALIDARGS="clean libraft pyraft docs -v -g --compilelibs --allgpuarch --nvtx --show_depr_warn -h --nogtest --buildfaiss"
HELP="$0 [<target> ...] [<flag> ...]
where <target> is:
clean - remove all existing build artifacts and configuration (start over)
Expand All @@ -33,6 +33,7 @@ HELP="$0 [<target> ...] [<flag> ...]
--compilelibs - compile shared libraries
--allgpuarch - build for all supported GPU architectures
--buildfaiss - build faiss statically into raft
--nogtest - do not build google tests for libraft
--nvtx - Enable nvtx for profiling support
--show_depr_warn - show cmake deprecation warnings
-h - print this text
Expand All @@ -49,14 +50,15 @@ BUILD_DIRS="${CPP_RAFT_BUILD_DIR} ${PY_RAFT_BUILD_DIR} ${PYTHON_DEPS_CLONE}"
CMAKE_LOG_LEVEL=""
VERBOSE_FLAG=""
BUILD_ALL_GPU_ARCH=0
BUILD_TESTS=ON
BUILD_TESTS=YES
BUILD_STATIC_FAISS=OFF
COMPILE_LIBRARIES=OFF
ENABLE_NN_DEPENDENCIES=ON
COMPILE_LIBRARIES=${BUILD_TESTS}
ENABLE_NN_DEPENDENCIES=${BUILD_TESTS}
SINGLEGPU=""
NVTX=OFF
CLEAN=0
BUILD_DISABLE_DEPRECATION_WARNING=ON
DISABLE_DEPRECATION_WARNINGS=ON
CMAKE_TARGET=""

# Set defaults for vars that may not have been defined externally
# FIXME: if INSTALL_PREFIX is not set, check PREFIX, then check
Expand Down Expand Up @@ -97,18 +99,20 @@ if hasArg -g; then
BUILD_TYPE=Debug
fi

if hasArg --compilelibs; then
COMPILE_LIBRARIES=ON
fi

if hasArg --allgpuarch; then
BUILD_ALL_GPU_ARCH=1
fi
if hasArg --buildgtest; then
BUILD_GTEST=ON
if hasArg --nogtest; then
BUILD_TESTS=OFF
COMPILE_LIBRARIES=OFF
ENABLE_NN_DEPENDENCIES=OFF
fi
if hasArg --compilelibs; then
COMPILE_LIBRARIES=ON
ENABLE_NN_DEPENDENCIES=ON
fi
if hasArg --buildfaiss; then
BUILD_STATIC_FAISS=ON
BUILD_STATIC_FAISS=ON
fi
if hasArg --singlegpu; then
SINGLEGPU="--singlegpu"
Expand All @@ -117,7 +121,7 @@ if hasArg --nvtx; then
NVTX=ON
fi
if hasArg --show_depr_warn; then
BUILD_DISABLE_DEPRECATION_WARNING=OFF
DISABLE_DEPRECATION_WARNINGS=OFF
fi
if hasArg clean; then
CLEAN=1
Expand Down Expand Up @@ -160,14 +164,14 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs; then
-DRAFT_COMPILE_LIBRARIES=${COMPILE_LIBRARIES} \
-DRAFT_ENABLE_NN_DEPENDENCIES=${ENABLE_NN_DEPENDENCIES} \
-DNVTX=${NVTX} \
-DDISABLE_DEPRECATION_WARNING=${BUILD_DISABLE_DEPRECATION_WARNING} \
-DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
-DBUILD_TESTS=${BUILD_TESTS} \
-DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS} \
..

if (( ${NUMARGS} == 0 )) || hasArg libraft; then
# Run all c++ targets at once
cmake --build ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG}
cmake --build ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG} ${CMAKE_TARGET}
fi
fi

Expand Down
5 changes: 4 additions & 1 deletion cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,10 @@ if(RAFT_COMPILE_LIBRARIES)
add_library(raft_nn_lib SHARED
src/nn/specializations/ball_cover.cu
src/nn/specializations/detail/ball_cover_lowdim.cu
src/nn/specializations/fused_l2_knn.cu
src/nn/specializations/fused_l2_knn_long_float_true.cu
src/nn/specializations/fused_l2_knn_long_float_false.cu
src/nn/specializations/fused_l2_knn_int_float_true.cu
src/nn/specializations/fused_l2_knn_int_float_false.cu
src/nn/specializations/knn.cu
)
set_target_properties(raft_nn_lib PROPERTIES OUTPUT_NAME raft_nn)
Expand Down
80 changes: 0 additions & 80 deletions cpp/src/nn/specializations/fused_l2_knn.cu

This file was deleted.

41 changes: 41 additions & 0 deletions cpp/src/nn/specializations/fused_l2_knn_int_float_false.cu
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/*
* Copyright (c) 2022, 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.
*/

#include <cstdint>
#include <raft/spatial/knn/detail/fused_l2_knn.cuh>

namespace raft {
namespace spatial {
namespace knn {
namespace detail {

template void fusedL2Knn<int, float, false>(size_t D,
int* out_inds,
float* out_dists,
const float* index,
const float* query,
size_t n_index_rows,
size_t n_query_rows,
int k,
bool rowMajorIndex,
bool rowMajorQuery,
cudaStream_t stream,
raft::distance::DistanceType metric);

}; // namespace detail
}; // namespace knn
}; // namespace spatial
}; // namespace raft
40 changes: 40 additions & 0 deletions cpp/src/nn/specializations/fused_l2_knn_int_float_true.cu
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright (c) 2022, 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.
*/

#include <cstdint>
#include <raft/spatial/knn/detail/fused_l2_knn.cuh>

namespace raft {
namespace spatial {
namespace knn {
namespace detail {
template void fusedL2Knn<int, float, true>(size_t D,
int* out_inds,
float* out_dists,
const float* index,
const float* query,
size_t n_index_rows,
size_t n_query_rows,
int k,
bool rowMajorIndex,
bool rowMajorQuery,
cudaStream_t stream,
raft::distance::DistanceType metric);

}; // namespace detail
}; // namespace knn
}; // namespace spatial
}; // namespace raft
40 changes: 40 additions & 0 deletions cpp/src/nn/specializations/fused_l2_knn_long_float_false.cu
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright (c) 2022, 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.
*/

#include <cstdint>
#include <raft/spatial/knn/detail/fused_l2_knn.cuh>

namespace raft {
namespace spatial {
namespace knn {
namespace detail {

template void fusedL2Knn<long, float, false>(size_t D,
long* out_inds,
float* out_dists,
const float* index,
const float* query,
size_t n_index_rows,
size_t n_query_rows,
int k,
bool rowMajorIndex,
bool rowMajorQuery,
cudaStream_t stream,
raft::distance::DistanceType metric);
}; // namespace detail
}; // namespace knn
}; // namespace spatial
}; // namespace raft
40 changes: 40 additions & 0 deletions cpp/src/nn/specializations/fused_l2_knn_long_float_true.cu
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
/*
* Copyright (c) 2022, 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.
*/

#include <cstdint>
#include <raft/spatial/knn/detail/fused_l2_knn.cuh>

namespace raft {
namespace spatial {
namespace knn {
namespace detail {

template void fusedL2Knn<long, float, true>(size_t D,
long* out_inds,
float* out_dists,
const float* index,
const float* query,
size_t n_index_rows,
size_t n_query_rows,
int k,
bool rowMajorIndex,
bool rowMajorQuery,
cudaStream_t stream,
raft::distance::DistanceType metric);
}; // namespace detail
}; // namespace knn
}; // namespace spatial
}; // namespace raft

0 comments on commit c39ecd8

Please sign in to comment.