wehub-resource-sync
273 lines
9.3 KiB
C++
273 lines
9.3 KiB
C++
/* ******************************************************************************
|
|
*
|
|
* 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.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* 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 DIRECTSHAPETRIE_H
|
|
#define DIRECTSHAPETRIE_H
|
|
|
|
#include <array/ConstantShapeBuffer.h>
|
|
#include <system/common.h>
|
|
#include "ShapeBufferPlatformHelper.h"
|
|
#include <array>
|
|
#include <atomic>
|
|
#include <memory>
|
|
#include <unordered_set>
|
|
#include <vector>
|
|
#include "./generic/StripedLocks.h"
|
|
namespace sd {
|
|
|
|
|
|
|
|
class SD_LIB_EXPORT ShapeTrieNode {
|
|
private:
|
|
std::vector<ShapeTrieNode*> _children;
|
|
LongType _value;
|
|
int _level;
|
|
bool _isShape;
|
|
ConstantShapeBuffer* _buffer = nullptr; // Now accessed atomically
|
|
int _shapeHash; // Store a hash of the shape for validation
|
|
|
|
#if defined(SD_GCC_FUNCTRACE)
|
|
backward::StackTrace storeStackTrace;
|
|
#endif
|
|
|
|
public:
|
|
ShapeTrieNode(LongType value = 0, int level = 0, bool isShape = true, sd::LongType shapeHash = 0)
|
|
: _value(value), _level(level), _isShape(isShape), _shapeHash(shapeHash) {
|
|
}
|
|
|
|
~ShapeTrieNode() {
|
|
// Delete children
|
|
for (auto* child : _children) {
|
|
delete child;
|
|
}
|
|
_children.clear();
|
|
|
|
// CRITICAL FIX: Respect reference counting before deleting buffers!
|
|
// The buffer should ONLY be deleted if refcount == 1 (cache is the only owner).
|
|
// If refcount > 1, other NDArrays still hold references - deleting would create
|
|
// dangling pointers that cause SIGSEGV crashes when accessed later.
|
|
//
|
|
// This is especially important during JVM shutdown where static destruction order
|
|
// is unpredictable and NDArrays may outlive the cache singleton.
|
|
if (_buffer != nullptr) {
|
|
int refCount = _buffer->getRefCount();
|
|
if (refCount <= 1) {
|
|
// Only cache owns it - safe to delete
|
|
delete _buffer;
|
|
}
|
|
// If refCount > 1, intentionally leak to prevent use-after-free crashes.
|
|
// This is the correct behavior for cache cleanup with external references.
|
|
_buffer = nullptr;
|
|
}
|
|
}
|
|
|
|
|
|
// Replace the current findOrCreateChild with this more defensive version
|
|
ShapeTrieNode* findOrCreateChild(LongType value, int level, bool isShape, int shapeHash) {
|
|
// First search for existing child
|
|
for (auto* child : _children) {
|
|
if (child != nullptr &&
|
|
child->value() == value &&
|
|
child->level() == level &&
|
|
child->isShape() == isShape &&
|
|
(shapeHash == 0 || child->shapeHash() == shapeHash)) {
|
|
return child;
|
|
}
|
|
}
|
|
|
|
// Create new node with standard allocation
|
|
try {
|
|
auto* newNode = new ShapeTrieNode(value, level, isShape, shapeHash);
|
|
// Add to children using standard vector operations
|
|
_children.push_back(newNode);
|
|
return newNode;
|
|
} catch (const std::exception& e) {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
const std::vector<ShapeTrieNode*>& children() const { return _children; }
|
|
LongType value() const { return _value; }
|
|
int level() const { return _level; }
|
|
bool isShape() const { return _isShape; }
|
|
int shapeHash() const { return _shapeHash; }
|
|
|
|
void setBuffer(ConstantShapeBuffer* buf);
|
|
ConstantShapeBuffer* buffer() const { return _buffer; }
|
|
|
|
#if defined(SD_GCC_FUNCTRACE)
|
|
void collectStoreStackTrace();
|
|
#endif
|
|
};
|
|
|
|
class SD_LIB_EXPORT DirectShapeTrie {
|
|
private:
|
|
static const size_t NUM_STRIPES = 256; // Increased from 32 to reduce collisions
|
|
std::array<ShapeTrieNode*, NUM_STRIPES> *_roots;
|
|
std::array<MUTEX_TYPE*, NUM_STRIPES> *_mutexes = nullptr;
|
|
|
|
// Initialization guards to prevent partially-constructed state from leaking
|
|
std::atomic<bool> _initialization_complete{false};
|
|
std::atomic<bool> _initialization_in_progress{false};
|
|
|
|
// Cache statistics tracking
|
|
mutable std::atomic<LongType> _current_entries{0};
|
|
mutable std::atomic<LongType> _current_bytes{0};
|
|
mutable std::atomic<LongType> _peak_entries{0};
|
|
mutable std::atomic<LongType> _peak_bytes{0};
|
|
|
|
// Helper method to create a fallback buffer when trie insertion fails
|
|
// Always returns a valid shape buffer or throws an exception
|
|
ConstantShapeBuffer* createFallbackBuffer(const LongType* shapeInfo, int rank);
|
|
|
|
// Internal helper to recursively count entries and bytes in a subtrie
|
|
void countEntriesAndBytes(const ShapeTrieNode* node, LongType& entries, LongType& bytes) const;
|
|
|
|
// Internal helper to build string representation recursively
|
|
void buildStringRepresentation(const ShapeTrieNode* node, std::stringstream& ss,
|
|
const std::string& indent, int currentDepth,
|
|
int maxDepth, int& entriesShown, int maxEntries) const;
|
|
|
|
public:
|
|
// Constructor
|
|
DirectShapeTrie() {
|
|
_initialization_in_progress.store(true, std::memory_order_release);
|
|
_roots = new std::array<ShapeTrieNode*, NUM_STRIPES>();
|
|
_mutexes = new std::array<MUTEX_TYPE*, NUM_STRIPES>();
|
|
|
|
for (size_t i = 0; i < NUM_STRIPES; i++) {
|
|
(*_roots)[i] = new ShapeTrieNode(0, 0, false);
|
|
// Allocate mutexes on the heap
|
|
(*_mutexes)[i] = new MUTEX_TYPE();
|
|
}
|
|
|
|
ShapeBufferPlatformHelper::initialize();
|
|
_initialization_in_progress.store(false, std::memory_order_release);
|
|
_initialization_complete.store(true, std::memory_order_release);
|
|
}
|
|
|
|
// Delete copy constructor and assignment
|
|
DirectShapeTrie(const DirectShapeTrie&) = delete;
|
|
DirectShapeTrie& operator=(const DirectShapeTrie&) = delete;
|
|
|
|
~DirectShapeTrie() {
|
|
// Clean up all mutexes
|
|
if (_mutexes != nullptr) {
|
|
for (size_t i = 0; i < NUM_STRIPES; i++) {
|
|
if ((*_mutexes)[i] != nullptr) {
|
|
delete (*_mutexes)[i];
|
|
}
|
|
}
|
|
delete _mutexes;
|
|
}
|
|
|
|
// Clean up all root nodes
|
|
if (_roots != nullptr) {
|
|
for (size_t i = 0; i < NUM_STRIPES; i++) {
|
|
if ((*_roots)[i] != nullptr) {
|
|
delete (*_roots)[i];
|
|
}
|
|
}
|
|
delete _roots;
|
|
}
|
|
}
|
|
|
|
// Delete move operations
|
|
DirectShapeTrie(DirectShapeTrie&&) = delete;
|
|
DirectShapeTrie& operator=(DirectShapeTrie&&) = delete;
|
|
|
|
// Improved thread-safe getOrCreate
|
|
// Always returns a valid shape buffer or throws an exception
|
|
ConstantShapeBuffer* getOrCreate(const LongType* shapeInfo);
|
|
|
|
// Check if a shape info already exists in the trie
|
|
bool exists(const LongType* shapeInfo) const;
|
|
|
|
// Wait until constructor finishes allocating internal state.
|
|
void waitForInitialization() const;
|
|
|
|
// Helper methods
|
|
size_t computeHash(const LongType* shapeInfo) const;
|
|
size_t getStripeIndex(const LongType* shapeInfo) const;
|
|
bool shapeInfoEqual(const LongType* a, const LongType* b) const;
|
|
void validateShapeInfo(const LongType* shapeInfo) const;
|
|
ConstantShapeBuffer* search(const LongType* shapeInfo, size_t stripeIdx) const;
|
|
ConstantShapeBuffer* insert(const LongType* shapeInfo, size_t stripeIdx);
|
|
const ShapeTrieNode* findChild(const ShapeTrieNode* node, LongType value, int level, bool isShape, int shapeHash = 0) const;
|
|
|
|
// Calculate a unique shape signature for additional validation
|
|
int calculateShapeSignature(const LongType* shapeInfo) const;
|
|
|
|
// Clear all cached shape buffers
|
|
void clearCache();
|
|
|
|
/**
|
|
* Get the total number of cached shape entries.
|
|
*
|
|
* @return Total number of cached shape buffers across all stripes
|
|
*/
|
|
LongType getCachedEntries() const;
|
|
|
|
/**
|
|
* Get the total memory used by cached shape buffers in bytes.
|
|
* This includes the shape_info buffer sizes across all cached entries.
|
|
*
|
|
* @return Total memory used in bytes
|
|
*/
|
|
LongType getCachedBytes() const;
|
|
|
|
/**
|
|
* Get the peak number of entries that were cached simultaneously.
|
|
*
|
|
* @return Peak number of cached entries
|
|
*/
|
|
LongType getPeakCachedEntries() const;
|
|
|
|
/**
|
|
* Get the peak memory usage by cached shape buffers in bytes.
|
|
*
|
|
* @return Peak memory usage in bytes
|
|
*/
|
|
LongType getPeakCachedBytes() const;
|
|
|
|
/**
|
|
* Generate a human-readable string representation of the trie structure.
|
|
* Shows the hierarchy of nodes and cached shape buffers for debugging.
|
|
*
|
|
* @param maxDepth Maximum depth to traverse (default: 10, -1 for unlimited)
|
|
* @param maxEntries Maximum number of entries to show (default: 100, -1 for unlimited)
|
|
* @return String representation of the trie
|
|
*/
|
|
std::string toString(int maxDepth = 10, int maxEntries = 100) const;
|
|
|
|
/**
|
|
* Get all ConstantShapeBuffer pointers currently in the cache.
|
|
* This is used by lifecycle tracking to distinguish cached entries from real leaks.
|
|
*
|
|
* @param out_pointers Set to fill with pointers to all cached ConstantShapeBuffer objects
|
|
*/
|
|
void getCachedPointers(std::unordered_set<void*>& out_pointers) const;
|
|
|
|
private:
|
|
// Internal helper to collect ConstantShapeBuffer pointers recursively
|
|
void collectCachedPointers(const ShapeTrieNode* node, std::unordered_set<void*>& out_pointers) const;
|
|
};
|
|
|
|
} // namespace sd
|
|
|
|
#endif // DIRECTSHAPETRIE_H
|