wehub-resource-sync
136 lines
4.4 KiB
C++
136 lines
4.4 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 DEV_TESTS_CONSTANTTADHELPER_H
|
|
#define DEV_TESTS_CONSTANTTADHELPER_H
|
|
|
|
#include <array/ConstantDataBuffer.h>
|
|
#include <array/ConstantDescriptor.h>
|
|
#include <array/ShapeDescriptor.h>
|
|
#include <array/TadDescriptor.h>
|
|
#include <array/TadPack.h>
|
|
#include <helpers/DirectTadTrie.h>
|
|
#include <system/op_boilerplate.h>
|
|
|
|
#include <map>
|
|
#include <mutex>
|
|
#include <unordered_set>
|
|
#include <vector>
|
|
|
|
namespace sd {
|
|
class SD_LIB_EXPORT ConstantTadHelper {
|
|
private:
|
|
DirectTadTrie _trie; // Single trie for device 0
|
|
|
|
ConstantTadHelper() = default;
|
|
|
|
public:
|
|
~ConstantTadHelper() = default;
|
|
|
|
static ConstantTadHelper &getInstance();
|
|
|
|
/**
|
|
* These methods calculate Tensor-Along-Dimension(s) shape and offsets
|
|
*
|
|
* @param originalShape
|
|
* @param dimensions
|
|
* @param keepUnitiesInShape
|
|
* @return
|
|
*/
|
|
|
|
std::shared_ptr<TadPack> tadForDimensions(LongType *originalShape, LongType *dimensions, LongType dimLength);
|
|
std::shared_ptr<TadPack> tadForDimensions(ShapeDescriptor &descriptor, std::vector<LongType> &dimensions,
|
|
const bool keepUnitiesInShape = false);
|
|
std::shared_ptr<TadPack> tadForDimensions(TadDescriptor *descriptor);
|
|
|
|
std::shared_ptr<TadPack> tadForDimensions(LongType *originalShape, LongType dimension);
|
|
std::shared_ptr<TadPack> tadForDimensions(LongType *originalShape, std::vector<LongType> *dimensions);
|
|
|
|
/**
|
|
* Clear all cached TAD packs to prevent memory leaks during testing.
|
|
* NOTE: Will return early without action if setShutdownInProgress(true) was called.
|
|
*/
|
|
void clearCache();
|
|
|
|
/**
|
|
* Mark that shutdown is in progress.
|
|
* When true, clearCache() will skip tree traversal to avoid SIGSEGV from corrupted memory
|
|
* during JVM/static destruction.
|
|
* @param inProgress true to mark shutdown in progress, false otherwise
|
|
*/
|
|
void setShutdownInProgress(bool inProgress) {
|
|
_trie.setShutdownInProgress(inProgress);
|
|
}
|
|
|
|
/**
|
|
* Check if shutdown is in progress.
|
|
* @return true if shutdown is marked as in progress
|
|
*/
|
|
bool isShutdownInProgress() const {
|
|
return _trie.isShutdownInProgress();
|
|
}
|
|
|
|
/**
|
|
* Get the total number of cached TAD pack entries.
|
|
*
|
|
* @return Total number of cached TAD packs across all stripes
|
|
*/
|
|
LongType getCachedEntries() const;
|
|
|
|
/**
|
|
* Get the total memory used by cached TAD packs in bytes.
|
|
* This includes both shape_info and offset buffer sizes.
|
|
*
|
|
* @return Total memory used in bytes
|
|
*/
|
|
LongType getCachedBytes() const;
|
|
|
|
/**
|
|
* Get the peak number of TAD pack entries that were cached simultaneously.
|
|
*
|
|
* @return Peak number of cached TAD packs
|
|
*/
|
|
LongType getPeakCachedEntries() const;
|
|
|
|
/**
|
|
* Get the peak memory usage by cached TAD packs in bytes.
|
|
*
|
|
* @return Peak memory usage in bytes
|
|
*/
|
|
LongType getPeakCachedBytes() const;
|
|
|
|
/**
|
|
* Generate a human-readable string representation of the TAD cache.
|
|
* Shows the trie structure with cached TAD packs 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 cache
|
|
*/
|
|
std::string toString(int maxDepth = 10, int maxEntries = 100) const;
|
|
|
|
/**
|
|
* Get all TadPack 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 TadPack objects
|
|
*/
|
|
void getCachedPointers(std::unordered_set<void*>& out_pointers) const;
|
|
};
|
|
} // namespace sd
|
|
|
|
#endif // DEV_TESTS_CONSTANTTADHELPER_H
|