diff --git a/nsparse/cluster/inverted_list_clusters.cpp b/nsparse/cluster/inverted_list_clusters.cpp index c8f637a..d0253fe 100644 --- a/nsparse/cluster/inverted_list_clusters.cpp +++ b/nsparse/cluster/inverted_list_clusters.cpp @@ -125,14 +125,16 @@ SparseVectors summarize_with_cpu_(const SparseVectors* vectors, #ifdef NSPARSE_WITH_GPU // GPU per-term max-pool: one kernel launch for the whole list, then the shared -// CPU sort/truncate. float (U32) only. Returns std::nullopt when the GPU -// declines (unavailable / empty list) so the caller falls back to the CPU path. +// CPU sort/truncate. float (U32) and 8-bit (uint8_t) only; a 0..255 code max is +// exact in the float summary value. Returns std::nullopt when the GPU declines +// (unavailable / empty list) or T is unsupported (e.g. uint16_t) so the caller +// falls back to the CPU path. template std::optional summarize_with_gpu_( const SparseVectors* vectors, std::span group_of_doc_ids, std::span offsets, float alpha) { - if constexpr (!std::is_same_v) { - return std::nullopt; // GPU max-pool is float-only + if constexpr (!std::is_same_v && !std::is_same_v) { + return std::nullopt; // GPU max-pool supports only float and 8-bit } else { const size_t n_clusters = offsets.size() - 1; std::vector gpu_clusters; diff --git a/nsparse/cluster/kmeans_utils.cpp b/nsparse/cluster/kmeans_utils.cpp index 521ec56..1291475 100644 --- a/nsparse/cluster/kmeans_utils.cpp +++ b/nsparse/cluster/kmeans_utils.cpp @@ -169,11 +169,14 @@ void map_docs_to_clusters(const SparseVectors* vectors, return; } #ifdef NSPARSE_WITH_GPU - // GPU assignment (cuSPARSE) for float (U32) weights. It matches the CPU - // path below (skips centroids, ties to the lowest cluster index); on any - // failure we fall through to CPU. No AVX-512 exclusion is needed: there is - // now a single CPU path and it skips centroids like the GPU one. - if (vectors->get_element_size() == U32 && + // GPU assignment (cuSPARSE) for float (U32) or 8-bit (U8) weights. It + // matches the CPU path below (skips centroids, ties to the lowest cluster + // index); on any failure we fall through to CPU. 8-bit uses int32 + // accumulation on the GPU, which is exact for byte-sized products and so + // yields the same argmax as the CPU int64 path. 16-bit (U16) is not + // offloaded: its products need int64, which the GPU path does not do. + const auto gpu_element_size = vectors->get_element_size(); + if ((gpu_element_size == U32 || gpu_element_size == U8) && should_offload_assignment_to_gpu(docs.size(), clusters.size())) { try { GpuClusterAssigner::instance().assign(vectors, docs, clusters); diff --git a/nsparse/gpu/gpu_cluster_assigner.cu b/nsparse/gpu/gpu_cluster_assigner.cu index 8a79f13..519b00c 100644 --- a/nsparse/gpu/gpu_cluster_assigner.cu +++ b/nsparse/gpu/gpu_cluster_assigner.cu @@ -27,15 +27,18 @@ namespace { // One thread per document row: argmax over the row's n_clusters scores, ties // broken to the lowest cluster index (strict-greater) to match the CPU path. -__global__ void row_argmax_kernel(const float* __restrict__ scores, int n_rows, +// TScore is float for the unquantized path, int64_t for the 8-bit path (the +// int64 dot mirrors the CPU reference); the argmax is identical either way. +template +__global__ void row_argmax_kernel(const TScore* __restrict__ scores, int n_rows, int n_clusters, int32_t* __restrict__ best_cluster) { int row = blockIdx.x * blockDim.x + threadIdx.x; if (row >= n_rows) { return; } - const float* row_scores = scores + static_cast(row) * n_clusters; - float best = row_scores[0]; + const TScore* row_scores = scores + static_cast(row) * n_clusters; + TScore best = row_scores[0]; int best_j = 0; for (int j = 1; j < n_clusters; ++j) { if (row_scores[j] > best) { @@ -48,13 +51,15 @@ __global__ void row_argmax_kernel(const float* __restrict__ scores, int n_rows, // Gather the CSR of the document sub-matrix A from the resident corpus, so the // bulk doc data never re-crosses PCIe per list. One thread per output row. +// TVal is the corpus value width (float or int8_t). +template __global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, const int32_t* __restrict__ corpus_indices, - const float* __restrict__ corpus_values, + const TVal* __restrict__ corpus_values, const int32_t* __restrict__ docs, int n_docs, const int32_t* __restrict__ a_row_ptr, int32_t* __restrict__ a_col, - float* __restrict__ a_val) { + TVal* __restrict__ a_val) { int row = blockIdx.x * blockDim.x + threadIdx.x; if (row >= n_docs) { return; @@ -71,13 +76,14 @@ __global__ void gather_csr_kernel(const int32_t* __restrict__ corpus_indptr, // Scatter centroid rows into the dense B matrix (dim x n_clusters, row-major, // ldb = n_clusters). One thread per centroid; each owns a disjoint column, so -// no races. B must be zeroed first. +// no races. B must be zeroed first. TVal is the corpus value width. +template __global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, const int32_t* __restrict__ corpus_indices, - const float* __restrict__ corpus_values, + const TVal* __restrict__ corpus_values, const int32_t* __restrict__ centroids, int n_clusters, int ldb, - float* __restrict__ b) { + TVal* __restrict__ b) { int j = blockIdx.x * blockDim.x + threadIdx.x; if (j >= n_clusters) { return; @@ -91,6 +97,41 @@ __global__ void scatter_dense_kernel(const int32_t* __restrict__ corpus_indptr, } } +// SpMM for 8-bit codes: C[row][j] = sum over doc row's corpus nonzeros of +// code * B[col][j], read straight from the resident corpus (no CSR gather). +// Used instead of cuSPARSE: its int8 SpMM is SIGNED and would misread the +// unsigned [0,255] codes, and it offers no unsigned-int8 path. The accumulator +// is int64 to match the CPU reference for every input: a single product fits +// int32 (<= 255*255), but a doc overlapping a centroid in more than +// INT32_MAX/65025 (~33k) high-value terms -- reachable once the vocabulary +// exceeds ~33k -- would overflow an int32 sum and diverge from the CPU int64 +// argmax. One block per doc row; threads stripe over the clusters. +__global__ void assign_u8_spmm_kernel( + const int32_t* __restrict__ corpus_indptr, + const int32_t* __restrict__ corpus_indices, + const uint8_t* __restrict__ corpus_values, + const int32_t* __restrict__ docs, int n_docs, + const uint8_t* __restrict__ b, int n_clusters, int ldb, + int64_t* __restrict__ c) { + const int row = blockIdx.x; + if (row >= n_docs) { + return; + } + const int32_t d = docs[row]; + const int32_t start = corpus_indptr[d]; + const int32_t end = corpus_indptr[d + 1]; + int64_t* c_row = c + static_cast(row) * n_clusters; + for (int j = threadIdx.x; j < n_clusters; j += blockDim.x) { + int64_t acc = 0; + for (int32_t p = start; p < end; ++p) { + const int32_t col = corpus_indices[p]; + acc += static_cast(corpus_values[p]) * + static_cast(b[static_cast(col) * ldb + j]); + } + c_row[j] = acc; + } +} + // Per-thread GPU resources: an independent cuSPARSE handle + stream (handles // are not shareable across threads) and scratch buffers reused across lists, // grown on demand so steady-state builds do no per-list cudaMalloc. @@ -103,9 +144,10 @@ struct ThreadCtx { int32_t* d_centroids = nullptr; size_t cent_cap = 0; int32_t* d_a_row_ptr = nullptr; size_t rowptr_cap = 0; int32_t* d_a_col = nullptr; size_t col_cap = 0; - float* d_a_val = nullptr; size_t val_cap = 0; - float* d_b = nullptr; size_t b_cap = 0; - float* d_c = nullptr; size_t c_cap = 0; + // Value/dense/output scratch are byte-capacity (float, or uint8 + int64). + void* d_a_val = nullptr; size_t val_cap = 0; + void* d_b = nullptr; size_t b_cap = 0; + void* d_c = nullptr; size_t c_cap = 0; int32_t* d_best = nullptr; size_t best_cap = 0; void* d_spmm = nullptr; size_t spmm_cap = 0; @@ -162,6 +204,120 @@ struct DnMatGuard { } }; +// cuSPARSE SpMM core for the unquantized (float) path: build dense B and CSR A, +// run C = A * B, and argmax each row into ctx.d_best (copied to h_best). +// Templated so the types are explicit, but instantiated only as ; +// 8-bit codes cannot use cuSPARSE (its int8 SpMM is signed) and go through +// assign_u8 instead. The d_docs, d_centroids and d_a_row_ptr scratch are +// uploaded by the caller. +template +void assign_spmm(const DeviceCorpus& corpus, ThreadCtx& ctx, int n_docs, + int n_clusters, int dim, int64_t nnz_a, + cudaDataType_t val_type, cudaDataType_t compute_type, + std::vector& h_best) { + cudaStream_t stream = ctx.stream; + constexpr int kBlock = 256; + + // Dense B (dim x n_clusters, row-major) from centroid rows. + const size_t b_bytes = static_cast(dim) * n_clusters * sizeof(TVal); + ensure_capacity(reinterpret_cast(&ctx.d_b), ctx.b_cap, b_bytes); + NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_b, 0, b_bytes, stream)); + const int cent_grid = (n_clusters + kBlock - 1) / kBlock; + scatter_dense_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_centroids, n_clusters, + n_clusters, static_cast(ctx.d_b)); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + // CSR A gathered from the resident corpus, then C = A * B. + ensure_capacity(reinterpret_cast(&ctx.d_a_val), ctx.val_cap, + static_cast(nnz_a) * sizeof(TVal)); + ensure_capacity(reinterpret_cast(&ctx.d_c), ctx.c_cap, + static_cast(n_docs) * n_clusters * sizeof(TScore)); + const int docs_grid = (n_docs + kBlock - 1) / kBlock; + gather_csr_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_docs, n_docs, + ctx.d_a_row_ptr, ctx.d_a_col, static_cast(ctx.d_a_val)); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + SpMatGuard a_guard; + DnMatGuard b_guard; + DnMatGuard c_guard; + NSPARSE_CUSPARSE_CHECK(cusparseCreateCsr( + &a_guard.desc, n_docs, dim, nnz_a, ctx.d_a_row_ptr, ctx.d_a_col, + ctx.d_a_val, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_BASE_ZERO, val_type)); + NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat(&b_guard.desc, dim, n_clusters, + n_clusters, ctx.d_b, val_type, + CUSPARSE_ORDER_ROW)); + NSPARSE_CUSPARSE_CHECK( + cusparseCreateDnMat(&c_guard.desc, n_docs, n_clusters, n_clusters, + ctx.d_c, compute_type, CUSPARSE_ORDER_ROW)); + + const TScore alpha_v = 1; + const TScore beta_v = 0; + size_t buffer_size = 0; + NSPARSE_CUSPARSE_CHECK(cusparseSpMM_bufferSize( + ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, a_guard.desc, b_guard.desc, + &beta_v, c_guard.desc, compute_type, CUSPARSE_SPMM_ALG_DEFAULT, + &buffer_size)); + ensure_capacity(reinterpret_cast(&ctx.d_spmm), ctx.spmm_cap, + buffer_size == 0 ? 1 : buffer_size); + NSPARSE_CUSPARSE_CHECK(cusparseSpMM( + ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, + CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, a_guard.desc, b_guard.desc, + &beta_v, c_guard.desc, compute_type, CUSPARSE_SPMM_ALG_DEFAULT, + ctx.d_spmm)); + + row_argmax_kernel<<>>( + static_cast(ctx.d_c), n_docs, n_clusters, ctx.d_best); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_best.data(), ctx.d_best, + static_cast(n_docs) * + sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +// 8-bit path: dense B (uint8) from centroids + the custom int32 SpMM kernel + +// argmax. Separate from assign_spmm because cuSPARSE cannot do unsigned int8. +void assign_u8(const DeviceCorpus& corpus, ThreadCtx& ctx, int n_docs, + int n_clusters, int dim, std::vector& h_best) { + cudaStream_t stream = ctx.stream; + constexpr int kBlock = 256; + + const size_t b_bytes = static_cast(dim) * n_clusters; + ensure_capacity(reinterpret_cast(&ctx.d_b), ctx.b_cap, b_bytes); + NSPARSE_CUDA_CHECK(cudaMemsetAsync(ctx.d_b, 0, b_bytes, stream)); + const int cent_grid = (n_clusters + kBlock - 1) / kBlock; + scatter_dense_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_centroids, n_clusters, + n_clusters, static_cast(ctx.d_b)); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + ensure_capacity(reinterpret_cast(&ctx.d_c), ctx.c_cap, + static_cast(n_docs) * n_clusters * sizeof(int64_t)); + assign_u8_spmm_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_docs, n_docs, + static_cast(ctx.d_b), n_clusters, n_clusters, + static_cast(ctx.d_c)); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + + const int docs_grid = (n_docs + kBlock - 1) / kBlock; + row_argmax_kernel<<>>( + static_cast(ctx.d_c), n_docs, n_clusters, ctx.d_best); + NSPARSE_CUDA_CHECK(cudaGetLastError()); + NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_best.data(), ctx.d_best, + static_cast(n_docs) * + sizeof(int32_t), + cudaMemcpyDeviceToHost, stream)); + NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); +} + } // namespace GpuClusterAssigner& GpuClusterAssigner::instance() { @@ -209,11 +365,16 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, ThreadCtx& ctx = thread_ctx(); cudaStream_t stream = ctx.stream; - // Upload only the small per-list metadata (doc/centroid ids). + // Upload the small per-list metadata (doc/centroid ids, row pointers) and + // size the index scratch; the value/dense/output buffers are sized per + // width inside assign_spmm. ensure_capacity(&ctx.d_docs, ctx.docs_cap, n_docs * sizeof(int32_t)); ensure_capacity(&ctx.d_centroids, ctx.cent_cap, n_clusters * sizeof(int32_t)); ensure_capacity(&ctx.d_best, ctx.best_cap, n_docs * sizeof(int32_t)); - ensure_capacity(&ctx.d_b, ctx.b_cap, dim * n_clusters * sizeof(float)); + ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, + (n_docs + 1) * sizeof(int32_t)); + ensure_capacity(&ctx.d_a_col, ctx.col_cap, + static_cast(nnz_a) * sizeof(int32_t)); NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_docs, docs.data(), n_docs * sizeof(int32_t), @@ -221,78 +382,23 @@ void GpuClusterAssigner::assign(const SparseVectors* vectors, NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_centroids, centroid_docs.data(), n_clusters * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - - // Build dense B (dim x n_clusters, row-major) from centroid rows. - constexpr int kBlock = 256; - NSPARSE_CUDA_CHECK( - cudaMemsetAsync(ctx.d_b, 0, dim * n_clusters * sizeof(float), stream)); - const int cent_grid = static_cast((n_clusters + kBlock - 1) / kBlock); - scatter_dense_kernel<<>>( - corpus.indptr, corpus.indices, corpus.values, ctx.d_centroids, - static_cast(n_clusters), static_cast(n_clusters), ctx.d_b); - NSPARSE_CUDA_CHECK(cudaGetLastError()); - - // Gather A's CSR, then C = A * B (SpMM) and per-row argmax. - ensure_capacity(&ctx.d_a_row_ptr, ctx.rowptr_cap, - (n_docs + 1) * sizeof(int32_t)); - ensure_capacity(&ctx.d_a_col, ctx.col_cap, - static_cast(nnz_a) * sizeof(int32_t)); - ensure_capacity(&ctx.d_a_val, ctx.val_cap, - static_cast(nnz_a) * sizeof(float)); - ensure_capacity(&ctx.d_c, ctx.c_cap, n_docs * n_clusters * sizeof(float)); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(ctx.d_a_row_ptr, h_a_row_ptr.data(), (n_docs + 1) * sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - const int docs_grid = static_cast((n_docs + kBlock - 1) / kBlock); - gather_csr_kernel<<>>( - corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, - static_cast(n_docs), ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val); - NSPARSE_CUDA_CHECK(cudaGetLastError()); - - SpMatGuard a_guard; - DnMatGuard b_guard; - DnMatGuard c_guard; - cusparseSpMatDescr_t& mat_a = a_guard.desc; - cusparseDnMatDescr_t& mat_b = b_guard.desc; - cusparseDnMatDescr_t& mat_c = c_guard.desc; - NSPARSE_CUSPARSE_CHECK(cusparseCreateCsr( - &mat_a, static_cast(n_docs), static_cast(dim), nnz_a, - ctx.d_a_row_ptr, ctx.d_a_col, ctx.d_a_val, CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); - NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_b, static_cast(dim), static_cast(n_clusters), - static_cast(n_clusters), ctx.d_b, CUDA_R_32F, - CUSPARSE_ORDER_ROW)); - NSPARSE_CUSPARSE_CHECK(cusparseCreateDnMat( - &mat_c, static_cast(n_docs), static_cast(n_clusters), - static_cast(n_clusters), ctx.d_c, CUDA_R_32F, - CUSPARSE_ORDER_ROW)); - - const float alpha_v = 1.0F; - const float beta_v = 0.0F; - size_t buffer_size = 0; - NSPARSE_CUSPARSE_CHECK(cusparseSpMM_bufferSize( - ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, - CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, mat_a, mat_b, &beta_v, - mat_c, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, &buffer_size)); - ensure_capacity(reinterpret_cast(&ctx.d_spmm), ctx.spmm_cap, - buffer_size == 0 ? 1 : buffer_size); - NSPARSE_CUSPARSE_CHECK(cusparseSpMM( - ctx.handle, CUSPARSE_OPERATION_NON_TRANSPOSE, - CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha_v, mat_a, mat_b, &beta_v, - mat_c, CUDA_R_32F, CUSPARSE_SPMM_ALG_DEFAULT, ctx.d_spmm)); - - row_argmax_kernel<<>>( - ctx.d_c, static_cast(n_docs), static_cast(n_clusters), - ctx.d_best); - NSPARSE_CUDA_CHECK(cudaGetLastError()); + // Float corpus -> cuSPARSE float SpMM; 8-bit codes -> the custom uint8 + // int64-accumulate kernel (matches the CPU int64 argmax for any input). + // Only these two widths reach here (ensure_resident rejects others). std::vector h_best(n_docs); - NSPARSE_CUDA_CHECK(cudaMemcpyAsync(h_best.data(), ctx.d_best, - n_docs * sizeof(int32_t), - cudaMemcpyDeviceToHost, stream)); - NSPARSE_CUDA_CHECK(cudaStreamSynchronize(stream)); + if (corpus.element_size == U32) { + assign_spmm(corpus, ctx, static_cast(n_docs), + static_cast(n_clusters), + static_cast(dim), nnz_a, CUDA_R_32F, + CUDA_R_32F, h_best); + } else { + assign_u8(corpus, ctx, static_cast(n_docs), + static_cast(n_clusters), static_cast(dim), h_best); + } // Append assignments, skipping docs that are themselves a centroid. for (size_t i = 0; i < n_docs; ++i) { diff --git a/nsparse/gpu/gpu_common.cuh b/nsparse/gpu/gpu_common.cuh index b736390..53e42f5 100644 --- a/nsparse/gpu/gpu_common.cuh +++ b/nsparse/gpu/gpu_common.cuh @@ -63,16 +63,25 @@ void ensure_capacity(T** ptr, size_t& cap_bytes, size_t need_bytes) { struct DeviceCorpus { int32_t* indptr = nullptr; // n_vectors + 1 int32_t* indices = nullptr; // nnz - float* values = nullptr; // nnz + // Raw code bytes at `element_size` width: float (U32) for the unquantized + // indexes, int8 (U8) for 8-bit scalar-quantized. Consumers cast per width. + void* values = nullptr; // nnz * element_size bytes + size_t element_size = 0; int64_t nnz = 0; size_t dim = 0; size_t n_vectors = 0; // Residency identity. A pointer alone is unsafe (a freed SparseVectors - // address can be reused), so also compare n_vectors and nnz. + // address can be reused), so also compare the shape (n_vectors, nnz, dim) + // and the code width (element_size). Width is part of the key because path + // selection and byte interpretation now depend on element_size: a + // same-address corpus of a different width (float vs 8-bit) with matching + // counts must not be reused as-is. const SparseVectors* owner = nullptr; - bool matches(const SparseVectors* v, size_t nv, int64_t nz) const { - return owner == v && n_vectors == nv && nnz == nz; + bool matches(const SparseVectors* v, size_t nv, int64_t nz, size_t es, + size_t d) const { + return owner == v && n_vectors == nv && nnz == nz && + element_size == es && dim == d; } }; @@ -91,7 +100,9 @@ public: const size_t n_vectors = vectors->num_vectors(); const offset_t* indptr = vectors->indptr_data(); const term_t* indices = vectors->indices_data(); - const float* values = vectors->values_data_float(); + const uint8_t* values = vectors->values_data(); + const size_t element_size = vectors->get_element_size(); + const size_t dim = vectors->get_dimension(); const int64_t nnz = indptr[n_vectors]; // The device stores CSR offsets as int32 (and cuSPARSE's CSR API is @@ -101,9 +112,17 @@ public: throw std::runtime_error( "GPU build path requires nnz <= INT32_MAX; use the CPU path"); } + // Only float (U32) and 8-bit (U8) codes have GPU kernels; other widths + // (e.g. 16-bit) must use the CPU path. The build-time callers guard on + // element_size, but the invariant is re-checked here so a direct caller + // cannot upload a width the kernels would misread. + if (element_size != U32 && element_size != U8) { + throw std::runtime_error( + "GPU build path supports only float (U32) or 8-bit (U8) codes"); + } std::lock_guard lock(mutex_); - if (corpus_.matches(vectors, n_vectors, nnz)) { + if (corpus_.matches(vectors, n_vectors, nnz, element_size, dim)) { return corpus_; } free_locked(); @@ -116,11 +135,12 @@ public: for (int64_t i = 0; i < nnz; ++i) { indices32[i] = static_cast(indices[i]); } + const size_t values_bytes = static_cast(nnz) * element_size; check_cuda(cudaMalloc(&corpus_.indptr, (n_vectors + 1) * sizeof(int32_t)), "cudaMalloc(corpus.indptr)"); check_cuda(cudaMalloc(&corpus_.indices, nnz * sizeof(int32_t)), "cudaMalloc(corpus.indices)"); - check_cuda(cudaMalloc(&corpus_.values, nnz * sizeof(float)), + check_cuda(cudaMalloc(&corpus_.values, values_bytes), "cudaMalloc(corpus.values)"); check_cuda(cudaMemcpy(corpus_.indptr, indptr32.data(), (n_vectors + 1) * sizeof(int32_t), @@ -129,11 +149,12 @@ public: check_cuda(cudaMemcpy(corpus_.indices, indices32.data(), nnz * sizeof(int32_t), cudaMemcpyHostToDevice), "cudaMemcpy(corpus.indices)"); - check_cuda(cudaMemcpy(corpus_.values, values, nnz * sizeof(float), + check_cuda(cudaMemcpy(corpus_.values, values, values_bytes, cudaMemcpyHostToDevice), "cudaMemcpy(corpus.values)"); + corpus_.element_size = element_size; corpus_.nnz = nnz; - corpus_.dim = vectors->get_dimension(); + corpus_.dim = dim; corpus_.n_vectors = n_vectors; corpus_.owner = vectors; return corpus_; diff --git a/nsparse/gpu/gpu_summarizer.cu b/nsparse/gpu/gpu_summarizer.cu index f89fe53..af760e8 100644 --- a/nsparse/gpu/gpu_summarizer.cu +++ b/nsparse/gpu/gpu_summarizer.cu @@ -12,6 +12,7 @@ #include #include #include +#include #include #include "cuda_runtime.h" @@ -35,10 +36,15 @@ namespace { // are restored to zero before return. // out_term / out_val : compacted results, per cluster at out_base[b]. // out_count[b] / out_sum[b] : distinct terms and sum of maxes for cluster b. +// TVal is the corpus value width: float (unquantized) or uint8_t (8-bit codes). +// Both accumulate a per-term max in an int32 slot -- float via its +// order-preserving int bit-pattern (SPLADE weights >= 0), uint8 as the plain +// unsigned integer. out_val is float either way (a 0..255 code is exact). +template __global__ void summarize_list_kernel( const int32_t* __restrict__ corpus_indptr, const int32_t* __restrict__ corpus_indices, - const float* __restrict__ corpus_values, + const TVal* __restrict__ corpus_values, const int32_t* __restrict__ docs, const int32_t* __restrict__ offsets, const int32_t* __restrict__ out_base, int dim, int32_t* __restrict__ acc, int32_t* __restrict__ seen, @@ -64,8 +70,12 @@ __global__ void summarize_list_kernel( const int32_t len = corpus_indptr[d + 1] - src; for (int32_t t = 0; t < len; ++t) { const int32_t term = corpus_indices[src + t]; - const float v = corpus_values[src + t]; - atomicMax(&acc[base + term], __float_as_int(v)); + const TVal v = corpus_values[src + t]; + if constexpr (std::is_same_v) { + atomicMax(&acc[base + term], __float_as_int(v)); + } else { + atomicMax(&acc[base + term], static_cast(v)); + } if (atomicCAS(&seen[base + term], 0, 1) == 0) { out_term[obase + atomicAdd(&s_count, 1)] = term; } @@ -84,7 +94,9 @@ __global__ void summarize_list_kernel( for (int k = threadIdx.x; k < s_count; k += blockDim.x) { const int32_t term = out_term[obase + k]; - const float v = __int_as_float(acc[base + term]); + const float v = std::is_same_v + ? __int_as_float(acc[base + term]) + : static_cast(acc[base + term]); out_val[obase + k] = v; atomicAdd(&s_sum, v); acc[base + term] = 0; @@ -243,10 +255,25 @@ bool summarize_list_impl(const SparseVectors* vectors, const idx_t* docs, cudaMemcpyHostToDevice, stream)); constexpr int kBlock = 128; - summarize_list_kernel<<(n_clusters), kBlock, 0, stream>>>( - corpus.indptr, corpus.indices, corpus.values, ctx.d_docs, ctx.d_offsets, - ctx.d_out_base, static_cast(dim), ctx.d_acc, ctx.d_seen, - ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, ctx.d_out_sum); + const int grid = static_cast(n_clusters); + // Float corpus -> float max-pool; 8-bit codes -> int8 max-pool. Only these + // two widths reach here (ensure_resident rejects others). + if (corpus.element_size == U32) { + summarize_list_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_docs, ctx.d_offsets, + ctx.d_out_base, static_cast(dim), ctx.d_acc, ctx.d_seen, + ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, ctx.d_out_sum); + } else { + // 8-bit codes are unsigned [0,255]; read as uint8_t so the max-pool + // int32 promotion is the true code value, not a sign-extended one. + summarize_list_kernel<<>>( + corpus.indptr, corpus.indices, + static_cast(corpus.values), ctx.d_docs, + ctx.d_offsets, ctx.d_out_base, static_cast(dim), ctx.d_acc, + ctx.d_seen, ctx.d_out_term, ctx.d_out_val, ctx.d_out_count, + ctx.d_out_sum); + } NSPARSE_CUDA_CHECK(cudaGetLastError()); std::vector h_term(total_nnz); diff --git a/tests/gpu_cluster_assigner_test.cpp b/tests/gpu_cluster_assigner_test.cpp index 8d88a98..19f4bf3 100644 --- a/tests/gpu_cluster_assigner_test.cpp +++ b/tests/gpu_cluster_assigner_test.cpp @@ -235,7 +235,8 @@ TEST(GpuClusterAssignerTest, SummarizeListMaxpoolMatchesCpu) { EXPECT_FLOAT_EQ(gc.values[order[i]], exp_vals[i]) << "cluster " << b << " value " << i; } - EXPECT_NEAR(gc.sum, exp_sum, exp_sum * 1e-5F + 1e-5F) << "cluster " << b; + EXPECT_NEAR(gc.sum, exp_sum, exp_sum * 1e-5F + 1e-5F) + << "cluster " << b; } } @@ -275,5 +276,247 @@ TEST(GpuClusterAssignerTest, SummarizeListMaxpoolReuseIsClean) { } } +// ---- 8-bit (scalar-quantized) parity ---- + +// A corpus of unsigned 8-bit codes (element_size U8). Codes span [1, 255], +// including values > 127, so the GPU path is exercised on codes that a signed +// int8 read would misinterpret. +SparseVectors make_random_corpus_u8(size_t n_docs, size_t dim, + size_t nnz_per_doc, unsigned seed) { + SparseVectors vectors({.element_size = U8, .dimension = dim}); + std::mt19937 gen(seed); + std::uniform_int_distribution term_dist(0, static_cast(dim - 1)); + std::uniform_int_distribution code_dist(1, 255); + for (size_t d = 0; d < n_docs; ++d) { + std::vector chosen; + for (size_t k = 0; k < nnz_per_doc; ++k) { + chosen.push_back(static_cast(term_dist(gen))); + } + std::ranges::sort(chosen); + chosen.erase(std::unique(chosen.begin(), chosen.end()), chosen.end()); + std::vector terms; + std::vector codes; + for (term_t t : chosen) { + terms.push_back(t); + codes.push_back(static_cast(code_dist(gen))); + } + vectors.add_vector(terms.data(), terms.size(), codes.data(), + codes.size()); + } + return vectors; +} + +// CPU reference for 8-bit assignment: int64 dot of unsigned codes (exact), +// strict-greater argmax, centroids skipped -- mirrors the CPU int64 path the +// GPU uint8 kernel must reproduce. +std::vector> cpu_reference_assign_u8( + const SparseVectors* vectors, const std::vector& docs, + std::vector> clusters) { + const offset_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const uint8_t* values = vectors->values_data(); + const size_t n_clusters = clusters.size(); + for (size_t i = 0; i < docs.size(); ++i) { + const auto dense = vectors->get_dense_vector(docs[i]); + int64_t best = std::numeric_limits::lowest(); + size_t best_j = 0; + bool is_centroid = false; + for (size_t j = 0; j < n_clusters; ++j) { + const idx_t c = clusters[j].front(); + if (docs[i] == c) { + is_centroid = true; + break; + } + const offset_t start = indptr[c]; + const size_t len = indptr[c + 1] - start; + int64_t score = 0; + for (size_t t = 0; t < len; ++t) { + score += static_cast(dense[indices[start + t]]) * + static_cast(values[start + t]); + } + if (score > best) { + best = score; + best_j = j; + } + } + if (!is_centroid) { + clusters[best_j].push_back(docs[i]); + } + } + return clusters; +} + +TEST(GpuClusterAssignerTest, MatchesCpuReference8Bit) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kNumDocs = 6000; + constexpr size_t kDim = 2000; + constexpr size_t kNnz = 40; + constexpr size_t kNumClusters = 32; + + SparseVectors vectors = make_random_corpus_u8(kNumDocs, kDim, kNnz, 1234); + std::vector docs(kNumDocs); + std::iota(docs.begin(), docs.end(), 0); + + auto expected = cpu_reference_assign_u8(&vectors, docs, + seed_clusters(docs, kNumClusters)); + auto actual = seed_clusters(docs, kNumClusters); + GpuClusterAssigner::instance().assign(&vectors, docs, actual); + + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < expected.size(); ++j) { + EXPECT_EQ(actual[j], expected[j]) << "cluster " << j << " differs"; + } +} + +TEST(GpuClusterAssignerTest, AutoPathViaMapDocsMatchesReference8Bit) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kNumDocs = 5000; + constexpr size_t kDim = 1500; + constexpr size_t kNnz = 30; + constexpr size_t kNumClusters = 16; + + SparseVectors vectors = make_random_corpus_u8(kNumDocs, kDim, kNnz, 99); + std::vector docs(kNumDocs); + std::iota(docs.begin(), docs.end(), 0); + + auto expected = cpu_reference_assign_u8(&vectors, docs, + seed_clusters(docs, kNumClusters)); + auto actual = seed_clusters(docs, kNumClusters); + map_docs_to_clusters(&vectors, docs, actual); // routes to GPU (U8) + + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < expected.size(); ++j) { + EXPECT_EQ(actual[j], expected[j]) << "cluster " << j << " differs"; + } +} + +// CPU reference for the 8-bit max-pool: per-term max of unsigned codes (as +// float, exact for 0..255), plus the sum of maxes, terms ascending. +void cpu_reference_maxpool_u8(const SparseVectors* vectors, + const std::vector& doc_ids, + std::vector& terms, + std::vector& values, float& sum) { + const offset_t* indptr = vectors->indptr_data(); + const term_t* indices = vectors->indices_data(); + const uint8_t* vals = vectors->values_data(); + std::map m; + for (idx_t d : doc_ids) { + for (offset_t j = indptr[d]; j < indptr[d + 1]; ++j) { + auto& v = m[indices[j]]; + v = std::max(v, vals[j]); + } + } + sum = 0.0F; + terms.clear(); + values.clear(); + for (auto& [t, v] : m) { + terms.push_back(t); + values.push_back(static_cast(v)); + sum += static_cast(v); + } +} + +TEST(GpuClusterAssignerTest, SummarizeListMaxpool8BitMatchesCpu) { + if (!GpuSummarizer::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kNumDocs = 8000; + constexpr size_t kDim = 3000; + constexpr size_t kNnz = 50; + SparseVectors vectors = make_random_corpus_u8(kNumDocs, kDim, kNnz, 7); + + std::vector flat_docs; + std::vector offsets = {0}; + const std::vector sizes = {5, 137, 1, 900, 42, 2000}; + idx_t d = 0; + for (size_t sz : sizes) { + for (size_t k = 0; k < sz && d < static_cast(kNumDocs); ++k) { + flat_docs.push_back(d++); + } + offsets.push_back(static_cast(flat_docs.size())); + } + const size_t n_clusters = offsets.size() - 1; + + std::vector gpu_out; + ASSERT_TRUE(GpuSummarizer::instance().summarize_list( + &vectors, flat_docs.data(), offsets.data(), n_clusters, gpu_out)); + ASSERT_EQ(gpu_out.size(), n_clusters); + + for (size_t b = 0; b < n_clusters; ++b) { + std::vector cluster(flat_docs.begin() + offsets[b], + flat_docs.begin() + offsets[b + 1]); + std::vector exp_terms; + std::vector exp_vals; + float exp_sum = 0.0F; + cpu_reference_maxpool_u8(&vectors, cluster, exp_terms, exp_vals, + exp_sum); + + const auto& gc = gpu_out[b]; + std::vector order(gc.terms.size()); + std::iota(order.begin(), order.end(), 0); + std::ranges::sort(order, [&](size_t a, size_t c) { + return gc.terms[a] < gc.terms[c]; + }); + ASSERT_EQ(gc.terms.size(), exp_terms.size()) << "cluster " << b; + for (size_t i = 0; i < order.size(); ++i) { + EXPECT_EQ(gc.terms[order[i]], exp_terms[i]) + << "cluster " << b << " term " << i; + EXPECT_FLOAT_EQ(gc.values[order[i]], exp_vals[i]) + << "cluster " << b << " value " << i; + } + EXPECT_NEAR(gc.sum, exp_sum, exp_sum * 1e-5F + 1e-5F) + << "cluster " << b; + } +} + +// Regression for the int64 accumulator. With a large vocabulary a single doc +// can overlap a centroid in enough max-value terms that the true int64 dot +// exceeds INT32_MAX (255*255*overlap): an int32 accumulator wraps -- here the +// larger score wraps negative -- and the argmax picks the wrong cluster, while +// int64 matches the CPU reference. dim=40000 with a ~34000-term all-255 row is +// what crosses the limit (34000*65025 = 2.21e9 > INT32_MAX). +TEST(GpuClusterAssignerTest, Assign8BitInt64AccumulatorNoOverflow) { + if (!GpuClusterAssigner::available()) { + GTEST_SKIP() << "No CUDA-capable GPU available"; + } + constexpr size_t kDim = 40000; + constexpr term_t kWide = 34000; // 34000 * 255*255 = 2.21e9 > INT32_MAX + constexpr term_t kNarrow = 100; + + SparseVectors vectors({.element_size = U8, .dimension = kDim}); + auto add_row = [&](term_t hi) { // terms [0, hi) all at code 255 + std::vector terms; + std::vector codes; + for (term_t t = 0; t < hi; ++t) { + terms.push_back(t); + codes.push_back(255); + } + vectors.add_vector(terms.data(), terms.size(), codes.data(), + codes.size()); + }; + add_row(kWide); // doc 0: centroid A -- wide overlap with doc 2 + add_row(kNarrow); // doc 1: centroid B -- narrow overlap with doc 2 + add_row(kWide); // doc 2: assigned; true argmax is A, but B under int32 + + std::vector docs = {0, 1, 2}; + auto expected = + cpu_reference_assign_u8(&vectors, docs, seed_clusters(docs, 2)); + auto actual = seed_clusters(docs, 2); + GpuClusterAssigner::instance().assign(&vectors, docs, actual); + + ASSERT_EQ(actual.size(), expected.size()); + for (size_t j = 0; j < expected.size(); ++j) { + EXPECT_EQ(actual[j], expected[j]) << "cluster " << j << " differs"; + } + // The int64 dot puts doc 2 with centroid A (cluster 0); an int32 sum would + // wrap A's score negative and mis-assign doc 2 to cluster 1. + EXPECT_EQ(actual[0], (std::vector{0, 2})); + EXPECT_EQ(actual[1], (std::vector{1})); +} + } // namespace } // namespace nsparse::detail