range_search on GPU

[tejaskhot]

Is range_search not implemented on GPUs for any index?

The tutorial examples with IndexFlatL2 don't work with range_search when using GPUs.

[chan4cc]

I facing the same problem which give the error message like this:

~/anaconda2/lib/python2.7/site-packages/faiss/swigfaiss_gpu.pyc in range_search(self, n, x, radius, result) 1457 1458 def range_search(self, n, x, radius, result): -> 1459 return _swigfaiss_gpu.Index_range_search(self, n, x, radius, result) 1460 1461 def assign(self, n, x, labels, k=1):

RuntimeError: Error in virtual void faiss::Index::range_search(faiss::Index::idx_t, const float*, float, faiss::RangeSearchResult*) const at Index.cpp:32: range search not implemented

[mdouze]

Range search is implemented only for a limited number of index types. https://github.com/facebookresearch/faiss/wiki/Special-operations-on-indexes#range-search Implementing it for more index types is not high-priority for us, but of course we welcome PRs in that direction. CC: @wickedfoo

[tejaskhot]

@mdouze Yes, but the wiki does not state if for those index types for which it is implemented (IndexFlat, IndexIVFFlat), it is compatible to run on GPU or not. It is possible to push these index types to the GPU using faiss.index_cpu_to_gpu and that works fine for a k nearest neighbors search, but doesn't for range_search. I am wondering if that particular part (GPU compatibility of range_search for those indices) is not implemented or if it needs code to be written in a different manner than the wiki example in order to work?

[mdouze]

It is not implemented and range search needs a very different memory management pattern, so it is not trivial to add it. CC: @wickedfoo

[wickedfoo]

This wouldn't be easy to implement, and I don't have the time to do it, but may be a long term feature to add.

[silcowitz]

+1 for this feature :)

[siddhantsomani]

Can someone help me understand the output of range_search? lims, D, I ? I understand that lims are used to fetch records but what are D, I?

[mdouze]

see here: https://github.com/facebookresearch/faiss/wiki/Special-operations-on-indexes#range-search

[ljstrnadiii]

Although I have yet to measure the negative effects, I have implemented range search for gpu index like this:

 def radius_search(index, x, thresh, nprobe, max_k=2048, init_k=100, gpu=False):

    if hasattr(index, 'nprobe'):
        index.nprobe = nprobe
    else:
        faiss.downcast_index(index).nprobe = nprobe 
        
    if not gpu:   
        lims, D, I = index.range_search(x, thresh=thresh)

        # faiss doesn't sort
        for i in range(len(lims)-1):
            sorted_idx = np.argsort(D[lims[i]: lims[i+1]])
            D[lims[i]: lims[i+1]] = D[lims[i]: lims[i+1]][sorted_idx]
            I[lims[i]: lims[i+1]] = I[lims[i]: lims[i+1]][sorted_idx]

    else:
        ii = defaultdict(list)
        dd = defaultdict(list)

        k = init_k
        D, I = index.search(x, k)
        n = len(D)
        r, c = np.where(D <= thresh)
        actual_r = r

        while True:
            for row, col, act_r in zip(r, c, actual_r):
                ii[act_r].append(I[row, col])
                dd[act_r].append(D[row, col])

            continue_idx = [rr for rr, v in dd.items() if (len(v) == k)]
            if len(continue_idx) == 0:
                break

            k *= 2
            if k >= max_k:
                break

            D, I = index.search(x[continue_idx], k=k)

            prev_k = int(k/2)
            D = D[:, prev_k:]
            I = I[:, prev_k:]
            r, c = np.where(D <= thresh)
            _, cts = np.unique(r, return_counts=True)
            actual_r = np.repeat(continue_idx, cts)

        sorted_rows = range(n)
        lims = np.concatenate([np.array([0]), np.cumsum([len(dd[i]) for i in sorted_rows])]).astype(np.uint64)
        D = np.array([sl for l in [dd[r] for r in sorted_rows] for sl in l])
        I = np.array([sl for l in [ii[r] for r in sorted_rows] for sl in l])

    return lims, D, I

I am seeing comparable speed on cpu between gpu=True and False, but have not tested it much. Hopefully this isn't too painful to look at for the OGs who actually right faiss code. Caution to anyone who might copy and paste this -- this is a hack.

[wickedfoo]

def radius_search(index, x, thresh, nprobe, max_k=3200, init_k=100, gpu=False):

note, the maximum k-select value on the GPU is 2048 at present

[ljstrnadiii]

Ah that's right. Thanks for the reminder.

[mdouze]

Still no GPU implementation but here is some code that runs k-NN search on GPU and only does range search on CPU if the k-nn search did not return enough results:

https://github.com/facebookresearch/faiss/blob/main/contrib/exhaustive_search.py#L54