From 04058d6e62ff7102a23b6cfde5c0b301f47d568a Mon Sep 17 00:00:00 2001 From: Alex Fallin Date: Fri, 14 Aug 2026 18:10:11 -0700 Subject: [PATCH 1/4] Replace the existing MST solver with SC'23 implementation --- .../raft/sparse/solver/detail/mst_kernels.cuh | 670 +++++++++++------- .../sparse/solver/detail/mst_solver_inl.cuh | 612 +++++++--------- .../raft/sparse/solver/detail/mst_utils.cuh | 25 - cpp/include/raft/sparse/solver/mst.cuh | 21 +- cpp/include/raft/sparse/solver/mst_solver.cuh | 41 +- cpp/tests/sparse/mst.cu | 421 ++++++++++- 6 files changed, 1115 insertions(+), 675 deletions(-) delete mode 100644 cpp/include/raft/sparse/solver/detail/mst_utils.cuh diff --git a/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh b/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh index 9d243c9e66..97eddd913b 100644 --- a/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh +++ b/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh @@ -1,4 +1,3 @@ - /* * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 @@ -7,318 +6,475 @@ #pragma once #include -#include -#include -#include +#include +#include + +/* + * Based on ECL-MST (Fallin, Gonzalez, Seo, Burtscher, SC'23). Per-component + * minima are selected on the packed key (order_key(weight), edge index): + * ties break deterministically by edge index. + */ + +// __CUDA_ARCH_LIST__ is sorted ascending and visible to all compile passes: +// its first element is the min target arch, so mixed fatbins stay consistent. +#if defined(__CUDA_ARCH_LIST__) +#define RAFT_MST_FIRST_ARCH_(a, ...) a +#define RAFT_MST_FIRST_ARCH(...) RAFT_MST_FIRST_ARCH_(__VA_ARGS__) +#define RAFT_MST_MIN_ARCH RAFT_MST_FIRST_ARCH(__CUDA_ARCH_LIST__) +#else +#define RAFT_MST_MIN_ARCH 0 // unknown toolchain: take the portable path +#endif +// RAFT_MST_FORCE_TWOPASS: testing knob, compiles the portable wide path on +// any target. Must be defined consistently across all TUs of a binary (ODR). +#if defined(RAFT_MST_FORCE_TWOPASS) +#define RAFT_MST_HAS_CAS128 0 +#elif RAFT_MST_MIN_ARCH >= 900 +#define RAFT_MST_HAS_CAS128 1 +#else +#define RAFT_MST_HAS_CAS128 0 +#endif +#undef RAFT_MST_FIRST_ARCH_ +#undef RAFT_MST_FIRST_ARCH +#undef RAFT_MST_MIN_ARCH namespace raft { namespace sparse::solver::detail { -template -RAFT_KERNEL kernel_min_edge_per_vertex(const edge_t* offsets, - const vertex_t* indices, - const alteration_t* weights, - const vertex_t* color, - const vertex_t* color_index, - edge_t* new_mst_edge, - const bool* mst_edge, - alteration_t* min_edge_color, - const vertex_t v) +using mst_ull = unsigned long long; +#if RAFT_MST_HAS_CAS128 +using mst_u128 = unsigned __int128; +#endif + +constexpr int mst_block_size = 512; + +// order-preserving unsigned reinterpretation (negatives included) +RAFT_DEVICE_INLINE_FUNCTION uint32_t mst_order_key(float w) { - edge_t tid = threadIdx.x + blockIdx.x * blockDim.x; - - unsigned warp_id = tid / 32; - unsigned lane_id = tid % 32; - - __shared__ edge_t min_edge_index[32]; - __shared__ alteration_t min_edge_weight[32]; - __shared__ vertex_t min_color[32]; - - min_edge_index[lane_id] = std::numeric_limits::max(); - min_edge_weight[lane_id] = std::numeric_limits::max(); - min_color[lane_id] = std::numeric_limits::max(); - - __syncthreads(); - - vertex_t self_color_idx = color_index[warp_id]; - vertex_t self_color = color[self_color_idx]; - - // find the minimum edge associated per row - // each thread in warp holds the minimum edge for - // only the edges that thread scanned - if (warp_id < v) { - // one row is associated with one warp - edge_t row_start = offsets[warp_id]; - edge_t row_end = offsets[warp_id + 1]; - - // assuming one warp per row - // find min for each thread in warp - for (edge_t e = row_start + lane_id; e < row_end; e += 32) { - alteration_t curr_edge_weight = weights[e]; - vertex_t successor_color_idx = color_index[indices[e]]; - vertex_t successor_color = color[successor_color_idx]; - - if (!mst_edge[e] && self_color != successor_color) { - if (curr_edge_weight < min_edge_weight[lane_id]) { - min_color[lane_id] = successor_color; - min_edge_weight[lane_id] = curr_edge_weight; - min_edge_index[lane_id] = e; - } - } - } + uint32_t k = __float_as_uint(w); + return (k & 0x80000000u) ? ~k : (k | 0x80000000u); +} +RAFT_DEVICE_INLINE_FUNCTION uint32_t mst_order_key(int32_t w) +{ + return static_cast(w) ^ 0x80000000u; +} +RAFT_DEVICE_INLINE_FUNCTION mst_ull mst_order_key(double w) +{ + mst_ull k = __double_as_longlong(w); + return (k & 0x8000000000000000ull) ? ~k : (k | 0x8000000000000000ull); +} +RAFT_DEVICE_INLINE_FUNCTION mst_ull mst_order_key(int64_t w) +{ + return static_cast(w) ^ 0x8000000000000000ull; +} + +// Wide worklist entry: 16-byte-aligned {x, y, z, w} of long long. +// CUDA 13 deprecates longlong4 in favor of longlong4_16a +#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 +using mst_entry64 = longlong4_16a; +#else +using mst_entry64 = longlong4; +#endif +static_assert(sizeof(mst_entry64) == 32 && alignof(mst_entry64) == 16, + "mst_entry64 layout differs between toolkit branches"); + +template +struct mst_traits { + static constexpr bool narrow = (sizeof(weight_t) == 4) && (sizeof(edge_t) == 4); + using key_t = std::conditional_t; + using entry_t = std::conditional_t; + using wl_size_t = std::conditional_t; +}; + +// atomics over possibly-signed types; all values here are non-negative +template +RAFT_DEVICE_INLINE_FUNCTION T mst_atomic_cas(T* addr, T compare, T val) +{ + if constexpr (sizeof(T) == 4) { + return static_cast(atomicCAS(reinterpret_cast(addr), + static_cast(compare), + static_cast(val))); + } else { + return static_cast(atomicCAS( + reinterpret_cast(addr), static_cast(compare), static_cast(val))); } - __syncthreads(); - - // reduce across threads in warp - // each thread in warp holds min edge scanned by itself - // reduce across all those warps - for (int offset = 16; offset > 0; offset >>= 1) { - if (lane_id < offset) { - if (min_edge_weight[lane_id] > min_edge_weight[lane_id + offset]) { - min_color[lane_id] = min_color[lane_id + offset]; - min_edge_weight[lane_id] = min_edge_weight[lane_id + offset]; - min_edge_index[lane_id] = min_edge_index[lane_id + offset]; - } - } - __syncthreads(); +} + +template +RAFT_DEVICE_INLINE_FUNCTION T mst_atomic_add(T* addr, T val) +{ + if constexpr (sizeof(T) == 4) { + return static_cast( + atomicAdd(reinterpret_cast(addr), static_cast(val))); + } else { + return static_cast(atomicAdd(reinterpret_cast(addr), static_cast(val))); } +} - // min edge may now be found in first thread - if (lane_id == 0) { - if (min_edge_weight[0] != std::numeric_limits::max()) { - new_mst_edge[warp_id] = min_edge_index[0]; +// Single-instruction word accesses: racing reads can never see a torn value +// (uniform-size races are defined, PTX ISA 8.7.2) and L1 is kept. Do NOT +// replace with atomics: they bypass L1, causing a large perf hit. +template +RAFT_DEVICE_INLINE_FUNCTION T mst_word_load(const T* addr) +{ + if constexpr (sizeof(T) == 4) { + unsigned int r; + asm volatile("ld.b32 %0, [%1];" : "=r"(r) : "l"(addr) : "memory"); + return static_cast(r); + } else { + unsigned long long r; + asm volatile("ld.b64 %0, [%1];" : "=l"(r) : "l"(addr) : "memory"); + return static_cast(r); + } +} - // atomically set min edge per color - // takes care of super vertex case - atomicMin(&min_edge_color[self_color], min_edge_weight[0]); - } +template +RAFT_DEVICE_INLINE_FUNCTION void mst_word_store(T* addr, T val) +{ + if constexpr (sizeof(T) == 4) { + asm volatile("st.b32 [%0], %1;" ::"l"(addr), "r"(static_cast(val)) : "memory"); + } else { + asm volatile("st.b64 [%0], %1;" ::"l"(addr), "l"(static_cast(val)) + : "memory"); } } -template -RAFT_KERNEL min_edge_per_supervertex(const vertex_t* color, - const vertex_t* color_index, - edge_t* new_mst_edge, - bool* mst_edge, - const vertex_t* indices, - const weight_t* weights, - const alteration_t* altered_weights, - vertex_t* temp_src, - vertex_t* temp_dst, - weight_t* temp_weights, - const alteration_t* min_edge_color, - const vertex_t v, - bool symmetrize_output) +// Find with path halving (without it, equal-weight tie chains go quadratic) +template +RAFT_DEVICE_INLINE_FUNCTION vertex_t mst_uf_find(vertex_t curr, vertex_t* const __restrict__ parent) { - auto tid = get_1D_idx(); - if (tid < v) { - vertex_t vertex_color_idx = color_index[tid]; - vertex_t vertex_color = color[vertex_color_idx]; - edge_t edge_idx = new_mst_edge[tid]; - - // check if valid outgoing edge was found - // find minimum edge is same as minimum edge of whole supervertex - // if yes, that is part of mst - if (edge_idx != std::numeric_limits::max()) { - alteration_t vertex_weight = altered_weights[edge_idx]; - - bool add_edge = false; - if (min_edge_color[vertex_color] == vertex_weight) { - add_edge = true; - - auto dst = indices[edge_idx]; - if (!symmetrize_output) { - auto dst_edge_idx = new_mst_edge[dst]; - auto dst_color = color[color_index[dst]]; - - // vertices added each other - // only if destination has found an edge - // the edge points back to source - // the edge is minimum edge found for dst color - if (dst_edge_idx != std::numeric_limits::max() && indices[dst_edge_idx] == tid && - min_edge_color[dst_color] == altered_weights[dst_edge_idx]) { - if (vertex_color > dst_color) { add_edge = false; } - } - } + vertex_t next; + while (curr != (next = mst_word_load(&parent[curr]))) { + const vertex_t grand = mst_word_load(&parent[next]); + if (grand != next) mst_word_store(&parent[curr], grand); + curr = next; + } + return curr; +} - if (add_edge) { - temp_src[tid] = tid; - temp_dst[tid] = dst; - temp_weights[tid] = weights[edge_idx]; - mst_edge[edge_idx] = true; - } - } +template +RAFT_DEVICE_INLINE_FUNCTION void mst_uf_join(vertex_t arep, + vertex_t brep, + vertex_t* const __restrict__ parent) +{ + vertex_t mrep; + do { + mrep = max(arep, brep); + arep = min(arep, brep); + } while ((brep = mst_atomic_cas(&parent[mrep], mrep, arep)) != mrep); +} - if (!add_edge) { new_mst_edge[tid] = std::numeric_limits::max(); } - } - } +RAFT_INLINE_FUNCTION unsigned int mst_sample_hash(unsigned int val) +{ + val = ((val >> 16) ^ val) * 0x45d9f3b; + val = ((val >> 16) ^ val) * 0x45d9f3b; + return (val >> 16) ^ val; } -template -RAFT_KERNEL add_reverse_edge(const edge_t* new_mst_edge, - const vertex_t* indices, - const weight_t* weights, - vertex_t* temp_src, - vertex_t* temp_dst, - weight_t* temp_weights, - const vertex_t v, - bool symmetrize_output) +RAFT_DEVICE_INLINE_FUNCTION long long mst_grid_idx() { - auto tid = get_1D_idx(); + return threadIdx.x + static_cast(blockIdx.x) * mst_block_size; +} - if (tid < v) { - bool reverse_needed = false; +template +RAFT_KERNEL mst_init_parent_kernel(const vertex_t v, + const vertex_t* const __restrict__ color, + vertex_t* const __restrict__ parent) +{ + const long long i = mst_grid_idx(); + if (i < v) parent[i] = color[i]; +} - edge_t edge_idx = new_mst_edge[tid]; - if (edge_idx != std::numeric_limits::max()) { - vertex_t neighbor_vertex = indices[edge_idx]; - edge_t neighbor_edge_idx = new_mst_edge[neighbor_vertex]; +template +RAFT_KERNEL mst_flatten_colors_kernel(const vertex_t v, + vertex_t* const __restrict__ parent, + vertex_t* const __restrict__ color) +{ + const long long i = mst_grid_idx(); + if (i < v) color[i] = mst_uf_find(static_cast(i), parent); +} - // if neighbor picked no vertex then reverse edge is - // definitely needed - if (neighbor_edge_idx == std::numeric_limits::max()) { - reverse_needed = true; - } else { - // check what vertex the neighbor vertex picked - if (symmetrize_output) { - vertex_t neighbor_vertex_neighbor = indices[neighbor_edge_idx]; +template +RAFT_KERNEL mst_sample_keys_kernel( + const int n_samples, + const edge_t e, + const weight_t* const __restrict__ weights, + typename mst_traits::key_t* const __restrict__ keys) +{ + const int i = threadIdx.x + blockIdx.x * blockDim.x; + if (i < n_samples) keys[i] = mst_order_key(weights[mst_sample_hash(i) % e]); +} - // if vertices did not pick each other - // add a reverse edge - if (tid != neighbor_vertex_neighbor) { reverse_needed = true; } - } +// gather flagged CSR edges into COO; optionally emit both directions +template +RAFT_KERNEL mst_extract_coo_kernel(const vertex_t v, + const edge_t e, + const edge_t* const __restrict__ offsets, + const vertex_t* const __restrict__ indices, + const weight_t* const __restrict__ weights, + const bool* const __restrict__ in_mst, + const bool symmetrize, + vertex_t* const __restrict__ out_src, + vertex_t* const __restrict__ out_dst, + weight_t* const __restrict__ out_w, + edge_t* const __restrict__ out_count) +{ + const long long j = mst_grid_idx(); + if (j < e && in_mst[j]) { + // row of edge j by binary search over the offsets + vertex_t lo = 0, hi = v; + while (lo + 1 < hi) { + const vertex_t mid = lo + (hi - lo) / 2; + if (offsets[mid] <= j) { + lo = mid; + } else { + hi = mid; } + } + const edge_t k = mst_atomic_add(out_count, static_cast(symmetrize ? 2 : 1)); + out_src[k] = lo; + out_dst[k] = indices[j]; + out_w[k] = weights[j]; + if (symmetrize) { + out_src[k + 1] = indices[j]; + out_dst[k + 1] = lo; + out_w[k + 1] = weights[j]; + } + } +} - // if reverse was needed, add the edge - if (reverse_needed) { - // it is assumed the each vertex only picks one valid min edge - // per cycle - // hence, we store at index tid + v for the reverse edge scenario - temp_src[tid + v] = neighbor_vertex; - temp_dst[tid + v] = tid; - temp_weights[tid + v] = weights[edge_idx]; +template +RAFT_KERNEL mst_init_worklist_kernel( + typename mst_traits::entry_t* const __restrict__ wl, + typename mst_traits::wl_size_t* const __restrict__ wl_size, + const typename mst_traits::wl_size_t wl_capacity, + const vertex_t v, + const edge_t e, + const edge_t* const __restrict__ offsets, + const vertex_t* const __restrict__ indices, + const weight_t* const __restrict__ weights, + vertex_t* const __restrict__ parent, + const typename mst_traits::key_t thr_key) +{ + const long long j = mst_grid_idx(); + if (j < e) { + const vertex_t n = indices[j]; + // row of edge j by binary search over the offsets + vertex_t lo = 0, hi = v; + while (lo + 1 < hi) { + const vertex_t mid = lo + (hi - lo) / 2; + if (offsets[mid] <= j) { + lo = mid; + } else { + hi = mid; + } + } + const vertex_t r = lo; + if (n > r) { + const typename mst_traits::key_t k = mst_order_key(weights[j]); + if (FIRST ? (k <= thr_key) : (k > thr_key)) { + const vertex_t arep = FIRST ? r : mst_uf_find(r, parent); + const vertex_t brep = FIRST ? n : mst_uf_find(n, parent); + if (FIRST || (arep != brep)) { + using wl_size_t = typename mst_traits::wl_size_t; + const wl_size_t slot = mst_atomic_add(wl_size, static_cast(1)); + // slot >= 0: counter wraparound defense for malformed inputs + if (slot >= 0 && slot < wl_capacity) { + if constexpr (mst_traits::narrow) { + wl[slot] = int4{static_cast(arep), + static_cast(brep), + static_cast(k), + static_cast(j)}; + } else { + wl[slot] = mst_entry64{static_cast(arep), + static_cast(brep), + static_cast(k), + j}; + } + } + } } } } } -// executes for newly added mst edges and updates the colors of both vertices to the lower color -template -RAFT_KERNEL min_pair_colors(const vertex_t v, - const vertex_t* indices, - const edge_t* new_mst_edge, - const vertex_t* color, - const vertex_t* color_index, - vertex_t* next_color) +// ---- narrow path (4-byte weight_t + edge_t): packed 64-bit (key, index) ---- +template +RAFT_KERNEL mst_filter_min_kernel(const int4* const __restrict__ wl1, + const int wl1_size, + int4* const __restrict__ wl2, + int* const __restrict__ wl2_size, + vertex_t* const __restrict__ parent, + volatile mst_ull* const __restrict__ minv, + volatile mst_ull* const __restrict__ minv_prev) { - auto i = get_1D_idx(); - - if (i < v) { - edge_t edge_idx = new_mst_edge[i]; - - if (edge_idx != std::numeric_limits::max()) { - vertex_t neighbor_vertex = indices[edge_idx]; - // vertex_t self_color = color[i]; - vertex_t self_color_idx = color_index[i]; - vertex_t self_color = color[self_color_idx]; - vertex_t neighbor_color_idx = color_index[neighbor_vertex]; - vertex_t neighbor_super_color = color[neighbor_color_idx]; - - // update my own color as source of edge - // update neighbour color index directly - // this will ensure v1 updates supervertex color - // while v2 will update the color of its supervertex - // thus, allowing the colors to progress towards 0 - atomicMin(&next_color[self_color_idx], neighbor_super_color); - atomicMin(&next_color[neighbor_color_idx], self_color); + const int idx = threadIdx.x + blockIdx.x * mst_block_size; + if (idx < wl1_size) { + int4 el = wl1[idx]; + const vertex_t arep = mst_uf_find(static_cast(el.x), parent); + const vertex_t brep = mst_uf_find(static_cast(el.y), parent); + if (arep != brep) { + minv_prev[arep] = ~0ull; // ping-pong reset + minv_prev[brep] = ~0ull; + el.x = arep; + el.y = brep; + wl2[atomicAdd(wl2_size, 1)] = el; + const mst_ull val = + ((static_cast(static_cast(el.z))) << 32) | static_cast(el.w); + if (minv[arep] > val) atomicMin(const_cast(&minv[arep]), val); + if (minv[brep] > val) atomicMin(const_cast(&minv[brep]), val); } } } -// for each vertex, update color if it was changed in min_pair_colors kernel template -RAFT_KERNEL update_colors(const vertex_t v, - vertex_t* color, - const vertex_t* color_index, - const vertex_t* next_color, - bool* done) +RAFT_KERNEL mst_select_join_kernel(const int4* const __restrict__ wl, + const int wl_size, + vertex_t* const __restrict__ parent, + const mst_ull* const __restrict__ minv, + bool* const __restrict__ in_mst) { - auto i = get_1D_idx(); + const int idx = threadIdx.x + blockIdx.x * mst_block_size; + if (idx < wl_size) { + const int4 el = wl[idx]; + const mst_ull val = + ((static_cast(static_cast(el.z))) << 32) | static_cast(el.w); + if ((val == minv[el.x]) || (val == minv[el.y])) { + mst_uf_join(static_cast(el.x), static_cast(el.y), parent); + in_mst[el.w] = true; + } + } +} - if (i < v) { - vertex_t self_color = color[i]; - vertex_t self_color_idx = color_index[i]; - vertex_t new_color = next_color[self_color_idx]; +#if RAFT_MST_HAS_CAS128 +// ---- wide path, sm_90+: single-pass packed 128-bit (key, edge index) ------- +RAFT_DEVICE_INLINE_FUNCTION void mst_atomic_min_u128(mst_u128* const addr, const mst_u128 val) +{ + mst_u128 old = atomicCAS(addr, val, val); + while (old > val) { + const mst_u128 assumed = old; + old = atomicCAS(addr, assumed, val); + if (old == assumed) break; + } +} - // update self color to new smaller color - if (self_color > new_color) { - color[i] = new_color; - *done = false; +template +RAFT_KERNEL mst_filter_min_kernel(const mst_entry64* const __restrict__ wl1, + const wl_size_t wl1_size, + mst_entry64* const __restrict__ wl2, + wl_size_t* const __restrict__ wl2_size, + vertex_t* const __restrict__ parent, + mst_u128* const __restrict__ minv, + mst_u128* const __restrict__ minv_prev) +{ + const long long idx = mst_grid_idx(); + if (idx < wl1_size) { + mst_entry64 el = wl1[idx]; + const vertex_t arep = mst_uf_find(static_cast(el.x), parent); + const vertex_t brep = mst_uf_find(static_cast(el.y), parent); + if (arep != brep) { + minv_prev[arep] = ~static_cast(0); + minv_prev[brep] = ~static_cast(0); + el.x = arep; + el.y = brep; + wl2[mst_atomic_add(wl2_size, static_cast(1))] = el; + const mst_u128 val = + ((static_cast(static_cast(el.z))) << 64) | static_cast(el.w); + const mst_ull key_a = reinterpret_cast(&minv[arep])[1]; + if (key_a >= static_cast(el.z)) mst_atomic_min_u128(&minv[arep], val); + const mst_ull key_b = reinterpret_cast(&minv[brep])[1]; + if (key_b >= static_cast(el.z)) mst_atomic_min_u128(&minv[brep], val); } } } -// point vertices to their final color index -template -RAFT_KERNEL final_color_indices(const vertex_t v, const vertex_t* color, vertex_t* color_index) +template +RAFT_KERNEL mst_select_join_kernel(const mst_entry64* const __restrict__ wl, + const wl_size_t wl_size, + vertex_t* const __restrict__ parent, + const mst_u128* const __restrict__ minv, + bool* const __restrict__ in_mst) { - auto i = get_1D_idx(); - - if (i < v) { - vertex_t self_color_idx = color_index[i]; - vertex_t self_color = color[self_color_idx]; - - // if self color is not equal to self color index, - // it means self is not supervertex - // in which case, iterate until we can find - // parent supervertex - while (self_color_idx != self_color) { - self_color_idx = color_index[self_color]; - self_color = color[self_color_idx]; + const long long idx = mst_grid_idx(); + if (idx < wl_size) { + const mst_entry64 el = wl[idx]; + const mst_u128 val = + ((static_cast(static_cast(el.z))) << 64) | static_cast(el.w); + if ((val == minv[el.x]) || (val == minv[el.y])) { // no concurrent writers + mst_uf_join(static_cast(el.x), static_cast(el.y), parent); + in_mst[el.w] = true; } - - // point to new supervertex - color_index[i] = self_color_idx; } } -// Alterate the weights, make all undirected edge weight unique while keeping Wuv == Wvu -// Consider using curand device API instead of precomputed random_values array -template -RAFT_KERNEL alteration_kernel(const vertex_t v, - const edge_t e, - const edge_t* offsets, - const vertex_t* indices, - const weight_t* weights, - alteration_t max, - alteration_t* random_values, - alteration_t* altered_weights) +#else +// ---- wide path, portable: two-pass min (weight key, then edge index) ------- + +template +RAFT_KERNEL mst_filter_min_kernel(const mst_entry64* const __restrict__ wl1, + const wl_size_t wl1_size, + mst_entry64* const __restrict__ wl2, + wl_size_t* const __restrict__ wl2_size, + vertex_t* const __restrict__ parent, + volatile mst_ull* const __restrict__ minw, + volatile mst_ull* const __restrict__ minw_prev, + volatile mst_ull* const __restrict__ mine_prev) { - auto row = get_1D_idx(); - if (row < v) { - auto row_begin = offsets[row]; - auto row_end = offsets[row + 1]; - for (auto i = row_begin; i < row_end; i++) { - auto column = indices[i]; - altered_weights[i] = weights[i] + max * (random_values[row] + random_values[column]); + const long long idx = mst_grid_idx(); + if (idx < wl1_size) { + mst_entry64 el = wl1[idx]; + const vertex_t arep = mst_uf_find(static_cast(el.x), parent); + const vertex_t brep = mst_uf_find(static_cast(el.y), parent); + if (arep != brep) { + minw_prev[arep] = ~0ull; // ping-pong resets + minw_prev[brep] = ~0ull; + mine_prev[arep] = ~0ull; + mine_prev[brep] = ~0ull; + el.x = arep; + el.y = brep; + wl2[mst_atomic_add(wl2_size, static_cast(1))] = el; + const mst_ull k = static_cast(el.z); + if (minw[arep] > k) atomicMin(const_cast(&minw[arep]), k); + if (minw[brep] > k) atomicMin(const_cast(&minw[brep]), k); } } } -template -RAFT_KERNEL kernel_count_new_mst_edges(const vertex_t* mst_src, - edge_t* mst_edge_count, - const vertex_t v) +// pass 2: min edge index among key-tied edges +template +RAFT_KERNEL mst_min_index_kernel(const entry_t* const __restrict__ wl, + const wl_size_t wl_size, + const mst_ull* const __restrict__ minw, + volatile mst_ull* const __restrict__ mine) { - auto tid = get_1D_idx(); - - // count number of new mst edges added - bool predicate = tid < v && (mst_src[tid] != std::numeric_limits::max()); - vertex_t block_count = __syncthreads_count(predicate); + const long long idx = mst_grid_idx(); + if (idx < wl_size) { + const entry_t el = wl[idx]; + const mst_ull k = static_cast(el.z); + const mst_ull id = static_cast(el.w); + if (k == minw[el.x] && mine[el.x] > id) atomicMin(const_cast(&mine[el.x]), id); + if (k == minw[el.y] && mine[el.y] > id) atomicMin(const_cast(&mine[el.y]), id); + } +} - if (threadIdx.x == 0 && block_count > 0) { atomicAdd(mst_edge_count, block_count); } +template +RAFT_KERNEL mst_select_join_kernel(const mst_entry64* const __restrict__ wl, + const wl_size_t wl_size, + vertex_t* const __restrict__ parent, + const mst_ull* const __restrict__ mine, + bool* const __restrict__ in_mst) +{ + const long long idx = mst_grid_idx(); + if (idx < wl_size) { + const mst_entry64 el = wl[idx]; + const mst_ull id = static_cast(el.w); + if ((id == mine[el.x]) || (id == mine[el.y])) { // edge ids globally unique + mst_uf_join(static_cast(el.x), static_cast(el.y), parent); + in_mst[el.w] = true; + } + } } +#endif // RAFT_MST_HAS_CAS128 } // namespace sparse::solver::detail } // namespace raft diff --git a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh index feaa6c0d21..900055819e 100644 --- a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh +++ b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh @@ -6,161 +6,253 @@ #pragma once #include -#include -#include #include -#include #include #include #include +#include -#include -#include -#include -#include -#include -#include -#include -#include -#include #include -#include -#include -#include -#include -#include - -#include +#include +#include +#include +#include namespace raft { namespace sparse::solver { -// curand generator uniform -inline curandStatus_t curand_generate_uniformX(curandGenerator_t generator, - float* outputPtr, - size_t n) -{ - return curandGenerateUniform(generator, outputPtr, n); -} -inline curandStatus_t curand_generate_uniformX(curandGenerator_t generator, - double* outputPtr, - size_t n) -{ - return curandGenerateUniformDouble(generator, outputPtr, n); -} +namespace detail { -template -MST_solver::MST_solver(raft::resources const& handle_, - const edge_t* offsets_, - const vertex_t* indices_, - const weight_t* weights_, - const vertex_t v_, - const edge_t e_, - vertex_t* color_, - cudaStream_t stream_, - bool symmetrize_output_, - bool initialize_colors_, - int iterations_) - : handle(handle_), - offsets(offsets_), - indices(indices_), - weights(weights_), - altered_weights(e_, stream_), - v(v_), - e(e_), - color_index(color_), - color(v_, stream_), - next_color(v_, stream_), - min_edge_color(v_, stream_), - new_mst_edge(v_, stream_), - mst_edge(e_, stream_), - temp_src(2 * v_, stream_), - temp_dst(2 * v_, stream_), - temp_weights(2 * v_, stream_), - mst_edge_count(1, stream_), - prev_mst_edge_count(1, stream_), - stream(stream_), - symmetrize_output(symmetrize_output_), - initialize_colors(initialize_colors_), - iterations(iterations_) +/* + * MST solver (algorithm: mst_kernels.cuh; user contract: mst.cuh). + * Key order (from float): -0.0 < +0.0; NaNs order by bit pattern. + * Colors = min vertex id per component. + * iterations != 0 disables the edge filter; > 0 bounds the rounds. + */ +template +Graph_COO mst_solve(raft::resources const& handle, + edge_t const* offsets, + vertex_t const* indices, + weight_t const* weights, + vertex_t const v, + edge_t const e, + vertex_t* color, + cudaStream_t stream, + bool symmetrize_output, + bool initialize_colors, + int iterations) { - max_blocks = resource::get_device_properties(handle_).maxGridSize[0]; - max_threads = resource::get_device_properties(handle_).maxThreadsPerBlock; - sm_count = resource::get_device_properties(handle_).multiProcessorCount; - - mst_edge_count.set_value_to_zero_async(stream); - prev_mst_edge_count.set_value_to_zero_async(stream); - RAFT_CUDA_TRY(cudaMemsetAsync(mst_edge.data(), 0, mst_edge.size() * sizeof(bool), stream)); - - // Initially, color holds the vertex id as color - auto policy = resource::get_thrust_policy(handle); - if (initialize_colors_) { - thrust::sequence(policy, color.begin(), color.end(), 0); - thrust::sequence(policy, color_index, color_index + v, 0); - } else { - raft::copy(color.data(), color_index, v, stream); - } - thrust::sequence(policy, next_color.begin(), next_color.end(), 0); -} + static_assert((sizeof(vertex_t) == 4 || sizeof(vertex_t) == 8) && + (sizeof(edge_t) == 4 || sizeof(edge_t) == 8) && + sizeof(vertex_t) <= sizeof(edge_t), + "raft::sparse::solver::mst supports 32- and 64-bit vertex_t/edge_t, with " + "edge_t at least as wide as vertex_t (worklist entries store vertex ids in " + "fields sized by the edge width)"); + constexpr bool narrow = mst_traits::narrow; + using key_t = typename mst_traits::key_t; + using entry_t = typename mst_traits::entry_t; + using wl_size_t = typename mst_traits::wl_size_t; -template -Graph_COO MST_solver::solve() -{ RAFT_EXPECTS(v > 0, "0 vertices"); RAFT_EXPECTS(e > 0, "0 edges"); RAFT_EXPECTS(offsets != nullptr, "Null offsets."); RAFT_EXPECTS(indices != nullptr, "Null indices."); RAFT_EXPECTS(weights != nullptr, "Null weights."); + // narrow packing uses signed 32-bit int4 fields; unsigned 32-bit types + // can exceed them (use a 64-bit edge_t for edge counts above INT_MAX) + if constexpr (narrow && std::is_unsigned_v) { + RAFT_EXPECTS(v <= static_cast(std::numeric_limits::max()), + "unsigned 32-bit vertex ids above INT_MAX are not supported"); + } + if constexpr (narrow && std::is_unsigned_v) { + RAFT_EXPECTS(e <= static_cast(std::numeric_limits::max()), + "unsigned 32-bit edge counts above INT_MAX are not supported"); + } - // Alterating the weights - // this is done by identifying the lowest cost edge weight gap that is not 0, call this theta. - // For each edge, add noise that is less than theta. That is, generate a random number in the - // range [0.0, theta) and add it to each edge weight. - if (e > 1) alteration(); - - auto max_mst_edges = symmetrize_output ? 2 * v - 2 : v - 1; - - Graph_COO mst_result(max_mst_edges, stream); - - // Boruvka original formulation says "while more than 1 supervertex remains" - // Here we adjust it to support disconnected components (spanning forest) - // track completion with mst_edge_found status and v as upper bound - auto mst_iterations = iterations > 0 ? iterations : v; - for (auto i = 0; i < mst_iterations; i++) { - // Finds the minimum edge from each vertex to the lowest color - // by working at each vertex of the supervertex - min_edge_per_vertex(); - - // Finds the minimum edge from each supervertex to the lowest color - min_edge_per_supervertex(); - - // check if msf/mst done, count new edges added - check_termination(); - - auto curr_mst_edge_count = mst_edge_count.value(stream); - RAFT_EXPECTS(curr_mst_edge_count <= max_mst_edges, - "Number of edges found by MST is invalid. This may be due to " - "loss in precision. Try increasing precision of weights."); + // one worklist entry per undirected edge + const wl_size_t wl_capacity = static_cast(e / 2 + 1); + + rmm::device_uvector parent(v, stream); + rmm::device_uvector in_mst(e, stream); + rmm::device_uvector wl1(static_cast(wl_capacity), stream); + rmm::device_uvector wl2(static_cast(wl_capacity), stream); + rmm::device_scalar wl_size_d(stream); + +#if RAFT_MST_HAS_CAS128 + const size_t minv_bytes = narrow ? 2 * static_cast(v) * sizeof(mst_ull) + : 2 * static_cast(v) * sizeof(mst_u128); +#else + const size_t minv_bytes = narrow ? 2 * static_cast(v) * sizeof(mst_ull) + : 4 * static_cast(v) * sizeof(mst_ull); +#endif + rmm::device_uvector minv_raw(minv_bytes, stream); + + const int vblocks = + static_cast((static_cast(v) + mst_block_size - 1) / mst_block_size); + if (initialize_colors) { + thrust::sequence(rmm::exec_policy(stream), parent.begin(), parent.end()); + } else { + mst_init_parent_kernel<<>>(v, color, parent.data()); + } + RAFT_CUDA_TRY(cudaMemsetAsync(minv_raw.data(), 0xFF, minv_bytes, stream)); + RAFT_CUDA_TRY(cudaMemsetAsync(in_mst.data(), 0, e * sizeof(bool), stream)); + + // Two-phase filter: solve a sampled light-edge prefix first, then the + // rest find-filtered. Disabled for bounded solves. Constants are empirical. + constexpr int filter_min_avg_degree = 4; + constexpr double filter_light_edges_per_vertex = 3.0; + constexpr int max_samples = 20; + + key_t thr_key = std::numeric_limits::max(); + bool filtered = false; + if (iterations == 0 && e / v >= filter_min_avg_degree) { + const int ns = static_cast(std::min(e, max_samples)); + rmm::device_uvector keys_d(ns, stream); + mst_sample_keys_kernel<<<1, 32, 0, stream>>>(ns, e, weights, keys_d.data()); + key_t keys[max_samples]; + raft::update_host(keys, keys_d.data(), ns, stream); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); + std::sort(keys, keys + ns); + thr_key = + keys[std::min(max_samples - 1, static_cast(filter_light_edges_per_vertex * v * ns / e))]; + filtered = true; + } - if (curr_mst_edge_count == prev_mst_edge_count.value(stream)) { - // exit here when reaching steady state - break; + int round = 0; + auto boruvka = [&](wl_size_t wl_size) { + entry_t* d1 = wl1.data(); + entry_t* d2 = wl2.data(); + while (wl_size > 0) { + if (iterations > 0 && round >= iterations) break; + wl_size_d.set_value_to_zero_async(stream); + const int wblocks = + static_cast((static_cast(wl_size) + mst_block_size - 1) / mst_block_size); + if constexpr (narrow) { + mst_ull* const base = reinterpret_cast(minv_raw.data()); + mst_ull* const cur = base + (round % 2) * static_cast(v); + mst_ull* const prev = base + ((round + 1) % 2) * static_cast(v); + mst_filter_min_kernel<<>>( + d1, wl_size, d2, wl_size_d.data(), parent.data(), cur, prev); + std::swap(d1, d2); + wl_size = wl_size_d.value(stream); + if (wl_size > 0) { + const int nblocks = static_cast( + (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); + mst_select_join_kernel<<>>( + d1, wl_size, parent.data(), cur, in_mst.data()); + } + } else { +#if RAFT_MST_HAS_CAS128 + mst_u128* const base = reinterpret_cast(minv_raw.data()); + mst_u128* const cur = base + (round % 2) * static_cast(v); + mst_u128* const prev = base + ((round + 1) % 2) * static_cast(v); + mst_filter_min_kernel<<>>( + d1, wl_size, d2, wl_size_d.data(), parent.data(), cur, prev); + std::swap(d1, d2); + wl_size = wl_size_d.value(stream); + if (wl_size > 0) { + const int nblocks = static_cast( + (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); + mst_select_join_kernel<<>>( + d1, wl_size, parent.data(), cur, in_mst.data()); + } +#else + mst_ull* const base = reinterpret_cast(minv_raw.data()); + mst_ull* const minw_cur = base + (round % 2) * static_cast(v); + mst_ull* const minw_prev = base + ((round + 1) % 2) * static_cast(v); + mst_ull* const mine_cur = + base + 2 * static_cast(v) + (round % 2) * static_cast(v); + mst_ull* const mine_prev = + base + 2 * static_cast(v) + ((round + 1) % 2) * static_cast(v); + mst_filter_min_kernel<<>>( + d1, wl_size, d2, wl_size_d.data(), parent.data(), minw_cur, minw_prev, mine_prev); + std::swap(d1, d2); + wl_size = wl_size_d.value(stream); + if (wl_size > 0) { + const int nblocks = static_cast( + (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); + mst_min_index_kernel<<>>( + d1, wl_size, minw_cur, mine_cur); + mst_select_join_kernel<<>>( + d1, wl_size, parent.data(), mine_cur, in_mst.data()); + } +#endif + } + round++; } - - // append the newly found MST edges to the final output - append_src_dst_pair(mst_result.src.data(), mst_result.dst.data(), mst_result.weights.data()); - - // updates colors of vertices by propagating the lower color to the higher - label_prop(mst_result.src.data(), mst_result.dst.data()); - - // copy this iteration's results and store - prev_mst_edge_count.set_value_async(curr_mst_edge_count, stream); + }; + + const int eblocks = + static_cast((static_cast(e) + mst_block_size - 1) / mst_block_size); + auto launch_init = [&](bool first) { + wl_size_d.set_value_to_zero_async(stream); + if (first) { + mst_init_worklist_kernel<<>>(wl1.data(), + wl_size_d.data(), + wl_capacity, + v, + e, + offsets, + indices, + weights, + parent.data(), + thr_key); + } else { + mst_init_worklist_kernel<<>>(wl1.data(), + wl_size_d.data(), + wl_capacity, + v, + e, + offsets, + indices, + weights, + parent.data(), + thr_key); + } + const wl_size_t wl_size = wl_size_d.value(stream); + // wl_size < 0 means the admission counter wrapped (malformed input) + RAFT_EXPECTS(wl_size >= 0 && wl_size <= wl_capacity, + "MST worklist overflow: the input CSR must be symmetric (each " + "undirected edge stored in both directions)."); + return wl_size; + }; + + boruvka(launch_init(true)); + + if (filtered) { + // clear straggler minima before admitting the remaining edges + RAFT_CUDA_TRY(cudaMemsetAsync(minv_raw.data(), 0xFF, minv_bytes, stream)); + boruvka(launch_init(false)); } - // result packaging - mst_result.n_edges = mst_edge_count.value(stream); + mst_flatten_colors_kernel<<>>(v, parent.data(), color); + + // symmetrized count can exceed 32-bit edge_t/vertex_t for v > 2^30: + // fail loudly rather than under-allocate + const int64_t max_out_wide = + symmetrize_output ? 2 * (static_cast(v) - 1) : (static_cast(v) - 1); + RAFT_EXPECTS(max_out_wide <= std::numeric_limits::max() && + max_out_wide <= std::numeric_limits::max(), + "MST output edge count can exceed edge_t/vertex_t; use fewer vertices or " + "symmetrize_output=false"); + const edge_t max_out = static_cast(max_out_wide); + Graph_COO mst_result(std::max(max_out, 1), stream); + rmm::device_scalar out_count(stream); + out_count.set_value_to_zero_async(stream); + mst_extract_coo_kernel<<>>(v, + e, + offsets, + indices, + weights, + in_mst.data(), + symmetrize_output, + mst_result.src.data(), + mst_result.dst.data(), + mst_result.weights.data(), + out_count.data()); + mst_result.n_edges = out_count.value(stream); mst_result.src.resize(mst_result.n_edges, stream); mst_result.dst.resize(mst_result.n_edges, stream); mst_result.weights.resize(mst_result.n_edges, stream); @@ -168,239 +260,49 @@ Graph_COO MST_solver -struct alteration_functor { - __host__ __device__ weight_t operator()(const cuda::std::tuple& t) - { - auto x = cuda::std::get<0>(t); - auto y = cuda::std::get<1>(t); - x = x < 0 ? -x : x; - y = y < 0 ? -y : y; - return x < y ? y - x : x - y; - } -}; - -// Compute the uper bound for the alteration -template -alteration_t MST_solver::alteration_max() -{ - auto policy = resource::get_thrust_policy(handle); - rmm::device_uvector tmp(e, stream); - thrust::device_ptr weights_ptr(weights); - thrust::copy(policy, weights_ptr, weights_ptr + e, tmp.begin()); - // sort tmp weights - thrust::sort(policy, tmp.begin(), tmp.end()); - - // remove duplicates - auto new_end = thrust::unique(policy, tmp.begin(), tmp.end()); - - // min(a[i+1]-a[i])/2 - auto begin = thrust::make_zip_iterator(cuda::std::make_tuple(tmp.begin(), tmp.begin() + 1)); - auto end = thrust::make_zip_iterator(cuda::std::make_tuple(new_end - 1, new_end)); - auto init = tmp.element(1, stream) - tmp.element(0, stream); - auto max = thrust::transform_reduce( - policy, begin, end, alteration_functor(), init, cuda::minimum()); - // Enforce distinct weights if initial edge weights are identical by returning - // a value of 1 - return max > 0 ? max / static_cast(2) : 1; -} - -// Compute the alteration to make all undirected edge weight unique -// Preserves weights order -template -void MST_solver::alteration() -{ - auto nthreads = std::min(v, max_threads); - auto nblocks = std::min((v + nthreads - 1) / nthreads, max_blocks); - - // maximum alteration that does not change relative weights order - // Note: The relative weights order will be altered if initial edge weights are identical - alteration_t max = alteration_max(); - - // pool of rand values - rmm::device_uvector rand_values(v, stream); - - // Random number generator - curandGenerator_t randGen; - curandCreateGenerator(&randGen, CURAND_RNG_PSEUDO_DEFAULT); - curandSetPseudoRandomGeneratorSeed(randGen, 1234567); - - // Initialize rand values - auto curand_status = curand_generate_uniformX(randGen, rand_values.data(), v); - RAFT_EXPECTS(curand_status == CURAND_STATUS_SUCCESS, "MST: CURAND failed"); - curand_status = curandDestroyGenerator(randGen); - RAFT_EXPECTS(curand_status == CURAND_STATUS_SUCCESS, "MST: CURAND cleanup failed"); - - // Alterate the weights, make all undirected edge weight unique while keeping Wuv == Wvu - detail::alteration_kernel<<>>( - v, e, offsets, indices, weights, max, rand_values.data(), altered_weights.data()); -} - -// updates colors of vertices by propagating the lower color to the higher -template -void MST_solver::label_prop(vertex_t* mst_src, - vertex_t* mst_dst) -{ - // update the colors of both ends its until there is no change in colors - edge_t curr_mst_edge_count = mst_edge_count.value(stream); - - auto min_pair_nthreads = std::min(v, (vertex_t)max_threads); - auto min_pair_nblocks = - std::min((v + min_pair_nthreads - 1) / min_pair_nthreads, (vertex_t)max_blocks); - - edge_t* new_mst_edge_ptr = new_mst_edge.data(); - vertex_t* color_ptr = color.data(); - vertex_t* next_color_ptr = next_color.data(); - - rmm::device_scalar done(stream); - done.set_value_to_zero_async(stream); - bool* done_ptr = done.data(); - const bool true_val = true; - - auto i = 0; - while (!done.value(stream)) { - done.set_value_async(true_val, stream); +} // namespace detail - detail::min_pair_colors<<>>( - v, indices, new_mst_edge_ptr, color_ptr, color_index, next_color_ptr); - - detail::update_colors<<>>( - v, color_ptr, color_index, next_color_ptr, done_ptr); - i++; - } - - detail::final_color_indices<<>>( - v, color_ptr, color_index); -} - -// Finds the minimum edge from each vertex to the lowest color -template -void MST_solver::min_edge_per_vertex() -{ - auto policy = resource::get_thrust_policy(handle); - thrust::fill( - policy, min_edge_color.begin(), min_edge_color.end(), std::numeric_limits::max()); - thrust::fill( - policy, new_mst_edge.begin(), new_mst_edge.end(), std::numeric_limits::max()); - - int n_threads = 32; - - vertex_t* color_ptr = color.data(); - edge_t* new_mst_edge_ptr = new_mst_edge.data(); - bool* mst_edge_ptr = mst_edge.data(); - alteration_t* min_edge_color_ptr = min_edge_color.data(); - alteration_t* altered_weights_ptr = altered_weights.data(); - - detail::kernel_min_edge_per_vertex<<>>(offsets, - indices, - altered_weights_ptr, - color_ptr, - color_index, - new_mst_edge_ptr, - mst_edge_ptr, - min_edge_color_ptr, - v); -} - -// Finds the minimum edge from each supervertex to the lowest color template -void MST_solver::min_edge_per_supervertex() +MST_solver::MST_solver(raft::resources const& handle_, + const edge_t* offsets_, + const vertex_t* indices_, + const weight_t* weights_, + const vertex_t v_, + const edge_t e_, + vertex_t* color_, + cudaStream_t stream_, + bool symmetrize_output_, + bool initialize_colors_, + int iterations_) + : handle(handle_), + stream(stream_), + symmetrize_output(symmetrize_output_), + initialize_colors(initialize_colors_), + iterations(iterations_), + offsets(offsets_), + indices(indices_), + weights(weights_), + v(v_), + e(e_), + color_index(color_) { - auto nthreads = std::min(v, max_threads); - auto nblocks = std::min((v + nthreads - 1) / nthreads, max_blocks); - - auto policy = resource::get_thrust_policy(handle); - thrust::fill(policy, temp_src.begin(), temp_src.end(), std::numeric_limits::max()); - - vertex_t* color_ptr = color.data(); - edge_t* new_mst_edge_ptr = new_mst_edge.data(); - bool* mst_edge_ptr = mst_edge.data(); - alteration_t* min_edge_color_ptr = min_edge_color.data(); - alteration_t* altered_weights_ptr = altered_weights.data(); - vertex_t* temp_src_ptr = temp_src.data(); - vertex_t* temp_dst_ptr = temp_dst.data(); - weight_t* temp_weights_ptr = temp_weights.data(); - - detail::min_edge_per_supervertex<<>>(color_ptr, - color_index, - new_mst_edge_ptr, - mst_edge_ptr, - indices, - weights, - altered_weights_ptr, - temp_src_ptr, - temp_dst_ptr, - temp_weights_ptr, - min_edge_color_ptr, - v, - symmetrize_output); - - // the above kernel only adds directed mst edges in the case where - // a pair of vertices don't pick the same min edge between them - // so, now we add the reverse edge to make it undirected - if (symmetrize_output) { - detail::add_reverse_edge<<>>(new_mst_edge_ptr, - indices, - weights, - temp_src_ptr, - temp_dst_ptr, - temp_weights_ptr, - v, - symmetrize_output); - } } template -void MST_solver::check_termination() +Graph_COO MST_solver::solve() { - vertex_t nthreads = std::min(2 * v, (vertex_t)max_threads); - vertex_t nblocks = std::min((2 * v + nthreads - 1) / nthreads, (vertex_t)max_blocks); - - // count number of new mst edges - edge_t* mst_edge_count_ptr = mst_edge_count.data(); - vertex_t* temp_src_ptr = temp_src.data(); - - detail::kernel_count_new_mst_edges<<>>( - temp_src_ptr, mst_edge_count_ptr, 2 * v); + return detail::mst_solve(handle, + offsets, + indices, + weights, + v, + e, + color_index, + stream, + symmetrize_output, + initialize_colors, + iterations); } -template -struct new_edges_functor { - __host__ __device__ bool operator()(const cuda::std::tuple& t) - { - auto src = cuda::std::get<0>(t); - - return src != std::numeric_limits::max() ? true : false; - } -}; - -template -void MST_solver::append_src_dst_pair( - vertex_t* mst_src, vertex_t* mst_dst, weight_t* mst_weights) -{ - auto policy = resource::get_thrust_policy(handle); - - edge_t curr_mst_edge_count = prev_mst_edge_count.value(stream); - - // iterator to end of mst edges added to final output in previous iteration - auto src_dst_zip_end = - thrust::make_zip_iterator(cuda::std::make_tuple(mst_src + curr_mst_edge_count, - mst_dst + curr_mst_edge_count, - mst_weights + curr_mst_edge_count)); - - // iterator to new mst edges found - auto temp_src_dst_zip_begin = thrust::make_zip_iterator( - cuda::std::make_tuple(temp_src.begin(), temp_dst.begin(), temp_weights.begin())); - auto temp_src_dst_zip_end = thrust::make_zip_iterator( - cuda::std::make_tuple(temp_src.end(), temp_dst.end(), temp_weights.end())); - - // copy new mst edges to final output - thrust::copy_if(policy, - temp_src_dst_zip_begin, - temp_src_dst_zip_end, - src_dst_zip_end, - new_edges_functor()); -} } // namespace sparse::solver } // namespace raft diff --git a/cpp/include/raft/sparse/solver/detail/mst_utils.cuh b/cpp/include/raft/sparse/solver/detail/mst_utils.cuh deleted file mode 100644 index bf1520710a..0000000000 --- a/cpp/include/raft/sparse/solver/detail/mst_utils.cuh +++ /dev/null @@ -1,25 +0,0 @@ - -/* - * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include - -#include - -#include - -namespace raft { -namespace sparse::solver::detail { - -template -__device__ idx_t get_1D_idx() -{ - return blockIdx.x * blockDim.x + threadIdx.x; -} - -} // namespace sparse::solver::detail -} // namespace raft diff --git a/cpp/include/raft/sparse/solver/mst.cuh b/cpp/include/raft/sparse/solver/mst.cuh index 10801482cb..9f5aacb68e 100644 --- a/cpp/include/raft/sparse/solver/mst.cuh +++ b/cpp/include/raft/sparse/solver/mst.cuh @@ -14,23 +14,32 @@ namespace sparse::solver { /** * Compute the minimum spanning tree (MST) or minimum spanning forest (MSF) depending on * the connected components of the given graph. + * Algorithm based on ECL-MST (Fallin, Gonzalez, Seo, Burtscher, SC'23). * * @tparam vertex_t integral type for precision of vertex indexing * @tparam edge_t integral type for precision of edge indexing * @tparam weight_t type of weights array - * @tparam alteration_t type to use for random alteration + * @tparam alteration_t unused; retained for source compatibility (the solver + * breaks weight ties deterministically by edge index and no longer alters + * weights) * * @param handle - * @param offsets csr inptr array of row offsets (size v+1) - * @param indices csr array of column indices (size e) + * @param offsets csr indptr array of row offsets (size v+1, symmetric input required: each + * undirected edge must be stored in both directions with equal weights) + * @param indices csr array of column indices (size e, each in [0, v)) * @param weights csr array of weights (size e) * @param v number of vertices in graph * @param e number of edges in graph - * @param color array to store resulting colors for MSF + * @param color array to store resulting colors for MSF; when initialize_colors is false it is + * also the input seeding and must hold a valid component labeling from a previous solve of the same + * graph * @param stream cuda stream for ordering operations - * @param symmetrize_output should the resulting output edge list should be symmetrized? + * @param symmetrize_output should the resulting output edge list be symmetrized? * @param initialize_colors should the colors array be initialized inside the MST? - * @param iterations maximum number of iterations to perform + * @param iterations maximum number of Boruvka rounds to perform (values <= 0 solve to + * completion). + * Bounded solves run textbook Boruvka rounds (the internal two-phase edge filter is disabled), + * so partial results are deterministic and can be resumed exactly via initialize_colors=false * @return a list of edges containing the mst (or a subset of the edges guaranteed to be in the mst * when an msf is encountered) */ diff --git a/cpp/include/raft/sparse/solver/mst_solver.cuh b/cpp/include/raft/sparse/solver/mst_solver.cuh index 5ce7917569..c8b3207bd2 100644 --- a/cpp/include/raft/sparse/solver/mst_solver.cuh +++ b/cpp/include/raft/sparse/solver/mst_solver.cuh @@ -9,7 +9,6 @@ #include #include -#include #include namespace raft { @@ -28,6 +27,18 @@ struct Graph_COO { } }; +/** + * @brief MST solver based on ECL-MST, with deterministic lexicographic + * (weight, edge index) tie-breaking. + * + * @tparam vertex_t integral type for vertex indexing (32- or 64-bit) + * @tparam edge_t integral type for edge indexing (32- or 64-bit, at least as + * wide as vertex_t) + * @tparam weight_t type of the weights array + * @tparam alteration_t unused; retained for source compatibility with + * existing callers (the solver no longer perturbs ("alters") weights to break + * ties) + */ template class MST_solver { public: @@ -60,33 +71,7 @@ class MST_solver { const vertex_t v; const edge_t e; - vertex_t max_blocks; - vertex_t max_threads; - vertex_t sm_count; - - vertex_t* color_index; // represent each supervertex as a color - rmm::device_uvector min_edge_color; // minimum incident edge weight per color - rmm::device_uvector new_mst_edge; // new minimum edge per vertex - rmm::device_uvector altered_weights; // weights to be used for mst - rmm::device_scalar mst_edge_count; // total number of edges added after every iteration - rmm::device_scalar - prev_mst_edge_count; // total number of edges up to the previous iteration - rmm::device_uvector mst_edge; // mst output - true if the edge belongs in mst - rmm::device_uvector next_color; // next iteration color - rmm::device_uvector color; // index of color that vertex points to - - // new src-dst pairs found per iteration - rmm::device_uvector temp_src; - rmm::device_uvector temp_dst; - rmm::device_uvector temp_weights; - - void label_prop(vertex_t* mst_src, vertex_t* mst_dst); - void min_edge_per_vertex(); - void min_edge_per_supervertex(); - void check_termination(); - void alteration(); - alteration_t alteration_max(); - void append_src_dst_pair(vertex_t* mst_src, vertex_t* mst_dst, weight_t* mst_weights); + vertex_t* color_index; // user-provided output colors array (one color per vertex) }; } // namespace sparse::solver diff --git a/cpp/tests/sparse/mst.cu b/cpp/tests/sparse/mst.cu index 9bf1c40a3b..f4968f7006 100644 --- a/cpp/tests/sparse/mst.cu +++ b/cpp/tests/sparse/mst.cu @@ -17,11 +17,17 @@ #include #include -#include #include +#include #include #include +#include +#include +#include +#include +#include +#include #include template @@ -55,7 +61,6 @@ weight_t prims(CSRHost& csr_h) std::size_t n_vertices = csr_h.offsets.size() - 1; bool active_vertex[n_vertices]; - // bool mst_set[csr_h.n_edges]; weight_t curr_edge[n_vertices]; for (std::size_t i = 0; i < n_vertices; i++) { @@ -108,6 +113,188 @@ weight_t prims(CSRHost& csr_h) return total_weight; } +// Build a symmetric CSR from an undirected edge list +template +CSRHost csr_from_undirected_edges( + vertex_t v, const std::vector>& edges) +{ + std::vector>> adj(v); + for (auto& [s, d, w] : edges) { + adj[s].push_back({d, w}); + adj[d].push_back({s, w}); + } + CSRHost csr; + csr.offsets.push_back(0); + for (vertex_t i = 0; i < v; i++) { + std::sort(adj[i].begin(), adj[i].end()); + for (auto& [d, w] : adj[i]) { + csr.indices.push_back(d); + csr.weights.push_back(w); + } + csr.offsets.push_back(static_cast(csr.indices.size())); + } + return csr; +} + +// Kruskal oracle with the solver's tie-break (weight, then CSR edge index); +// also produces the expected MSF colors (min vertex id per component). +// Weights must not contain NaN (the sort comparator requires a total order). +struct KruskalResult { + double weight; + int n_edges; + std::vector colors; +}; + +template +KruskalResult kruskal_mst(const CSRHost& csr_h) +{ + const vertex_t v = static_cast(csr_h.offsets.size() - 1); + const edge_t e = static_cast(csr_h.indices.size()); + + std::vector row_of(e); + std::vector order; + for (vertex_t r = 0; r < v; r++) { + for (edge_t j = csr_h.offsets[r]; j < csr_h.offsets[r + 1]; j++) { + row_of[j] = r; + if (csr_h.indices[j] > r) order.push_back(j); // canonical direction, skips self-loops + } + } + std::sort(order.begin(), order.end(), [&](edge_t a, edge_t b) { + if (csr_h.weights[a] != csr_h.weights[b]) return csr_h.weights[a] < csr_h.weights[b]; + return a < b; + }); + + std::vector parent(v); + std::iota(parent.begin(), parent.end(), 0); + auto find = [&](vertex_t x) { + while (parent[x] != x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; + + KruskalResult res{0.0, 0, {}}; + for (edge_t j : order) { + vertex_t a = find(row_of[j]); + vertex_t b = find(csr_h.indices[j]); + if (a != b) { + parent[std::max(a, b)] = std::min(a, b); // root of every set is its min vertex id + res.weight += static_cast(csr_h.weights[j]); + res.n_edges++; + } + } + res.colors.resize(v); + for (vertex_t i = 0; i < v; i++) { + res.colors[i] = find(i); + } + return res; +} + +template +struct MSTResultHost { + std::vector src, dst; + std::vector weights; + std::vector colors; + int n_edges; + + double total_weight() const + { + double sum = 0.0; + for (int i = 0; i < n_edges; i++) { + sum += static_cast(weights[i]); + } + return sum; + } + + std::vector> sorted_edges() const + { + std::vector> out; + for (int i = 0; i < n_edges; i++) { + out.push_back({src[i], dst[i], weights[i]}); + } + std::sort(out.begin(), out.end()); + return out; + } +}; + +// colors_in seeds resume runs (initialize_colors = false) +template +MSTResultHost run_mst_gpu(raft::resources const& handle, + const CSRHost& csr_h, + bool symmetrize_output, + bool initialize_colors, + int iterations, + const std::vector* colors_in = nullptr) +{ + auto stream = resource::get_cuda_stream(handle); + const int v = static_cast(csr_h.offsets.size() - 1); + const int e = static_cast(csr_h.indices.size()); + + rmm::device_uvector offsets_d(v + 1, stream); + rmm::device_uvector indices_d(e, stream); + rmm::device_uvector weights_d(e, stream); + rmm::device_uvector colors_d(v, stream); + raft::update_device(offsets_d.data(), csr_h.offsets.data(), v + 1, stream); + raft::update_device(indices_d.data(), csr_h.indices.data(), e, stream); + raft::update_device(weights_d.data(), csr_h.weights.data(), e, stream); + if (colors_in != nullptr) { raft::update_device(colors_d.data(), colors_in->data(), v, stream); } + + raft::sparse::solver::MST_solver solver(handle, + offsets_d.data(), + indices_d.data(), + weights_d.data(), + v, + e, + colors_d.data(), + stream, + symmetrize_output, + initialize_colors, + iterations); + auto result = solver.solve(); + + MSTResultHost out; + out.n_edges = result.n_edges; + out.src.resize(result.n_edges); + out.dst.resize(result.n_edges); + out.weights.resize(result.n_edges); + out.colors.resize(v); + raft::update_host(out.src.data(), result.src.data(), result.n_edges, stream); + raft::update_host(out.dst.data(), result.dst.data(), result.n_edges, stream); + raft::update_host(out.weights.data(), result.weights.data(), result.n_edges, stream); + raft::update_host(out.colors.data(), colors_d.data(), v, stream); + resource::sync_stream(handle, stream); + return out; +} + +// Connected random graph: ring + deterministic pseudo-random chords +template +CSRHost make_ring_plus_chords(int v, + int chords_per_vertex, + weight_fn_t weight_of) +{ + std::vector> edges; + std::set> seen; + std::mt19937 rng(42); + auto add = [&](int a, int b) { + if (a == b) return; + auto p = std::minmax(a, b); + if (seen.insert({p.first, p.second}).second) { + edges.push_back({p.first, p.second, weight_of(static_cast(edges.size()))}); + } + }; + for (int i = 0; i < v; i++) { + add(i, (i + 1) % v); + } + std::uniform_int_distribution pick(0, v - 1); + for (int i = 0; i < v; i++) { + for (int c = 0; c < chords_per_vertex; c++) { + add(i, pick(rng)); + } + } + return csr_from_undirected_edges(v, edges); +} + template class MSTTest : public ::testing::TestWithParam> { protected: @@ -366,8 +553,6 @@ TEST_P(MSTTestSequential, Sequential) auto& symmetric_result = results_pair.first; auto& non_symmetric_result = results_pair.second; - // do assertions here - // in this case, running sequential MST auto prims_result = prims(mst_input.csr_h); auto symmetric_sum = thrust::reduce(thrust::device, @@ -384,5 +569,233 @@ TEST_P(MSTTestSequential, Sequential) INSTANTIATE_TEST_SUITE_P(MSTTests, MSTTestSequential, ::testing::ValuesIn(csr_in_h)); +void expect_forest(int v, const std::vector& src, const std::vector& dst, int n_edges) +{ + std::vector parent(v); + std::iota(parent.begin(), parent.end(), 0); + auto find = [&](int x) { + while (parent[x] != x) { + parent[x] = parent[parent[x]]; + x = parent[x]; + } + return x; + }; + for (int i = 0; i < n_edges; i++) { + int a = find(src[i]); + int b = find(dst[i]); + ASSERT_NE(a, b) << "cycle introduced by edge " << src[i] << " -> " << dst[i]; + parent[std::max(a, b)] = std::min(a, b); + } +} + +// Regression: heavily-tied integer-valued float weights at ~1e7, where the +// previous solver's alteration rounded away (float ulp there is 1-2) and +// produced cycles or non-minimal forests past its own guard. +TEST(MST, FloatMagnitudeTies) +{ + raft::resources handle; + const int v = 4096; + auto csr_h = make_ring_plus_chords( + v, 4, [](int idx) { return 1.0e7f + static_cast(idx % 2); }); + auto truth = kruskal_mst(csr_h); + + auto result = run_mst_gpu(handle, csr_h, false, true, 0); + + ASSERT_EQ(truth.n_edges, result.n_edges); + ASSERT_EQ(v - 1, result.n_edges); // ring keeps it connected + // integer-valued weights, sum < 2^53: exact compare + ASSERT_EQ(truth.weight, result.total_weight()); + expect_forest(v, result.src, result.dst, result.n_edges); + for (int i = 0; i < v; i++) { + ASSERT_EQ(0, result.colors[i]); // connected: min vertex id everywhere + } +} + +// symmetrize_output must emit every edge in both directions and agree with +// the non-symmetrized solve +TEST(MST, SymmetrizeOutput) +{ + raft::resources handle; + const int v = 1024; + auto csr_h = make_ring_plus_chords( + v, 4, [v](int idx) { return static_cast(1 + (idx * 7919) % v); }); + auto truth = kruskal_mst(csr_h); + + auto sym = run_mst_gpu(handle, csr_h, true, true, 0); + auto non_sym = run_mst_gpu(handle, csr_h, false, true, 0); + + ASSERT_EQ(non_sym.n_edges, truth.n_edges); + ASSERT_EQ(sym.n_edges, 2 * truth.n_edges); + ASSERT_EQ(truth.weight, non_sym.total_weight()); + ASSERT_EQ(2 * truth.weight, sym.total_weight()); + + std::vector> expected; + for (int i = 0; i < non_sym.n_edges; i++) { + expected.push_back({non_sym.src[i], non_sym.dst[i], non_sym.weights[i]}); + expected.push_back({non_sym.dst[i], non_sym.src[i], non_sym.weights[i]}); + } + std::sort(expected.begin(), expected.end()); + ASSERT_EQ(expected, sym.sorted_edges()); +} + +// bounded solve + resume must reproduce the one-shot result exactly +TEST(MST, ResumeMatchesOneShot) +{ + raft::resources handle; + const int v = 2048; + auto csr_h = make_ring_plus_chords( + v, 4, [v](int idx) { return static_cast(1 + (idx * 7919) % v); }); + + auto one_shot = run_mst_gpu(handle, csr_h, false, true, 0); + + auto partial = run_mst_gpu(handle, csr_h, false, true, 2); + ASSERT_LT(partial.n_edges, one_shot.n_edges); // 2 rounds must not finish this graph + auto resumed = run_mst_gpu(handle, csr_h, false, false, 0, &partial.colors); + + ASSERT_EQ(one_shot.n_edges, partial.n_edges + resumed.n_edges); + auto combined = partial.sorted_edges(); + auto rest = resumed.sorted_edges(); + combined.insert(combined.end(), rest.begin(), rest.end()); + std::sort(combined.begin(), combined.end()); + ASSERT_EQ(one_shot.sorted_edges(), combined); + ASSERT_EQ(one_shot.colors, resumed.colors); +} + +// disconnected MSF: n_edges = v - #components, colors = min id per component +TEST(MST, DisconnectedColors) +{ + raft::resources handle; + for (const auto& csr_h : {csr_in4_h[0], csr_in5_h[0]}) { + auto truth = kruskal_mst(csr_h); + auto result = run_mst_gpu(handle, csr_h, false, true, 0); + + ASSERT_EQ(truth.n_edges, result.n_edges); + ASSERT_EQ(truth.weight, result.total_weight()); + ASSERT_EQ(truth.colors, result.colors); + expect_forest( + static_cast(csr_h.offsets.size() - 1), result.src, result.dst, result.n_edges); + } +} + +// Regression: an equal-weight path graph chain-joins into an O(v)-deep +// parent chain; quadratic (minutes at v=1M) without path halving. +TEST(MST, UniformWeightPathGraph) +{ + raft::resources handle; + const int v = 1000000; + std::vector> edges; + for (int i = 0; i + 1 < v; i++) { + edges.push_back({i, i + 1, 1.0f}); + } + auto csr_h = csr_from_undirected_edges(v, edges); + + auto result = run_mst_gpu(handle, csr_h, false, true, 0); + + ASSERT_EQ(v - 1, result.n_edges); + ASSERT_EQ(static_cast(v - 1), result.total_weight()); + for (int i = 0; i < v; i++) { + ASSERT_EQ(0, result.colors[i]); + } +} + +// double weights exercise the wide (8-byte key) kernel path +TEST(MST, DoubleWeights) +{ + raft::resources handle; + const int v = 2048; + auto csr_h = make_ring_plus_chords( + v, 4, [v](int idx) { return static_cast(1 + (idx * 7919) % v); }); + auto truth = kruskal_mst(csr_h); + + auto result = run_mst_gpu(handle, csr_h, false, true, 0); + + ASSERT_EQ(truth.n_edges, result.n_edges); + ASSERT_EQ(truth.weight, result.total_weight()); + expect_forest(v, result.src, result.dst, result.n_edges); +} + +// int32 weights exercise the integer order-key overload +TEST(MST, IntegerWeights) +{ + raft::resources handle; + const int v = 1024; + auto csr_f = make_ring_plus_chords( + v, 4, [v](int idx) { return static_cast(1 + (idx * 7919) % v); }); + auto truth = kruskal_mst(csr_f); + + CSRHost csr_i; + csr_i.offsets = csr_f.offsets; + csr_i.indices = csr_f.indices; + csr_i.weights.assign(csr_f.weights.begin(), csr_f.weights.end()); + auto result = run_mst_gpu(handle, csr_i, false, true, 0); + + ASSERT_EQ(truth.n_edges, result.n_edges); + ASSERT_EQ(truth.weight, result.total_weight()); + ASSERT_EQ(truth.colors, result.colors); +} + +// wide, mixed, and unsigned instantiations must match the 32-bit forest exactly +TEST(MST, Int64Indices) +{ + raft::resources handle; + auto stream = resource::get_cuda_stream(handle); + const int v = 4096; + auto csr32 = make_ring_plus_chords( + v, 4, [v](int idx) { return static_cast(1 + (idx * 7919) % v); }); + auto truth = kruskal_mst(csr32); + const int e = static_cast(csr32.indices.size()); + + auto run_typed = [&](auto vertex_tag, auto edge_tag, auto weight_tag) { + using vertex2_t = decltype(vertex_tag); + using edge2_t = decltype(edge_tag); + using weight_t = decltype(weight_tag); + std::vector off_h(csr32.offsets.begin(), csr32.offsets.end()); + std::vector ind_h(csr32.indices.begin(), csr32.indices.end()); + std::vector w_h(csr32.weights.begin(), csr32.weights.end()); + + rmm::device_uvector off_d(v + 1, stream); + rmm::device_uvector ind_d(e, stream); + rmm::device_uvector w_d(e, stream); + rmm::device_uvector col_d(v, stream); + raft::update_device(off_d.data(), off_h.data(), v + 1, stream); + raft::update_device(ind_d.data(), ind_h.data(), e, stream); + raft::update_device(w_d.data(), w_h.data(), e, stream); + + auto res = + raft::sparse::solver::mst(handle, + off_d.data(), + ind_d.data(), + w_d.data(), + static_cast(v), + static_cast(e), + col_d.data(), + stream, + false, + true, + 0); + + std::vector w_out(res.n_edges); + std::vector col_out(v); + raft::update_host(w_out.data(), res.weights.data(), res.n_edges, stream); + raft::update_host(col_out.data(), col_d.data(), v, stream); + resource::sync_stream(handle, stream); + + ASSERT_EQ(truth.n_edges, static_cast(res.n_edges)); + double sum = 0.0; + for (auto w : w_out) + sum += static_cast(w); + ASSERT_EQ(truth.weight, sum); + for (int i = 0; i < v; i++) { + ASSERT_EQ(static_cast(truth.colors[i]), col_out[i]); + } + }; + + run_typed(int64_t{}, int64_t{}, float{}); // uniform 64-bit, widened-key wide path + run_typed(int64_t{}, int64_t{}, double{}); // uniform 64-bit, full 128-bit key + run_typed(int32_t{}, int64_t{}, float{}); // mixed 32-bit vertex_t, 64-bit edge_t + run_typed(uint32_t{}, int64_t{}, float{}); // unsigned vertex_t on the wide path + run_typed(uint32_t{}, uint32_t{}, float{}); // unsigned narrow path (below INT_MAX) +} + } // namespace mst } // namespace raft From 7e3df5e21879b2a756f7dab561dcc252e65c37af Mon Sep 17 00:00:00 2001 From: Alex Fallin Date: Fri, 14 Aug 2026 18:42:05 -0700 Subject: [PATCH 2/4] Use raft::launch_kernel for kernel launches (#3104) --- .../sparse/solver/detail/mst_solver_inl.cuh | 167 +++++++++++++----- 1 file changed, 119 insertions(+), 48 deletions(-) diff --git a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh index 0d08e23bb2..68d4a16914 100644 --- a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh +++ b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -94,7 +95,8 @@ Graph_COO mst_solve(raft::resources const& handle, if (initialize_colors) { thrust::sequence(rmm::exec_policy(stream), parent.begin(), parent.end()); } else { - mst_init_parent_kernel<<>>(v, color, parent.data()); + raft::launch_kernel( + stream, vblocks, mst_block_size, mst_init_parent_kernel, v, color, parent.data()); } RAFT_CUDA_TRY(cudaMemsetAsync(minv_raw.data(), 0xFF, minv_bytes, stream)); RAFT_CUDA_TRY(cudaMemsetAsync(in_mst.data(), 0, e * sizeof(bool), stream)); @@ -110,7 +112,8 @@ Graph_COO mst_solve(raft::resources const& handle, if (iterations == 0 && e / v >= filter_min_avg_degree) { const int ns = static_cast(std::min(e, max_samples)); rmm::device_uvector keys_d(ns, stream); - mst_sample_keys_kernel<<<1, 32, 0, stream>>>(ns, e, weights, keys_d.data()); + raft::launch_kernel( + stream, 1, 32, mst_sample_keys_kernel, ns, e, weights, keys_d.data()); key_t keys[max_samples]; raft::update_host(keys, keys_d.data(), ns, stream); RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); @@ -133,30 +136,62 @@ Graph_COO mst_solve(raft::resources const& handle, mst_ull* const base = reinterpret_cast(minv_raw.data()); mst_ull* const cur = base + (round % 2) * static_cast(v); mst_ull* const prev = base + ((round + 1) % 2) * static_cast(v); - mst_filter_min_kernel<<>>( - d1, wl_size, d2, wl_size_d.data(), parent.data(), cur, prev); + raft::launch_kernel(stream, + wblocks, + mst_block_size, + mst_filter_min_kernel, + d1, + wl_size, + d2, + wl_size_d.data(), + parent.data(), + cur, + prev); std::swap(d1, d2); wl_size = wl_size_d.value(stream); if (wl_size > 0) { const int nblocks = static_cast( (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); - mst_select_join_kernel<<>>( - d1, wl_size, parent.data(), cur, in_mst.data()); + raft::launch_kernel(stream, + nblocks, + mst_block_size, + mst_select_join_kernel, + d1, + wl_size, + parent.data(), + cur, + in_mst.data()); } } else { #if RAFT_MST_HAS_CAS128 mst_u128* const base = reinterpret_cast(minv_raw.data()); mst_u128* const cur = base + (round % 2) * static_cast(v); mst_u128* const prev = base + ((round + 1) % 2) * static_cast(v); - mst_filter_min_kernel<<>>( - d1, wl_size, d2, wl_size_d.data(), parent.data(), cur, prev); + raft::launch_kernel(stream, + wblocks, + mst_block_size, + mst_filter_min_kernel, + d1, + wl_size, + d2, + wl_size_d.data(), + parent.data(), + cur, + prev); std::swap(d1, d2); wl_size = wl_size_d.value(stream); if (wl_size > 0) { const int nblocks = static_cast( (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); - mst_select_join_kernel<<>>( - d1, wl_size, parent.data(), cur, in_mst.data()); + raft::launch_kernel(stream, + nblocks, + mst_block_size, + mst_select_join_kernel, + d1, + wl_size, + parent.data(), + cur, + in_mst.data()); } #else mst_ull* const base = reinterpret_cast(minv_raw.data()); @@ -166,17 +201,40 @@ Graph_COO mst_solve(raft::resources const& handle, base + 2 * static_cast(v) + (round % 2) * static_cast(v); mst_ull* const mine_prev = base + 2 * static_cast(v) + ((round + 1) % 2) * static_cast(v); - mst_filter_min_kernel<<>>( - d1, wl_size, d2, wl_size_d.data(), parent.data(), minw_cur, minw_prev, mine_prev); + raft::launch_kernel(stream, + wblocks, + mst_block_size, + mst_filter_min_kernel, + d1, + wl_size, + d2, + wl_size_d.data(), + parent.data(), + minw_cur, + minw_prev, + mine_prev); std::swap(d1, d2); wl_size = wl_size_d.value(stream); if (wl_size > 0) { const int nblocks = static_cast( (static_cast(wl_size) + mst_block_size - 1) / mst_block_size); - mst_min_index_kernel<<>>( - d1, wl_size, minw_cur, mine_cur); - mst_select_join_kernel<<>>( - d1, wl_size, parent.data(), mine_cur, in_mst.data()); + raft::launch_kernel(stream, + nblocks, + mst_block_size, + mst_min_index_kernel, + d1, + wl_size, + minw_cur, + mine_cur); + raft::launch_kernel(stream, + nblocks, + mst_block_size, + mst_select_join_kernel, + d1, + wl_size, + parent.data(), + mine_cur, + in_mst.data()); } #endif } @@ -189,27 +247,35 @@ Graph_COO mst_solve(raft::resources const& handle, auto launch_init = [&](bool first) { wl_size_d.set_value_to_zero_async(stream); if (first) { - mst_init_worklist_kernel<<>>(wl1.data(), - wl_size_d.data(), - wl_capacity, - v, - e, - offsets, - indices, - weights, - parent.data(), - thr_key); + raft::launch_kernel(stream, + eblocks, + mst_block_size, + mst_init_worklist_kernel, + wl1.data(), + wl_size_d.data(), + wl_capacity, + v, + e, + offsets, + indices, + weights, + parent.data(), + thr_key); } else { - mst_init_worklist_kernel<<>>(wl1.data(), - wl_size_d.data(), - wl_capacity, - v, - e, - offsets, - indices, - weights, - parent.data(), - thr_key); + raft::launch_kernel(stream, + eblocks, + mst_block_size, + mst_init_worklist_kernel, + wl1.data(), + wl_size_d.data(), + wl_capacity, + v, + e, + offsets, + indices, + weights, + parent.data(), + thr_key); } const wl_size_t wl_size = wl_size_d.value(stream); // wl_size < 0 means the admission counter wrapped (malformed input) @@ -227,7 +293,8 @@ Graph_COO mst_solve(raft::resources const& handle, boruvka(launch_init(false)); } - mst_flatten_colors_kernel<<>>(v, parent.data(), color); + raft::launch_kernel( + stream, vblocks, mst_block_size, mst_flatten_colors_kernel, v, parent.data(), color); // symmetrized count can exceed 32-bit edge_t/vertex_t for v > 2^30: // fail loudly rather than under-allocate @@ -241,17 +308,21 @@ Graph_COO mst_solve(raft::resources const& handle, Graph_COO mst_result(std::max(max_out, 1), stream); rmm::device_scalar out_count(stream); out_count.set_value_to_zero_async(stream); - mst_extract_coo_kernel<<>>(v, - e, - offsets, - indices, - weights, - in_mst.data(), - symmetrize_output, - mst_result.src.data(), - mst_result.dst.data(), - mst_result.weights.data(), - out_count.data()); + raft::launch_kernel(stream, + eblocks, + mst_block_size, + mst_extract_coo_kernel, + v, + e, + offsets, + indices, + weights, + in_mst.data(), + symmetrize_output, + mst_result.src.data(), + mst_result.dst.data(), + mst_result.weights.data(), + out_count.data()); mst_result.n_edges = out_count.value(stream); mst_result.src.resize(mst_result.n_edges, stream); mst_result.dst.resize(mst_result.n_edges, stream); From 25bff4ce295db8c33a606ecf757e60f65aec53d2 Mon Sep 17 00:00:00 2001 From: Alex Fallin Date: Mon, 17 Aug 2026 10:05:40 -0700 Subject: [PATCH 3/4] Add usage example to mst() documentation --- cpp/include/raft/sparse/solver/mst.cuh | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/cpp/include/raft/sparse/solver/mst.cuh b/cpp/include/raft/sparse/solver/mst.cuh index 326164fa9a..c32bc167a6 100644 --- a/cpp/include/raft/sparse/solver/mst.cuh +++ b/cpp/include/raft/sparse/solver/mst.cuh @@ -16,6 +16,22 @@ namespace sparse::solver { * the connected components of the given graph. * Algorithm based on ECL-MST (Fallin, Gonzalez, Seo, Burtscher, SC'23). * + * Usage example: + * @code{.cpp} + * #include + * #include + * #include + * + * raft::resources res; + * auto stream = raft::resource::get_cuda_stream(res); + * // device CSR of a symmetric graph: offsets (size v+1), indices and weights (size e) + * rmm::device_uvector colors(v, stream); + * auto forest = raft::sparse::solver::mst( + * res, offsets, indices, weights, v, e, colors.data(), stream); + * // forest.src/dst/weights hold forest.n_edges edges (both directions when + * // symmetrize_output); colors[i] = component id of vertex i + * @endcode + * * @tparam vertex_t integral type for precision of vertex indexing * @tparam edge_t integral type for precision of edge indexing * @tparam weight_t type of weights array From 746ba6c78b3299b9f9b95cbd94e17a7deddcf080 Mon Sep 17 00:00:00 2001 From: Alex Fallin Date: Mon, 17 Aug 2026 15:56:46 -0700 Subject: [PATCH 4/4] Tighten comments and clarify resume-colors documentation --- .../raft/sparse/solver/detail/mst_kernels.cuh | 16 ++++------------ .../raft/sparse/solver/detail/mst_solver_inl.cuh | 3 --- cpp/include/raft/sparse/solver/mst.cuh | 4 ++-- cpp/include/raft/sparse/solver/mst_solver.cuh | 6 ++---- cpp/tests/sparse/mst.cu | 8 ++++---- 5 files changed, 12 insertions(+), 25 deletions(-) diff --git a/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh b/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh index 18bf372820..74a2f8623d 100644 --- a/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh +++ b/cpp/include/raft/sparse/solver/detail/mst_kernels.cuh @@ -25,7 +25,7 @@ #else #define RAFT_MST_MIN_ARCH 0 // unknown toolchain: take the portable path #endif -// RAFT_MST_FORCE_TWOPASS: testing knob, compiles the portable wide path on +// RAFT_MST_FORCE_TWOPASS: for testing, compiles the portable wide path on // any target. Must be defined consistently across all TUs of a binary (ODR). #if defined(RAFT_MST_FORCE_TWOPASS) #define RAFT_MST_HAS_CAS128 0 @@ -68,7 +68,6 @@ RAFT_DEVICE_INLINE_FUNCTION mst_ull mst_order_key(int64_t w) return static_cast(w) ^ 0x8000000000000000ull; } -// Wide worklist entry: 16-byte-aligned {x, y, z, w} of long long. // CUDA 13 deprecates longlong4 in favor of longlong4_16a #if defined(CUDART_VERSION) && CUDART_VERSION >= 13000 using mst_entry64 = longlong4_16a; @@ -86,7 +85,6 @@ struct mst_traits { using wl_size_t = std::conditional_t; }; -// atomics over possibly-signed types; all values here are non-negative template RAFT_DEVICE_INLINE_FUNCTION T mst_atomic_cas(T* addr, T compare, T val) { @@ -139,7 +137,7 @@ RAFT_DEVICE_INLINE_FUNCTION void mst_word_store(T* addr, T val) } } -// Find with path halving (without it, equal-weight tie chains go quadratic) +// Find with path halving (without it, equal-weight tie chains go quadratic, see tests for example) template RAFT_DEVICE_INLINE_FUNCTION vertex_t mst_uf_find(vertex_t curr, vertex_t* const __restrict__ parent) { @@ -205,7 +203,6 @@ RAFT_KERNEL mst_sample_keys_kernel( if (i < n_samples) keys[i] = mst_order_key(weights[mst_sample_hash(i) % e]); } -// gather flagged CSR edges into COO; optionally emit both directions template RAFT_KERNEL mst_extract_coo_kernel(const vertex_t v, const edge_t e, @@ -221,7 +218,6 @@ RAFT_KERNEL mst_extract_coo_kernel(const vertex_t v, { const long long j = mst_grid_idx(); if (j < e && in_mst[j]) { - // row of edge j by binary search over the offsets vertex_t lo = 0, hi = v; while (lo + 1 < hi) { const vertex_t mid = lo + (hi - lo) / 2; @@ -259,7 +255,6 @@ RAFT_KERNEL mst_init_worklist_kernel( const long long j = mst_grid_idx(); if (j < e) { const vertex_t n = indices[j]; - // row of edge j by binary search over the offsets vertex_t lo = 0, hi = v; while (lo + 1 < hi) { const vertex_t mid = lo + (hi - lo) / 2; @@ -298,7 +293,6 @@ RAFT_KERNEL mst_init_worklist_kernel( } } -// ---- narrow path (4-byte weight_t + edge_t): packed 64-bit (key, index) ---- template RAFT_KERNEL mst_filter_min_kernel(const int4* const __restrict__ wl1, const int wl1_size, @@ -347,7 +341,7 @@ RAFT_KERNEL mst_select_join_kernel(const int4* const __restrict__ wl, } #if RAFT_MST_HAS_CAS128 -// ---- wide path, sm_90+: single-pass packed 128-bit (key, edge index) ------- +// sm_90+: single-pass RAFT_DEVICE_INLINE_FUNCTION void mst_atomic_min_u128(mst_u128* const addr, const mst_u128 val) { mst_u128 old = atomicCAS(addr, val, val); @@ -408,8 +402,7 @@ RAFT_KERNEL mst_select_join_kernel(const mst_entry64* const __restrict__ wl, } #else -// ---- wide path, portable: two-pass min (weight key, then edge index) ------- - +// lower than sm_90 portable two-pass template RAFT_KERNEL mst_filter_min_kernel(const mst_entry64* const __restrict__ wl1, const wl_size_t wl1_size, @@ -440,7 +433,6 @@ RAFT_KERNEL mst_filter_min_kernel(const mst_entry64* const __restrict__ wl1, } } -// pass 2: min edge index among key-tied edges template RAFT_KERNEL mst_min_index_kernel(const entry_t* const __restrict__ wl, const wl_size_t wl_size, diff --git a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh index 68d4a16914..d934b519be 100644 --- a/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh +++ b/cpp/include/raft/sparse/solver/detail/mst_solver_inl.cuh @@ -72,7 +72,6 @@ Graph_COO mst_solve(raft::resources const& handle, "unsigned 32-bit edge counts above INT_MAX are not supported"); } - // one worklist entry per undirected edge const wl_size_t wl_capacity = static_cast(e / 2 + 1); rmm::device_uvector parent(v, stream); @@ -278,7 +277,6 @@ Graph_COO mst_solve(raft::resources const& handle, thr_key); } const wl_size_t wl_size = wl_size_d.value(stream); - // wl_size < 0 means the admission counter wrapped (malformed input) RAFT_EXPECTS(wl_size >= 0 && wl_size <= wl_capacity, "MST worklist overflow: the input CSR must be symmetric (each " "undirected edge stored in both directions)."); @@ -288,7 +286,6 @@ Graph_COO mst_solve(raft::resources const& handle, boruvka(launch_init(true)); if (filtered) { - // clear straggler minima before admitting the remaining edges RAFT_CUDA_TRY(cudaMemsetAsync(minv_raw.data(), 0xFF, minv_bytes, stream)); boruvka(launch_init(false)); } diff --git a/cpp/include/raft/sparse/solver/mst.cuh b/cpp/include/raft/sparse/solver/mst.cuh index c32bc167a6..83e2949ee0 100644 --- a/cpp/include/raft/sparse/solver/mst.cuh +++ b/cpp/include/raft/sparse/solver/mst.cuh @@ -47,8 +47,8 @@ namespace sparse::solver { * @param v number of vertices in graph * @param e number of edges in graph * @param color array to store resulting colors for MSF; when initialize_colors is false it is - * also the input seeding and must hold a valid component labeling from a previous solve of the same - * graph + * also the input seeding and must hold a valid component labeling from a previous solve over the + * same vertex set (the edge set may differ, e.g. when reconnecting components) * @param stream cuda stream for ordering operations * @param symmetrize_output should the resulting output edge list be symmetrized? * @param initialize_colors should the colors array be initialized inside the MST? diff --git a/cpp/include/raft/sparse/solver/mst_solver.cuh b/cpp/include/raft/sparse/solver/mst_solver.cuh index 77274e9dfa..76727a947b 100644 --- a/cpp/include/raft/sparse/solver/mst_solver.cuh +++ b/cpp/include/raft/sparse/solver/mst_solver.cuh @@ -28,16 +28,14 @@ struct Graph_COO { }; /** - * @brief MST solver based on ECL-MST, with deterministic lexicographic - * (weight, edge index) tie-breaking. + * @brief MST solver with deterministic tie-breaking. * * @tparam vertex_t integral type for vertex indexing (32- or 64-bit) * @tparam edge_t integral type for edge indexing (32- or 64-bit, at least as * wide as vertex_t) * @tparam weight_t type of the weights array * @tparam alteration_t unused; retained for source compatibility with - * existing callers (the solver no longer perturbs ("alters") weights to break - * ties) + * existing callers (the solver no longer alters weights to break ties) */ template class MST_solver { diff --git a/cpp/tests/sparse/mst.cu b/cpp/tests/sparse/mst.cu index d057c4a6a2..edcc020a4e 100644 --- a/cpp/tests/sparse/mst.cu +++ b/cpp/tests/sparse/mst.cu @@ -588,8 +588,8 @@ void expect_forest(int v, const std::vector& src, const std::vector& d } } -// Regression: heavily-tied integer-valued float weights at ~1e7, where the -// previous solver's alteration rounded away (float ulp there is 1-2) and +// heavily-tied integer-valued float weights at ~1e7, where the +// previous solver's alteration rounded away and // produced cycles or non-minimal forests past its own guard. TEST(MST, FloatMagnitudeTies) { @@ -677,8 +677,8 @@ TEST(MST, DisconnectedColors) } } -// Regression: an equal-weight path graph chain-joins into an O(v)-deep -// parent chain; quadratic (minutes at v=1M) without path halving. +// an equal-weight path graph chain-joins into an O(v)-deep parent chain +// quadratic (minutes at v=1M) w/o path halving (change from the SC'23 approach) TEST(MST, UniformWeightPathGraph) { raft::resources handle;