Faiss(14)：IndexIVFPQ的CPU search过程分析

1. 说明

之前分析过了faiss 在GPU中的search过程，这里分析一下IndexIVFPQ在CPU中的search过程，即不将index拷贝到GPU中。

2. 过程分析

2.1 python接口

CPU search的python接口与GPU的完全一致，没有差别。

D, I = gpu_index.search(xq_t[x],top_k)

2.2 faiss core

IndexIVF::search

因为IndexIVFPQ没有override search，所以在实际运行过程中会调用父类IndexIVF中的实现。

void IndexIVF::search (idx_t n, const float *x, idx_t k,
                         float *distances, idx_t *labels) const
{
    std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
    std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);

    // 1. quantizer中的粗量搜索
    quantizer->search (n, x, nprobe, coarse_dis.get(), idx.get());

    // 2. prefetch_lists函数为空
    invlists->prefetch_lists (idx.get(), n * nprobe);

    // 3. 预分配搜索
    search_preassigned (n, x, k, idx.get(), coarse_dis.get(),
                        distances, labels, false);
}

search过程主要包含三个部分：

1. quantizer->search
   量化器搜索聚类中心，该过程在quantizer实例中进行，计算所有(nlist个)聚类中心与每一条搜索向量的距离，并从中找出最近的nprobe个聚类中心，输出到idx和coarse_dis中。
   所以这里的idx和coarse_dis是大小为 nnprobe 的二维数组，分别存放nnprobe个向量label和与原向量的距离。

2. invlists->prefetch_lists
   这里没有用到这类功能，所以实际调用的该函数为空。

3. search_preassigned
   经过第一步计算出粗聚类中心向量后，在这里进行二次计算，即在选定的聚类里再次计算出top_k个近邻向量。由于在add时已经对向量进行了预计算形成残差，所以这里只要进行向量和运算就可以了。
   在实际测试时发现，search的主要时间消耗是在这里，占比达到96%以上，所以针对这一过程进行进一步分析。
   
   

IndexIVF::search_preassigned

    /** search a set of vectors, that are pre-quantized by the IVF
     *  quantizer. Fill in the corresponding heaps with the query
     *  results. The default implementation uses InvertedListScanners
     *  to do the search.
     *
     * @param n      nb of vectors to query
     * @param x      query vectors, size nx * d
     * @param assign coarse quantization indices, size nx * nprobe
     * @param centroid_dis
     *               distances to coarse centroids, size nx * nprobe
     * @param distance
     *               output distances, size n * k
     * @param labels output labels, size n * k
     * @param store_pairs store inv list index + inv list offset
     *                     instead in upper/lower 32 bit of result,
     *                     instead of ids (used for reranking).
     * @param params used to override the object's search parameters
     */
