Fix launch_kernel handling of restricted and partially specified kernel arguments

cpp/include/raft/sparse/linalg/detail/symmetrize.cuh

@@ -145,7 +145,7 @@ void coo_symmetrize(COO<T, IdxT, nnz_t>* in,
   raft::launch_kernel(stream,
                       grid,
                       blk,
-                      coo_symmetrize_kernel<TPB_X, T>,
+                      coo_symmetrize_kernel<TPB_X, T, Lambda, IdxT, nnz_t>,
                       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<value_idx, value_t>,
                       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<value_idx>,
                       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<value_idx, value_t>,
                       edges,
                       knn_dists,
                       knn_indices,

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<std::remove_cvref_t<Args>...>(
+  // 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<void*>(kernel), where.location, std::forward<Args>(args)...);
 }
 
@@ -222,8 +225,12 @@ requires(sizeof...(Params) == sizeof...(Args) &&
 
   where.config.gridDim  = grid;
   where.config.blockDim = block;
-  detail::dispatch<Params...>(
-    where.config, reinterpret_cast<void*>(kernel), where.location, std::forward<Args>(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<void*>(kernel),
+                   where.location,
+                   static_cast<Params>(std::forward<Args>(args))...);
 }
 
 /** @} */  // end group kernel_launch

cpp/tests/sparse/symmetrize.cu

@@ -212,6 +212,57 @@ TEST_P(COOSymmetrizeView, ResultView)
                                        raft::Compare<float>()));
 }
 
+TEST_P(COOSymmetrizeView, ResultLegacy)
+{
+  raft::resources handle;
+  auto stream = resource::get_cuda_stream(handle);
+
+  raft::sparse::COO<float> in(stream, params.nnz, params.n_rows, params.n_cols, false);
+  raft::sparse::COO<float> 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<int>(out.rows(), params.exp_rows_h.data(), out.nnz, raft::Compare<int>()));
+  ASSERT_TRUE(
+    raft::devArrMatch<int>(out.cols(), params.exp_cols_h.data(), out.nnz, raft::Compare<int>()));
+  ASSERT_TRUE(raft::devArrMatch<float>(
+    out.vals(), params.exp_vals_h.data(), out.nnz, raft::Compare<float>()));
+}
+
+TEST(FromKnnSymmetrizeTest, RestrictedPointerArguments)
+{
+  raft::resources handle;
+  auto stream = resource::get_cuda_stream(handle);
+
+  constexpr int n = 2;
+  constexpr int k = 1;
+  std::vector<int> indices_h{1, 0};
+  std::vector<float> distances_h{0.5f, 0.5f};
+  rmm::device_uvector<int> indices(indices_h.size(), stream);
+  rmm::device_uvector<float> distances(distances_h.size(), stream);
+  raft::sparse::COO<float> 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<COOSymmetrizeInputs<float>> inputsf = {
   // first test fails without fix in #2582
   {

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<int> 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<int> in(1, stream);
+  rmm::device_uvector<int> 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;