Files
2026-07-13 12:47:05 +08:00

542 lines
20 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 NDARRAY_CPP
#define NDARRAY_CPP
#include <array/NDArray.h>
#include <helpers/ArrayUtils.h>
#include <helpers/ConstantTadHelper.h>
#include <helpers/ShapeUtils.h>
#include <helpers/logger.h>
#include <indexing/NDIndex.h>
#include <loops/BroadcastPairwiseConverter.h>
#include <loops/random.h>
#include <ops/ops.h>
#include <array/NDArray.hXX>
#include <memory>
#include <sstream>
#include <stdexcept>
#include <system/selective_rendering.h>
namespace sd {
// NDArray implementation for .cpp file
void NDArray::printBufferDebug(const char* msg, sd::LongType offset, sd::LongType limit) {
if (msg) sd_printf("%s:\n", msg);
if(limit < 0) limit = lengthOf();
// Print array info
sd_printf("NDArray: Shape=[", 0);
for (int i = 0; i < rankOf(); i++) {
sd_printf("%lld", (long long)sizeAt(i));
if (i < rankOf() - 1) sd_printf(",", 0);
}
sd_printf("], DataType=%s, EWS=%lld, Order=%c\n",
DataTypeUtils::asString(dataType()).c_str(), (long long)ews(), ordering());
// Print buffer state
if (_buffer != nullptr) {
_buffer->printBufferDebug("Buffer contents", offset, limit);
} else {
sd_printf("Buffer is nullptr\n", 0);
}
}
////////////////////////////////////////////////////////////////////////
void* NDArray::platformBuffer() { return buffer(); }
////////////////////////////////////////////////////////////////////////
template <typename T>
void NDArray::fillAsTriangular(const float val, int lower, int upper, NDArray& target, const char direction, const bool includeEdges) {
if (isS()) THROW_EXCEPTION("NDArray::fillArrayAsTriangular: you can't use this method on String array!");
if (!isSameShape(target) &&
!(rankOf() == 1 && target.rankOf() == 2 && sizeAt(0) == target.sizeAt(0) && sizeAt(0) == target.sizeAt(1)))
THROW_EXCEPTION("NDArray::fillArrayAsTriangular method: wrong shape of target array !");
const T value = static_cast<T>(val);
const auto x = reinterpret_cast<const T*>(buffer());
auto z = reinterpret_cast<T*>(target.buffer());
const int xRank = rankOf();
const int zRank = target.rankOf();
const auto zLen = target.lengthOf();
const bool areSameOffsets = shape::haveSameShapeAndStrides(shapeInfo(), target.shapeInfo());
sd::LongType *targetShape = shape::shapeOf(target.shapeInfo());
sd::LongType *targetStride = shape::stride(target.shapeInfo());
sd::LongType targetRank = target.rankOf();
sd::LongType *xShape = shape::shapeOf(shapeInfo());
sd::LongType *xStride = shape::stride(shapeInfo());
sd::LongType thisRank = this->rankOf();
auto func = PRAGMA_THREADS_FOR {
sd::LongType coords[SD_MAX_RANK], temp;
sd::LongType vectorCoord[1] = {0};
bool notVectorScalar = targetRank == 2 && thisRank == 2;
bool thisNotVectorScalar = !shape::isScalar(this->shapeInfo()) && !shape::isVector(this->shapeInfo());
bool targetNotVectorScalar = !shape::isScalar(target.shapeInfo()) && !shape::isVector(target.shapeInfo());
for (sd::LongType i = start; i < stop; i++) {
INDEX2COORDS(i, targetRank,targetShape, coords);
sd::LongType row = targetNotVectorScalar ? coords[zRank - 2] : 0;
sd::LongType col = targetNotVectorScalar ? coords[zRank - 1] : 1;
sd::LongType zOffset, xOffset;
if (target.rankOf() < 2) {
COORDS2INDEX(targetRank, targetStride, vectorCoord, zOffset);
} else {
COORDS2INDEX(targetRank, targetStride, coords, zOffset);
}
if (!areSameOffsets && rankOf() < 2) {
COORDS2INDEX(thisRank, xStride, vectorCoord, xOffset);
} else if (areSameOffsets) {
xOffset = zOffset;
} else {
COORDS2INDEX(thisRank, xStride, coords, xOffset);
}
bool rowExclusive = this->rankOf() == target.rankOf();
bool colExclusive = this->rankOf() == target.rankOf();
auto lCompare = includeEdges ? row <= (col - lower) : row < (col - lower);
auto uCompare = includeEdges ? row >= (col - upper) : row > (col - upper);
if ((direction == 'u' && lCompare) || (direction == 'l' && uCompare)) {
z[zOffset] = value;
} else {
z[zOffset] = x[xOffset];
}
if (this != &target) {
if (xRank != zRank) {
temp = coords[0];
coords[0] = coords[1];
}
if (xRank != zRank) // restore first coordinate
coords[0] = temp;
}
if (vectorCoord[0] == this->lengthOf() - 1) {
vectorCoord[0] = 0;
} else {
vectorCoord[0] = vectorCoord[0] + 1;
}
}
};
samediff::Threads::parallel_for(func, 0, zLen);
}
BUILD_SINGLE_TEMPLATE( void NDArray::fillAsTriangular,
(const float val, int lower, int upper, NDArray& target, const char direction,const bool includeEdges), SD_COMMON_TYPES);
////////////////////////////////////////////////////////////////////////
void NDArray::setIdentity() {
if (isS()) THROW_EXCEPTION("NDArray::setIdentity: you can't use this method on String array!");
this->nullify();
int rank = rankOf();
auto shape = shapeOf();
int minDim = SD_MAX_INT;
sd::LongType indices[SD_MAX_RANK];
for (int j = 0; j < rank; ++j) indices[j] = 1;
sd::LongType offset;
COORDS2INDEX(rank, shape::stride(shapeInfo()), indices, offset);
for (int i = 0; i < rank; ++i)
if (minDim > shape[i]) minDim = shape[i];
float v = 1.0f;
for (int i = 0; i < minDim; ++i) templatedSet<float,float>(buffer(), i * offset, this->dataType(), &v);
}
////////////////////////////////////////////////////////////////////////
template <typename T>
static void templatedSwap(void* xBuffer, void* yBuffer, const sd::LongType* xShapeInfo, const sd::LongType* yShapeInfo,
sd::LongType length) {
auto x = reinterpret_cast<T*>(xBuffer);
auto y = reinterpret_cast<T*>(yBuffer);
const bool isSameOrders = shape::order(xShapeInfo) == shape::order(xShapeInfo);
sd::LongType xRank = shape::rank(xShapeInfo);
sd::LongType yRank = shape::rank(yShapeInfo);
sd::LongType *xShape = shape::shapeOf(xShapeInfo);
sd::LongType *yShape = shape::shapeOf(yShapeInfo);
sd::LongType *xStride = shape::stride(xShapeInfo);
sd::LongType *yStride = shape::stride(yShapeInfo);
auto func = PRAGMA_THREADS_FOR {
if (isSameOrders) {
for (sd::LongType i = start; i < stop; i++) {
LongType xCoords[SD_MAX_RANK];
LongType yCoords[SD_MAX_RANK];
LongType xOffset;
LongType yOffset;
INDEX2COORDS(i, xRank, xShape, xCoords);
COORDS2INDEX(xRank, shape::stride(xShapeInfo), xCoords, xOffset);
INDEX2COORDS(i, yRank,yShape, yCoords);
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
sd::math::sd_swap(x[xOffset], y[yOffset]);
}
} else if (shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo)) {
for (sd::LongType i = start; i < stop; i++) {
LongType coords[SD_MAX_RANK];
LongType ind;
INDEX2COORDS(i, xRank, xShape, coords);
COORDS2INDEX(xRank, xStride, coords, ind);
sd::math::sd_swap(x[ind], y[ind]);
}
} else {
for (sd::LongType i = start; i < stop; i++) {
LongType xCoords[SD_MAX_RANK];
LongType yCoords[SD_MAX_RANK];
LongType xInd;
LongType yInd;
INDEX2COORDS(i, xRank, xShape, xCoords);
COORDS2INDEX(xRank, xStride, xCoords, xInd);
INDEX2COORDS(i, yRank, yShape, yCoords);
COORDS2INDEX(yRank, yStride, yCoords, yInd);
sd::math::sd_swap(x[xInd], y[yInd]);
}
}
};
samediff::Threads::parallel_for(func, 0, length);
}
BUILD_SINGLE_TEMPLATE( void templatedSwap,
(void* xBuffer, void* yBuffer, const sd::LongType* xShapeInfo, const sd::LongType* yShapeInfo,
sd::LongType length),
SD_COMMON_TYPES);
////////////////////////////////////////////////////////////////////////
void NDArray::swapUnsafe(NDArray& other) {
auto xType = this->dataType();
if (xType != other.dataType())
THROW_EXCEPTION("NDArray::swapUnsage method: both arrays must have the same data type");
if (buffer() == nullptr || other.buffer() == nullptr)
THROW_EXCEPTION("NDArray::swapUnsafe method: input array should not be empty!");
if (lengthOf() != other.lengthOf())
THROW_EXCEPTION("NDArray::swapUnsafe method: input arrays should have the same length!");
BUILD_SINGLE_SELECTOR(xType, templatedSwap,
(buffer(), other.buffer(), shapeInfo(), other.shapeInfo(), this->lengthOf()), SD_COMMON_TYPES);
}
////////////////////////////////////////////////////////////////////////
void NDArray::synchronize(const char* msg) {
// no-op
}
void NDArray::syncToDevice() {}
void NDArray::syncToHost() {}
void NDArray::tickWriteHost() {}
void NDArray::tickWriteDevice() {}
void NDArray::tickReadHost() {}
void NDArray::tickReadDevice() {}
void NDArray::tickBothActual() {}
bool NDArray::isActualOnHostSide() { return true; }
bool NDArray::isActualOnDeviceSide() { return true; }
void NDArray::makeBothBuffersActual() {}
void NDArray::preparePrimaryUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList, bool synchronizeWritables) {
// no-op
}
void NDArray::registerPrimaryUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList) {
// no-op
}
void NDArray::prepareSpecialUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList, bool synchronizeWritables) {
// no-op
}
void NDArray::registerSpecialUse(const std::vector<NDArray*>& writeList,
const std::vector<NDArray*>& readList) {
// no-op
}
void NDArray::syncShape() {
// no-op
}
//////////////////////////////////////////////////////////////////////////
template <typename T>
void NDArray::printCurrentBuffer(const bool host, const char* msg, const int precision) {}
template void NDArray::printCurrentBuffer<int>(const bool host, const char* msg, const int precision) ;
template void NDArray::printCurrentBuffer<float>(const bool host, const char* msg, const int precision) ;
template void NDArray::printCurrentBuffer<double>(const bool host, const char* msg, const int precision);
template void NDArray::printCurrentBuffer<sd::LongType>(const bool host, const char* msg, const int precision) ;
////////////////////////////////////////////////////////////////////////
void* NDArray::specialBuffer() {
if (_buffer == nullptr) {
THROW_EXCEPTION("NDArray::specialBuffer(): _buffer is nullptr - array not properly initialized");
}
void* specialBuf = _buffer->special();
// On CPU, special buffer is nullptr (only used for GPU/CUDA) - this is expected and normal
if (specialBuf == nullptr) {
return nullptr;
}
return static_cast<int8_t*>(specialBuf) + (_offset * sizeOfT());
}
//////////////////////////////////////////////////////////////////////////
// change an array by repeating it the number of times given by reps.
NDArray NDArray::tile(const std::vector<sd::LongType>& reps) {
const int repsSize = reps.size();
sd::LongType product = 1;
for (const auto& item : reps) product *= item;
if (product == 0) THROW_EXCEPTION("NDArray::tile method: one of the elements in reps array is zero !");
int rankOld = rankOf();
int diff = rankOld - repsSize;
if (product == 1) { // in this case 2 possibilities are present: just reshape or nothing to do
NDArray result(*this);
if (diff < 0) { // reshape to higher dimension
std::vector<sd::LongType> shapeNew =
reps; // there is requirement to have unities at first "diff" positions of new shape
memcpy(&shapeNew[-diff], result.shapeInfo() + 1,
rankOld * sizeof(sd::LongType)); // put old shape numbers at rest of positions
result.reshapei(ordering(), shapeNew);
}
return result; // nothing to do, if diff >= 0 -> identity tile
}
// evaluate shapeInfo for resulting array
auto newShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
// create new buffer, in any case the memory amount new buffer points to is bigger then those for old _buffer
DataBuffer * newBuff =
new DataBuffer(shape::length(newShapeInfo) * sizeOfT(), dataType(), getContext()->getWorkspace());
// assign new shape and new buffer to resulting array
NDArray result(newBuff,newShapeInfo , getContext());
// fill newBuff, loop through all elements of newBuff
// looping through _buffer goes automatically by means of getSubArrayIndex applying
const auto resultLen = result.lengthOf();
auto xType = this->dataType();
auto func = PRAGMA_THREADS_FOR {
for (auto i = start; i < stop; i++) {
auto xOffset = result.getOffset(i);
auto yOffset = shape::subArrayOffset(i, newShapeInfo, shapeInfo());
BUILD_SINGLE_SELECTOR(xType, this->template templatedAssign,
(result.buffer(), xOffset, this->buffer(), yOffset), SD_COMMON_TYPES);
}
};
samediff::Threads::parallel_for(func, 0, resultLen);
result.tickWriteHost();
return result;
}
//////////////////////////////////////////////////////////////////////////
// change an array by repeating it the number of times given by reps.
void NDArray::tile(const std::vector<LongType>& reps, NDArray& target) {
// Validate the tile operation
auto repProd = shape::prodLong(reps.data(), reps.size());
if (repProd < 1)
THROW_EXCEPTION("NDArray::tile: reps can't contain 0s");
// Validate the target shape
auto correctShapeInfo = ShapeUtils::evalTileShapeInfo(*this, reps, getContext()->getWorkspace());
if (!shape::equalsSoft(correctShapeInfo, target.shapeInfo())) {
THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation!");
}
const auto targetLen = target.lengthOf();
// Safely calculate source array offset
for (LongType i = 0; i < targetLen; ++i) {
// Calculate target array offset
auto xOffset = target.getOffset(i);
// Calculate source coordinates based on target coordinates
LongType targetCoords[SD_MAX_RANK];
INDEX2COORDS(i, shape::rank(target.shapeInfo()), shape::shapeOf(target.shapeInfo()), targetCoords);
// Map target coordinates to source coordinates manually
LongType sourceCoords[SD_MAX_RANK];
for (int d = 0; d < shape::rank(shapeInfo()); d++) {
// Apply modulo for each dimension
sourceCoords[d] = targetCoords[d] % shape::sizeAt(shapeInfo(), d);
}
// Calculate source offset from source coordinates
LongType sourceOffset;
COORDS2INDEX(shape::rank(shapeInfo()), shape::stride(shapeInfo()), sourceCoords, sourceOffset);
auto targetDataType = target.dataType();
auto selfDType = dataType();
// Copy the value
BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
(target.buffer(), xOffset, buffer(), sourceOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
}
}
//////////////////////////////////////////////////////////////////////////
void NDArray::tile(NDArray& target) {
if (rankOf() > target.rankOf())
THROW_EXCEPTION(
"NDArray::tile method - rank of target array must be bigger or equal to the rank of this array !");
if (!ShapeUtils::areShapesBroadcastable(*this, target))
THROW_EXCEPTION("NDArray::tile method - shapeInfo of target array is not suitable for tile operation !");
// fill newBuff, loop through all elements of newBuff
// looping through _buffer goes automatically by means of getSubArrayIndex applying
const auto ews = target.ews();
const auto targetLen = target.lengthOf();
if (target.ordering() == 'c' && ews >= 1) {
for (sd::LongType i = 0; i < targetLen; ++i) {
auto yOffset = shape::subArrayOffset(i, target.shapeInfo(), shapeInfo());
auto targetDataType = target.dataType();
auto selfDType = dataType();
BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
(target.buffer(), i * ews, buffer(), yOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
}
} else {
for (sd::LongType i = 0; i < targetLen; ++i) {
auto xOffset = target.getOffset(i);
auto yOffset = shape::subArrayOffset(i, target.shapeInfo(), shapeInfo());
auto targetDataType = target.dataType();
auto selfDType = dataType();
BUILD_DOUBLE_SELECTOR(target.dataType(), dataType(), templatedDoubleAssign,
(target.buffer(), xOffset, buffer(), yOffset), SD_COMMON_TYPES, SD_COMMON_TYPES);
}
}
}
////////////////////////////////////////////////////////////////////////
template <typename X, typename Z>
static void repeat_(NDArray& input, NDArray& output, const std::vector<LongType>& repeats, const LongType axis) {
const X* x = input.bufferAsT<X>();
Z* z = output.bufferAsT<Z>();
const sd::LongType rank = input.rankOf(); // xRank = zRank
const sd::LongType zLen = output.lengthOf(); // xLen <= zLen
const sd::LongType repSize = repeats.size();
sd::LongType outputRank = output.rankOf();
sd::LongType* outputShape = shape::shapeOf(output.shapeInfo());
sd::LongType* outputStride = shape::stride(output.shapeInfo());
sd::LongType inputRank = input.rankOf();
sd::LongType* inputShape = shape::shapeOf(input.shapeInfo());
sd::LongType* inputStride = shape::stride(input.shapeInfo());
// loop through input array
auto func = PRAGMA_THREADS_FOR {
sd::LongType coords[SD_MAX_RANK], temp;
for (sd::LongType i = start; i < stop; i++) {
INDEX2COORDS(i, outputRank, outputShape, coords);
sd::LongType zOffset;
COORDS2INDEX(outputRank, outputStride, coords, zOffset);
temp = coords[axis];
if (repSize > 1) {
for (sd::LongType j = 0; j < repSize; ++j) {
coords[axis] -= repeats[j];
if (coords[axis] < 0) {
coords[axis] = j;
break;
}
}
} else
coords[axis] /= repeats[0];
sd::LongType xOffset;
COORDS2INDEX(inputRank,inputStride, coords, xOffset);
z[zOffset] = x[xOffset];
coords[axis] = temp;
}
};
samediff::Threads::parallel_for(func, 0, zLen);
}
//////////////////////////////////////////////////////////////////////////
// create new array by repeating it the number of times given by repeats
NDArray NDArray::repeat(const int axis, const std::vector<LongType>& repeats) {
NDArray *thisArr = const_cast<NDArray*>(this);
std::vector<sd::LongType> repeatShape = ShapeUtils::evalRepeatShape(axis, repeats, *thisArr);
NDArray output('c',repeatShape, dataType(), getContext());
BUILD_SINGLE_SELECTOR_TWICE(dataType(), repeat_, (*this, output, repeats, axis), SD_COMMON_TYPES);
return output;
}
//////////////////////////////////////////////////////////////////////////
// fill array by repeating it the number of times given by reps
void NDArray::repeat(const int axis, const std::vector<LongType>& repeats, NDArray& target) {
NDArray *thisArr = const_cast<NDArray*>(this);
if (!target.isSameShape(ShapeUtils::evalRepeatShape(axis, repeats, *thisArr)))
THROW_EXCEPTION(
"NDArray::repeat(const int axis, const std::vector<int>& repeats, NDArray& target) method: wrong shape of "
"target array!");
auto targetDataType = target.dataType();
auto selfDType = dataType();
BUILD_DOUBLE_SELECTOR(dataType(), target.dataType(), repeat_, (*this, target, repeats, axis), SD_COMMON_TYPES,
SD_COMMON_TYPES);
}
//////////////////////////////////////////////////////////////////////////
#ifndef __JAVACPP_HACK__
#include "NDArrayLambda.hpp"
#endif
} // namespace sd
#endif