diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index f7b80fbf7f..5d7038f712 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -307,6 +307,39 @@ RAFT_KERNEL compute_l2_norms_kernel(const Data_t* data, int dim, DistData_t* l2_ if (lane_id == 0) { l2_norms[list_id] = l2_norm; } } +template +size_t compute_l2_norms_smem_size(size_t dim) +{ + return sizeof(Data_t) * raft::ceildiv(dim, static_cast(raft::warp_size())) * + static_cast(raft::warp_size()); +} + +template +size_t configure_compute_l2_norms_kernel_smem(size_t dim) +{ + size_t smem = compute_l2_norms_smem_size(dim); + + int device_id = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device_id)); + + int max_smem = 0; + RAFT_CUDA_TRY( + cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id)); + RAFT_EXPECTS(smem <= static_cast(max_smem), + "NN_DESCENT L2 norm kernel requires %zu bytes of dynamic shared memory " + "(dim=%zu, sizeof(Data_t)=%zu), but the device supports at most %d bytes " + "per block via opt-in.", + smem, + dim, + sizeof(Data_t), + max_smem); + + auto kernel = compute_l2_norms_kernel; + RAFT_CUDA_TRY(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem))); + return smem; +} + template RAFT_KERNEL convert_copy_kernel(const Src_t* src, Dst_t* dst, size_t n) { @@ -1547,19 +1580,16 @@ void GNND::build(Data_t* data, batch_size, stream); constexpr int TPB = 256; + size_t l2_norms_smem = + needs_l2_norms ? configure_compute_l2_norms_kernel_smem(build_config_.dataset_dim) + : 0; for (auto const& batch : vec_batches) { size_t n_elems = batch.size() * build_config_.dataset_dim; int num_blocks = raft::ceildiv(n_elems, static_cast(TPB)); size_t dst_offset = batch.offset() * build_config_.dataset_dim; if (needs_l2_norms) { // Compute l2 norms on the fp32 batches before they're downcast to fp16. - compute_l2_norms_kernel<<(raft::warp_size())) * - raft::warp_size(), - stream>>>( + compute_l2_norms_kernel<<>>( batch.data(), build_config_.dataset_dim, l2_norms_.data_handle() + batch.offset()); RAFT_CUDA_TRY(cudaPeekAtLastError()); } @@ -1583,13 +1613,8 @@ void GNND::build(Data_t* data, } if (needs_l2_norms && !downcast_host_data) { - compute_l2_norms_kernel<<(raft::warp_size())) * - raft::warp_size(), - stream>>>( + auto smem = configure_compute_l2_norms_kernel_smem(build_config_.dataset_dim); + compute_l2_norms_kernel<<>>( static_cast(d_data_ptr_), build_config_.dataset_dim, l2_norms_.data_handle()); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(res);