chore: import upstream snapshot with attribution

graphbolt/src/cache_policy.cc

@@ -0,0 +1,287 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cache_policy.cc
+ * @brief Cache policy implementation on the CPU.
+ */
+#include "./cache_policy.h"
+
+#include "./utils.h"
+
+namespace graphbolt {
+namespace storage {
+
+template <typename CachePolicy>
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+BaseCachePolicy::QueryImpl(CachePolicy& policy, torch::Tensor keys) {
+  auto positions = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto indices = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto found_ptr_tensor = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto missing_keys = torch::empty_like(
+      keys, keys.options().pinned_memory(utils::is_pinned(keys)));
+  int64_t found_cnt = 0;
+  int64_t missing_cnt = keys.size(0);
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "BaseCachePolicy::Query::DispatchForKeys", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        auto positions_ptr = positions.data_ptr<int64_t>();
+        auto indices_ptr = indices.data_ptr<int64_t>();
+        static_assert(
+            sizeof(CacheKey*) == sizeof(int64_t), "You need 64 bit pointers.");
+        auto found_ptr =
+            reinterpret_cast<CacheKey**>(found_ptr_tensor.data_ptr<int64_t>());
+        auto missing_keys_ptr = missing_keys.data_ptr<index_t>();
+        for (int64_t i = 0; i < keys.size(0); i++) {
+          const auto key = keys_ptr[i];
+          auto cache_key_ptr = policy.Read(key);
+          if (cache_key_ptr) {
+            positions_ptr[found_cnt] = cache_key_ptr->getPos();
+            found_ptr[found_cnt] = cache_key_ptr;
+            indices_ptr[found_cnt++] = i;
+          } else {
+            indices_ptr[--missing_cnt] = i;
+            missing_keys_ptr[missing_cnt] = key;
+          }
+        }
+      }));
+  return {
+      positions.slice(0, 0, found_cnt), indices,
+      missing_keys.slice(0, found_cnt),
+      found_ptr_tensor.slice(0, 0, found_cnt)};
+}
+
+template <typename CachePolicy>
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+BaseCachePolicy::QueryAndReplaceImpl(CachePolicy& policy, torch::Tensor keys) {
+  auto positions = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto indices = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto pointers = torch::empty_like(keys, keys.options().dtype(torch::kInt64));
+  auto missing_keys = torch::empty_like(
+      keys, keys.options().pinned_memory(utils::is_pinned(keys)));
+  int64_t found_cnt = 0;
+  int64_t missing_cnt = keys.size(0);
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "BaseCachePolicy::Replace", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        auto positions_ptr = positions.data_ptr<int64_t>();
+        auto indices_ptr = indices.data_ptr<int64_t>();
+        static_assert(
+            sizeof(CacheKey*) == sizeof(int64_t), "You need 64 bit pointers.");
+        auto pointers_ptr =
+            reinterpret_cast<CacheKey**>(pointers.data_ptr<int64_t>());
+        auto missing_keys_ptr = missing_keys.data_ptr<index_t>();
+        set_t<int64_t> position_set;
+        position_set.reserve(keys.size(0));
+        // Query and Replace combined.
+        for (int64_t i = 0; i < keys.size(0); i++) {
+          const auto key = keys_ptr[i];
+          const auto [it, can_read] = policy.Emplace(key);
+          if (can_read) {
+            auto& cache_key = *it->second;
+            positions_ptr[found_cnt] = cache_key.getPos();
+            pointers_ptr[found_cnt] = &cache_key;
+            indices_ptr[found_cnt++] = i;
+          } else {
+            indices_ptr[--missing_cnt] = i;
+            missing_keys_ptr[missing_cnt] = key;
+            // Ensure that even if an offset is added, it stays negative.
+            auto position = std::numeric_limits<int64_t>::min();
+            CacheKey* cache_key_ptr = nullptr;
+            if (it->second == policy.getMapSentinelValue()) {
+              cache_key_ptr = policy.Insert(it);
+              position = cache_key_ptr->getPos();
+              TORCH_CHECK(
+                  // We check for the uniqueness of the positions.
+                  std::get<1>(position_set.insert(position)),
+                  "Can't insert all, larger cache capacity is needed.");
+            }
+            positions_ptr[missing_cnt] = position;
+            pointers_ptr[missing_cnt] = cache_key_ptr;
+          }
+        }
+      }));
+  return {positions, indices, pointers, missing_keys.slice(0, found_cnt)};
+}
+
+template <typename CachePolicy>
+std::tuple<torch::Tensor, torch::Tensor> BaseCachePolicy::ReplaceImpl(
+    CachePolicy& policy, torch::Tensor keys) {
+  auto positions = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto pointers = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "BaseCachePolicy::Replace", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        auto positions_ptr = positions.data_ptr<int64_t>();
+        static_assert(
+            sizeof(CacheKey*) == sizeof(int64_t), "You need 64 bit pointers.");
+        auto pointers_ptr =
+            reinterpret_cast<CacheKey**>(pointers.data_ptr<int64_t>());
+        set_t<int64_t> position_set;
+        position_set.reserve(keys.size(0));
+        for (int64_t i = 0; i < keys.size(0); i++) {
+          const auto key = keys_ptr[i];
+          // Ensure that even if an offset is added, it stays negative.
+          auto position = std::numeric_limits<int64_t>::min();
+          CacheKey* cache_key_ptr = nullptr;
+          const auto [it, _] = policy.Emplace(key);
+          if (it->second == policy.getMapSentinelValue()) {
+            cache_key_ptr = policy.Insert(it);
+            position = cache_key_ptr->getPos();
+            TORCH_CHECK(
+                // We check for the uniqueness of the positions.
+                std::get<1>(position_set.insert(position)),
+                "Can't insert all, larger cache capacity is needed.");
+          }
+          positions_ptr[i] = position;
+          pointers_ptr[i] = cache_key_ptr;
+        }
+      }));
+  return {positions, pointers};
+}
+
+template <bool write>
+void BaseCachePolicy::ReadingWritingCompletedImpl(torch::Tensor pointers) {
+  static_assert(
+      sizeof(CacheKey*) == sizeof(int64_t), "You need 64 bit pointers.");
+  auto pointers_ptr =
+      reinterpret_cast<CacheKey**>(pointers.data_ptr<int64_t>());
+  for (int64_t i = 0; i < pointers.size(0); i++) {
+    const auto pointer = pointers_ptr[i];
+    if (!write || pointer) {
+      pointer->EndUse<write>();
+    }
+  }
+}
+
+void BaseCachePolicy::ReadingCompleted(torch::Tensor pointers) {
+  ReadingWritingCompletedImpl<false>(pointers);
+}
+
+void BaseCachePolicy::WritingCompleted(torch::Tensor pointers) {
+  ReadingWritingCompletedImpl<true>(pointers);
+}
+
+S3FifoCachePolicy::S3FifoCachePolicy(int64_t capacity)
+    : BaseCachePolicy(capacity),
+      ghost_queue_(capacity - capacity / 10),
+      small_queue_size_target_(capacity / 10),
+      small_queue_size_(0) {
+  TORCH_CHECK(small_queue_size_target_ > 0, "Capacity is not large enough.");
+  ghost_set_.reserve(ghost_queue_.Capacity());
+  key_to_cache_key_.reserve(kCapacityFactor * (capacity + 1));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+S3FifoCachePolicy::Query(torch::Tensor keys) {
+  return QueryImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+S3FifoCachePolicy::QueryAndReplace(torch::Tensor keys) {
+  return QueryAndReplaceImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> S3FifoCachePolicy::Replace(
+    torch::Tensor keys) {
+  return ReplaceImpl(*this, keys);
+}
+
+SieveCachePolicy::SieveCachePolicy(int64_t capacity)
+    // Ensure that queue_ is constructed first before accessing its `.end()`.
+    : BaseCachePolicy(capacity), queue_(), hand_(queue_.end()) {
+  TORCH_CHECK(capacity > 0, "Capacity needs to be positive.");
+  key_to_cache_key_.reserve(kCapacityFactor * (capacity + 1));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+SieveCachePolicy::Query(torch::Tensor keys) {
+  return QueryImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+SieveCachePolicy::QueryAndReplace(torch::Tensor keys) {
+  return QueryAndReplaceImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> SieveCachePolicy::Replace(
+    torch::Tensor keys) {
+  return ReplaceImpl(*this, keys);
+}
+
+LruCachePolicy::LruCachePolicy(int64_t capacity) : BaseCachePolicy(capacity) {
+  TORCH_CHECK(capacity > 0, "Capacity needs to be positive.");
+  key_to_cache_key_.reserve(kCapacityFactor * (capacity + 1));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+LruCachePolicy::Query(torch::Tensor keys) {
+  return QueryImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+LruCachePolicy::QueryAndReplace(torch::Tensor keys) {
+  return QueryAndReplaceImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> LruCachePolicy::Replace(
+    torch::Tensor keys) {
+  return ReplaceImpl(*this, keys);
+}
+
+ClockCachePolicy::ClockCachePolicy(int64_t capacity)
+    : BaseCachePolicy(capacity) {
+  TORCH_CHECK(capacity > 0, "Capacity needs to be positive.");
+  key_to_cache_key_.reserve(kCapacityFactor * (capacity + 1));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+ClockCachePolicy::Query(torch::Tensor keys) {
+  return QueryImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+ClockCachePolicy::QueryAndReplace(torch::Tensor keys) {
+  return QueryAndReplaceImpl(*this, keys);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> ClockCachePolicy::Replace(
+    torch::Tensor keys) {
+  return ReplaceImpl(*this, keys);
+}
+
+}  // namespace storage
+}  // namespace graphbolt

graphbolt/src/cache_policy.h

@@ -0,0 +1,673 @@
+/**
+ *   Copyright (c) 2024, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cache_policy.h
+ * @brief Cache policy implementation on the CPU.
+ */
+#ifndef GRAPHBOLT_CACHE_POLICY_H_
+#define GRAPHBOLT_CACHE_POLICY_H_
+
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+#include <tsl/robin_map.h>
+#include <tsl/robin_set.h>
+
+#include <cuda/std/atomic>
+#include <limits>
+
+#include "./circular_queue.h"
+
+namespace graphbolt {
+namespace storage {
+
+struct CacheKey {
+  auto getKey() const {
+    return (static_cast<int64_t>(key_higher_16_bits_) << 32) +
+           key_lower_32_bits_;
+  }
+
+  CacheKey(int64_t key) : CacheKey(key, std::numeric_limits<int64_t>::min()) {}
+
+  CacheKey(int64_t key, int64_t position)
+      : freq_(0),
+        // EndUse<true>() should be called to reset the reference count.
+        reference_count_(-1),
+        key_higher_16_bits_(key >> 32),
+        key_lower_32_bits_(key),
+        position_in_cache_(position) {
+    TORCH_CHECK(key == getKey());
+    static_assert(sizeof(CacheKey) == 2 * sizeof(int64_t));
+  }
+
+  CacheKey() = default;
+
+  auto getFreq() const { return freq_; }
+
+  auto getPos() const { return position_in_cache_; }
+
+  CacheKey& setPos(int64_t pos) {
+    position_in_cache_ = pos;
+    return *this;
+  }
+
+  CacheKey& Increment() {
+    freq_ = std::min(3, static_cast<int>(freq_ + 1));
+    return *this;
+  }
+
+  CacheKey& Decrement() {
+    freq_ = std::max(0, static_cast<int>(freq_ - 1));
+    return *this;
+  }
+
+  CacheKey& SetFreq() {
+    freq_ = 1;
+    return *this;
+  }
+
+  CacheKey& ResetFreq() {
+    freq_ = 0;
+    return *this;
+  }
+
+  CacheKey& StartRead() {
+    ::cuda::std::atomic_ref ref(reference_count_);
+    // StartRead runs concurrently only with EndUse. EndUse does not need to see
+    // this modification at all. So we can use the relaxed memory order.
+    const auto old_val = ref.fetch_add(1, ::cuda::std::memory_order_relaxed);
+    TORCH_CHECK(
+        old_val < std::numeric_limits<int8_t>::max(),
+        "There are too many in-flight read requests to the same cache entry!");
+    return *this;
+  }
+
+  template <bool write>
+  CacheKey& EndUse() {
+    ::cuda::std::atomic_ref ref(reference_count_);
+    // The EndUse operation needs to synchronize with the InUse operation. So we
+    // have an release-acquire ordering between the two.
+    // https://en.cppreference.com/w/cpp/atomic/memory_order#Release-Acquire_ordering
+    if constexpr (write) {
+      ref.fetch_add(1, ::cuda::std::memory_order_release);
+    } else {
+      ref.fetch_add(-1, ::cuda::std::memory_order_release);
+    }
+    return *this;
+  }
+
+  bool InUse() const {
+    ::cuda::std::atomic_ref ref(reference_count_);
+    // The operations after a call to this function need to happen after the
+    // load operation. Hence the acquire order.
+    return ref.load(::cuda::std::memory_order_acquire);
+  }
+
+  bool BeingWritten() const {
+    ::cuda::std::atomic_ref ref(reference_count_);
+    // The only operation coming after this op is the StartRead operation. Since
+    // StartRead is a refcount increment operation, it is fine if we don't
+    // synchronize with EndUse ops.
+    return ref.load(::cuda::std::memory_order_relaxed) < 0;
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const CacheKey& key_ref) {
+    ::cuda::std::atomic_ref ref(key_ref.reference_count_);
+    return os << '(' << key_ref.getKey() << ", " << key_ref.freq_ << ", "
+              << key_ref.position_in_cache_ << ", " << ref.load() << ")";
+  }
+
+ private:
+  int8_t freq_;
+  // Negative values indicate writing while positive values indicate reading.
+  // Access only through an std::atomic_ref instance atomically.
+  int8_t reference_count_;
+  // Keys are restricted to be 48-bit unsigned integers.
+  uint16_t key_higher_16_bits_;
+  uint32_t key_lower_32_bits_;
+  int64_t position_in_cache_;
+};
+
+class BaseCachePolicy {
+ public:
+  BaseCachePolicy(int64_t capacity) : capacity_(capacity), cache_usage_(0) {}
+
+  BaseCachePolicy() = default;
+
+  /**
+   * @brief A virtual base class constructor ensures that the derived class
+   * destructor gets called.
+   */
+  virtual ~BaseCachePolicy() = default;
+
+  /**
+   * @brief The policy query function.
+   * @param keys The keys to query the cache.
+   *
+   * @return (positions, indices, missing_keys, found_ptrs), where positions has
+   * the locations of the keys which were found in the cache, missing_keys has
+   * the keys that were not found and indices is defined such that
+   * keys[indices[:positions.size(0)]] gives us the keys for the found pointers
+   * and keys[indices[positions.size(0):]] is identical to missing_keys.
+   */
+  virtual std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  Query(torch::Tensor keys) = 0;
+
+  /**
+   * @brief The policy query function.
+   * @param keys The keys to query the cache.
+   *
+   * @return (positions, indices, pointers, missing_keys), where positions has
+   * the locations of the keys which were emplaced into the cache, pointers
+   * point to the emplaced CacheKey pointers in the cache, missing_keys has the
+   * keys that were not found and just inserted and indices is defined such that
+   * keys[indices[:keys.size(0) - missing_keys.size(0)]] gives us the keys for
+   * the found keys and keys[indices[keys.size(0) - missing_keys.size(0):]] is
+   * identical to missing_keys.
+   */
+  virtual std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplace(torch::Tensor keys) = 0;
+
+  /**
+   * @brief The policy replace function.
+   * @param keys The keys to query the cache.
+   *
+   * @return (positions, pointers), where positions has the locations of the
+   * replaced entries and pointers point to their CacheKey pointers in the
+   * cache.
+   */
+  virtual std::tuple<torch::Tensor, torch::Tensor> Replace(
+      torch::Tensor keys) = 0;
+
+  /**
+   * @brief A reader has finished reading these keys, so they can be evicted.
+   * @param pointers The CacheKey pointers in the cache to unmark.
+   */
+  static void ReadingCompleted(torch::Tensor pointers);
+
+  /**
+   * @brief A writer has finished writing these keys, so they can be evicted.
+   * @param pointers The CacheKey pointers in the cache to unmark.
+   */
+  static void WritingCompleted(torch::Tensor pointers);
+
+ protected:
+  template <typename K, typename V>
+  using map_t = tsl::robin_map<K, V>;
+  template <typename K>
+  using set_t = tsl::robin_set<K>;
+  template <typename iterator>
+  static auto& mutable_value_ref(iterator it) {
+    return it.value();
+  }
+  static constexpr int kCapacityFactor = 2;
+
+  template <typename CachePolicy>
+  static std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryImpl(CachePolicy& policy, torch::Tensor keys);
+
+  template <typename CachePolicy>
+  static std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplaceImpl(CachePolicy& policy, torch::Tensor keys);
+
+  template <typename CachePolicy>
+  static std::tuple<torch::Tensor, torch::Tensor> ReplaceImpl(
+      CachePolicy& policy, torch::Tensor keys);
+
+  template <typename T>
+  static void MoveToFront(
+      std::list<T>& from, std::list<T>& to,
+      typename std::list<T>::iterator it) {
+    std::list<T> temp;
+    // Transfer the element to temp to keep references valid.
+    auto next_it = it;
+    std::advance(next_it, 1);
+    temp.splice(temp.begin(), from, it, next_it);
+    // Move the element to the beginning of the queue.
+    to.splice(to.begin(), temp);
+    // The iterators and references are not invalidated.
+    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(it == to.begin());
+  }
+
+  int64_t capacity_;
+  int64_t cache_usage_;
+
+ private:
+  template <bool write>
+  static void ReadingWritingCompletedImpl(torch::Tensor pointers);
+};
+
+/**
+ * @brief S3FIFO is a simple, scalable FIFObased algorithm with three static
+ * queues (S3-FIFO). https://dl.acm.org/doi/pdf/10.1145/3600006.3613147
+ **/
+class S3FifoCachePolicy : public BaseCachePolicy {
+ public:
+  using map_iterator = map_t<int64_t, CacheKey*>::iterator;
+  /**
+   * @brief Constructor for the S3FifoCachePolicy class.
+   *
+   * @param capacity The capacity of the cache in terms of # elements.
+   */
+  S3FifoCachePolicy(int64_t capacity);
+
+  S3FifoCachePolicy() = default;
+
+  S3FifoCachePolicy(S3FifoCachePolicy&&) = default;
+
+  virtual ~S3FifoCachePolicy() = default;
+
+  /**
+   * @brief See BaseCachePolicy::Query.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> Query(
+      torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::QueryAndReplace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplace(torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::Replace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor> Replace(torch::Tensor keys);
+
+  CacheKey* Read(int64_t key) {
+    auto it = key_to_cache_key_.find(key);
+    if (it != key_to_cache_key_.end()) {
+      auto& cache_key = it->second->Increment();
+      if (!cache_key.BeingWritten()) {
+        return &cache_key.StartRead();
+      }
+    }
+    return nullptr;
+  }
+
+  auto getMapSentinelValue() const { return nullptr; }
+
+  std::pair<map_iterator, bool> Emplace(int64_t key) {
+    auto [it, inserted] = key_to_cache_key_.emplace(key, getMapSentinelValue());
+    bool readable = false;
+    if (!inserted) {
+      auto& cache_key = it->second->Increment();
+      if (!cache_key.BeingWritten()) {
+        cache_key.StartRead();
+        readable = true;
+      }
+    }
+    return {it, readable};
+  }
+
+  CacheKey* Insert(map_iterator it) {
+    const auto key = it->first;
+    const auto in_ghost_queue = ghost_set_.erase(key);
+    auto& queue = in_ghost_queue ? main_queue_ : small_queue_;
+    queue.push_front(CacheKey(key));
+    small_queue_size_ += 1 - in_ghost_queue;
+    auto cache_key_ptr = &queue.front();
+    mutable_value_ref(it) = cache_key_ptr;
+    return &cache_key_ptr->setPos(Evict());
+  }
+
+ private:
+  int64_t EvictMainQueue() {
+    while (true) {
+      auto& evicted = main_queue_.back();
+      if (evicted.getFreq() > 0 || evicted.InUse()) {
+        evicted.Decrement();
+        auto it = main_queue_.end();
+        std::advance(it, -1);
+        MoveToFront(main_queue_, main_queue_, it);
+      } else {
+        key_to_cache_key_.erase(evicted.getKey());
+        const auto evicted_pos = evicted.getPos();
+        main_queue_.pop_back();
+        return evicted_pos;
+      }
+    }
+  }
+
+  int64_t EvictSmallQueue() {
+    while (small_queue_size_ > small_queue_size_target_) {
+      --small_queue_size_;
+      auto& evicted = small_queue_.back();
+      if (evicted.getFreq() > 0 || evicted.InUse()) {
+        evicted.ResetFreq();
+        auto it = small_queue_.end();
+        std::advance(it, -1);
+        MoveToFront(small_queue_, main_queue_, it);
+      } else {
+        const auto evicted_key = evicted.getKey();
+        key_to_cache_key_.erase(evicted_key);
+        const auto evicted_pos = evicted.getPos();
+        small_queue_.pop_back();
+        if (ghost_queue_.IsFull()) {
+          ghost_set_.erase(ghost_queue_.Pop());
+        }
+        ghost_set_.insert(evicted_key);
+        ghost_queue_.Push(evicted_key);
+        return evicted_pos;
+      }
+    }
+    return -1;
+  }
+
+  int64_t Evict() {
+    // If the cache has space, get an unused slot otherwise perform eviction.
+    if (cache_usage_ < capacity_) return cache_usage_++;
+    const auto pos = EvictSmallQueue();
+    return pos >= 0 ? pos : EvictMainQueue();
+  }
+
+  std::list<CacheKey> small_queue_, main_queue_;
+  CircularQueue<int64_t> ghost_queue_;
+  size_t small_queue_size_target_;
+  // std::list<>::size() is O(N) before the CXX11 ABI which torch enforces.
+  size_t small_queue_size_;
+  set_t<int64_t> ghost_set_;
+  map_t<int64_t, CacheKey*> key_to_cache_key_;
+};
+
+/**
+ * @brief SIEVE is a simple, scalable FIFObased algorithm with a single static
+ * queue. https://www.usenix.org/system/files/nsdi24-zhang-yazhuo.pdf
+ **/
+class SieveCachePolicy : public BaseCachePolicy {
+ public:
+  using map_iterator = map_t<int64_t, CacheKey*>::iterator;
+  /**
+   * @brief Constructor for the SieveCachePolicy class.
+   *
+   * @param capacity The capacity of the cache in terms of # elements.
+   */
+  SieveCachePolicy(int64_t capacity);
+
+  SieveCachePolicy() = default;
+
+  virtual ~SieveCachePolicy() = default;
+
+  /**
+   * @brief See BaseCachePolicy::Query.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> Query(
+      torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::QueryAndReplace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplace(torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::Replace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor> Replace(torch::Tensor keys);
+
+  CacheKey* Read(int64_t key) {
+    auto it = key_to_cache_key_.find(key);
+    if (it != key_to_cache_key_.end()) {
+      auto& cache_key = it->second->SetFreq();
+      if (!cache_key.BeingWritten()) {
+        return &cache_key.StartRead();
+      }
+    }
+    return nullptr;
+  }
+
+  auto getMapSentinelValue() const { return nullptr; }
+
+  std::pair<map_iterator, bool> Emplace(int64_t key) {
+    auto [it, inserted] = key_to_cache_key_.emplace(key, getMapSentinelValue());
+    bool readable = false;
+    if (!inserted) {
+      auto& cache_key = it->second->SetFreq();
+      if (!cache_key.BeingWritten()) {
+        cache_key.StartRead();
+        readable = true;
+      }
+    }
+    return {it, readable};
+  }
+
+  CacheKey* Insert(map_iterator it) {
+    const auto key = it->first;
+    queue_.push_front(CacheKey(key));
+    auto cache_key_ptr = &queue_.front();
+    mutable_value_ref(it) = cache_key_ptr;
+    return &cache_key_ptr->setPos(Evict());
+  }
+
+ private:
+  int64_t Evict() {
+    // If the cache has space, get an unused slot otherwise perform eviction.
+    if (cache_usage_ < capacity_) return cache_usage_++;
+    --hand_;
+    while (hand_->getFreq() || hand_->InUse()) {
+      hand_->ResetFreq();
+      if (hand_ == queue_.begin()) hand_ = queue_.end();
+      --hand_;
+    }
+    key_to_cache_key_.erase(hand_->getKey());
+    const auto pos = hand_->getPos();
+    const auto temp = hand_;
+    if (hand_ == queue_.begin()) {
+      hand_ = queue_.end();
+    } else {
+      ++hand_;
+    }
+    queue_.erase(temp);
+    return pos;
+  }
+
+  std::list<CacheKey> queue_;
+  decltype(queue_)::iterator hand_;
+  map_t<int64_t, CacheKey*> key_to_cache_key_;
+};
+
+/**
+ * @brief LeastRecentlyUsed is a simple, scalable FIFObased algorithm with a
+ * single static queue.
+ **/
+class LruCachePolicy : public BaseCachePolicy {
+ public:
+  using map_iterator = map_t<int64_t, std::list<CacheKey>::iterator>::iterator;
+  /**
+   * @brief Constructor for the LruCachePolicy class.
+   *
+   * @param capacity The capacity of the cache in terms of # elements.
+   */
+  LruCachePolicy(int64_t capacity);
+
+  LruCachePolicy() = default;
+
+  virtual ~LruCachePolicy() = default;
+
+  /**
+   * @brief See BaseCachePolicy::Query.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> Query(
+      torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::QueryAndReplace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplace(torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::Replace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor> Replace(torch::Tensor keys);
+
+  CacheKey* Read(int64_t key) {
+    auto it = key_to_cache_key_.find(key);
+    if (it != key_to_cache_key_.end()) {
+      auto& cache_key = *it->second;
+      MoveToFront(queue_, queue_, it->second);
+      if (!cache_key.BeingWritten()) {
+        return &cache_key.StartRead();
+      }
+    }
+    return nullptr;
+  }
+
+  auto getMapSentinelValue() { return queue_.end(); }
+
+  std::pair<map_iterator, bool> Emplace(int64_t key) {
+    auto [it, inserted] = key_to_cache_key_.emplace(key, getMapSentinelValue());
+    bool readable = false;
+    if (!inserted) {
+      auto& cache_key = *it->second;
+      MoveToFront(queue_, queue_, it->second);
+      if (!cache_key.BeingWritten()) {
+        cache_key.StartRead();
+        readable = true;
+      }
+    }
+    return {it, readable};
+  }
+
+  CacheKey* Insert(map_iterator it) {
+    const auto key = it->first;
+    queue_.push_front(CacheKey(key));
+    mutable_value_ref(it) = queue_.begin();
+    auto cache_key_ptr = &queue_.front();
+    return &cache_key_ptr->setPos(Evict());
+  }
+
+ private:
+  int64_t Evict() {
+    // If the cache has space, get an unused slot otherwise perform eviction.
+    if (cache_usage_ < capacity_) return cache_usage_++;
+    // Do not evict items that are still in use.
+    while (queue_.back().InUse()) {
+      auto it = queue_.end();
+      std::advance(it, -1);
+      // Move the last element to the front without invalidating references.
+      MoveToFront(queue_, queue_, it);
+    }
+    const auto& cache_key = queue_.back();
+    key_to_cache_key_.erase(cache_key.getKey());
+    const auto pos = cache_key.getPos();
+    queue_.pop_back();
+    return pos;
+  }
+
+  std::list<CacheKey> queue_;
+  map_t<int64_t, decltype(queue_)::iterator> key_to_cache_key_;
+};
+
+/**
+ * @brief Clock (FIFO-Reinsertion) is a simple, scalable FIFObased algorithm
+ * with a single static queue.
+ * https://people.csail.mit.edu/saltzer/Multics/MHP-Saltzer-060508/bookcases/M00s/M0104%20074-12).PDF
+ **/
+class ClockCachePolicy : public BaseCachePolicy {
+ public:
+  using map_iterator = map_t<int64_t, CacheKey*>::iterator;
+  /**
+   * @brief Constructor for the ClockCachePolicy class.
+   *
+   * @param capacity The capacity of the cache in terms of # elements.
+   */
+  ClockCachePolicy(int64_t capacity);
+
+  ClockCachePolicy() = default;
+
+  ClockCachePolicy(ClockCachePolicy&&) = default;
+
+  virtual ~ClockCachePolicy() = default;
+
+  /**
+   * @brief See BaseCachePolicy::Query.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> Query(
+      torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::QueryAndReplace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+  QueryAndReplace(torch::Tensor keys);
+
+  /**
+   * @brief See BaseCachePolicy::Replace.
+   */
+  std::tuple<torch::Tensor, torch::Tensor> Replace(torch::Tensor keys);
+
+  CacheKey* Read(int64_t key) {
+    auto it = key_to_cache_key_.find(key);
+    if (it != key_to_cache_key_.end()) {
+      auto& cache_key = it->second->SetFreq();
+      if (!cache_key.BeingWritten()) {
+        return &cache_key.StartRead();
+      }
+    }
+    return nullptr;
+  }
+
+  auto getMapSentinelValue() const { return nullptr; }
+
+  std::pair<map_iterator, bool> Emplace(int64_t key) {
+    auto [it, inserted] = key_to_cache_key_.emplace(key, getMapSentinelValue());
+    bool readable = false;
+    if (!inserted) {
+      auto& cache_key = it->second->SetFreq();
+      if (!cache_key.BeingWritten()) {
+        cache_key.StartRead();
+        readable = true;
+      }
+    }
+    return {it, readable};
+  }
+
+  CacheKey* Insert(map_iterator it) {
+    const auto key = it->first;
+    queue_.push_front(CacheKey(key));
+    auto cache_key_ptr = &queue_.front();
+    mutable_value_ref(it) = cache_key_ptr;
+    return &cache_key_ptr->setPos(Evict());
+  }
+
+ private:
+  int64_t Evict() {
+    // If the cache has space, get an unused slot otherwise perform eviction.
+    if (cache_usage_ < capacity_) return cache_usage_++;
+    while (true) {
+      auto& cache_key = queue_.back();
+      if (cache_key.getFreq() || cache_key.InUse()) {
+        cache_key.ResetFreq();
+        auto it = queue_.end();
+        std::advance(it, -1);
+        MoveToFront(queue_, queue_, it);
+      } else {
+        key_to_cache_key_.erase(cache_key.getKey());
+        const auto evicted_pos = cache_key.getPos();
+        queue_.pop_back();
+        return evicted_pos;
+      }
+    }
+  }
+
+  std::list<CacheKey> queue_;
+  map_t<int64_t, CacheKey*> key_to_cache_key_;
+};
+
+}  // namespace storage
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CACHE_POLICY_H_

graphbolt/src/circular_queue.h

@@ -0,0 +1,93 @@
+/**
+ *   Copyright (c) 2024, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file circular_queue.h
+ * @brief Circular queue implementation.
+ */
+#ifndef GRAPHBOLT_CIRCULAR_QUEUE_H_
+#define GRAPHBOLT_CIRCULAR_QUEUE_H_
+
+#include <memory>
+
+namespace graphbolt {
+
+template <typename T>
+struct CircularQueue {
+  CircularQueue(const int64_t capacity)
+      : tail_(0),
+        head_(0),
+        // + 1 is needed to be able to differentiate empty and full states.
+        capacity_(capacity + 1),
+        data_{new T[capacity + 1]} {}
+
+  CircularQueue() = default;
+
+  T* Push(const T& x) {
+    auto insert_ptr = &data_[PostIncrement(tail_)];
+    *insert_ptr = x;
+    return insert_ptr;
+  }
+
+  T Pop() { return data_[PostIncrement(head_)]; }
+
+  void PopN(int64_t N) {
+    head_ += N;
+    if (head_ >= capacity_) head_ -= capacity_;
+  }
+
+  auto Clear() { head_ = tail_; }
+
+  T& Front() const { return data_[head_]; }
+
+  bool IsFull() const {
+    const auto diff = tail_ + 1 - head_;
+    return diff == 0 || diff == capacity_;
+  }
+
+  auto Size() const {
+    auto diff = tail_ - head_;
+    if (diff < 0) diff += capacity_;
+    return diff;
+  }
+
+  friend std::ostream& operator<<(
+      std::ostream& os, const CircularQueue& queue) {
+    for (auto i = queue.head_; i != queue.tail_; queue.PostIncrement(i)) {
+      os << queue.data_[i] << ", ";
+    }
+    return os << "\n";
+  }
+
+  bool IsEmpty() const { return tail_ == head_; }
+
+  auto Capacity() const { return capacity_ - 1; }
+
+ private:
+  int64_t PostIncrement(int64_t& i) const {
+    const auto ret = i++;
+    if (i >= capacity_) i -= capacity_;
+    return ret;
+  }
+
+  int64_t tail_;
+  int64_t head_;
+  int64_t capacity_;
+  std::unique_ptr<T[]> data_;
+};
+
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CIRCULAR_QUEUE_H_

graphbolt/src/cnumpy.cc

@@ -0,0 +1,320 @@
+/**
+ *  Copyright (c) 2024, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *  Copyright (c) 2023 by Contributors
+ * @file cnumpy.cc
+ * @brief Numpy File Fetecher class.
+ */
+
+#include "./cnumpy.h"
+
+#include "./io_uring.h"
+
+#ifdef HAVE_LIBRARY_LIBURING
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#endif
+
+#include <graphbolt/async.h>
+#include <torch/torch.h>
+
+#include <atomic>
+#include <cstring>
+#include <fstream>
+#include <memory>
+#include <numeric>
+#include <stdexcept>
+#include <vector>
+
+#include "./circular_queue.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace storage {
+
+OnDiskNpyArray::OnDiskNpyArray(
+    std::string filename, torch::ScalarType dtype,
+    const std::vector<int64_t> &shape, torch::optional<int64_t> num_threads)
+    : filename_(filename),
+      feature_dim_(shape),
+      dtype_(dtype),
+      feature_size_(std::accumulate(
+          shape.begin() + 1, shape.end(), c10::elementSize(dtype),
+          std::multiplies<int64_t>())) {
+#ifndef __linux__
+  throw std::runtime_error(
+      "OnDiskNpyArray is not supported on non-Linux systems.");
+#endif
+#ifdef HAVE_LIBRARY_LIBURING
+  ParseNumpyHeader();
+  file_description_ = ::open(filename.c_str(), O_RDONLY | O_DIRECT);
+  if (file_description_ < 0) {
+    throw std::runtime_error("npy_load: Unable to open file " + filename);
+  }
+  struct stat st;
+  TORCH_CHECK(::fstat(file_description_, &st) == 0);
+  const auto file_size = st.st_size;
+  block_size_ = st.st_blksize;
+  TORCH_CHECK(file_size - prefix_len_ >= feature_dim_[0] * feature_size_);
+
+  // The minimum page size to contain one feature.
+  aligned_length_ = (feature_size_ + block_size_ - 1) & ~(block_size_ - 1);
+
+  std::call_once(call_once_flag_, [&] {
+    // Get system max interop thread count.
+    num_queues_ =
+        io_uring::num_threads.value_or(torch::get_num_interop_threads());
+    TORCH_CHECK(num_queues_ > 0, "A positive # queues is required.");
+    io_uring_queue_ = std::unique_ptr<::io_uring[], io_uring_queue_destroyer>(
+        new ::io_uring[num_queues_], io_uring_queue_destroyer{num_queues_});
+    TORCH_CHECK(num_queues_ <= counting_semaphore_t::max());
+    semaphore_.release(num_queues_);
+    available_queues_.reserve(num_queues_);
+    // Init io_uring queue.
+    for (int64_t t = 0; t < num_queues_; t++) {
+      available_queues_.push_back(t);
+      TORCH_CHECK(
+          ::io_uring_queue_init(2 * kGroupSize, &io_uring_queue_[t], 0) == 0);
+      // We have allocated 2 * kGroupSize submission queue entries and
+      // 4 * kGroupSize completion queue entries after this call.
+    }
+  });
+
+  num_thread_ = std::min(
+      static_cast<int64_t>(num_queues_), num_threads.value_or(num_queues_));
+  TORCH_CHECK(num_thread_ > 0, "A positive # threads is required.");
+
+  // We allocate buffers for each existing queue because we might get assigned
+  // any queue in range [0, num_queues_).
+  read_tensor_ = torch::empty(
+      ReadBufferSizePerThread() * num_queues_ + block_size_ - 1,
+      torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU));
+#else
+  throw std::runtime_error("DiskBasedFeature is not available now.");
+#endif  // HAVE_LIBRARY_LIBURING
+}
+
+c10::intrusive_ptr<OnDiskNpyArray> OnDiskNpyArray::Create(
+    std::string path, torch::ScalarType dtype,
+    const std::vector<int64_t> &shape, torch::optional<int64_t> num_threads) {
+  return c10::make_intrusive<OnDiskNpyArray>(path, dtype, shape, num_threads);
+}
+
+OnDiskNpyArray::~OnDiskNpyArray() {
+#ifdef HAVE_LIBRARY_LIBURING
+  TORCH_CHECK(::close(file_description_) == 0);
+#endif  // HAVE_LIBRARY_LIBURING
+}
+
+void OnDiskNpyArray::ParseNumpyHeader() {
+  // Parse numpy file header to get basic info of feature.
+  // Get file prefix length.
+  std::ifstream file(filename_);
+  if (!file.is_open()) {
+    throw std::runtime_error(
+        "ParseNumpyHeader: Unable to open file " + filename_);
+  }
+  std::string header;
+  std::getline(file, header);
+  // Get prefix length for computing feature offset,
+  // add one for new-line character.
+  prefix_len_ = header.size() + 1;
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> OnDiskNpyArray::IndexSelect(
+    torch::Tensor index) {
+#ifdef HAVE_LIBRARY_LIBURING
+  return IndexSelectIOUring(index);
+#else
+  TORCH_CHECK(false, "OnDiskNpyArray is not supported on non-Linux systems.");
+  return {};
+#endif  // HAVE_LIBRARY_LIBURING
+}
+
+class ReadRequest {
+ public:
+  char *destination_;
+  int64_t read_len_;
+  int64_t offset_;
+  int64_t block_size_;
+  char *aligned_read_buffer_;
+
+  auto AlignedOffset() const { return offset_ & ~(block_size_ - 1); }
+
+  auto ReadBuffer() const {
+    return aligned_read_buffer_ + offset_ - AlignedOffset();
+  }
+
+  auto AlignedReadSize() const {
+    const int64_t end_offset = offset_ + read_len_;
+    const int64_t aligned_end_offset =
+        (end_offset + block_size_ - 1) & ~(block_size_ - 1);
+    return aligned_end_offset - AlignedOffset();
+  }
+
+  auto MinimumReadSize() const { return offset_ + read_len_ - AlignedOffset(); }
+};
+
+#ifdef HAVE_LIBRARY_LIBURING
+torch::Tensor OnDiskNpyArray::IndexSelectIOUringImpl(torch::Tensor index) {
+  std::vector<int64_t> shape(index.sizes().begin(), index.sizes().end());
+  shape.insert(shape.end(), feature_dim_.begin() + 1, feature_dim_.end());
+  auto result = torch::empty(
+      shape, index.options()
+                 .dtype(dtype_)
+                 .layout(torch::kStrided)
+                 .pinned_memory(utils::is_pinned(index))
+                 .requires_grad(false));
+  auto result_buffer = reinterpret_cast<char *>(result.data_ptr());
+
+  // Indicator for index error.
+  std::atomic<int> error_flag{};
+  std::atomic<int64_t> work_queue{};
+  // Construct a QueueAndBufferAcquirer object so that the worker threads can
+  // share the available queues and buffers.
+  QueueAndBufferAcquirer queue_source(this);
+  graphbolt::parallel_for_each_interop(0, num_thread_, 1, [&](int) {
+    // The completion queue might contain 4 * kGroupSize while we may submit
+    // 4 * kGroupSize more. No harm in overallocation here.
+    CircularQueue<ReadRequest> read_queue(8 * kGroupSize);
+    int64_t num_submitted = 0;
+    int64_t num_completed = 0;
+    auto [acquired_queue_handle, read_buffer_source2] = queue_source.get();
+    auto &io_uring_queue = acquired_queue_handle.get();
+    // Capturing structured binding is available only in C++20, so we rename.
+    auto read_buffer_source = read_buffer_source2;
+    auto submit_fn = [&](int64_t submission_minimum_batch_size) {
+      if (read_queue.Size() < submission_minimum_batch_size) return;
+      TORCH_CHECK(  // Check for sqe overflow.
+          read_queue.Size() <= 2 * kGroupSize);
+      TORCH_CHECK(  // Check for cqe overflow.
+          read_queue.Size() + num_submitted - num_completed <= 4 * kGroupSize);
+      // Submit and wait for the reads.
+      while (!read_queue.IsEmpty()) {
+        const auto submitted = ::io_uring_submit(&io_uring_queue);
+        TORCH_CHECK(submitted >= 0);
+        num_submitted += submitted;
+        // Pop the submitted entries from the queue.
+        read_queue.PopN(submitted);
+      }
+    };
+    for (int64_t read_buffer_slot = 0; true;) {
+      auto request_read_buffer = [&]() {
+        return read_buffer_source + (aligned_length_ + block_size_) *
+                                        (read_buffer_slot++ % (8 * kGroupSize));
+      };
+      const auto num_requested_items = std::max(
+          std::min(
+              // The condition not to overflow the completion queue.
+              2 * kGroupSize -
+                  (read_queue.Size() + num_submitted - num_completed),
+              // The condition not to overflow the submission queue.
+              kGroupSize - read_queue.Size()),
+          int64_t{});
+      const auto begin =
+          work_queue.fetch_add(num_requested_items, std::memory_order_relaxed);
+      if ((begin >= index.numel() && read_queue.IsEmpty() &&
+           num_completed >= num_submitted) ||
+          // Even when we encounter out of bounds index (error_flag == 1), we
+          // continue. We want to ensure the reads in flight successfully
+          // complete to avoid the instability due to incompleted reads.
+          error_flag.load(std::memory_order_relaxed) > 1)
+        break;
+      const auto end = std::min(begin + num_requested_items, index.numel());
+      AT_DISPATCH_INDEX_TYPES(
+          index.scalar_type(), "IndexSelectIOUring", ([&] {
+            auto index_data = index.data_ptr<index_t>();
+            for (int64_t i = begin; i < end; ++i) {
+              int64_t feature_id = index_data[i];
+              if (feature_id < 0) feature_id += feature_dim_[0];
+              if (feature_id < 0 || feature_id >= feature_dim_[0]) {
+                error_flag.store(1, std::memory_order_relaxed);
+                // Simply skip the out of bounds index.
+                continue;
+              }
+              // calculate offset of the feature.
+              const int64_t offset = feature_id * feature_size_ + prefix_len_;
+
+              ReadRequest req{
+                  result_buffer + feature_size_ * i, feature_size_, offset,
+                  block_size_, request_read_buffer()};
+
+              // Put requests into io_uring queue.
+              struct io_uring_sqe *sqe = io_uring_get_sqe(&io_uring_queue);
+              TORCH_CHECK(sqe);
+              io_uring_sqe_set_data(sqe, read_queue.Push(req));
+              io_uring_prep_read(
+                  sqe, file_description_, req.aligned_read_buffer_,
+                  req.AlignedReadSize(), req.AlignedOffset());
+              submit_fn(kGroupSize);
+            }
+          }));
+
+      submit_fn(1);  // Submit all sqes.
+      // Wait for the reads; completion queue entries.
+      struct io_uring_cqe *cqe;
+      TORCH_CHECK(num_submitted - num_completed <= 2 * kGroupSize);
+      TORCH_CHECK(
+          ::io_uring_wait_cqe_nr(
+              &io_uring_queue, &cqe, num_submitted - num_completed) == 0);
+      // Check the reads and abort on failure.
+      int num_cqes_seen = 0;
+      unsigned head;
+      io_uring_for_each_cqe(&io_uring_queue, head, cqe) {
+        const auto &req =
+            *reinterpret_cast<ReadRequest *>(io_uring_cqe_get_data(cqe));
+        auto actual_read_len = cqe->res;
+        if (actual_read_len < 0) {
+          error_flag.store(actual_read_len, std::memory_order_relaxed);
+          break;
+        }
+        const auto remaining_read_len =
+            std::max(req.MinimumReadSize() - actual_read_len, int64_t{});
+        const auto remaining_useful_read_len =
+            std::min(remaining_read_len, req.read_len_);
+        const auto useful_read_len = req.read_len_ - remaining_useful_read_len;
+        if (remaining_read_len) {
+          // Remaining portion will be read as part of the next batch.
+          ReadRequest rest{
+              req.destination_ + useful_read_len, remaining_useful_read_len,
+              req.offset_ + useful_read_len, block_size_,
+              request_read_buffer()};
+          // Put requests into io_uring queue.
+          struct io_uring_sqe *sqe = io_uring_get_sqe(&io_uring_queue);
+          TORCH_CHECK(sqe);
+          io_uring_sqe_set_data(sqe, read_queue.Push(rest));
+          io_uring_prep_read(
+              sqe, file_description_, rest.aligned_read_buffer_,
+              rest.AlignedReadSize(), rest.AlignedOffset());
+          submit_fn(kGroupSize);
+        }
+        // Copy results into result_buffer.
+        std::memcpy(req.destination_, req.ReadBuffer(), useful_read_len);
+        num_cqes_seen++;
+      }
+
+      // Move the head pointer of completion queue.
+      io_uring_cq_advance(&io_uring_queue, num_cqes_seen);
+      num_completed += num_cqes_seen;
+    }
+  });
+  const auto ret_val = error_flag.load(std::memory_order_relaxed);
+  switch (ret_val) {
+    case 0:  // Successful.
+      return result;
+    case 1:
+      throw std::out_of_range("IndexError: Index out of range.");
+    default:
+      throw std::runtime_error(
+          "io_uring error with errno: " + std::to_string(-ret_val));
+  }
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> OnDiskNpyArray::IndexSelectIOUring(
+    torch::Tensor index) {
+  return async([=, this] { return IndexSelectIOUringImpl(index); });
+}
+
+#endif  // HAVE_LIBRARY_LIBURING
+}  // namespace storage
+}  // namespace graphbolt

graphbolt/src/cnumpy.h

@@ -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

graphbolt/src/concurrent_id_hash_map.cc

@@ -0,0 +1,219 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file concurrent_id_hash_map.cc
+ * @brief Class about id hash map.
+ */
+
+#include "concurrent_id_hash_map.h"
+
+#ifdef _MSC_VER
+#include <intrin.h>
+#endif  // _MSC_VER
+
+#include <cmath>
+#include <cuda/std/atomic>
+#include <numeric>
+
+namespace {
+static constexpr int64_t kEmptyKey = -1;
+static constexpr int kGrainSize = 256;
+
+// The formula is established from experience which is used to get the hashmap
+// size from the input array size.
+inline size_t GetMapSize(size_t num) {
+  size_t capacity = 1;
+  return capacity << static_cast<size_t>(1 + std::log2(num * 3));
+}
+}  // namespace
+
+namespace graphbolt {
+namespace sampling {
+
+template <typename IdType>
+ConcurrentIdHashMap<IdType>::ConcurrentIdHashMap(
+    const torch::Tensor& ids, size_t num_seeds) {
+  const IdType* ids_data = ids.data_ptr<IdType>();
+  const size_t num_ids = static_cast<size_t>(ids.size(0));
+  size_t capacity = GetMapSize(num_ids);
+  mask_ = static_cast<IdType>(capacity - 1);
+
+  hash_map_ =
+      torch::full({static_cast<int64_t>(capacity * 2)}, -1, ids.options());
+
+  // This code block is to fill the ids into hash_map_.
+  unique_ids_ = torch::empty_like(ids);
+  IdType* unique_ids_data = unique_ids_.data_ptr<IdType>();
+  // Insert all ids into the hash map.
+  torch::parallel_for(0, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
+    for (int64_t i = s; i < e; i++) {
+      InsertAndSetMin(ids_data[i], static_cast<IdType>(i));
+    }
+  });
+  // Place the first `num_seeds` ids.
+  unique_ids_.slice(0, 0, num_seeds) = ids.slice(0, 0, num_seeds);
+
+  auto valid_tensor = torch::empty(num_ids, ids.options().dtype(torch::kInt8));
+  auto valid = valid_tensor.data_ptr<int8_t>();
+
+  const int64_t num_threads = torch::get_num_threads();
+  std::vector<size_t> block_offset(num_threads + 1, 0);
+
+  // Count the valid numbers in each thread.
+  torch::parallel_for(
+      num_seeds, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
+        size_t count = 0;
+        for (int64_t i = s; i < e; i++) {
+          if (MapId(ids_data[i]) == i) {
+            count++;
+            valid[i] = 1;
+          } else {
+            valid[i] = 0;
+          }
+        }
+        auto thread_id = torch::get_thread_num();
+        block_offset[thread_id + 1] = count;
+      });
+
+  // Get ExclusiveSum of each block.
+  std::partial_sum(
+      block_offset.begin() + 1, block_offset.end(), block_offset.begin() + 1);
+  unique_ids_ = unique_ids_.slice(0, 0, num_seeds + block_offset.back());
+
+  // Get unique array from ids and set value for hash map.
+  torch::parallel_for(
+      num_seeds, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
+        auto thread_id = torch::get_thread_num();
+        auto pos = block_offset[thread_id] + num_seeds;
+        for (int64_t i = s; i < e; i++) {
+          if (valid[i]) {
+            unique_ids_data[pos] = ids_data[i];
+            Set(ids_data[i], pos);
+            pos = pos + 1;
+          }
+        }
+      });
+}
+
+template <typename IdType>
+torch::Tensor ConcurrentIdHashMap<IdType>::MapIds(
+    const torch::Tensor& ids) const {
+  const IdType* ids_data = ids.data_ptr<IdType>();
+
+  torch::Tensor new_ids = torch::empty_like(ids);
+  auto num_ids = new_ids.size(0);
+  IdType* values_data = new_ids.data_ptr<IdType>();
+
+  torch::parallel_for(0, num_ids, kGrainSize, [&](int64_t s, int64_t e) {
+    for (int64_t i = s; i < e; i++) {
+      values_data[i] = MapId(ids_data[i]);
+    }
+  });
+  return new_ids;
+}
+
+template <typename IdType>
+constexpr IdType getKeyIndex(IdType pos) {
+  return 2 * pos;
+}
+
+template <typename IdType>
+constexpr IdType getValueIndex(IdType pos) {
+  return 2 * pos + 1;
+}
+
+template <typename IdType>
+inline void ConcurrentIdHashMap<IdType>::Next(
+    IdType* pos, IdType* delta) const {
+  // Use Quadric probing.
+  *pos = (*pos + (*delta) * (*delta)) & mask_;
+  *delta = *delta + 1;
+}
+
+template <typename IdType>
+inline IdType ConcurrentIdHashMap<IdType>::MapId(IdType id) const {
+  IdType pos = (id & mask_), delta = 1;
+  IdType empty_key = static_cast<IdType>(kEmptyKey);
+  IdType* hash_map_data = hash_map_.data_ptr<IdType>();
+  IdType key = hash_map_data[getKeyIndex(pos)];
+  while (key != empty_key && key != id) {
+    Next(&pos, &delta);
+    key = hash_map_data[getKeyIndex(pos)];
+  }
+  if (key == empty_key) {
+    throw std::out_of_range("Id not found: " + std::to_string(id));
+  }
+  return hash_map_data[getValueIndex(pos)];
+}
+
+template <typename IdType>
+bool ConcurrentIdHashMap<IdType>::Insert(IdType id) {
+  IdType pos = (id & mask_), delta = 1;
+  InsertState state = AttemptInsertAt(pos, id);
+  while (state == InsertState::OCCUPIED) {
+    Next(&pos, &delta);
+    state = AttemptInsertAt(pos, id);
+  }
+
+  return state == InsertState::INSERTED;
+}
+
+template <typename IdType>
+inline void ConcurrentIdHashMap<IdType>::Set(IdType key, IdType value) {
+  IdType pos = (key & mask_), delta = 1;
+  IdType* hash_map_data = hash_map_.data_ptr<IdType>();
+  while (hash_map_data[getKeyIndex(pos)] != key) {
+    Next(&pos, &delta);
+  }
+
+  hash_map_data[getValueIndex(pos)] = value;
+}
+
+template <typename IdType>
+inline void ConcurrentIdHashMap<IdType>::InsertAndSet(IdType id, IdType value) {
+  IdType pos = (id & mask_), delta = 1;
+  while (AttemptInsertAt(pos, id) == InsertState::OCCUPIED) {
+    Next(&pos, &delta);
+  }
+
+  hash_map_.data_ptr<IdType>()[getValueIndex(pos)] = value;
+}
+
+template <typename IdType>
+void ConcurrentIdHashMap<IdType>::InsertAndSetMin(IdType id, IdType value) {
+  IdType pos = (id & mask_), delta = 1;
+  InsertState state = AttemptInsertAt(pos, id);
+  while (state == InsertState::OCCUPIED) {
+    Next(&pos, &delta);
+    state = AttemptInsertAt(pos, id);
+  }
+
+  IdType empty_key = static_cast<IdType>(kEmptyKey);
+  IdType val_pos = getValueIndex(pos);
+  ::cuda::std::atomic_ref value_ref(
+      reinterpret_cast<IdType*>(hash_map_.data_ptr())[val_pos]);
+  for (auto old_val = empty_key; old_val == empty_key || old_val > value;) {
+    // It is more efficient to use weak variant in a loop.
+    if (value_ref.compare_exchange_weak(old_val, value)) break;
+  }
+}
+
+template <typename IdType>
+inline typename ConcurrentIdHashMap<IdType>::InsertState
+ConcurrentIdHashMap<IdType>::AttemptInsertAt(int64_t pos, IdType key) {
+  auto expected = static_cast<IdType>(kEmptyKey);
+  ::cuda::std::atomic_ref key_ref(
+      reinterpret_cast<IdType*>(hash_map_.data_ptr())[getKeyIndex(pos)]);
+  if (key_ref.compare_exchange_strong(expected, key)) {
+    return InsertState::INSERTED;
+  } else if (expected == key) {
+    return InsertState::EXISTED;
+  } else {
+    return InsertState::OCCUPIED;
+  }
+}
+
+template class ConcurrentIdHashMap<int32_t>;
+template class ConcurrentIdHashMap<int64_t>;
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/concurrent_id_hash_map.h

@@ -0,0 +1,190 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file concurrent_id_hash_map.h
+ * @brief Class about concurrent id hash map.
+ */
+
+#ifndef GRAPHBOLT_CONCURRENT_ID_HASH_MAP_H_
+#define GRAPHBOLT_CONCURRENT_ID_HASH_MAP_H_
+
+#include <torch/torch.h>
+
+#include <functional>
+#include <memory>
+#include <vector>
+
+namespace graphbolt {
+namespace sampling {
+
+/**
+ * @brief A CPU targeted hashmap for mapping duplicate and non-consecutive ids
+ * in the provided array to unique and consecutive ones. It utilizes
+ * multi-threading to accelerate the insert and search speed. Currently it is
+ * only designed to be used in `ToBlockCpu` for optimizing, so it only support
+ * key insertions once with Init function, and it does not support key deletion.
+ *
+ * The hash map should be prepared in two phases before using. With the first
+ * being creating the hashmap, and then initialize it with an id array which is
+ * divided into 2 parts: [`seed ids`, `sampled ids`]. `Seed ids` refer to
+ * a set ids chosen as the input for sampling process and `sampled ids` are the
+ * ids new sampled from the process (note the the `seed ids` might also be
+ * sampled in the process and included in the `sampled ids`). In result `seed
+ * ids` are mapped to [0, num_seed_ids) and `sampled ids` to [num_seed_ids,
+ * num_unique_ids). Notice that mapping order is stable for `seed ids` while not
+ * for the `sampled ids`.
+ *
+ * For example, for an array `A` having 4 seed ids with following entries:
+ * [99, 98, 100, 97, 97, 101, 101, 102, 101]
+ * Create the hashmap `H` with:
+ * `H = ConcurrentIdHashMap()` (1)
+ * And Init it with:
+ * `U = H.Init(A)` (2)  (U is an id array used to store the unqiue
+ * ids in A).
+ * Then `U` should be (U is not exclusive as the overall mapping is not stable):
+ * [99, 98, 100, 97, 102, 101]
+ * And the hashmap should generate following mappings:
+ *  * [
+ *   {key: 99, value: 0},
+ *   {key: 98, value: 1},
+ *   {key: 100, value: 2},
+ *   {key: 97, value: 3},
+ *   {key: 102, value: 4},
+ *   {key: 101, value: 5}
+ * ]
+ * Search the hashmap with array `I`=[98, 99, 102]:
+ * R = H.Map(I) (3)
+ * R should be:
+ * [1, 0, 4]
+ **/
+template <typename IdType>
+class ConcurrentIdHashMap {
+ private:
+  /**
+   * @brief The result state of an attempt to insert.
+   */
+  enum class InsertState {
+    OCCUPIED,  // Indicates that the space where an insertion is being
+               // attempted is already occupied by another element.
+    EXISTED,  // Indicates that the element being inserted already exists in the
+              // map, and thus no insertion is performed.
+    INSERTED  // Indicates that the insertion was successful and a new element
+              // was added to the map.
+  };
+
+ public:
+  /**
+   * @brief Initialize the hashmap with an array of ids. The first `num_seeds`
+   * ids are unique and must be mapped to a contiguous array starting
+   * from 0. The left can be duplicated and the mapping result is not stable.
+   * The unique'ified ids can be accessed through calling `GetUniqueIds()`;
+   *
+   * @param ids The array of the ids to be inserted.
+   * @param num_seeds The number of seed ids.
+   */
+  ConcurrentIdHashMap(const torch::Tensor& ids, size_t num_seeds);
+
+  ConcurrentIdHashMap(const ConcurrentIdHashMap& other) = delete;
+  ConcurrentIdHashMap& operator=(const ConcurrentIdHashMap& other) = delete;
+
+  /**
+   * @brief Get the unique ids for the keys given in the constructor.
+   */
+  const torch::Tensor& GetUniqueIds() const { return unique_ids_; }
+
+  /**
+   * @brief Find mappings of given keys.
+   *
+   * @param ids The keys to map for.
+   *
+   * @return Mapping results corresponding to `ids`.
+   */
+  torch::Tensor MapIds(const torch::Tensor& ids) const;
+
+ private:
+  /**
+   * @brief Get the next position and delta for probing.
+   *
+   * @param[in,out] pos Calculate the next position with quadric probing.
+   * @param[in,out] delta Calculate the next delta by adding 1.
+   */
+  inline void Next(IdType* pos, IdType* delta) const;
+
+  /**
+   * @brief Find the mapping of a given key.
+   *
+   * @param id The key to map for.
+   *
+   * @return Mapping result corresponding to `id`.
+   */
+  inline IdType MapId(const IdType id) const;
+
+  /**
+   * @brief Insert an id into the hash map.
+   *
+   * @param id The id to be inserted.
+   *
+   * @return Whether the `id` is inserted or not.
+   */
+  inline bool Insert(IdType id);
+
+  /**
+   * @brief Set the value for the key in the hash map.
+   *
+   * @param key The key to set for.
+   * @param value The value to be set for the `key`.
+   *
+   * @warning Key must exist.
+   */
+  inline void Set(IdType key, IdType value);
+
+  /**
+   * @brief Insert a key into the hash map.
+   *
+   * @param id The key to be inserted.
+   * @param value The value to be set for the `key`.
+   *
+   */
+  inline void InsertAndSet(IdType key, IdType value);
+
+  /**
+   * @brief Insert a key into the hash map. If the key exists, set the value
+   * with the smaller value.
+   *
+   * @param id The key to be inserted.
+   * @param value The value to be set for the `key`.
+   *
+   */
+  inline void InsertAndSetMin(IdType id, IdType value);
+
+  /**
+   * @brief Attempt to insert the key into the hash map at the given position.
+   *
+   * @param pos The position in the hash map to be inserted at.
+   * @param key The key to be inserted.
+   *
+   * @return The state of the insertion.
+   */
+  inline InsertState AttemptInsertAt(int64_t pos, IdType key);
+
+ private:
+  /**
+   * @brief Hash maps which is used to store all elements.
+   */
+  torch::Tensor hash_map_;
+
+  /**
+   * @brief Holds the ids that are made unique in the constructor.
+   */
+  torch::Tensor unique_ids_;
+
+  /**
+   * @brief Mask which is assisted to get the position in the table
+   * for a key by performing `&` operation with it.
+   */
+  IdType mask_;
+};
+
+}  // namespace sampling
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CONCURRENT_ID_HASH_MAP_H_

graphbolt/src/cuda/common.h

@@ -0,0 +1,231 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/common.h
+ * @brief Common utilities for CUDA
+ */
+#ifndef GRAPHBOLT_CUDA_COMMON_H_
+#define GRAPHBOLT_CUDA_COMMON_H_
+
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+#include <c10/cuda/CUDAException.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime.h>
+#include <thrust/execution_policy.h>
+#include <torch/script.h>
+
+#include <memory>
+#include <unordered_map>
+
+namespace graphbolt {
+namespace cuda {
+
+/**
+ * @brief This class is designed to allocate workspace storage
+ * and to get a nonblocking thrust execution policy
+ * that uses torch's CUDA memory pool and the current cuda stream:
+ *
+ * cuda::CUDAWorkspaceAllocator allocator;
+ * const auto stream = torch::cuda::getDefaultCUDAStream();
+ * const auto exec_policy = thrust::cuda::par_nosync(allocator).on(stream);
+ *
+ * Now, one can pass exec_policy to thrust functions
+ *
+ * To get an integer array of size 1000 whose lifetime is managed by unique_ptr,
+ * use:
+ *
+ * auto int_array = allocator.AllocateStorage<int>(1000);
+ *
+ * int_array.get() gives the raw pointer.
+ */
+template <typename value_t = char>
+struct CUDAWorkspaceAllocator {
+  static_assert(sizeof(char) == 1, "sizeof(char) == 1 should hold.");
+  // Required by thrust to satisfy allocator requirements.
+  using value_type = value_t;
+
+  explicit CUDAWorkspaceAllocator() {
+    at::globalContext().lazyInitDevice(at::kCUDA);
+  }
+
+  template <class U>
+  CUDAWorkspaceAllocator(CUDAWorkspaceAllocator<U> const&) noexcept {}
+
+  CUDAWorkspaceAllocator& operator=(const CUDAWorkspaceAllocator&) = default;
+
+  void operator()(void* ptr) const {
+    c10::cuda::CUDACachingAllocator::raw_delete(ptr);
+  }
+
+  // Required by thrust to satisfy allocator requirements.
+  value_type* allocate(std::ptrdiff_t size) const {
+    return reinterpret_cast<value_type*>(
+        c10::cuda::CUDACachingAllocator::raw_alloc(size * sizeof(value_type)));
+  }
+
+  // Required by thrust to satisfy allocator requirements.
+  void deallocate(value_type* ptr, std::size_t) const { operator()(ptr); }
+
+  template <typename T>
+  std::unique_ptr<T, CUDAWorkspaceAllocator> AllocateStorage(
+      std::size_t size) const {
+    return std::unique_ptr<T, CUDAWorkspaceAllocator>(
+        reinterpret_cast<T*>(
+            c10::cuda::CUDACachingAllocator::raw_alloc(sizeof(T) * size)),
+        *this);
+  }
+};
+
+inline auto GetAllocator() { return CUDAWorkspaceAllocator{}; }
+
+inline auto GetCurrentStream() { return c10::cuda::getCurrentCUDAStream(); }
+
+template <typename T>
+inline bool is_zero(T size) {
+  return size == 0;
+}
+
+template <>
+inline bool is_zero<dim3>(dim3 size) {
+  return size.x == 0 || size.y == 0 || size.z == 0;
+}
+
+#define CUDA_RUNTIME_CHECK(EXPR)                           \
+  do {                                                     \
+    cudaError_t __err = EXPR;                              \
+    if (__err != cudaSuccess) {                            \
+      auto get_error_str_err = cudaGetErrorString(__err);  \
+      AT_ERROR("CUDA runtime error: ", get_error_str_err); \
+    }                                                      \
+  } while (0)
+
+#define CUDA_CALL(func) C10_CUDA_CHECK((func))
+
+#define CUDA_KERNEL_CALL(kernel, nblks, nthrs, shmem, ...)          \
+  {                                                                 \
+    if (!graphbolt::cuda::is_zero((nblks)) &&                       \
+        !graphbolt::cuda::is_zero((nthrs))) {                       \
+      auto stream = graphbolt::cuda::GetCurrentStream();            \
+      (kernel)<<<(nblks), (nthrs), (shmem), stream>>>(__VA_ARGS__); \
+      C10_CUDA_KERNEL_LAUNCH_CHECK();                               \
+    }                                                               \
+  }
+
+#define CUB_CALL(fn, ...)                                                     \
+  {                                                                           \
+    auto allocator = graphbolt::cuda::GetAllocator();                         \
+    auto stream = graphbolt::cuda::GetCurrentStream();                        \
+    size_t workspace_size = 0;                                                \
+    CUDA_CALL(cub::fn(nullptr, workspace_size, __VA_ARGS__, stream));         \
+    auto workspace = allocator.AllocateStorage<char>(workspace_size);         \
+    CUDA_CALL(cub::fn(workspace.get(), workspace_size, __VA_ARGS__, stream)); \
+  }
+
+#define THRUST_CALL(fn, ...)                                                 \
+  [&] {                                                                      \
+    auto allocator = graphbolt::cuda::GetAllocator();                        \
+    auto stream = graphbolt::cuda::GetCurrentStream();                       \
+    const auto exec_policy = thrust::cuda::par_nosync(allocator).on(stream); \
+    return thrust::fn(exec_policy, __VA_ARGS__);                             \
+  }()
+
+/**
+ * @brief This class is designed to handle the copy operation of a single
+ * scalar_t item from a given CUDA device pointer. Later, if the object is cast
+ * into scalar_t, the value can be read.
+ *
+ * auto num_edges = cuda::CopyScalar(indptr.data_ptr<scalar_t>() +
+ *     indptr.size(0) - 1);
+ * // Perform many operations here, they will run as normal.
+ * // We finally need to read num_edges.
+ * auto indices = torch::empty(static_cast<scalar_t>(num_edges));
+ */
+template <typename scalar_t>
+struct CopyScalar {
+  CopyScalar() : is_ready_(true) { init_pinned_storage(); }
+
+  void record(at::cuda::CUDAStream stream = GetCurrentStream()) {
+    copy_event_.record(stream);
+    is_ready_ = false;
+  }
+
+  scalar_t* get() {
+    return reinterpret_cast<scalar_t*>(pinned_scalar_.data_ptr());
+  }
+
+  CopyScalar(const scalar_t* device_ptr) {
+    init_pinned_storage();
+    auto stream = GetCurrentStream();
+    CUDA_CALL(cudaMemcpyAsync(
+        reinterpret_cast<scalar_t*>(pinned_scalar_.data_ptr()), device_ptr,
+        sizeof(scalar_t), cudaMemcpyDeviceToHost, stream));
+    record(stream);
+  }
+
+  operator scalar_t() {
+    if (!is_ready_) {
+      copy_event_.synchronize();
+      is_ready_ = true;
+    }
+    return *get();
+  }
+
+ private:
+  void init_pinned_storage() {
+    pinned_scalar_ = torch::empty(
+        sizeof(scalar_t),
+        c10::TensorOptions().dtype(torch::kBool).pinned_memory(true));
+  }
+
+  torch::Tensor pinned_scalar_;
+  at::cuda::CUDAEvent copy_event_;
+  bool is_ready_;
+};
+
+#define GRAPHBOLT_DISPATCH_ELEMENT_SIZES(element_size, name, ...)             \
+  [&] {                                                                       \
+    switch (element_size) {                                                   \
+      case 1: {                                                               \
+        using element_size_t = uint8_t;                                       \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      case 2: {                                                               \
+        using element_size_t = uint16_t;                                      \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      case 4: {                                                               \
+        using element_size_t = uint32_t;                                      \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      case 8: {                                                               \
+        using element_size_t = uint64_t;                                      \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      case 16: {                                                              \
+        using element_size_t = float4;                                        \
+        return __VA_ARGS__();                                                 \
+      }                                                                       \
+      default:                                                                \
+        TORCH_CHECK(false, name, " with the element_size is not supported!"); \
+        using element_size_t = uint8_t;                                       \
+        return __VA_ARGS__();                                                 \
+    }                                                                         \
+  }()
+
+}  // namespace cuda
+}  // namespace graphbolt
+#endif  // GRAPHBOLT_CUDA_COMMON_H_

graphbolt/src/cuda/cooperative_minibatching_utils.cu

@@ -0,0 +1,159 @@
+/**
+ *   Copyright (c) 2024, mfbalin (Muhammed Fatih Balin)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/cooperative_minibatching_utils.cu
+ * @brief Cooperative Minibatching (arXiv:2310.12403) utility function
+ * implementations in CUDA.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/scatter.h>
+#include <thrust/transform.h>
+
+#include <cub/cub.cuh>
+#include <cuda/functional>
+
+#include "../utils.h"
+#include "./common.h"
+#include "./cooperative_minibatching_utils.cuh"
+#include "./cooperative_minibatching_utils.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace cuda {
+
+torch::Tensor RankAssignment(
+    torch::Tensor nodes, const int64_t rank, const int64_t world_size) {
+  auto part_ids = torch::empty_like(nodes, nodes.options().dtype(kPartDType));
+  auto part_ids_ptr = part_ids.data_ptr<part_t>();
+  AT_DISPATCH_INDEX_TYPES(
+      nodes.scalar_type(), "RankAssignment", ([&] {
+        auto nodes_ptr = nodes.data_ptr<index_t>();
+        THRUST_CALL(
+            transform, nodes_ptr, nodes_ptr + nodes.numel(), part_ids_ptr,
+            ::cuda::proclaim_return_type<part_t>(
+                [rank = static_cast<uint32_t>(rank),
+                 world_size = static_cast<uint32_t>(
+                     world_size)] __device__(index_t id) -> part_t {
+                  return rank_assignment(id, rank, world_size);
+                }));
+      }));
+  return part_ids;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, at::cuda::CUDAEvent>
+RankSortImpl(
+    torch::Tensor nodes, torch::Tensor part_ids, torch::Tensor offsets_dev,
+    const int64_t world_size) {
+  const int num_bits = cuda::NumberOfBits(world_size);
+  const auto num_batches = offsets_dev.numel() - 1;
+  auto offsets_dev_ptr = offsets_dev.data_ptr<int64_t>();
+  auto part_ids_sorted = torch::empty_like(part_ids);
+  auto part_ids2 = part_ids.clone();
+  auto part_ids2_sorted = torch::empty_like(part_ids2);
+  auto nodes_sorted = torch::empty_like(nodes);
+  auto index = torch::arange(nodes.numel(), nodes.options());
+  auto index_sorted = torch::empty_like(index);
+  return AT_DISPATCH_INDEX_TYPES(
+      nodes.scalar_type(), "RankSortImpl", ([&] {
+        CUB_CALL(
+            DeviceSegmentedRadixSort::SortPairs,
+            part_ids.data_ptr<cuda::part_t>(),
+            part_ids_sorted.data_ptr<cuda::part_t>(), nodes.data_ptr<index_t>(),
+            nodes_sorted.data_ptr<index_t>(), nodes.numel(), num_batches,
+            offsets_dev_ptr, offsets_dev_ptr + 1, 0, num_bits);
+        auto offsets = torch::empty(
+            num_batches * world_size + 1, c10::TensorOptions()
+                                              .dtype(offsets_dev.scalar_type())
+                                              .pinned_memory(true));
+        CUB_CALL(
+            DeviceFor::Bulk, num_batches * world_size + 1,
+            [=, part_ids = part_ids_sorted.data_ptr<cuda::part_t>(),
+             offsets = offsets.data_ptr<int64_t>()] __device__(int64_t i) {
+              const auto batch_id = i / world_size;
+              const auto rank = i % world_size;
+              const auto offset_begin = offsets_dev_ptr[batch_id];
+              const auto offset_end =
+                  offsets_dev_ptr[::cuda::std::min(batch_id + 1, num_batches)];
+              offsets[i] = cub::LowerBound(
+                               part_ids + offset_begin,
+                               offset_end - offset_begin, rank) +
+                           offset_begin;
+            });
+        at::cuda::CUDAEvent offsets_event;
+        offsets_event.record();
+        CUB_CALL(
+            DeviceSegmentedRadixSort::SortPairs,
+            part_ids2.data_ptr<cuda::part_t>(),
+            part_ids2_sorted.data_ptr<cuda::part_t>(),
+            index.data_ptr<index_t>(), index_sorted.data_ptr<index_t>(),
+            nodes.numel(), num_batches, offsets_dev_ptr, offsets_dev_ptr + 1, 0,
+            num_bits);
+        auto values = ops::IndptrEdgeIdsImpl(
+            offsets_dev, nodes.scalar_type(), torch::nullopt, nodes.numel());
+        THRUST_CALL(
+            scatter, values.data_ptr<index_t>(),
+            values.data_ptr<index_t>() + values.numel(),
+            index_sorted.data_ptr<index_t>(), index.data_ptr<index_t>());
+        return std::make_tuple(
+            nodes_sorted, index, offsets, std::move(offsets_event));
+      }));
+}
+
+std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> RankSort(
+    const std::vector<torch::Tensor>& nodes_list, const int64_t rank,
+    const int64_t world_size) {
+  const auto num_batches = nodes_list.size();
+  auto nodes = torch::cat(nodes_list, 0);
+  auto offsets = torch::empty(
+      num_batches + 1,
+      c10::TensorOptions().dtype(torch::kInt64).pinned_memory(true));
+  auto offsets_ptr = offsets.data_ptr<int64_t>();
+  offsets_ptr[0] = 0;
+  for (int64_t i = 0; i < num_batches; i++) {
+    offsets_ptr[i + 1] = offsets_ptr[i] + nodes_list[i].numel();
+  }
+  auto part_ids = RankAssignment(nodes, rank, world_size);
+  auto offsets_dev =
+      torch::empty_like(offsets, nodes.options().dtype(offsets.scalar_type()));
+  CUDA_CALL(cudaMemcpyAsync(
+      offsets_dev.data_ptr<int64_t>(), offsets_ptr,
+      sizeof(int64_t) * offsets.numel(), cudaMemcpyHostToDevice,
+      cuda::GetCurrentStream()));
+  auto [nodes_sorted, index_sorted, rank_offsets, rank_offsets_event] =
+      RankSortImpl(nodes, part_ids, offsets_dev, world_size);
+  std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> results;
+  rank_offsets_event.synchronize();
+  for (int64_t i = 0; i < num_batches; i++) {
+    results.emplace_back(
+        nodes_sorted.slice(0, offsets_ptr[i], offsets_ptr[i + 1]),
+        index_sorted.slice(0, offsets_ptr[i], offsets_ptr[i + 1]),
+        rank_offsets.slice(0, i * world_size, (i + 1) * world_size + 1));
+  }
+  return results;
+}
+
+c10::intrusive_ptr<Future<
+    std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>>>
+RankSortAsync(
+    const std::vector<torch::Tensor>& nodes_list, const int64_t rank,
+    const int64_t world_size) {
+  return async(
+      [=] { return RankSort(nodes_list, rank, world_size); },
+      utils::is_on_gpu(nodes_list.at(0)));
+}
+
+}  // namespace cuda
+}  // namespace graphbolt

graphbolt/src/cuda/cooperative_minibatching_utils.cuh

@@ -0,0 +1,55 @@
+/**
+ *   Copyright (c) 2024, mfbalin (Muhammed Fatih Balin)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/cooperative_minibatching_utils.cuh
+ * @brief Cooperative Minibatching (arXiv:2310.12403) utility device functions
+ * in CUDA.
+ */
+#ifndef GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_CUH_
+#define GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_CUH_
+
+#include <curand_kernel.h>
+
+namespace graphbolt {
+namespace cuda {
+
+using part_t = uint8_t;
+constexpr auto kPartDType = torch::kUInt8;
+
+/**
+ * @brief Given a vertex id, the rank of current GPU and the world size, returns
+ * the rank that this id belongs in a deterministic manner.
+ *
+ * @param id         The node id that will mapped to a rank in [0, world_size).
+ * @param rank       The rank of the current GPU.
+ * @param world_size The world size, the total number of cooperating GPUs.
+ *
+ * @return The rank of the GPU the given id is mapped to.
+ */
+template <typename index_t>
+__device__ inline auto rank_assignment(
+    index_t id, uint32_t rank, uint32_t world_size) {
+  // Consider using a faster implementation in the future.
+  constexpr uint64_t kCurandSeed = 999961;  // Any random number.
+  curandStatePhilox4_32_10_t rng;
+  curand_init(kCurandSeed, 0, id, &rng);
+  return (curand(&rng) - rank) % world_size;
+}
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_CUH_

graphbolt/src/cuda/cooperative_minibatching_utils.h

@@ -0,0 +1,96 @@
+/**
+ *   Copyright (c) 2024, mfbalin (Muhammed Fatih Balin)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/cooperative_minibatching_utils.h
+ * @brief Cooperative Minibatching (arXiv:2310.12403) utility function headers
+ * in CUDA.
+ */
+#ifndef GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_H_
+#define GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_H_
+
+#include <ATen/cuda/CUDAEvent.h>
+#include <graphbolt/async.h>
+#include <torch/script.h>
+
+namespace graphbolt {
+namespace cuda {
+
+/**
+ * @brief Given node ids, the rank of current GPU and the world size, returns
+ * the ranks that the given ids belong in a deterministic manner.
+ *
+ * @param nodes      Node id tensor to be mapped to a rank in [0, world_size).
+ * @param rank       Rank of the current GPU.
+ * @param world_size World size, the total number of cooperating GPUs.
+ *
+ * @return The rank tensor of the GPU the given id tensor is mapped to.
+ */
+torch::Tensor RankAssignment(
+    torch::Tensor nodes, int64_t rank, int64_t world_size);
+
+/**
+ * @brief Given node ids, the ranks they belong, the offsets to separate
+ * different node types and world size, returns node ids sorted w.r.t. the ranks
+ * that the given ids belong along with their new positions.
+ *
+ * @param nodes        Node id tensor to be mapped to a rank in [0, world_size).
+ * @param part_ids     Rank tensor the nodes belong to.
+ * @param offsets_dev  Offsets to separate different node types.
+ * @param world_size   World size, the total number of cooperating GPUs.
+ *
+ * @return (sorted_nodes, new_positions, rank_offsets, rank_offsets_event),
+ * where the first one includes sorted nodes, the second contains new positions
+ * of the given nodes, so that sorted_nodes[new_positions] == nodes, and the
+ * third contains the offsets of the sorted_nodes indicating
+ * sorted_nodes[rank_offsets[i]: rank_offsets[i + 1]] contains nodes that
+ * belongs to the `i`th rank. Before accessing rank_offsets on the CPU,
+ * `rank_offsets_event.synchronize()` is required.
+ */
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, at::cuda::CUDAEvent>
+RankSortImpl(
+    torch::Tensor nodes, torch::Tensor part_ids, torch::Tensor offsets_dev,
+    int64_t world_size);
+
+/**
+ * @brief Given a vector of node ids, the rank of current GPU and the world
+ * size, returns node ids sorted w.r.t. the ranks that the given ids belong
+ * along with the original positions.
+ *
+ * @param nodes_list   Node id tensor to be mapped to a rank in [0, world_size).
+ * @param rank         Rank of the current GPU.
+ * @param world_size   World size, the total number of cooperating GPUs.
+ *
+ * @return vector of (sorted_nodes, new_positions, rank_offsets), where the
+ * first one includes sorted nodes, the second contains new positions of the
+ * given nodes, so that sorted_nodes[new_positions] == nodes, and the third
+ * contains the offsets of the sorted_nodes indicating
+ * sorted_nodes[rank_offsets[i]: rank_offsets[i + 1]] contains nodes that
+ * belongs to the `i`th rank.
+ */
+std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>> RankSort(
+    const std::vector<torch::Tensor>& nodes_list, int64_t rank,
+    int64_t world_size);
+
+c10::intrusive_ptr<Future<
+    std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>>>
+RankSortAsync(
+    const std::vector<torch::Tensor>& nodes_list, const int64_t rank,
+    const int64_t world_size);
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CUDA_COOPERATIVE_MINIBATCHING_UTILS_H_

graphbolt/src/cuda/cumsum.cu

@@ -0,0 +1,40 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/cumsum.cu
+ * @brief Cumsum operators implementation on CUDA.
+ */
+#include <cub/cub.cuh>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+torch::Tensor ExclusiveCumSum(torch::Tensor input) {
+  auto result = torch::empty_like(input);
+
+  AT_DISPATCH_INTEGRAL_TYPES(input.scalar_type(), "ExclusiveCumSum", ([&] {
+                               CUB_CALL(
+                                   DeviceScan::ExclusiveSum,
+                                   input.data_ptr<scalar_t>(),
+                                   result.data_ptr<scalar_t>(), input.size(0));
+                             }));
+  return result;
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/expand_indptr.cu

@@ -0,0 +1,145 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/expand_indptr.cu
+ * @brief ExpandIndptr operator implementation on CUDA.
+ */
+#include <thrust/iterator/constant_iterator.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+
+#include <cub/cub.cuh>
+#include <limits>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+template <typename indices_t, typename nodes_t>
+struct RepeatIndex {
+  const nodes_t* nodes;
+  __host__ __device__ auto operator()(indices_t i) {
+    return thrust::make_constant_iterator(nodes ? nodes[i] : i);
+  }
+};
+
+template <typename indices_t, typename nodes_t>
+struct IotaIndex {
+  const nodes_t* nodes;
+  __host__ __device__ auto operator()(indices_t i) {
+    return thrust::make_counting_iterator(nodes ? nodes[i] : 0);
+  }
+};
+
+template <typename indptr_t, typename indices_t>
+struct OutputBufferIndexer {
+  const indptr_t* indptr;
+  indices_t* buffer;
+  __host__ __device__ auto operator()(int64_t i) { return buffer + indptr[i]; }
+};
+
+template <typename indptr_t>
+struct AdjacentDifference {
+  const indptr_t* indptr;
+  __host__ __device__ auto operator()(int64_t i) {
+    return indptr[i + 1] - indptr[i];
+  }
+};
+
+torch::Tensor ExpandIndptrImpl(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> nodes, torch::optional<int64_t> output_size,
+    const bool is_edge_ids_variant) {
+  if (!output_size.has_value()) {
+    output_size = AT_DISPATCH_INTEGRAL_TYPES(
+        indptr.scalar_type(), "ExpandIndptrIndptr[-1]", ([&]() -> int64_t {
+          auto indptr_ptr = indptr.data_ptr<scalar_t>();
+          auto output_size = cuda::CopyScalar{indptr_ptr + indptr.size(0) - 1};
+          return static_cast<scalar_t>(output_size);
+        }));
+  }
+  auto csc_rows =
+      torch::empty(output_size.value(), indptr.options().dtype(dtype));
+
+  AT_DISPATCH_INTEGRAL_TYPES(
+      indptr.scalar_type(), "ExpandIndptrIndptr", ([&] {
+        using indptr_t = scalar_t;
+        auto indptr_ptr = indptr.data_ptr<indptr_t>();
+        AT_DISPATCH_INTEGRAL_TYPES(
+            dtype, "ExpandIndptrIndices", ([&] {
+              using indices_t = scalar_t;
+              auto csc_rows_ptr = csc_rows.data_ptr<indices_t>();
+
+              auto nodes_dtype = nodes ? nodes.value().scalar_type() : dtype;
+              AT_DISPATCH_INTEGRAL_TYPES(
+                  nodes_dtype, "ExpandIndptrNodes", ([&] {
+                    using nodes_t = scalar_t;
+                    auto nodes_ptr =
+                        nodes ? nodes.value().data_ptr<nodes_t>() : nullptr;
+
+                    thrust::counting_iterator<int64_t> iota(0);
+                    auto output_buffer = thrust::make_transform_iterator(
+                        iota, OutputBufferIndexer<indptr_t, indices_t>{
+                                  indptr_ptr, csc_rows_ptr});
+                    auto buffer_sizes = thrust::make_transform_iterator(
+                        iota, AdjacentDifference<indptr_t>{indptr_ptr});
+
+                    const auto num_rows = indptr.size(0) - 1;
+                    constexpr int64_t max_copy_at_once =
+                        std::numeric_limits<int32_t>::max();
+
+                    if (is_edge_ids_variant) {
+                      auto input_buffer = thrust::make_transform_iterator(
+                          iota, IotaIndex<indices_t, nodes_t>{nodes_ptr});
+                      for (int64_t i = 0; i < num_rows; i += max_copy_at_once) {
+                        CUB_CALL(
+                            DeviceCopy::Batched, input_buffer + i,
+                            output_buffer + i, buffer_sizes + i,
+                            std::min(num_rows - i, max_copy_at_once));
+                      }
+                    } else {
+                      auto input_buffer = thrust::make_transform_iterator(
+                          iota, RepeatIndex<indices_t, nodes_t>{nodes_ptr});
+                      for (int64_t i = 0; i < num_rows; i += max_copy_at_once) {
+                        CUB_CALL(
+                            DeviceCopy::Batched, input_buffer + i,
+                            output_buffer + i, buffer_sizes + i,
+                            std::min(num_rows - i, max_copy_at_once));
+                      }
+                    }
+                  }));
+            }));
+      }));
+  return csc_rows;
+}
+
+torch::Tensor ExpandIndptrImpl(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> nodes,
+    torch::optional<int64_t> output_size) {
+  return ExpandIndptrImpl(indptr, dtype, nodes, output_size, false);
+}
+
+torch::Tensor IndptrEdgeIdsImpl(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> offset,
+    torch::optional<int64_t> output_size) {
+  return ExpandIndptrImpl(indptr, dtype, offset, output_size, true);
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/extension/gpu_cache.cu

@@ -0,0 +1,133 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/gpu_cache.cu
+ * @brief GPUCache implementation on CUDA.
+ */
+#include <numeric>
+
+#include "../common.h"
+#include "./gpu_cache.h"
+
+namespace graphbolt {
+namespace cuda {
+
+GpuCache::GpuCache(const std::vector<int64_t> &shape, torch::ScalarType dtype) {
+  TORCH_CHECK(shape.size() >= 2, "Shape must at least have 2 dimensions.");
+  const auto num_items = shape[0];
+  TORCH_CHECK(
+      num_items > 0, "The capacity of GpuCache needs to be a positive.");
+  const int64_t num_feats =
+      std::accumulate(shape.begin() + 1, shape.end(), 1ll, std::multiplies<>());
+  const int element_size =
+      torch::empty(1, torch::TensorOptions().dtype(dtype)).element_size();
+  num_bytes_ = num_feats * element_size;
+  num_float_feats_ = (num_bytes_ + sizeof(float) - 1) / sizeof(float);
+  cache_ = std::make_unique<gpu_cache_t>(
+      (num_items + bucket_size - 1) / bucket_size, num_float_feats_);
+  shape_ = shape;
+  shape_[0] = -1;
+  dtype_ = dtype;
+  device_id_ = cuda::GetCurrentStream().device_index();
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> GpuCache::Query(
+    torch::Tensor keys) {
+  TORCH_CHECK(keys.device().is_cuda(), "Keys should be on a CUDA device.");
+  TORCH_CHECK(
+      keys.device().index() == device_id_,
+      "Keys should be on the correct CUDA device.");
+  TORCH_CHECK(keys.sizes().size() == 1, "Keys should be a 1D tensor.");
+  keys = keys.to(torch::kLong);
+  auto values = torch::empty(
+      {keys.size(0), num_float_feats_}, keys.options().dtype(torch::kFloat));
+  auto missing_index =
+      torch::empty(keys.size(0), keys.options().dtype(torch::kLong));
+  auto missing_keys =
+      torch::empty(keys.size(0), keys.options().dtype(torch::kLong));
+  auto allocator = cuda::GetAllocator();
+  auto missing_len_device = allocator.AllocateStorage<size_t>(1);
+  cache_->Query(
+      reinterpret_cast<const key_t *>(keys.data_ptr()), keys.size(0),
+      values.data_ptr<float>(),
+      reinterpret_cast<uint64_t *>(missing_index.data_ptr()),
+      reinterpret_cast<key_t *>(missing_keys.data_ptr()),
+      missing_len_device.get(), cuda::GetCurrentStream());
+  values = values.view(torch::kByte)
+               .slice(1, 0, num_bytes_)
+               .view(dtype_)
+               .view(shape_);
+  cuda::CopyScalar<size_t> missing_len(missing_len_device.get());
+  missing_index = missing_index.slice(0, 0, static_cast<size_t>(missing_len));
+  missing_keys = missing_keys.slice(0, 0, static_cast<size_t>(missing_len));
+  return std::make_tuple(values, missing_index, missing_keys);
+}
+
+c10::intrusive_ptr<Future<std::vector<torch::Tensor>>> GpuCache::QueryAsync(
+    torch::Tensor keys) {
+  return async(
+      [=] {
+        auto [values, missing_index, missing_keys] = Query(keys);
+        return std::vector{values, missing_index, missing_keys};
+      },
+      true);
+}
+
+void GpuCache::Replace(torch::Tensor keys, torch::Tensor values) {
+  TORCH_CHECK(keys.device().is_cuda(), "Keys should be on a CUDA device.");
+  TORCH_CHECK(
+      keys.device().index() == device_id_,
+      "Keys should be on the correct CUDA device.");
+  TORCH_CHECK(values.device().is_cuda(), "Keys should be on a CUDA device.");
+  TORCH_CHECK(
+      values.device().index() == device_id_,
+      "Values should be on the correct CUDA device.");
+  TORCH_CHECK(
+      keys.size(0) == values.size(0),
+      "The first dimensions of keys and values must match.");
+  TORCH_CHECK(
+      std::equal(shape_.begin() + 1, shape_.end(), values.sizes().begin() + 1),
+      "Values should have the correct dimensions.");
+  TORCH_CHECK(
+      values.scalar_type() == dtype_, "Values should have the correct dtype.");
+  if (keys.numel() == 0) return;
+  keys = keys.to(torch::kLong);
+  torch::Tensor float_values;
+  if (num_bytes_ % sizeof(float) != 0) {
+    float_values = torch::empty(
+        {values.size(0), num_float_feats_},
+        values.options().dtype(torch::kFloat));
+    float_values.view(torch::kByte)
+        .slice(1, 0, num_bytes_)
+        .copy_(values.view(torch::kByte).view({values.size(0), -1}));
+  } else {
+    float_values = values.view(torch::kByte)
+                       .view({values.size(0), -1})
+                       .view(torch::kFloat)
+                       .contiguous();
+  }
+  cache_->Replace(
+      reinterpret_cast<const key_t *>(keys.data_ptr()), keys.size(0),
+      float_values.data_ptr<float>(), cuda::GetCurrentStream());
+}
+
+c10::intrusive_ptr<GpuCache> GpuCache::Create(
+    const std::vector<int64_t> &shape, torch::ScalarType dtype) {
+  return c10::make_intrusive<GpuCache>(shape, dtype);
+}
+
+}  // namespace cuda
+}  // namespace graphbolt

graphbolt/src/cuda/extension/gpu_cache.h

@@ -0,0 +1,83 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/gpu_cache.h
+ * @brief Header file of HugeCTR gpu_cache wrapper.
+ */
+
+#ifndef GRAPHBOLT_GPU_CACHE_H_
+#define GRAPHBOLT_GPU_CACHE_H_
+
+#include <graphbolt/async.h>
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+
+#include <limits>
+#include <nv_gpu_cache.hpp>
+
+namespace graphbolt {
+namespace cuda {
+
+class GpuCache : public torch::CustomClassHolder {
+  using key_t = long long;
+  constexpr static int set_associativity = 2;
+  constexpr static int WARP_SIZE = 32;
+  constexpr static int bucket_size = WARP_SIZE * set_associativity;
+  using gpu_cache_t = ::gpu_cache::gpu_cache<
+      key_t, uint64_t, std::numeric_limits<key_t>::max(), set_associativity,
+      WARP_SIZE>;
+
+ public:
+  /**
+   * @brief Constructor for the GpuCache struct.
+   *
+   * @param shape The shape of the GPU cache.
+   * @param dtype The datatype of items to be stored.
+   */
+  GpuCache(const std::vector<int64_t>& shape, torch::ScalarType dtype);
+
+  GpuCache() = default;
+
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Query(
+      torch::Tensor keys);
+
+  c10::intrusive_ptr<Future<std::vector<torch::Tensor>>> QueryAsync(
+      torch::Tensor keys);
+
+  void Replace(torch::Tensor keys, torch::Tensor values);
+
+  static c10::intrusive_ptr<GpuCache> Create(
+      const std::vector<int64_t>& shape, torch::ScalarType dtype);
+
+ private:
+  std::vector<int64_t> shape_;
+  torch::ScalarType dtype_;
+  std::unique_ptr<gpu_cache_t> cache_;
+  int64_t num_bytes_;
+  int64_t num_float_feats_;
+  torch::DeviceIndex device_id_;
+};
+
+// The cu file in HugeCTR gpu cache uses unsigned int and long long.
+// Changing to int64_t results in a mismatch of template arguments.
+static_assert(
+    sizeof(long long) == sizeof(int64_t),
+    "long long and int64_t needs to have the same size.");  // NOLINT
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_GPU_CACHE_H_

graphbolt/src/cuda/extension/gpu_graph_cache.cu

@@ -0,0 +1,546 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/gpu_graph_cache.cu
+ * @brief GPU graph cache implementation on CUDA.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/gather.h>
+#include <thrust/transform.h>
+
+#include <cstddef>
+#include <cub/cub.cuh>
+#include <cuco/static_map.cuh>
+#include <cuda/std/atomic>
+#include <cuda/stream_ref>
+#include <limits>
+#include <numeric>
+#include <type_traits>
+
+#include "../common.h"
+#include "../utils.h"
+#include "./gpu_graph_cache.h"
+
+namespace graphbolt {
+namespace cuda {
+
+namespace {
+
+constexpr int cg_size = 1;
+template <typename index_t>
+using probing_t =
+    cuco::linear_probing<cg_size, cuco::default_hash_function<index_t>>;
+template <typename index_t>
+using allocator_t = cuda::CUDAWorkspaceAllocator<cuco::pair<index_t, index_t>>;
+template <typename index_t>
+using map_t = cuco::static_map<
+    index_t, index_t, cuco::extent<int64_t>, ::cuda::thread_scope_device,
+    thrust::equal_to<index_t>, probing_t<index_t>, allocator_t<index_t>>;
+
+template <typename index_t, typename map_t>
+__global__ void _Insert(
+    const int64_t num_nodes, const index_t num_existing, const index_t* seeds,
+    const index_t* missing_indices, const index_t* indices, map_t map) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  while (i < num_nodes) {
+    const auto key = seeds[missing_indices[indices[i]]];
+
+    auto slot = map.find(key);
+    slot->second = num_existing + i;
+
+    i += stride;
+  }
+}
+
+/**
+ * @brief For node ids not in the cache, it keeps their access count inside
+ * a hash table as (v, -c) where v is the node id and c is the access count.
+ * When c == -threshold, it means that v will be inserted into the cache
+ * during the call to the replace method. Once v is inserted into the cache,
+ * c is assigned to a nonnegative value and indicates the local id of vertex
+ * v in the cache.
+ *
+ * @param num_nodes The number of node ids.
+ * @param seeds The node ids the cache is being queried with.
+ * @param positions Holds the values found in the hash table.
+ * @param map The hash table holding (v, -c) or (v, local_id).
+ *
+ */
+template <typename index_t, typename map_t>
+__global__ void _QueryAndIncrement(
+    const int64_t num_nodes, const index_t* seeds, index_t* positions,
+    map_t map) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  while (i < num_nodes) {
+    const auto key = seeds[i];
+
+    constexpr index_t minusONE = -1;
+    auto [slot, is_new_key] = map.insert_and_find(cuco::pair{key, minusONE});
+
+    int64_t position = -1;
+
+    if (!is_new_key) {
+      auto ref = ::cuda::atomic_ref<index_t, ::cuda::thread_scope_device>{
+          slot->second};
+      position = ref.load(::cuda::memory_order_relaxed);
+      if (position < 0) {
+        position = ref.fetch_add(-1, ::cuda::memory_order_relaxed) - 1;
+      }
+    }
+
+    positions[i] = position;
+
+    i += stride;
+  }
+}
+
+constexpr int kIntBlockSize = 512;
+}  // namespace
+
+c10::intrusive_ptr<GpuGraphCache> GpuGraphCache::Create(
+    const int64_t num_edges, const int64_t threshold,
+    torch::ScalarType indptr_dtype, std::vector<torch::ScalarType> dtypes,
+    bool has_original_edge_ids) {
+  return c10::make_intrusive<GpuGraphCache>(
+      num_edges, threshold, indptr_dtype, dtypes, has_original_edge_ids);
+}
+
+GpuGraphCache::GpuGraphCache(
+    const int64_t num_edges, const int64_t threshold,
+    torch::ScalarType indptr_dtype, std::vector<torch::ScalarType> dtypes,
+    bool has_original_edge_ids) {
+  const int64_t initial_node_capacity = 1024;
+  AT_DISPATCH_INDEX_TYPES(
+      dtypes.at(0), "GpuGraphCache::GpuGraphCache", ([&] {
+        auto map_temp = map_t<index_t>{
+            initial_node_capacity,
+            kDoubleLoadFactor,
+            cuco::empty_key{static_cast<index_t>(-1)},
+            cuco::empty_value{std::numeric_limits<index_t>::lowest()},
+            {},
+            probing_t<index_t>{},
+            {},
+            {},
+            allocator_t<index_t>{},
+            ::cuda::stream_ref{cuda::GetCurrentStream()}};
+        map_ = new map_t<index_t>{std::move(map_temp)};
+      }));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
+  const auto options = torch::TensorOptions().device(c10::DeviceType::CUDA);
+  TORCH_CHECK(threshold > 0, "Threshold should be a position integer.");
+  threshold_ = threshold;
+  device_id_ = cuda::GetCurrentStream().device_index();
+  map_size_ = 0;
+  num_nodes_ = 0;
+  num_edges_ = 0;
+  indptr_ =
+      torch::zeros(initial_node_capacity + 1, options.dtype(indptr_dtype));
+  if (!has_original_edge_ids) {
+    offset_ = torch::empty(indptr_.size(0) - 1, indptr_.options());
+  }
+  for (auto dtype : dtypes) {
+    cached_edge_tensors_.push_back(
+        torch::empty(num_edges, options.dtype(dtype)));
+  }
+}
+
+GpuGraphCache::~GpuGraphCache() {
+  AT_DISPATCH_INDEX_TYPES(
+      cached_edge_tensors_.at(0).scalar_type(), "GpuGraphCache::GpuGraphCache",
+      ([&] { delete reinterpret_cast<map_t<index_t>*>(map_); }));
+}
+
+std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t> GpuGraphCache::Query(
+    torch::Tensor seeds) {
+  TORCH_CHECK(seeds.device().is_cuda(), "Seeds should be on a CUDA device.");
+  TORCH_CHECK(
+      seeds.device().index() == device_id_,
+      "Seeds should be on the correct CUDA device.");
+  TORCH_CHECK(seeds.sizes().size() == 1, "Keys should be a 1D tensor.");
+  std::lock_guard lock(mtx_);
+  auto allocator = cuda::GetAllocator();
+  auto index_dtype = cached_edge_tensors_.at(0).scalar_type();
+  const dim3 block(kIntBlockSize);
+  const dim3 grid((seeds.size(0) + kIntBlockSize - 1) / kIntBlockSize);
+  return AT_DISPATCH_INDEX_TYPES(
+      index_dtype, "GpuGraphCache::Query", ([&] {
+        auto map = reinterpret_cast<map_t<index_t>*>(map_);
+        while ((
+            map_size_ + seeds.size(0) >= map->capacity() * kDoubleLoadFactor)) {
+          map->rehash_async(
+              map->capacity() * kIntGrowthFactor,
+              ::cuda::stream_ref{cuda::GetCurrentStream()});
+        }
+        auto positions = torch::empty_like(seeds);
+        CUDA_KERNEL_CALL(
+            _QueryAndIncrement, grid, block, 0,
+            static_cast<int64_t>(seeds.size(0)), seeds.data_ptr<index_t>(),
+            positions.data_ptr<index_t>(), map->ref(cuco::insert_and_find));
+        auto num_threshold_new_hit =
+            allocator.AllocateStorage<thrust::tuple<int64_t, int64_t, int64_t>>(
+                1);
+        // Since threshold_ is a class member, we want the lambda functions
+        // below to only capture this particular variable by reassigning it to a
+        // local variable.
+        const auto threshold = -threshold_;
+        auto is_threshold_new_hit = thrust::make_transform_iterator(
+            positions.data_ptr<index_t>(), [=] __host__ __device__(index_t x) {
+              int64_t is_threshold = x == threshold;
+              int64_t is_new = x == -1;
+              int64_t is_hit = x >= 0;
+              return thrust::make_tuple(is_threshold, is_new, is_hit);
+            });
+        CUB_CALL(
+            DeviceReduce::Reduce, is_threshold_new_hit,
+            num_threshold_new_hit.get(), positions.size(0),
+            [] __host__ __device__(
+                const thrust::tuple<int64_t, int64_t, int64_t>& a,
+                const thrust::tuple<int64_t, int64_t, int64_t>& b) {
+              return thrust::make_tuple(
+                  thrust::get<0>(a) + thrust::get<0>(b),
+                  thrust::get<1>(a) + thrust::get<1>(b),
+                  thrust::get<2>(a) + thrust::get<2>(b));
+            },
+            thrust::tuple<int64_t, int64_t, int64_t>{});
+        CopyScalar num_threshold_new_hit_cpu{num_threshold_new_hit.get()};
+        thrust::counting_iterator<index_t> iota{0};
+        auto position_and_index =
+            thrust::make_zip_iterator(positions.data_ptr<index_t>(), iota);
+        auto output_positions = torch::empty_like(seeds);
+        auto output_indices = torch::empty_like(seeds);
+        auto output_position_and_index = thrust::make_zip_iterator(
+            output_positions.data_ptr<index_t>(),
+            output_indices.data_ptr<index_t>());
+        CUB_CALL(
+            DevicePartition::If, position_and_index, output_position_and_index,
+            cub::DiscardOutputIterator{}, seeds.size(0),
+            [] __device__(thrust::tuple<index_t, index_t> & x) {
+              return thrust::get<0>(x) >= 0;
+            });
+        const auto [num_threshold, num_new, num_hit] =
+            static_cast<thrust::tuple<int64_t, int64_t, int64_t>>(
+                num_threshold_new_hit_cpu);
+        map_size_ += num_new;
+
+        return std::make_tuple(
+            output_indices, output_positions, num_hit, num_threshold);
+      }));
+}
+
+c10::intrusive_ptr<
+    Future<std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t>>>
+GpuGraphCache::QueryAsync(torch::Tensor seeds) {
+  return async([=] { return Query(seeds); }, true);
+}
+
+std::tuple<torch::Tensor, std::vector<torch::Tensor>> GpuGraphCache::Replace(
+    torch::Tensor seeds, torch::Tensor indices, torch::Tensor positions,
+    int64_t num_hit, int64_t num_threshold, torch::Tensor indptr,
+    std::vector<torch::Tensor> edge_tensors) {
+  const auto with_edge_ids = offset_.has_value();
+  // The last element of edge_tensors has the edge ids.
+  const auto num_tensors = edge_tensors.size() - with_edge_ids;
+  TORCH_CHECK(
+      num_tensors == cached_edge_tensors_.size(),
+      "Same number of tensors need to be passed!");
+  const auto num_nodes = seeds.size(0);
+  TORCH_CHECK(
+      indptr.size(0) == num_nodes - num_hit + 1,
+      "(indptr.size(0) == seeds.size(0) - num_hit + 1) failed.");
+  std::lock_guard lock(mtx_);
+  const int64_t num_buffers = num_nodes * num_tensors;
+  auto allocator = cuda::GetAllocator();
+  auto index_dtype = cached_edge_tensors_.at(0).scalar_type();
+  return AT_DISPATCH_INDEX_TYPES(
+      index_dtype, "GpuGraphCache::Replace", ([&] {
+        using indices_t = index_t;
+        return AT_DISPATCH_INDEX_TYPES(
+            indptr_.scalar_type(), "GpuGraphCache::Replace::copy_prep", ([&] {
+              using indptr_t = index_t;
+              static_assert(
+                  sizeof(int64_t) == sizeof(void*),
+                  "Pointers have to be 64-bit.");
+              static_assert(
+                  sizeof(std::byte) == 1, "Byte needs to have a size of 1.");
+              auto cache_missing_dtype = torch::empty(
+                  // Below, we use this storage to store a tuple of 4 elements,
+                  // since each element is 64-bit, we need 4x int64 storage.
+                  4 * num_tensors, c10::TensorOptions()
+                                       .dtype(torch::kInt64)
+                                       .pinned_memory(true));
+              auto cache_missing_dtype_ptr =
+                  reinterpret_cast<::cuda::std::tuple<
+                      std::byte*, std::byte*, int64_t, int64_t>*>(
+                      cache_missing_dtype.data_ptr());
+              int64_t total_size = 0;
+              for (size_t i = 0; i < num_tensors; i++) {
+                TORCH_CHECK(
+                    cached_edge_tensors_[i].scalar_type() ==
+                        edge_tensors[i].scalar_type(),
+                    "The dtypes of edge tensors must match.");
+                if (i > 0) {
+                  TORCH_CHECK(
+                      edge_tensors[i - 1].size(0) == edge_tensors[i].size(0),
+                      "The missing edge tensors should have identical size.");
+                }
+                const int64_t element_size = edge_tensors[i].element_size();
+                cache_missing_dtype_ptr[i] = {
+                    reinterpret_cast<std::byte*>(
+                        cached_edge_tensors_[i].data_ptr()),
+                    reinterpret_cast<std::byte*>(edge_tensors[i].data_ptr()),
+                    element_size, total_size};
+                total_size += element_size;
+              }
+              auto cache_missing_dtype_dev = allocator.AllocateStorage<
+                  ::cuda::std::tuple<std::byte*, std::byte*, int64_t, int64_t>>(
+                  num_tensors);
+              THRUST_CALL(
+                  copy_n, cache_missing_dtype_ptr, num_tensors,
+                  cache_missing_dtype_dev.get());
+
+              auto input = allocator.AllocateStorage<std::byte*>(num_buffers);
+              auto input_size =
+                  allocator.AllocateStorage<size_t>(num_buffers + 1);
+              torch::optional<torch::Tensor> edge_id_offsets;
+              if (with_edge_ids) {
+                edge_id_offsets = torch::empty(
+                    num_nodes,
+                    seeds.options().dtype(offset_.value().scalar_type()));
+              }
+              const auto cache_missing_dtype_dev_ptr =
+                  cache_missing_dtype_dev.get();
+              const auto indices_ptr = indices.data_ptr<indices_t>();
+              const auto positions_ptr = positions.data_ptr<indices_t>();
+              const auto input_ptr = input.get();
+              const auto input_size_ptr = input_size.get();
+              const auto edge_id_offsets_ptr =
+                  edge_id_offsets ? edge_id_offsets->data_ptr<indptr_t>()
+                                  : nullptr;
+              const auto cache_indptr = indptr_.data_ptr<indptr_t>();
+              const auto missing_indptr = indptr.data_ptr<indptr_t>();
+              const auto cache_offset =
+                  offset_ ? offset_->data_ptr<indptr_t>() : nullptr;
+              const auto missing_edge_ids =
+                  edge_id_offsets ? edge_tensors.back().data_ptr<indptr_t>()
+                                  : nullptr;
+              CUB_CALL(DeviceFor::Bulk, num_buffers, [=] __device__(int64_t i) {
+                const auto tensor_idx = i / num_nodes;
+                const auto idx = i % num_nodes;
+                const auto pos = positions_ptr[idx];
+                const auto original_idx = indices_ptr[idx];
+                const auto [cache_ptr, missing_ptr, size, cum_size] =
+                    cache_missing_dtype_dev_ptr[tensor_idx];
+                const auto is_cached = pos >= 0;
+                const auto offset = is_cached ? cache_indptr[pos]
+                                              : missing_indptr[idx - num_hit];
+                const auto offset_end = is_cached
+                                            ? cache_indptr[pos + 1]
+                                            : missing_indptr[idx - num_hit + 1];
+                const auto out_idx = tensor_idx * num_nodes + original_idx;
+
+                input_ptr[out_idx] =
+                    (is_cached ? cache_ptr : missing_ptr) + offset * size;
+                input_size_ptr[out_idx] = size * (offset_end - offset);
+                if (edge_id_offsets_ptr && i < num_nodes) {
+                  const auto edge_id =
+                      is_cached ? cache_offset[pos] : missing_edge_ids[offset];
+                  edge_id_offsets_ptr[out_idx] = edge_id;
+                }
+              });
+              auto output_indptr = torch::empty(
+                  num_nodes + 1, seeds.options().dtype(indptr_.scalar_type()));
+              auto output_indptr_ptr = output_indptr.data_ptr<indptr_t>();
+              const auto element_size =
+                  ::cuda::std::get<2>(cache_missing_dtype_ptr[0]);
+              auto input_indegree = thrust::make_transform_iterator(
+                  input_size_ptr, [=] __host__ __device__(size_t x) {
+                    return x / element_size;
+                  });
+              CUB_CALL(
+                  DeviceScan::ExclusiveSum, input_indegree, output_indptr_ptr,
+                  num_nodes + 1);
+              CopyScalar output_size{output_indptr_ptr + num_nodes};
+
+              if (num_threshold > 0) {
+                // Insert the vertices whose access count equal threshold.
+                auto missing_positions = positions.slice(0, num_hit);
+                auto missing_indices = indices.slice(0, num_hit);
+
+                thrust::counting_iterator<indices_t> iota{0};
+                auto threshold = -threshold_;
+                auto is_threshold = thrust::make_transform_iterator(
+                    missing_positions.data_ptr<indices_t>(),
+                    [=] __host__ __device__(indices_t x) {
+                      return x == threshold;
+                    });
+                auto output_indices =
+                    torch::empty(num_threshold, seeds.options());
+                CUB_CALL(
+                    DeviceSelect::Flagged, iota, is_threshold,
+                    output_indices.data_ptr<indices_t>(),
+                    cub::DiscardOutputIterator{}, missing_positions.size(0));
+                auto [in_degree, sliced_indptr] =
+                    ops::SliceCSCIndptr(indptr, output_indices);
+                while (num_nodes_ + num_threshold >= indptr_.size(0)) {
+                  auto new_indptr = torch::empty(
+                      indptr_.size(0) * kIntGrowthFactor, indptr_.options());
+                  new_indptr.slice(0, 0, indptr_.size(0)) = indptr_;
+                  indptr_ = new_indptr;
+                  if (offset_) {
+                    auto new_offset =
+                        torch::empty(indptr_.size(0) - 1, offset_->options());
+                    new_offset.slice(0, 0, offset_->size(0)) = *offset_;
+                    offset_ = new_offset;
+                  }
+                }
+                torch::Tensor sindptr;
+                bool enough_space;
+                torch::optional<int64_t> cached_output_size;
+                for (size_t i = 0; i < num_tensors; i++) {
+                  torch::Tensor sindices;
+                  std::tie(sindptr, sindices) = ops::IndexSelectCSCImpl(
+                      in_degree, sliced_indptr, edge_tensors[i], output_indices,
+                      indptr.size(0) - 2, cached_output_size);
+                  cached_output_size = sindices.size(0);
+                  enough_space = num_edges_ + *cached_output_size <=
+                                 cached_edge_tensors_[i].size(0);
+                  if (enough_space) {
+                    cached_edge_tensors_[i].slice(
+                        0, num_edges_, num_edges_ + *cached_output_size) =
+                        sindices;
+                  } else
+                    break;
+                }
+                if (enough_space) {
+                  auto num_edges = num_edges_;
+                  if (offset_) {
+                    auto transform_input_it = thrust::make_zip_iterator(
+                        sindptr.data_ptr<indptr_t>() + 1,
+                        sliced_indptr.data_ptr<indptr_t>());
+                    auto transform_output_it = thrust::make_zip_iterator(
+                        indptr_.data_ptr<indptr_t>() + num_nodes_ + 1,
+                        offset_->data_ptr<indptr_t>() + num_nodes_);
+                    THRUST_CALL(
+                        transform, transform_input_it,
+                        transform_input_it + sindptr.size(0) - 1,
+                        transform_output_it,
+                        [=] __host__ __device__(
+                            const thrust::tuple<indptr_t, indptr_t>& x) {
+                          return thrust::make_tuple(
+                              thrust::get<0>(x) + num_edges,
+                              missing_edge_ids[thrust::get<1>(x)]);
+                        });
+                  } else {
+                    THRUST_CALL(
+                        transform, sindptr.data_ptr<indptr_t>() + 1,
+                        sindptr.data_ptr<indptr_t>() + sindptr.size(0),
+                        indptr_.data_ptr<indptr_t>() + num_nodes_ + 1,
+                        [=] __host__ __device__(const indptr_t& x) {
+                          return x + num_edges;
+                        });
+                  }
+                  auto map = reinterpret_cast<map_t<indices_t>*>(map_);
+                  const dim3 block(kIntBlockSize);
+                  const dim3 grid(
+                      (num_threshold + kIntBlockSize - 1) / kIntBlockSize);
+                  CUDA_KERNEL_CALL(
+                      _Insert, grid, block, 0, output_indices.size(0),
+                      static_cast<indices_t>(num_nodes_),
+                      seeds.data_ptr<indices_t>(),
+                      missing_indices.data_ptr<indices_t>(),
+                      output_indices.data_ptr<indices_t>(),
+                      map->ref(cuco::find));
+                  num_edges_ += *cached_output_size;
+                  num_nodes_ += num_threshold;
+                }
+              }
+
+              constexpr int alignment = 128;
+              const auto output_allocation_count =
+                  (static_cast<indptr_t>(output_size) + alignment - 1) /
+                  alignment * alignment;
+              auto output_allocation = torch::empty(
+                  output_allocation_count * total_size,
+                  seeds.options().dtype(torch::kInt8));
+              const auto output_allocation_ptr =
+                  output_allocation.data_ptr<int8_t>();
+
+              std::vector<torch::Tensor> output_edge_tensors;
+              for (size_t i = 0; i < num_tensors; i++) {
+                const auto cum_size =
+                    ::cuda::std::get<3>(cache_missing_dtype_ptr[i]);
+                output_edge_tensors.push_back(
+                    output_allocation
+                        .slice(0, cum_size * output_allocation_count)
+                        .view(edge_tensors[i].scalar_type())
+                        .slice(0, 0, static_cast<indptr_t>(output_size)));
+              }
+              if (edge_id_offsets) {
+                // Append the edge ids as the last element of the output.
+                output_edge_tensors.push_back(ops::IndptrEdgeIdsImpl(
+                    output_indptr, output_indptr.scalar_type(),
+                    *edge_id_offsets,
+                    static_cast<int64_t>(static_cast<indptr_t>(output_size))));
+              }
+
+              {
+                thrust::counting_iterator<int64_t> iota{0};
+                auto output_buffer_it = thrust::make_transform_iterator(
+                    iota, [=] __host__ __device__(int64_t i) {
+                      const auto tensor_idx = i / num_nodes;
+                      const auto idx = i % num_nodes;
+                      const auto offset = output_indptr_ptr[idx];
+                      const auto [_0, _1, size, cum_size] =
+                          cache_missing_dtype_dev_ptr[tensor_idx];
+                      return output_allocation_ptr +
+                             cum_size * output_allocation_count + offset * size;
+                    });
+                constexpr int64_t max_copy_at_once =
+                    std::numeric_limits<int32_t>::max();
+                for (int64_t i = 0; i < num_buffers; i += max_copy_at_once) {
+                  CUB_CALL(
+                      DeviceMemcpy::Batched, input.get() + i,
+                      output_buffer_it + i, input_size_ptr + i,
+                      std::min(num_buffers - i, max_copy_at_once));
+                }
+              }
+
+              return std::make_tuple(output_indptr, output_edge_tensors);
+            }));
+      }));
+}
+
+c10::intrusive_ptr<
+    Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>>
+GpuGraphCache::ReplaceAsync(
+    torch::Tensor seeds, torch::Tensor indices, torch::Tensor positions,
+    int64_t num_hit, int64_t num_threshold, torch::Tensor indptr,
+    std::vector<torch::Tensor> edge_tensors) {
+  return async(
+      [=] {
+        return Replace(
+            seeds, indices, positions, num_hit, num_threshold, indptr,
+            edge_tensors);
+      },
+      true);
+}
+
+}  // namespace cuda
+}  // namespace graphbolt

graphbolt/src/cuda/extension/gpu_graph_cache.h

@@ -0,0 +1,136 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/gpu_graph_cache.h
+ * @brief Header file of GPU graph cache.
+ */
+
+#ifndef GRAPHBOLT_GPU_GRAPH_CACHE_H_
+#define GRAPHBOLT_GPU_GRAPH_CACHE_H_
+
+#include <graphbolt/async.h>
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+
+#include <mutex>
+
+namespace graphbolt {
+namespace cuda {
+
+class GpuGraphCache : public torch::CustomClassHolder {
+  // The load factor of the constructed hash table.
+  static constexpr double kDoubleLoadFactor = 0.8;
+  // The growth factor of the hash table and the dynamically sized indptr
+  // tensor.
+  static constexpr int kIntGrowthFactor = 2;
+
+ public:
+  /**
+   * @brief Constructor for the GpuGraphCache struct.
+   *
+   * @param num_edges The edge capacity of GPU cache.
+   * @param threshold The access threshold before a vertex neighborhood is
+   * cached.
+   * @param indptr_dtype The node id datatype.
+   * @param dtypes The dtypes of the edge tensors to be cached. dtypes[0] is
+   * reserved for the indices edge tensor holding node ids.
+   * @param has_original_edge_ids Whether the graph to be cached has original
+   * edge ids.
+   */
+  GpuGraphCache(
+      const int64_t num_edges, const int64_t threshold,
+      torch::ScalarType indptr_dtype, std::vector<torch::ScalarType> dtypes,
+      bool has_original_edge_ids);
+
+  GpuGraphCache() = default;
+
+  ~GpuGraphCache();
+
+  /**
+   * @brief Queries the cache. Returns tensors indicating which elements are
+   * missing.
+   *
+   * @param seeds The node ids to query the cache with.
+   *
+   * @return
+   * (torch::Tensor, torch::Tensor, int64_t, int64_t) index, position,
+   * number of cache hits and number of ids that will enter the cache.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t> Query(
+      torch::Tensor seeds);
+
+  c10::intrusive_ptr<
+      Future<std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t>>>
+  QueryAsync(torch::Tensor seeds);
+
+  /**
+   * @brief After the graph structure for the missing node ids are fetched, it
+   * inserts the node ids which passes the threshold and returns the final
+   * output graph structure, combining the information in the cache with the
+   * graph structure for the missing node ids.
+   *
+   * @param seeds The node ids that the cache was queried with.
+   * @param indices seeds[indices[:num_hit]] gives us the node ids that were
+   * found in the cache
+   * @param positions positions[:num_hit] gives where the node ids can be found
+   * in the cache.
+   * @param num_hit The number of seeds that are already in the cache.
+   * @param num_threshold The number of seeds among the missing node ids that
+   * will be inserted into the cache.
+   * @param indptr The indptr for the missing seeds fetched from remote.
+   * @param edge_tensors The edge tensors for the missing seeds. The last
+   * element of edge_tensors is treated as the edge ids tensor with
+   * indptr_dtype.
+   *
+   * @return (torch::Tensor, std::vector<torch::Tensor>) The final indptr and
+   * edge_tensors, directly corresponding to the seeds tensor.
+   */
+  std::tuple<torch::Tensor, std::vector<torch::Tensor>> Replace(
+      torch::Tensor seeds, torch::Tensor indices, torch::Tensor positions,
+      int64_t num_hit, int64_t num_threshold, torch::Tensor indptr,
+      std::vector<torch::Tensor> edge_tensors);
+
+  c10::intrusive_ptr<
+      Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>>
+  ReplaceAsync(
+      torch::Tensor seeds, torch::Tensor indices, torch::Tensor positions,
+      int64_t num_hit, int64_t num_threshold, torch::Tensor indptr,
+      std::vector<torch::Tensor> edge_tensors);
+
+  static c10::intrusive_ptr<GpuGraphCache> Create(
+      const int64_t num_edges, const int64_t threshold,
+      torch::ScalarType indptr_dtype, std::vector<torch::ScalarType> dtypes,
+      bool has_original_edge_ids);
+
+ private:
+  void* map_;                     // pointer to the hash table.
+  int64_t threshold_;             // A positive threshold value.
+  torch::DeviceIndex device_id_;  // Which GPU the cache resides in.
+  int64_t map_size_;              // The number of nodes inside the hash table.
+  int64_t num_nodes_;             // The number of cached nodes in the cache.
+  int64_t num_edges_;             // The number of cached edges in the cache.
+  torch::Tensor indptr_;          // The cached graph structure indptr tensor.
+  torch::optional<torch::Tensor>
+      offset_;  // The original graph's sliced_indptr tensor.
+  std::vector<torch::Tensor> cached_edge_tensors_;  // The cached graph
+                                                    // structure edge tensors.
+  std::mutex mtx_;  // Protects the data structure and makes it threadsafe.
+};
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_GPU_CACHE_H_

graphbolt/src/cuda/extension/unique_and_compact.h

@@ -0,0 +1,42 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/unique_and_compact.h
+ * @brief Unique and compact operator utilities on CUDA using hash table.
+ */
+
+#ifndef GRAPHBOLT_CUDA_UNIQUE_AND_COMPACT_H_
+#define GRAPHBOLT_CUDA_UNIQUE_AND_COMPACT_H_
+
+#include <torch/script.h>
+
+#include <vector>
+
+namespace graphbolt {
+namespace ops {
+
+std::vector<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+UniqueAndCompactBatchedHashMapBased(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor>& unique_dst_ids, const int64_t rank,
+    const int64_t world_size);
+
+}  // namespace ops
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CUDA_UNIQUE_AND_COMPACT_H_

graphbolt/src/cuda/extension/unique_and_compact_map.cu

@@ -0,0 +1,335 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/unique_and_compact_map.cu
+ * @brief Unique and compact operator implementation on CUDA using hash table.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/iterator/reverse_iterator.h>
+#include <thrust/iterator/tabulate_output_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/transform_output_iterator.h>
+
+#include <cub/cub.cuh>
+#include <cuco/static_map.cuh>
+#include <cuda/functional>
+#include <cuda/std/atomic>
+#include <cuda/std/utility>
+#include <cuda/stream_ref>
+#include <limits>
+#include <numeric>
+
+#include "../common.h"
+#include "../cooperative_minibatching_utils.cuh"
+#include "../cooperative_minibatching_utils.h"
+#include "../utils.h"
+#include "./unique_and_compact.h"
+
+namespace graphbolt {
+namespace ops {
+
+// Support graphs with up to 2^kNodeIdBits nodes.
+constexpr int kNodeIdBits = 40;
+
+template <typename index_t, typename map_t>
+__global__ void _InsertAndSetMinBatched(
+    const int64_t num_edges, const int32_t* const indexes, index_t** pointers,
+    const int64_t* const offsets, map_t map) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  while (i < num_edges) {
+    const auto tensor_index = indexes[i];
+    const auto tensor_offset = i - offsets[tensor_index];
+    const int64_t node_id = pointers[tensor_index][tensor_offset];
+    const int64_t batch_index = tensor_index / 2;
+    const int64_t key = node_id | (batch_index << kNodeIdBits);
+
+    auto [slot, is_new_key] = map.insert_and_find(cuco::pair{key, i});
+
+    if (!is_new_key) {
+      auto ref = ::cuda::atomic_ref<int64_t, ::cuda::thread_scope_device>{
+          slot->second};
+      ref.fetch_min(i, ::cuda::memory_order_relaxed);
+    }
+
+    i += stride;
+  }
+}
+
+template <typename index_t, typename map_t>
+__global__ void _MapIdsBatched(
+    const int num_batches, const int64_t num_edges,
+    const int32_t* const indexes, index_t** pointers,
+    const int64_t* const offsets, const int64_t* const unique_ids_offsets,
+    const index_t* const index, map_t map, index_t* mapped_ids) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  while (i < num_edges) {
+    const auto tensor_index = indexes[i];
+    int64_t batch_index;
+
+    if (tensor_index >= 2 * num_batches) {
+      batch_index = tensor_index - 2 * num_batches;
+    } else if (tensor_index & 1) {
+      batch_index = tensor_index / 2;
+    } else {
+      batch_index = -1;
+    }
+
+    // Only map src or dst ids.
+    if (batch_index >= 0) {
+      const auto tensor_offset = i - offsets[tensor_index];
+      const int64_t node_id = pointers[tensor_index][tensor_offset];
+      const int64_t key = node_id | (batch_index << kNodeIdBits);
+
+      auto slot = map.find(key);
+      auto new_id = slot->second;
+      if (index) {
+        new_id = index[new_id];
+      } else {
+        new_id -= unique_ids_offsets[batch_index];
+      }
+      mapped_ids[i] = new_id;
+    }
+
+    i += stride;
+  }
+}
+
+std::vector<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+UniqueAndCompactBatchedHashMapBased(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor>& unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  TORCH_CHECK(
+      rank < world_size, "rank needs to be smaller than the world_size.");
+  TORCH_CHECK(world_size <= std::numeric_limits<uint32_t>::max());
+  auto allocator = cuda::GetAllocator();
+  auto stream = cuda::GetCurrentStream();
+  auto scalar_type = src_ids.at(0).scalar_type();
+  constexpr int BLOCK_SIZE = 512;
+  const auto num_batches = src_ids.size();
+  static_assert(
+      sizeof(std::ptrdiff_t) == sizeof(int64_t),
+      "Need to be compiled on a 64-bit system.");
+  constexpr int batch_id_bits = sizeof(int64_t) * 8 - 1 - kNodeIdBits;
+  TORCH_CHECK(
+      num_batches <= (1 << batch_id_bits),
+      "UniqueAndCompactBatched supports a batch size of up to ",
+      1 << batch_id_bits);
+  return AT_DISPATCH_INDEX_TYPES(
+      scalar_type, "unique_and_compact", ([&] {
+        // For 2 batches of inputs, stores the input tensor pointers in the
+        // unique_dst, src, unique_dst, src, dst, dst order. Since there are
+        // 3 * num_batches input tensors, we need the first 3 * num_batches to
+        // store the input tensor pointers. Then, we store offsets in the rest
+        // of the 3 * num_batches + 1 space as if they were stored contiguously.
+        auto pointers_and_offsets = torch::empty(
+            6 * num_batches + 1,
+            c10::TensorOptions().dtype(torch::kInt64).pinned_memory(true));
+        // Points to the input tensor pointers.
+        auto pointers_ptr =
+            reinterpret_cast<index_t**>(pointers_and_offsets.data_ptr());
+        // Points to the input tensor storage logical offsets.
+        auto offsets_ptr =
+            pointers_and_offsets.data_ptr<int64_t>() + 3 * num_batches;
+        for (std::size_t i = 0; i < num_batches; i++) {
+          pointers_ptr[2 * i] = unique_dst_ids.at(i).data_ptr<index_t>();
+          offsets_ptr[2 * i] = unique_dst_ids[i].size(0);
+          pointers_ptr[2 * i + 1] = src_ids.at(i).data_ptr<index_t>();
+          offsets_ptr[2 * i + 1] = src_ids[i].size(0);
+          pointers_ptr[2 * num_batches + i] = dst_ids.at(i).data_ptr<index_t>();
+          offsets_ptr[2 * num_batches + i] = dst_ids[i].size(0);
+        }
+        // Finish computing the offsets by taking a cumulative sum.
+        std::exclusive_scan(
+            offsets_ptr, offsets_ptr + 3 * num_batches + 1, offsets_ptr, 0ll);
+        // Device version of the tensors defined above. We store the information
+        // initially on the CPU, which are later copied to the device.
+        auto pointers_and_offsets_dev = torch::empty(
+            pointers_and_offsets.size(0),
+            src_ids[0].options().dtype(pointers_and_offsets.scalar_type()));
+        auto offsets_dev = pointers_and_offsets_dev.slice(0, 3 * num_batches);
+        auto pointers_dev_ptr =
+            reinterpret_cast<index_t**>(pointers_and_offsets_dev.data_ptr());
+        auto offsets_dev_ptr = offsets_dev.data_ptr<int64_t>();
+        CUDA_CALL(cudaMemcpyAsync(
+            pointers_dev_ptr, pointers_ptr,
+            sizeof(int64_t) * pointers_and_offsets.size(0),
+            cudaMemcpyHostToDevice, stream));
+        auto indexes = ExpandIndptrImpl(
+            offsets_dev, torch::kInt32, torch::nullopt,
+            offsets_ptr[3 * num_batches]);
+        cuco::static_map map{
+            offsets_ptr[2 * num_batches],
+            0.5,  // load_factor
+            cuco::empty_key{static_cast<int64_t>(-1)},
+            cuco::empty_value{static_cast<int64_t>(-1)},
+            {},
+            cuco::linear_probing<1, cuco::default_hash_function<int64_t>>{},
+            {},
+            {},
+            cuda::CUDAWorkspaceAllocator<cuco::pair<int64_t, int64_t>>{},
+            ::cuda::stream_ref{stream},
+        };
+        C10_CUDA_KERNEL_LAUNCH_CHECK();  // Check the map constructor's success.
+        const dim3 block(BLOCK_SIZE);
+        const dim3 grid(
+            (offsets_ptr[2 * num_batches] + BLOCK_SIZE - 1) / BLOCK_SIZE);
+        CUDA_KERNEL_CALL(
+            _InsertAndSetMinBatched, grid, block, 0,
+            offsets_ptr[2 * num_batches], indexes.data_ptr<int32_t>(),
+            pointers_dev_ptr, offsets_dev_ptr, map.ref(cuco::insert_and_find));
+        cub::ArgIndexInputIterator index_it(indexes.data_ptr<int32_t>());
+        auto input_it = thrust::make_transform_iterator(
+            index_it,
+            ::cuda::proclaim_return_type<
+                ::cuda::std::tuple<int64_t*, index_t, int32_t, bool>>(
+                [=, map = map.ref(cuco::find)] __device__(auto it)
+                    -> ::cuda::std::tuple<int64_t*, index_t, int32_t, bool> {
+                  const auto i = it.key;
+                  const auto tensor_index = it.value;
+                  const auto tensor_offset = i - offsets_dev_ptr[tensor_index];
+                  const int64_t node_id =
+                      pointers_dev_ptr[tensor_index][tensor_offset];
+                  const auto batch_index = tensor_index / 2;
+                  const int64_t key =
+                      node_id |
+                      (static_cast<int64_t>(batch_index) << kNodeIdBits);
+                  const auto batch_offset = offsets_dev_ptr[batch_index * 2];
+
+                  auto slot = map.find(key);
+                  const auto valid = slot->second == i;
+
+                  return {&slot->second, node_id, batch_index, valid};
+                }));
+        torch::optional<torch::Tensor> part_ids;
+        if (world_size > 1) {
+          part_ids = torch::empty(
+              offsets_ptr[2 * num_batches],
+              src_ids[0].options().dtype(cuda::kPartDType));
+        }
+        auto unique_ids =
+            torch::empty(offsets_ptr[2 * num_batches], src_ids[0].options());
+        auto unique_ids_offsets_dev = torch::full(
+            num_batches + 1, std::numeric_limits<int64_t>::max(),
+            src_ids[0].options().dtype(torch::kInt64));
+        auto unique_ids_offsets_dev_ptr =
+            unique_ids_offsets_dev.data_ptr<int64_t>();
+        auto output_it = thrust::make_tabulate_output_iterator(
+            ::cuda::proclaim_return_type<void>(
+                [=, unique_ids_ptr = unique_ids.data_ptr<index_t>(),
+                 part_ids_ptr =
+                     part_ids ? part_ids->data_ptr<cuda::part_t>() : nullptr,
+                 rank = static_cast<uint32_t>(rank),
+                 world_size = static_cast<uint32_t>(
+                     world_size)] __device__(const int64_t i, const auto& t) {
+                  *::cuda::std::get<0>(t) = i;
+                  const auto node_id = ::cuda::std::get<1>(t);
+                  unique_ids_ptr[i] = node_id;
+                  if (part_ids_ptr) {
+                    part_ids_ptr[i] =
+                        cuda::rank_assignment(node_id, rank, world_size);
+                  }
+                  const auto batch_index = ::cuda::std::get<2>(t);
+                  auto ref =
+                      ::cuda::atomic_ref<int64_t, ::cuda::thread_scope_device>{
+                          unique_ids_offsets_dev_ptr[batch_index]};
+                  ref.fetch_min(i, ::cuda::memory_order_relaxed);
+                }));
+        CUB_CALL(
+            DeviceSelect::If, input_it, output_it,
+            unique_ids_offsets_dev_ptr + num_batches,
+            offsets_ptr[2 * num_batches],
+            ::cuda::proclaim_return_type<bool>([] __device__(const auto& t) {
+              return ::cuda::std::get<3>(t);
+            }));
+        auto unique_ids_offsets = torch::empty(
+            num_batches + 1,
+            c10::TensorOptions().dtype(torch::kInt64).pinned_memory(true));
+        {
+          auto unique_ids_offsets_dev2 =
+              torch::empty_like(unique_ids_offsets_dev);
+          CUB_CALL(
+              DeviceScan::InclusiveScan,
+              thrust::make_reverse_iterator(
+                  num_batches + 1 + unique_ids_offsets_dev_ptr),
+              thrust::make_reverse_iterator(
+                  num_batches + 1 +
+                  thrust::make_transform_output_iterator(
+                      thrust::make_zip_iterator(
+                          unique_ids_offsets_dev2.data_ptr<int64_t>(),
+                          unique_ids_offsets.data_ptr<int64_t>()),
+                      ::cuda::proclaim_return_type<
+                          thrust::tuple<int64_t, int64_t>>(
+                          [=] __device__(const auto x) {
+                            return thrust::make_tuple(x, x);
+                          }))),
+              cub::Min{}, num_batches + 1);
+          unique_ids_offsets_dev = unique_ids_offsets_dev2;
+          unique_ids_offsets_dev_ptr =
+              unique_ids_offsets_dev.data_ptr<int64_t>();
+        }
+        at::cuda::CUDAEvent unique_ids_offsets_event;
+        unique_ids_offsets_event.record();
+        torch::optional<torch::Tensor> index;
+        if (part_ids) {
+          unique_ids_offsets_event.synchronize();
+          const auto num_unique =
+              unique_ids_offsets.data_ptr<int64_t>()[num_batches];
+          unique_ids = unique_ids.slice(0, 0, num_unique);
+          part_ids = part_ids->slice(0, 0, num_unique);
+          std::tie(
+              unique_ids, index, unique_ids_offsets, unique_ids_offsets_event) =
+              cuda::RankSortImpl(
+                  unique_ids, *part_ids, unique_ids_offsets_dev, world_size);
+        }
+        auto mapped_ids =
+            torch::empty(offsets_ptr[3 * num_batches], unique_ids.options());
+        CUDA_KERNEL_CALL(
+            _MapIdsBatched, grid, block, 0, num_batches,
+            offsets_ptr[3 * num_batches], indexes.data_ptr<int32_t>(),
+            pointers_dev_ptr, offsets_dev_ptr, unique_ids_offsets_dev_ptr,
+            index ? index->data_ptr<index_t>() : nullptr, map.ref(cuco::find),
+            mapped_ids.data_ptr<index_t>());
+        std::vector<std::tuple<
+            torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+            results;
+        unique_ids_offsets_event.synchronize();
+        auto unique_ids_offsets_ptr = unique_ids_offsets.data_ptr<int64_t>();
+        for (int64_t i = 0; i < num_batches; i++) {
+          results.emplace_back(
+              unique_ids.slice(
+                  0, unique_ids_offsets_ptr[i * world_size],
+                  unique_ids_offsets_ptr[(i + 1) * world_size]),
+              mapped_ids.slice(
+                  0, offsets_ptr[2 * i + 1], offsets_ptr[2 * i + 2]),
+              mapped_ids.slice(
+                  0, offsets_ptr[2 * num_batches + i],
+                  offsets_ptr[2 * num_batches + i + 1]),
+              unique_ids_offsets.slice(
+                  0, i * world_size, (i + 1) * world_size + 1));
+        }
+        return results;
+      }));
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/gather.cu

@@ -0,0 +1,50 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/gather.cu
+ * @brief Gather operators implementation on CUDA.
+ */
+#include <thrust/gather.h>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+torch::Tensor Gather(
+    torch::Tensor input, torch::Tensor index,
+    torch::optional<torch::ScalarType> dtype) {
+  if (!dtype.has_value()) dtype = input.scalar_type();
+  auto output = torch::empty(index.sizes(), index.options().dtype(*dtype));
+  AT_DISPATCH_INDEX_TYPES(
+      index.scalar_type(), "GatherIndexType", ([&] {
+        AT_DISPATCH_INTEGRAL_TYPES(
+            input.scalar_type(), "GatherInputType", ([&] {
+              using input_t = scalar_t;
+              AT_DISPATCH_INTEGRAL_TYPES(*dtype, "GatherOutputType", ([&] {
+                using output_t = scalar_t;
+                THRUST_CALL(
+                    gather, index.data_ptr<index_t>(),
+                    index.data_ptr<index_t>() + index.size(0),
+                    input.data_ptr<input_t>(), output.data_ptr<output_t>());
+              }));
+            }));
+      }));
+  return output;
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/index_select_csc_impl.cu

@@ -0,0 +1,335 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/index_select_csc_impl.cu
+ * @brief Index select csc operator implementation on CUDA.
+ */
+#include <c10/core/ScalarType.h>
+#include <graphbolt/cuda_ops.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/zip_iterator.h>
+
+#include <cstdint>
+#include <cub/cub.cuh>
+#include <numeric>
+
+#include "./common.h"
+#include "./max_uva_threads.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+constexpr int BLOCK_SIZE = CUDA_MAX_NUM_THREADS;
+
+// Given the in_degree array and a permutation, returns in_degree of the output
+// and the permuted and modified in_degree of the input. The modified in_degree
+// is modified so that there is slack to be able to align as needed.
+template <typename indptr_t, typename indices_t>
+struct AlignmentFunc {
+  static_assert(GPU_CACHE_LINE_SIZE % sizeof(indices_t) == 0);
+  const indptr_t* in_degree;
+  const int64_t* perm;
+  int64_t num_nodes;
+  __host__ __device__ auto operator()(int64_t row) {
+    constexpr int num_elements = GPU_CACHE_LINE_SIZE / sizeof(indices_t);
+    return thrust::make_tuple(
+        in_degree[row],
+        // A single cache line has num_elements items, we add num_elements - 1
+        // to ensure there is enough slack to move forward or backward by
+        // num_elements - 1 items if the performed access is not aligned.
+        static_cast<indptr_t>(
+            in_degree[perm ? perm[row % num_nodes] : row] + num_elements - 1));
+  }
+};
+
+template <typename indptr_t, typename indices_t, typename coo_rows_t>
+__global__ void _CopyIndicesAlignedKernel(
+    const indptr_t edge_count, const indptr_t* const indptr,
+    const indptr_t* const output_indptr,
+    const indptr_t* const output_indptr_aligned, const indices_t* const indices,
+    const coo_rows_t* const coo_aligned_rows, indices_t* const output_indices,
+    const int64_t* const perm) {
+  indptr_t idx = static_cast<indptr_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  const int stride_x = gridDim.x * blockDim.x;
+
+  while (idx < edge_count) {
+    const auto permuted_row_pos = coo_aligned_rows[idx];
+    const auto row_pos = perm ? perm[permuted_row_pos] : permuted_row_pos;
+    const auto out_row = output_indptr[row_pos];
+    const auto d = output_indptr[row_pos + 1] - out_row;
+    const int offset = (reinterpret_cast<std::uintptr_t>(
+                            indices + indptr[row_pos] -
+                            output_indptr_aligned[permuted_row_pos]) %
+                        GPU_CACHE_LINE_SIZE) /
+                       sizeof(indices_t);
+    const auto rofs = idx - output_indptr_aligned[permuted_row_pos] - offset;
+    if (rofs >= 0 && rofs < d) {
+      const auto in_idx = indptr[row_pos] + rofs;
+      assert(
+          reinterpret_cast<std::uintptr_t>(indices + in_idx - idx) %
+              GPU_CACHE_LINE_SIZE ==
+          0);
+      const auto u = indices[in_idx];
+      output_indices[out_row + rofs] = u;
+    }
+    idx += stride_x;
+  }
+}
+
+struct PairSum {
+  template <typename indptr_t>
+  __host__ __device__ auto operator()(
+      const thrust::tuple<indptr_t, indptr_t> a,
+      const thrust::tuple<indptr_t, indptr_t> b) {
+    return thrust::make_tuple(
+        thrust::get<0>(a) + thrust::get<0>(b),
+        thrust::get<1>(a) + thrust::get<1>(b));
+  };
+};
+
+template <typename indptr_t, typename indices_t>
+std::tuple<torch::Tensor, torch::Tensor> UVAIndexSelectCSCCopyIndices(
+    torch::Tensor indices, const int64_t num_nodes,
+    const indptr_t* const in_degree, const indptr_t* const sliced_indptr,
+    const int64_t* const perm, torch::TensorOptions options,
+    torch::ScalarType indptr_scalar_type,
+    torch::optional<int64_t> output_size) {
+  auto allocator = cuda::GetAllocator();
+  thrust::counting_iterator<int64_t> iota(0);
+
+  // Output indptr for the slice indexed by nodes.
+  auto output_indptr =
+      torch::empty(num_nodes + 1, options.dtype(indptr_scalar_type));
+
+  auto output_indptr_aligned =
+      torch::empty(num_nodes + 1, options.dtype(indptr_scalar_type));
+  auto output_indptr_aligned_ptr = output_indptr_aligned.data_ptr<indptr_t>();
+
+  {
+    // Returns the actual and modified_indegree as a pair, the
+    // latter overestimates the actual indegree for alignment
+    // purposes.
+    auto modified_in_degree = thrust::make_transform_iterator(
+        iota, AlignmentFunc<indptr_t, indices_t>{in_degree, perm, num_nodes});
+    auto output_indptr_pair = thrust::make_zip_iterator(
+        output_indptr.data_ptr<indptr_t>(), output_indptr_aligned_ptr);
+    thrust::tuple<indptr_t, indptr_t> zero_value{};
+    // Compute the prefix sum over actual and modified indegrees.
+    CUB_CALL(
+        DeviceScan::ExclusiveScan, modified_in_degree, output_indptr_pair,
+        PairSum{}, zero_value, num_nodes + 1);
+  }
+
+  // Copy the actual total number of edges.
+  if (!output_size.has_value()) {
+    auto edge_count =
+        cuda::CopyScalar{output_indptr.data_ptr<indptr_t>() + num_nodes};
+    output_size = static_cast<indptr_t>(edge_count);
+  }
+  // Copy the modified number of edges.
+  auto edge_count_aligned_ =
+      cuda::CopyScalar{output_indptr_aligned_ptr + num_nodes};
+  const int64_t edge_count_aligned = static_cast<indptr_t>(edge_count_aligned_);
+
+  // Allocate output array with actual number of edges.
+  torch::Tensor output_indices =
+      torch::empty(output_size.value(), options.dtype(indices.scalar_type()));
+  const dim3 block(BLOCK_SIZE);
+  const dim3 grid(
+      (std::min(edge_count_aligned, cuda::max_uva_threads.value_or(1 << 20)) +
+       BLOCK_SIZE - 1) /
+      BLOCK_SIZE);
+
+  // Find the smallest integer type to store the coo_aligned_rows tensor.
+  const int num_bits = cuda::NumberOfBits(num_nodes);
+  std::array<int, 4> type_bits = {8, 15, 31, 63};
+  const auto type_index =
+      std::lower_bound(type_bits.begin(), type_bits.end(), num_bits) -
+      type_bits.begin();
+  std::array<torch::ScalarType, 5> types = {
+      torch::kByte, torch::kInt16, torch::kInt32, torch::kLong, torch::kLong};
+  auto coo_dtype = types[type_index];
+
+  auto coo_aligned_rows = ExpandIndptrImpl(
+      output_indptr_aligned, coo_dtype, torch::nullopt, edge_count_aligned);
+
+  AT_DISPATCH_INTEGRAL_TYPES(
+      coo_dtype, "UVAIndexSelectCSCCopyIndicesCOO", ([&] {
+        using coo_rows_t = scalar_t;
+        // Perform the actual copying, of the indices array into
+        // output_indices in an aligned manner.
+        CUDA_KERNEL_CALL(
+            _CopyIndicesAlignedKernel, grid, block, 0,
+            static_cast<indptr_t>(edge_count_aligned_), sliced_indptr,
+            output_indptr.data_ptr<indptr_t>(), output_indptr_aligned_ptr,
+            reinterpret_cast<indices_t*>(indices.data_ptr()),
+            coo_aligned_rows.data_ptr<coo_rows_t>(),
+            reinterpret_cast<indices_t*>(output_indices.data_ptr()), perm);
+      }));
+  return {output_indptr, output_indices};
+}
+
+std::tuple<torch::Tensor, torch::Tensor> UVAIndexSelectCSCImpl(
+    torch::Tensor in_degree, torch::Tensor sliced_indptr, torch::Tensor indices,
+    torch::Tensor nodes, int num_bits, torch::optional<int64_t> output_size) {
+  // Sorting nodes so that accesses over PCI-e are more regular.
+  const auto sorted_idx = Sort(nodes, num_bits).second;
+  const int64_t num_nodes = nodes.size(0);
+
+  return AT_DISPATCH_INTEGRAL_TYPES(
+      sliced_indptr.scalar_type(), "UVAIndexSelectCSCIndptr", ([&] {
+        using indptr_t = scalar_t;
+        return GRAPHBOLT_DISPATCH_ELEMENT_SIZES(
+            indices.element_size(), "UVAIndexSelectCSCCopyIndices", ([&] {
+              return UVAIndexSelectCSCCopyIndices<indptr_t, element_size_t>(
+                  indices, num_nodes, in_degree.data_ptr<indptr_t>(),
+                  sliced_indptr.data_ptr<indptr_t>(),
+                  sorted_idx.data_ptr<int64_t>(), nodes.options(),
+                  sliced_indptr.scalar_type(), output_size);
+            }));
+      }));
+}
+
+template <typename indptr_t, typename indices_t>
+struct IteratorFunc {
+  indptr_t* indptr;
+  indices_t* indices;
+  __host__ __device__ auto operator()(int64_t i) { return indices + indptr[i]; }
+};
+
+template <typename indptr_t, typename indices_t>
+struct ConvertToBytes {
+  const indptr_t* in_degree;
+  __host__ __device__ indptr_t operator()(int64_t i) {
+    return in_degree[i] * sizeof(indices_t);
+  }
+};
+
+template <typename indptr_t, typename indices_t>
+void IndexSelectCSCCopyIndices(
+    const int64_t num_nodes, indices_t* const indices,
+    indptr_t* const sliced_indptr, const indptr_t* const in_degree,
+    indptr_t* const output_indptr, indices_t* const output_indices) {
+  thrust::counting_iterator<int64_t> iota(0);
+
+  auto input_buffer_it = thrust::make_transform_iterator(
+      iota, IteratorFunc<indptr_t, indices_t>{sliced_indptr, indices});
+  auto output_buffer_it = thrust::make_transform_iterator(
+      iota, IteratorFunc<indptr_t, indices_t>{output_indptr, output_indices});
+  auto buffer_sizes = thrust::make_transform_iterator(
+      iota, ConvertToBytes<indptr_t, indices_t>{in_degree});
+  constexpr int64_t max_copy_at_once = std::numeric_limits<int32_t>::max();
+
+  // Performs the copy from indices into output_indices.
+  for (int64_t i = 0; i < num_nodes; i += max_copy_at_once) {
+    CUB_CALL(
+        DeviceMemcpy::Batched, input_buffer_it + i, output_buffer_it + i,
+        buffer_sizes + i, std::min(num_nodes - i, max_copy_at_once));
+  }
+}
+
+std::tuple<torch::Tensor, torch::Tensor> DeviceIndexSelectCSCImpl(
+    torch::Tensor in_degree, torch::Tensor sliced_indptr, torch::Tensor indices,
+    torch::TensorOptions options, torch::optional<int64_t> output_size) {
+  const int64_t num_nodes = sliced_indptr.size(0);
+  return AT_DISPATCH_INTEGRAL_TYPES(
+      sliced_indptr.scalar_type(), "IndexSelectCSCIndptr", ([&] {
+        using indptr_t = scalar_t;
+        auto in_degree_ptr = in_degree.data_ptr<indptr_t>();
+        auto sliced_indptr_ptr = sliced_indptr.data_ptr<indptr_t>();
+        // Output indptr for the slice indexed by nodes.
+        torch::Tensor output_indptr = torch::empty(
+            num_nodes + 1, options.dtype(sliced_indptr.scalar_type()));
+
+        // Compute the output indptr, output_indptr.
+        CUB_CALL(
+            DeviceScan::ExclusiveSum, in_degree_ptr,
+            output_indptr.data_ptr<indptr_t>(), num_nodes + 1);
+
+        // Number of edges being copied.
+        if (!output_size.has_value()) {
+          auto edge_count =
+              cuda::CopyScalar{output_indptr.data_ptr<indptr_t>() + num_nodes};
+          output_size = static_cast<indptr_t>(edge_count);
+        }
+        // Allocate output array of size number of copied edges.
+        torch::Tensor output_indices = torch::empty(
+            output_size.value(), options.dtype(indices.scalar_type()));
+        GRAPHBOLT_DISPATCH_ELEMENT_SIZES(
+            indices.element_size(), "IndexSelectCSCCopyIndices", ([&] {
+              using indices_t = element_size_t;
+              IndexSelectCSCCopyIndices<indptr_t, indices_t>(
+                  num_nodes, reinterpret_cast<indices_t*>(indices.data_ptr()),
+                  sliced_indptr_ptr, in_degree_ptr,
+                  output_indptr.data_ptr<indptr_t>(),
+                  reinterpret_cast<indices_t*>(output_indices.data_ptr()));
+            }));
+        return std::make_tuple(output_indptr, output_indices);
+      }));
+}
+
+std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSCImpl(
+    torch::Tensor in_degree, torch::Tensor sliced_indptr, torch::Tensor indices,
+    torch::Tensor nodes, int64_t nodes_max,
+    torch::optional<int64_t> output_size) {
+  if (indices.is_pinned()) {
+    int num_bits = cuda::NumberOfBits(nodes_max + 1);
+    return UVAIndexSelectCSCImpl(
+        in_degree, sliced_indptr, indices, nodes, num_bits, output_size);
+  } else {
+    return DeviceIndexSelectCSCImpl(
+        in_degree, sliced_indptr, indices, nodes.options(), output_size);
+  }
+}
+
+std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSCImpl(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes,
+    torch::optional<int64_t> output_size) {
+  auto [in_degree, sliced_indptr] = SliceCSCIndptr(indptr, nodes);
+  return IndexSelectCSCImpl(
+      in_degree, sliced_indptr, indices, nodes, indptr.size(0) - 2,
+      output_size);
+}
+
+std::tuple<torch::Tensor, std::vector<torch::Tensor>> IndexSelectCSCBatchedImpl(
+    torch::Tensor indptr, std::vector<torch::Tensor> indices_list,
+    torch::Tensor nodes, bool with_edge_ids,
+    torch::optional<int64_t> output_size) {
+  auto [in_degree, sliced_indptr] = SliceCSCIndptr(indptr, nodes);
+  std::vector<torch::Tensor> results;
+  results.reserve(indices_list.size());
+  torch::Tensor output_indptr;
+  for (auto& indices : indices_list) {
+    torch::Tensor output_indices;
+    std::tie(output_indptr, output_indices) = IndexSelectCSCImpl(
+        in_degree, sliced_indptr, indices, nodes, indptr.size(0) - 2,
+        output_size);
+    if (!output_size.has_value()) output_size = output_indices.size(0);
+    TORCH_CHECK(*output_size == output_indices.size(0));
+    results.push_back(output_indices);
+  }
+  if (with_edge_ids) {
+    results.push_back(IndptrEdgeIdsImpl(
+        output_indptr, sliced_indptr.scalar_type(), sliced_indptr,
+        output_size));
+  }
+  return {output_indptr, results};
+}
+
+}  //  namespace ops
+}  //  namespace graphbolt

graphbolt/src/cuda/index_select_impl.cu

@@ -0,0 +1,208 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/index_select_impl.cu
+ * @brief Index select operator implementation on CUDA.
+ */
+#include <c10/core/ScalarType.h>
+#include <graphbolt/cuda_ops.h>
+
+#include <numeric>
+
+#include "./common.h"
+#include "./max_uva_threads.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+/** @brief Index select operator implementation for feature size 1. */
+template <typename DType, typename IdType>
+__global__ void IndexSelectSingleKernel(
+    const DType* input, const int64_t input_len, const IdType* index,
+    const int64_t output_len, DType* output,
+    const int64_t* permutation = nullptr) {
+  int64_t out_row_index = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  while (out_row_index < output_len) {
+    assert(index[out_row_index] >= 0 && index[out_row_index] < input_len);
+    const auto out_row =
+        permutation ? permutation[out_row_index] : out_row_index;
+    output[out_row] = input[index[out_row_index]];
+    out_row_index += stride;
+  }
+}
+
+/**
+ * @brief Index select operator implementation for feature size > 1.
+ */
+template <typename DType, typename IdType>
+__global__ void IndexSelectMultiKernel(
+    const DType* const input, const int64_t input_len,
+    const int64_t feature_size, const IdType* const index,
+    const int64_t output_len, DType* const output,
+    const int64_t* permutation = nullptr) {
+  int64_t out_row_index = blockIdx.x * blockDim.y + threadIdx.y;
+
+  const int64_t stride = blockDim.y * gridDim.x;
+
+  while (out_row_index < output_len) {
+    int64_t column = threadIdx.x;
+    const int64_t in_row = index[out_row_index];
+    assert(in_row >= 0 && in_row < input_len);
+    const auto out_row =
+        permutation ? permutation[out_row_index] : out_row_index;
+    while (column < feature_size) {
+      output[out_row * feature_size + column] =
+          input[in_row * feature_size + column];
+      column += blockDim.x;
+    }
+    out_row_index += stride;
+  }
+}
+
+/**
+ * @brief Index select operator implementation for feature size > 1.
+ *
+ * @note This is a cross-device access version of IndexSelectMultiKernel. Since
+ * the memory access over PCIe is more sensitive to the data access aligment
+ * (cacheline), we need a separate version here.
+ */
+template <typename DType, typename IdType>
+__global__ void IndexSelectMultiKernelAligned(
+    const DType* const input, const int64_t input_len,
+    const int64_t feature_size, const IdType* const index,
+    const int64_t output_len, DType* const output,
+    const int64_t* permutation = nullptr) {
+  int64_t out_row_index = blockIdx.x * blockDim.y + threadIdx.y;
+
+  const int64_t stride = blockDim.y * gridDim.x;
+
+  while (out_row_index < output_len) {
+    int64_t col = threadIdx.x;
+    const int64_t in_row = index[out_row_index];
+    assert(in_row >= 0 && in_row < input_len);
+    const int64_t idx_offset =
+        ((uint64_t)(&input[in_row * feature_size]) % GPU_CACHE_LINE_SIZE) /
+        sizeof(DType);
+    col = col - idx_offset;
+    const auto out_row =
+        permutation ? permutation[out_row_index] : out_row_index;
+    while (col < feature_size) {
+      if (col >= 0)
+        output[out_row * feature_size + col] =
+            input[in_row * feature_size + col];
+      col += blockDim.x;
+    }
+    out_row_index += stride;
+  }
+}
+
+template <typename DType, typename IdType>
+torch::Tensor UVAIndexSelectImpl_(torch::Tensor input, torch::Tensor index) {
+  const int64_t input_len = input.size(0);
+  const int64_t return_len = index.size(0);
+  const int64_t original_feature_size = std::accumulate(
+      input.sizes().begin() + 1, input.sizes().end(), 1ll, std::multiplies<>());
+  const auto aligned_feature_size =
+      input.element_size() * original_feature_size / sizeof(DType);
+  torch::Tensor ret = torch::empty(
+      {return_len, original_feature_size}, torch::TensorOptions()
+                                               .dtype(input.dtype())
+                                               .device(c10::DeviceType::CUDA));
+  DType* input_ptr = reinterpret_cast<DType*>(input.data_ptr());
+  DType* ret_ptr = reinterpret_cast<DType*>(ret.data_ptr());
+
+  // Sort the index to improve the memory access pattern.
+  torch::Tensor sorted_index, permutation;
+  std::tie(sorted_index, permutation) =
+      Sort(index, cuda::NumberOfBits(input_len));
+  const IdType* index_sorted_ptr = sorted_index.data_ptr<IdType>();
+  const int64_t* permutation_ptr = permutation.data_ptr<int64_t>();
+
+  if (aligned_feature_size == 1) {
+    // Use a single thread to process each output row to avoid wasting threads.
+    const int num_threads = cuda::FindNumThreads(return_len);
+    const int num_blocks =
+        (std::min(return_len, cuda::max_uva_threads.value_or(1 << 20)) +
+         num_threads - 1) /
+        num_threads;
+    CUDA_KERNEL_CALL(
+        IndexSelectSingleKernel, num_blocks, num_threads, 0, input_ptr,
+        input_len, index_sorted_ptr, return_len, ret_ptr, permutation_ptr);
+  } else {
+    constexpr int BLOCK_SIZE = CUDA_MAX_NUM_THREADS;
+    dim3 block(BLOCK_SIZE, 1);
+    while (static_cast<int64_t>(block.x) >= 2 * aligned_feature_size) {
+      block.x >>= 1;
+      block.y <<= 1;
+    }
+    const dim3 grid(std::min(
+        (return_len + block.y - 1) / block.y,
+        cuda::max_uva_threads.value_or(1 << 20) / BLOCK_SIZE));
+    if (aligned_feature_size * sizeof(DType) <= GPU_CACHE_LINE_SIZE) {
+      // When feature size is smaller than GPU cache line size, use unaligned
+      // version for less SM usage, which is more resource efficient.
+      CUDA_KERNEL_CALL(
+          IndexSelectMultiKernel, grid, block, 0, input_ptr, input_len,
+          aligned_feature_size, index_sorted_ptr, return_len, ret_ptr,
+          permutation_ptr);
+    } else {
+      // Use aligned version to improve the memory access pattern.
+      CUDA_KERNEL_CALL(
+          IndexSelectMultiKernelAligned, grid, block, 0, input_ptr, input_len,
+          aligned_feature_size, index_sorted_ptr, return_len, ret_ptr,
+          permutation_ptr);
+    }
+  }
+
+  auto return_shape = std::vector<int64_t>({return_len});
+  return_shape.insert(
+      return_shape.end(), input.sizes().begin() + 1, input.sizes().end());
+  ret = ret.reshape(return_shape);
+  return ret;
+}
+
+/**
+ * @brief UVA index select operator implementation on CUDA.
+ *
+ * All basic torch types are supported for input.
+ * The supporting index types are: int, int64_t.
+ */
+torch::Tensor UVAIndexSelectImpl(torch::Tensor input, torch::Tensor index) {
+  return AT_DISPATCH_INDEX_TYPES(
+      index.scalar_type(), "UVAIndexSelectImpl", ([&] {
+        const auto ptr = (size_t)input.data_ptr();
+        const int64_t feature_size = std::accumulate(
+            input.sizes().begin() + 1, input.sizes().end(), 1ll,
+            std::multiplies<>());
+        // We perform the copy with datatype of size powers of 2, and the
+        // maximum data type we use has 16 bytes. We check the alignment of the
+        // pointer and the feature dimensionality to determine the largest
+        // type to use for the copy to minimize the number of CUDA threads used.
+        // Alignment denotes the maximum suitable alignment and datatype size
+        // for the copies.
+        const int aligned_access_size =
+            std::gcd(16, std::gcd(ptr, input.element_size() * feature_size));
+        return GRAPHBOLT_DISPATCH_ELEMENT_SIZES(
+            aligned_access_size, "UVAIndexSelectImplElementSize", ([&] {
+              return UVAIndexSelectImpl_<element_size_t, index_t>(input, index);
+            }));
+      }));
+}
+
+}  //  namespace ops
+}  //  namespace graphbolt

graphbolt/src/cuda/insubgraph.cu

@@ -0,0 +1,51 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/insubgraph.cu
+ * @brief InSubgraph operator implementation on CUDA.
+ */
+
+#include <graphbolt/cuda_ops.h>
+#include <graphbolt/cuda_sampling_ops.h>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+c10::intrusive_ptr<sampling::FusedSampledSubgraph> InSubgraph(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes,
+    torch::optional<torch::Tensor> type_per_edge) {
+  auto [in_degree, sliced_indptr] = SliceCSCIndptr(indptr, nodes);
+  auto [output_indptr, output_indices] = IndexSelectCSCImpl(
+      in_degree, sliced_indptr, indices, nodes, indptr.size(0) - 2);
+  const int64_t num_edges = output_indices.size(0);
+  torch::optional<torch::Tensor> output_type_per_edge;
+  if (type_per_edge) {
+    output_type_per_edge = std::get<1>(IndexSelectCSCImpl(
+        in_degree, sliced_indptr, type_per_edge.value(), nodes,
+        indptr.size(0) - 2, num_edges));
+  }
+  auto edge_ids = IndptrEdgeIdsImpl(
+      output_indptr, sliced_indptr.scalar_type(), sliced_indptr, num_edges);
+
+  return c10::make_intrusive<sampling::FusedSampledSubgraph>(
+      output_indptr, output_indices, edge_ids, nodes, torch::nullopt,
+      output_type_per_edge);
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/isin.cu

@@ -0,0 +1,68 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/isin.cu
+ * @brief IsIn operator implementation on CUDA.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/binary_search.h>
+
+#include <cub/cub.cuh>
+
+#include "./common.h"
+
+namespace graphbolt {
+namespace ops {
+
+torch::Tensor IsIn(torch::Tensor elements, torch::Tensor test_elements) {
+  auto sorted_test_elements = Sort<false>(test_elements);
+  auto result = torch::empty_like(elements, torch::kBool);
+
+  AT_DISPATCH_INTEGRAL_TYPES(
+      elements.scalar_type(), "IsInOperation", ([&] {
+        THRUST_CALL(
+            binary_search, sorted_test_elements.data_ptr<scalar_t>(),
+            sorted_test_elements.data_ptr<scalar_t>() +
+                sorted_test_elements.size(0),
+            elements.data_ptr<scalar_t>(),
+            elements.data_ptr<scalar_t>() + elements.size(0),
+            result.data_ptr<bool>());
+      }));
+  return result;
+}
+
+torch::Tensor Nonzero(torch::Tensor mask, bool logical_not) {
+  thrust::counting_iterator<int64_t> iota(0);
+  auto result = torch::empty_like(mask, torch::kInt64);
+  auto mask_ptr = mask.data_ptr<bool>();
+  auto result_ptr = result.data_ptr<int64_t>();
+  auto allocator = cuda::GetAllocator();
+  auto num_copied = allocator.AllocateStorage<int64_t>(1);
+  if (logical_not) {
+    CUB_CALL(
+        DeviceSelect::FlaggedIf, iota, mask_ptr, result_ptr, num_copied.get(),
+        mask.numel(), thrust::logical_not<bool>{});
+  } else {
+    CUB_CALL(
+        DeviceSelect::Flagged, iota, mask_ptr, result_ptr, num_copied.get(),
+        mask.numel());
+  }
+  cuda::CopyScalar num_copied_cpu(num_copied.get());
+  return result.slice(0, 0, static_cast<int64_t>(num_copied_cpu));
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/max_uva_threads.cc

@@ -0,0 +1,28 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/max_uva_threads.cc
+ * @brief Max uva threads variable setter function.
+ */
+#include "./max_uva_threads.h"
+
+namespace graphbolt {
+namespace cuda {
+
+void set_max_uva_threads(int64_t count) { max_uva_threads = count; }
+
+}  // namespace cuda
+}  // namespace graphbolt

graphbolt/src/cuda/max_uva_threads.h

@@ -0,0 +1,37 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/max_uva_threads.h
+ * @brief Max uva threads variable declaration.
+ */
+#ifndef GRAPHBOLT_MAX_UVA_THREADS_H_
+#define GRAPHBOLT_MAX_UVA_THREADS_H_
+
+#include <cstdint>
+#include <optional>
+
+namespace graphbolt {
+namespace cuda {
+
+/** @brief Set a limit on the number of CUDA threads for UVA accesses. */
+inline std::optional<int64_t> max_uva_threads;
+
+void set_max_uva_threads(int64_t count);
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_MAX_UVA_THREADS_H_

graphbolt/src/cuda/neighbor_sampler.cu

@@ -0,0 +1,807 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/index_select_impl.cu
+ * @brief Index select operator implementation on CUDA.
+ */
+#include <c10/core/ScalarType.h>
+#include <curand_kernel.h>
+#include <graphbolt/continuous_seed.h>
+#include <graphbolt/cuda_ops.h>
+#include <graphbolt/cuda_sampling_ops.h>
+#include <thrust/copy.h>
+#include <thrust/gather.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+#include <thrust/iterator/transform_output_iterator.h>
+
+#include <algorithm>
+#include <array>
+#include <cub/cub.cuh>
+#if __CUDA_ARCH__ >= 700
+#include <cuda/atomic>
+#endif  // __CUDA_ARCH__ >= 700
+#include <limits>
+#include <numeric>
+#include <type_traits>
+
+#include "../macro.h"
+#include "../random.h"
+#include "../utils.h"
+#include "./common.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+constexpr int BLOCK_SIZE = 128;
+
+inline __device__ int64_t AtomicMax(int64_t* const address, const int64_t val) {
+  // To match the type of "::atomicCAS", ignore lint warning.
+  using Type = unsigned long long int;  // NOLINT
+
+  static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
+
+  return atomicMax(reinterpret_cast<Type*>(address), static_cast<Type>(val));
+}
+
+inline __device__ int32_t AtomicMax(int32_t* const address, const int32_t val) {
+  // To match the type of "::atomicCAS", ignore lint warning.
+  using Type = int;  // NOLINT
+
+  static_assert(sizeof(Type) == sizeof(*address), "Type width must match");
+
+  return atomicMax(reinterpret_cast<Type*>(address), static_cast<Type>(val));
+}
+
+/**
+ * @brief Performs neighbor sampling and fills the edge_ids array with
+ * original edge ids if sliced_indptr is valid. If not, then it fills the edge
+ * ids array with numbers upto the node degree.
+ */
+template <typename indptr_t, typename indices_t>
+__global__ void _ComputeRandomsNS(
+    const int64_t num_edges, const indptr_t* const sliced_indptr,
+    const indptr_t* const sub_indptr, const indptr_t* const output_indptr,
+    const indices_t* const csr_rows, const uint64_t random_seed,
+    indptr_t* edge_ids) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  curandStatePhilox4_32_10_t rng;
+  curand_init(random_seed, i, 0, &rng);
+
+  while (i < num_edges) {
+    const auto row_position = csr_rows[i];
+    const auto row_offset = i - sub_indptr[row_position];
+    const auto output_offset = output_indptr[row_position];
+    const auto fanout = output_indptr[row_position + 1] - output_offset;
+    const auto rnd =
+        row_offset < fanout ? row_offset : curand(&rng) % (row_offset + 1);
+    if (rnd < fanout) {
+      const indptr_t edge_id =
+          row_offset + (sliced_indptr ? sliced_indptr[row_position] : 0);
+#if __CUDA_ARCH__ >= 700
+      ::cuda::atomic_ref<indptr_t, ::cuda::thread_scope_device> a(
+          edge_ids[output_offset + rnd]);
+      a.fetch_max(edge_id, ::cuda::std::memory_order_relaxed);
+#else
+      AtomicMax(edge_ids + output_offset + rnd, edge_id);
+#endif  // __CUDA_ARCH__
+    }
+
+    i += stride;
+  }
+}
+
+/**
+ * @brief Fills the random_arr with random numbers and the edge_ids array with
+ * original edge ids. When random_arr is sorted along with edge_ids, the first
+ * fanout elements of each row gives us the sampled edges.
+ */
+template <
+    typename float_t, typename indptr_t, typename indices_t, typename weights_t,
+    typename edge_id_t>
+__global__ void _ComputeRandoms(
+    const int64_t num_edges, const indptr_t* const sliced_indptr,
+    const indptr_t* const sub_indptr, const indices_t* const csr_rows,
+    const weights_t* const sliced_weights, const indices_t* const indices,
+    const continuous_seed random_seed, float_t* random_arr,
+    edge_id_t* edge_ids) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+  const auto labor = indices != nullptr;
+  const float_t inf =
+      static_cast<float_t>(std::numeric_limits<float>::infinity());
+
+  while (i < num_edges) {
+    const auto row_position = csr_rows[i];
+    const auto row_offset = i - sub_indptr[row_position];
+    const auto in_idx = sliced_indptr[row_position] + row_offset;
+    const auto rnd = random_seed.uniform(labor ? indices[in_idx] : i);
+    const auto prob =
+        sliced_weights ? sliced_weights[i] : static_cast<weights_t>(1);
+    const auto exp_rnd = -__logf(rnd);
+    const float_t adjusted_rnd =
+        prob > 0 ? static_cast<float_t>(exp_rnd / prob) : inf;
+    random_arr[i] = adjusted_rnd;
+    edge_ids[i] = row_offset;
+
+    i += stride;
+  }
+}
+
+struct IsPositive {
+  template <typename probs_t>
+  __host__ __device__ auto operator()(probs_t x) {
+    return x > 0;
+  }
+};
+
+template <typename indptr_t>
+struct MinInDegreeFanout {
+  const indptr_t* in_degree;
+  const int64_t* fanouts;
+  size_t num_fanouts;
+  __host__ __device__ auto operator()(int64_t i) {
+    return static_cast<indptr_t>(
+        min(static_cast<int64_t>(in_degree[i]), fanouts[i % num_fanouts]));
+  }
+};
+
+template <typename indptr_t, typename indices_t>
+struct IteratorFunc {
+  indptr_t* indptr;
+  indices_t* indices;
+  __host__ __device__ auto operator()(int64_t i) { return indices + indptr[i]; }
+};
+
+template <typename indptr_t>
+struct AddOffset {
+  indptr_t offset;
+  template <typename edge_id_t>
+  __host__ __device__ indptr_t operator()(edge_id_t x) {
+    return x + offset;
+  }
+};
+
+template <typename indptr_t, typename indices_t>
+struct IteratorFuncAddOffset {
+  indptr_t* indptr;
+  indptr_t* sliced_indptr;
+  indices_t* indices;
+  __host__ __device__ auto operator()(int64_t i) {
+    return thrust::transform_output_iterator{
+        indices + indptr[i], AddOffset<indptr_t>{sliced_indptr[i]}};
+  }
+};
+
+template <typename indptr_t, typename in_degree_iterator_t>
+struct SegmentEndFunc {
+  indptr_t* indptr;
+  in_degree_iterator_t in_degree;
+  __host__ __device__ auto operator()(int64_t i) {
+    return indptr[i] + in_degree[i];
+  }
+};
+
+c10::intrusive_ptr<sampling::FusedSampledSubgraph> SampleNeighbors(
+    torch::Tensor indptr, torch::Tensor indices,
+    torch::optional<torch::Tensor> seeds,
+    torch::optional<std::vector<int64_t>> seed_offsets,
+    const std::vector<int64_t>& fanouts, bool replace, bool layer,
+    bool returning_indices_is_optional,
+    torch::optional<torch::Tensor> type_per_edge,
+    torch::optional<torch::Tensor> probs_or_mask,
+    torch::optional<torch::Tensor> node_type_offset,
+    torch::optional<torch::Dict<std::string, int64_t>> node_type_to_id,
+    torch::optional<torch::Dict<std::string, int64_t>> edge_type_to_id,
+    torch::optional<torch::Tensor> random_seed_tensor, float seed2_contribution,
+    // Optional temporal sampling arguments begin.
+    torch::optional<torch::Tensor> seeds_timestamp,
+    torch::optional<torch::Tensor> seeds_pre_time_window,
+    torch::optional<torch::Tensor> node_timestamp,
+    torch::optional<torch::Tensor> edge_timestamp
+    // Optional temporal sampling arguments end.
+) {
+  // When seed_offsets.has_value() in the hetero case, we compute the output of
+  // sample_neighbors _convert_to_sampled_subgraph in a fused manner so that
+  // _convert_to_sampled_subgraph only has to perform slices over the returned
+  // indptr and indices tensors to form CSC outputs for each edge type.
+  TORCH_CHECK(!replace, "Sampling with replacement is not supported yet!");
+  // Assume that indptr, indices, seeds, type_per_edge and probs_or_mask
+  // are all resident on the GPU. If not, it is better to first extract them
+  // before calling this function.
+  auto allocator = cuda::GetAllocator();
+  auto num_rows =
+      seeds.has_value() ? seeds.value().size(0) : indptr.size(0) - 1;
+  auto fanouts_pinned = torch::empty(
+      fanouts.size(),
+      c10::TensorOptions().dtype(torch::kLong).pinned_memory(true));
+  auto fanouts_pinned_ptr = fanouts_pinned.data_ptr<int64_t>();
+  for (size_t i = 0; i < fanouts.size(); i++) {
+    fanouts_pinned_ptr[i] =
+        fanouts[i] >= 0 ? fanouts[i] : std::numeric_limits<int64_t>::max();
+  }
+  // Finally, copy the adjusted fanout values to the device memory.
+  auto fanouts_device = allocator.AllocateStorage<int64_t>(fanouts.size());
+  CUDA_CALL(cudaMemcpyAsync(
+      fanouts_device.get(), fanouts_pinned_ptr,
+      sizeof(int64_t) * fanouts.size(), cudaMemcpyHostToDevice,
+      cuda::GetCurrentStream()));
+  auto in_degree_and_sliced_indptr = SliceCSCIndptr(indptr, seeds);
+  auto in_degree = std::get<0>(in_degree_and_sliced_indptr);
+  auto sliced_indptr = std::get<1>(in_degree_and_sliced_indptr);
+  const auto homo_in_degree = in_degree;
+  const auto homo_sliced_indptr = sliced_indptr;
+  auto max_in_degree = torch::empty(
+      1,
+      c10::TensorOptions().dtype(in_degree.scalar_type()).pinned_memory(true));
+  AT_DISPATCH_INDEX_TYPES(
+      indptr.scalar_type(), "SampleNeighborsMaxInDegree", ([&] {
+        CUB_CALL(
+            DeviceReduce::Max, in_degree.data_ptr<index_t>(),
+            max_in_degree.data_ptr<index_t>(), num_rows);
+      }));
+  // Protect access to max_in_degree with a CUDAEvent
+  at::cuda::CUDAEvent max_in_degree_event;
+  max_in_degree_event.record();
+  torch::optional<int64_t> num_edges;
+  torch::Tensor sub_indptr;
+  if (!seeds.has_value()) {
+    num_edges = indices.size(0);
+    sub_indptr = indptr;
+  }
+  torch::optional<torch::Tensor> sliced_probs_or_mask;
+  if (probs_or_mask.has_value()) {
+    if (seeds.has_value()) {
+      torch::Tensor sliced_probs_or_mask_tensor;
+      std::tie(sub_indptr, sliced_probs_or_mask_tensor) = IndexSelectCSCImpl(
+          in_degree, sliced_indptr, probs_or_mask.value(), seeds.value(),
+          indptr.size(0) - 2, num_edges);
+      sliced_probs_or_mask = sliced_probs_or_mask_tensor;
+      num_edges = sliced_probs_or_mask_tensor.size(0);
+    } else {
+      sliced_probs_or_mask = probs_or_mask;
+    }
+  }
+  if (fanouts.size() > 1) {
+    torch::Tensor sliced_type_per_edge;
+    if (seeds.has_value()) {
+      std::tie(sub_indptr, sliced_type_per_edge) = IndexSelectCSCImpl(
+          in_degree, sliced_indptr, type_per_edge.value(), seeds.value(),
+          indptr.size(0) - 2, num_edges);
+    } else {
+      sliced_type_per_edge = type_per_edge.value();
+    }
+    std::tie(sub_indptr, in_degree, sliced_indptr) = SliceCSCIndptrHetero(
+        sub_indptr, sliced_type_per_edge, sliced_indptr, fanouts.size());
+    num_rows = sliced_indptr.size(0);
+    num_edges = sliced_type_per_edge.size(0);
+  }
+  // If sub_indptr was not computed in the two code blocks above:
+  if (seeds.has_value() && !probs_or_mask.has_value() && fanouts.size() <= 1) {
+    sub_indptr = ExclusiveCumSum(in_degree);
+  }
+  torch::optional<torch::Tensor> homo_coo_rows;
+  if (seeds_timestamp.has_value()) {
+    // Temporal sampling is enabled.
+    const auto homo_sub_indptr =
+        fanouts.size() > 1 ? ExclusiveCumSum(homo_in_degree) : sub_indptr;
+    homo_coo_rows = ExpandIndptrImpl(
+        homo_sub_indptr, indices.scalar_type(), torch::nullopt, num_edges);
+    num_edges = homo_coo_rows->size(0);
+    const auto is_probs_initialized = sliced_probs_or_mask.has_value();
+    if (!is_probs_initialized) {
+      sliced_probs_or_mask =
+          torch::empty(*num_edges, sub_indptr.options().dtype(torch::kBool));
+    }
+    GRAPHBOLT_DISPATCH_ALL_TYPES(
+        sliced_probs_or_mask->scalar_type(),
+        "SampleNeighborsTemporalProbsOrMask", ([&] {
+          const scalar_t* input_probs_ptr =
+              is_probs_initialized ? sliced_probs_or_mask->data_ptr<scalar_t>()
+                                   : nullptr;
+          auto output_probs_ptr = sliced_probs_or_mask->data_ptr<scalar_t>();
+          using timestamp_t = int64_t;
+          const auto seeds_timestamp_ptr =
+              seeds_timestamp->data_ptr<timestamp_t>();
+          const timestamp_t* seeds_pre_time_window_ptr =
+              seeds_pre_time_window.has_value()
+                  ? seeds_pre_time_window->data_ptr<timestamp_t>()
+                  : nullptr;
+          const timestamp_t* node_timestamp_ptr =
+              node_timestamp.has_value()
+                  ? node_timestamp->data_ptr<timestamp_t>()
+                  : nullptr;
+          const timestamp_t* edge_timestamp_ptr =
+              edge_timestamp.has_value()
+                  ? edge_timestamp->data_ptr<timestamp_t>()
+                  : nullptr;
+          AT_DISPATCH_INDEX_TYPES(
+              homo_coo_rows->scalar_type(),
+              "SampleNeighborsTemporalMaskIndices", ([&] {
+                const auto coo_rows_ptr = homo_coo_rows->data_ptr<index_t>();
+                const auto indices_ptr = indices.data_ptr<index_t>();
+                AT_DISPATCH_INDEX_TYPES(
+                    homo_sliced_indptr.scalar_type(),
+                    "SampleNeighborsTemporalMaskIndptr", ([&] {
+                      const auto sliced_indptr_data =
+                          homo_sliced_indptr.data_ptr<index_t>();
+                      const auto sub_indptr_data =
+                          homo_sub_indptr.data_ptr<index_t>();
+                      CUB_CALL(
+                          DeviceFor::Bulk, *num_edges,
+                          [=] __device__(int64_t i) {
+                            const auto row = coo_rows_ptr[i];
+                            const auto seed_timestamp =
+                                seeds_timestamp_ptr[row];
+                            const auto row_offset = i - sub_indptr_data[row];
+                            const auto in_idx =
+                                sliced_indptr_data[row] + row_offset;
+                            bool mask = true;
+                            if (node_timestamp_ptr) {
+                              const auto index = indices_ptr[in_idx];
+                              const auto neighbor_timestamp =
+                                  node_timestamp_ptr[index];
+                              mask &= neighbor_timestamp < seed_timestamp;
+                              if (seeds_pre_time_window_ptr) {
+                                mask &= neighbor_timestamp >
+                                        seed_timestamp -
+                                            seeds_pre_time_window_ptr[row];
+                              }
+                            }
+                            if (edge_timestamp_ptr) {
+                              const auto edge_timestamp =
+                                  edge_timestamp_ptr[in_idx];
+                              mask &= edge_timestamp < seed_timestamp;
+                              if (seeds_pre_time_window_ptr) {
+                                mask &= edge_timestamp >
+                                        seed_timestamp -
+                                            seeds_pre_time_window_ptr[row];
+                              }
+                            }
+                            const scalar_t prob = input_probs_ptr
+                                                      ? input_probs_ptr[i]
+                                                      : scalar_t{1};
+                            output_probs_ptr[i] =
+                                prob * static_cast<scalar_t>(mask);
+                          });
+                    }));
+              }));
+        }));
+  }
+  const continuous_seed random_seed = [&] {
+    if (random_seed_tensor.has_value()) {
+      return continuous_seed(random_seed_tensor.value(), seed2_contribution);
+    } else {
+      return continuous_seed{RandomEngine::ThreadLocal()->RandInt(
+          static_cast<int64_t>(0), std::numeric_limits<int64_t>::max())};
+    }
+  }();
+  auto output_indptr = torch::empty_like(sub_indptr);
+  torch::Tensor picked_eids;
+  torch::optional<torch::Tensor> output_indices;
+
+  AT_DISPATCH_INDEX_TYPES(
+      indptr.scalar_type(), "SampleNeighborsIndptr", ([&] {
+        using indptr_t = index_t;
+        if (sliced_probs_or_mask.has_value()) {
+          // Count nonzero probs into in_degree.
+          GRAPHBOLT_DISPATCH_ALL_TYPES(
+              sliced_probs_or_mask->scalar_type(),
+              "SampleNeighborsPositiveProbs", ([&] {
+                using probs_t = scalar_t;
+                auto is_nonzero = thrust::make_transform_iterator(
+                    sliced_probs_or_mask->data_ptr<probs_t>(), IsPositive{});
+                CUB_CALL(
+                    DeviceSegmentedReduce::Sum, is_nonzero,
+                    in_degree.data_ptr<indptr_t>(), num_rows,
+                    sub_indptr.data_ptr<indptr_t>(),
+                    sub_indptr.data_ptr<indptr_t>() + 1);
+              }));
+        }
+        thrust::counting_iterator<int64_t> iota(0);
+        auto sampled_degree = thrust::make_transform_iterator(
+            iota, MinInDegreeFanout<indptr_t>{
+                      in_degree.data_ptr<indptr_t>(), fanouts_device.get(),
+                      fanouts.size()});
+
+        // Compute output_indptr.
+        CUB_CALL(
+            DeviceScan::ExclusiveSum, sampled_degree,
+            output_indptr.data_ptr<indptr_t>(), num_rows + 1);
+
+        auto num_sampled_edges =
+            cuda::CopyScalar{output_indptr.data_ptr<indptr_t>() + num_rows};
+
+        // This operation is placed after num_sampled_edges copy is started to
+        // hide the latency of copy synchronization later.
+        torch::Tensor coo_rows;
+        if (!homo_coo_rows.has_value() || fanouts.size() > 1) {
+          coo_rows = ExpandIndptrImpl(
+              sub_indptr, indices.scalar_type(), torch::nullopt, num_edges);
+          num_edges = coo_rows.size(0);
+        } else {
+          coo_rows = *homo_coo_rows;
+        }
+
+        // Find the smallest integer type to store the edge id offsets. We synch
+        // the CUDAEvent so that the access is safe.
+        auto compute_num_bits = [&] {
+          max_in_degree_event.synchronize();
+          return cuda::NumberOfBits(max_in_degree.data_ptr<indptr_t>()[0]);
+        };
+        if (layer || sliced_probs_or_mask.has_value()) {
+          const int num_bits = compute_num_bits();
+          std::array<int, 4> type_bits = {8, 16, 32, 64};
+          const auto type_index =
+              std::lower_bound(type_bits.begin(), type_bits.end(), num_bits) -
+              type_bits.begin();
+          std::array<torch::ScalarType, 5> types = {
+              torch::kByte, torch::kInt16, torch::kInt32, torch::kLong,
+              torch::kLong};
+          auto edge_id_dtype = types[type_index];
+          AT_DISPATCH_INTEGRAL_TYPES(
+              edge_id_dtype, "SampleNeighborsEdgeIDs", ([&] {
+                using edge_id_t = std::make_unsigned_t<scalar_t>;
+                TORCH_CHECK(
+                    num_bits <= sizeof(edge_id_t) * 8,
+                    "Selected edge_id_t must be capable of storing edge_ids.");
+                // Using bfloat16 for random numbers works just as reliably as
+                // float32 and provides around 30% speedup.
+                using rnd_t = nv_bfloat16;
+                auto randoms =
+                    allocator.AllocateStorage<rnd_t>(num_edges.value());
+                auto randoms_sorted =
+                    allocator.AllocateStorage<rnd_t>(num_edges.value());
+                auto edge_id_segments =
+                    allocator.AllocateStorage<edge_id_t>(num_edges.value());
+                auto sorted_edge_id_segments =
+                    allocator.AllocateStorage<edge_id_t>(num_edges.value());
+                AT_DISPATCH_INDEX_TYPES(
+                    indices.scalar_type(), "SampleNeighborsIndices", ([&] {
+                      using indices_t = index_t;
+                      auto probs_or_mask_scalar_type = torch::kFloat32;
+                      if (sliced_probs_or_mask.has_value()) {
+                        probs_or_mask_scalar_type =
+                            sliced_probs_or_mask->scalar_type();
+                      }
+                      GRAPHBOLT_DISPATCH_ALL_TYPES(
+                          probs_or_mask_scalar_type, "SampleNeighborsProbs",
+                          ([&] {
+                            using probs_t = scalar_t;
+                            probs_t* sliced_probs_ptr = nullptr;
+                            if (sliced_probs_or_mask.has_value()) {
+                              sliced_probs_ptr =
+                                  sliced_probs_or_mask->data_ptr<probs_t>();
+                            }
+                            const indices_t* indices_ptr =
+                                layer ? indices.data_ptr<indices_t>() : nullptr;
+                            const dim3 block(BLOCK_SIZE);
+                            const dim3 grid(
+                                (num_edges.value() + BLOCK_SIZE - 1) /
+                                BLOCK_SIZE);
+                            // Compute row and random number pairs.
+                            CUDA_KERNEL_CALL(
+                                _ComputeRandoms, grid, block, 0,
+                                num_edges.value(),
+                                sliced_indptr.data_ptr<indptr_t>(),
+                                sub_indptr.data_ptr<indptr_t>(),
+                                coo_rows.data_ptr<indices_t>(),
+                                sliced_probs_ptr, indices_ptr, random_seed,
+                                randoms.get(), edge_id_segments.get());
+                          }));
+                    }));
+
+                // Sort the random numbers along with edge ids, after
+                // sorting the first fanout elements of each row will
+                // give us the sampled edges.
+                CUB_CALL(
+                    DeviceSegmentedSort::SortPairs, randoms.get(),
+                    randoms_sorted.get(), edge_id_segments.get(),
+                    sorted_edge_id_segments.get(), num_edges.value(), num_rows,
+                    sub_indptr.data_ptr<indptr_t>(),
+                    sub_indptr.data_ptr<indptr_t>() + 1);
+
+                picked_eids = torch::empty(
+                    static_cast<indptr_t>(num_sampled_edges),
+                    sub_indptr.options());
+
+                // Need to sort the sampled edges only when fanouts.size() == 1
+                // since multiple fanout sampling case is automatically going to
+                // be sorted.
+                if (type_per_edge && fanouts.size() == 1) {
+                  // Ensuring sort result still ends up in
+                  // sorted_edge_id_segments
+                  std::swap(edge_id_segments, sorted_edge_id_segments);
+                  auto sampled_segment_end_it = thrust::make_transform_iterator(
+                      iota,
+                      SegmentEndFunc<indptr_t, decltype(sampled_degree)>{
+                          sub_indptr.data_ptr<indptr_t>(), sampled_degree});
+                  CUB_CALL(
+                      DeviceSegmentedSort::SortKeys, edge_id_segments.get(),
+                      sorted_edge_id_segments.get(), picked_eids.size(0),
+                      num_rows, sub_indptr.data_ptr<indptr_t>(),
+                      sampled_segment_end_it);
+                }
+
+                auto input_buffer_it = thrust::make_transform_iterator(
+                    iota, IteratorFunc<indptr_t, edge_id_t>{
+                              sub_indptr.data_ptr<indptr_t>(),
+                              sorted_edge_id_segments.get()});
+                auto output_buffer_it = thrust::make_transform_iterator(
+                    iota, IteratorFuncAddOffset<indptr_t, indptr_t>{
+                              output_indptr.data_ptr<indptr_t>(),
+                              sliced_indptr.data_ptr<indptr_t>(),
+                              picked_eids.data_ptr<indptr_t>()});
+                constexpr int64_t max_copy_at_once =
+                    std::numeric_limits<int32_t>::max();
+
+                // Copy the sampled edge ids into picked_eids tensor.
+                for (int64_t i = 0; i < num_rows; i += max_copy_at_once) {
+                  CUB_CALL(
+                      DeviceCopy::Batched, input_buffer_it + i,
+                      output_buffer_it + i, sampled_degree + i,
+                      std::min(num_rows - i, max_copy_at_once));
+                }
+              }));
+        } else {  // Non-weighted neighbor sampling.
+          picked_eids = torch::zeros(num_edges.value(), sub_indptr.options());
+          const auto sort_needed = type_per_edge && fanouts.size() == 1;
+          const auto sliced_indptr_ptr =
+              sort_needed ? nullptr : sliced_indptr.data_ptr<indptr_t>();
+
+          const dim3 block(BLOCK_SIZE);
+          const dim3 grid(
+              (std::min(num_edges.value(), static_cast<int64_t>(1 << 20)) +
+               BLOCK_SIZE - 1) /
+              BLOCK_SIZE);
+          AT_DISPATCH_INDEX_TYPES(
+              indices.scalar_type(), "SampleNeighborsIndices", ([&] {
+                using indices_t = index_t;
+                // Compute row and random number pairs.
+                CUDA_KERNEL_CALL(
+                    _ComputeRandomsNS, grid, block, 0, num_edges.value(),
+                    sliced_indptr_ptr, sub_indptr.data_ptr<indptr_t>(),
+                    output_indptr.data_ptr<indptr_t>(),
+                    coo_rows.data_ptr<indices_t>(), random_seed.get_seed(0),
+                    picked_eids.data_ptr<indptr_t>());
+              }));
+
+          picked_eids =
+              picked_eids.slice(0, 0, static_cast<indptr_t>(num_sampled_edges));
+
+          // Need to sort the sampled edges only when fanouts.size() == 1
+          // since multiple fanout sampling case is automatically going to
+          // be sorted.
+          if (sort_needed) {
+            const int num_bits = compute_num_bits();
+            std::array<int, 4> type_bits = {8, 15, 31, 63};
+            const auto type_index =
+                std::lower_bound(type_bits.begin(), type_bits.end(), num_bits) -
+                type_bits.begin();
+            std::array<torch::ScalarType, 5> types = {
+                torch::kByte, torch::kInt16, torch::kInt32, torch::kLong,
+                torch::kLong};
+            auto edge_id_dtype = types[type_index];
+            AT_DISPATCH_INTEGRAL_TYPES(
+                edge_id_dtype, "SampleNeighborsEdgeIDs", ([&] {
+                  using edge_id_t = scalar_t;
+                  TORCH_CHECK(
+                      num_bits <= sizeof(edge_id_t) * 8,
+                      "Selected edge_id_t must be capable of storing "
+                      "edge_ids.");
+                  auto picked_offsets = picked_eids.to(edge_id_dtype);
+                  auto sorted_offsets = torch::empty_like(picked_offsets);
+                  CUB_CALL(
+                      DeviceSegmentedSort::SortKeys,
+                      picked_offsets.data_ptr<edge_id_t>(),
+                      sorted_offsets.data_ptr<edge_id_t>(), picked_eids.size(0),
+                      num_rows, output_indptr.data_ptr<indptr_t>(),
+                      output_indptr.data_ptr<indptr_t>() + 1);
+                  auto edge_id_offsets = ExpandIndptrImpl(
+                      output_indptr, picked_eids.scalar_type(), sliced_indptr,
+                      picked_eids.size(0));
+                  picked_eids = sorted_offsets.to(picked_eids.scalar_type()) +
+                                edge_id_offsets;
+                }));
+          }
+        }
+
+        if (!returning_indices_is_optional || utils::is_on_gpu(indices)) {
+          output_indices = Gather(indices, picked_eids);
+        }
+      }));
+
+  torch::optional<torch::Tensor> output_type_per_edge;
+  torch::optional<torch::Tensor> edge_offsets;
+  if (type_per_edge && seed_offsets) {
+    const int64_t num_etypes =
+        edge_type_to_id.has_value() ? edge_type_to_id->size() : 1;
+    // If we performed homogenous sampling on hetero graph, we have to look at
+    // type_per_edge of sampled edges and determine the offsets of different
+    // sampled etypes and convert to fused hetero indptr representation.
+    if (fanouts.size() == 1) {
+      output_type_per_edge = Gather(*type_per_edge, picked_eids);
+      torch::Tensor output_in_degree, sliced_output_indptr;
+      sliced_output_indptr =
+          output_indptr.slice(0, 0, output_indptr.size(0) - 1);
+      std::tie(output_indptr, output_in_degree, sliced_output_indptr) =
+          SliceCSCIndptrHetero(
+              output_indptr, output_type_per_edge.value(), sliced_output_indptr,
+              num_etypes);
+      // We use num_rows to hold num_seeds * num_etypes. So, it needs to be
+      // updated when sampling with a single fanout value when the graph is
+      // heterogenous.
+      num_rows = sliced_output_indptr.size(0);
+    }
+    // Here, we check what are the dst node types for the given seeds so that
+    // we can compute the output indptr space later.
+    std::vector<int64_t> etype_id_to_dst_ntype_id(num_etypes);
+    // Here, we check what are the src node types for the given seeds so that
+    // we can subtract source node offset from indices later.
+    auto etype_id_to_src_ntype_id = torch::empty(
+        2 * num_etypes,
+        c10::TensorOptions().dtype(torch::kLong).pinned_memory(true));
+    auto etype_id_to_src_ntype_id_ptr =
+        etype_id_to_src_ntype_id.data_ptr<int64_t>();
+    for (auto& etype_and_id : edge_type_to_id.value()) {
+      auto etype = etype_and_id.key();
+      auto id = etype_and_id.value();
+      auto [src_type, dst_type] = utils::parse_src_dst_ntype_from_etype(etype);
+      etype_id_to_dst_ntype_id[id] = node_type_to_id->at(dst_type);
+      etype_id_to_src_ntype_id_ptr[2 * id] =
+          etype_id_to_src_ntype_id_ptr[2 * id + 1] =
+              node_type_to_id->at(src_type);
+    }
+    auto indices_offsets_device = torch::empty(
+        etype_id_to_src_ntype_id.size(0),
+        picked_eids.options().dtype(torch::kLong));
+    AT_DISPATCH_INDEX_TYPES(
+        node_type_offset->scalar_type(), "SampleNeighborsNodeTypeOffset", ([&] {
+          THRUST_CALL(
+              gather, etype_id_to_src_ntype_id_ptr,
+              etype_id_to_src_ntype_id_ptr + etype_id_to_src_ntype_id.size(0),
+              node_type_offset->data_ptr<index_t>(),
+              indices_offsets_device.data_ptr<int64_t>());
+        }));
+    // For each edge type, we compute the start and end offsets to index into
+    // indptr to form the final output_indptr.
+    auto indptr_offsets = torch::empty(
+        num_etypes * 2,
+        c10::TensorOptions().dtype(torch::kLong).pinned_memory(true));
+    auto indptr_offsets_ptr = indptr_offsets.data_ptr<int64_t>();
+    // We compute the indptr offsets here, right now, output_indptr is of size
+    // # seeds * num_etypes + 1. We can simply take slices to get correct output
+    // indptr. The final output_indptr is same as current indptr except that
+    // some intermediate values are removed to change the node ids space from
+    // all of the seed vertices to the node id space of the dst node type of
+    // each edge type.
+    for (int i = 0; i < num_etypes; i++) {
+      indptr_offsets_ptr[2 * i] = num_rows / num_etypes * i +
+                                  seed_offsets->at(etype_id_to_dst_ntype_id[i]);
+      indptr_offsets_ptr[2 * i + 1] =
+          num_rows / num_etypes * i +
+          seed_offsets->at(etype_id_to_dst_ntype_id[i] + 1);
+    }
+    auto permutation = torch::arange(
+        0, num_rows * num_etypes, num_etypes, output_indptr.options());
+    permutation =
+        permutation.remainder(num_rows) + permutation.div(num_rows, "floor");
+    // This permutation, when applied sorts the sampled edges with respect to
+    // edge types.
+    auto [output_in_degree, sliced_output_indptr] =
+        SliceCSCIndptr(output_indptr, permutation);
+    std::tie(output_indptr, picked_eids) = IndexSelectCSCImpl(
+        output_in_degree, sliced_output_indptr, picked_eids, permutation,
+        num_rows - 1, picked_eids.size(0));
+    edge_offsets = torch::empty(
+        num_etypes * 2, c10::TensorOptions()
+                            .dtype(output_indptr.scalar_type())
+                            .pinned_memory(true));
+    auto edge_offsets_device =
+        torch::empty(num_etypes * 2, output_indptr.options());
+    at::cuda::CUDAEvent edge_offsets_event;
+    AT_DISPATCH_INDEX_TYPES(
+        indptr.scalar_type(), "SampleNeighborsEdgeOffsets", ([&] {
+          auto edge_offsets_pinned_device_pair =
+              thrust::make_transform_output_iterator(
+                  thrust::make_zip_iterator(
+                      edge_offsets->data_ptr<index_t>(),
+                      edge_offsets_device.data_ptr<index_t>()),
+                  [=] __device__(index_t x) {
+                    return thrust::make_tuple(x, x);
+                  });
+          THRUST_CALL(
+              gather, indptr_offsets_ptr,
+              indptr_offsets_ptr + indptr_offsets.size(0),
+              output_indptr.data_ptr<index_t>(),
+              edge_offsets_pinned_device_pair);
+        }));
+    edge_offsets_event.record();
+    if (output_indices.has_value()) {
+      auto indices_offset_subtract = ExpandIndptrImpl(
+          edge_offsets_device, indices.scalar_type(), indices_offsets_device,
+          output_indices->size(0));
+      // The output_indices is permuted here.
+      std::tie(output_indptr, output_indices) = IndexSelectCSCImpl(
+          output_in_degree, sliced_output_indptr, *output_indices, permutation,
+          num_rows - 1, output_indices->size(0));
+      *output_indices -= indices_offset_subtract;
+    }
+    auto output_indptr_offsets = torch::empty(
+        num_etypes * 2,
+        c10::TensorOptions().dtype(torch::kLong).pinned_memory(true));
+    auto output_indptr_offsets_ptr = output_indptr_offsets.data_ptr<int64_t>();
+    std::vector<torch::Tensor> indptr_list;
+    for (int i = 0; i < num_etypes; i++) {
+      indptr_list.push_back(output_indptr.slice(
+          0, indptr_offsets_ptr[2 * i], indptr_offsets_ptr[2 * i + 1] + 1));
+      output_indptr_offsets_ptr[2 * i] =
+          i == 0 ? 0 : output_indptr_offsets_ptr[2 * i - 1];
+      output_indptr_offsets_ptr[2 * i + 1] =
+          output_indptr_offsets_ptr[2 * i] + indptr_list.back().size(0);
+    }
+    auto output_indptr_offsets_device = torch::empty(
+        output_indptr_offsets.size(0),
+        output_indptr.options().dtype(torch::kLong));
+    THRUST_CALL(
+        copy_n, output_indptr_offsets_ptr, output_indptr_offsets.size(0),
+        output_indptr_offsets_device.data_ptr<int64_t>());
+    // We form the final output indptr by concatenating pieces for different
+    // edge types.
+    output_indptr = torch::cat(indptr_list);
+    auto indptr_offset_subtract = ExpandIndptrImpl(
+        output_indptr_offsets_device, indptr.scalar_type(), edge_offsets_device,
+        output_indptr.size(0));
+    output_indptr -= indptr_offset_subtract;
+    edge_offsets_event.synchronize();
+    // We read the edge_offsets here, they are in pairs but we don't need it to
+    // be in pairs. So we remove the duplicate information from it and turn it
+    // into a real offsets array.
+    AT_DISPATCH_INDEX_TYPES(
+        indptr.scalar_type(), "SampleNeighborsEdgeOffsetsCheck", ([&] {
+          auto edge_offsets_ptr = edge_offsets->data_ptr<index_t>();
+          TORCH_CHECK(edge_offsets_ptr[0] == 0, "edge_offsets is incorrect.");
+          for (int i = 1; i < num_etypes; i++) {
+            TORCH_CHECK(
+                edge_offsets_ptr[2 * i - 1] == edge_offsets_ptr[2 * i],
+                "edge_offsets is incorrect.");
+          }
+          TORCH_CHECK(
+              edge_offsets_ptr[2 * num_etypes - 1] == picked_eids.size(0),
+              "edge_offsets is incorrect.");
+          for (int i = 0; i < num_etypes; i++) {
+            edge_offsets_ptr[i + 1] = edge_offsets_ptr[2 * i + 1];
+          }
+        }));
+    edge_offsets = edge_offsets->slice(0, 0, num_etypes + 1);
+  } else {
+    // Convert output_indptr back to homo by discarding intermediate offsets.
+    output_indptr =
+        output_indptr.slice(0, 0, output_indptr.size(0), fanouts.size());
+    if (type_per_edge)
+      output_type_per_edge = Gather(*type_per_edge, picked_eids);
+  }
+
+  return c10::make_intrusive<sampling::FusedSampledSubgraph>(
+      output_indptr, output_indices, picked_eids, seeds, torch::nullopt,
+      output_type_per_edge, edge_offsets);
+}
+
+}  //  namespace ops
+}  //  namespace graphbolt

graphbolt/src/cuda/sampling_utils.cu

@@ -0,0 +1,153 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/sampling_utils.cu
+ * @brief Sampling utility function implementations on CUDA.
+ */
+#include <thrust/for_each.h>
+#include <thrust/iterator/counting_iterator.h>
+
+#include <cub/cub.cuh>
+
+#include "./common.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+// Given rows and indptr, computes:
+// inrow_indptr[i] = indptr[rows[i]];
+// in_degree[i] = indptr[rows[i] + 1] - indptr[rows[i]];
+template <typename indptr_t, typename nodes_t>
+struct SliceFunc {
+  const nodes_t* rows;
+  const indptr_t* indptr;
+  indptr_t* in_degree;
+  indptr_t* inrow_indptr;
+  __host__ __device__ auto operator()(int64_t tIdx) {
+    const auto out_row = rows[tIdx];
+    const auto indptr_val = indptr[out_row];
+    const auto degree = indptr[out_row + 1] - indptr_val;
+    in_degree[tIdx] = degree;
+    inrow_indptr[tIdx] = indptr_val;
+  }
+};
+
+// Returns (indptr[nodes + 1] - indptr[nodes], indptr[nodes])
+std::tuple<torch::Tensor, torch::Tensor> SliceCSCIndptr(
+    torch::Tensor indptr, torch::optional<torch::Tensor> nodes_optional) {
+  if (nodes_optional.has_value()) {
+    auto nodes = nodes_optional.value();
+    const int64_t num_nodes = nodes.size(0);
+    // Read indptr only once in case it is pinned and access is slow.
+    auto sliced_indptr =
+        torch::empty(num_nodes, nodes.options().dtype(indptr.scalar_type()));
+    // compute in-degrees
+    auto in_degree = torch::empty(
+        num_nodes + 1, nodes.options().dtype(indptr.scalar_type()));
+    thrust::counting_iterator<int64_t> iota(0);
+    AT_DISPATCH_INTEGRAL_TYPES(
+        indptr.scalar_type(), "IndexSelectCSCIndptr", ([&] {
+          using indptr_t = scalar_t;
+          AT_DISPATCH_INDEX_TYPES(
+              nodes.scalar_type(), "IndexSelectCSCNodes", ([&] {
+                using nodes_t = index_t;
+                THRUST_CALL(
+                    for_each, iota, iota + num_nodes,
+                    SliceFunc<indptr_t, nodes_t>{
+                        nodes.data_ptr<nodes_t>(), indptr.data_ptr<indptr_t>(),
+                        in_degree.data_ptr<indptr_t>(),
+                        sliced_indptr.data_ptr<indptr_t>()});
+              }));
+        }));
+    return {in_degree, sliced_indptr};
+  } else {
+    const int64_t num_nodes = indptr.size(0) - 1;
+    auto sliced_indptr = indptr.slice(0, 0, num_nodes);
+    auto in_degree = torch::empty(
+        num_nodes + 2, indptr.options().dtype(indptr.scalar_type()));
+    AT_DISPATCH_INTEGRAL_TYPES(
+        indptr.scalar_type(), "IndexSelectCSCIndptr", ([&] {
+          using indptr_t = scalar_t;
+          CUB_CALL(
+              DeviceAdjacentDifference::SubtractLeftCopy,
+              indptr.data_ptr<indptr_t>(), in_degree.data_ptr<indptr_t>(),
+              num_nodes + 1, cub::Difference{});
+        }));
+    in_degree = in_degree.slice(0, 1);
+    return {in_degree, sliced_indptr};
+  }
+}
+
+template <typename indptr_t, typename etype_t>
+struct EdgeTypeSearch {
+  const indptr_t* sub_indptr;
+  const indptr_t* sliced_indptr;
+  const etype_t* etypes;
+  int64_t num_fanouts;
+  int64_t num_rows;
+  indptr_t* new_sub_indptr;
+  indptr_t* new_sliced_indptr;
+  __host__ __device__ auto operator()(int64_t i) {
+    const auto homo_i = i / num_fanouts;
+    const auto indptr_i = sub_indptr[homo_i];
+    const auto degree = sub_indptr[homo_i + 1] - indptr_i;
+    const etype_t etype = i % num_fanouts;
+    auto offset = cub::LowerBound(etypes + indptr_i, degree, etype);
+    new_sub_indptr[i] = indptr_i + offset;
+    new_sliced_indptr[i] = sliced_indptr[homo_i] + offset;
+    if (i == num_rows - 1) new_sub_indptr[num_rows] = indptr_i + degree;
+  }
+};
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> SliceCSCIndptrHetero(
+    torch::Tensor sub_indptr, torch::Tensor etypes, torch::Tensor sliced_indptr,
+    int64_t num_fanouts) {
+  auto num_rows = (sub_indptr.size(0) - 1) * num_fanouts;
+  auto new_sub_indptr = torch::empty(num_rows + 1, sub_indptr.options());
+  auto new_indegree = torch::empty(num_rows + 2, sub_indptr.options());
+  auto new_sliced_indptr = torch::empty(num_rows, sliced_indptr.options());
+  thrust::counting_iterator<int64_t> iota(0);
+  AT_DISPATCH_INTEGRAL_TYPES(
+      sub_indptr.scalar_type(), "SliceCSCIndptrHeteroIndptr", ([&] {
+        using indptr_t = scalar_t;
+        AT_DISPATCH_INTEGRAL_TYPES(
+            etypes.scalar_type(), "SliceCSCIndptrHeteroTypePerEdge", ([&] {
+              using etype_t = scalar_t;
+              THRUST_CALL(
+                  for_each, iota, iota + num_rows,
+                  EdgeTypeSearch<indptr_t, etype_t>{
+                      sub_indptr.data_ptr<indptr_t>(),
+                      sliced_indptr.data_ptr<indptr_t>(),
+                      etypes.data_ptr<etype_t>(), num_fanouts, num_rows,
+                      new_sub_indptr.data_ptr<indptr_t>(),
+                      new_sliced_indptr.data_ptr<indptr_t>()});
+            }));
+        CUB_CALL(
+            DeviceAdjacentDifference::SubtractLeftCopy,
+            new_sub_indptr.data_ptr<indptr_t>(),
+            new_indegree.data_ptr<indptr_t>(), num_rows + 1, cub::Difference{});
+      }));
+  // Discard the first element of the SubtractLeftCopy result and ensure that
+  // new_indegree tensor has size num_rows + 1 so that its ExclusiveCumSum is
+  // directly equivalent to new_sub_indptr.
+  // Equivalent to new_indegree = new_indegree[1:] in Python.
+  new_indegree = new_indegree.slice(0, 1);
+  return {new_sub_indptr, new_indegree, new_sliced_indptr};
+}
+
+}  //  namespace ops
+}  //  namespace graphbolt

graphbolt/src/cuda/sort_impl.cu

@@ -0,0 +1,78 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/sort_impl.cu
+ * @brief Sort implementation on CUDA.
+ */
+#include <c10/core/ScalarType.h>
+
+#include <cub/cub.cuh>
+
+#include "./common.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+template <bool return_original_positions, typename scalar_t>
+std::conditional_t<
+    return_original_positions, std::pair<torch::Tensor, torch::Tensor>,
+    torch::Tensor>
+Sort(const scalar_t* input_keys, int64_t num_items, int num_bits) {
+  const auto options = torch::TensorOptions().device(c10::DeviceType::CUDA);
+  constexpr c10::ScalarType dtype = c10::CppTypeToScalarType<scalar_t>::value;
+  auto sorted_array = torch::empty(num_items, options.dtype(dtype));
+  auto sorted_keys = sorted_array.data_ptr<scalar_t>();
+  if (num_bits == 0) {
+    num_bits = sizeof(scalar_t) * 8;
+  }
+
+  if constexpr (return_original_positions) {
+    // We utilize int64_t for the values array. (torch::kLong == int64_t)
+    auto original_idx = torch::arange(num_items, options.dtype(torch::kLong));
+    auto sorted_idx = torch::empty_like(original_idx);
+    const int64_t* input_values = original_idx.data_ptr<int64_t>();
+    int64_t* sorted_values = sorted_idx.data_ptr<int64_t>();
+    CUB_CALL(
+        DeviceRadixSort::SortPairs, input_keys, sorted_keys, input_values,
+        sorted_values, num_items, 0, num_bits);
+    return std::make_pair(sorted_array, sorted_idx);
+  } else {
+    CUB_CALL(
+        DeviceRadixSort::SortKeys, input_keys, sorted_keys, num_items, 0,
+        num_bits);
+    return sorted_array;
+  }
+}
+
+template <bool return_original_positions>
+std::conditional_t<
+    return_original_positions, std::pair<torch::Tensor, torch::Tensor>,
+    torch::Tensor>
+Sort(torch::Tensor input, int num_bits) {
+  return AT_DISPATCH_INTEGRAL_TYPES(input.scalar_type(), "SortImpl", ([&] {
+                                      return Sort<return_original_positions>(
+                                          input.data_ptr<scalar_t>(),
+                                          input.size(0), num_bits);
+                                    }));
+}
+
+template torch::Tensor Sort<false>(torch::Tensor input, int num_bits);
+template std::pair<torch::Tensor, torch::Tensor> Sort<true>(
+    torch::Tensor input, int num_bits);
+
+}  //  namespace ops
+}  //  namespace graphbolt

graphbolt/src/cuda/unique_and_compact_impl.cu

@@ -0,0 +1,323 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file cuda/unique_and_compact_impl.cu
+ * @brief Unique and compact operator implementation on CUDA.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <thrust/binary_search.h>
+#include <thrust/functional.h>
+#include <thrust/gather.h>
+#include <thrust/logical.h>
+
+#include <cub/cub.cuh>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+#include "./common.h"
+#include "./extension/unique_and_compact.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+template <typename scalar_t>
+struct EqualityFunc {
+  const scalar_t* sorted_order;
+  const scalar_t* found_locations;
+  const scalar_t* searched_items;
+  __host__ __device__ auto operator()(int64_t i) {
+    return sorted_order[found_locations[i]] == searched_items[i];
+  }
+};
+
+#define DefineCubReductionFunction(cub_reduce_fn, name)           \
+  template <typename scalar_iterator_t>                           \
+  auto name(const scalar_iterator_t input, int64_t size) {        \
+    using scalar_t = std::remove_reference_t<decltype(input[0])>; \
+    cuda::CopyScalar<scalar_t> result;                            \
+    CUB_CALL(cub_reduce_fn, input, result.get(), size);           \
+    return result;                                                \
+  }
+
+DefineCubReductionFunction(DeviceReduce::Max, Max);
+DefineCubReductionFunction(DeviceReduce::Min, Min);
+
+std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
+UniqueAndCompactBatchedSortBased(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor>& unique_dst_ids, int num_bits = 0) {
+  auto allocator = cuda::GetAllocator();
+  auto stream = cuda::GetCurrentStream();
+  auto scalar_type = src_ids.at(0).scalar_type();
+  return AT_DISPATCH_INDEX_TYPES(
+      scalar_type, "unique_and_compact", ([&] {
+        std::vector<index_t*> src_ids_ptr, dst_ids_ptr, unique_dst_ids_ptr;
+        for (std::size_t i = 0; i < src_ids.size(); i++) {
+          src_ids_ptr.emplace_back(src_ids[i].data_ptr<index_t>());
+          dst_ids_ptr.emplace_back(dst_ids[i].data_ptr<index_t>());
+          unique_dst_ids_ptr.emplace_back(
+              unique_dst_ids[i].data_ptr<index_t>());
+        }
+
+        // If num_bits is not given, compute maximum vertex ids to compute
+        // num_bits later to speedup the expensive sort operations.
+        std::vector<cuda::CopyScalar<index_t>> max_id_src;
+        std::vector<cuda::CopyScalar<index_t>> max_id_dst;
+        for (std::size_t i = 0; num_bits == 0 && i < src_ids.size(); i++) {
+          max_id_src.emplace_back(Max(src_ids_ptr[i], src_ids[i].size(0)));
+          max_id_dst.emplace_back(
+              Max(unique_dst_ids_ptr[i], unique_dst_ids[i].size(0)));
+        }
+
+        // Sort the unique_dst_ids tensor.
+        std::vector<torch::Tensor> sorted_unique_dst_ids;
+        std::vector<index_t*> sorted_unique_dst_ids_ptr;
+        for (std::size_t i = 0; i < unique_dst_ids.size(); i++) {
+          sorted_unique_dst_ids.emplace_back(Sort<false>(
+              unique_dst_ids_ptr[i], unique_dst_ids[i].size(0), num_bits));
+          sorted_unique_dst_ids_ptr.emplace_back(
+              sorted_unique_dst_ids[i].data_ptr<index_t>());
+        }
+
+        // Mark dst nodes in the src_ids tensor.
+        std::vector<decltype(allocator.AllocateStorage<bool>(0))> is_dst;
+        for (std::size_t i = 0; i < src_ids.size(); i++) {
+          is_dst.emplace_back(
+              allocator.AllocateStorage<bool>(src_ids[i].size(0)));
+          THRUST_CALL(
+              binary_search, sorted_unique_dst_ids_ptr[i],
+              sorted_unique_dst_ids_ptr[i] + unique_dst_ids[i].size(0),
+              src_ids_ptr[i], src_ids_ptr[i] + src_ids[i].size(0),
+              is_dst[i].get());
+        }
+
+        // Filter the non-dst nodes in the src_ids tensor, hence only_src.
+        std::vector<torch::Tensor> only_src;
+        {
+          std::vector<cuda::CopyScalar<int64_t>> only_src_size;
+          for (std::size_t i = 0; i < src_ids.size(); i++) {
+            only_src.emplace_back(torch::empty(
+                src_ids[i].size(0), sorted_unique_dst_ids[i].options()));
+            auto is_src = thrust::make_transform_iterator(
+                is_dst[i].get(), thrust::logical_not<bool>{});
+            only_src_size.emplace_back(cuda::CopyScalar<int64_t>{});
+            CUB_CALL(
+                DeviceSelect::Flagged, src_ids_ptr[i], is_src,
+                only_src[i].data_ptr<index_t>(), only_src_size[i].get(),
+                src_ids[i].size(0));
+          }
+          stream.synchronize();
+          for (std::size_t i = 0; i < only_src.size(); i++) {
+            only_src[i] =
+                only_src[i].slice(0, 0, static_cast<int64_t>(only_src_size[i]));
+          }
+        }
+
+        // The code block above synchronizes, ensuring safe access to
+        // max_id_src and max_id_dst.
+        if (num_bits == 0) {
+          index_t max_id = 0;
+          for (std::size_t i = 0; i < max_id_src.size(); i++) {
+            max_id = std::max(max_id, static_cast<index_t>(max_id_src[i]));
+            max_id = std::max(max_id, static_cast<index_t>(max_id_dst[i]));
+          }
+          num_bits = cuda::NumberOfBits(1ll + max_id);
+        }
+
+        // Sort the only_src tensor so that we can unique it later.
+        std::vector<torch::Tensor> sorted_only_src;
+        for (auto& only_src_i : only_src) {
+          sorted_only_src.emplace_back(Sort<false>(
+              only_src_i.data_ptr<index_t>(), only_src_i.size(0), num_bits));
+        }
+
+        std::vector<torch::Tensor> unique_only_src;
+        std::vector<index_t*> unique_only_src_ptr;
+
+        std::vector<cuda::CopyScalar<int64_t>> unique_only_src_size;
+        for (std::size_t i = 0; i < src_ids.size(); i++) {
+          // Compute the unique operation on the only_src tensor.
+          unique_only_src.emplace_back(
+              torch::empty(only_src[i].size(0), src_ids[i].options()));
+          unique_only_src_ptr.emplace_back(
+              unique_only_src[i].data_ptr<index_t>());
+          unique_only_src_size.emplace_back(cuda::CopyScalar<int64_t>{});
+          CUB_CALL(
+              DeviceSelect::Unique, sorted_only_src[i].data_ptr<index_t>(),
+              unique_only_src_ptr[i], unique_only_src_size[i].get(),
+              only_src[i].size(0));
+        }
+        stream.synchronize();
+        for (std::size_t i = 0; i < unique_only_src.size(); i++) {
+          unique_only_src[i] = unique_only_src[i].slice(
+              0, 0, static_cast<int64_t>(unique_only_src_size[i]));
+        }
+
+        std::vector<torch::Tensor> real_order;
+        for (std::size_t i = 0; i < unique_dst_ids.size(); i++) {
+          real_order.emplace_back(
+              torch::cat({unique_dst_ids[i], unique_only_src[i]}));
+        }
+        // Sort here so that binary search can be used to lookup new_ids.
+        std::vector<torch::Tensor> sorted_order, new_ids;
+        std::vector<index_t*> sorted_order_ptr;
+        std::vector<int64_t*> new_ids_ptr;
+        for (std::size_t i = 0; i < real_order.size(); i++) {
+          auto [sorted_order_i, new_ids_i] = Sort(real_order[i], num_bits);
+          sorted_order_ptr.emplace_back(sorted_order_i.data_ptr<index_t>());
+          new_ids_ptr.emplace_back(new_ids_i.data_ptr<int64_t>());
+          sorted_order.emplace_back(std::move(sorted_order_i));
+          new_ids.emplace_back(std::move(new_ids_i));
+        }
+        // Holds the found locations of the src and dst ids in the
+        // sorted_order. Later is used to lookup the new ids of the src_ids
+        // and dst_ids tensors.
+        std::vector<decltype(allocator.AllocateStorage<index_t>(0))>
+            new_dst_ids_loc;
+        for (std::size_t i = 0; i < sorted_order.size(); i++) {
+          new_dst_ids_loc.emplace_back(
+              allocator.AllocateStorage<index_t>(dst_ids[i].size(0)));
+          THRUST_CALL(
+              lower_bound, sorted_order_ptr[i],
+              sorted_order_ptr[i] + sorted_order[i].size(0), dst_ids_ptr[i],
+              dst_ids_ptr[i] + dst_ids[i].size(0), new_dst_ids_loc[i].get());
+        }
+
+        std::vector<cuda::CopyScalar<bool>> all_exist;
+        at::cuda::CUDAEvent all_exist_event;
+        bool should_record = false;
+        // Check if unique_dst_ids includes all dst_ids.
+        for (std::size_t i = 0; i < dst_ids.size(); i++) {
+          if (dst_ids[i].size(0) > 0) {
+            thrust::counting_iterator<int64_t> iota(0);
+            auto equal_it = thrust::make_transform_iterator(
+                iota, EqualityFunc<index_t>{
+                          sorted_order_ptr[i], new_dst_ids_loc[i].get(),
+                          dst_ids_ptr[i]});
+            all_exist.emplace_back(Min(equal_it, dst_ids[i].size(0)));
+            should_record = true;
+          } else {
+            all_exist.emplace_back(cuda::CopyScalar<bool>{});
+          }
+        }
+        if (should_record) all_exist_event.record();
+
+        std::vector<decltype(allocator.AllocateStorage<index_t>(0))>
+            new_src_ids_loc;
+        for (std::size_t i = 0; i < sorted_order.size(); i++) {
+          new_src_ids_loc.emplace_back(
+              allocator.AllocateStorage<index_t>(src_ids[i].size(0)));
+          THRUST_CALL(
+              lower_bound, sorted_order_ptr[i],
+              sorted_order_ptr[i] + sorted_order[i].size(0), src_ids_ptr[i],
+              src_ids_ptr[i] + src_ids[i].size(0), new_src_ids_loc[i].get());
+        }
+
+        // Finally, lookup the new compact ids of the src and dst tensors
+        // via gather operations.
+        std::vector<torch::Tensor> new_src_ids;
+        for (std::size_t i = 0; i < src_ids.size(); i++) {
+          new_src_ids.emplace_back(torch::empty_like(src_ids[i]));
+          THRUST_CALL(
+              gather, new_src_ids_loc[i].get(),
+              new_src_ids_loc[i].get() + src_ids[i].size(0),
+              new_ids[i].data_ptr<int64_t>(),
+              new_src_ids[i].data_ptr<index_t>());
+        }
+        // Perform check before we gather for the dst indices.
+        for (std::size_t i = 0; i < dst_ids.size(); i++) {
+          if (dst_ids[i].size(0) > 0) {
+            if (should_record) {
+              all_exist_event.synchronize();
+              should_record = false;
+            }
+            if (!static_cast<bool>(all_exist[i])) {
+              throw std::out_of_range("Some ids not found.");
+            }
+          }
+        }
+        std::vector<torch::Tensor> new_dst_ids;
+        for (std::size_t i = 0; i < dst_ids.size(); i++) {
+          new_dst_ids.emplace_back(torch::empty_like(dst_ids[i]));
+          THRUST_CALL(
+              gather, new_dst_ids_loc[i].get(),
+              new_dst_ids_loc[i].get() + dst_ids[i].size(0),
+              new_ids[i].data_ptr<int64_t>(),
+              new_dst_ids[i].data_ptr<index_t>());
+        }
+        std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>
+            results;
+        for (std::size_t i = 0; i < src_ids.size(); i++) {
+          results.emplace_back(
+              std::move(real_order[i]), std::move(new_src_ids[i]),
+              std::move(new_dst_ids[i]));
+        }
+        return results;
+      }));
+}
+
+std::vector<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+UniqueAndCompactBatched(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor>& unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  if (cuda::compute_capability() >= 70) {
+    // Utilizes a hash table based implementation, the mapped id of a vertex
+    // will be monotonically increasing as the first occurrence index of it in
+    // torch.cat([unique_dst_ids, src_ids]). Thus, it is deterministic.
+    return UniqueAndCompactBatchedHashMapBased(
+        src_ids, dst_ids, unique_dst_ids, rank, world_size);
+  }
+  TORCH_CHECK(
+      world_size <= 1,
+      "Cooperative Minibatching (arXiv:2310.12403) is not supported on "
+      "pre-Volta generation GPUs.");
+  // Utilizes a sort based algorithm, the mapped id of a vertex part of the
+  // src_ids but not part of the unique_dst_ids will be monotonically increasing
+  // as the actual vertex id increases. Thus, it is deterministic.
+  auto results3 =
+      UniqueAndCompactBatchedSortBased(src_ids, dst_ids, unique_dst_ids);
+  std::vector<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+      results4;
+  auto offsets = torch::zeros(
+      2 * results3.size(),
+      c10::TensorOptions().dtype(torch::kInt64).pinned_memory(true));
+  for (const auto& [a, b, c] : results3) {
+    auto d = offsets.slice(0, 0, 2);
+    d.data_ptr<int64_t>()[1] = a.size(0);
+    results4.emplace_back(a, b, c, d);
+    offsets = offsets.slice(0, 2);
+  }
+  return results4;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+UniqueAndCompact(
+    const torch::Tensor src_ids, const torch::Tensor dst_ids,
+    const torch::Tensor unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  return UniqueAndCompactBatched(
+      {src_ids}, {dst_ids}, {unique_dst_ids}, rank, world_size)[0];
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/cuda/utils.h

@@ -0,0 +1,67 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file utils.h
+ * @brief CUDA utilities.
+ */
+
+#ifndef GRAPHBOLT_CUDA_UTILS_H_
+#define GRAPHBOLT_CUDA_UTILS_H_
+
+// The cache line size of GPU.
+constexpr int GPU_CACHE_LINE_SIZE = 128;
+// The max number of threads per block.
+constexpr int CUDA_MAX_NUM_THREADS = 1024;
+
+namespace graphbolt {
+namespace cuda {
+
+/**
+ * @brief Returns the compute capability of the cuda device, e.g. 70 for Volta.
+ */
+inline int compute_capability(
+    int device = cuda::GetCurrentStream().device_index()) {
+  int sm_version;
+  CUDA_RUNTIME_CHECK(cub::SmVersion(sm_version, device));
+  return sm_version / 10;
+};
+
+/**
+ * @brief Calculate the number of threads needed given the size of the dimension
+ * to be processed.
+ *
+ * It finds the largest power of two that is less than or equal to the minimum
+ * of size and CUDA_MAX_NUM_THREADS.
+ */
+inline int FindNumThreads(int size) {
+  int ret = 1;
+  while ((ret << 1) <= std::min(size, CUDA_MAX_NUM_THREADS)) {
+    ret <<= 1;
+  }
+  return ret;
+}
+
+/**
+ * @brief Calculate the smallest number of bits needed to represent a given
+ * range of integers [0, range).
+ */
+template <typename T>
+int NumberOfBits(const T& range) {
+  if (range <= 1) {
+    // ranges of 0 or 1 require no bits to store
+    return 0;
+  }
+
+  int bits = 1;
+  const auto urange = static_cast<std::make_unsigned_t<T>>(range);
+  while (bits < static_cast<int>(sizeof(T) * 8) && (1ull << bits) < urange) {
+    ++bits;
+  }
+
+  return bits;
+}
+
+}  // namespace cuda
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_CUDA_UTILS_H_

graphbolt/src/expand_indptr.cc

@@ -0,0 +1,75 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *  Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ * @file expand_indptr.cc
+ * @brief ExpandIndptr operators.
+ */
+#include <graphbolt/cuda_ops.h>
+#include <torch/autograd.h>
+
+#include "./macro.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+torch::Tensor ExpandIndptr(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> node_ids,
+    torch::optional<int64_t> output_size) {
+  if (utils::is_on_gpu(indptr) &&
+      (!node_ids.has_value() || utils::is_on_gpu(node_ids.value()))) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(c10::DeviceType::CUDA, "ExpandIndptr", {
+      return ExpandIndptrImpl(indptr, dtype, node_ids, output_size);
+    });
+  }
+  if (!node_ids.has_value()) {
+    return torch::repeat_interleave(indptr.diff(), output_size).to(dtype);
+  }
+  return node_ids.value().to(dtype).repeat_interleave(
+      indptr.diff(), 0, output_size);
+}
+
+torch::Tensor IndptrEdgeIds(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> offset,
+    torch::optional<int64_t> output_size) {
+  if (utils::is_on_gpu(indptr) &&
+      (!offset.has_value() || utils::is_on_gpu(offset.value()))) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "IndptrEdgeIds",
+        { return IndptrEdgeIdsImpl(indptr, dtype, offset, output_size); });
+  }
+  TORCH_CHECK(false, "CPU implementation of IndptrEdgeIds is not available.");
+}
+
+TORCH_LIBRARY_IMPL(graphbolt, CPU, m) {
+  m.impl("expand_indptr", &ExpandIndptr);
+}
+
+#ifdef GRAPHBOLT_USE_CUDA
+TORCH_LIBRARY_IMPL(graphbolt, CUDA, m) {
+  m.impl("expand_indptr", &ExpandIndptrImpl);
+}
+#endif
+
+TORCH_LIBRARY_IMPL(graphbolt, Autograd, m) {
+  m.impl("expand_indptr", torch::autograd::autogradNotImplementedFallback());
+}
+
+TORCH_LIBRARY_IMPL(graphbolt, CPU, m) {
+  m.impl("indptr_edge_ids", &IndptrEdgeIds);
+}
+
+#ifdef GRAPHBOLT_USE_CUDA
+TORCH_LIBRARY_IMPL(graphbolt, CUDA, m) {
+  m.impl("indptr_edge_ids", &IndptrEdgeIdsImpl);
+}
+#endif
+
+TORCH_LIBRARY_IMPL(graphbolt, Autograd, m) {
+  m.impl("indptr_edge_ids", torch::autograd::autogradNotImplementedFallback());
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/expand_indptr.h

@@ -0,0 +1,55 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *  Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ * @file expand_indptr.h
+ * @brief ExpandIndptr operators.
+ */
+#ifndef GRAPHBOLT_EXPAND_INDPTR_H_
+#define GRAPHBOLT_EXPAND_INDPTR_H_
+
+#include <torch/script.h>
+
+namespace graphbolt {
+namespace ops {
+
+/**
+ * @brief ExpandIndptr implements conversion from a given indptr offset
+ * tensor to a COO format tensor. If node_ids is not given, it is assumed to be
+ * equal to torch::arange(indptr.size(0) - 1, dtype=dtype).
+ *
+ * @param indptr       The indptr offset tensor.
+ * @param dtype        The dtype of the returned output tensor.
+ * @param node_ids     1D tensor represents the node ids.
+ * @param output_size  Optional, value of indptr[-1]. Passing it eliminates CPU
+ * GPU synchronization.
+ *
+ * @return The resulting tensor.
+ */
+torch::Tensor ExpandIndptr(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> node_ids = torch::nullopt,
+    torch::optional<int64_t> output_size = torch::nullopt);
+
+/**
+ * @brief IndptrEdgeIdsImpl implements conversion from a given indptr offset
+ * tensor to a COO edge ids tensor. For a given indptr [0, 2, 5, 7] and offset
+ * tensor [0, 100, 200], the output will be [0, 1, 100, 101, 102, 201, 202]. If
+ * offset was not provided, the output would be [0, 1, 0, 1, 2, 0, 1].
+ *
+ * @param indptr       The indptr offset tensor.
+ * @param dtype        The dtype of the returned output tensor.
+ * @param offset       The offset tensor.
+ * @param output_size  Optional value of indptr[-1]. Passing it eliminates CPU
+ * GPU synchronization.
+ *
+ * @return The resulting tensor.
+ */
+torch::Tensor IndptrEdgeIds(
+    torch::Tensor indptr, torch::ScalarType dtype,
+    torch::optional<torch::Tensor> offset,
+    torch::optional<int64_t> output_size);
+
+}  // namespace ops
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_EXPAND_INDPTR_H_

graphbolt/src/feature_cache.cc

@@ -0,0 +1,101 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file feature_cache.cc
+ * @brief Feature cache implementation on the CPU.
+ */
+#include "./feature_cache.h"
+
+#include "./index_select.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace storage {
+
+constexpr int kIntGrainSize = 64;
+
+FeatureCache::FeatureCache(
+    const std::vector<int64_t>& shape, torch::ScalarType dtype, bool pin_memory)
+    : tensor_(torch::empty(
+          shape, c10::TensorOptions().dtype(dtype).pinned_memory(pin_memory))) {
+}
+
+torch::Tensor FeatureCache::Query(
+    torch::Tensor positions, torch::Tensor indices, int64_t size) {
+  const bool pin_memory =
+      utils::is_pinned(positions) || utils::is_pinned(indices);
+  std::vector<int64_t> output_shape{
+      tensor_.sizes().begin(), tensor_.sizes().end()};
+  output_shape[0] = size;
+  auto values =
+      torch::empty(output_shape, tensor_.options().pinned_memory(pin_memory));
+  const auto row_bytes = values.slice(0, 0, 1).numel() * values.element_size();
+  auto values_ptr = reinterpret_cast<std::byte*>(values.data_ptr());
+  const auto tensor_ptr = reinterpret_cast<std::byte*>(tensor_.data_ptr());
+  const auto positions_ptr = positions.data_ptr<int64_t>();
+  const auto indices_ptr = indices.data_ptr<int64_t>();
+  graphbolt::parallel_for_each(
+      0, positions.size(0), kIntGrainSize, [&](const int64_t i) {
+        std::memcpy(
+            values_ptr + indices_ptr[i] * row_bytes,
+            tensor_ptr + positions_ptr[i] * row_bytes, row_bytes);
+      });
+  return values;
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> FeatureCache::QueryAsync(
+    torch::Tensor positions, torch::Tensor indices, int64_t size) {
+  return async([=] { return Query(positions, indices, size); });
+}
+
+torch::Tensor FeatureCache::IndexSelect(torch::Tensor positions) {
+  return ops::IndexSelect(tensor_, positions);
+}
+
+void FeatureCache::Replace(torch::Tensor positions, torch::Tensor values) {
+  TORCH_CHECK(positions.size(0) == values.size(0));
+  if (values.numel() == 0) return;
+  const auto row_bytes = values.slice(0, 0, 1).numel() * values.element_size();
+  TORCH_CHECK(
+      row_bytes == tensor_.slice(0, 0, 1).numel() * tensor_.element_size(),
+      "The # bytes of a single row should match the cache's.");
+  auto values_ptr = reinterpret_cast<std::byte*>(values.data_ptr());
+  const auto tensor_ptr = reinterpret_cast<std::byte*>(tensor_.data_ptr());
+  const auto positions_ptr = positions.data_ptr<int64_t>();
+  graphbolt::parallel_for_each(
+      0, positions.size(0), kIntGrainSize, [&](const int64_t i) {
+        const auto position = positions_ptr[i];
+        if (position >= 0) {
+          std::memcpy(
+              tensor_ptr + position * row_bytes, values_ptr + i * row_bytes,
+              row_bytes);
+        }
+      });
+}
+
+c10::intrusive_ptr<Future<void>> FeatureCache::ReplaceAsync(
+    torch::Tensor positions, torch::Tensor values) {
+  return async([=] { return Replace(positions, values); });
+}
+
+c10::intrusive_ptr<FeatureCache> FeatureCache::Create(
+    const std::vector<int64_t>& shape, torch::ScalarType dtype,
+    bool pin_memory) {
+  return c10::make_intrusive<FeatureCache>(shape, dtype, pin_memory);
+}
+
+}  // namespace storage
+}  // namespace graphbolt

graphbolt/src/feature_cache.h

@@ -0,0 +1,101 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file feature_cache.h
+ * @brief Feature cache implementation on the CPU.
+ */
+#ifndef GRAPHBOLT_FEATURE_CACHE_H_
+#define GRAPHBOLT_FEATURE_CACHE_H_
+
+#include <graphbolt/async.h>
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+namespace graphbolt {
+namespace storage {
+
+struct FeatureCache : public torch::CustomClassHolder {
+  /**
+   * @brief Constructor for the FeatureCache struct.
+   *
+   * @param shape The shape of the cache.
+   * @param dtype The dtype of elements stored in the cache.
+   * @param pin_memory Whether to pin the memory of the cache storage tensor.
+   */
+  FeatureCache(
+      const std::vector<int64_t>& shape, torch::ScalarType dtype,
+      bool pin_memory);
+
+  bool IsPinned() const { return tensor_.is_pinned(); }
+
+  int64_t NumBytes() const { return tensor_.numel() * tensor_.element_size(); }
+
+  /**
+   * @brief The cache query function. Allocates an empty tensor `values` with
+   * size as the first dimension and runs
+   * values[indices[:positions.size(0)]] = cache_tensor[positions] before
+   * returning it.
+   *
+   * @param positions The positions of the queried items.
+   * @param indices The indices of the queried items among the original keys.
+   * Only the first portion corresponding to the provided positions tensor is
+   * used, e.g. indices[:positions.size(0)].
+   * @param size The size of the original keys, hence the first dimension of
+   * the output shape.
+   *
+   * @return The values tensor is returned. Its memory is pinned if pin_memory
+   * is true.
+   */
+  torch::Tensor Query(
+      torch::Tensor positions, torch::Tensor indices, int64_t size);
+
+  c10::intrusive_ptr<Future<torch::Tensor>> QueryAsync(
+      torch::Tensor positions, torch::Tensor indices, int64_t size);
+
+  /**
+   * @brief The cache tensor index_select returns cache_tensor[positions].
+   *
+   * @param positions The positions of the queried items.
+   *
+   * @return The values tensor is returned on the same device as positions.
+   */
+  torch::Tensor IndexSelect(torch::Tensor positions);
+
+  /**
+   * @brief The cache replace function.
+   *
+   * @param positions The positions to replace in the cache.
+   * @param values The values to be inserted into the cache.
+   */
+  void Replace(torch::Tensor positions, torch::Tensor values);
+
+  c10::intrusive_ptr<Future<void>> ReplaceAsync(
+      torch::Tensor positions, torch::Tensor values);
+
+  static c10::intrusive_ptr<FeatureCache> Create(
+      const std::vector<int64_t>& shape, torch::ScalarType dtype,
+      bool pin_memory);
+
+ private:
+  torch::Tensor tensor_;
+};
+
+}  // namespace storage
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_FEATURE_CACHE_H_

graphbolt/src/index_select.cc

@@ -0,0 +1,225 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file index_select.cc
+ * @brief Index select operators.
+ */
+#include "./index_select.h"
+
+#include <graphbolt/cuda_ops.h>
+#include <graphbolt/fused_csc_sampling_graph.h>
+
+#include <cstring>
+#include <numeric>
+
+#include "./macro.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace ops {
+
+constexpr int kIntGrainSize = 64;
+
+torch::Tensor IndexSelect(torch::Tensor input, torch::Tensor index) {
+  if (utils::is_on_gpu(index)) {
+    if (input.is_pinned()) {
+      GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+          c10::DeviceType::CUDA, "UVAIndexSelect",
+          { return UVAIndexSelectImpl(input, index); });
+    } else {
+      return torch::index_select(input, 0, index);
+    }
+  }
+  auto output_shape = input.sizes().vec();
+  output_shape[0] = index.numel();
+  auto result = torch::empty(
+      output_shape, index.options()
+                        .dtype(input.dtype())
+                        .pinned_memory(utils::is_pinned(index)));
+  auto result_ptr = reinterpret_cast<std::byte*>(result.data_ptr());
+  const auto input_ptr = reinterpret_cast<std::byte*>(input.data_ptr());
+  const auto row_bytes = input.slice(0, 0, 1).numel() * input.element_size();
+  const auto stride = input.stride(0) * input.element_size();
+  const auto num_input_rows = input.size(0);
+  AT_DISPATCH_INDEX_TYPES(
+      index.scalar_type(), "IndexSelect::index::scalar_type()", ([&] {
+        const auto index_ptr = index.data_ptr<index_t>();
+        graphbolt::parallel_for(
+            0, index.size(0), kIntGrainSize, [&](int64_t begin, int64_t end) {
+              for (int64_t i = begin; i < end; i++) {
+                auto idx = index_ptr[i];
+                if (idx < 0) idx += num_input_rows;
+                if (idx < 0 || idx >= num_input_rows) {
+                  // Throw IndexError via torch.
+                  idx += input[num_input_rows].item<index_t>();
+                }
+                std::memcpy(
+                    result_ptr + i * row_bytes, input_ptr + idx * stride,
+                    row_bytes);
+              }
+            });
+      }));
+  return result;
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> IndexSelectAsync(
+    torch::Tensor input, torch::Tensor index) {
+  TORCH_CHECK(!utils::is_on_gpu(index) && !utils::is_on_gpu(input));
+  return async([=] { return IndexSelect(input, index); });
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> ScatterAsync(
+    torch::Tensor input, torch::Tensor index, torch::Tensor src) {
+  TORCH_CHECK(
+      !utils::is_on_gpu(input) && !utils::is_on_gpu(index) &&
+      !utils::is_on_gpu(src));
+  TORCH_CHECK(index.sizes().size() == 1, "index tensor needs to be 1d.");
+  for (size_t i = 1; i < input.sizes().size(); i++) {
+    TORCH_CHECK(
+        input.size(i) == src.size(i),
+        "dimension mismatch between input and src at ", i,
+        "th dimension: ", input.size(i), " != ", src.size(i), ".");
+  }
+  return async([=] {
+    const auto row_bytes = src.slice(0, 0, 1).numel() * src.element_size();
+    const auto src_ptr = reinterpret_cast<std::byte*>(src.data_ptr());
+    auto input_ptr = reinterpret_cast<std::byte*>(input.data_ptr());
+    AT_DISPATCH_INDEX_TYPES(
+        index.scalar_type(), "ScatterAsync::index::scalar_type()", ([&] {
+          const auto index_ptr = index.data_ptr<index_t>();
+          graphbolt::parallel_for(
+              0, index.size(0), kIntGrainSize, [&](int64_t begin, int64_t end) {
+                for (int64_t i = begin; i < end; i++) {
+                  std::memcpy(
+                      input_ptr + index_ptr[i] * row_bytes,
+                      src_ptr + i * row_bytes, row_bytes);
+                }
+              });
+        }));
+    return input;
+  });
+}
+
+std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSC(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes,
+    torch::optional<int64_t> output_size) {
+  TORCH_CHECK(
+      indices.sizes().size() == 1, "IndexSelectCSC only supports 1d tensors");
+  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr) &&
+      utils::is_accessible_from_gpu(indices)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "IndexSelectCSCImpl",
+        { return IndexSelectCSCImpl(indptr, indices, nodes, output_size); });
+  }
+  auto [output_indptr, results] = IndexSelectCSCBatched(
+      indptr, std::vector{indices}, nodes, false, output_size);
+  return std::make_tuple(output_indptr, results.at(0));
+}
+
+std::tuple<torch::Tensor, std::vector<torch::Tensor>> IndexSelectCSCBatched(
+    torch::Tensor indptr, std::vector<torch::Tensor> indices_list,
+    torch::Tensor nodes, bool with_edge_ids,
+    torch::optional<int64_t> output_size) {
+  for (auto& indices : indices_list) {
+    TORCH_CHECK(
+        indices.sizes().size() == 1,
+        "IndexSelectCSCBatched only supports 1d tensors");
+  }
+  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr) &&
+      utils::are_accessible_from_gpu(indices_list)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "IndexSelectCSCImpl", {
+          return IndexSelectCSCBatchedImpl(
+              indptr, indices_list, nodes, with_edge_ids, output_size);
+        });
+  }
+  constexpr int kDefaultGrainSize = 128;
+  const auto num_nodes = nodes.size(0);
+  torch::Tensor output_indptr = torch::empty(
+      {num_nodes + 1}, nodes.options().dtype(indptr.scalar_type()));
+  std::vector<torch::Tensor> results;
+  torch::optional<torch::Tensor> edge_ids;
+  AT_DISPATCH_INDEX_TYPES(
+      indptr.scalar_type(), "IndexSelectCSCBatched::indptr", ([&] {
+        using indptr_t = index_t;
+        const auto indptr_data = indptr.data_ptr<indptr_t>();
+        auto out_indptr_data = output_indptr.data_ptr<indptr_t>();
+        out_indptr_data[0] = 0;
+        AT_DISPATCH_INDEX_TYPES(
+            nodes.scalar_type(), "IndexSelectCSCBatched::nodes", ([&] {
+              const auto nodes_data = nodes.data_ptr<index_t>();
+              torch::parallel_for(
+                  0, num_nodes, kDefaultGrainSize,
+                  [&](int64_t begin, int64_t end) {
+                    for (int64_t i = begin; i < end; i++) {
+                      const auto node_id = nodes_data[i];
+                      const auto degree =
+                          indptr_data[node_id + 1] - indptr_data[node_id];
+                      out_indptr_data[i + 1] = degree;
+                    }
+                  });
+              output_indptr = output_indptr.cumsum(0, indptr.scalar_type());
+              out_indptr_data = output_indptr.data_ptr<indptr_t>();
+              TORCH_CHECK(
+                  !output_size.has_value() ||
+                      out_indptr_data[num_nodes] == *output_size,
+                  "An incorrect output_size argument was provided.");
+              output_size = out_indptr_data[num_nodes];
+              for (const auto& indices : indices_list) {
+                results.push_back(torch::empty(
+                    *output_size,
+                    nodes.options().dtype(indices.scalar_type())));
+              }
+              if (with_edge_ids) {
+                edge_ids = torch::empty(
+                    *output_size, nodes.options().dtype(indptr.scalar_type()));
+              }
+              torch::parallel_for(
+                  0, num_nodes, kDefaultGrainSize,
+                  [&](int64_t begin, int64_t end) {
+                    for (int64_t i = begin; i < end; i++) {
+                      const auto output_offset = out_indptr_data[i];
+                      const auto numel = out_indptr_data[i + 1] - output_offset;
+                      const auto input_offset = indptr_data[nodes_data[i]];
+                      for (size_t tensor_id = 0;
+                           tensor_id < indices_list.size(); tensor_id++) {
+                        auto output = reinterpret_cast<std::byte*>(
+                            results[tensor_id].data_ptr());
+                        const auto input = reinterpret_cast<std::byte*>(
+                            indices_list[tensor_id].data_ptr());
+                        const auto element_size =
+                            indices_list[tensor_id].element_size();
+                        std::memcpy(
+                            output + output_offset * element_size,
+                            input + input_offset * element_size,
+                            element_size * numel);
+                      }
+                      if (edge_ids.has_value()) {
+                        auto output = edge_ids->data_ptr<indptr_t>();
+                        std::iota(
+                            output + output_offset,
+                            output + output_offset + numel, input_offset);
+                      }
+                    }
+                  });
+            }));
+      }));
+  if (edge_ids) results.push_back(*edge_ids);
+  return std::make_tuple(output_indptr, results);
+}
+
+c10::intrusive_ptr<
+    Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>>
+IndexSelectCSCBatchedAsync(
+    torch::Tensor indptr, std::vector<torch::Tensor> indices_list,
+    torch::Tensor nodes, bool with_edge_ids,
+    torch::optional<int64_t> output_size) {
+  return async(
+      [=] {
+        return IndexSelectCSCBatched(
+            indptr, indices_list, nodes, with_edge_ids, output_size);
+      },
+      utils::is_on_gpu(nodes));
+}
+
+}  // namespace ops
+}  // namespace graphbolt

graphbolt/src/index_select.h

@@ -0,0 +1,105 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file index_select.h
+ * @brief Index select operators.
+ */
+#ifndef GRAPHBOLT_INDEX_SELECT_H_
+#define GRAPHBOLT_INDEX_SELECT_H_
+
+#include <graphbolt/async.h>
+#include <torch/script.h>
+
+namespace graphbolt {
+namespace ops {
+
+/**
+ * @brief Select columns for a sparse matrix in a CSC format according to nodes
+ * tensor.
+ *
+ * NOTE:
+ * 1. The shape of all tensors must be 1-D.
+ * 2. If indices is on pinned memory and nodes is on pinned memory or GPU
+ * memory, then UVAIndexSelectCSCImpl will be called. If indices is on GPU
+ * memory, then IndexSelectCSCImpl will be called. Otherwise,
+ * FusedCSCSamplingGraph::InSubgraph will be called.
+ *
+ * @param indptr Indptr tensor containing offsets with shape (N,).
+ * @param indices Indices tensor with edge information of shape (indptr[N],).
+ * @param nodes Nodes tensor with shape (M,).
+ * @param output_size The total number of edges being copied.
+ * @return (torch::Tensor, torch::Tensor) Output indptr and indices tensors of
+ * shapes (M + 1,) and ((indptr[nodes + 1] - indptr[nodes]).sum(),).
+ */
+std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSC(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes,
+    torch::optional<int64_t> output_size = torch::nullopt);
+
+/**
+ * @brief Select rows from input tensor according to index tensor.
+ *
+ * NOTE:
+ * 1. The shape of input tensor can be multi-dimensional, but the index tensor
+ * must be 1-D.
+ * 2. If input is on pinned memory and index is on pinned memory or GPU memory,
+ * then UVAIndexSelectImpl will be called. Otherwise, torch::index_select will
+ * be called.
+ *
+ * @param input Input tensor with shape (N, ...).
+ * @param index Index tensor with shape (M,).
+ * @return torch::Tensor Output tensor with shape (M, ...).
+ */
+torch::Tensor IndexSelect(torch::Tensor input, torch::Tensor index);
+
+/**
+ * @brief The async version of IndexSelect, available for only CPU tensors.
+ *
+ * @return Returns a future containing a torch::Tensor.
+ */
+c10::intrusive_ptr<Future<torch::Tensor>> IndexSelectAsync(
+    torch::Tensor input, torch::Tensor index);
+
+/**
+ * @brief The async version of operation input[index] = src.
+ * @param input The input tensor.
+ * @param index The index tensor into input.
+ * @param src The src tensor being assigned into input.
+ *
+ * @return Returns a future containing input, a torch::Tensor.
+ */
+c10::intrusive_ptr<Future<torch::Tensor>> ScatterAsync(
+    torch::Tensor input, torch::Tensor index, torch::Tensor src);
+
+/**
+ * @brief Select columns for a sparse matrix in a CSC format according to nodes
+ * tensor.
+ *
+ * NOTE: The shape of all tensors must be 1-D.
+ *
+ * @param indptr Indptr tensor containing offsets with shape (N,).
+ * @param indices_list Vector of indices tensor with edge information of shape
+ * (indptr[N],).
+ * @param nodes Nodes tensor with shape (M,).
+ * @param with_edge_ids Whether to return edge ids tensor corresponding to
+ * sliced edges as the last element of the output.
+ * @param output_size The total number of edges being copied.
+ *
+ * @return (torch::Tensor, std::vector<torch::Tensor>) Output indptr and vector
+ * of indices tensors of shapes (M + 1,) and ((indptr[nodes + 1] -
+ * indptr[nodes]).sum(),).
+ */
+std::tuple<torch::Tensor, std::vector<torch::Tensor>> IndexSelectCSCBatched(
+    torch::Tensor indptr, std::vector<torch::Tensor> indices_list,
+    torch::Tensor nodes, bool with_edge_ids,
+    torch::optional<int64_t> output_size);
+
+c10::intrusive_ptr<
+    Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>>
+IndexSelectCSCBatchedAsync(
+    torch::Tensor indptr, std::vector<torch::Tensor> indices_list,
+    torch::Tensor nodes, bool with_edge_ids,
+    torch::optional<int64_t> output_size);
+
+}  // namespace ops
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_INDEX_SELECT_H_

graphbolt/src/io_uring.cc

@@ -0,0 +1,82 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file io_uring.cc
+ * @brief io_uring related functions.
+ */
+#include "./io_uring.h"
+
+#ifdef HAVE_LIBRARY_LIBURING
+
+#include <errno.h>
+#include <liburing.h>
+#include <liburing/io_uring.h>
+#include <stddef.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+#include <memory>
+#include <mutex>
+
+struct io_uring_probe_destroyer {
+  void operator()(struct io_uring_probe* p) {
+    if (p) io_uring_free_probe(p);
+  }
+};
+#endif
+
+namespace graphbolt {
+namespace io_uring {
+
+bool IsAvailable() {
+#ifdef HAVE_LIBRARY_LIBURING
+  /** @brief The cached value of whether io_uring is available. */
+  static bool cached_is_available;
+
+  /** @brief Ensure cached_is_available is initialized once and thread-safe. */
+  static std::once_flag initialization_flag;
+
+  std::call_once(initialization_flag, []() {
+    // https://unix.stackexchange.com/a/596284/314554
+    cached_is_available =
+        !(syscall(
+              __NR_io_uring_register, 0, IORING_UNREGISTER_BUFFERS, NULL, 0) &&
+          errno == ENOSYS);
+
+    std::unique_ptr<struct io_uring_probe, io_uring_probe_destroyer> probe(
+        io_uring_get_probe(), io_uring_probe_destroyer());
+    if (probe.get()) {
+      cached_is_available =
+          cached_is_available &&
+          io_uring_opcode_supported(probe.get(), IORING_OP_READ);
+      cached_is_available =
+          cached_is_available &&
+          io_uring_opcode_supported(probe.get(), IORING_OP_READV);
+    } else {
+      cached_is_available = false;
+    }
+  });
+
+  return cached_is_available;
+#else
+  return false;
+#endif
+}
+
+void SetNumThreads(int64_t count) { num_threads = count; }
+
+}  // namespace io_uring
+}  // namespace graphbolt

graphbolt/src/io_uring.h

@@ -0,0 +1,42 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file io_uring.h
+ * @brief io_uring related functions.
+ */
+#ifndef GRAPHBOLT_IO_URING_H_
+#define GRAPHBOLT_IO_URING_H_
+
+#include <cstdint>
+#include <optional>
+
+namespace graphbolt {
+namespace io_uring {
+
+bool IsAvailable();
+
+/** @brief Set a limit on # background io_uring threads. */
+inline std::optional<int64_t> num_threads;
+
+/**
+ * @brief Set the number of background io_uring threads.
+ */
+void SetNumThreads(int64_t count);
+
+}  // namespace io_uring
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_IO_URING_H_

graphbolt/src/isin.cc

@@ -0,0 +1,77 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file isin.cc
+ * @brief Isin op.
+ */
+
+#include <graphbolt/cuda_ops.h>
+#include <graphbolt/isin.h>
+
+#include "./macro.h"
+#include "./utils.h"
+
+namespace {
+static constexpr int kSearchGrainSize = 4096;
+}  // namespace
+
+namespace graphbolt {
+namespace sampling {
+
+torch::Tensor IsInCPU(
+    const torch::Tensor& elements, const torch::Tensor& test_elements) {
+  torch::Tensor sorted_test_elements;
+  std::tie(sorted_test_elements, std::ignore) = test_elements.sort(
+      /*stable=*/false, /*dim=*/0, /*descending=*/false);
+  torch::Tensor result = torch::empty_like(elements, torch::kBool);
+  size_t num_test_elements = test_elements.size(0);
+  size_t num_elements = elements.size(0);
+
+  AT_DISPATCH_INTEGRAL_TYPES(
+      elements.scalar_type(), "IsInOperation", ([&] {
+        const scalar_t* elements_ptr = elements.data_ptr<scalar_t>();
+        const scalar_t* sorted_test_elements_ptr =
+            sorted_test_elements.data_ptr<scalar_t>();
+        bool* result_ptr = result.data_ptr<bool>();
+        torch::parallel_for(
+            0, num_elements, kSearchGrainSize, [&](size_t start, size_t end) {
+              for (auto i = start; i < end; i++) {
+                result_ptr[i] = std::binary_search(
+                    sorted_test_elements_ptr,
+                    sorted_test_elements_ptr + num_test_elements,
+                    elements_ptr[i]);
+              }
+            });
+      }));
+  return result;
+}
+
+torch::Tensor IsIn(
+    const torch::Tensor& elements, const torch::Tensor& test_elements) {
+  if (utils::is_on_gpu(elements) && utils::is_on_gpu(test_elements)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "IsInOperation",
+        { return ops::IsIn(elements, test_elements); });
+  } else {
+    return IsInCPU(elements, test_elements);
+  }
+}
+
+torch::Tensor IsNotInIndex(
+    const torch::Tensor& elements, const torch::Tensor& test_elements) {
+  auto mask = IsIn(elements, test_elements);
+  if (utils::is_on_gpu(mask)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "NonzeroOperation",
+        { return ops::Nonzero(mask, true); });
+  }
+  return torch::nonzero(torch::logical_not(mask)).squeeze(1);
+}
+
+c10::intrusive_ptr<Future<torch::Tensor>> IsNotInIndexAsync(
+    const torch::Tensor& elements, const torch::Tensor& test_elements) {
+  return async([=] { return IsNotInIndex(elements, test_elements); });
+}
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/macro.h

@@ -0,0 +1,40 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file macro.h
+ * @brief Graphbolt macros.
+ */
+
+#ifndef GRAPHBOLT_MACRO_H_
+#define GRAPHBOLT_MACRO_H_
+
+#include <torch/script.h>
+
+namespace graphbolt {
+
+// Dispatch operator implementation function to CUDA device only.
+#ifdef GRAPHBOLT_USE_CUDA
+#define GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(device_type, name, ...) \
+  if (device_type == c10::DeviceType::CUDA) {                       \
+    [[maybe_unused]] auto XPU = c10::DeviceType::CUDA;              \
+    __VA_ARGS__                                                     \
+  } else {                                                          \
+    TORCH_CHECK(false, name, " is only available on CUDA device."); \
+  }
+#else
+#define GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(device_type, name, ...) \
+  TORCH_CHECK(false, name, " is only available on CUDA device.");
+#endif
+
+// This includes all integer, float and boolean types.
+#define GRAPHBOLT_DISPATCH_CASE_ALL_TYPES(...)            \
+  AT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__)                 \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)     \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Bool, __VA_ARGS__)
+
+#define GRAPHBOLT_DISPATCH_ALL_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, GRAPHBOLT_DISPATCH_CASE_ALL_TYPES(__VA_ARGS__))
+
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_MACRO_H_

graphbolt/src/partitioned_cache_policy.cc

@@ -0,0 +1,523 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file partitioned_cache_policy.cc
+ * @brief Partitioned cache policy implementation on the CPU.
+ */
+#include "./partitioned_cache_policy.h"
+
+#include <algorithm>
+#include <limits>
+#include <numeric>
+
+#include "./utils.h"
+
+namespace graphbolt {
+namespace storage {
+
+constexpr int kIntGrainSize = 256;
+
+torch::Tensor AddOffset(torch::Tensor keys, int64_t offset) {
+  if (offset == 0) return keys;
+  auto output = torch::empty_like(
+      keys, keys.options().pinned_memory(utils::is_pinned(keys)));
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "AddOffset", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        auto output_ptr = output.data_ptr<index_t>();
+        graphbolt::parallel_for_each(
+            0, keys.numel(), kIntGrainSize, [&](int64_t i) {
+              const auto result = keys_ptr[i] + offset;
+              if constexpr (!std::is_same_v<index_t, int64_t>) {
+                TORCH_CHECK(
+                    std::numeric_limits<index_t>::min() <= result &&
+                    result <= std::numeric_limits<index_t>::max());
+              }
+              output_ptr[i] = static_cast<index_t>(result);
+            });
+      }));
+  return output;
+}
+
+template <typename CachePolicy>
+PartitionedCachePolicy::PartitionedCachePolicy(
+    CachePolicy, int64_t capacity, int64_t num_partitions)
+    : capacity_(capacity) {
+  TORCH_CHECK(num_partitions >= 1, "# partitions need to be positive.");
+  for (int64_t i = 0; i < num_partitions; i++) {
+    const auto begin = i * capacity / num_partitions;
+    const auto end = (i + 1) * capacity / num_partitions;
+    policies_.emplace_back(std::make_unique<CachePolicy>(end - begin));
+  }
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+PartitionedCachePolicy::Partition(torch::Tensor keys) {
+  const int64_t num_parts = policies_.size();
+  torch::Tensor offsets = torch::empty(
+      num_parts * num_parts + 1, keys.options().dtype(torch::kInt64));
+  auto offsets_ptr = offsets.data_ptr<int64_t>();
+  std::fill_n(offsets_ptr, offsets.size(0), int64_t{});
+  auto indices = torch::empty_like(keys, keys.options().dtype(torch::kInt64));
+  auto part_id = torch::empty_like(keys, keys.options().dtype(torch::kInt32));
+  const auto num_keys = keys.size(0);
+  auto part_id_ptr = part_id.data_ptr<int32_t>();
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "PartitionedCachePolicy::partition", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        namespace gb = graphbolt;
+        gb::parallel_for_each(0, num_parts, 1, [&](int64_t tid) {
+          const auto begin = tid * num_keys / num_parts;
+          const auto end = (tid + 1) * num_keys / num_parts;
+          for (int64_t i = begin; i < end; i++) {
+            const auto part_id = PartAssignment(keys_ptr[i]);
+            offsets_ptr[tid * num_parts + part_id]++;
+            part_id_ptr[i] = part_id;
+          }
+        });
+      }));
+
+  // Transpose the offsets tensor, take cumsum and transpose back.
+  auto offsets_permuted = torch::empty_like(offsets);
+  auto offsets_permuted_ptr = offsets_permuted.data_ptr<int64_t>();
+  graphbolt::parallel_for_each(
+      0, num_parts * num_parts, kIntGrainSize, [&](int64_t i) {
+        const auto part_id = i % num_parts;
+        const auto tid = i / num_parts;
+        // + 1 so that we have exclusive_scan after torch.cumsum().
+        offsets_permuted_ptr[part_id * num_parts + tid + 1] = offsets_ptr[i];
+      });
+  offsets_permuted_ptr[0] = 0;
+  // offsets = offsets_permuted.cumsum(0); @TODO implement this in parallel.
+  std::inclusive_scan(
+      offsets_permuted_ptr, offsets_permuted_ptr + num_parts * num_parts + 1,
+      offsets_ptr);
+  offsets_ptr = offsets.data_ptr<int64_t>();
+  graphbolt::parallel_for_each(
+      0, num_parts * num_parts, kIntGrainSize, [&](int64_t i) {
+        const auto part_id = i % num_parts;
+        const auto tid = i / num_parts;
+        offsets_permuted_ptr[i] = offsets_ptr[part_id * num_parts + tid];
+      });
+  auto indices_ptr = indices.data_ptr<int64_t>();
+  auto permuted_keys = torch::empty_like(keys);
+  auto offsets_sliced = torch::empty(num_parts + 1, offsets.options());
+  auto offsets_sliced_ptr = offsets_sliced.data_ptr<int64_t>();
+  offsets_sliced_ptr[0] = 0;
+  AT_DISPATCH_INDEX_TYPES(
+      keys.scalar_type(), "PartitionedCachePolicy::partition", ([&] {
+        auto keys_ptr = keys.data_ptr<index_t>();
+        auto permuted_keys_ptr = permuted_keys.data_ptr<index_t>();
+        namespace gb = graphbolt;
+        gb::parallel_for_each(0, num_parts, 1, [&](int64_t tid) {
+          const auto begin = tid * num_keys / num_parts;
+          const auto end = (tid + 1) * num_keys / num_parts;
+          for (int64_t i = begin; i < end; i++) {
+            const auto part_id = part_id_ptr[i];
+            auto& offset = offsets_permuted_ptr[tid * num_parts + part_id];
+            indices_ptr[offset] = i;
+            permuted_keys_ptr[offset++] = keys_ptr[i];
+          }
+          offsets_sliced_ptr[tid + 1] = offsets_ptr[(tid + 1) * num_parts];
+        });
+      }));
+  return {offsets_sliced, indices, permuted_keys};
+}
+
+std::tuple<
+    torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+    torch::Tensor>
+PartitionedCachePolicy::Query(torch::Tensor keys, const int64_t offset) {
+  keys = AddOffset(keys, offset);
+  if (policies_.size() == 1) {
+    std::lock_guard lock(mtx_);
+    auto [positions, output_indices, missing_keys, found_pointers] =
+        policies_[0]->Query(keys);
+    auto found_and_missing_offsets = torch::empty(4, found_pointers.options());
+    auto found_and_missing_offsets_ptr =
+        found_and_missing_offsets.data_ptr<int64_t>();
+    // Found offsets part.
+    found_and_missing_offsets_ptr[0] = 0;
+    found_and_missing_offsets_ptr[1] = found_pointers.size(0);
+    // Missing offsets part.
+    found_and_missing_offsets_ptr[2] = 0;
+    found_and_missing_offsets_ptr[3] = missing_keys.size(0);
+    auto found_offsets = found_and_missing_offsets.slice(0, 0, 2);
+    auto missing_offsets = found_and_missing_offsets.slice(0, 2);
+    missing_keys = AddOffset(missing_keys, -offset);
+    return {positions,      output_indices, missing_keys,
+            found_pointers, found_offsets,  missing_offsets};
+  };
+  torch::Tensor offsets, indices, permuted_keys;
+  std::tie(offsets, indices, permuted_keys) = Partition(keys);
+  auto offsets_ptr = offsets.data_ptr<int64_t>();
+  auto indices_ptr = indices.data_ptr<int64_t>();
+  std::vector<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+      results(policies_.size());
+  torch::Tensor result_offsets_tensor =
+      torch::empty(policies_.size() * 2 + 1, offsets.options());
+  auto result_offsets = result_offsets_tensor.data_ptr<int64_t>();
+  namespace gb = graphbolt;
+  {
+    std::lock_guard lock(mtx_);
+    gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+      const auto begin = offsets_ptr[tid];
+      const auto end = offsets_ptr[tid + 1];
+      results[tid] =
+          policies_.at(tid)->Query(permuted_keys.slice(0, begin, end));
+      result_offsets[tid] = std::get<0>(results[tid]).size(0);
+      result_offsets[tid + policies_.size()] =
+          std::get<2>(results[tid]).size(0);
+    });
+  }
+  std::exclusive_scan(
+      result_offsets, result_offsets + result_offsets_tensor.size(0),
+      result_offsets, 0);
+  torch::Tensor positions = torch::empty(
+      result_offsets[policies_.size()],
+      std::get<0>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor output_indices = torch::empty_like(
+      indices, indices.options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor missing_keys = torch::empty(
+      indices.size(0) - positions.size(0),
+      std::get<2>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor found_pointers = torch::empty(
+      positions.size(0),
+      std::get<3>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  auto missing_offsets =
+      torch::empty(policies_.size() + 1, result_offsets_tensor.options());
+  auto output_indices_ptr = output_indices.data_ptr<int64_t>();
+  auto missing_offsets_ptr = missing_offsets.data_ptr<int64_t>();
+  missing_offsets_ptr[0] = 0;
+  gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+    auto out_index_ptr = indices_ptr + offsets_ptr[tid];
+    auto begin = result_offsets[tid];
+    auto end = result_offsets[tid + 1];
+    const auto num_selected = end - begin;
+    auto indices_ptr = std::get<1>(results[tid]).data_ptr<int64_t>();
+    for (int64_t i = 0; i < num_selected; i++) {
+      output_indices_ptr[begin + i] = out_index_ptr[indices_ptr[i]];
+    }
+    auto selected_positions_ptr = std::get<0>(results[tid]).data_ptr<int64_t>();
+    std::transform(
+        selected_positions_ptr, selected_positions_ptr + num_selected,
+        positions.data_ptr<int64_t>() + begin,
+        [off = tid * capacity_ / policies_.size()](auto x) { return x + off; });
+    auto selected_pointers_ptr = std::get<3>(results[tid]).data_ptr<int64_t>();
+    std::copy(
+        selected_pointers_ptr, selected_pointers_ptr + num_selected,
+        found_pointers.data_ptr<int64_t>() + begin);
+    begin = result_offsets[policies_.size() + tid];
+    end = result_offsets[policies_.size() + tid + 1];
+    missing_offsets[tid + 1] = end - result_offsets[policies_.size()];
+    const auto num_missing = end - begin;
+    for (int64_t i = 0; i < num_missing; i++) {
+      output_indices_ptr[begin + i] =
+          out_index_ptr[indices_ptr[i + num_selected]];
+    }
+    std::memcpy(
+        reinterpret_cast<std::byte*>(missing_keys.data_ptr()) +
+            (begin - positions.size(0)) * missing_keys.element_size(),
+        std::get<2>(results[tid]).data_ptr(),
+        num_missing * missing_keys.element_size());
+  });
+  auto found_offsets = result_offsets_tensor.slice(0, 0, policies_.size() + 1);
+  missing_keys = AddOffset(missing_keys, -offset);
+  return std::make_tuple(
+      positions, output_indices, missing_keys, found_pointers, found_offsets,
+      missing_offsets);
+}
+
+c10::intrusive_ptr<Future<std::vector<torch::Tensor>>>
+PartitionedCachePolicy::QueryAsync(torch::Tensor keys, const int64_t offset) {
+  return async([=] {
+    auto
+        [positions, output_indices, missing_keys, found_pointers, found_offsets,
+         missing_offsets] = Query(keys, offset);
+    return std::vector{positions,      output_indices, missing_keys,
+                       found_pointers, found_offsets,  missing_offsets};
+  });
+}
+
+std::tuple<
+    torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+    torch::Tensor>
+PartitionedCachePolicy::QueryAndReplace(
+    torch::Tensor keys, const int64_t offset) {
+  keys = AddOffset(keys, offset);
+  if (policies_.size() == 1) {
+    std::lock_guard lock(mtx_);
+    auto [positions, output_indices, pointers, missing_keys] =
+        policies_[0]->QueryAndReplace(keys);
+    auto found_and_missing_offsets = torch::empty(4, pointers.options());
+    auto found_and_missing_offsets_ptr =
+        found_and_missing_offsets.data_ptr<int64_t>();
+    // Found offsets part.
+    found_and_missing_offsets_ptr[0] = 0;
+    found_and_missing_offsets_ptr[1] = keys.size(0) - missing_keys.size(0);
+    // Missing offsets part.
+    found_and_missing_offsets_ptr[2] = 0;
+    found_and_missing_offsets_ptr[3] = missing_keys.size(0);
+    auto found_offsets = found_and_missing_offsets.slice(0, 0, 2);
+    auto missing_offsets = found_and_missing_offsets.slice(0, 2);
+    missing_keys = AddOffset(missing_keys, -offset);
+    return {positions,    output_indices, pointers,
+            missing_keys, found_offsets,  missing_offsets};
+  }
+  torch::Tensor offsets, indices, permuted_keys;
+  std::tie(offsets, indices, permuted_keys) = Partition(keys);
+  auto offsets_ptr = offsets.data_ptr<int64_t>();
+  auto indices_ptr = indices.data_ptr<int64_t>();
+  std::vector<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+      results(policies_.size());
+  torch::Tensor result_offsets_tensor =
+      torch::empty(policies_.size() * 2 + 1, offsets.options());
+  auto result_offsets = result_offsets_tensor.data_ptr<int64_t>();
+  namespace gb = graphbolt;
+  {
+    std::lock_guard lock(mtx_);
+    gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+      const auto begin = offsets_ptr[tid];
+      const auto end = offsets_ptr[tid + 1];
+      results[tid] = policies_.at(tid)->QueryAndReplace(
+          permuted_keys.slice(0, begin, end));
+      const auto missing_cnt = std::get<3>(results[tid]).size(0);
+      result_offsets[tid] = end - begin - missing_cnt;
+      result_offsets[tid + policies_.size()] = missing_cnt;
+    });
+  }
+  std::exclusive_scan(
+      result_offsets, result_offsets + result_offsets_tensor.size(0),
+      result_offsets, 0);
+  torch::Tensor positions = torch::empty(
+      keys.size(0),
+      std::get<0>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor output_indices = torch::empty_like(
+      indices, indices.options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor pointers = torch::empty(
+      keys.size(0),
+      std::get<2>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  torch::Tensor missing_keys = torch::empty(
+      result_offsets[2 * policies_.size()] - result_offsets[policies_.size()],
+      std::get<3>(results[0]).options().pinned_memory(utils::is_pinned(keys)));
+  auto missing_offsets =
+      torch::empty(policies_.size() + 1, result_offsets_tensor.options());
+  auto positions_ptr = positions.data_ptr<int64_t>();
+  auto output_indices_ptr = output_indices.data_ptr<int64_t>();
+  auto pointers_ptr = pointers.data_ptr<int64_t>();
+  auto missing_offsets_ptr = missing_offsets.data_ptr<int64_t>();
+  missing_offsets_ptr[0] = 0;
+  gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+    auto out_index_ptr = indices_ptr + offsets_ptr[tid];
+    auto begin = result_offsets[tid];
+    auto end = result_offsets[tid + 1];
+    const auto num_selected = end - begin;
+    auto indices_ptr = std::get<1>(results[tid]).data_ptr<int64_t>();
+    for (int64_t i = 0; i < num_selected; i++) {
+      output_indices_ptr[begin + i] = out_index_ptr[indices_ptr[i]];
+    }
+    auto selected_positions_ptr = std::get<0>(results[tid]).data_ptr<int64_t>();
+    std::transform(
+        selected_positions_ptr, selected_positions_ptr + num_selected,
+        positions_ptr + begin,
+        [off = tid * capacity_ / policies_.size()](auto x) { return x + off; });
+    auto selected_pointers_ptr = std::get<2>(results[tid]).data_ptr<int64_t>();
+    std::copy(
+        selected_pointers_ptr, selected_pointers_ptr + num_selected,
+        pointers_ptr + begin);
+    begin = result_offsets[policies_.size() + tid];
+    end = result_offsets[policies_.size() + tid + 1];
+    missing_offsets[tid + 1] = end - result_offsets[policies_.size()];
+    const auto num_missing = end - begin;
+    for (int64_t i = 0; i < num_missing; i++) {
+      output_indices_ptr[begin + i] =
+          out_index_ptr[indices_ptr[i + num_selected]];
+    }
+    auto missing_positions_ptr = selected_positions_ptr + num_selected;
+    std::transform(
+        missing_positions_ptr, missing_positions_ptr + num_missing,
+        positions_ptr + begin,
+        [off = tid * capacity_ / policies_.size()](auto x) { return x + off; });
+    auto missing_pointers_ptr = selected_pointers_ptr + num_selected;
+    std::copy(
+        missing_pointers_ptr, missing_pointers_ptr + num_missing,
+        pointers_ptr + begin);
+    std::memcpy(
+        reinterpret_cast<std::byte*>(missing_keys.data_ptr()) +
+            (begin - result_offsets[policies_.size()]) *
+                missing_keys.element_size(),
+        std::get<3>(results[tid]).data_ptr(),
+        num_missing * missing_keys.element_size());
+  });
+  auto found_offsets = result_offsets_tensor.slice(0, 0, policies_.size() + 1);
+  missing_keys = AddOffset(missing_keys, -offset);
+  return std::make_tuple(
+      positions, output_indices, pointers, missing_keys, found_offsets,
+      missing_offsets);
+}
+
+c10::intrusive_ptr<Future<std::vector<torch::Tensor>>>
+PartitionedCachePolicy::QueryAndReplaceAsync(
+    torch::Tensor keys, const int64_t offset) {
+  return async([=] {
+    auto
+        [positions, output_indices, pointers, missing_keys, found_offsets,
+         missing_offsets] = QueryAndReplace(keys, offset);
+    return std::vector{positions,    output_indices, pointers,
+                       missing_keys, found_offsets,  missing_offsets};
+  });
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+PartitionedCachePolicy::Replace(
+    torch::Tensor keys, torch::optional<torch::Tensor> offsets,
+    const int64_t offset) {
+  keys = AddOffset(keys, offset);
+  if (policies_.size() == 1) {
+    std::lock_guard lock(mtx_);
+    auto [positions, pointers] = policies_[0]->Replace(keys);
+    if (!offsets.has_value()) {
+      offsets = torch::empty(2, pointers.options());
+      auto offsets_ptr = offsets->data_ptr<int64_t>();
+      offsets_ptr[0] = 0;
+      offsets_ptr[1] = pointers.size(0);
+    }
+    return {positions, pointers, *offsets};
+  }
+  const auto offsets_provided = offsets.has_value();
+  torch::Tensor indices, permuted_keys;
+  if (!offsets_provided) {
+    std::tie(offsets, indices, permuted_keys) = Partition(keys);
+  } else {
+    permuted_keys = keys;
+  }
+  auto output_positions = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto output_pointers = torch::empty_like(
+      keys, keys.options()
+                .dtype(torch::kInt64)
+                .pinned_memory(utils::is_pinned(keys)));
+  auto offsets_ptr = offsets->data_ptr<int64_t>();
+  auto indices_ptr = offsets_provided ? nullptr : indices.data_ptr<int64_t>();
+  auto output_positions_ptr = output_positions.data_ptr<int64_t>();
+  auto output_pointers_ptr = output_pointers.data_ptr<int64_t>();
+  namespace gb = graphbolt;
+  std::unique_lock lock(mtx_);
+  std::atomic<size_t> semaphore = policies_.size();
+  gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+    const auto begin = offsets_ptr[tid];
+    const auto end = offsets_ptr[tid + 1];
+    auto [positions, pointers] =
+        policies_.at(tid)->Replace(permuted_keys.slice(0, begin, end));
+    const auto ticket = semaphore.fetch_add(-1, std::memory_order_release) - 1;
+    if (ticket == 0) {
+      // This thread was the last thread in the critical region.
+      lock.unlock();
+    }
+    auto positions_ptr = positions.data_ptr<int64_t>();
+    const auto off = tid * capacity_ / policies_.size();
+    if (indices_ptr) {
+      for (int64_t i = 0; i < positions.size(0); i++) {
+        output_positions_ptr[indices_ptr[begin + i]] = positions_ptr[i] + off;
+      }
+    } else {
+      std::transform(
+          positions_ptr, positions_ptr + positions.size(0),
+          output_positions_ptr + begin, [off](auto x) { return x + off; });
+    }
+    auto pointers_ptr = pointers.data_ptr<int64_t>();
+    std::copy(
+        pointers_ptr, pointers_ptr + pointers.size(0),
+        output_pointers_ptr + begin);
+  });
+  return {output_positions, output_pointers, *offsets};
+}
+
+c10::intrusive_ptr<Future<std::vector<torch::Tensor>>>
+PartitionedCachePolicy::ReplaceAsync(
+    torch::Tensor keys, torch::optional<torch::Tensor> offsets,
+    const int64_t offset) {
+  return async([=] {
+    auto [positions, pointers, offsets_out] = Replace(keys, offsets, offset);
+    return std::vector{positions, pointers, offsets_out};
+  });
+}
+
+template <bool write>
+void PartitionedCachePolicy::ReadingWritingCompletedImpl(
+    torch::Tensor pointers, torch::Tensor offsets) {
+  if (policies_.size() == 1) {
+    if constexpr (write)
+      policies_[0]->WritingCompleted(pointers);
+    else
+      policies_[0]->ReadingCompleted(pointers);
+    return;
+  }
+  auto offsets_ptr = offsets.data_ptr<int64_t>();
+  namespace gb = graphbolt;
+  gb::parallel_for_each(0, policies_.size(), 1, [&](int64_t tid) {
+    const auto begin = offsets_ptr[tid];
+    const auto end = offsets_ptr[tid + 1];
+    if constexpr (write)
+      policies_.at(tid)->WritingCompleted(pointers.slice(0, begin, end));
+    else
+      policies_.at(tid)->ReadingCompleted(pointers.slice(0, begin, end));
+  });
+}
+
+void PartitionedCachePolicy::ReadingCompleted(
+    torch::Tensor pointers, torch::Tensor offsets) {
+  ReadingWritingCompletedImpl<false>(pointers, offsets);
+}
+
+void PartitionedCachePolicy::WritingCompleted(
+    torch::Tensor pointers, torch::Tensor offsets) {
+  ReadingWritingCompletedImpl<true>(pointers, offsets);
+}
+
+c10::intrusive_ptr<Future<void>> PartitionedCachePolicy::ReadingCompletedAsync(
+    torch::Tensor pointers, torch::Tensor offsets) {
+  return async([=] { return ReadingCompleted(pointers, offsets); });
+}
+
+c10::intrusive_ptr<Future<void>> PartitionedCachePolicy::WritingCompletedAsync(
+    torch::Tensor pointers, torch::Tensor offsets) {
+  return async([=] { return WritingCompleted(pointers, offsets); });
+}
+
+template <typename CachePolicy>
+c10::intrusive_ptr<PartitionedCachePolicy> PartitionedCachePolicy::Create(
+    int64_t capacity, int64_t num_partitions) {
+  static_assert(std::is_base_of_v<BaseCachePolicy, CachePolicy>);
+  return c10::make_intrusive<PartitionedCachePolicy>(
+      CachePolicy(), capacity, num_partitions);
+}
+
+template c10::intrusive_ptr<PartitionedCachePolicy>
+    PartitionedCachePolicy::Create<S3FifoCachePolicy>(int64_t, int64_t);
+template c10::intrusive_ptr<PartitionedCachePolicy>
+    PartitionedCachePolicy::Create<SieveCachePolicy>(int64_t, int64_t);
+template c10::intrusive_ptr<PartitionedCachePolicy>
+    PartitionedCachePolicy::Create<LruCachePolicy>(int64_t, int64_t);
+template c10::intrusive_ptr<PartitionedCachePolicy>
+    PartitionedCachePolicy::Create<ClockCachePolicy>(int64_t, int64_t);
+
+}  // namespace storage
+}  // namespace graphbolt

graphbolt/src/partitioned_cache_policy.h

@@ -0,0 +1,181 @@
+/**
+ *   Copyright (c) 2023, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file partitioned_cache_policy.h
+ * @brief Partitioned cache policy implementation on the CPU.
+ */
+#ifndef GRAPHBOLT_PARTITIONED_CACHE_H_
+#define GRAPHBOLT_PARTITIONED_CACHE_H_
+
+#include <graphbolt/async.h>
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+
+#include <mutex>
+#include <pcg_random.hpp>
+#include <random>
+#include <type_traits>
+#include <vector>
+
+#include "./cache_policy.h"
+
+namespace graphbolt {
+namespace storage {
+
+/**
+ * @brief PartitionedCachePolicy works by partitioning the key space to a set
+ * number of partitions that is provided as the second argument of its
+ * constructor. Since the partitioning is random but deterministic, the caching
+ * policy performance is not affected as the key distribution stays the same in
+ * each partition.
+ **/
+class PartitionedCachePolicy : public torch::CustomClassHolder {
+ public:
+  /**
+   * @brief The policy query function.
+   * @param capacity The capacity of the cache.
+   * @param num_partitions The number of caching policies instantiated in a
+   * one-to-one mapping to each partition.
+   */
+  template <typename CachePolicy>
+  PartitionedCachePolicy(CachePolicy, int64_t capacity, int64_t num_partitions);
+
+  /**
+   * @brief The policy query function.
+   * @param keys The keys to query the cache.
+   * @param offset The offset to be added to the keys.
+   *
+   * @return (positions, indices, missing_keys, found_ptrs, found_offsets,
+   * missing_offsets), where positions has the locations of the keys which were
+   * found in the cache, missing_keys has the keys that were not found and
+   * indices is defined such that keys[indices[:positions.size(0)]] gives us the
+   * keys for the found pointers and keys[indices[positions.size(0):]] is
+   * identical to missing_keys. The found_offsets tensor holds the partition
+   * offsets for the found pointers. The missing_offsets holds the partition
+   * offsets for the missing_keys.
+   */
+  std::tuple<
+      torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+      torch::Tensor>
+  Query(torch::Tensor keys, int64_t offset);
+
+  c10::intrusive_ptr<Future<std::vector<torch::Tensor>>> QueryAsync(
+      torch::Tensor keys, int64_t offset);
+
+  /**
+   * @brief The policy query and then replace function.
+   * @param keys The keys to query the cache.
+   * @param offset The offset to be added to the keys.
+   *
+   * @return (positions, indices, pointers, missing_keys, found_offsets,
+   * missing_offsets), where positions has the locations of the keys which were
+   * emplaced into the cache, pointers point to the emplaced CacheKey pointers
+   * in the cache, missing_keys has the keys that were not found and just
+   * inserted and indices is defined such that keys[indices[:keys.size(0) -
+   * missing_keys.size(0)]] gives us the keys for the found keys and
+   * keys[indices[keys.size(0) - missing_keys.size(0):]] is identical to
+   * missing_keys. The found_offsets tensor holds the partition offsets for the
+   * found pointers. The missing_offsets holds the partition offsets for the
+   * missing_keys and missing pointers.
+   */
+  std::tuple<
+      torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor,
+      torch::Tensor>
+  QueryAndReplace(torch::Tensor keys, int64_t offset);
+
+  c10::intrusive_ptr<Future<std::vector<torch::Tensor>>> QueryAndReplaceAsync(
+      torch::Tensor keys, int64_t offset);
+
+  /**
+   * @brief The policy replace function.
+   * @param keys The keys to query the cache.
+   * @param offsets The partition offsets for the keys.
+   * @param offset The offset to be added to the keys.
+   *
+   * @return (positions, pointers, offsets), where positions holds the locations
+   * of the replaced entries in the cache, pointers holds the CacheKey pointers
+   * for the inserted keys and offsets holds the partition offsets for pointers.
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Replace(
+      torch::Tensor keys, torch::optional<torch::Tensor> offsets,
+      int64_t offset);
+
+  c10::intrusive_ptr<Future<std::vector<torch::Tensor>>> ReplaceAsync(
+      torch::Tensor keys, torch::optional<torch::Tensor> offsets,
+      int64_t offset);
+
+  template <bool write>
+  void ReadingWritingCompletedImpl(
+      torch::Tensor pointers, torch::Tensor offsets);
+
+  /**
+   * @brief A reader has finished reading these keys, so they can be
+   * evicted.
+   * @param pointers The CacheKey pointers in the cache to unmark.
+   * @param offsets The partition offsets for the pointers.
+   */
+  void ReadingCompleted(torch::Tensor pointers, torch::Tensor offsets);
+
+  /**
+   * @brief A writer has finished writing these keys, so they can be evicted.
+   * @param pointers The CacheKey pointers in the cache to unmark.
+   * @param offsets The partition offsets for the pointers.
+   */
+  void WritingCompleted(torch::Tensor pointers, torch::Tensor offsets);
+
+  c10::intrusive_ptr<Future<void>> ReadingCompletedAsync(
+      torch::Tensor pointers, torch::Tensor offsets);
+
+  c10::intrusive_ptr<Future<void>> WritingCompletedAsync(
+      torch::Tensor pointers, torch::Tensor offsets);
+
+  template <typename CachePolicy>
+  static c10::intrusive_ptr<PartitionedCachePolicy> Create(
+      int64_t capacity, int64_t num_partitions);
+
+ private:
+  static constexpr uint64_t seed = 1e9 + 7;
+
+  /**
+   * @brief Deterministic assignment of keys to different parts.
+   */
+  int32_t PartAssignment(int64_t key) {
+    pcg32 rng(seed, key);
+    std::uniform_int_distribution<int32_t> dist(0, policies_.size() - 1);
+    return dist(rng);
+  }
+
+  /**
+   * @brief The partition function for a given keys tensor.
+   * @param keys The keys to query the cache.
+   *
+   * @return (offsets, indices, permuted_keys), the returned tensors have the
+   * following properties:
+   * permuted_keys[offsets[i]: offsets[i + 1]] belong to part i and
+   * keys[indices] == permuted_keys
+   */
+  std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> Partition(
+      torch::Tensor keys);
+
+  int64_t capacity_;
+  std::vector<std::unique_ptr<BaseCachePolicy>> policies_;
+  std::mutex mtx_;
+};
+
+}  // namespace storage
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_PARTITIONED_CACHE_H_

graphbolt/src/python_binding.cc

@@ -0,0 +1,233 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file python_binding.cc
+ * @brief Graph bolt library Python binding.
+ */
+
+#include <graphbolt/fused_csc_sampling_graph.h>
+#include <graphbolt/isin.h>
+#include <graphbolt/serialize.h>
+#include <graphbolt/unique_and_compact.h>
+
+#ifdef GRAPHBOLT_USE_CUDA
+#include "./cuda/cooperative_minibatching_utils.h"
+#include "./cuda/max_uva_threads.h"
+#endif
+#include "./cnumpy.h"
+#include "./feature_cache.h"
+#include "./index_select.h"
+#include "./io_uring.h"
+#include "./partitioned_cache_policy.h"
+#include "./random.h"
+#include "./utils.h"
+
+#ifdef GRAPHBOLT_USE_CUDA
+#include "./cuda/extension/gpu_cache.h"
+#include "./cuda/extension/gpu_graph_cache.h"
+#endif
+
+namespace graphbolt {
+namespace sampling {
+
+TORCH_LIBRARY(graphbolt, m) {
+  m.class_<FusedSampledSubgraph>("FusedSampledSubgraph")
+      .def(torch::init<>())
+      .def_readwrite("indptr", &FusedSampledSubgraph::indptr)
+      .def_readwrite("indices", &FusedSampledSubgraph::indices)
+      .def_readwrite(
+          "original_row_node_ids", &FusedSampledSubgraph::original_row_node_ids)
+      .def_readwrite(
+          "original_column_node_ids",
+          &FusedSampledSubgraph::original_column_node_ids)
+      .def_readwrite(
+          "original_edge_ids", &FusedSampledSubgraph::original_edge_ids)
+      .def_readwrite("type_per_edge", &FusedSampledSubgraph::type_per_edge)
+      .def_readwrite("etype_offsets", &FusedSampledSubgraph::etype_offsets);
+  m.class_<Future<void>>("VoidFuture").def("wait", &Future<void>::Wait);
+  m.class_<Future<torch::Tensor>>("TensorFuture")
+      .def("wait", &Future<torch::Tensor>::Wait);
+  m.class_<Future<std::vector<torch::Tensor>>>("TensorListFuture")
+      .def("wait", &Future<std::vector<torch::Tensor>>::Wait);
+  m.class_<Future<c10::intrusive_ptr<FusedSampledSubgraph>>>(
+       "FusedSampledSubgraphFuture")
+      .def("wait", &Future<c10::intrusive_ptr<FusedSampledSubgraph>>::Wait);
+  m.class_<Future<std::vector<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>>>(
+       "UniqueAndCompactBatchedFuture")
+      .def(
+          "wait",
+          &Future<std::vector<std::tuple<
+              torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>>::
+              Wait);
+  m.class_<Future<
+      std::vector<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>>>(
+       "RankSortFuture")
+      .def(
+          "wait",
+          &Future<std::vector<
+              std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>>>::Wait);
+  m.class_<Future<std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t>>>(
+       "GpuGraphCacheQueryFuture")
+      .def(
+          "wait",
+          &Future<std::tuple<torch::Tensor, torch::Tensor, int64_t, int64_t>>::
+              Wait);
+  m.class_<Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>>(
+       "GpuGraphCacheReplaceFuture")
+      .def(
+          "wait",
+          &Future<std::tuple<torch::Tensor, std::vector<torch::Tensor>>>::Wait);
+  m.class_<storage::OnDiskNpyArray>("OnDiskNpyArray")
+      .def("index_select", &storage::OnDiskNpyArray::IndexSelect);
+  m.class_<FusedCSCSamplingGraph>("FusedCSCSamplingGraph")
+      .def("num_nodes", &FusedCSCSamplingGraph::NumNodes)
+      .def("num_edges", &FusedCSCSamplingGraph::NumEdges)
+      .def("csc_indptr", &FusedCSCSamplingGraph::CSCIndptr)
+      .def("indices", &FusedCSCSamplingGraph::Indices)
+      .def("node_type_offset", &FusedCSCSamplingGraph::NodeTypeOffset)
+      .def("type_per_edge", &FusedCSCSamplingGraph::TypePerEdge)
+      .def("node_type_to_id", &FusedCSCSamplingGraph::NodeTypeToID)
+      .def("edge_type_to_id", &FusedCSCSamplingGraph::EdgeTypeToID)
+      .def("node_attributes", &FusedCSCSamplingGraph::NodeAttributes)
+      .def("edge_attributes", &FusedCSCSamplingGraph::EdgeAttributes)
+      .def("node_attribute", &FusedCSCSamplingGraph::NodeAttribute)
+      .def("edge_attribute", &FusedCSCSamplingGraph::EdgeAttribute)
+      .def("set_csc_indptr", &FusedCSCSamplingGraph::SetCSCIndptr)
+      .def("set_indices", &FusedCSCSamplingGraph::SetIndices)
+      .def("set_node_type_offset", &FusedCSCSamplingGraph::SetNodeTypeOffset)
+      .def("set_type_per_edge", &FusedCSCSamplingGraph::SetTypePerEdge)
+      .def("set_node_type_to_id", &FusedCSCSamplingGraph::SetNodeTypeToID)
+      .def("set_edge_type_to_id", &FusedCSCSamplingGraph::SetEdgeTypeToID)
+      .def("set_node_attributes", &FusedCSCSamplingGraph::SetNodeAttributes)
+      .def("set_edge_attributes", &FusedCSCSamplingGraph::SetEdgeAttributes)
+      .def("add_node_attribute", &FusedCSCSamplingGraph::AddNodeAttribute)
+      .def("add_edge_attribute", &FusedCSCSamplingGraph::AddEdgeAttribute)
+      .def("in_subgraph", &FusedCSCSamplingGraph::InSubgraph)
+      .def("sample_neighbors", &FusedCSCSamplingGraph::SampleNeighbors)
+      .def(
+          "sample_neighbors_async",
+          &FusedCSCSamplingGraph::SampleNeighborsAsync)
+      .def(
+          "temporal_sample_neighbors",
+          &FusedCSCSamplingGraph::TemporalSampleNeighbors)
+      .def("copy_to_shared_memory", &FusedCSCSamplingGraph::CopyToSharedMemory)
+      .def_pickle(
+          // __getstate__
+          [](const c10::intrusive_ptr<FusedCSCSamplingGraph>& self)
+              -> torch::Dict<
+                  std::string, torch::Dict<std::string, torch::Tensor>> {
+            return self->GetState();
+          },
+          // __setstate__
+          [](torch::Dict<std::string, torch::Dict<std::string, torch::Tensor>>
+                 state) -> c10::intrusive_ptr<FusedCSCSamplingGraph> {
+            auto g = c10::make_intrusive<FusedCSCSamplingGraph>();
+            g->SetState(state);
+            return g;
+          });
+#ifdef GRAPHBOLT_USE_CUDA
+  m.class_<cuda::GpuCache>("GpuCache")
+      .def("query", &cuda::GpuCache::Query)
+      .def("query_async", &cuda::GpuCache::QueryAsync)
+      .def("replace", &cuda::GpuCache::Replace);
+  m.def("gpu_cache", &cuda::GpuCache::Create);
+  m.class_<cuda::GpuGraphCache>("GpuGraphCache")
+      .def("query", &cuda::GpuGraphCache::Query)
+      .def("query_async", &cuda::GpuGraphCache::QueryAsync)
+      .def("replace", &cuda::GpuGraphCache::Replace)
+      .def("replace_async", &cuda::GpuGraphCache::ReplaceAsync);
+  m.def("gpu_graph_cache", &cuda::GpuGraphCache::Create);
+#endif
+  m.def("fused_csc_sampling_graph", &FusedCSCSamplingGraph::Create);
+  m.class_<storage::PartitionedCachePolicy>("PartitionedCachePolicy")
+      .def("query", &storage::PartitionedCachePolicy::Query)
+      .def("query_async", &storage::PartitionedCachePolicy::QueryAsync)
+      .def(
+          "query_and_replace",
+          &storage::PartitionedCachePolicy::QueryAndReplace)
+      .def(
+          "query_and_replace_async",
+          &storage::PartitionedCachePolicy::QueryAndReplaceAsync)
+      .def("replace", &storage::PartitionedCachePolicy::Replace)
+      .def("replace_async", &storage::PartitionedCachePolicy::ReplaceAsync)
+      .def(
+          "reading_completed",
+          &storage::PartitionedCachePolicy::ReadingCompleted)
+      .def(
+          "reading_completed_async",
+          &storage::PartitionedCachePolicy::ReadingCompletedAsync)
+      .def(
+          "writing_completed",
+          &storage::PartitionedCachePolicy::WritingCompleted)
+      .def(
+          "writing_completed_async",
+          &storage::PartitionedCachePolicy::WritingCompletedAsync);
+  m.def(
+      "s3_fifo_cache_policy",
+      &storage::PartitionedCachePolicy::Create<storage::S3FifoCachePolicy>);
+  m.def(
+      "sieve_cache_policy",
+      &storage::PartitionedCachePolicy::Create<storage::SieveCachePolicy>);
+  m.def(
+      "lru_cache_policy",
+      &storage::PartitionedCachePolicy::Create<storage::LruCachePolicy>);
+  m.def(
+      "clock_cache_policy",
+      &storage::PartitionedCachePolicy::Create<storage::ClockCachePolicy>);
+  m.class_<storage::FeatureCache>("FeatureCache")
+      .def("is_pinned", &storage::FeatureCache::IsPinned)
+      .def_property("nbytes", &storage::FeatureCache::NumBytes)
+      .def("index_select", &storage::FeatureCache::IndexSelect)
+      .def("query", &storage::FeatureCache::Query)
+      .def("query_async", &storage::FeatureCache::QueryAsync)
+      .def("replace", &storage::FeatureCache::Replace)
+      .def("replace_async", &storage::FeatureCache::ReplaceAsync);
+  m.def("feature_cache", &storage::FeatureCache::Create);
+  m.def(
+      "load_from_shared_memory", &FusedCSCSamplingGraph::LoadFromSharedMemory);
+  m.def("unique_and_compact", &UniqueAndCompact);
+  m.def("unique_and_compact_batched", &UniqueAndCompactBatched);
+  m.def("unique_and_compact_batched_async", &UniqueAndCompactBatchedAsync);
+  m.def("isin", &IsIn);
+  m.def("is_not_in_index", &IsNotInIndex);
+  m.def("is_not_in_index_async", &IsNotInIndexAsync);
+  m.def("index_select", &ops::IndexSelect);
+  m.def("index_select_async", &ops::IndexSelectAsync);
+  m.def("scatter_async", &ops::ScatterAsync);
+  m.def("index_select_csc", &ops::IndexSelectCSC);
+  m.def("index_select_csc_batched", &ops::IndexSelectCSCBatched);
+  m.def("index_select_csc_batched_async", &ops::IndexSelectCSCBatchedAsync);
+  m.def("ondisk_npy_array", &storage::OnDiskNpyArray::Create);
+  m.def("detect_io_uring", &io_uring::IsAvailable);
+  m.def("set_num_io_uring_threads", &io_uring::SetNumThreads);
+  m.def("set_worker_id", &utils::SetWorkerId);
+  m.def("set_seed", &RandomEngine::SetManualSeed);
+#ifdef GRAPHBOLT_USE_CUDA
+  m.def("set_max_uva_threads", &cuda::set_max_uva_threads);
+  m.def("rank_sort", &cuda::RankSort);
+  m.def("rank_sort_async", &cuda::RankSortAsync);
+#endif
+#ifdef HAS_IMPL_ABSTRACT_PYSTUB
+  m.impl_abstract_pystub("dgl.graphbolt.base", "//dgl.graphbolt.base");
+#endif
+  m.def(
+      "expand_indptr(Tensor indptr, ScalarType dtype, Tensor? node_ids, "
+      "SymInt? output_size) -> Tensor"
+#ifdef HAS_PT2_COMPLIANT_TAG
+      ,
+      {at::Tag::pt2_compliant_tag}
+#endif
+  );
+  m.def(
+      "indptr_edge_ids(Tensor indptr, ScalarType dtype, Tensor? offset, "
+      "SymInt? output_size) -> "
+      "Tensor"
+#ifdef HAS_PT2_COMPLIANT_TAG
+      ,
+      {at::Tag::pt2_compliant_tag}
+#endif
+  );
+}
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/random.cc

@@ -0,0 +1,72 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file random.cc
+ * @brief Random Engine.
+ */
+
+#include "./random.h"
+
+#include <torch/torch.h>
+
+namespace graphbolt {
+
+namespace {
+
+// Get a unique integer ID representing this thread.
+inline uint32_t GetThreadId() {
+  static int num_threads = 0;
+  static std::mutex mutex;
+  static thread_local int id = -1;
+
+  if (id == -1) {
+    std::lock_guard<std::mutex> guard(mutex);
+    id = num_threads;
+    num_threads++;
+  }
+  return id;
+}
+
+};  // namespace
+
+std::mutex RandomEngine::manual_seed_mutex;
+std::optional<uint64_t> RandomEngine::manual_seed;
+
+/** @brief Constructor with default seed. */
+RandomEngine::RandomEngine() {
+  std::random_device rd;
+  std::lock_guard lock(manual_seed_mutex);
+  if (!manual_seed.has_value()) manual_seed = rd();
+  SetSeed(manual_seed.value());
+}
+
+/** @brief Constructor with given seed. */
+RandomEngine::RandomEngine(uint64_t seed) : RandomEngine(seed, GetThreadId()) {}
+
+/** @brief Constructor with given seed. */
+RandomEngine::RandomEngine(uint64_t seed, uint64_t stream) {
+  SetSeed(seed, stream);
+}
+
+/** @brief Get the thread-local random number generator instance. */
+RandomEngine* RandomEngine::ThreadLocal() {
+  static thread_local RandomEngine engine;
+  return &engine;
+}
+
+/** @brief Set the seed. */
+void RandomEngine::SetSeed(uint64_t seed) { SetSeed(seed, GetThreadId()); }
+
+/** @brief Set the seed. */
+void RandomEngine::SetSeed(uint64_t seed, uint64_t stream) {
+  rng_.seed(seed, stream);
+}
+
+/** @brief Manually fix the seed. */
+void RandomEngine::SetManualSeed(int64_t seed) {
+  // Intentionally set the seed for current thread also.
+  RandomEngine::ThreadLocal()->SetSeed(seed);
+  std::lock_guard lock(manual_seed_mutex);
+  manual_seed = seed;
+}
+
+}  // namespace graphbolt

graphbolt/src/random.h

@@ -0,0 +1,79 @@
+
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file random.h
+ * @brief Random Engine class.
+ */
+#ifndef GRAPHBOLT_RANDOM_H_
+#define GRAPHBOLT_RANDOM_H_
+
+#include <mutex>
+#include <optional>
+#include <pcg_random.hpp>
+#include <random>
+#include <thread>
+
+namespace graphbolt {
+
+/**
+ * @brief Thread-local Random Number Generator class.
+ */
+class RandomEngine {
+ public:
+  /** @brief Constructor with default seed. */
+  RandomEngine();
+
+  /** @brief Constructor with given seed. */
+  explicit RandomEngine(uint64_t seed);
+  explicit RandomEngine(uint64_t seed, uint64_t stream);
+
+  /** @brief Get the thread-local random number generator instance. */
+  static RandomEngine* ThreadLocal();
+
+  /** @brief Set the seed. */
+  void SetSeed(uint64_t seed);
+  void SetSeed(uint64_t seed, uint64_t stream);
+
+  /** @brief Protect manual seed accesses. */
+  static std::mutex manual_seed_mutex;
+
+  /** @brief Manually fix the seed. */
+  static std::optional<uint64_t> manual_seed;
+  static void SetManualSeed(int64_t seed);
+
+  /**
+   * @brief Generate a uniform random integer in [low, high).
+   */
+  template <typename T>
+  T RandInt(T lower, T upper) {
+    std::uniform_int_distribution<T> dist(lower, upper - 1);
+    return dist(rng_);
+  }
+
+  /**
+   * @brief Generate a uniform random real number in [low, high).
+   */
+  template <typename T>
+  T Uniform(T lower, T upper) {
+    std::uniform_real_distribution<T> dist(lower, upper);
+    return dist(rng_);
+  }
+
+  /**
+   * @brief Generate random non-negative floating-point values according to
+   * exponential distribution. Probability density function: P(x|λ) = λe^(-λx).
+   */
+  template <typename T>
+  T Exponential(T lambda) {
+    std::exponential_distribution<T> dist(lambda);
+    return dist(rng_);
+  }
+
+ private:
+  pcg32 rng_;
+};
+
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_RANDOM_H_

graphbolt/src/serialize.cc

@@ -0,0 +1,26 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file graphbolt/src/serialize.cc
+ * @brief Source file of serialize.
+ */
+
+#include <graphbolt/serialize.h>
+#include <torch/torch.h>
+
+namespace torch {
+
+serialize::InputArchive& operator>>(
+    serialize::InputArchive& archive,
+    graphbolt::sampling::FusedCSCSamplingGraph& graph) {
+  graph.Load(archive);
+  return archive;
+}
+
+serialize::OutputArchive& operator<<(
+    serialize::OutputArchive& archive,
+    const graphbolt::sampling::FusedCSCSamplingGraph& graph) {
+  graph.Save(archive);
+  return archive;
+}
+
+}  // namespace torch

graphbolt/src/shared_memory.cc

@@ -0,0 +1,167 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file shared_memory.cc
+ * @brief Source file of graphbolt shared memory.
+ */
+#ifndef _WIN32
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+#endif  // !_WIN32
+
+#include <graphbolt/shared_memory.h>
+#include <stdio.h>
+#include <string.h>
+#include <torch/torch.h>
+
+namespace graphbolt {
+namespace sampling {
+
+// Two processes opening the same path are guaranteed to access the same shared
+// memory object if and only if path begins with a slash ('/') character.
+constexpr char kSharedMemNamePrefix[] = "/dgl.graphbolt.";
+constexpr char kSharedMemNameSuffix[] = ".lock";
+
+// A prefix and a suffix are added to the name of the shared memory to create
+// the name of the shared memory object.
+inline std::string DecorateName(const std::string& name) {
+  return kSharedMemNamePrefix + name + kSharedMemNameSuffix;
+}
+
+SharedMemory::SharedMemory(const std::string& name)
+    : name_(name), size_(0), ptr_(nullptr) {
+#ifdef _WIN32
+  this->handle_ = nullptr;
+#else   // _WIN32
+  this->file_descriptor_ = -1;
+  this->is_creator_ = false;
+#endif  // _WIN32
+}
+
+#ifdef _WIN32
+
+SharedMemory::~SharedMemory() {
+  if (ptr_) CHECK(UnmapViewOfFile(ptr_)) << "Win32 Error: " << GetLastError();
+  if (handle_) CloseHandle(handle_);
+}
+
+void* SharedMemory::Create(size_t size) {
+  size_ = size;
+
+  std::string decorated_name = DecorateName(name_);
+  handle_ = CreateFileMapping(
+      INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE,
+      static_cast<DWORD>(size >> 32), static_cast<DWORD>(size & 0xFFFFFFFF),
+      decorated_name.c_str());
+  TORCH_CHECK(
+      handle_ != nullptr, "Failed to open ", decorated_name,
+      ", Win32 error: ", GetLastError());
+
+  ptr_ = MapViewOfFile(handle_, FILE_MAP_ALL_ACCESS, 0, 0, size);
+  TORCH_CHECK(
+      ptr_ != nullptr, "Memory mapping failed, Win32 error: ", GetLastError());
+  return ptr_;
+}
+
+void* SharedMemory::Open() {
+  std::string decorated_name = DecorateName(name_);
+  handle_ = OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, decorated_name.c_str());
+  TORCH_CHECK(
+      handle_ != nullptr, "Failed to open ", decorated_name,
+      ", Win32 Error: ", GetLastError());
+
+  ptr_ = MapViewOfFile(handle_, FILE_MAP_ALL_ACCESS, 0, 0, 0);
+  TORCH_CHECK(
+      ptr_ != nullptr, "Memory mapping failed, Win32 error: ", GetLastError());
+
+  // Obtain the size of the memory-mapped file.
+  MEMORY_BASIC_INFORMATION memInfo;
+  TORCH_CHECK(
+      VirtualQuery(ptr_, &memInfo, sizeof(memInfo)) != 0,
+      "Failed to get the size of shared memory: ", GetLastError());
+  size_ = static_cast<size_t>(memInfo.RegionSize);
+
+  return ptr_;
+}
+
+bool SharedMemory::Exists(const std::string& name) {
+  std::string decorated_name = DecorateName(name);
+  HANDLE handle =
+      OpenFileMapping(FILE_MAP_ALL_ACCESS, FALSE, decorated_name.c_str());
+  bool exists = handle != nullptr;
+  if (exists) {
+    CloseHandle(handle);
+  }
+  return exists;
+}
+
+#else  // _WIN32
+
+SharedMemory::~SharedMemory() {
+  if (ptr_ && size_ != 0) CHECK(munmap(ptr_, size_) != -1) << strerror(errno);
+  if (file_descriptor_ != -1) close(file_descriptor_);
+
+  std::string decorated_name = DecorateName(name_);
+  if (is_creator_ && decorated_name != "") shm_unlink(decorated_name.c_str());
+}
+
+void *SharedMemory::Create(size_t size) {
+  size_ = size;
+  is_creator_ = true;
+
+  // TODO(zhenkun): handle the error properly if the shared memory object
+  // already exists.
+  std::string decorated_name = DecorateName(name_);
+  file_descriptor_ =
+      shm_open(decorated_name.c_str(), O_RDWR | O_CREAT, S_IRUSR | S_IWUSR);
+  TORCH_CHECK(file_descriptor_ != -1, "Failed to open: ", strerror(errno));
+
+  auto status = ftruncate(file_descriptor_, size);
+  TORCH_CHECK(status != -1, "Failed to truncate the file: ", strerror(errno));
+
+  ptr_ =
+      mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, file_descriptor_, 0);
+  TORCH_CHECK(
+      ptr_ != MAP_FAILED,
+      "Failed to map shared memory, mmap failed with error: ", strerror(errno));
+  return ptr_;
+}
+
+void *SharedMemory::Open() {
+  std::string decorated_name = DecorateName(name_);
+  file_descriptor_ =
+      shm_open(decorated_name.c_str(), O_RDWR, S_IRUSR | S_IWUSR);
+  TORCH_CHECK(
+      file_descriptor_ != -1, "Failed to open ", decorated_name, ": ",
+      strerror(errno));
+
+  struct stat shm_stat;
+  TORCH_CHECK(
+      fstat(file_descriptor_, &shm_stat) == 0,
+      "Failed to get the size of shared memory: ", strerror(errno));
+  size_ = shm_stat.st_size;
+
+  ptr_ = mmap(
+      NULL, size_, PROT_READ | PROT_WRITE, MAP_SHARED, file_descriptor_, 0);
+  TORCH_CHECK(
+      ptr_ != MAP_FAILED,
+      "Failed to map shared memory, mmap failed with error: ", strerror(errno));
+  return ptr_;
+}
+
+bool SharedMemory::Exists(const std::string &name) {
+  std::string decorated_name = DecorateName(name);
+  int file_descriptor =
+      shm_open(decorated_name.c_str(), O_RDONLY, S_IRUSR | S_IWUSR);
+  bool exists = file_descriptor > 0;
+  if (exists) {
+    close(file_descriptor);
+  }
+  return exists;
+}
+
+#endif  // _WIN32
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/shared_memory_helper.cc

@@ -0,0 +1,223 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file shared_memory_helper.cc
+ * @brief Share memory helper implementation.
+ */
+#include "./shared_memory_helper.h"
+
+#include <graphbolt/serialize.h>
+#include <graphbolt/shared_memory.h>
+#include <torch/torch.h>
+
+#include <cstring>
+#include <string>
+#include <tuple>
+#include <vector>
+
+namespace graphbolt {
+namespace sampling {
+
+static std::string GetSharedMemoryMetadataName(const std::string& name) {
+  return name + "_metadata";
+}
+
+static std::string GetSharedMemoryDataName(const std::string& name) {
+  return name + "_data";
+}
+
+// To avoid unaligned memory access, we round the size of the binary buffer to
+// the nearest multiple of 8 bytes.
+inline static int64_t GetRoundedSize(int64_t size) {
+  constexpr int64_t ALIGNED_SIZE = 8;
+  return (size + ALIGNED_SIZE - 1) / ALIGNED_SIZE * ALIGNED_SIZE;
+}
+
+SharedMemoryHelper::SharedMemoryHelper(const std::string& name)
+    : name_(name),
+      metadata_size_(0),
+      data_size_(0),
+      metadata_shared_memory_(nullptr),
+      data_shared_memory_(nullptr),
+      metadata_offset_(0),
+      data_offset_(0) {}
+
+void SharedMemoryHelper::InitializeRead() {
+  metadata_offset_ = 0;
+  data_offset_ = 0;
+  if (metadata_shared_memory_ == nullptr) {
+    // Reader process opens the shared memory.
+    metadata_shared_memory_ =
+        std::make_unique<SharedMemory>(GetSharedMemoryMetadataName(name_));
+    metadata_shared_memory_->Open();
+    metadata_size_ = metadata_shared_memory_->GetSize();
+    data_shared_memory_ =
+        std::make_unique<SharedMemory>(GetSharedMemoryDataName(name_));
+    data_shared_memory_->Open();
+    data_size_ = data_shared_memory_->GetSize();
+  }
+}
+
+void SharedMemoryHelper::WriteTorchArchive(
+    torch::serialize::OutputArchive&& archive) {
+  metadata_to_write_.emplace_back(std::move(archive));
+}
+
+torch::serialize::InputArchive SharedMemoryHelper::ReadTorchArchive() {
+  auto metadata_ptr = this->GetCurrentMetadataPtr();
+  int64_t metadata_size = static_cast<int64_t*>(metadata_ptr)[0];
+  torch::serialize::InputArchive archive;
+  archive.load_from(
+      static_cast<const char*>(metadata_ptr) + sizeof(int64_t), metadata_size);
+  auto rounded_size = GetRoundedSize(metadata_size);
+  this->MoveMetadataPtr(sizeof(int64_t) + rounded_size);
+  return archive;
+}
+
+void SharedMemoryHelper::WriteTorchTensor(
+    torch::optional<torch::Tensor> tensor) {
+  torch::serialize::OutputArchive archive;
+  archive.write("has_value", tensor.has_value());
+  if (tensor.has_value()) {
+    archive.write("shape", tensor.value().sizes());
+    archive.write("dtype", tensor.value().scalar_type());
+  }
+  this->WriteTorchArchive(std::move(archive));
+  tensors_to_write_.push_back(tensor);
+}
+
+torch::optional<torch::Tensor> SharedMemoryHelper::ReadTorchTensor() {
+  auto archive = this->ReadTorchArchive();
+  bool has_value = read_from_archive<bool>(archive, "has_value");
+  if (has_value) {
+    auto shape = read_from_archive<std::vector<int64_t>>(archive, "shape");
+    auto dtype = read_from_archive<torch::ScalarType>(archive, "dtype");
+    auto data_ptr = this->GetCurrentDataPtr();
+    auto tensor = torch::from_blob(data_ptr, shape, dtype);
+    auto rounded_size = GetRoundedSize(tensor.numel() * tensor.element_size());
+    this->MoveDataPtr(rounded_size);
+    return tensor;
+  } else {
+    return torch::nullopt;
+  }
+}
+
+void SharedMemoryHelper::WriteTorchTensorDict(
+    torch::optional<torch::Dict<std::string, torch::Tensor>> tensor_dict) {
+  torch::serialize::OutputArchive archive;
+  if (!tensor_dict.has_value()) {
+    archive.write("has_value", false);
+    this->WriteTorchArchive(std::move(archive));
+    return;
+  }
+  archive.write("has_value", true);
+  auto dict_value = tensor_dict.value();
+  archive.write("num_tensors", static_cast<int64_t>(dict_value.size()));
+  int counter = 0;
+  for (auto it = dict_value.begin(); it != dict_value.end(); ++it) {
+    archive.write(std::string("key_") + std::to_string(counter), it->key());
+    counter++;
+  }
+  this->WriteTorchArchive(std::move(archive));
+  for (auto it = dict_value.begin(); it != dict_value.end(); ++it) {
+    this->WriteTorchTensor(it->value());
+  }
+}
+
+torch::optional<torch::Dict<std::string, torch::Tensor>>
+SharedMemoryHelper::ReadTorchTensorDict() {
+  auto archive = this->ReadTorchArchive();
+  if (!read_from_archive<bool>(archive, "has_value")) {
+    return torch::nullopt;
+  }
+  int64_t num_tensors = read_from_archive<int64_t>(archive, "num_tensors");
+  torch::Dict<std::string, torch::Tensor> tensor_dict;
+  for (int64_t i = 0; i < num_tensors; ++i) {
+    auto key = read_from_archive<std::string>(
+        archive, std::string("key_") + std::to_string(i));
+    auto tensor = this->ReadTorchTensor();
+    tensor_dict.insert(key, tensor.value());
+  }
+  return tensor_dict;
+}
+
+void SharedMemoryHelper::SerializeMetadata() {
+  for (auto& archive : metadata_to_write_) {
+    std::stringstream serialized;
+    archive.save_to(serialized);
+    metadata_strings_to_write_.push_back(std::move(serialized.str()));
+  }
+  metadata_to_write_.clear();
+}
+
+void SharedMemoryHelper::WriteMetadataToSharedMemory() {
+  metadata_offset_ = 0;
+  for (const auto& str : metadata_strings_to_write_) {
+    auto metadata_ptr = this->GetCurrentMetadataPtr();
+    static_cast<int64_t*>(metadata_ptr)[0] = str.size();
+    memcpy(
+        static_cast<char*>(metadata_ptr) + sizeof(int64_t), str.data(),
+        str.size());
+    int64_t rounded_size = GetRoundedSize(str.size());
+    this->MoveMetadataPtr(sizeof(int64_t) + rounded_size);
+  }
+  metadata_strings_to_write_.clear();
+}
+
+void SharedMemoryHelper::WriteTorchTensorInternal(
+    torch::optional<torch::Tensor> tensor) {
+  if (tensor.has_value()) {
+    size_t memory_size = tensor.value().numel() * tensor.value().element_size();
+    auto data_ptr = this->GetCurrentDataPtr();
+    auto contiguous_tensor = tensor.value().contiguous();
+    memcpy(data_ptr, contiguous_tensor.data_ptr(), memory_size);
+    this->MoveDataPtr(GetRoundedSize(memory_size));
+  }
+}
+
+void SharedMemoryHelper::Flush() {
+  size_t data_size = 0;
+  for (auto tensor : tensors_to_write_) {
+    if (tensor.has_value()) {
+      auto tensor_size = tensor.value().numel() * tensor.value().element_size();
+      data_size += GetRoundedSize(tensor_size);
+    }
+  }
+
+  // Serialize the metadata archives.
+  SerializeMetadata();
+
+  // Create the shared memory objects.
+  const size_t metadata_size = std::accumulate(
+      metadata_strings_to_write_.begin(), metadata_strings_to_write_.end(), 0,
+      [](size_t sum, const std::string& str) {
+        return sum + sizeof(int64_t) + GetRoundedSize(str.size());
+      });
+  metadata_shared_memory_ =
+      std::make_unique<SharedMemory>(GetSharedMemoryMetadataName(name_));
+  metadata_shared_memory_->Create(metadata_size);
+  metadata_size_ = metadata_size;
+
+  // Write the metadata and tensor data to the shared memory.
+  WriteMetadataToSharedMemory();
+  data_shared_memory_ =
+      std::make_unique<SharedMemory>(GetSharedMemoryDataName(name_));
+  data_shared_memory_->Create(data_size);
+  data_size_ = data_size;
+  data_offset_ = 0;
+  for (auto tensor : tensors_to_write_) {
+    this->WriteTorchTensorInternal(tensor);
+  }
+
+  metadata_to_write_.clear();
+  tensors_to_write_.clear();
+}
+
+std::pair<SharedMemoryPtr, SharedMemoryPtr>
+SharedMemoryHelper::ReleaseSharedMemory() {
+  return std::make_pair(
+      std::move(metadata_shared_memory_), std::move(data_shared_memory_));
+}
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/shared_memory_helper.h

@@ -0,0 +1,153 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file shared_memory_helper.h
+ * @brief Share memory helper.
+ */
+#ifndef GRAPHBOLT_SHARED_MEMORY_HELPER_H_
+#define GRAPHBOLT_SHARED_MEMORY_HELPER_H_
+
+#include <graphbolt/shared_memory.h>
+#include <torch/torch.h>
+
+#include <memory>
+#include <sstream>
+#include <string>
+#include <tuple>
+#include <vector>
+
+namespace graphbolt {
+namespace sampling {
+
+/**
+ * @brief SharedMemoryHelper is a helper class to write/read data structures
+ * to/from shared memory.
+ *
+ * In order to write data structure to shared memory, we need to serialize the
+ * data structure to a binary buffer and then write the buffer to the shared
+ * memory. However, the size of the binary buffer is not known in advance. To
+ * solve this problem, we use two shared memory objects: one for storing the
+ * metadata and the other for storing the binary buffer. The metadata includes
+ * the metadata of data structures such as size and shape. The size of the
+ * metadata is decided by the size of metadata. The size of the binary buffer is
+ * decided by the size of the data structures.
+ *
+ * To avoid repeated shared memory allocation, this helper class uses lazy data
+ * structure writing. The data structures are written to the shared memory only
+ * when `Flush` is called. The data structures are written in the order of
+ * calling `WriteTorchArchive`, `WriteTorchTensor` and `WriteTorchTensorDict`,
+ * and also read in the same order.
+ *
+ * The usage of this class as a writer is as follows:
+ * @code{.cpp}
+ * SharedMemoryHelper shm_helper("shm_name", 1024, true);
+ * shm_helper.WriteTorchArchive(archive);
+ * shm_helper.WriteTorchTensor(tensor);
+ * shm_helper.WriteTorchTensorDict(tensor_dict);
+ * shm_helper.Flush();
+ * // After `Flush`, the data structures are written to the shared memory.
+ * // Then the helper class can be used as a reader.
+ * shm_helper.InitializeRead();
+ * auto archive = shm_helper.ReadTorchArchive();
+ * auto tensor = shm_helper.ReadTorchTensor();
+ * auto tensor_dict = shm_helper.ReadTorchTensorDict();
+ * @endcode
+ *
+ * The usage of this class as a reader is as follows:
+ * @code{.cpp}
+ * SharedMemoryHelper shm_helper("shm_name", 1024, false);
+ * shm_helper.InitializeRead();
+ * auto archive = shm_helper.ReadTorchArchive();
+ * auto tensor = shm_helper.ReadTorchTensor();
+ * auto tensor_dict = shm_helper.ReadTorchTensorDict();
+ * @endcode
+ *
+ *
+ */
+class SharedMemoryHelper {
+ public:
+  /**
+   * @brief Constructor of the shared memory helper.
+   * @param name The name of the shared memory.
+   */
+  SharedMemoryHelper(const std::string& name);
+
+  /** @brief Initialize this helper class before reading. */
+  void InitializeRead();
+
+  void WriteTorchArchive(torch::serialize::OutputArchive&& archive);
+  torch::serialize::InputArchive ReadTorchArchive();
+
+  void WriteTorchTensor(torch::optional<torch::Tensor> tensor);
+  torch::optional<torch::Tensor> ReadTorchTensor();
+
+  void WriteTorchTensorDict(
+      torch::optional<torch::Dict<std::string, torch::Tensor>> tensor_dict);
+  torch::optional<torch::Dict<std::string, torch::Tensor>>
+  ReadTorchTensorDict();
+
+  /** @brief Flush the data structures to the shared memory. */
+  void Flush();
+
+  /** @brief Release the shared memory and return their left values. */
+  std::pair<SharedMemoryPtr, SharedMemoryPtr> ReleaseSharedMemory();
+
+ private:
+  /**
+   * @brief Serialize metadata to string.
+   */
+  void SerializeMetadata();
+  /**
+   * @brief Write the metadata to the shared memory. This function is
+   * called by `Flush`.
+   */
+  void WriteMetadataToSharedMemory();
+  /**
+   * @brief Write the tensor data to the shared memory. This function is
+   * called by `Flush`.
+   */
+  void WriteTorchTensorInternal(torch::optional<torch::Tensor> tensor);
+
+  inline void* GetCurrentMetadataPtr() const {
+    return static_cast<char*>(metadata_shared_memory_->GetMemory()) +
+           metadata_offset_;
+  }
+  inline void* GetCurrentDataPtr() const {
+    return static_cast<char*>(data_shared_memory_->GetMemory()) + data_offset_;
+  }
+  inline void MoveMetadataPtr(int64_t offset) {
+    TORCH_CHECK(
+        metadata_offset_ + offset <= metadata_size_,
+        "The size of metadata exceeds the maximum size of shared memory.");
+    metadata_offset_ += offset;
+  }
+  inline void MoveDataPtr(int64_t offset) {
+    TORCH_CHECK(
+        data_offset_ + offset <= data_size_,
+        "The size of data exceeds the maximum size of shared memory.");
+    data_offset_ += offset;
+  }
+
+  std::string name_;
+  bool is_creator_;
+
+  size_t metadata_size_;
+  size_t data_size_;
+
+  // The shared memory objects for storing metadata and tensor data.
+  SharedMemoryPtr metadata_shared_memory_, data_shared_memory_;
+
+  // The read/write offsets of the metadata and tensor data.
+  size_t metadata_offset_, data_offset_;
+
+  // The data structures to write to the shared memory. They are written to the
+  // shared memory only when `Flush` is called.
+  std::vector<torch::serialize::OutputArchive> metadata_to_write_;
+  std::vector<std::string> metadata_strings_to_write_;
+  std::vector<torch::optional<torch::Tensor>> tensors_to_write_;
+};
+
+}  // namespace sampling
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_SHARED_MEMORY_HELPER_H_

graphbolt/src/unique_and_compact.cc

@@ -0,0 +1,99 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ *
+ * @file unique_and_compact.cc
+ * @brief Unique and compact op.
+ */
+
+#include <graphbolt/cuda_ops.h>
+#include <graphbolt/unique_and_compact.h>
+
+#include "./concurrent_id_hash_map.h"
+#include "./macro.h"
+#include "./utils.h"
+
+namespace graphbolt {
+namespace sampling {
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+UniqueAndCompact(
+    const torch::Tensor& src_ids, const torch::Tensor& dst_ids,
+    const torch::Tensor unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  if (utils::is_on_gpu(src_ids) && utils::is_on_gpu(dst_ids) &&
+      utils::is_on_gpu(unique_dst_ids)) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "unique_and_compact", {
+          return ops::UniqueAndCompact(
+              src_ids, dst_ids, unique_dst_ids, rank, world_size);
+        });
+  }
+  TORCH_CHECK(
+      world_size <= 1,
+      "Cooperative Minibatching (arXiv:2310.12403) is supported only on GPUs.");
+  auto num_dst = unique_dst_ids.size(0);
+  torch::Tensor ids = torch::cat({unique_dst_ids, src_ids});
+  auto [unique_ids, compacted_src, compacted_dst] = AT_DISPATCH_INDEX_TYPES(
+      ids.scalar_type(), "unique_and_compact", ([&] {
+        ConcurrentIdHashMap<index_t> id_map(ids, num_dst);
+        return std::make_tuple(
+            id_map.GetUniqueIds(), id_map.MapIds(src_ids),
+            id_map.MapIds(dst_ids));
+      }));
+  auto offsets = torch::zeros(2, c10::TensorOptions().dtype(torch::kInt64));
+  offsets.data_ptr<int64_t>()[1] = unique_ids.size(0);
+  return {unique_ids, compacted_src, compacted_dst, offsets};
+}
+
+std::vector<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+UniqueAndCompactBatched(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor> unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  TORCH_CHECK(
+      src_ids.size() == dst_ids.size() &&
+          dst_ids.size() == unique_dst_ids.size(),
+      "The batch dimension of the parameters need to be identical.");
+  bool all_on_gpu = true;
+  for (std::size_t i = 0; i < src_ids.size(); i++) {
+    all_on_gpu = all_on_gpu && utils::is_on_gpu(src_ids[i]) &&
+                 utils::is_on_gpu(dst_ids[i]) &&
+                 utils::is_on_gpu(unique_dst_ids[i]);
+    if (!all_on_gpu) break;
+  }
+  if (all_on_gpu) {
+    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
+        c10::DeviceType::CUDA, "unique_and_compact", {
+          return ops::UniqueAndCompactBatched(
+              src_ids, dst_ids, unique_dst_ids, rank, world_size);
+        });
+  }
+  std::vector<
+      std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
+      results;
+  results.reserve(src_ids.size());
+  for (std::size_t i = 0; i < src_ids.size(); i++) {
+    results.emplace_back(UniqueAndCompact(
+        src_ids[i], dst_ids[i], unique_dst_ids[i], rank, world_size));
+  }
+  return results;
+}
+
+c10::intrusive_ptr<Future<std::vector<
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>>>
+UniqueAndCompactBatchedAsync(
+    const std::vector<torch::Tensor>& src_ids,
+    const std::vector<torch::Tensor>& dst_ids,
+    const std::vector<torch::Tensor> unique_dst_ids, const int64_t rank,
+    const int64_t world_size) {
+  return async(
+      [=] {
+        return UniqueAndCompactBatched(
+            src_ids, dst_ids, unique_dst_ids, rank, world_size);
+      },
+      utils::is_on_gpu(src_ids.at(0)));
+}
+
+}  // namespace sampling
+}  // namespace graphbolt

graphbolt/src/utils.cc

@@ -0,0 +1,36 @@
+/**
+ *   Copyright (c) 2024, GT-TDAlab (Muhammed Fatih Balin & Umit V. Catalyurek)
+ *   All rights reserved.
+ *
+ *   Licensed under the Apache License, Version 2.0 (the "License");
+ *   you may not use this file except in compliance with the License.
+ *   You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *   Unless required by applicable law or agreed to in writing, software
+ *   distributed under the License is distributed on an "AS IS" BASIS,
+ *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *   See the License for the specific language governing permissions and
+ *   limitations under the License.
+ *
+ * @file utils.cc
+ * @brief Graphbolt utils implementations.
+ */
+#include "./utils.h"
+
+#include <optional>
+
+namespace graphbolt {
+namespace utils {
+
+namespace {
+std::optional<int64_t> worker_id;
+}
+
+std::optional<int64_t> GetWorkerId() { return worker_id; }
+
+void SetWorkerId(int64_t worker_id_value) { worker_id = worker_id_value; }
+
+}  // namespace utils
+}  // namespace graphbolt

graphbolt/src/utils.h

@@ -0,0 +1,103 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file utils.h
+ * @brief Graphbolt utils.
+ */
+
+#ifndef GRAPHBOLT_UTILS_H_
+#define GRAPHBOLT_UTILS_H_
+
+#include <torch/script.h>
+
+#include <optional>
+
+namespace graphbolt {
+namespace utils {
+
+/**
+ * @brief If this process is a worker part as part of a DataLoader, then returns
+ * the assigned worker id less than the # workers.
+ */
+std::optional<int64_t> GetWorkerId();
+
+/**
+ * @brief If this process is a worker part as part of a DataLoader, then this
+ * function is called to initialize its worked id to be less than the # workers.
+ */
+void SetWorkerId(int64_t worker_id_value);
+
+/**
+ * @brief Checks whether the tensor is stored on the GPU.
+ */
+inline bool is_on_gpu(const torch::Tensor& tensor) {
+  return tensor.device().is_cuda();
+}
+
+/**
+ * @brief Checks whether the tensor is stored on the GPU or the pinned memory.
+ */
+inline bool is_accessible_from_gpu(const torch::Tensor& tensor) {
+  return is_on_gpu(tensor) || tensor.is_pinned();
+}
+
+/**
+ * @brief Checks whether the tensor is stored on the pinned memory.
+ */
+inline bool is_pinned(const torch::Tensor& tensor) {
+  // If this process is a worker, we should avoid initializing the CUDA context.
+  return !GetWorkerId() && tensor.is_pinned();
+}
+
+/**
+ * @brief Checks whether the tensors are all stored on the GPU or the pinned
+ * memory.
+ */
+template <typename TensorContainer>
+inline bool are_accessible_from_gpu(const TensorContainer& tensors) {
+  for (auto& tensor : tensors) {
+    if (!is_accessible_from_gpu(tensor)) return false;
+  }
+  return true;
+}
+
+/**
+ * @brief Parses the source and destination node type from a given edge type
+ * triple seperated with ":".
+ */
+inline std::pair<std::string, std::string> parse_src_dst_ntype_from_etype(
+    std::string etype) {
+  auto first_seperator_it = std::find(etype.begin(), etype.end(), ':');
+  auto second_seperator_pos =
+      std::find(first_seperator_it + 1, etype.end(), ':') - etype.begin();
+  return {
+      etype.substr(0, first_seperator_it - etype.begin()),
+      etype.substr(second_seperator_pos + 1)};
+}
+
+/**
+ * @brief Retrieves the value of the tensor at the given index.
+ *
+ * @note If the tensor is not contiguous, it will be copied to a contiguous
+ * tensor.
+ *
+ * @tparam T The type of the tensor.
+ * @param tensor The tensor.
+ * @param index The index.
+ *
+ * @return T The value of the tensor at the given index.
+ */
+template <typename T>
+T GetValueByIndex(const torch::Tensor& tensor, int64_t index) {
+  TORCH_CHECK(
+      index >= 0 && index < tensor.numel(),
+      "The index should be within the range of the tensor, but got index ",
+      index, " and tensor size ", tensor.numel());
+  auto contiguous_tensor = tensor.contiguous();
+  auto data_ptr = contiguous_tensor.data_ptr<T>();
+  return data_ptr[index];
+}
+
+}  // namespace utils
+}  // namespace graphbolt
+
+#endif  // GRAPHBOLT_UTILS_H_