chore: import upstream snapshot with attribution

libnd4j/include/helpers/generic/ResourceManager.h

@@ -0,0 +1,93 @@
+/* ******************************************************************************
+ *
+ * Copyright (c) 2024 Konduit K.K.
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+#ifndef SD_GENERIC_RESOURCE_MANAGER_H_
+#define SD_GENERIC_RESOURCE_MANAGER_H_
+
+#include <system/common.h>
+#include <atomic>
+#include <chrono>
+#include <memory>
+#include <mutex>
+
+namespace sd {
+namespace generic {
+
+class ResourceManager {
+private:
+    mutable std::atomic<size_t> _totalNodes{0};
+    std::atomic<size_t> _activeOperations{0};
+    std::chrono::steady_clock::time_point _lastCleanup;
+    std::mutex _cleanupMutex;
+    static constexpr size_t CLEANUP_THRESHOLD = 1000;
+    static constexpr auto CLEANUP_INTERVAL = std::chrono::minutes(5);
+
+public:
+    class OperationScope {
+    private:
+        ResourceManager& _manager;
+    public:
+        explicit OperationScope(ResourceManager& m) : _manager(m) {
+            _manager._activeOperations.fetch_add(1, std::memory_order_acquire);
+        }
+        ~OperationScope() {
+            _manager._activeOperations.fetch_sub(1, std::memory_order_release);
+        }
+    };
+
+    void registerNode() const {
+        _totalNodes.fetch_add(1, std::memory_order_relaxed);
+    }
+
+    void unregisterNode() const {
+        _totalNodes.fetch_sub(1, std::memory_order_relaxed);
+    }
+
+    OperationScope createScope() const {
+        ResourceManager &m = const_cast<ResourceManager&>(*this);
+        return OperationScope(m);
+    }
+
+    size_t activeOperations() const {
+        return _activeOperations.load(std::memory_order_relaxed);
+    }
+
+    size_t totalNodes() const {
+        return _totalNodes.load(std::memory_order_relaxed);
+    }
+
+    void resetCleanupTimer() {
+        std::lock_guard<std::mutex> lock(_cleanupMutex);
+        _lastCleanup = std::chrono::steady_clock::now();
+    }
+
+    bool shouldCleanup() {
+        if (_totalNodes.load(std::memory_order_relaxed) < CLEANUP_THRESHOLD)
+            return false;
+
+        auto now = std::chrono::steady_clock::now();
+        std::lock_guard<std::mutex> lock(_cleanupMutex);
+        if (now - _lastCleanup >= CLEANUP_INTERVAL) {
+            _lastCleanup = now;
+            return true;
+        }
+        return false;
+    }
+};
+
+}} // namespace sd::generic
+
+#endif // SD_GENERIC_RESOURCE_MANAGER_H_

libnd4j/include/helpers/generic/StripedLocks.h

@@ -0,0 +1,228 @@
+/* ******************************************************************************
+*
+* Copyright (c) 2024 Konduit K.K.
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://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.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+#ifndef SD_STRIPED_LOCKS_H_
+#define SD_STRIPED_LOCKS_H_
+
+#include <algorithm>
+#include <array>
+#include <atomic>
+#if __cplusplus >= 201703L && (!defined(__APPLE__))
+#include <shared_mutex>
+#define SD_HAS_SHARED_MUTEX 1
+#else
+#include <mutex>
+#define SD_HAS_SHARED_MUTEX 0
+#endif
+
+#include <stdexcept>
+#include <thread>
+#include <vector>
+
+#include "system/op_boilerplate.h"
+
+namespace sd {
+namespace generic {
+
+#if SD_HAS_SHARED_MUTEX
+#define MUTEX_TYPE std::shared_mutex
+#define SHARED_LOCK_TYPE std::shared_lock
+#define EXCLUSIVE_LOCK_TYPE std::unique_lock
+#else
+#define MUTEX_TYPE std::mutex
+#define SHARED_LOCK_TYPE std::lock_guard
+#define EXCLUSIVE_LOCK_TYPE std::lock_guard
+#endif
+
+template<size_t NUM_STRIPES = 32>
+class StripedLocks {
+ private:
+  mutable std::array<MUTEX_TYPE, NUM_STRIPES> _mutexes;
+  mutable std::array<std::atomic<uint32_t>, NUM_STRIPES> _stripeCounts{};
+  static constexpr int MAX_RETRIES = 500;
+  static constexpr auto RETRY_DELAY = std::chrono::microseconds(50);
+  static constexpr int MAX_BACKOFF_SHIFT = 10;
+  static constexpr auto MIN_DELAY = std::chrono::microseconds(50);
+
+  static_assert((NUM_STRIPES & (NUM_STRIPES - 1)) == 0, "NUM_STRIPES must be a power of 2");
+
+  bool acquireLockWithTimeout(size_t stripe, bool exclusive, int maxRetries) const {
+    auto startTime = std::chrono::steady_clock::now();
+
+    for (int attempt = 0; attempt < maxRetries; ++attempt) {
+      bool locked;
+
+#if SD_HAS_SHARED_MUTEX
+      locked = exclusive ?
+                         _mutexes[stripe].try_lock() :
+                         _mutexes[stripe].try_lock_shared();
+#else
+      // With std::mutex, we can only do exclusive locks
+      locked = _mutexes[stripe].try_lock();
+#endif
+
+      if (locked) {
+        _stripeCounts[stripe].fetch_add(1, std::memory_order_acq_rel);
+        return true;
+      }
+
+      if (std::chrono::steady_clock::now() - startTime > std::chrono::seconds(1)) {
+        return false;
+      }
+
+      auto backoff = RETRY_DELAY * (1 << std::min(attempt, MAX_BACKOFF_SHIFT));
+      std::this_thread::sleep_for(backoff);
+    }
+    return false;
+  }
+
+ public:
+  StripedLocks() {
+    initializeCounts();
+  }
+
+  class MultiLockGuard {
+   private:
+    StripedLocks& _locks;
+    std::vector<size_t> _stripes;
+    bool _exclusive;
+    bool _acquired{false};
+    static constexpr auto LOCK_TIMEOUT = std::chrono::seconds(2);
+
+    bool acquireAllLocksWithTimeout(const std::chrono::milliseconds& timeout) {
+      auto startTime = std::chrono::steady_clock::now();
+      std::vector<size_t> acquired;
+      acquired.reserve(_stripes.size());
+
+      while (std::chrono::steady_clock::now() - startTime < timeout) {
+        bool allLocked = true;
+        for (size_t stripe : _stripes) {
+          if (!_locks.acquireLockWithTimeout(stripe, _exclusive, 1)) {
+            allLocked = false;
+            for (auto& s : acquired) {
+              _locks.unlockStripe(s, _exclusive);
+            }
+            acquired.clear();
+            break;
+          }
+          acquired.push_back(stripe);
+        }
+        if (allLocked) {
+          _acquired = true;
+          return true;
+        }
+        std::this_thread::sleep_for(MIN_DELAY);
+      }
+      return false;
+    }
+
+   public:
+    MultiLockGuard(StripedLocks& locks,
+                   const std::vector<size_t>& stripes,
+                   bool exclusive,
+                   const std::chrono::milliseconds& timeout = std::chrono::seconds(2))
+        : _locks(locks), _stripes(stripes), _exclusive(exclusive) {
+      std::sort(_stripes.begin(), _stripes.end());
+      _stripes.erase(std::unique(_stripes.begin(), _stripes.end()), _stripes.end());
+      acquireAllLocksWithTimeout(timeout);
+    }
+
+    bool acquired() const { return _acquired; }
+
+    void release() {
+      if (!_acquired) return;
+      for (auto it = _stripes.rbegin(); it != _stripes.rend(); ++it) {
+        _locks.unlockStripe(*it, _exclusive);
+      }
+      _acquired = false;
+    }
+
+    ~MultiLockGuard() {
+      if (_acquired) {
+        release();
+      }
+    }
+
+    MultiLockGuard(const MultiLockGuard&) = delete;
+    MultiLockGuard& operator=(const MultiLockGuard&) = delete;
+    MultiLockGuard(MultiLockGuard&&) noexcept = default;
+  };
+
+  MultiLockGuard acquireMultiLockWithTimeout(const std::vector<size_t>& stripes,
+                                             bool exclusive,
+                                             const std::chrono::milliseconds& timeout) {
+    return MultiLockGuard(*this, stripes, exclusive, timeout);
+  }
+
+  void lockStripe(size_t stripe, bool exclusive = false) const {
+    if (stripe >= NUM_STRIPES) {
+      THROW_EXCEPTION("Invalid stripe index");
+    }
+
+    if (!acquireLockWithTimeout(stripe, exclusive, MAX_RETRIES)) {
+      auto currentCount = _stripeCounts[stripe].load(std::memory_order_relaxed);
+      std::string msg = "Failed to acquire " +
+                        std::string(exclusive ? "exclusive" : "shared") +
+                        " lock for stripe " + std::to_string(stripe) +
+                        " after " + std::to_string(MAX_RETRIES) +
+                        " attempts. Current count: " + std::to_string(currentCount);
+      THROW_EXCEPTION(msg.c_str());
+    }
+  }
+
+  void unlockStripe(size_t stripe, bool exclusive = false) const {
+    if (stripe >= NUM_STRIPES) {
+      THROW_EXCEPTION("Invalid stripe index");
+    }
+    _stripeCounts[stripe].fetch_sub(1, std::memory_order_relaxed);
+
+#if SD_HAS_SHARED_MUTEX
+    if (exclusive) {
+      _mutexes[stripe].unlock();
+    } else {
+      _mutexes[stripe].unlock_shared();
+    }
+#else
+    // With std::mutex, we can only do exclusive unlocks
+    _mutexes[stripe].unlock();
+#endif
+  }
+
+  MultiLockGuard acquireMultiLock(const std::vector<size_t>& stripes, bool exclusive = false) {
+    return MultiLockGuard(*this, stripes, exclusive);
+  }
+
+  uint32_t getStripeCount(size_t stripe) const {
+    return _stripeCounts[stripe].load(std::memory_order_relaxed);
+  }
+
+  template<typename T>
+  size_t getStripeIndex(const T& value) const {
+    return std::hash<T>{}(value) & (NUM_STRIPES - 1);
+  }
+
+  void initializeCounts() {
+    for (auto& count : _stripeCounts) {
+      count.store(0, std::memory_order_relaxed);
+    }
+  }
+
+  static constexpr size_t getNumStripes() { return NUM_STRIPES; }
+};
+
+}} // namespace sd::generic
+
+#endif // SD_STRIPED_LOCKS_H_

libnd4j/include/helpers/generic/TypedTrie.h

@@ -0,0 +1,82 @@
+/* ******************************************************************************
+*
+* Copyright (c) 2024 Konduit K.K.
+* This program and the accompanying materials are made available under the
+* terms of the Apache License, Version 2.0 which is available at
+* https://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.
+*
+* SPDX-License-Identifier: Apache-2.0
+******************************************************************************/
+
+#ifndef SD_GENERIC_TYPED_TRIE_H_
+#define SD_GENERIC_TYPED_TRIE_H_
+
+#include <system/common.h>
+
+#include <array>
+#include <atomic>
+#include <cassert>
+#include <chrono>
+#include <memory>
+#include <mutex>
+#if __cplusplus >= 201703L && (!defined(__APPLE__))
+#include <shared_mutex>
+#endif
+#include <thread>
+#include <type_traits>
+#include <unordered_map>
+#include <vector>
+
+#include "StripedLocks.h"
+#include "ResourceManager.h"
+namespace sd {
+namespace generic {
+
+template<typename KeyType, typename ValueType, size_t NUM_STRIPES = 32>
+class SD_LIB_EXPORT TypedTrie {
+private:
+ struct TrieNode;
+
+ std::array<std::unique_ptr<TrieNode>, NUM_STRIPES> _roots;
+ mutable StripedLocks<NUM_STRIPES> _locks;
+ ResourceManager _resourceManager;
+
+ size_t getStripeIndex(const KeyType& key) const;
+ TrieNode* findOrCreateNode(TrieNode* root, const KeyType& key, bool createIfMissing) const;
+ void cleanupNode(TrieNode* node);
+
+public:
+ ~TypedTrie();
+ TypedTrie();
+
+ TypedTrie(const TypedTrie&) = delete;
+ TypedTrie& operator=(const TypedTrie&) = delete;
+
+ // Allow move operations
+ TypedTrie(TypedTrie&&) = default;
+ TypedTrie& operator=(TypedTrie&&) = default;
+
+ std::shared_ptr<ValueType> get(const KeyType& key) const;
+ bool insert(const KeyType& key, std::shared_ptr<ValueType> value);
+ bool remove(const KeyType& key);
+ void cleanup();
+
+ struct Stats {
+   size_t totalNodes;
+   size_t activeOperations;
+   std::array<uint32_t, NUM_STRIPES> stripeCounts;
+ };
+
+ Stats getStats() const;
+};
+
+}  // namespace generic
+}  // namespace sd
+
+#endif  // SD_GENERIC_TYPED_TRIE_H_