chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,228 @@
+/**
+ *  Copyright (c) 2024, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *  Copyright (c) 2023 by Contributors
+ * @file cnumpy.h
+ * @brief Numpy File Fetecher class.
+ */
+
+#ifdef HAVE_LIBRARY_LIBURING
+#include <liburing.h>
+#endif  // HAVE_LIBRARY_LIBURING
+
+#include <graphbolt/async.h>
+#include <torch/script.h>
+
+#include <cassert>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <cuda/std/semaphore>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace graphbolt {
+namespace storage {
+
+namespace {
+#ifdef HAVE_LIBRARY_LIBURING
+struct io_uring_queue_destroyer {
+  int num_thread_;
+  void operator()(::io_uring* queues) {
+    if (!queues) return;
+    for (int t = 0; t < num_thread_; t++) {
+      // IO queue exit.
+      ::io_uring_queue_exit(&queues[t]);
+    }
+    delete[] queues;
+  }
+};
+#endif  // HAVE_LIBRARY_LIBURING
+}  // namespace
+
+/**
+ * @brief Disk Numpy Fetecher class.
+ */
+class OnDiskNpyArray : public torch::CustomClassHolder {
+  // No user will need more than 1024 io_uring queues.
+  using counting_semaphore_t = ::cuda::std::counting_semaphore<1024>;
+
+ public:
+  static constexpr int kGroupSize = 256;
+
+  /** @brief Default constructor. */
+  OnDiskNpyArray() = default;
+
+  /**
+   * @brief Constructor with given file path and data type.
+   * @param path Path to the on disk numpy file.
+   * @param dtype Data type of numpy array.
+   *
+   * @return OnDiskNpyArray
+   */
+  OnDiskNpyArray(
+      std::string filename, torch::ScalarType dtype,
+      const std::vector<int64_t>& shape, torch::optional<int64_t> num_threads);
+
+  /** @brief Create a disk feature fetcher from numpy file. */
+  static c10::intrusive_ptr<OnDiskNpyArray> Create(
+      std::string path, torch::ScalarType dtype,
+      const std::vector<int64_t>& shape, torch::optional<int64_t> num_threads);
+
+  /** @brief Deconstructor. */
+  ~OnDiskNpyArray();
+
+  /**
+   * @brief Parses the header of a numpy file to extract feature information.
+   **/
+  void ParseNumpyHeader();
+
+  /**
+   * @brief Read disk numpy file based on given index and transform to
+   * tensor.
+   */
+  c10::intrusive_ptr<Future<torch::Tensor>> IndexSelect(torch::Tensor index);
+
+#ifdef HAVE_LIBRARY_LIBURING
+  /**
+   * @brief Index-select operation on an on-disk numpy array using IO Uring for
+   * asynchronous I/O.
+   *
+   * This function performs index-select operation on an on-disk numpy array. It
+   * uses IO Uring for asynchronous I/O to efficiently read data from disk. The
+   * input tensor 'index' specifies the indices of features to select. The
+   * function reads features corresponding to the indices from the disk and
+   * returns a new tensor containing the selected features.
+   *
+   * @param index A 1D tensor containing the indices of features to select.
+   * @return A tensor containing the selected features.
+   * @throws std::runtime_error If index is out of range.
+   */
+  c10::intrusive_ptr<Future<torch::Tensor>> IndexSelectIOUring(
+      torch::Tensor index);
+
+  torch::Tensor IndexSelectIOUringImpl(torch::Tensor index);
+
+#endif  // HAVE_LIBRARY_LIBURING
+ private:
+  int64_t ReadBufferSizePerThread() const {
+    return (aligned_length_ + block_size_) * kGroupSize * 8;
+  }
+
+  char* ReadBuffer(int thread_id) const {
+    auto read_buffer_void_ptr = read_tensor_.data_ptr();
+    size_t read_buffer_size = read_tensor_.numel();
+    auto read_buffer = reinterpret_cast<char*>(std::align(
+        block_size_, ReadBufferSizePerThread() * num_thread_,
+        read_buffer_void_ptr, read_buffer_size));
+    TORCH_CHECK(read_buffer, "read_buffer allocation failed!");
+    return read_buffer + ReadBufferSizePerThread() * thread_id;
+  }
+
+  const std::string filename_;  // Path to numpy file.
+  int file_description_;        // File description.
+  int64_t block_size_;          // Block size of the opened file.
+  int64_t prefix_len_;          // Length of head data in numpy file.
+  const std::vector<int64_t>
+      feature_dim_;                // Shape of features, e.g. {N,M,K,L}.
+  const torch::ScalarType dtype_;  // Feature data type.
+  const int64_t feature_size_;     // Number of bytes of feature size.
+  int64_t aligned_length_;         // Aligned feature_size.
+  int num_thread_;                 // Default thread number.
+  torch::Tensor read_tensor_;      // Provides temporary read buffer.
+
+#ifdef HAVE_LIBRARY_LIBURING
+
+  static inline std::once_flag
+      call_once_flag_;            // Protect initialization of below.
+  static inline int num_queues_;  // Number of queues.
+  static inline std::unique_ptr<::io_uring[], io_uring_queue_destroyer>
+      io_uring_queue_;  // io_uring queue.
+  static inline counting_semaphore_t semaphore_{
+      0};  // Control access to the io_uring queues.
+  static inline std::mutex available_queues_mtx_;  // available_queues_ mutex.
+  static inline std::vector<int> available_queues_;
+
+  /**
+   * @brief This class is meant to distribute the available read buffers and the
+   * statically declared io_uring queues among the worker threads.
+   */
+  class QueueAndBufferAcquirer {
+   public:
+    class UniqueQueue {
+     public:
+      UniqueQueue(int thread_id) : thread_id_(thread_id) {}
+      UniqueQueue(const UniqueQueue&) = delete;
+      UniqueQueue& operator=(const UniqueQueue&) = delete;
+
+      /**
+       * @brief Returns the queue back to the pool.
+       */
+      ~UniqueQueue() {
+        {
+          // We give back the slot we used.
+          std::lock_guard lock(available_queues_mtx_);
+          available_queues_.push_back(thread_id_);
+        }
+        semaphore_.release();
+      }
+
+      /**
+       * @brief Returns the raw io_uring queue.
+       */
+      ::io_uring& get() const { return io_uring_queue_[thread_id_]; }
+
+     private:
+      int thread_id_;
+    };
+
+    QueueAndBufferAcquirer(OnDiskNpyArray* array) : array_(array) {
+      semaphore_.acquire();
+    }
+
+    ~QueueAndBufferAcquirer() {
+      // If none of the worker threads acquire the semaphore, we make sure to
+      // release the ticket taken in the constructor.
+      if (!entering_first_.test_and_set(std::memory_order_relaxed)) {
+        semaphore_.release();
+      }
+    }
+
+    /**
+     * @brief Returns the secured io_uring queue and the read buffer as a pair.
+     * The raw io_uring queue can be accessed by calling `.get()` on the
+     * returned UniqueQueue object.
+     *
+     * @note The returned UniqueQueue object manages the lifetime of the
+     * io_uring queue. Its destructor returns the queue back to the pool.
+     */
+    std::pair<UniqueQueue, char*> get() {
+      // We consume a slot from the semaphore to use a queue.
+      if (entering_first_.test_and_set(std::memory_order_relaxed)) {
+        semaphore_.acquire();
+      }
+      const auto thread_id = [&] {
+        std::lock_guard lock(available_queues_mtx_);
+        TORCH_CHECK(!available_queues_.empty());
+        const auto thread_id = available_queues_.back();
+        available_queues_.pop_back();
+        return thread_id;
+      }();
+      return {
+          std::piecewise_construct, std::make_tuple(thread_id),
+          std::make_tuple(array_->ReadBuffer(thread_id))};
+    }
+
+   private:
+    const OnDiskNpyArray* array_;
+    std::atomic_flag entering_first_ = ATOMIC_FLAG_INIT;
+  };
+
+#endif  // HAVE_LIBRARY_LIBURING
+};
+
+}  // namespace storage
+}  // namespace graphbolt