Skip to content

Add binary_mult and const fixes for ratio and sqrt apis #2940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 4 commits into from
Feb 11, 2026
Merged

Add binary_mult and const fixes for ratio and sqrt apis #2940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e8d4790
fix-const-sqrt
aamijar Feb 3, 2026
8903546
Merge branch 'main' into fix-const-sqrt
aamijar Feb 3, 2026
65d6de9
add gtests binary_mult
aamijar Feb 10, 2026
46c2c6d
Merge branch 'main' into fix-const-sqrt
aamijar Feb 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 30 additions & 1 deletion cpp/include/raft/linalg/matrix_vector.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand All @@ -18,6 +18,35 @@ namespace raft::linalg {
* @{
*/

/**
* @brief multiply each row or column of matrix with vector
* @param[in] handle: raft handle for managing library resources
* @param[inout] data: input matrix, results are in-place
* @param[in] vec: input vector
*/
template <Apply apply, typename math_t, typename idx_t, typename layout_t>
void binary_mult(raft::resources const& handle,
raft::device_matrix_view<math_t, idx_t, layout_t> data,
raft::device_vector_view<const math_t, idx_t> vec)
{
constexpr auto row_major = std::is_same_v<layout_t, raft::row_major>;
constexpr auto bcast_along_rows = apply == Apply::ALONG_ROWS;

idx_t vec_size = bcast_along_rows ? data.extent(1) : data.extent(0);

RAFT_EXPECTS(
vec.extent(0) == vec_size,
"If `bcast_along_rows==true`, vector size must equal number of columns in the matrix."
"If `bcast_along_rows==false`, vector size must equal number of rows in the matrix.");

matrix::detail::matrixVectorBinaryMult<row_major, bcast_along_rows>(
data.data_handle(),
vec.data_handle(),
data.extent(0),
data.extent(1),
resource::get_cuda_stream(handle));
}

/**
* @brief multiply each row or column of matrix with vector, skipping zeros in vector
* @param [in] handle: raft handle for managing library resources
Expand Down
8 changes: 4 additions & 4 deletions cpp/include/raft/matrix/detail/math.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -58,7 +58,7 @@ void power(math_t* in, math_t* out, int len, cudaStream_t stream)
}

template <typename math_t, typename IdxType = int>
void seqRoot(math_t* in,
void seqRoot(const math_t* in,
math_t* out,
math_t scalar,
IdxType len,
Expand Down Expand Up @@ -94,7 +94,7 @@ void seqRoot(
}

template <typename math_t, typename IdxType = int>
void seqRoot(math_t* in, math_t* out, IdxType len, cudaStream_t stream)
void seqRoot(const math_t* in, math_t* out, IdxType len, cudaStream_t stream)
{
math_t scalar = 1.0;
seqRoot(in, out, scalar, len, stream);
Expand Down Expand Up @@ -184,7 +184,7 @@ void setValue(math_t* out, const math_t* in, math_t scalar, int len, cudaStream_

template <typename math_t, typename IdxType = int>
void ratio(
raft::resources const& handle, math_t* src, math_t* dest, IdxType len, cudaStream_t stream)
raft::resources const& handle, const math_t* src, math_t* dest, IdxType len, cudaStream_t stream)
{
auto d_src = src;
auto d_dest = dest;
Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/matrix/ratio.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

#pragma once

#include <raft/core/device_mdspan.hpp>
#include <raft/core/resource/cuda_stream.hpp>
#include <raft/matrix/detail/matrix.cuh>
#include <raft/matrix/detail/math.cuh>

namespace raft::matrix {

Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/matrix/sqrt.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2023, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand All @@ -8,7 +8,7 @@
#include <raft/core/device_mdspan.hpp>
#include <raft/core/host_mdspan.hpp>
#include <raft/core/resource/cuda_stream.hpp>
#include <raft/matrix/detail/matrix.cuh>
#include <raft/matrix/detail/math.cuh>

namespace raft::matrix {

Expand Down
79 changes: 44 additions & 35 deletions cpp/tests/linalg/matrix_vector.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION.
* SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
* SPDX-License-Identifier: Apache-2.0
*/

Expand Down Expand Up @@ -124,6 +124,20 @@ void matrix_vector_op_launch(const raft::resources& handle,
binary_sub<Apply::ALONG_COLUMNS>(handle, in_col_major, vec1_view);
}
}
} else if (operation_type == 5) {
if (row_major) {
if (apply == Apply::ALONG_ROWS) {
binary_mult<Apply::ALONG_ROWS>(handle, in_row_major, vec1_view);
} else {
binary_mult<Apply::ALONG_COLUMNS>(handle, in_row_major, vec1_view);
}
} else {
if (apply == Apply::ALONG_ROWS) {
binary_mult<Apply::ALONG_ROWS>(handle, in_col_major, vec1_view);
} else {
binary_mult<Apply::ALONG_COLUMNS>(handle, in_col_major, vec1_view);
}
}
} else {
THROW("Unknown operation type '%d'!", (int)operation_type);
}
Expand Down Expand Up @@ -166,6 +180,8 @@ void naive_matrix_vector_op_launch(const raft::resources& handle,
naiveMatVec(in, in, vec1, D, N, row_major, bcast_along_rows, raft::add_op{}, stream);
} else if (operation_type == 4) {
naiveMatVec(in, in, vec1, D, N, row_major, bcast_along_rows, raft::sub_op{}, stream);
} else if (operation_type == 5) {
naiveMatVec(in, in, vec1, D, N, row_major, bcast_along_rows, raft::mul_op{}, stream);
} else {
THROW("Unknown operation type '%d'!", (int)operation_type);
}
Expand Down Expand Up @@ -223,23 +239,18 @@ class MatrixVectorTest : public ::testing::TestWithParam<MatrixVectorInputs<T, I
};

