faiss/impl/lattice_Zn.cpp

/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

// -*- c++ -*-

#include <faiss/impl/lattice_Zn.h>

#include <cassert>
#include <cmath>
#include <cstdlib>
#include <cstring>

#include <algorithm>
#include <queue>
#include <unordered_map>
#include <unordered_set>

#include <faiss/impl/platform_macros.h>
#include <faiss/utils/distances.h>

namespace faiss {

/********************************************
 * small utility functions
 ********************************************/

namespace {

inline float sqr(float x) {
    return x * x;
}

typedef std::vector<float> point_list_t;

struct Comb {
    std::vector<uint64_t> tab; // Pascal's triangle
    int nmax;

    explicit Comb(int nmax) : nmax(nmax) {
        tab.resize(nmax * nmax, 0);
        tab[0] = 1;
        for (int i = 1; i < nmax; i++) {
            tab[i * nmax] = 1;
            for (int j = 1; j <= i; j++) {
                tab[i * nmax + j] =
                        tab[(i - 1) * nmax + j] + tab[(i - 1) * nmax + (j - 1)];
            }
        }
    }

    uint64_t operator()(int n, int p) const {
        assert(n < nmax && p < nmax);
        if (p > n)
            return 0;
        return tab[n * nmax + p];
    }
};

Comb comb(100);

// compute combinations of n integer values <= v that sum up to total (squared)
point_list_t sum_of_sq(float total, int v, int n, float add = 0) {
    if (total < 0) {
        return point_list_t();
    } else if (n == 1) {
        while (sqr(v + add) > total)
            v--;
        if (sqr(v + add) == total) {
            return point_list_t(1, v + add);
        } else {
            return point_list_t();
        }
    } else {
        point_list_t res;
        while (v >= 0) {
            point_list_t sub_points =
                    sum_of_sq(total - sqr(v + add), v, n - 1, add);
            for (size_t i = 0; i < sub_points.size(); i += n - 1) {
                res.push_back(v + add);
                for (int j = 0; j < n - 1; j++) {
                    res.push_back(sub_points[i + j]);
                }
            }
            v--;
        }
        return res;
    }
}

int decode_comb_1(uint64_t* n, int k1, int r) {
    while (comb(r, k1) > *n) {
        r--;
    }
    *n -= comb(r, k1);
    return r;
}

// optimized version for < 64 bits
uint64_t repeats_encode_64(
        const std::vector<Repeat>& repeats,
        int dim,
        const float* c) {
    uint64_t coded = 0;
    int nfree = dim;
    uint64_t code = 0, shift = 1;
    for (auto r = repeats.begin(); r != repeats.end(); ++r) {
        int rank = 0, occ = 0;
        uint64_t code_comb = 0;
        uint64_t tosee = ~coded;
        for (;;) {
            // directly jump to next available slot.
            int i = __builtin_ctzll(tosee);
            tosee &= ~(uint64_t{1} << i);
            if (c[i] == r->val) {
                code_comb += comb(rank, occ + 1);
                occ++;
                coded |= uint64_t{1} << i;
                if (occ == r->n)
                    break;
            }
            rank++;
        }
        uint64_t max_comb = comb(nfree, r->n);
        code += shift * code_comb;
        shift *= max_comb;
        nfree -= r->n;
    }
    return code;
}

void repeats_decode_64(
        const std::vector<Repeat>& repeats,
        int dim,
        uint64_t code,
        float* c) {
    uint64_t decoded = 0;
    int nfree = dim;
    for (auto r = repeats.begin(); r != repeats.end(); ++r) {
        uint64_t max_comb = comb(nfree, r->n);
        uint64_t code_comb = code % max_comb;
        code /= max_comb;

        int occ = 0;
        int rank = nfree;
        int next_rank = decode_comb_1(&code_comb, r->n, rank);
        uint64_t tosee = ((uint64_t{1} << dim) - 1) ^ decoded;
        for (;;) {
            int i = 63 - __builtin_clzll(tosee);
            tosee &= ~(uint64_t{1} << i);
            rank--;
            if (rank == next_rank) {
                decoded |= uint64_t{1} << i;
                c[i] = r->val;
                occ++;
                if (occ == r->n)
                    break;
                next_rank = decode_comb_1(&code_comb, r->n - occ, next_rank);
            }
        }
        nfree -= r->n;
    }
}

} // anonymous namespace

Repeats::Repeats(int dim, const float* c) : dim(dim) {
    for (int i = 0; i < dim; i++) {
        int j = 0;
        for (;;) {
            if (j == repeats.size()) {
                repeats.push_back(Repeat{c[i], 1});
                break;
            }
            if (repeats[j].val == c[i]) {
                repeats[j].n++;
                break;
            }
            j++;
        }
    }
}

uint64_t Repeats::count() const {
    uint64_t accu = 1;
    int remain = dim;
    for (int i = 0; i < repeats.size(); i++) {
        accu *= comb(remain, repeats[i].n);
        remain -= repeats[i].n;
    }
    return accu;
}

// version with a bool vector that works for > 64 dim
uint64_t Repeats::encode(const float* c) const {
    if (dim < 64) {
        return repeats_encode_64(repeats, dim, c);
    }
    std::vector<bool> coded(dim, false);
    int nfree = dim;
    uint64_t code = 0, shift = 1;
    for (auto r = repeats.begin(); r != repeats.end(); ++r) {
        int rank = 0, occ = 0;
        uint64_t code_comb = 0;
        for (int i = 0; i < dim; i++) {
            if (!coded[i]) {
                if (c[i] == r->val) {
                    code_comb += comb(rank, occ + 1);
                    occ++;
                    coded[i] = true;
                    if (occ == r->n)
                        break;
                }
                rank++;
            }
        }
        uint64_t max_comb = comb(nfree, r->n);
        code += shift * code_comb;
        shift *= max_comb;
        nfree -= r->n;
    }
    return code;
}

void Repeats::decode(uint64_t code, float* c) const {
    if (dim < 64) {
        repeats_decode_64(repeats, dim, code, c);
        return;
    }

    std::vector<bool> decoded(dim, false);
    int nfree = dim;
    for (auto r = repeats.begin(); r != repeats.end(); ++r) {
        uint64_t max_comb = comb(nfree, r->n);
        uint64_t code_comb = code % max_comb;
        code /= max_comb;

        int occ = 0;
        int rank = nfree;
        int next_rank = decode_comb_1(&code_comb, r->n, rank);
        for (int i = dim - 1; i >= 0; i--) {
            if (!decoded[i]) {
                rank--;
                if (rank == next_rank) {
                    decoded[i] = true;
                    c[i] = r->val;
                    occ++;
                    if (occ == r->n)
                        break;
                    next_rank =
                            decode_comb_1(&code_comb, r->n - occ, next_rank);
                }
            }
        }
        nfree -= r->n;
    }
}

/********************************************
 * EnumeratedVectors functions
 ********************************************/

void EnumeratedVectors::encode_multi(size_t n, const float* c, uint64_t* codes)
        const {
#pragma omp parallel if (n > 1000)
    {
#pragma omp for
        for (int i = 0; i < n; i++) {
            codes[i] = encode(c + i * dim);
        }
    }
}

void EnumeratedVectors::decode_multi(size_t n, const uint64_t* codes, float* c)
        const {
#pragma omp parallel if (n > 1000)
    {
#pragma omp for
        for (int i = 0; i < n; i++) {
            decode(codes[i], c + i * dim);
        }
    }
}

void EnumeratedVectors::find_nn(
        size_t nc,
        const uint64_t* codes,
        size_t nq,
        const float* xq,
        int64_t* labels,
        float* distances) {
    for (size_t i = 0; i < nq; i++) {
        distances[i] = -1e20;
        labels[i] = -1;
    }

    std::vector<float> c(dim);
    for (size_t i = 0; i < nc; i++) {
        uint64_t code = codes[nc];
        decode(code, c.data());
        for (size_t j = 0; j < nq; j++) {
            const float* x = xq + j * dim;
            float dis = fvec_inner_product(x, c.data(), dim);
            if (dis > distances[j]) {
                distances[j] = dis;
                labels[j] = i;
            }
        }
    }
}

/**********************************************************
 * ZnSphereSearch
 **********************************************************/

ZnSphereSearch::ZnSphereSearch(int dim, int r2) : dimS(dim), r2(r2) {
    voc = sum_of_sq(r2, int(ceil(sqrt(r2)) + 1), dim);
    natom = voc.size() / dim;
}

float ZnSphereSearch::search(const float* x, float* c) const {
    std::vector<float> tmp(dimS * 2);
    std::vector<int> tmp_int(dimS);
    return search(x, c, tmp.data(), tmp_int.data());
}

float ZnSphereSearch::search(
        const float* x,
        float* c,
        float* tmp,   // size 2 *dim
        int* tmp_int, // size dim
        int* ibest_out) const {
    int dim = dimS;
    assert(natom > 0);
    int* o = tmp_int;
    float* xabs = tmp;
    float* xperm = tmp + dim;

    // argsort
    for (int i = 0; i < dim; i++) {
        o[i] = i;
        xabs[i] = fabsf(x[i]);
    }
    std::sort(o, o + dim, [xabs](int a, int b) { return xabs[a] > xabs[b]; });
    for (int i = 0; i < dim; i++) {
        xperm[i] = xabs[o[i]];
    }
    // find best
    int ibest = -1;
    float dpbest = -100;
    for (int i = 0; i < natom; i++) {
        float dp = fvec_inner_product(voc.data() + i * dim, xperm, dim);
        if (dp > dpbest) {
            dpbest = dp;
            ibest = i;
        }
    }
    // revert sort
    const float* cin = voc.data() + ibest * dim;
    for (int i = 0; i < dim; i++) {
        c[o[i]] = copysignf(cin[i], x[o[i]]);
    }
    if (ibest_out) {
        *ibest_out = ibest;
    }
    return dpbest;
}

void ZnSphereSearch::search_multi(
        int n,
        const float* x,
        float* c_out,
        float* dp_out) {
#pragma omp parallel if (n > 1000)
    {
#pragma omp for
        for (int i = 0; i < n; i++) {
            dp_out[i] = search(x + i * dimS, c_out + i * dimS);
        }
    }
}

/**********************************************************
 * ZnSphereCodec
 **********************************************************/

ZnSphereCodec::ZnSphereCodec(int dim, int r2)
        : ZnSphereSearch(dim, r2), EnumeratedVectors(dim) {
    nv = 0;
    for (int i = 0; i < natom; i++) {
        Repeats repeats(dim, &voc[i * dim]);
        CodeSegment cs(repeats);
        cs.c0 = nv;
        Repeat& br = repeats.repeats.back();
        cs.signbits = br.val == 0 ? dim - br.n : dim;
        code_segments.push_back(cs);
        nv += repeats.count() << cs.signbits;
    }

    uint64_t nvx = nv;
    code_size = 0;
    while (nvx > 0) {
        nvx >>= 8;
        code_size++;
    }
}

uint64_t ZnSphereCodec::search_and_encode(const float* x) const {
    std::vector<float> tmp(dim * 2);
    std::vector<int> tmp_int(dim);
    int ano; // atom number
    std::vector<float> c(dim);
    search(x, c.data(), tmp.data(), tmp_int.data(), &ano);
    uint64_t signs = 0;
    std::vector<float> cabs(dim);
    int nnz = 0;
    for (int i = 0; i < dim; i++) {
        cabs[i] = fabs(c[i]);
        if (c[i] != 0) {
            if (c[i] < 0) {
                signs |= uint64_t{1} << nnz;
            }
            nnz++;
        }
    }
    const CodeSegment& cs = code_segments[ano];
    assert(nnz == cs.signbits);
    uint64_t code = cs.c0 + signs;
    code += cs.encode(cabs.data()) << cs.signbits;
    return code;
}

uint64_t ZnSphereCodec::encode(const float* x) const {
    return search_and_encode(x);
}

void ZnSphereCodec::decode(uint64_t code, float* c) const {
    int i0 = 0, i1 = natom;
    while (i0 + 1 < i1) {
        int imed = (i0 + i1) / 2;
        if (code_segments[imed].c0 <= code)
            i0 = imed;
        else
            i1 = imed;
    }
    const CodeSegment& cs = code_segments[i0];
    code -= cs.c0;
    uint64_t signs = code;
    code >>= cs.signbits;
    cs.decode(code, c);

    int nnz = 0;
    for (int i = 0; i < dim; i++) {
        if (c[i] != 0) {
            if (signs & (1UL << nnz)) {
                c[i] = -c[i];
            }
            nnz++;
        }
    }
}

/**************************************************************
 * ZnSphereCodecRec
 **************************************************************/

uint64_t ZnSphereCodecRec::get_nv(int ld, int r2a) const {
    return all_nv[ld * (r2 + 1) + r2a];
}

uint64_t ZnSphereCodecRec::get_nv_cum(int ld, int r2t, int r2a) const {
    return all_nv_cum[(ld * (r2 + 1) + r2t) * (r2 + 1) + r2a];
}

void ZnSphereCodecRec::set_nv_cum(int ld, int r2t, int r2a, uint64_t cum) {
    all_nv_cum[(ld * (r2 + 1) + r2t) * (r2 + 1) + r2a] = cum;
}

ZnSphereCodecRec::ZnSphereCodecRec(int dim, int r2)
        : EnumeratedVectors(dim), r2(r2) {
    log2_dim = 0;
    while (dim > (1 << log2_dim)) {
        log2_dim++;
    }
    assert(dim == (1 << log2_dim) || !"dimension must be a power of 2");

    all_nv.resize((log2_dim + 1) * (r2 + 1));
    all_nv_cum.resize((log2_dim + 1) * (r2 + 1) * (r2 + 1));

    for (int r2a = 0; r2a <= r2; r2a++) {
        int r = int(sqrt(r2a));
        if (r * r == r2a) {
            all_nv[r2a] = r == 0 ? 1 : 2;
        } else {
            all_nv[r2a] = 0;
        }
    }

    for (int ld = 1; ld <= log2_dim; ld++) {
        for (int r2sub = 0; r2sub <= r2; r2sub++) {
            uint64_t nv = 0;
            for (int r2a = 0; r2a <= r2sub; r2a++) {
                int r2b = r2sub - r2a;
                set_nv_cum(ld, r2sub, r2a, nv);
                nv += get_nv(ld - 1, r2a) * get_nv(ld - 1, r2b);
            }
            all_nv[ld * (r2 + 1) + r2sub] = nv;
        }
    }
    nv = get_nv(log2_dim, r2);

    uint64_t nvx = nv;
    code_size = 0;
    while (nvx > 0) {
        nvx >>= 8;
        code_size++;
    }

    int cache_level = std::min(3, log2_dim - 1);
    decode_cache_ld = 0;
    assert(cache_level <= log2_dim);
    decode_cache.resize((r2 + 1));

    for (int r2sub = 0; r2sub <= r2; r2sub++) {
        int ld = cache_level;
        uint64_t nvi = get_nv(ld, r2sub);
        std::vector<float>& cache = decode_cache[r2sub];
        int dimsub = (1 << cache_level);
        cache.resize(nvi * dimsub);
        std::vector<float> c(dim);
        uint64_t code0 = get_nv_cum(cache_level + 1, r2, r2 - r2sub);
        for (int i = 0; i < nvi; i++) {
            decode(i + code0, c.data());
            memcpy(&cache[i * dimsub],
                   c.data() + dim - dimsub,
                   dimsub * sizeof(*c.data()));
        }
    }
    decode_cache_ld = cache_level;
}

uint64_t ZnSphereCodecRec::encode(const float* c) const {
    return encode_centroid(c);
}

uint64_t ZnSphereCodecRec::encode_centroid(const float* c) const {
    std::vector<uint64_t> codes(dim);
    std::vector<int> norm2s(dim);
    for (int i = 0; i < dim; i++) {
        if (c[i] == 0) {
            codes[i] = 0;
            norm2s[i] = 0;
        } else {
            int r2i = int(c[i] * c[i]);
            norm2s[i] = r2i;
            codes[i] = c[i] >= 0 ? 0 : 1;
        }
    }
    int dim2 = dim / 2;
    for (int ld = 1; ld <= log2_dim; ld++) {
        for (int i = 0; i < dim2; i++) {
            int r2a = norm2s[2 * i];
            int r2b = norm2s[2 * i + 1];

            uint64_t code_a = codes[2 * i];
            uint64_t code_b = codes[2 * i + 1];

            codes[i] = get_nv_cum(ld, r2a + r2b, r2a) +
                    code_a * get_nv(ld - 1, r2b) + code_b;
            norm2s[i] = r2a + r2b;
        }
        dim2 /= 2;
    }
    return codes[0];
}

void ZnSphereCodecRec::decode(uint64_t code, float* c) const {
    std::vector<uint64_t> codes(dim);
    std::vector<int> norm2s(dim);
    codes[0] = code;
    norm2s[0] = r2;

    int dim2 = 1;
    for (int ld = log2_dim; ld > decode_cache_ld; ld--) {
        for (int i = dim2 - 1; i >= 0; i--) {
            int r2sub = norm2s[i];
            int i0 = 0, i1 = r2sub + 1;
            uint64_t codei = codes[i];
            const uint64_t* cum =
                    &all_nv_cum[(ld * (r2 + 1) + r2sub) * (r2 + 1)];
            while (i1 > i0 + 1) {
                int imed = (i0 + i1) / 2;
                if (cum[imed] <= codei)
                    i0 = imed;
                else
                    i1 = imed;
            }
            int r2a = i0, r2b = r2sub - i0;
            codei -= cum[r2a];
            norm2s[2 * i] = r2a;
            norm2s[2 * i + 1] = r2b;

            uint64_t code_a = codei / get_nv(ld - 1, r2b);
            uint64_t code_b = codei % get_nv(ld - 1, r2b);

            codes[2 * i] = code_a;
            codes[2 * i + 1] = code_b;
        }
        dim2 *= 2;
    }

    if (decode_cache_ld == 0) {
        for (int i = 0; i < dim; i++) {
            if (norm2s[i] == 0) {
                c[i] = 0;
            } else {
                float r = sqrt(norm2s[i]);
                assert(r * r == norm2s[i]);
                c[i] = codes[i] == 0 ? r : -r;
            }
        }
    } else {
        int subdim = 1 << decode_cache_ld;
        assert((dim2 * subdim) == dim);

        for (int i = 0; i < dim2; i++) {
            const std::vector<float>& cache = decode_cache[norm2s[i]];
            assert(codes[i] < cache.size());
            memcpy(c + i * subdim,
                   &cache[codes[i] * subdim],
                   sizeof(*c) * subdim);
        }
    }
}

// if not use_rec, instanciate an arbitrary harmless znc_rec
ZnSphereCodecAlt::ZnSphereCodecAlt(int dim, int r2)
        : ZnSphereCodec(dim, r2),
          use_rec((dim & (dim - 1)) == 0),
          znc_rec(use_rec ? dim : 8, use_rec ? r2 : 14) {}

uint64_t ZnSphereCodecAlt::encode(const float* x) const {
    if (!use_rec) {
        // it's ok if the vector is not normalized
        return ZnSphereCodec::encode(x);
    } else {
        // find nearest centroid
        std::vector<float> centroid(dim);
        search(x, centroid.data());
        return znc_rec.encode(centroid.data());
    }
}

void ZnSphereCodecAlt::decode(uint64_t code, float* c) const {
    if (!use_rec) {
        ZnSphereCodec::decode(code, c);
    } else {
        znc_rec.decode(code, c);
    }
}

} // namespace faiss