diff --git a/cpp/include/raft/sparse/linalg/detail/symmetrize.cuh b/cpp/include/raft/sparse/linalg/detail/symmetrize.cuh index 0423211c87..6b369057d9 100644 --- a/cpp/include/raft/sparse/linalg/detail/symmetrize.cuh +++ b/cpp/include/raft/sparse/linalg/detail/symmetrize.cuh @@ -145,7 +145,7 @@ void coo_symmetrize(COO* in, raft::launch_kernel(stream, grid, blk, - coo_symmetrize_kernel, + coo_symmetrize_kernel, in_row_ind.data(), in->rows(), in->cols(), @@ -358,7 +358,7 @@ void from_knn_symmetrize_matrix(const value_idx* __restrict__ knn_indices, raft::launch_kernel(stream, numBlocks, threadsPerBlock, - symmetric_find_size, + symmetric_find_size, knn_dists, knn_indices, n, @@ -369,7 +369,7 @@ void from_knn_symmetrize_matrix(const value_idx* __restrict__ knn_indices, raft::launch_kernel(stream, raft::ceildiv(n, (value_idx)1024), 1024, - reduce_find_size, + reduce_find_size, n, k, row_sizes.data(), @@ -397,7 +397,7 @@ void from_knn_symmetrize_matrix(const value_idx* __restrict__ knn_indices, raft::launch_kernel(stream, numBlocks, threadsPerBlock, - symmetric_sum, + symmetric_sum, edges, knn_dists, knn_indices, diff --git a/cpp/include/raft/util/kernel_launch.hpp b/cpp/include/raft/util/kernel_launch.hpp index 834d1a5982..7fb1a064a2 100644 --- a/cpp/include/raft/util/kernel_launch.hpp +++ b/cpp/include/raft/util/kernel_launch.hpp @@ -191,7 +191,10 @@ void launch_kernel(launch_on where, { where.config.gridDim = grid; where.config.blockDim = block; - detail::dispatch...>( + // Let dispatch deduce its by-value parameter types instead of explicitly forwarding Args. + // In particular, this drops outermost extended qualifiers such as __restrict__ before dispatch + // takes the address of each parameter copy for cudaLaunchKernelExC. + detail::dispatch( where.config, reinterpret_cast(kernel), where.location, std::forward(args)...); } @@ -222,8 +225,12 @@ requires(sizeof...(Params) == sizeof...(Args) && where.config.gridDim = grid; where.config.blockDim = block; - detail::dispatch( - where.config, reinterpret_cast(kernel), where.location, std::forward(args)...); + // Convert to the kernel parameter types before dispatch, then let dispatch deduce its by-value + // parameters so outermost extended qualifiers such as __restrict__ are not preserved. + detail::dispatch(where.config, + reinterpret_cast(kernel), + where.location, + static_cast(std::forward(args))...); } /** @} */ // end group kernel_launch diff --git a/cpp/tests/sparse/symmetrize.cu b/cpp/tests/sparse/symmetrize.cu index 47132ff8eb..b4f7c84b1d 100644 --- a/cpp/tests/sparse/symmetrize.cu +++ b/cpp/tests/sparse/symmetrize.cu @@ -212,6 +212,57 @@ TEST_P(COOSymmetrizeView, ResultView) raft::Compare())); } +TEST_P(COOSymmetrizeView, ResultLegacy) +{ + raft::resources handle; + auto stream = resource::get_cuda_stream(handle); + + raft::sparse::COO in(stream, params.nnz, params.n_rows, params.n_cols, false); + raft::sparse::COO out(stream); + + raft::update_device(in.rows(), params.in_rows_h.data(), params.nnz, stream); + raft::update_device(in.cols(), params.in_cols_h.data(), params.nnz, stream); + raft::update_device(in.vals(), params.in_vals_h.data(), params.nnz, stream); + + linalg::coo_symmetrize( + &in, + &out, + [] __device__(int row, int col, float val, float trans) { return val + trans; }, + stream); + + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + + ASSERT_EQ(out.nnz, params.nnz * 2); + ASSERT_TRUE( + raft::devArrMatch(out.rows(), params.exp_rows_h.data(), out.nnz, raft::Compare())); + ASSERT_TRUE( + raft::devArrMatch(out.cols(), params.exp_cols_h.data(), out.nnz, raft::Compare())); + ASSERT_TRUE(raft::devArrMatch( + out.vals(), params.exp_vals_h.data(), out.nnz, raft::Compare())); +} + +TEST(FromKnnSymmetrizeTest, RestrictedPointerArguments) +{ + raft::resources handle; + auto stream = resource::get_cuda_stream(handle); + + constexpr int n = 2; + constexpr int k = 1; + std::vector indices_h{1, 0}; + std::vector distances_h{0.5f, 0.5f}; + rmm::device_uvector indices(indices_h.size(), stream); + rmm::device_uvector distances(distances_h.size(), stream); + raft::sparse::COO out(stream); + + raft::update_device(indices.data(), indices_h.data(), indices_h.size(), stream); + raft::update_device(distances.data(), distances_h.data(), distances_h.size(), stream); + + linalg::from_knn_symmetrize_matrix(indices.data(), distances.data(), n, k, &out, stream); + + resource::sync_stream(handle); + EXPECT_EQ(out.nnz, 2 * n * k); +} + const std::vector> inputsf = { // first test fails without fix in #2582 { diff --git a/cpp/tests/util/kernel_launch.cu b/cpp/tests/util/kernel_launch.cu index 733c986737..f83ce7e668 100644 --- a/cpp/tests/util/kernel_launch.cu +++ b/cpp/tests/util/kernel_launch.cu @@ -28,6 +28,16 @@ RAFT_KERNEL write_one_kernel(int* out) if (threadIdx.x == 0 && blockIdx.x == 0) { *out = 1; } } +RAFT_KERNEL copy_restricted_kernel(int const* __restrict__ in, int* out) +{ + if (threadIdx.x == 0 && blockIdx.x == 0) { *out = *in; } +} + +void launch_write_one_with_restricted_pointer(raft::resources const& res, int* __restrict__ out) +{ + raft::launch_kernel(res, 1, 32, write_one_kernel, out); +} + RAFT_KERNEL smem_kernel(int* out) { extern __shared__ int shared[]; // NOLINT(modernize-avoid-c-arrays) @@ -81,6 +91,37 @@ TEST(KernelLaunch, SuccessfulLaunch) EXPECT_EQ(host_out, 1); } +TEST(KernelLaunch, RestrictedPointerArgument) +{ + raft::resources res; + rmm::device_uvector out(1, resource::get_cuda_stream(res)); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data(), 0, sizeof(int), resource::get_cuda_stream(res))); + + launch_write_one_with_restricted_pointer(res, out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + +TEST(KernelLaunch, ConvertedRestrictedPointerArgument) +{ + raft::resources res; + auto stream = resource::get_cuda_stream(res); + rmm::device_uvector in(1, stream); + rmm::device_uvector out(1, stream); + int host_in = 1; + RAFT_CUDA_TRY(cudaMemcpyAsync(in.data(), &host_in, sizeof(int), cudaMemcpyHostToDevice, stream)); + + raft::launch_kernel(res, 1, 32, copy_restricted_kernel, in.data(), out.data()); + resource::sync_stream(res); + + int host_out = 0; + RAFT_CUDA_TRY(cudaMemcpy(&host_out, out.data(), sizeof(int), cudaMemcpyDeviceToHost)); + EXPECT_EQ(host_out, 1); +} + TEST(KernelLaunch, StreamOverload) { raft::resources res;