diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 9c42cda720..caf46f5d6a 100755 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -34,9 +34,9 @@ Here are some guidelines to help the review process go smoothly. features or make changes out of the scope of those requested by the reviewer (doing this just add delays as already reviewed code ends up having to be re-reviewed/it is hard to tell what is new etc!). Further, please do not - rebase your branch on master/force push/rewrite history, doing any of these + rebase your branch/force push/rewrite history, doing any of these causes the context of any comments made by reviewers to be lost. If - conflicts occur against master they should be resolved by merging master + conflicts occur they should be resolved by merging the target branch into the branch used for making the pull request. Many thanks in advance for your cooperation! diff --git a/cpp/include/raft/sparse/matrix/detail/preprocessing.cuh b/cpp/include/raft/sparse/matrix/detail/preprocessing.cuh new file mode 100644 index 0000000000..26703d692c --- /dev/null +++ b/cpp/include/raft/sparse/matrix/detail/preprocessing.cuh @@ -0,0 +1,524 @@ +/* + * 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. + */ +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace raft::sparse::matrix::detail { + +/** + * @brief Calculates the BM25 values for a target matrix. + * @param num_feats: The total number of features in the matrix + * @param avg_feat_len: The avg length of all features combined. + * @param k_param: K value required by BM25 algorithm. + * @param b_param: B value required by BM25 algorithm. + */ +template +struct bm25 { + bm25(T1 num_feats, T2 avg_feat_len, T2 k_param, T2 b_param) + { + total_feats = num_feats; + avg_feat_length = avg_feat_len; + k = k_param; + b = b_param; + } + + float __device__ operator()(const T2& value, const T2& num_feats_id_occ, const T2& feat_length) + { + T2 tf = T2(value / feat_length); + T2 idf = raft::log(total_feats / num_feats_id_occ); + T2 bm = ((k + 1) * tf) / (k * ((1.0f - b) + b * (feat_length / avg_feat_length)) + tf); + + return idf * bm; + } + T2 avg_feat_length; + T1 total_feats; + T2 k; + T2 b; +}; + +/** + * @brief Calculates the tfidf values for a target matrix. Term frequency is calculate using + * logrithmically scaled frequency. + * @param total_feats_param: The total number of features in the matrix + */ +template +struct tfidf { + tfidf(T1 total_feats_param) { total_feats = total_feats_param; } + + float __device__ operator()(const T2& value, const T2& num_feats_id_occ, const T2& feat_length) + { + T2 tf = T2(value / feat_length); + T2 idf = raft::log(total_feats / num_feats_id_occ); + return tf * idf; + } + T1 total_feats; +}; + +template +struct mapper { + mapper(raft::device_vector_view map) : map(map) {} + + float __device__ operator()(const T& value) + { + T new_value = map[value]; + if (new_value) { + return new_value; + } else { + return 0.0; + } + } + + raft::device_vector_view map; +}; + +template +struct map_to { + map_to(raft::device_vector_view map) : map(map) {} + + float __device__ operator()(const T1& key, const T2& count) + { + map[key] = count; + return 0.0f; + } + + raft::device_vector_view map; +}; + +template +struct map_add { + map_add(raft::device_vector_view map) : map(map) {} + + float __device__ operator()(const T1& key, const T2& count) + { + map[key] = map[key] + count; + return 0.0f; + } + + raft::device_vector_view map; +}; + +template +struct map_inc { + map_inc(raft::device_vector_view map) : map(map) {} + + float __device__ operator()(const T1& key) + { + T1 value = map[key]; + map[key] = value + 1; + return 0.0f; + } + + raft::device_vector_view map; +}; + +/** + * @brief Get unique counts + * @param handle: raft resource handle + * @param sort_vector: Input COO array that contains the keys. + * @param secondary_vector: Input with secondary keys of COO, (columns or rows). + * @param data: Input COO values array. + * @param itr_vals: Input array used to calculate counts. + * @param keys_out: Output array with one entry for each key. (same size as counts_out) + * @param counts_out: Output array with cumulative sum for each key. (same size as keys_out) + */ +template +void get_uniques_counts(raft::resources& handle, + raft::device_vector_view sort_vector, + raft::device_vector_view secondary_vector, + raft::device_vector_view data, + raft::device_vector_view itr_vals, + raft::device_vector_view keys_out, + raft::device_vector_view counts_out) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + raft::sparse::op::coo_sort(int(sort_vector.size()), + int(secondary_vector.size()), + int(data.size()), + sort_vector.data_handle(), + secondary_vector.data_handle(), + data.data_handle(), + stream); + + thrust::reduce_by_key(raft::resource::get_thrust_policy(handle), + sort_vector.data_handle(), + sort_vector.data_handle() + sort_vector.size(), + itr_vals.data_handle(), + keys_out.data_handle(), + counts_out.data_handle()); +} + +/** + * @brief Compute cumulative sum for each unique value in the origin array + * @param handle: raft resource handle + * @param origin: Input array that has values to use for computation + * @param keys: Output array that has keys, should be the size of unique + * @param counts: Output array that contains the computed counts + * @param results: Output array that scatters the counts to origin value positions. Same size as + * origin array. + */ +template +void create_mapped_vector(raft::resources& handle, + const raft::device_vector_view origin, + const raft::device_vector_view keys, + const raft::device_vector_view counts, + raft::device_vector_view result, + T1 key_size) +{ + // index into the last element and then add 1 to it. + auto origin_map = raft::make_device_vector(handle, key_size + 1); + raft::matrix::fill(handle, origin_map.view(), 0.0f); + + auto dummy_vec = raft::make_device_vector(handle, keys.size()); + raft::linalg::map(handle, + dummy_vec.view(), + map_to(origin_map.view()), + raft::make_const_mdspan(keys), + raft::make_const_mdspan(counts)); + + raft::linalg::map(handle, result, raft::cast_op{}, raft::make_const_mdspan(origin)); + raft::linalg::map(handle, result, mapper(origin_map.view()), raft::make_const_mdspan(result)); +} + +/** + * @brief Compute row(id) counts + * @param handle: raft resource handle + * @param rows: Input COO rows array + * @param columns: Input COO columns array + * @param values: Input COO values array + * @param id_counts: Output array that stores counts per row, scattered to same shape as rows. + * @param n_rows: Number of rows in matrix + */ +template +void get_id_counts(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view id_counts, + T1 n_rows) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + raft::sparse::op::coo_sort(int(rows.size()), + int(columns.size()), + int(values.size()), + rows.data_handle(), + columns.data_handle(), + values.data_handle(), + stream); + + // auto row_keys = raft::make_device_vector(handle, uniq_rows); + auto rows_counts = raft::make_device_vector(handle, n_rows); + raft::matrix::fill(handle, rows_counts.view(), 0); + + raft::sparse::linalg::coo_degree(raft::make_const_mdspan(rows).data_handle(), + int(rows.size()), + rows_counts.data_handle(), + stream); + + raft::linalg::map( + handle, id_counts, mapper(rows_counts.view()), raft::make_const_mdspan(rows)); +} + +/** + * @brief Gather per feature mean values, returns the cumulative avg feature length. + * @param handle: raft resource handle + * @param rows: Input COO rows array + * @param columns: Input COO columns array + * @param values: Input COO values array + * @param feat_lengths: Output array that stores mean per feature value + * @param n_cols: Number of columns in matrix + */ +template +float get_feature_data(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view feat_lengths, + T1 n_cols) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto preserved_columns = raft::make_device_vector(handle, columns.size()); + raft::copy(preserved_columns.data_handle(), columns.data_handle(), columns.size(), stream); + int uniq_cols = + raft::sparse::neighbors::get_n_components(columns.data_handle(), columns.size(), stream); + auto col_keys = raft::make_device_vector(handle, uniq_cols); + auto col_counts = raft::make_device_vector(handle, uniq_cols); + + get_uniques_counts(handle, columns, rows, values, values, col_keys.view(), col_counts.view()); + + auto total_feature_lengths = raft::make_device_scalar(handle, 0); + + raft::linalg::mapReduce(total_feature_lengths.data_handle(), + col_counts.size(), + 0, + raft::identity_op(), + raft::add_op(), + stream, + col_counts.data_handle()); + auto total_feature_lengths_host = raft::make_host_scalar(handle, 0); + raft::copy(total_feature_lengths_host.data_handle(), + total_feature_lengths.data_handle(), + total_feature_lengths.size(), + stream); + T2 avg_feat_length = T2(total_feature_lengths_host(0)) / n_cols; + create_mapped_vector( + handle, preserved_columns.view(), col_keys.view(), col_counts.view(), feat_lengths, n_cols); + return avg_feat_length; +} + +/** + * @brief Gather per feature mean values and id counts, returns the cumulative avg feature length. + * @param handle: raft resource handle + * @param rows: Input COO rows array + * @param columns: Input COO columns array + * @param values: Input COO values array + * @param feat_lengths: Output array that stores mean per feature value + * @param id_counts: Output array that stores id(row) counts for nz values + * @param n_rows: Number of rows in matrix + * @param n_cols: Number of columns in matrix + */ +template +float sparse_search_preprocess(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view feat_lengths, + raft::device_vector_view id_counts, + T1 n_rows, + T1 n_cols) +{ + auto avg_feature_len = get_feature_data(handle, rows, columns, values, feat_lengths, n_cols); + + get_id_counts(handle, rows, columns, values, id_counts, n_rows); + + return avg_feature_len; +} + +/** + * @brief Use TFIDF algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param rows: Input COO rows array + * @param columns: Input COO columns array + * @param values: Input COO values array + * @param values_out: Output COO values array + * @param n_rows: Number of rows in matrix + * @param n_cols: Number of columns in matrix + */ +template +void base_encode_tfidf(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view values_out, + T1 n_rows, + T1 n_cols) +{ + auto feat_lengths = raft::make_device_vector(handle, values.size()); + auto id_counts = raft::make_device_vector(handle, values.size()); + auto col_counts = raft::make_device_vector(handle, n_cols); + auto avg_feat_length = sparse_search_preprocess( + handle, rows, columns, values, feat_lengths.view(), id_counts.view(), n_rows, n_cols); + + raft::linalg::map(handle, + values_out, + tfidf(n_cols), + raft::make_const_mdspan(values), + raft::make_const_mdspan(id_counts.view()), + raft::make_const_mdspan(feat_lengths.view())); +} + +/** + * @brief Use TFIDF algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param coo_in: Input COO matrix + * @param values_out: Output COO values array + */ +template +void encode_tfidf(raft::resources& handle, + raft::device_coo_matrix_view coo_in, + raft::device_vector_view values_out) +{ + auto rows = raft::make_device_vector_view(coo_in.structure_view().get_rows().data(), + coo_in.structure_view().get_rows().size()); + auto columns = raft::make_device_vector_view(coo_in.structure_view().get_cols().data(), + coo_in.structure_view().get_cols().size()); + auto values = raft::make_device_vector_view(coo_in.get_elements().data(), + coo_in.get_elements().size()); + + base_encode_tfidf(handle, + rows, + columns, + values, + values_out, + coo_in.structure_view().get_n_rows(), + coo_in.structure_view().get_n_cols()); +} + +/** + * @brief Use TFIDF algorithm to encode features in CSR sparse matrix + * @param handle: raft resource handle + * @param csr_in: Input CSR matrix + * @param values_out: Output values array + */ +template +void encode_tfidf(raft::resources& handle, + raft::device_csr_matrix_view csr_in, + raft::device_vector_view values_out) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + auto indptr = raft::make_device_vector_view( + csr_in.structure_view().get_indptr().data(), csr_in.structure_view().get_indptr().size()); + auto indices = raft::make_device_vector_view( + csr_in.structure_view().get_indices().data(), csr_in.structure_view().get_indices().size()); + auto values = raft::make_device_vector_view(csr_in.get_elements().data(), + csr_in.get_elements().size()); + + auto rows = raft::make_device_vector(handle, values.size()); + raft::sparse::convert::csr_to_coo(indptr.data_handle(), + csr_in.structure_view().get_n_rows(), + rows.data_handle(), + rows.size(), + stream); + + base_encode_tfidf(handle, + rows.view(), + indices, + values, + values_out, + csr_in.structure_view().get_n_rows(), + csr_in.structure_view().get_n_cols()); +} + +/** + * @brief Use BM25 algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param rows: Input COO rows array + * @param columns: Input COO columns array + * @param values: Input COO values array + * @param values_out: Output COO values array + * @param n_rows: Number of rows in matrix + * @param n_cols: Number of columns in matrix + * @param k_param: K value to use for BM25 algorithm + * @param b_param: B value to use for BM25 algorithm + */ +template +void base_encode_bm25(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view values_out, + T1 n_rows, + T1 n_cols, + T2 k_param = 1.6f, + T2 b_param = 0.75f) +{ + auto feat_lengths = raft::make_device_vector(handle, values.size()); + auto id_counts = raft::make_device_vector(handle, values.size()); + auto col_counts = raft::make_device_vector(handle, n_cols); + + auto avg_feat_length = sparse_search_preprocess( + handle, rows, columns, values, feat_lengths.view(), id_counts.view(), n_rows, n_cols); + + raft::linalg::map(handle, + values_out, + bm25(n_cols, avg_feat_length, k_param, b_param), + raft::make_const_mdspan(values), + raft::make_const_mdspan(id_counts.view()), + raft::make_const_mdspan(feat_lengths.view())); +} + +/** + * @brief Use BM25 algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param coo_in: Input COO matrix + * @param values_out: Output values array + * @param k_param: K value to use for BM25 algorithm + * @param b_param: B value to use for BM25 algorithm + */ +template +void encode_bm25(raft::resources& handle, + raft::device_coo_matrix_view coo_in, + raft::device_vector_view values_out, + T2 k_param = 1.6f, + T2 b_param = 0.75f) +{ + auto rows = raft::make_device_vector_view(coo_in.structure_view().get_rows().data(), + coo_in.structure_view().get_rows().size()); + auto columns = raft::make_device_vector_view(coo_in.structure_view().get_cols().data(), + coo_in.structure_view().get_cols().size()); + auto values = raft::make_device_vector_view(coo_in.get_elements().data(), + coo_in.get_elements().size()); + + base_encode_bm25(handle, + rows, + columns, + values, + values_out, + coo_in.structure_view().get_n_rows(), + coo_in.structure_view().get_n_cols()); +} + +/** + * @brief Use BM25 algorithm to encode features in CSR sparse matrix + * @param handle: raft resource handle + * @param csr_in: Input CSR matrix + * @param values_out: Output values array + * @param k_param: K value to use for BM25 algorithm + * @param b_param: B value to use for BM25 algorithm + */ +template +void encode_bm25(raft::resources& handle, + raft::device_csr_matrix_view csr_in, + raft::device_vector_view values_out, + T2 k_param = 1.6f, + T2 b_param = 0.75f) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + auto indptr = raft::make_device_vector_view( + csr_in.structure_view().get_indptr().data(), csr_in.structure_view().get_indptr().size()); + auto indices = raft::make_device_vector_view( + csr_in.structure_view().get_indices().data(), csr_in.structure_view().get_indices().size()); + auto values = raft::make_device_vector_view(csr_in.get_elements().data(), + csr_in.get_elements().size()); + + auto rows = raft::make_device_vector(handle, values.size()); + + raft::sparse::convert::csr_to_coo(indptr.data_handle(), + csr_in.structure_view().get_n_rows(), + rows.data_handle(), + rows.size(), + stream); + + base_encode_bm25(handle, + rows.view(), + indices, + values, + values_out, + csr_in.structure_view().get_n_rows(), + csr_in.structure_view().get_n_cols()); +} + +} // namespace raft::sparse::matrix::detail \ No newline at end of file diff --git a/cpp/include/raft/sparse/matrix/preprocessing.cuh b/cpp/include/raft/sparse/matrix/preprocessing.cuh new file mode 100644 index 0000000000..e4b3edd64b --- /dev/null +++ b/cpp/include/raft/sparse/matrix/preprocessing.cuh @@ -0,0 +1,93 @@ +/* + * 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 +#include +#include +#include +#include + +#include + +namespace raft::sparse::matrix { + +/** + * @brief Use BM25 algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param coo_in: Input COO matrix + * @param values_out: Output values array + * @param k_param: K value to use for BM25 algorithm + * @param b_param: B value to use for BM25 algorithm + */ +template +void encode_bm25(raft::resources& handle, + raft::device_coo_matrix_view coo_in, + raft::device_vector_view values_out, + float k_param = 1.6f, + float b_param = 0.75) +{ + return matrix::detail::encode_bm25(handle, coo_in, values_out, k_param, b_param); +} + +/** + * @brief Use BM25 algorithm to encode features in CSR sparse matrix + * @param handle: raft resource handle + * @param csr_in: Input CSR matrix + * @param values_out: Output values array + * @param k_param: K value to use for BM25 algorithm + * @param b_param: B value to use for BM25 algorithm + */ +template +void encode_bm25(raft::resources& handle, + raft::device_csr_matrix_view csr_in, + raft::device_vector_view values_out, + float k_param = 1.6f, + float b_param = 0.75) +{ + return matrix::detail::encode_bm25(handle, csr_in, values_out, k_param, b_param); +} + +/** + * @brief Use TFIDF algorithm to encode features in COO sparse matrix + * @param handle: raft resource handle + * @param coo_in: Input COO matrix + * @param values_out: Output COO values array + */ +template +void encode_tfidf(raft::resources& handle, + raft::device_coo_matrix_view coo_in, + raft::device_vector_view values_out) +{ + return matrix::detail::encode_tfidf(handle, coo_in, values_out); +} + +/** + * @brief Use TFIDF algorithm to encode features in CSR sparse matrix + * @param handle: raft resource handle + * @param csr_in: Input CSR matrix + * @param values_out: Output values array + */ +template +void encode_tfidf(raft::resources& handle, + raft::device_csr_matrix_view csr_in, + raft::device_vector_view values_out) +{ + return matrix::detail::encode_tfidf(handle, csr_in, values_out); +} + +} // namespace raft::sparse::matrix diff --git a/cpp/include/raft/sparse/neighbors/knn.cuh b/cpp/include/raft/sparse/neighbors/knn.cuh index 2cf68818aa..bffbf6c943 100644 --- a/cpp/include/raft/sparse/neighbors/knn.cuh +++ b/cpp/include/raft/sparse/neighbors/knn.cuh @@ -30,8 +30,11 @@ " Please use the sparse/spatial version instead.") #endif +#include #include +#include #include +#include namespace raft::sparse::neighbors { @@ -103,4 +106,171 @@ void brute_force_knn(const value_idx* idxIndptr, metricArg); } +/** + * Search the sparse kNN for the k-nearest neighbors of a set of sparse query vectors + * using some distance implementation + * @param[in] csr_idx index csr matrix + * @param[in] csr_query query csr matrix + * @param[out] output_indices dense matrix for output indices (size n_query_rows * k) + * @param[out] output_dists dense matrix for output distances (size n_query_rows * k) + * @param[in] k the number of neighbors to query + * @param[in] handle CUDA resource::get_cuda_stream(handle) to order operations with respect to + * @param[in] batch_size_index maximum number of rows to use from index matrix per batch + * @param[in] batch_size_query maximum number of rows to use from query matrix per batch + * @param[in] metric distance metric/measure to use + * @param[in] metricArg potential argument for metric (currently unused) + */ +template +void brute_force_knn(raft::device_csr_matrix csr_idx, + raft::device_csr_matrix csr_query, + device_vector_view output_indices, + device_vector_view output_dists, + int k, + raft::resources const& handle, + size_t batch_size_index = 2 << 14, // approx 1M + size_t batch_size_query = 2 << 14, + raft::distance::DistanceType metric = raft::distance::DistanceType::L2Expanded, + float metricArg = 0) +{ + auto idxIndptr = csr_idx.structure_view().get_indptr(); + auto idxIndices = csr_idx.structure_view().get_indices(); + auto idxData = csr_idx.view().get_elements(); + + auto queryIndptr = csr_query.structure_view().get_indptr(); + auto queryIndices = csr_query.structure_view().get_indices(); + auto queryData = csr_query.view().get_elements(); + + brute_force::knn(idxIndptr.data(), + idxIndices.data(), + idxData.data(), + idxIndices.size(), + idxIndptr.size() - 1, + csr_idx.structure_view().get_n_cols(), + queryIndptr.data(), + queryIndices.data(), + queryData.data(), + queryIndices.size(), + queryIndptr.size() - 1, + csr_query.structure_view().get_n_cols(), + output_indices.data_handle(), + output_dists.data_handle(), + k, + handle, + batch_size_index, + batch_size_query, + metric, + metricArg); +} + +/** + * Search the sparse kNN for the k-nearest neighbors of a set of sparse query vectors + * using some distance implementation + * @param[in] coo_idx index coo matrix + * @param[in] coo_query query coo matrix + * @param[out] output_indices dense matrix for output indices (size n_query_rows * k) + * @param[out] output_dists dense matrix for output distances (size n_query_rows * k) + * @param[in] k the number of neighbors to query + * @param[in] handle CUDA resource::get_cuda_stream(handle) to order operations with respect to + * @param[in] batch_size_index maximum number of rows to use from index matrix per batch + * @param[in] batch_size_query maximum number of rows to use from query matrix per batch + * @param[in] metric distance metric/measure to use + * @param[in] metricArg potential argument for metric (currently unused) + */ +template +void brute_force_knn(raft::device_coo_matrix coo_idx, + raft::device_coo_matrix coo_query, + device_vector_view output_indices, + device_vector_view output_dists, + int k, + raft::resources const& handle, + size_t batch_size_index = 2 << 14, // approx 1M + size_t batch_size_query = 2 << 14, + raft::distance::DistanceType metric = raft::distance::DistanceType::L2Expanded, + float metricArg = 0) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + auto idxRows = coo_idx.structure_view().get_rows(); + auto idxCols = coo_idx.structure_view().get_cols(); + auto idxData = coo_idx.view().get_elements(); + + auto queryRows = coo_query.structure_view().get_rows(); + auto queryCols = coo_query.structure_view().get_cols(); + auto queryData = coo_query.view().get_elements(); + + raft::sparse::op::coo_sort(int(idxRows.size()), + int(idxCols.size()), + int(idxData.size()), + idxRows.data(), + idxCols.data(), + idxRows.data(), + stream); + + raft::sparse::op::coo_sort(int(queryRows.size()), + int(queryCols.size()), + int(queryData.size()), + queryRows.data(), + queryCols.data(), + queryData.data(), + stream); + // + 1 is to account for the 0 at the beginning of the csr representation + auto idxRowsCsr = raft::make_device_vector( + handle, coo_query.structure_view().get_n_rows() + 1); + auto queryRowsCsr = raft::make_device_vector( + handle, coo_query.structure_view().get_n_rows() + 1); + + raft::sparse::convert::sorted_coo_to_csr(idxRows.data(), + int(idxRows.size()), + idxRowsCsr.data_handle(), + coo_idx.structure_view().get_n_rows() + 1, + stream); + + raft::sparse::convert::sorted_coo_to_csr(queryRows.data(), + int(queryRows.size()), + queryRowsCsr.data_handle(), + coo_query.structure_view().get_n_rows() + 1, + stream); + + brute_force::knn(idxRowsCsr.data_handle(), + idxCols.data(), + idxData.data(), + idxCols.size(), + idxRowsCsr.size() - 1, + coo_idx.structure_view().get_n_cols(), + queryRowsCsr.data_handle(), + queryCols.data(), + queryData.data(), + queryCols.size(), + queryRowsCsr.size() - 1, + coo_query.structure_view().get_n_cols(), + output_indices.data_handle(), + output_dists.data_handle(), + k, + handle, + batch_size_index, + batch_size_query, + metric, + metricArg); +} + }; // namespace raft::sparse::neighbors diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 5d504d2100..fe1128622f 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -319,6 +319,8 @@ if(BUILD_TESTS) sparse/spgemmi.cu sparse/spmm.cu sparse/symmetrize.cu + sparse/preprocess_csr.cu + sparse/preprocess_coo.cu ) ConfigureTest( diff --git a/cpp/test/preprocess_utils.cu b/cpp/test/preprocess_utils.cu new file mode 100644 index 0000000000..5734128373 --- /dev/null +++ b/cpp/test/preprocess_utils.cu @@ -0,0 +1,283 @@ +/* + * 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. + */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace raft::util { + +template +struct check_zeroes { + float __device__ operator()(const T1& value, const T2& idx) + { + if (value == 0) { + return 0.f; + } else { + return 1.f; + } + } +}; + +template +void preproc_kernel(raft::resources& handle, + raft::host_vector_view h_rows, + raft::host_vector_view h_cols, + raft::host_vector_view h_elems, + raft::device_vector_view results, + int num_rows, + int num_cols, + bool tf_idf) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + int rows_size = h_rows.size(); + int cols_size = h_cols.size(); + int elements_size = h_elems.size(); + auto device_matrix = raft::make_device_matrix(handle, num_rows, num_cols); + raft::matrix::fill(handle, device_matrix.view(), 0.0f); + auto host_matrix = raft::make_host_matrix(handle, num_rows, num_cols); + raft::copy(host_matrix.data_handle(), device_matrix.data_handle(), device_matrix.size(), stream); + + for (int i = 0; i < elements_size; i++) { + int row = h_rows(i); + int col = h_cols(i); + float element = h_elems(i); + host_matrix(row, col) = element; + } + + raft::copy(device_matrix.data_handle(), host_matrix.data_handle(), host_matrix.size(), stream); + auto output_cols_lengths = raft::make_device_matrix(handle, 1, num_cols); + raft::linalg::reduce(output_cols_lengths.data_handle(), + device_matrix.data_handle(), + num_rows, + num_cols, + 0.0f, + false, + true, + stream); + auto h_output_cols_lengths = raft::make_host_matrix(handle, 1, num_cols); + raft::copy(h_output_cols_lengths.data_handle(), + output_cols_lengths.data_handle(), + output_cols_lengths.size(), + stream); + + auto output_cols_length_sum = raft::make_device_scalar(handle, 0); + raft::linalg::mapReduce(output_cols_length_sum.data_handle(), + num_cols, + 0, + raft::identity_op(), + raft::add_op(), + stream, + output_cols_lengths.data_handle()); + auto h_output_cols_length_sum = raft::make_host_scalar(handle, 0); + raft::copy(h_output_cols_length_sum.data_handle(), + output_cols_length_sum.data_handle(), + output_cols_length_sum.size(), + stream); + float avg_col_length = float(h_output_cols_length_sum(0)) / num_cols; + + auto output_rows_freq = raft::make_device_matrix(handle, 1, num_rows); + raft::linalg::reduce(output_rows_freq.data_handle(), + device_matrix.data_handle(), + num_rows, + num_cols, + 0.0f, + false, + false, + stream); + + auto output_rows_cnt = raft::make_device_matrix(handle, 1, num_rows); + raft::linalg::reduce(output_rows_cnt.data_handle(), + device_matrix.data_handle(), + num_rows, + num_cols, + 0.0f, + false, + false, + stream, + false, + check_zeroes()); + auto h_output_rows_cnt = raft::make_host_matrix(handle, 1, num_rows); + raft::copy( + h_output_rows_cnt.data_handle(), output_rows_cnt.data_handle(), output_rows_cnt.size(), stream); + + auto out_device_matrix = raft::make_device_matrix(handle, num_rows, num_cols); + raft::matrix::fill(handle, out_device_matrix.view(), 0.0f); + auto out_host_matrix = raft::make_host_matrix(handle, num_rows, num_cols); + auto out_host_vector = raft::make_host_vector(handle, results.size()); + + float k1 = 1.6f; + float b = 0.75f; + int count = 0; + float result; + for (int row = 0; row < num_rows; row++) { + for (int col = 0; col < num_cols; col++) { + float val = host_matrix(row, col); + if (val == 0) { + out_host_matrix(row, col) = 0.0f; + } else { + float tf = float(val / h_output_cols_lengths(0, col)); + float idf = raft::log(num_cols / h_output_rows_cnt(0, row)); + if (tf_idf) { + result = tf * idf; + } else { + float bm25 = ((k1 + 1) * tf) / + (k1 * ((1 - b) + b * (h_output_cols_lengths(0, col) / avg_col_length)) + tf); + result = idf * bm25; + } + out_host_matrix(row, col) = result; + out_host_vector(count) = result; + count++; + } + } + } + raft::copy(results.data_handle(), out_host_vector.data_handle(), out_host_vector.size(), stream); +} + +template +int get_dupe_mask_count(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + const raft::device_vector_view& mask) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + raft::sparse::op::coo_sort(int(rows.size()), + int(columns.size()), + int(values.size()), + rows.data_handle(), + columns.data_handle(), + values.data_handle(), + stream); + + raft::sparse::op::compute_duplicates_mask( + mask.data_handle(), rows.data_handle(), columns.data_handle(), rows.size(), stream); + + int col_nnz_count = thrust::reduce(raft::resource::get_thrust_policy(handle), + mask.data_handle(), + mask.data_handle() + mask.size()); + return col_nnz_count; +} + +template +void remove_dupes(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + raft::device_vector_view mask, + const raft::device_vector_view& out_rows, + const raft::device_vector_view& out_cols, + const raft::device_vector_view& out_vals, + int num_rows = 128) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + auto col_counts = raft::make_device_vector(handle, columns.size()); + + thrust::fill(raft::resource::get_thrust_policy(handle), + col_counts.data_handle(), + col_counts.data_handle() + col_counts.size(), + 1.0f); + + auto keys_out = raft::make_device_vector(handle, num_rows); + auto counts_out = raft::make_device_vector(handle, num_rows); + + thrust::reduce_by_key(raft::resource::get_thrust_policy(handle), + rows.data_handle(), + rows.data_handle() + rows.size(), + col_counts.data_handle(), + keys_out.data_handle(), + counts_out.data_handle()); + + auto mask_out = raft::make_device_vector(handle, rows.size()); + + raft::linalg::map(handle, mask_out.view(), raft::cast_op{}, raft::make_const_mdspan(mask)); + + auto values_c = raft::make_device_vector(handle, values.size()); + raft::linalg::map(handle, + values_c.view(), + raft::mul_op{}, + raft::make_const_mdspan(values), + raft::make_const_mdspan(mask_out.view())); + + auto keys_nnz_out = raft::make_device_vector(handle, num_rows); + auto counts_nnz_out = raft::make_device_vector(handle, num_rows); + + thrust::reduce_by_key(raft::resource::get_thrust_policy(handle), + rows.data_handle(), + rows.data_handle() + rows.size(), + mask.data_handle(), + keys_nnz_out.data_handle(), + counts_nnz_out.data_handle()); + + raft::sparse::op::coo_remove_scalar(rows.data_handle(), + columns.data_handle(), + values_c.data_handle(), + values_c.size(), + out_rows.data_handle(), + out_cols.data_handle(), + out_vals.data_handle(), + counts_nnz_out.data_handle(), + counts_out.data_handle(), + 0, + num_rows, + stream); +} + +template +void create_dataset(raft::resources& handle, + raft::device_vector_view rows, + raft::device_vector_view columns, + raft::device_vector_view values, + int max_term_occurence_doc = 5, + int num_rows_unique = 7, + int num_cols_unique = 7, + int seed = 12345) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + raft::random::RngState rng(seed); + + auto d_out = raft::make_device_vector(handle, rows.size() * 2); + + int theta_guide = max(num_rows_unique, num_cols_unique); + auto theta = raft::make_device_vector(handle, theta_guide * 4); + + raft::random::uniform(handle, rng, theta.view(), 0.0f, 1.0f); + + raft::random::rmat_rectangular_gen(d_out.data_handle(), + rows.data_handle(), + columns.data_handle(), + theta.data_handle(), + num_rows_unique, + num_cols_unique, + int(values.size()), + stream, + rng); + + auto vals = raft::make_device_vector(handle, rows.size()); + raft::random::uniformInt(handle, rng, vals.view(), 1, max_term_occurence_doc); + raft::linalg::map(handle, values, raft::cast_op{}, raft::make_const_mdspan(vals.view())); +} + +}; // namespace raft::util \ No newline at end of file diff --git a/cpp/test/sparse/neighbors/brute_force_coo.cu b/cpp/test/sparse/neighbors/brute_force_coo.cu new file mode 100644 index 0000000000..f1ebd6b578 --- /dev/null +++ b/cpp/test/sparse/neighbors/brute_force_coo.cu @@ -0,0 +1,190 @@ +/* + * Copyright (c) 2018-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. + */ + +#include "../../test_utils.cuh" + +#include +#include +#include +#include +#include + +#include +#include + +namespace raft { +namespace sparse { +namespace selection { + +using namespace raft; +using namespace raft::sparse; + +template +struct SparseKNNInputs { + value_idx n_cols; + + std::vector indptr_h; + std::vector indices_h; + std::vector data_h; + + std::vector out_dists_ref_h; + std::vector out_indices_ref_h; + + int k; + + int batch_size_index = 2; + int batch_size_query = 2; + + raft::distance::DistanceType metric = raft::distance::DistanceType::L2SqrtExpanded; +}; + +template +::std::ostream& operator<<(::std::ostream& os, const SparseKNNInputs& dims) +{ + return os; +} + +template +class SparseKNNCOOTest : public ::testing::TestWithParam> { + public: + SparseKNNCOOTest() + : params(::testing::TestWithParam>::GetParam()), + indptr(0, resource::get_cuda_stream(handle)), + indices(0, resource::get_cuda_stream(handle)), + data(0, resource::get_cuda_stream(handle)), + out_indices(0, resource::get_cuda_stream(handle)), + out_dists(0, resource::get_cuda_stream(handle)), + out_indices_ref(0, resource::get_cuda_stream(handle)), + out_dists_ref(0, resource::get_cuda_stream(handle)) + { + } + + protected: + void SetUp() override + { + n_rows = params.indptr_h.size() - 1; + nnz = params.indices_h.size(); + k = params.k; + + auto out_indices_dev = raft::make_device_vector(handle, n_rows * k); + auto out_dists_dev = raft::make_device_vector(handle, n_rows * k); + + auto rows = raft::make_device_vector(handle, nnz); + + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + make_data(); + + raft::sparse::convert::csr_to_coo( + indptr.data(), int(indptr.size()), rows.data_handle(), nnz, stream); + + auto coo_struct_view = raft::make_device_coordinate_structure_view( + rows.data_handle(), indices.data(), n_rows, params.n_cols, int(data.size())); + auto c_matrix = raft::make_device_coo_matrix( + handle, coo_struct_view); + raft::update_device( + c_matrix.view().get_elements().data(), data.data(), data.size(), stream); + + raft::sparse::neighbors::brute_force_knn(c_matrix, + c_matrix, + out_indices_dev.view(), + out_dists_dev.view(), + k, + handle, + params.batch_size_index, + params.batch_size_query, + params.metric); + + raft::copy(out_indices.data(), out_indices_dev.data_handle(), out_indices_dev.size(), stream); + raft::copy(out_dists.data(), out_dists_dev.data_handle(), out_dists_dev.size(), stream); + + RAFT_CUDA_TRY(cudaStreamSynchronize(resource::get_cuda_stream(handle))); + } + + void compare() + { + ASSERT_TRUE(devArrMatch( + out_dists_ref.data(), out_dists.data(), n_rows * k, CompareApprox(1e-4))); + ASSERT_TRUE( + devArrMatch(out_indices_ref.data(), out_indices.data(), n_rows * k, Compare())); + } + + protected: + void make_data() + { + std::vector indptr_h = params.indptr_h; + std::vector indices_h = params.indices_h; + std::vector data_h = params.data_h; + + auto stream = resource::get_cuda_stream(handle); + indptr.resize(indptr_h.size(), stream); + indices.resize(indices_h.size(), stream); + data.resize(data_h.size(), stream); + + update_device(indptr.data(), indptr_h.data(), indptr_h.size(), stream); + update_device(indices.data(), indices_h.data(), indices_h.size(), stream); + update_device(data.data(), data_h.data(), data_h.size(), stream); + + std::vector out_dists_ref_h = params.out_dists_ref_h; + std::vector out_indices_ref_h = params.out_indices_ref_h; + + out_indices_ref.resize(out_indices_ref_h.size(), stream); + out_dists_ref.resize(out_dists_ref_h.size(), stream); + + update_device( + out_indices_ref.data(), out_indices_ref_h.data(), out_indices_ref_h.size(), stream); + update_device(out_dists_ref.data(), out_dists_ref_h.data(), out_dists_ref_h.size(), stream); + + out_dists.resize(n_rows * k, stream); + out_indices.resize(n_rows * k, stream); + } + + raft::resources handle; + + int n_rows, nnz, k; + + // input data + rmm::device_uvector indptr, indices; + rmm::device_uvector data; + + // output data + rmm::device_uvector out_indices; + rmm::device_uvector out_dists; + + rmm::device_uvector out_indices_ref; + rmm::device_uvector out_dists_ref; + + SparseKNNInputs params; +}; + +const std::vector> inputs_i32_f = { + {9, // ncols + {0, 2, 4, 6, 8}, // indptr + {0, 4, 0, 3, 0, 2, 0, 8}, // indices + {0.0f, 1.0f, 5.0f, 6.0f, 5.0f, 6.0f, 0.0f, 1.0f}, // data + {0, 1.41421, 0, 7.87401, 0, 7.87401, 0, 1.41421}, // dists + {0, 3, 1, 0, 2, 0, 3, 0}, // inds + 2, + 2, + 2, + raft::distance::DistanceType::L2SqrtExpanded}}; +typedef SparseKNNCOOTest SparseKNNCOOTestF; +TEST_P(SparseKNNCOOTestF, Result) { compare(); } +INSTANTIATE_TEST_CASE_P(SparseKNNCOOTest, SparseKNNCOOTestF, ::testing::ValuesIn(inputs_i32_f)); + +}; // end namespace selection +}; // end namespace sparse +}; // end namespace raft diff --git a/cpp/test/sparse/neighbors/brute_force_csr.cu b/cpp/test/sparse/neighbors/brute_force_csr.cu new file mode 100644 index 0000000000..dec1914e09 --- /dev/null +++ b/cpp/test/sparse/neighbors/brute_force_csr.cu @@ -0,0 +1,183 @@ +/* + * Copyright (c) 2018-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. + */ + +#include "../../test_utils.cuh" + +#include +#include +#include +#include + +#include +#include + +namespace raft { +namespace sparse { +namespace selection { + +using namespace raft; +using namespace raft::sparse; + +template +struct SparseKNNInputs { + value_idx n_cols; + + std::vector indptr_h; + std::vector indices_h; + std::vector data_h; + + std::vector out_dists_ref_h; + std::vector out_indices_ref_h; + + int k; + + int batch_size_index = 2; + int batch_size_query = 2; + + raft::distance::DistanceType metric = raft::distance::DistanceType::L2SqrtExpanded; +}; + +template +::std::ostream& operator<<(::std::ostream& os, const SparseKNNInputs& dims) +{ + return os; +} + +template +class SparseKNNCSRTest : public ::testing::TestWithParam> { + public: + SparseKNNCSRTest() + : params(::testing::TestWithParam>::GetParam()), + indptr(0, resource::get_cuda_stream(handle)), + indices(0, resource::get_cuda_stream(handle)), + data(0, resource::get_cuda_stream(handle)), + out_indices(0, resource::get_cuda_stream(handle)), + out_dists(0, resource::get_cuda_stream(handle)), + out_indices_ref(0, resource::get_cuda_stream(handle)), + out_dists_ref(0, resource::get_cuda_stream(handle)) + { + } + + protected: + void SetUp() override + { + n_rows = params.indptr_h.size() - 1; + nnz = params.indices_h.size(); + k = params.k; + auto out_indices_dev = raft::make_device_vector(handle, n_rows * k); + auto out_dists_dev = raft::make_device_vector(handle, n_rows * k); + + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + make_data(); + auto csr_struct_view = raft::make_device_compressed_structure_view( + indptr.data(), indices.data(), n_rows, params.n_cols, int(data.size())); + auto c_matrix = raft::make_device_csr_matrix(handle, csr_struct_view); + + raft::update_device( + c_matrix.view().get_elements().data(), data.data(), data.size(), stream); + + raft::sparse::neighbors::brute_force_knn(c_matrix, + c_matrix, + out_indices_dev.view(), + out_dists_dev.view(), + k, + handle, + params.batch_size_index, + params.batch_size_query, + params.metric); + + raft::copy(out_indices.data(), out_indices_dev.data_handle(), out_indices_dev.size(), stream); + raft::copy(out_dists.data(), out_dists_dev.data_handle(), out_dists_dev.size(), stream); + std::cout << "finished copy" << std::endl; + + RAFT_CUDA_TRY(cudaStreamSynchronize(resource::get_cuda_stream(handle))); + } + + void compare() + { + ASSERT_TRUE(devArrMatch( + out_dists_ref.data(), out_dists.data(), n_rows * k, CompareApprox(1e-4))); + ASSERT_TRUE( + devArrMatch(out_indices_ref.data(), out_indices.data(), n_rows * k, Compare())); + } + + protected: + void make_data() + { + std::vector indptr_h = params.indptr_h; + std::vector indices_h = params.indices_h; + std::vector data_h = params.data_h; + + auto stream = resource::get_cuda_stream(handle); + indptr.resize(indptr_h.size(), stream); + indices.resize(indices_h.size(), stream); + data.resize(data_h.size(), stream); + + update_device(indptr.data(), indptr_h.data(), indptr_h.size(), stream); + update_device(indices.data(), indices_h.data(), indices_h.size(), stream); + update_device(data.data(), data_h.data(), data_h.size(), stream); + + std::vector out_dists_ref_h = params.out_dists_ref_h; + std::vector out_indices_ref_h = params.out_indices_ref_h; + + out_indices_ref.resize(out_indices_ref_h.size(), stream); + out_dists_ref.resize(out_dists_ref_h.size(), stream); + + update_device( + out_indices_ref.data(), out_indices_ref_h.data(), out_indices_ref_h.size(), stream); + update_device(out_dists_ref.data(), out_dists_ref_h.data(), out_dists_ref_h.size(), stream); + + out_dists.resize(n_rows * k, stream); + out_indices.resize(n_rows * k, stream); + } + + raft::resources handle; + + int n_rows, nnz, k; + + // input data + rmm::device_uvector indptr, indices; + rmm::device_uvector data; + + // output data + rmm::device_uvector out_indices; + rmm::device_uvector out_dists; + + rmm::device_uvector out_indices_ref; + rmm::device_uvector out_dists_ref; + + SparseKNNInputs params; +}; + +const std::vector> inputs_i32_f = { + {9, // ncols + {0, 2, 4, 6, 8}, // indptr + {0, 4, 0, 3, 0, 2, 0, 8}, // indices + {0.0f, 1.0f, 5.0f, 6.0f, 5.0f, 6.0f, 0.0f, 1.0f}, // data + {0, 1.41421, 0, 7.87401, 0, 7.87401, 0, 1.41421}, // dists + {0, 3, 1, 0, 2, 0, 3, 0}, // inds + 2, + 2, + 2, + raft::distance::DistanceType::L2SqrtExpanded}}; +typedef SparseKNNCSRTest SparseKNNCSRTestF; +TEST_P(SparseKNNCSRTestF, Result) { compare(); } +INSTANTIATE_TEST_CASE_P(SparseKNNCSRTest, SparseKNNCSRTestF, ::testing::ValuesIn(inputs_i32_f)); + +}; // end namespace selection +}; // end namespace sparse +}; // end namespace raft diff --git a/cpp/test/sparse/preprocess_coo.cu b/cpp/test/sparse/preprocess_coo.cu new file mode 100644 index 0000000000..b26e5122d7 --- /dev/null +++ b/cpp/test/sparse/preprocess_coo.cu @@ -0,0 +1,181 @@ +/* + * 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. + */ + +#include "../preprocess_utils.cu" +#include "../test_utils.cuh" + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace raft { +namespace sparse { + +template +void calc_tfidf_bm25(raft::resources& handle, + raft::device_coo_matrix_view coo_in, + raft::device_vector_view results, + bool tf_idf = false) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + int num_rows = coo_in.structure_view().get_n_rows(); + int num_cols = coo_in.structure_view().get_n_cols(); + int rows_size = coo_in.structure_view().get_cols().size(); + int cols_size = coo_in.structure_view().get_rows().size(); + int elements_size = coo_in.get_elements().size(); + + auto h_rows = raft::make_host_vector(handle, rows_size); + auto h_cols = raft::make_host_vector(handle, cols_size); + auto h_elems = raft::make_host_vector(handle, elements_size); + + raft::copy(h_rows.data_handle(), + coo_in.structure_view().get_rows().data(), + coo_in.structure_view().get_rows().size(), + stream); + raft::copy(h_cols.data_handle(), + coo_in.structure_view().get_cols().data(), + coo_in.structure_view().get_cols().size(), + stream); + raft::copy( + h_elems.data_handle(), coo_in.get_elements().data(), coo_in.get_elements().size(), stream); + raft::util::preproc_kernel( + handle, h_rows.view(), h_cols.view(), h_elems.view(), results, num_rows, num_cols, tf_idf); +} + +template +struct SparsePreprocessInputs { + int n_rows; + int n_cols; + int nnz_edges; +}; + +template +class SparsePreprocessCoo + : public ::testing::TestWithParam> { + public: + SparsePreprocessCoo() + : params(::testing::TestWithParam>::GetParam()), + stream(resource::get_cuda_stream(handle)) + { + } + + protected: + void SetUp() override {} + + void Run(bool bm25_on) + { + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + int num_rows = pow(2, params.n_rows); + int num_cols = pow(2, params.n_cols); + + auto rows = raft::make_device_vector(handle, params.nnz_edges); + auto columns = raft::make_device_vector(handle, params.nnz_edges); + auto values = raft::make_device_vector(handle, params.nnz_edges); + auto mask = raft::make_device_vector(handle, params.nnz_edges); + + raft::util::create_dataset( + handle, rows.view(), columns.view(), values.view(), 5, params.n_rows, params.n_cols); + int non_dupe_nnz_count = raft::util::get_dupe_mask_count( + handle, rows.view(), columns.view(), values.view(), mask.view()); + + auto rows_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + auto columns_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + auto values_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + raft::util::remove_dupes(handle, + rows.view(), + columns.view(), + values.view(), + mask.view(), + rows_nnz.view(), + columns_nnz.view(), + values_nnz.view(), + num_rows); + + auto coo_struct_view = raft::make_device_coordinate_structure_view(rows_nnz.data_handle(), + columns_nnz.data_handle(), + num_rows, + num_cols, + int(values_nnz.size())); + auto c_matrix = + raft::make_device_coo_matrix(handle, coo_struct_view); + raft::update_device( + c_matrix.view().get_elements().data(), values_nnz.data_handle(), values_nnz.size(), stream); + + auto result = raft::make_device_vector(handle, values_nnz.size()); + auto bm25_vals = raft::make_device_vector(handle, values_nnz.size()); + auto tfidf_vals = raft::make_device_vector(handle, values_nnz.size()); + + if (bm25_on) { + sparse::matrix::encode_bm25(handle, c_matrix.view(), result.view()); + calc_tfidf_bm25(handle, c_matrix.view(), bm25_vals.view()); + ASSERT_TRUE(raft::devArrMatch(bm25_vals.data_handle(), + result.data_handle(), + result.size(), + raft::CompareApprox(2e-5), + stream)); + } else { + sparse::matrix::encode_tfidf(handle, c_matrix.view(), result.view()); + calc_tfidf_bm25(handle, c_matrix.view(), tfidf_vals.view(), true); + ASSERT_TRUE(raft::devArrMatch(tfidf_vals.data_handle(), + result.data_handle(), + result.size(), + raft::CompareApprox(2e-5), + stream)); + } + + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + } + + protected: + raft::resources handle; + cudaStream_t stream; + + SparsePreprocessInputs params; +}; + +using SparsePreprocessTfidfCoo = SparsePreprocessCoo; +TEST_P(SparsePreprocessTfidfCoo, Result) { Run(false); } + +using SparsePreprocessBm25Coo = SparsePreprocessCoo; +TEST_P(SparsePreprocessBm25Coo, Result) { Run(true); } + +const std::vector> sparse_preprocess_inputs = { + { + 10, // n_rows_factor + 10, // n_cols_factor + 1000 // nnz_edges + }, +}; + +INSTANTIATE_TEST_CASE_P(SparsePreprocessCoo, + SparsePreprocessTfidfCoo, + ::testing::ValuesIn(sparse_preprocess_inputs)); +INSTANTIATE_TEST_CASE_P(SparsePreprocessCoo, + SparsePreprocessBm25Coo, + ::testing::ValuesIn(sparse_preprocess_inputs)); + +} // namespace sparse +} // namespace raft \ No newline at end of file diff --git a/cpp/test/sparse/preprocess_csr.cu b/cpp/test/sparse/preprocess_csr.cu new file mode 100644 index 0000000000..eab270ce79 --- /dev/null +++ b/cpp/test/sparse/preprocess_csr.cu @@ -0,0 +1,188 @@ +/* + * 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. + */ + +#include "../preprocess_utils.cu" +#include "../test_utils.cuh" + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace raft { +namespace sparse { + +template +void calc_tfidf_bm25(raft::resources& handle, + raft::device_csr_matrix_view csr_in, + raft::device_vector_view results, + bool tf_idf = false) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + int num_rows = csr_in.structure_view().get_n_rows(); + int num_cols = csr_in.structure_view().get_n_cols(); + int rows_size = csr_in.structure_view().get_indptr().size(); + int cols_size = csr_in.structure_view().get_indices().size(); + int elements_size = csr_in.get_elements().size(); + + auto h_rows = raft::make_host_vector(handle, rows_size); + auto h_cols = raft::make_host_vector(handle, cols_size); + auto h_elems = raft::make_host_vector(handle, elements_size); + + auto indptr = raft::make_device_vector_view( + csr_in.structure_view().get_indptr().data(), csr_in.structure_view().get_indptr().size()); + auto indices = raft::make_device_vector_view( + csr_in.structure_view().get_indices().data(), csr_in.structure_view().get_indices().size()); + auto values = raft::make_device_vector_view(csr_in.get_elements().data(), + csr_in.get_elements().size()); + auto rows = raft::make_device_vector(handle, values.size()); + + raft::sparse::convert::csr_to_coo( + indptr.data_handle(), num_rows, rows.data_handle(), rows.size(), stream); + + raft::copy(h_rows.data_handle(), rows.data_handle(), rows.size(), stream); + raft::copy(h_cols.data_handle(), indices.data_handle(), cols_size, stream); + raft::copy(h_elems.data_handle(), values.data_handle(), values.size(), stream); + raft::util::preproc_kernel( + handle, h_rows.view(), h_cols.view(), h_elems.view(), results, num_rows, num_cols, tf_idf); +} + +template +struct SparsePreprocessInputs { + int n_rows; + int n_cols; + int nnz_edges; +}; + +template +class SparsePreprocessCSR + : public ::testing::TestWithParam> { + public: + SparsePreprocessCSR() + : params(::testing::TestWithParam>::GetParam()), + stream(resource::get_cuda_stream(handle)) + { + } + + protected: + void SetUp() override {} + + void Run(bool bm25_on) + { + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + int num_rows = pow(2, params.n_rows); + int num_cols = pow(2, params.n_cols); + + auto rows = raft::make_device_vector(handle, params.nnz_edges); + auto columns = raft::make_device_vector(handle, params.nnz_edges); + auto values = raft::make_device_vector(handle, params.nnz_edges); + auto mask = raft::make_device_vector(handle, params.nnz_edges); + + raft::util::create_dataset( + handle, rows.view(), columns.view(), values.view(), 5, params.n_rows, params.n_cols); + int non_dupe_nnz_count = raft::util::get_dupe_mask_count( + handle, rows.view(), columns.view(), values.view(), mask.view()); + + auto rows_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + auto columns_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + auto values_nnz = raft::make_device_vector(handle, non_dupe_nnz_count); + raft::util::remove_dupes(handle, + rows.view(), + columns.view(), + values.view(), + mask.view(), + rows_nnz.view(), + columns_nnz.view(), + values_nnz.view(), + num_rows); + auto rows_csr = raft::make_device_vector(handle, non_dupe_nnz_count); + raft::sparse::convert::sorted_coo_to_csr( + rows_nnz.data_handle(), int(rows_nnz.size()), rows_csr.data_handle(), num_rows, stream); + + auto csr_struct_view = raft::make_device_compressed_structure_view(rows_csr.data_handle(), + columns_nnz.data_handle(), + num_rows, + num_cols, + int(values_nnz.size())); + auto c_matrix = + raft::make_device_csr_matrix(handle, csr_struct_view); + + raft::update_device( + c_matrix.view().get_elements().data(), values_nnz.data_handle(), values_nnz.size(), stream); + + auto result = raft::make_device_vector(handle, values_nnz.size()); + auto bm25_vals = raft::make_device_vector(handle, values_nnz.size()); + auto tfidf_vals = raft::make_device_vector(handle, values_nnz.size()); + + if (bm25_on) { + sparse::matrix::encode_bm25(handle, c_matrix.view(), result.view()); + calc_tfidf_bm25(handle, c_matrix.view(), bm25_vals.view()); + ASSERT_TRUE(raft::devArrMatch(bm25_vals.data_handle(), + result.data_handle(), + result.size(), + raft::CompareApprox(2e-5), + stream)); + } else { + sparse::matrix::encode_tfidf(handle, c_matrix.view(), result.view()); + calc_tfidf_bm25(handle, c_matrix.view(), tfidf_vals.view(), true); + ASSERT_TRUE(raft::devArrMatch(tfidf_vals.data_handle(), + result.data_handle(), + result.size(), + raft::CompareApprox(2e-5), + stream)); + } + + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + } + + protected: + raft::resources handle; + cudaStream_t stream; + + SparsePreprocessInputs params; +}; + +using SparsePreprocessTfidfCsr = SparsePreprocessCSR; +TEST_P(SparsePreprocessTfidfCsr, Result) { Run(false); } + +using SparsePreprocessBm25Csr = SparsePreprocessCSR; +TEST_P(SparsePreprocessBm25Csr, Result) { Run(true); } + +const std::vector> sparse_preprocess_inputs = { + { + 7, // n_rows_factor + 5, // n_cols_factor + 10 // num nnz values + }, +}; + +INSTANTIATE_TEST_CASE_P(SparsePreprocessCSR, + SparsePreprocessTfidfCsr, + ::testing::ValuesIn(sparse_preprocess_inputs)); +INSTANTIATE_TEST_CASE_P(SparsePreprocessCSR, + SparsePreprocessBm25Csr, + ::testing::ValuesIn(sparse_preprocess_inputs)); + +} // namespace sparse +} // namespace raft \ No newline at end of file diff --git a/docs/source/contributing.md b/docs/source/contributing.md index 1b4071d0a5..47eb88c429 100755 --- a/docs/source/contributing.md +++ b/docs/source/contributing.md @@ -88,6 +88,4 @@ others know you are working on it. If you have any questions related to the implementation of the issue, ask them in the issue instead of the PR. ## Attribution -Portions adopted from https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md - - +Portions adopted from https://github.com/pytorch/pytorch/blob/main/CONTRIBUTING.md