faiss 0.3.4 → 0.4.0

data/vendor/faiss/faiss/impl/mapped_io.cpp

@@ -0,0 +1,313 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <stdio.h>
+#include <string.h>
+
+#ifdef __linux__
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+#include <unistd.h>
+
+#elif defined(_WIN32)
+
+#include <Windows.h> // @manual
+#include <io.h>      // @manual
+
+#endif
+
+#include <cstring>
+
+#include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/mapped_io.h>
+
+namespace faiss {
+
+#ifdef __linux__
+
+struct MmappedFileMappingOwner::PImpl {
+    void* ptr = nullptr;
+    size_t ptr_size = 0;
+
+    PImpl(const std::string& filename) {
+        auto f = std::unique_ptr<FILE, decltype(&fclose)>(
+                fopen(filename.c_str(), "r"), &fclose);
+        FAISS_THROW_IF_NOT_FMT(
+                f.get(),
+                "could not open %s for reading: %s",
+                filename.c_str(),
+                strerror(errno));
+
+        // get the size
+        struct stat s;
+        int status = fstat(fileno(f.get()), &s);
+        FAISS_THROW_IF_NOT_FMT(
+                status >= 0, "fstat() failed: %s", strerror(errno));
+
+        const size_t filesize = s.st_size;
+
+        void* address = mmap(
+                nullptr, filesize, PROT_READ, MAP_SHARED, fileno(f.get()), 0);
+        FAISS_THROW_IF_NOT_FMT(
+                address != nullptr, "could not mmap(): %s", strerror(errno));
+
+        // btw, fd can be closed here
+
+        madvise(address, filesize, MADV_RANDOM);
+
+        // save it
+        ptr = address;
+        ptr_size = filesize;
+    }
+
+    PImpl(FILE* f) {
+        // get the size
+        struct stat s;
+        int status = fstat(fileno(f), &s);
+        FAISS_THROW_IF_NOT_FMT(
+                status >= 0, "fstat() failed: %s", strerror(errno));
+
+        const size_t filesize = s.st_size;
+
+        void* address =
+                mmap(nullptr, filesize, PROT_READ, MAP_SHARED, fileno(f), 0);
+        FAISS_THROW_IF_NOT_FMT(
+                address != nullptr, "could not mmap(): %s", strerror(errno));
+
+        // btw, fd can be closed here
+
+        madvise(address, filesize, MADV_RANDOM);
+
+        // save it
+        ptr = address;
+        ptr_size = filesize;
+    }
+
+    ~PImpl() {
+        // todo: check for an error
+        munmap(ptr, ptr_size);
+    }
+};
+
+#elif defined(_WIN32)
+
+struct MmappedFileMappingOwner::PImpl {
+    void* ptr = nullptr;
+    size_t ptr_size = 0;
+    HANDLE mapping_handle = INVALID_HANDLE_VALUE;
+
+    PImpl(const std::string& filename) {
+        HANDLE file_handle = CreateFile(
+                filename.c_str(),
+                GENERIC_READ,
+                FILE_SHARE_READ,
+                nullptr,
+                OPEN_EXISTING,
+                0,
+                nullptr);
+        if (file_handle == INVALID_HANDLE_VALUE) {
+            const auto error = GetLastError();
+            FAISS_THROW_FMT(
+                    "could not open the file, %s (error %d)",
+                    filename.c_str(),
+                    error);
+        }
+
+        // get the size of the file
+        LARGE_INTEGER len_li;
+        if (GetFileSizeEx(file_handle, &len_li) == 0) {
+            const auto error = GetLastError();
+
+            CloseHandle(file_handle);
+
+            FAISS_THROW_FMT(
+                    "could not get the file size, %s (error %d)",
+                    filename.c_str(),
+                    error);
+        }
+
+        // create a mapping
+        mapping_handle = CreateFileMapping(
+                file_handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
+        if (mapping_handle == 0) {
+            const auto error = GetLastError();
+
+            CloseHandle(file_handle);
+
+            FAISS_THROW_FMT(
+                    "could not create a file mapping, %s (error %d)",
+                    filename.c_str(),
+                    error);
+        }
+        CloseHandle(file_handle);
+
+        char* data =
+                (char*)MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, 0);
+        if (data == nullptr) {
+            const auto error = GetLastError();
+
+            CloseHandle(mapping_handle);
+            mapping_handle = INVALID_HANDLE_VALUE;
+
+            FAISS_THROW_FMT(
+                    "could not get map the file, %s (error %d)",
+                    filename.c_str(),
+                    error);
+        }
+
+        ptr = data;
+        ptr_size = len_li.QuadPart;
+    }
+
+    PImpl(FILE* f) {
+        // obtain a HANDLE from a FILE
+        const int fd = _fileno(f);
+        if (fd == -1) {
+            // no good
+            FAISS_THROW_FMT("could not get a HANDLE");
+        }
+
+        HANDLE file_handle = (HANDLE)_get_osfhandle(fd);
+        if (file_handle == INVALID_HANDLE_VALUE) {
+            FAISS_THROW_FMT("could not get an OS HANDLE");
+        }
+
+        // get the size of the file
+        LARGE_INTEGER len_li;
+        if (GetFileSizeEx(file_handle, &len_li) == 0) {
+            const auto error = GetLastError();
+            FAISS_THROW_FMT("could not get the file size (error %d)", error);
+        }
+
+        // create a mapping
+        mapping_handle = CreateFileMapping(
+                file_handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
+        if (mapping_handle == 0) {
+            const auto error = GetLastError();
+            FAISS_THROW_FMT(
+                    "could not create a file mapping, (error %d)", error);
+        }
+
+        // the handle is provided externally, so this is not our business
+        //   to close file_handle.
+
+        char* data =
+                (char*)MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, 0);
+        if (data == nullptr) {
+            const auto error = GetLastError();
+
+            CloseHandle(mapping_handle);
+            mapping_handle = INVALID_HANDLE_VALUE;
+
+            FAISS_THROW_FMT("could not get map the file, (error %d)", error);
+        }
+
+        ptr = data;
+        ptr_size = len_li.QuadPart;
+    }
+
+    ~PImpl() {
+        if (mapping_handle != INVALID_HANDLE_VALUE) {
+            UnmapViewOfFile(ptr);
+            CloseHandle(mapping_handle);
+
+            mapping_handle = INVALID_HANDLE_VALUE;
+            ptr = nullptr;
+        }
+    }
+};
+
+#else
+
+struct MmappedFileMappingOwner::PImpl {
+    void* ptr = nullptr;
+    size_t ptr_size = 0;
+
+    PImpl(const std::string& filename) {
+        FAISS_THROW_MSG("Not implemented");
+    }
+
+    PImpl(FILE* f) {
+        FAISS_THROW_MSG("Not implemented");
+    }
+};
+
+#endif
+
+MmappedFileMappingOwner::MmappedFileMappingOwner(const std::string& filename) {
+    p_impl = std::make_unique<MmappedFileMappingOwner::PImpl>(filename);
+}
+
+MmappedFileMappingOwner::MmappedFileMappingOwner(FILE* f) {
+    p_impl = std::make_unique<MmappedFileMappingOwner::PImpl>(f);
+}
+
+MmappedFileMappingOwner::~MmappedFileMappingOwner() = default;
+
+//
+void* MmappedFileMappingOwner::data() const {
+    return p_impl->ptr;
+}
+
+size_t MmappedFileMappingOwner::size() const {
+    return p_impl->ptr_size;
+}
+
+MappedFileIOReader::MappedFileIOReader(
+        const std::shared_ptr<MmappedFileMappingOwner>& owner)
+        : mmap_owner(owner) {}
+
+// this operation performs a copy
+size_t MappedFileIOReader::operator()(void* ptr, size_t size, size_t nitems) {
+    if (size * nitems == 0) {
+        return 0;
+    }
+
+    char* ptr_c = nullptr;
+
+    const size_t actual_nitems = this->mmap((void**)&ptr_c, size, nitems);
+    if (actual_nitems > 0) {
+        memcpy(ptr, ptr_c, size * actual_nitems);
+    }
+
+    return actual_nitems;
+}
+
+// this operation returns a mmapped address, owned by mmap_owner
+size_t MappedFileIOReader::mmap(void** ptr, size_t size, size_t nitems) {
+    if (size == 0) {
+        return nitems;
+    }
+
+    size_t actual_size = size * nitems;
+    if (pos + size * nitems > mmap_owner->size()) {
+        actual_size = mmap_owner->size() - pos;
+    }
+
+    size_t actual_nitems = (actual_size + size - 1) / size;
+    if (actual_nitems == 0) {
+        return 0;
+    }
+
+    // get an address
+    *ptr = (void*)(reinterpret_cast<const char*>(mmap_owner->data()) + pos);
+
+    // alter pos
+    pos += size * actual_nitems;
+
+    return actual_nitems;
+}
+
+int MappedFileIOReader::filedescriptor() {
+    // todo
+    return -1;
+}
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/mapped_io.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include <faiss/impl/io.h>
+#include <faiss/impl/maybe_owned_vector.h>
+
+namespace faiss {
+
+// holds a memory-mapped region over a file
+struct MmappedFileMappingOwner : public MaybeOwnedVectorOwner {
+    MmappedFileMappingOwner(const std::string& filename);
+    MmappedFileMappingOwner(FILE* f);
+    ~MmappedFileMappingOwner();
+
+    void* data() const;
+    size_t size() const;
+
+    struct PImpl;
+    std::unique_ptr<PImpl> p_impl;
+};
+
+// A deserializer that supports memory-mapped files.
+// All de-allocations should happen as soon as the index gets destroyed,
+//   after all underlying the MaybeOwnerVector objects are destroyed.
+struct MappedFileIOReader : IOReader {
+    std::shared_ptr<MmappedFileMappingOwner> mmap_owner;
+
+    size_t pos = 0;
+
+    MappedFileIOReader(const std::shared_ptr<MmappedFileMappingOwner>& owner);
+
+    // perform a copy
+    size_t operator()(void* ptr, size_t size, size_t nitems) override;
+    // perform a quasi-read that returns a mmapped address, owned by mmap_owner,
+    //   and updates the position
+    size_t mmap(void** ptr, size_t size, size_t nitems);
+
+    int filedescriptor() override;
+};
+
+} // namespace faiss

data/vendor/faiss/faiss/impl/maybe_owned_vector.h

@@ -0,0 +1,316 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <vector>
+
+#include <faiss/impl/FaissAssert.h>
+
+namespace faiss {
+
+// An interface for an owner of a MaybeOwnedVector.
+struct MaybeOwnedVectorOwner {
+    virtual ~MaybeOwnedVectorOwner() = default;
+};
+
+// a container that either works as std::vector<T> that owns its own memory,
+//    or as a view of a memory buffer, with a known size
+template <typename T>
+struct MaybeOwnedVector {
+    using value_type = T;
+    using self_type = MaybeOwnedVector<T>;
+    using iterator = typename std::vector<T>::iterator;
+    using const_iterator = typename std::vector<T>::const_iterator;
+    using size_type = typename std::vector<T>::size_type;
+
+    bool is_owned = true;
+
+    // this one is used if is_owned == true
+    std::vector<T> owned_data;
+
+    // these three are used if is_owned == false
+    T* view_data = nullptr;
+    // the number of T elements
+    size_t view_size = 0;
+    // who owns the data.
+    // This field can be nullptr, and it is present ONLY in order
+    //   to avoid possible tricky memory / resource leaks.
+    std::shared_ptr<MaybeOwnedVectorOwner> owner;
+
+    // points either to view_data, or to owned.data()
+    T* c_ptr = nullptr;
+    // uses either view_size, or owned.size();
+    size_t c_size = 0;
+
+    MaybeOwnedVector() = default;
+    MaybeOwnedVector(const size_t initial_size) {
+        is_owned = true;
+
+        owned_data.resize(initial_size);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+    }
+
+    explicit MaybeOwnedVector(const std::vector<T>& vec)
+            : faiss::MaybeOwnedVector<T>(vec.size()) {
+        if (vec.size() > 0) {
+            memcpy(owned_data.data(), vec.data(), sizeof(T) * vec.size());
+        }
+    }
+
+    MaybeOwnedVector(const MaybeOwnedVector& other) {
+        is_owned = other.is_owned;
+        owned_data = other.owned_data;
+
+        view_data = other.view_data;
+        view_size = other.view_size;
+        owner = other.owner;
+
+        if (is_owned) {
+            c_ptr = owned_data.data();
+            c_size = owned_data.size();
+        } else {
+            c_ptr = view_data;
+            c_size = view_size;
+        }
+    }
+
+    MaybeOwnedVector(MaybeOwnedVector&& other) {
+        is_owned = other.is_owned;
+        owned_data = std::move(other.owned_data);
+
+        view_data = other.view_data;
+        view_size = other.view_size;
+        owner = std::move(other.owner);
+        other.owner = nullptr;
+
+        if (is_owned) {
+            c_ptr = owned_data.data();
+            c_size = owned_data.size();
+        } else {
+            c_ptr = view_data;
+            c_size = view_size;
+        }
+    }
+
+    MaybeOwnedVector& operator=(const MaybeOwnedVector& other) {
+        if (this == &other) {
+            return *this;
+        }
+
+        // create a copy
+        MaybeOwnedVector cloned(other);
+        // swap
+        swap(*this, cloned);
+
+        return *this;
+    }
+
+    MaybeOwnedVector& operator=(MaybeOwnedVector&& other) {
+        if (this == &other) {
+            return *this;
+        }
+
+        // moved
+        MaybeOwnedVector moved(std::move(other));
+        // swap
+        swap(*this, moved);
+
+        return *this;
+    }
+
+    MaybeOwnedVector(std::vector<T>&& other) {
+        is_owned = true;
+
+        owned_data = std::move(other);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+    }
+
+    static MaybeOwnedVector create_view(
+            void* address,
+            const size_t n_elements,
+            const std::shared_ptr<MaybeOwnedVectorOwner>& owner) {
+        MaybeOwnedVector vec;
+        vec.is_owned = false;
+        vec.view_data = reinterpret_cast<T*>(address);
+        vec.view_size = n_elements;
+        vec.owner = owner;
+
+        vec.c_ptr = vec.view_data;
+        vec.c_size = vec.view_size;
+
+        return vec;
+    }
+
+    const T* data() const {
+        return c_ptr;
+    }
+
+    T* data() {
+        return c_ptr;
+    }
+
+    size_t size() const {
+        return c_size;
+    }
+
+    size_t byte_size() const {
+        return c_size * sizeof(T);
+    }
+
+    T& operator[](const size_t idx) {
+        return c_ptr[idx];
+    }
+
+    const T& operator[](const size_t idx) const {
+        return c_ptr[idx];
+    }
+
+    T& at(size_type pos) {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.at(pos);
+    }
+
+    const T& at(size_type pos) const {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.at(pos);
+    }
+
+    iterator begin() {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.begin();
+    }
+
+    const_iterator begin() const {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.begin();
+    }
+
+    iterator end() {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.end();
+    }
+
+    const_iterator end() const {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        return owned_data.end();
+    }
+
+    iterator erase(const_iterator begin, const_iterator end) {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        auto result = owned_data.erase(begin, end);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+
+        return result;
+    }
+
+    template <class InputIt>
+    iterator insert(const_iterator pos, InputIt first, InputIt last) {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        auto result = owned_data.insert(pos, first, last);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+
+        return result;
+    }
+
+    void clear() {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        owned_data.clear();
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+    }
+
+    void resize(const size_t new_size) {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        owned_data.resize(new_size);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+    }
+
+    void resize(const size_t new_size, const value_type v) {
+        FAISS_ASSERT_MSG(
+                is_owned,
+                "This operation cannot be performed on a viewed vector");
+
+        owned_data.resize(new_size, v);
+        c_ptr = owned_data.data();
+        c_size = owned_data.size();
+    }
+
+    friend void swap(self_type& a, self_type& b) {
+        std::swap(a.is_owned, b.is_owned);
+        std::swap(a.owned_data, b.owned_data);
+        std::swap(a.view_data, b.view_data);
+        std::swap(a.view_size, b.view_size);
+        std::swap(a.owner, b.owner);
+        std::swap(a.c_ptr, b.c_ptr);
+        std::swap(a.c_size, b.c_size);
+    }
+};
+
+template <typename T>
+struct is_maybe_owned_vector : std::false_type {};
+
+template <typename T>
+struct is_maybe_owned_vector<MaybeOwnedVector<T>> : std::true_type {};
+
+template <typename T>
+inline constexpr bool is_maybe_owned_vector_v = is_maybe_owned_vector<T>::value;
+
+template <typename T>
+bool operator==(
+        const MaybeOwnedVector<T>& lhs,
+        const MaybeOwnedVector<T>& rhs) {
+    return lhs.size() == rhs.size() &&
+            !memcmp(lhs.data(), rhs.data(), lhs.byte_size());
+}
+
+template <typename T>
+bool operator!=(
+        const MaybeOwnedVector<T>& lhs,
+        const MaybeOwnedVector<T>& rhs) {
+    return !(lhs == rhs);
+}
+
+} // namespace faiss