void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
                                   const idx_t *keys,
                                   const float *coarse_dis ,
                                   float *distances, idx_t *labels,
                                   bool store_pairs,
                                   const IVFSearchParameters *params) const
{
    // max_codes是默认值0
    long nprobe = params ? params->nprobe : this->nprobe;
    long max_codes = params ? params->max_codes : this->max_codes;

    size_t nlistv = 0, ndis = 0, nheap = 0;

    // 根据计算类型定义堆
    using HeapForIP = CMin<float, idx_t>;
    using HeapForL2 = CMax<float, idx_t>;

    bool interrupt = false;

    // don't start parallel section if single query
    bool do_parallel =
        parallel_mode == 0 ? n > 1 :
        parallel_mode == 1 ? nprobe > 1 :
        nprobe * n > 1;

#pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
    {
        InvertedListScanner *scanner = get_InvertedListScanner(store_pairs);
        ScopeDeleter1<InvertedListScanner> del(scanner);

        /****************************************************
         * Actual loops, depending on parallel_mode
         ****************************************************/

        if (parallel_mode == 0) {

#pragma omp for
            for (size_t i = 0; i < n; i++) {

                if (interrupt) {
                    continue;
                }
                
                // 在inverted_list中设置搜索的起始点
                scanner->set_query (x + i * d);
                // 根据i设置distances和labels的地址
                float * simi = distances + i * k;
                idx_t * idxi = labels + i * k;

                init_result (simi, idxi);

                long nscan = 0;

                // 依次在nprobe个聚类中进行搜索
                for (size_t ik = 0; ik < nprobe; ik++) {

                    nscan += scan_one_list (
                         keys [i * nprobe + ik],
                         coarse_dis[i * nprobe + ik],
                         simi, idxi
                    );

                    if (max_codes && nscan >= max_codes) {
                        break;
                    }
                }

                ndis += nscan;
                //对搜索结果进行排序
                reorder_result (simi, idxi);

                if (InterruptCallback::is_interrupted ()) {
                    interrupt = true;
                }

            } // parallel for
        } else if (parallel_mode == 1) {
            std::vector <idx_t> local_idx (k);
            std::vector <float> local_dis (k);

            for (size_t i = 0; i < n; i++) {
                scanner->set_query (x + i * d);
                init_result (local_dis.data(), local_idx.data());

#pragma omp for schedule(dynamic)
                for (size_t ik = 0; ik < nprobe; ik++) {
                    ndis += scan_one_list
                        (keys [i * nprobe + ik],
                         coarse_dis[i * nprobe + ik],
                         local_dis.data(), local_idx.data());

                    // can't do the test on max_codes
                }
                // merge thread-local results

                float * simi = distances + i * k;
                idx_t * idxi = labels + i * k;
#pragma omp single
                init_result (simi, idxi);

#pragma omp barrier
// 将各个线程产生的堆合并到结果堆中，临界访问
#pragma omp critical
                {
                    if (metric_type == METRIC_INNER_PRODUCT) {
                        heap_addn<HeapForIP>
                            (k, simi, idxi,
                             local_dis.data(), local_idx.data(), k);
                    } else {
                        heap_addn<HeapForL2>
                            (k, simi, idxi,
                             local_dis.data(), local_idx.data(), k);
                    }
                }
#pragma omp barrier
#pragma omp single
                reorder_result (simi, idxi);
            }
        } else {
            FAISS_THROW_FMT ("parallel_mode %d not supported\n",
                             parallel_mode);
        }
    } // parallel section

    if (interrupt) {
        FAISS_THROW_MSG ("computation interrupted");
    }

    indexIVF_stats.nq += n;
    indexIVF_stats.nlist += nlistv;
    indexIVF_stats.ndis += ndis;
    indexIVF_stats.nheap_updates += nheap;

}

这里的params使用默认值nullptr.

流程
从代码中可以看出，虽然根据parallel_mode值的不同，程序处理上会有部分差异，但差异主要体现在并行运算的时间点和内容，主要流程是一致的：

1. 首先根据原向量下标确定要输出的distances和labels的地址；
2. 在nprobe个倒序列表中进行搜索，并对搜索结果进行排序

堆的使用
Faiss使用堆对搜索结果进行排序，不同的搜索类型可能使用不同的堆，在L2的范式搜索中使用大根堆进行排序。

do_parallel
Faiss可以直接使用OpemMP的并行运算指令，do_parallel是打开并行运算的标志值，在以下三种情况下为1：

• parallel_mode == 0 && n > 1
• parallel_mode == 1 && nprobe > 1
• parallel_mode == 2 && nprobe * n > 1

parallel_mode
该值确定采用何种并行模式进行查询。0表示在查询时开启并行，1表示在inverted_list计算残差时开启并行，2表示在上述两个阶段都使用并行模式。

并行运算
Faiss默认添加了对OpenMP的支持，具体命令待添加

intialize + reorder resule heap
这部分原本是定义在search_preassigned 函数体内的函数，这里把它们单拎出来。顾名思义，这里是分别对堆进行初始化和排序，因为搜索结果通常是两个列表，distances和labels，所以这里也是两个堆。

 auto init_result = [&](float *simi, idx_t *idxi) {
            if (metric_type == METRIC_INNER_PRODUCT) {
                heap_heapify<HeapForIP> (k, simi, idxi);
            } else {
                heap_heapify<HeapForL2> (k, simi, idxi);
            }
        };

auto reorder_result = [&] (float *simi, idx_t *idxi) {
    if (metric_type == METRIC_INNER_PRODUCT) {
        heap_reorder<HeapForIP> (k, simi, idxi);
    } else {
        heap_reorder<HeapForL2> (k, simi, idxi);
    }
};

loop probes分析
每个线程在这一过程中会循环遍历nprobe个聚类，计算残差，以最后得出top个最近邻向量。代码内容如下：

// loop over probes
                for (size_t ik = 0; ik < nprobe; ik++) {
                    nscan += scan_one_list (
                         keys [i * nprobe + ik],
                         coarse_dis[i * nprobe + ik],
                         simi, idxi
                    );

                    if (max_codes && nscan >= max_codes) {
                        break;
                    }
                }

