wehub-resource-sync
448 lines
16 KiB
C++
448 lines
16 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 "helpers/ConstantShapeHelper.h"
|
|
|
|
#include "array/ConstantShapeBuffer.h"
|
|
#include "system/common.h"
|
|
#include <array/PrimaryPointerDeallocator.h>
|
|
#include <helpers/ConstantShapeHelper.h>
|
|
#include <helpers/ShapeBuilders.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <helpers/shape.h>
|
|
#include <system/Environment.h>
|
|
|
|
namespace sd {
|
|
|
|
ConstantShapeHelper::~ConstantShapeHelper() {
|
|
|
|
}
|
|
|
|
ConstantShapeHelper::ConstantShapeHelper() {
|
|
}
|
|
|
|
ConstantShapeHelper& ConstantShapeHelper::getInstance() {
|
|
static ConstantShapeHelper instance;
|
|
return instance;
|
|
}
|
|
|
|
void ConstantShapeHelper::initializeEarly() {
|
|
ConstantShapeHelper& instance = getInstance();
|
|
instance._shapeTrie.waitForInitialization();
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::createConstBuffFromExisting(sd::LongType* shapeInfo) {
|
|
auto result = bufferForShapeInfo(shapeInfo);
|
|
return result;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(LongType* shapeInfo) {
|
|
if(shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
if(shape::rank(shapeInfo) < 0 || shape::rank(shapeInfo) > SD_MAX_RANK) {
|
|
THROW_EXCEPTION("shapeInfo is not a valid rank.");
|
|
}
|
|
|
|
auto buffer = _shapeTrie.getOrCreate(shapeInfo);
|
|
if (buffer == nullptr || buffer->primary() == nullptr) {
|
|
THROW_EXCEPTION("Failed to get/create shape buffer");
|
|
}
|
|
return buffer;
|
|
}
|
|
ConstantShapeBuffer* ConstantShapeHelper::createSubArrShapeInfo( sd::LongType* inShapeInfo, LongType* dims,
|
|
sd::LongType dimsSize) {
|
|
sd::LongType* newShapeInfo = ShapeBuilders::createSubArrShapeInfo(inShapeInfo, dims, dimsSize, nullptr);
|
|
auto ret = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return ret;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(DataType dataType, char order,
|
|
const std::vector<LongType>& shape) {
|
|
auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, shape);
|
|
auto result = bufferForShapeInfo(descriptor);
|
|
delete[] descriptor;
|
|
return result;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfo(DataType dataType, char order,
|
|
int rank, LongType* shape) {
|
|
auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, rank, shape, nullptr, false);
|
|
auto result = bufferForShapeInfo(descriptor);
|
|
delete[] descriptor;
|
|
return result;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::emptyShapeInfoWithShape(DataType dataType, std::vector<LongType>& shape) {
|
|
auto descriptor = ShapeBuilders::createShapeInfo(dataType, 'c', shape, nullptr);
|
|
ArrayOptions::setPropertyBit(descriptor, ARRAY_EMPTY);
|
|
auto existing = createFromExisting(descriptor);
|
|
delete[] descriptor;
|
|
return existing;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, char order,
|
|
const std::vector<LongType>& shape) {
|
|
auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, shape);
|
|
auto result = bufferForShapeInfo(descriptor)->primary();
|
|
delete[] descriptor;
|
|
return result;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, char order, int rank,
|
|
LongType* shape, LongType extraProperties) {
|
|
if (extraProperties < 0) {
|
|
extraProperties = ArrayOptions::flagForDataType(dataType);
|
|
}
|
|
|
|
std::unique_ptr<LongType[]> strides(order == 'c' ? shape::calcStrides(shape, rank)
|
|
: shape::calcStridesFortran(shape, rank));
|
|
|
|
auto descriptor = ShapeBuilders::createShapeInfo(dataType, order, rank, shape, strides.get(),
|
|
nullptr, extraProperties);
|
|
auto ret = bufferForShapeInfo(descriptor)->primary();
|
|
ArrayOptions::validateSingleDataType(ArrayOptions::dataType(ret));
|
|
|
|
delete[] descriptor;
|
|
return ret;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::createShapeInfo(DataType dataType, LongType* shapeInfo) {
|
|
auto result = createShapeInfo(dataType, shape::order(shapeInfo), shape::rank(shapeInfo),
|
|
shape::shapeOf(const_cast<LongType*>(shapeInfo)), -1);
|
|
return result;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::emptyShapeInfo(DataType dataType) {
|
|
auto descriptor = ShapeBuilders::emptyShapeInfo(dataType);
|
|
auto result = bufferForShapeInfo(descriptor)->primary();
|
|
delete[] descriptor;
|
|
return result;
|
|
}
|
|
|
|
|
|
LongType* ConstantShapeHelper::scalarShapeInfo(DataType dataType) {
|
|
auto descriptor = ShapeBuilders::createScalarShapeInfo(dataType);
|
|
return bufferForShapeInfo(descriptor)->primary();
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::vectorShapeInfo(LongType length, DataType dataType) {
|
|
auto descriptor = ShapeBuilders::createVectorShapeInfo(dataType, length);
|
|
auto result = bufferForShapeInfo(descriptor)->primary();
|
|
delete[] descriptor;
|
|
return result;
|
|
}
|
|
|
|
|
|
LongType* ConstantShapeHelper::createShapeInfo(ShapeDescriptor* descriptor) {
|
|
auto shapeInfo = descriptor->toShapeInfo();
|
|
auto result = bufferForShapeInfo(shapeInfo)->primary();
|
|
delete[] shapeInfo;
|
|
return result;
|
|
}
|
|
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithView(LongType* shapeInfo) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
|
|
|
|
|
|
ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
|
|
delete[] newShapeInfo;
|
|
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutView(LongType* shapeInfo) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
ArrayOptions::unsetPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithNeedsCopy(LongType* shapeInfo) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_NEEDS_COPY);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutNeedsCopy(LongType* shapeInfo) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
ArrayOptions::unsetPropertyBit(newShapeInfo, ARRAY_NEEDS_COPY);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithCopyOffset(LongType* shapeInfo, int inputIndex) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
if (inputIndex < 0 || inputIndex > 10) {
|
|
THROW_EXCEPTION("Input index out of range [0-10]");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
|
|
ArrayOptions::setPropertyBit(newShapeInfo, flag);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutCopyOffset(LongType* shapeInfo, int inputIndex) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
if (inputIndex < 0 || inputIndex > 10) {
|
|
THROW_EXCEPTION("Input index out of range [0-10]");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
|
|
ArrayOptions::unsetPropertyBit(newShapeInfo, flag);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithoutAllCopyOffsets(LongType* shapeInfo) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
ArrayOptions::clearAllCopyOffsets(newShapeInfo);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoWithFlags(LongType* shapeInfo,
|
|
LongType flagsToSet,
|
|
LongType flagsToUnset) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
|
|
// Unset flags first
|
|
if (flagsToUnset != 0) {
|
|
LongType extraIdx = ArrayOptions::extraIndex(newShapeInfo);
|
|
newShapeInfo[extraIdx] = newShapeInfo[extraIdx] & ~flagsToUnset;
|
|
}
|
|
|
|
// Then set flags
|
|
if (flagsToSet != 0) {
|
|
LongType extraIdx = ArrayOptions::extraIndex(newShapeInfo);
|
|
newShapeInfo[extraIdx] = newShapeInfo[extraIdx] | flagsToSet;
|
|
}
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::bufferForShapeInfoAsViewWithOffset(LongType* shapeInfo,
|
|
int inputIndex) {
|
|
if (shapeInfo == nullptr) {
|
|
THROW_EXCEPTION("shapeInfo is nullptr");
|
|
}
|
|
|
|
if (inputIndex < 0 || inputIndex > 10) {
|
|
THROW_EXCEPTION("Input index out of range [0-10]");
|
|
}
|
|
|
|
LongType* newShapeInfo = ShapeBuilders::copyShapeInfo(shapeInfo, false, nullptr);
|
|
|
|
// Set view flag
|
|
ArrayOptions::setPropertyBit(newShapeInfo, ARRAY_IS_VIEW);
|
|
|
|
// Set copy offset flag for specified input
|
|
LongType flag = ArrayOptions::copyOffsetFlagForInput(inputIndex);
|
|
ArrayOptions::setPropertyBit(newShapeInfo, flag);
|
|
|
|
auto buffer = bufferForShapeInfo(newShapeInfo);
|
|
delete[] newShapeInfo;
|
|
return buffer;
|
|
}
|
|
|
|
LongType* ConstantShapeHelper::createFromExisting(LongType* shapeInfo) {
|
|
if (!shapeInfo) {
|
|
THROW_EXCEPTION("Null shape info");
|
|
}
|
|
auto buffer = bufferForShapeInfo(shapeInfo);
|
|
return buffer->primary();
|
|
}
|
|
|
|
|
|
LongType* ConstantShapeHelper::castToDataType(LongType* shapeInfo, DataType newType) {
|
|
if (!shapeInfo) {
|
|
THROW_EXCEPTION("Null shape info");
|
|
}
|
|
if (ArrayOptions::dataType(shapeInfo) == newType) {
|
|
return shapeInfo;
|
|
}
|
|
|
|
auto tempShapeInfo = ShapeBuilders::copyShapeInfoWithNewType(shapeInfo, newType);
|
|
if (!tempShapeInfo) {
|
|
THROW_EXCEPTION("Failed to create temp shape info");
|
|
}
|
|
|
|
auto buffer = bufferForShapeInfo(tempShapeInfo);
|
|
auto result = buffer->primary();
|
|
delete[] tempShapeInfo;
|
|
if(ArrayOptions::dataType(result) != newType) {
|
|
std::string errorMessage;
|
|
errorMessage += "castToDataType: new data type is ";
|
|
errorMessage += DataTypeUtils::asString(newType);
|
|
errorMessage += " data type from new constant created data type ";
|
|
errorMessage += DataTypeUtils::asString(ArrayOptions::dataType(result));
|
|
errorMessage += "\n";
|
|
THROW_EXCEPTION(errorMessage.c_str());
|
|
}
|
|
return result;
|
|
}
|
|
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::createShapeInfoWithUnitiesForBroadcast(sd::LongType* maxShapeInfo,
|
|
sd::LongType* minShapeInfo,
|
|
sd::memory::Workspace* workspace,
|
|
const std::vector<LongType>& dimensions) {
|
|
sd::LongType* newShapeInfo = nullptr;
|
|
ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(shape::rank(maxShapeInfo)), sd::LongType);
|
|
|
|
newShapeInfo[0] = shape::rank(maxShapeInfo);
|
|
newShapeInfo[2 * shape::rank(maxShapeInfo) + 1] = 0;
|
|
sd::ArrayOptions::copyDataType(newShapeInfo, minShapeInfo); // type
|
|
newShapeInfo[2 * newShapeInfo[0] + 2] = shape::elementWiseStride(minShapeInfo); // ews
|
|
newShapeInfo[2 * newShapeInfo[0] + 3] = shape::order(minShapeInfo); // order
|
|
|
|
if (!dimensions.empty()) {
|
|
for (sd::LongType k = 0, j = 0, i = 0; i < shape::rank(maxShapeInfo); ++i) {
|
|
if (j < static_cast<sd::LongType>(dimensions.size()) && dimensions[j] == i) {
|
|
shape::shapeOf(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[k];
|
|
shape::stride(newShapeInfo)[i] = shape::stride(minShapeInfo)[k++];
|
|
++j;
|
|
} else {
|
|
shape::shapeOf(newShapeInfo)[i] = 1;
|
|
shape::stride(newShapeInfo)[i] = 0;
|
|
if (shape::sizeAt(minShapeInfo, k) == 1 && static_cast<sd::LongType>(dimensions.size()) != shape::rank(minShapeInfo)) ++k;
|
|
}
|
|
}
|
|
} else {
|
|
for (int j = shape::rank(minShapeInfo) - 1, i = shape::rank(maxShapeInfo) - 1; i >= 0; --i) {
|
|
if (j >= 0) {
|
|
shape::shapeOf(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[j];
|
|
shape::stride(newShapeInfo)[i] = shape::shapeOf(minShapeInfo)[j] == 1 ? 0 : shape::stride(minShapeInfo)[j];
|
|
--j;
|
|
} else {
|
|
shape::shapeOf(newShapeInfo)[i] = 1;
|
|
shape::stride(newShapeInfo)[i] = 0;
|
|
}
|
|
}
|
|
}
|
|
|
|
auto ret = bufferForShapeInfo(newShapeInfo);
|
|
RELEASE(newShapeInfo, workspace);
|
|
return ret;
|
|
}
|
|
|
|
ConstantShapeBuffer* ConstantShapeHelper::createShapeInfoWithNoUnitiesForReduce(const sd::LongType* maxShapeInfo,
|
|
const std::vector<LongType>* dimsWithUnities,
|
|
sd::memory::Workspace* workspace) {
|
|
sd::LongType* newShapeInfo = nullptr;
|
|
ALLOCATE(newShapeInfo, workspace, shape::shapeInfoLength(shape::rank(maxShapeInfo) - dimsWithUnities->size()),
|
|
sd::LongType);
|
|
|
|
sd::LongType temp;
|
|
if (dimsWithUnities->size() == 1 && shape::isCommonVector(maxShapeInfo, temp) && temp == dimsWithUnities->at(0)) {
|
|
auto dims = ShapeUtils::evalDimsToExclude(shape::rank(maxShapeInfo), 1,&temp);
|
|
shape::excludeUnitiesFromShapeInfo(maxShapeInfo, dims->data(), dims->size(), newShapeInfo);
|
|
delete dims;
|
|
} else {
|
|
shape::excludeUnitiesFromShapeInfo(maxShapeInfo, dimsWithUnities->data(), dimsWithUnities->size(), newShapeInfo);
|
|
}
|
|
|
|
auto ret = bufferForShapeInfo(newShapeInfo);
|
|
RELEASE(newShapeInfo, workspace);
|
|
return ret;
|
|
}
|
|
|
|
void ConstantShapeHelper::clearCache() {
|
|
std::lock_guard<std::mutex> lock(_mutex);
|
|
_shapeTrie.clearCache();
|
|
}
|
|
|
|
LongType ConstantShapeHelper::getCachedEntries() const {
|
|
return _shapeTrie.getCachedEntries();
|
|
}
|
|
|
|
LongType ConstantShapeHelper::getCachedBytes() const {
|
|
return _shapeTrie.getCachedBytes();
|
|
}
|
|
|
|
LongType ConstantShapeHelper::getPeakCachedEntries() const {
|
|
return _shapeTrie.getPeakCachedEntries();
|
|
}
|
|
|
|
LongType ConstantShapeHelper::getPeakCachedBytes() const {
|
|
return _shapeTrie.getPeakCachedBytes();
|
|
}
|
|
|
|
std::string ConstantShapeHelper::toString(int maxDepth, int maxEntries) const {
|
|
return _shapeTrie.toString(maxDepth, maxEntries);
|
|
}
|
|
|
|
void ConstantShapeHelper::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
|
|
_shapeTrie.getCachedPointers(out_pointers);
|
|
}
|
|
|
|
} // namespace sd
|
|
|