Initial commit

packages/leann-backend-diskann/third_party/DiskANN/include/aligned_file_reader.h

@@ -0,0 +1,138 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#define MAX_IO_DEPTH 128
+
+#include <vector>
+#include <atomic>
+
+#ifdef __linux__
+#include <fcntl.h>
+#include <libaio.h>
+#include <unistd.h>
+#include <malloc.h>
+typedef io_context_t IOContext;
+#elif __APPLE__
+#include <dispatch/dispatch.h>
+#include <fcntl.h>
+#include <unistd.h>
+
+struct IOContext
+{
+    int fd;
+    dispatch_io_t channel;
+    dispatch_queue_t queue;
+    dispatch_group_t grp;
+};
+#elif _WINDOWS
+#include <Windows.h>
+#include <minwinbase.h>
+#include <malloc.h>
+
+#ifndef USE_BING_INFRA
+struct IOContext
+{
+    HANDLE fhandle = NULL;
+    HANDLE iocp = NULL;
+    std::vector<OVERLAPPED> reqs;
+};
+#else
+#include "IDiskPriorityIO.h"
+#include <atomic>
+// TODO: Caller code is very callous about copying IOContext objects
+// all over the place. MUST verify that it won't cause leaks/logical
+// errors.
+// Because of such callous copying, we have to use ptr->atomic instead
+// of atomic, as atomic is not copyable.
+struct IOContext
+{
+    enum Status
+    {
+        READ_WAIT = 0,
+        READ_SUCCESS,
+        READ_FAILED,
+        PROCESS_COMPLETE
+    };
+
+    std::shared_ptr<ANNIndex::IDiskPriorityIO> m_pDiskIO = nullptr;
+    std::shared_ptr<std::vector<ANNIndex::AsyncReadRequest>> m_pRequests;
+    std::shared_ptr<std::vector<Status>> m_pRequestsStatus;
+
+    // waitonaddress on this memory to wait for IO completion signal
+    // reader should signal this memory after IO completion
+    // TODO: WindowsAlignedFileReader can be modified to take advantage of this
+    //   and can largely share code with the file reader for Bing.
+    mutable volatile long m_completeCount = 0;
+
+    IOContext()
+        : m_pRequestsStatus(new std::vector<Status>()), m_pRequests(new std::vector<ANNIndex::AsyncReadRequest>())
+    {
+        (*m_pRequestsStatus).reserve(MAX_IO_DEPTH);
+        (*m_pRequests).reserve(MAX_IO_DEPTH);
+    }
+};
+#endif
+
+#endif
+
+#include <cstdio>
+#include <mutex>
+#include <thread>
+#include "tsl/robin_map.h"
+#include "utils.h"
+
+// NOTE :: all 3 fields must be 512-aligned
+struct AlignedRead
+{
+    uint64_t offset; // where to read from
+    uint64_t len;    // how much to read
+    void *buf;       // where to read into
+
+    AlignedRead() : offset(0), len(0), buf(nullptr)
+    {
+    }
+
+    AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf)
+    {
+        assert(IS_512_ALIGNED(offset));
+        assert(IS_512_ALIGNED(len));
+        assert(IS_512_ALIGNED(buf));
+        // assert(malloc_usable_size(buf) >= len);
+    }
+};
+
+class AlignedFileReader
+{
+  protected:
+    tsl::robin_map<std::thread::id, IOContext> ctx_map;
+    std::mutex ctx_mut;
+
+  public:
+    // returns the thread-specific context
+    // returns (io_context_t)(-1) if thread is not registered
+    virtual IOContext &get_ctx() = 0;
+
+    virtual ~AlignedFileReader(){};
+
+    // register thread-id for a context
+    virtual void register_thread() = 0;
+    // de-register thread-id for a context
+    virtual void deregister_thread() = 0;
+    virtual void deregister_all_threads() = 0;
+
+    // Open & close ops
+    // Blocking calls
+    virtual void open(const std::string &fname) = 0;
+    virtual void close() = 0;
+
+    // process batch of aligned requests in parallel
+    // NOTE :: blocking call
+    virtual void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx, bool async = false) = 0;
+
+#ifdef USE_BING_INFRA
+    // wait for completion of one request in a batch of requests
+    virtual void wait(IOContext &ctx, int &completedIndex) = 0;
+#endif
+};