wehub-resource-sync
140 lines
4.6 KiB
C++
140 lines
4.6 KiB
C++
/* ******************************************************************************
|
|
*
|
|
* 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
|
|
******************************************************************************/
|
|
|
|
#include <array/TadDescriptor.h>
|
|
#include <array/TadPack.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
|
|
namespace sd {
|
|
|
|
ConstantTadHelper& ConstantTadHelper::getInstance() {
|
|
static ConstantTadHelper instance;
|
|
return instance;
|
|
}
|
|
|
|
std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType dimension) {
|
|
return tadForDimensions(originalShape, &dimension, 1);
|
|
}
|
|
|
|
std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, std::vector<LongType>* dimensions) {
|
|
if (dimensions == nullptr) {
|
|
THROW_EXCEPTION("Dimensions vector is null");
|
|
}
|
|
return tadForDimensions(originalShape, const_cast<LongType*>(dimensions->data()), dimensions->size());
|
|
}
|
|
|
|
std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(TadDescriptor* descriptor) {
|
|
if (descriptor == nullptr) {
|
|
THROW_EXCEPTION("TadDescriptor is null");
|
|
}
|
|
return tadForDimensions(descriptor->originalShape(), descriptor->axis().data(),
|
|
descriptor->axis().size());
|
|
}
|
|
|
|
std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType* dimensions, LongType dimLength) {
|
|
if (originalShape == nullptr) {
|
|
THROW_EXCEPTION("Original shape is null");
|
|
}
|
|
|
|
if (dimensions == nullptr && dimLength > 0) {
|
|
THROW_EXCEPTION("Dimensions array is null but dimLength > 0");
|
|
}
|
|
|
|
// Check for empty array
|
|
if (shape::isEmptyConst(originalShape)) {
|
|
THROW_EXCEPTION("Cannot create TADs for empty array");
|
|
}
|
|
|
|
sd::LongType rank = shape::rank(originalShape);
|
|
if (rank < 0) {
|
|
THROW_EXCEPTION("Invalid shape rank");
|
|
}
|
|
|
|
// Check for zero-sized dimensions
|
|
for (LongType i = 0; i < rank; i++) {
|
|
if (shape::sizeAt(originalShape, i) == 0) {
|
|
THROW_EXCEPTION("Cannot create TADs for array with zero-sized dimensions");
|
|
}
|
|
}
|
|
|
|
// Handle zero dimension length case - treat entire array as single TAD
|
|
if (dimLength <= 0) {
|
|
// When no dimensions specified, create TAD along all dimensions
|
|
// This means the entire array is treated as a single TAD
|
|
std::vector<LongType> allDims;
|
|
for (LongType i = 0; i < rank; i++) {
|
|
allDims.push_back(i);
|
|
}
|
|
|
|
// Recursively call with all dimensions
|
|
return tadForDimensions(originalShape, allDims.data(), rank);
|
|
}
|
|
|
|
// Additional validation: check if dimensions are within valid range
|
|
for (LongType i = 0; i < dimLength; i++) {
|
|
LongType dim = dimensions[i];
|
|
if (dim < 0) dim += rank; // Handle negative dimensions
|
|
if (dim < 0 || dim >= rank) {
|
|
THROW_EXCEPTION("Dimension index is out of bounds");
|
|
}
|
|
}
|
|
|
|
// Create non-temporary vector to satisfy the reference requirement
|
|
std::vector<LongType> dims(dimensions, dimensions + dimLength);
|
|
|
|
// The shared_ptr keeps the TadPack alive even if the cache tries to clear it
|
|
std::shared_ptr<TadPack> result = nullptr;
|
|
try {
|
|
result = _trie.getOrCreate(dims, originalShape);
|
|
} catch (const std::exception& e) {
|
|
THROW_EXCEPTION("Failed to create or retrieve TAD pack");
|
|
}
|
|
|
|
// DO NOT call checkAndCleanupCaches() here - would delete the pack we just created!
|
|
// Cleanup happens at NativeOps layer AFTER operations complete.
|
|
return result;
|
|
}
|
|
|
|
void ConstantTadHelper::clearCache() {
|
|
_trie.clear();
|
|
}
|
|
|
|
LongType ConstantTadHelper::getCachedEntries() const {
|
|
return _trie.getCachedEntries();
|
|
}
|
|
|
|
LongType ConstantTadHelper::getCachedBytes() const {
|
|
return _trie.getCachedBytes();
|
|
}
|
|
|
|
LongType ConstantTadHelper::getPeakCachedEntries() const {
|
|
return _trie.getPeakCachedEntries();
|
|
}
|
|
|
|
LongType ConstantTadHelper::getPeakCachedBytes() const {
|
|
return _trie.getPeakCachedBytes();
|
|
}
|
|
|
|
std::string ConstantTadHelper::toString(int maxDepth, int maxEntries) const {
|
|
return _trie.toString(maxDepth, maxEntries);
|
|
}
|
|
|
|
void ConstantTadHelper::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
|
|
_trie.getCachedPointers(out_pointers);
|
|
}
|
|
|
|
} // namespace sd
|