GPU-based brute force computation with better memory management
(vectors, query_vectors, dimension, top_k=10)
| 263 | |
| 264 | |
| 265 | def _compute_brute_force_gpu(vectors, query_vectors, dimension, top_k=10): |
| 266 | """GPU-based brute force computation with better memory management""" |
| 267 | results = [] |
| 268 | n_vectors = len(vectors) |
| 269 | |
| 270 | # Check available GPU memory |
| 271 | mempool = cp.get_default_memory_pool() |
| 272 | available_memory = mempool.free_bytes() + mempool.total_bytes() |
| 273 | print(f"Available GPU memory: {available_memory / 1e9:.2f} GB") |
| 274 | |
| 275 | # Process in smaller chunks if needed |
| 276 | max_vectors_per_chunk = min(100000, n_vectors) # Limit chunk size |
| 277 | |
| 278 | print("Building dataset sparse matrix on GPU in chunks...") |
| 279 | |
| 280 | for chunk_start in tqdm(range(0, n_vectors, max_vectors_per_chunk)): |
| 281 | chunk_end = min(chunk_start + max_vectors_per_chunk, n_vectors) |
| 282 | chunk_vectors = vectors[chunk_start:chunk_end] |
| 283 | |
| 284 | # Build chunk matrix |
| 285 | data_list = [] |
| 286 | indices_list = [] |
| 287 | indptr = [0] |
| 288 | |
| 289 | for vec in chunk_vectors: |
| 290 | data_list.append(vec["values"]) |
| 291 | indices_list.append(vec["indices"]) |
| 292 | indptr.append(indptr[-1] + len(vec["indices"])) |
| 293 | |
| 294 | # Use smaller data types to save memory |
| 295 | data_gpu = cp.concatenate([cp.array(v, dtype=cp.float32) for v in data_list]) |
| 296 | indices_gpu = cp.concatenate( |
| 297 | [cp.array(v, dtype=cp.int32) for v in indices_list] |
| 298 | ) |
| 299 | indptr_gpu = cp.array(indptr, dtype=cp.int32) |
| 300 | |
| 301 | # Create CSR matrix for this chunk |
| 302 | A_chunk = csr_matrix( |
| 303 | (data_gpu, indices_gpu, indptr_gpu), shape=(len(chunk_vectors), dimension) |
| 304 | ) |
| 305 | |
| 306 | # Process queries against this chunk |
| 307 | for query in query_vectors: |
| 308 | # Build query vector |
| 309 | q_data = cp.array(query["values"], dtype=cp.float32) |
| 310 | q_indices = cp.array(query["indices"], dtype=cp.int32) |
| 311 | q_indptr = cp.array([0, len(q_data)], dtype=cp.int32) |
| 312 | |
| 313 | Q = csr_matrix((q_data, q_indices, q_indptr), shape=(1, dimension)) |
| 314 | |
| 315 | # Compute dot products for this chunk |
| 316 | scores_chunk = A_chunk.dot(Q.T).toarray().flatten() |
| 317 | |
| 318 | # Update results for this query |
| 319 | if chunk_start == 0: |
| 320 | # Initialize results for this query |
| 321 | query_idx = query_vectors.index(query) |
| 322 | if query_idx >= len(results): |
no test coverage detected