其中scan_one_list是函数内定义的函数，主要完成下列工作:

1. invlist->list_size：计算倒序列表的大小，为空则直接跳过；
2. scanner->set_list：根据key值，从coarse_dis_i列表中找到入口地址；
3. ScopedCodes：根据key值和invlists的内容生成ccode；
4. sids->get：重置并获取id号；
5. scanner->scan_codes：扫描一组代码，计算到当前查询的距离，并更新结果堆。

scan_one_list
这个函数的功能是在单个的聚类中进行搜索。

/*
* key: nprobe中的invetred list编号
* coarse_dis_i: key对应的聚类中心的distance
* simi：存放distances结果的堆
* idxi： 存放labels结果的堆
*/
auto scan_one_list = [&] (idx_t key, float coarse_dis_i,
                          float *simi, idx_t *idxi) {

    if (key < 0) {
        // not enough centroids for multiprobe
        return (size_t)0;
    }
    FAISS_THROW_IF_NOT_FMT (key < (idx_t) nlist,
                            "Invalid key=%ld nlist=%ld\n",
                            key, nlist);

    size_t list_size = invlists->list_size(key);

    // don't waste time on empty lists
    if (list_size == 0) {
        return (size_t)0;
    }

    scanner->set_list (key, coarse_dis_i);

    nlistv++;

    InvertedLists::ScopedCodes scodes (invlists, key);

    std::unique_ptr<InvertedLists::ScopedIds> sids;
    const Index::idx_t * ids = nullptr;

    if (!store_pairs)  {
        sids.reset (new InvertedLists::ScopedIds (invlists, key));
        ids = sids->get();
    }

    nheap += scanner->scan_codes (list_size, scodes.get(),
                                  ids, simi, idxi, k);

    return list_size;
}

这个函数的绝大部分时间消耗在scanner->scan_codes内。

scanner
从代码看，scanner是一个搜索引擎的实例，用于在InvertedList中进行搜索的具体实现。值得单独分析。
所有的线程都会单独生成一个自己的scanner。

InvertedListScanner *scanner = get_InvertedListScanner(store_pairs);

scanner->scan_codes
scanner是struct IVFPQScanner结构体的实例，定义在IndexIVFPQ.h中，scan_codes函数内容如下：

size_t scan_codes (size_t ncode,
                       const uint8_t *codes,
                       const idx_t *ids,
                       float *heap_sim, idx_t *heap_ids,
                       size_t k) const override
    {
        KnnSearchResults<C> res = {
            /* key */      this->key,
            /* ids */      this->store_pairs ? nullptr : ids,
            /* k */        k,
            /* heap_sim */ heap_sim,
            /* heap_ids */ heap_ids,
            /* nup */      0
        };

        if (this->polysemous_ht > 0) {
            assert(precompute_mode == 2);
            this->scan_list_polysemous (ncode, codes, res);
        } else if (precompute_mode == 2) {
            this->scan_list_with_table (ncode, codes, res);
        } else if (precompute_mode == 1) {
            this->scan_list_with_pointer (ncode, codes, res);
        } else if (precompute_mode == 0) {
            this->scan_on_the_fly_dist (ncode, codes, res);
        } else {
            FAISS_THROW_MSG("bad precomp mode");
        }
        return res.nup;
    }

实际运行中polysemous_ht为0，precompute_mode为2，故之后调用scan_list_with_table函数。

scan_list_with_table

函数内容如下：

template<class SearchResultType>
    void scan_list_with_table (size_t ncode, const uint8_t *codes,
                               SearchResultType & res) const
    {
        for (size_t j = 0; j < ncode; j++) {
            float dis = dis0;
            const float *tab = sim_table;
            for (size_t m = 0; m < pq.M; m++) {
                dis += tab[*codes++];
                tab += pq.ksub;
            }
            res.add(j, dis);
        }
}

程序运行到这里已经进入了IndexIVFPQ的聚类里面进行残差计算（浮点加）。

3. 总结

虽然不够细致，但经过本文的梳理，可以大致看出IndexIVF的搜索的过程，这其中最重要的两个步骤分别与上一篇文档中提到的Product Quantizer和Inverted File System对应，所以说只要搞清楚这两个实例的过程，便能完全了解整个search的流程了。