const std::vector<MatrixVectorInputs<float, int>> inputsf_i32 = {
{0.00001f, 1024, 32, 0, true, true, 1234ULL},
{0.00001f, 1024, 64, 1, true, true, 1234ULL},
{0.00001f, 1024, 32, 2, true, false, 1234ULL},
{0.00001f, 1024, 64, 3, true, false, 1234ULL},
{0.00001f, 1024, 32, 4, false, true, 1234ULL},
{0.00001f, 1024, 64, 0, false, true, 1234ULL},
{0.00001f, 1024, 32, 1, false, false, 1234ULL},
{0.00001f, 1024, 64, 2, false, false, 1234ULL},
{0.00001f, 1024, 32, 0, true, true, 1234ULL}, {0.00001f, 1024, 64, 1, true, true, 1234ULL},
{0.00001f, 1024, 32, 2, true, false, 1234ULL}, {0.00001f, 1024, 64, 3, true, false, 1234ULL},
{0.00001f, 1024, 32, 4, false, true, 1234ULL}, {0.00001f, 1024, 64, 0, false, true, 1234ULL},
{0.00001f, 1024, 32, 1, false, false, 1234ULL}, {0.00001f, 1024, 64, 2, false, false, 1234ULL},

{0.00001f, 1024, 32, 3, true, true, 1234ULL},
{0.00001f, 1024, 64, 4, true, true, 1234ULL},
{0.00001f, 1024, 32, 0, true, false, 1234ULL},
{0.00001f, 1024, 64, 1, true, false, 1234ULL},
{0.00001f, 1024, 32, 2, false, true, 1234ULL},
{0.00001f, 1024, 64, 3, false, true, 1234ULL},
{0.00001f, 1024, 32, 4, false, false, 1234ULL},
{0.00001f, 1024, 64, 0, false, false, 1234ULL}};
{0.00001f, 1024, 32, 3, true, true, 1234ULL}, {0.00001f, 1024, 64, 4, true, true, 1234ULL},
{0.00001f, 1024, 32, 0, true, false, 1234ULL}, {0.00001f, 1024, 64, 1, true, false, 1234ULL},
{0.00001f, 1024, 32, 2, false, true, 1234ULL}, {0.00001f, 1024, 64, 3, false, true, 1234ULL},
{0.00001f, 1024, 32, 4, false, false, 1234ULL}, {0.00001f, 1024, 64, 0, false, false, 1234ULL},

