wehub-resource-sync
112 lines
3.8 KiB
C++
112 lines
3.8 KiB
C++
#include <array/TadCalculator.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <helpers/ConstantHelper.h>
|
|
#include <helpers/ConstantShapeHelper.h>
|
|
|
|
namespace sd {
|
|
|
|
TadCalculator::TadCalculator(LongType* originalShape)
|
|
: _originalShape(originalShape), _numTads(0), _tadShape(nullptr), _tadOffsets(nullptr) {}
|
|
|
|
TadCalculator::~TadCalculator() {
|
|
// Only delete _tadOffsets if we still own it (hasn't been released)
|
|
if (_tadOffsets != nullptr) {
|
|
delete _tadOffsets;
|
|
_tadOffsets = nullptr;
|
|
}
|
|
// _tadShape comes from ConstantShapeHelper cache - DON'T delete
|
|
}
|
|
|
|
void TadCalculator::createTadPack(const std::vector<LongType>& dimensions) {
|
|
if (!_originalShape) {
|
|
THROW_EXCEPTION("Original shape is null");
|
|
}
|
|
|
|
// Check for empty array
|
|
if (shape::isEmptyConst(_originalShape)) {
|
|
THROW_EXCEPTION("Cannot create TADs for empty array");
|
|
}
|
|
|
|
auto shapeInfo = ConstantShapeHelper::getInstance().createFromExisting(_originalShape);
|
|
const LongType rank = shape::rank(shapeInfo);
|
|
|
|
// Check for zero-sized dimensions
|
|
for (LongType i = 0; i < rank; i++) {
|
|
if (shape::sizeAt(shapeInfo, i) == 0) {
|
|
THROW_EXCEPTION("Cannot create TADs for array with zero-sized dimensions");
|
|
}
|
|
}
|
|
|
|
const std::vector<LongType>* dimsToExclude = ShapeUtils::evalDimsToExclude(rank, dimensions.size(), dimensions.data());
|
|
if (!dimsToExclude) {
|
|
THROW_EXCEPTION("Failed to evaluate dimensions to exclude");
|
|
}
|
|
|
|
if (dimsToExclude->size() == 0 || dimsToExclude->size() == rank) {
|
|
const LongType totalElements = shape::length(shapeInfo);
|
|
|
|
auto scalarShapeInfo = ConstantShapeHelper::getInstance().scalarShapeInfo(ArrayOptions::dataType(shapeInfo));
|
|
auto scalarShapeBuffer = ConstantShapeHelper::getInstance().bufferForShapeInfo(scalarShapeInfo);
|
|
|
|
auto oPtr = std::make_shared<PointerWrapper>(new LongType[totalElements]);
|
|
LongType* offsets = oPtr->pointerAsT<LongType>();
|
|
|
|
for (LongType i = 0; i < totalElements; ++i) {
|
|
offsets[i] = i;
|
|
}
|
|
|
|
_tadShape = scalarShapeBuffer;
|
|
_tadOffsets = new ConstantOffsetsBuffer(oPtr);
|
|
_numTads = totalElements;
|
|
|
|
delete dimsToExclude;
|
|
return;
|
|
}
|
|
|
|
const LongType numOfSubArrs = ShapeUtils::getNumOfSubArrs(shapeInfo, *dimsToExclude);
|
|
|
|
if (numOfSubArrs > 0) {
|
|
const LongType subArrRank = (static_cast<size_t>(rank) == dimsToExclude->size() || false) ? rank : rank - dimsToExclude->size();
|
|
|
|
auto sPtr = std::make_shared<PointerWrapper>(new LongType[shape::shapeInfoLength(subArrRank)]);
|
|
auto oPtr = std::make_shared<PointerWrapper>(new LongType[numOfSubArrs]);
|
|
|
|
shape::calcSubArrsShapeInfoAndOffsets(
|
|
shapeInfo,
|
|
numOfSubArrs,
|
|
dimsToExclude->size(),
|
|
dimsToExclude->data(),
|
|
sPtr->pointerAsT<LongType>(),
|
|
oPtr->pointerAsT<LongType>(),
|
|
false);
|
|
|
|
auto shapesBuffer = ConstantShapeHelper::getInstance().bufferForShapeInfo(sPtr->pointerAsT<LongType>());
|
|
|
|
_tadOffsets = new ConstantOffsetsBuffer(oPtr);
|
|
_tadShape = shapesBuffer;
|
|
_numTads = numOfSubArrs;
|
|
} else {
|
|
const LongType subArrRank = rank;
|
|
|
|
auto sPtr = std::make_shared<PointerWrapper>(new LongType[shape::shapeInfoLength(subArrRank)]);
|
|
LongType* shapeInfo2 = sPtr->pointerAsT<LongType>();
|
|
|
|
auto nonConstant = const_cast<LongType*>(shapeInfo);
|
|
auto nonConst2 = const_cast<LongType*>(shapeInfo2);
|
|
shape::copyTo<LongType>(shape::shapeInfoLength(subArrRank), nonConstant, nonConst2);
|
|
|
|
LongType* baseOffset = new LongType[1];
|
|
baseOffset[0] = 0;
|
|
auto oPtr = std::make_shared<PointerWrapper>(baseOffset);
|
|
|
|
auto shapesBuffer = ConstantShapeHelper::getInstance().bufferForShapeInfo(sPtr->pointerAsT<LongType>());
|
|
_tadOffsets = new ConstantOffsetsBuffer(oPtr);
|
|
_tadShape = shapesBuffer;
|
|
_numTads = 1;
|
|
}
|
|
|
|
delete dimsToExclude;
|
|
}
|
|
|
|
} // namespace sd
|