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