{0.00001f, 1024, 32, 5, true, true, 1234ULL}, {0.00001f, 1024, 64, 5, true, false, 1234ULL},
{0.00001f, 1024, 32, 5, false, true, 1234ULL}, {0.00001f, 1024, 64, 5, false, false, 1234ULL}};
typedef MatrixVectorTest<float, int> MatrixVectorTestF_i32;
TEST_P(MatrixVectorTestF_i32, Result)
{
Expand All @@ -251,7 +262,9 @@ INSTANTIATE_TEST_SUITE_P(MatrixVectorTests,
::testing::ValuesIn(inputsf_i32));

const std::vector<MatrixVectorInputs<float, size_t>> inputsf_i64 = {
{0.00001f, 2500, 250, 0, false, false, 1234ULL}, {0.00001f, 2500, 250, 1, false, false, 1234ULL}};
{0.00001f, 2500, 250, 0, false, false, 1234ULL},
{0.00001f, 2500, 250, 1, false, false, 1234ULL},
{0.00001f, 2500, 250, 5, false, false, 1234ULL}};
typedef MatrixVectorTest<float, size_t> MatrixVectorTestF_i64;
TEST_P(MatrixVectorTestF_i64, Result)
{
Expand All @@ -263,23 +276,18 @@ INSTANTIATE_TEST_SUITE_P(MatrixVectorTests,
::testing::ValuesIn(inputsf_i64));

const std::vector<MatrixVectorInputs<double, int>> inputsd_i32 = {
{0.0000001, 1024, 32, 0, true, true, 1234ULL},
{0.0000001, 1024, 64, 1, true, true, 1234ULL},
{0.0000001, 1024, 32, 2, true, false, 1234ULL},
{0.0000001, 1024, 64, 3, true, false, 1234ULL},
{0.0000001, 1024, 32, 4, false, true, 1234ULL},
{0.0000001, 1024, 64, 0, false, true, 1234ULL},
{0.0000001, 1024, 32, 1, false, false, 1234ULL},
{0.0000001, 1024, 64, 2, false, false, 1234ULL},
{0.0000001, 1024, 32, 0, true, true, 1234ULL}, {0.0000001, 1024, 64, 1, true, true, 1234ULL},
{0.0000001, 1024, 32, 2, true, false, 1234ULL}, {0.0000001, 1024, 64, 3, true, false, 1234ULL},
{0.0000001, 1024, 32, 4, false, true, 1234ULL}, {0.0000001, 1024, 64, 0, false, true, 1234ULL},
{0.0000001, 1024, 32, 1, false, false, 1234ULL}, {0.0000001, 1024, 64, 2, false, false, 1234ULL},

{0.0000001, 1024, 32, 3, true, true, 1234ULL}, {0.0000001, 1024, 64, 4, true, true, 1234ULL},
{0.0000001, 1024, 32, 0, true, false, 1234ULL}, {0.0000001, 1024, 64, 1, true, false, 1234ULL},
{0.0000001, 1024, 32, 2, false, true, 1234ULL}, {0.0000001, 1024, 64, 3, false, true, 1234ULL},
{0.0000001, 1024, 32, 4, false, false, 1234ULL}, {0.0000001, 1024, 64, 0, false, false, 1234ULL},

{0.0000001, 1024, 32, 3, true, true, 1234ULL},
{0.0000001, 1024, 64, 4, true, true, 1234ULL},
{0.0000001, 1024, 32, 0, true, false, 1234ULL},
{0.0000001, 1024, 64, 1, true, false, 1234ULL},
{0.0000001, 1024, 32, 2, false, true, 1234ULL},
{0.0000001, 1024, 64, 3, false, true, 1234ULL},
{0.0000001, 1024, 32, 4, false, false, 1234ULL},
{0.0000001, 1024, 64, 0, false, false, 1234ULL}};
{0.0000001, 1024, 32, 5, true, true, 1234ULL}, {0.0000001, 1024, 64, 5, true, false, 1234ULL},
{0.0000001, 1024, 32, 5, false, true, 1234ULL}, {0.0000001, 1024, 64, 5, false, false, 1234ULL}};
typedef MatrixVectorTest<double, int> MatrixVectorTestD_i32;
TEST_P(MatrixVectorTestD_i32, Result)
{
Expand All @@ -294,7 +302,8 @@ INSTANTIATE_TEST_SUITE_P(MatrixVectorTests,

const std::vector<MatrixVectorInputs<double, size_t>> inputsd_i64 = {
{0.0000001, 2500, 250, 0, false, false, 1234ULL},
{0.0000001, 2500, 250, 1, false, false, 1234ULL}};
{0.0000001, 2500, 250, 1, false, false, 1234ULL},
{0.0000001, 2500, 250, 5, false, false, 1234ULL}};
typedef MatrixVectorTest<double, size_t> MatrixVectorTestD_i64;
TEST_P(MatrixVectorTestD_i64, Result)
{
Expand Down