chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,441 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// @author raver119@gmail.com, created on 07.10.2017.
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
|
||||
#include <array/NDArray.h>
|
||||
#include <helpers/Loops.h>
|
||||
|
||||
#include <helpers/shape.h>
|
||||
#include <ops/declarable/CustomOperations.h>
|
||||
#include <ops/specials.h>
|
||||
#include <types/types.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
/**
|
||||
* @brief Checks if the shape of NDArray contains 1 before(order c) or after(order f) the specified axis
|
||||
*
|
||||
* @param input
|
||||
* @param axis
|
||||
* @return int
|
||||
*/
|
||||
SD_INLINE int isShapeExtendedWithOnes(NDArray&input, LongType axis) {
|
||||
bool isAllOne = true;
|
||||
auto shapes = shape::shapeOf(input.shapeInfo());
|
||||
auto rank = input.rankOf();
|
||||
if (rank > axis) {
|
||||
if (input.ordering() == 'c') {
|
||||
// check before the axis
|
||||
for (sd::LongType i = 0; i < axis; i++) {
|
||||
isAllOne = isAllOne && (shapes[i] == 1);
|
||||
}
|
||||
} else {
|
||||
// check after the axis
|
||||
for (int i = axis + 1; i < rank; i++) {
|
||||
isAllOne = isAllOne && (shapes[i] == 1);
|
||||
}
|
||||
}
|
||||
return isAllOne;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct InputArgsCase2 {
|
||||
const T *ptr;
|
||||
int size;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::concatCpuGeneric(const std::vector<NDArray *> &inArrs, NDArray &output,
|
||||
const LongType axis) {
|
||||
const sd::LongType numOfInArrs = inArrs.size();
|
||||
const auto sizeofT = output.sizeOfT();
|
||||
|
||||
T *zBuff = output.bufferAsT<T>();
|
||||
|
||||
bool shapeExtendedWithOnes = isShapeExtendedWithOnes(output, axis);
|
||||
bool followEws1 = false;
|
||||
bool matchesOutputOrdering = true;
|
||||
for (int i = 0; i < numOfInArrs; ++i) {
|
||||
shapeExtendedWithOnes = shapeExtendedWithOnes && isShapeExtendedWithOnes(*inArrs[i], axis);
|
||||
matchesOutputOrdering = matchesOutputOrdering && inArrs[i]->ordering() == output.ordering();
|
||||
}
|
||||
|
||||
bool copyCaseEws1 = followEws1 & matchesOutputOrdering;
|
||||
bool copyCase1 = numOfInArrs > 1 ? copyCaseEws1 & shapeExtendedWithOnes : copyCaseEws1;
|
||||
|
||||
if (copyCase1) {
|
||||
// copyCase1:
|
||||
// When NdArrays follow the same order and unit elementwise stride and
|
||||
// the concantneation axis is 0th or has only 1 before it {1, 1, ..., axis} for "c"
|
||||
// or axis is (rank-1)th or has only 1 after it {axis, 1, 1, ..., 1} for "f"
|
||||
// we will concatenate them by sequential copying of the whole buffers
|
||||
|
||||
std::vector<T *> zPtrList;
|
||||
T *z = output.bufferAsT<T>();
|
||||
for (sd::LongType i = 0; i < numOfInArrs; i++) {
|
||||
zPtrList.push_back(z);
|
||||
z += inArrs[i]->lengthOf();
|
||||
}
|
||||
auto func = [&inArrs, &zPtrList](sd::LongType thread_id, sd::LongType start, sd::LongType stop,
|
||||
sd::LongType increment) -> void {
|
||||
for (sd::LongType i = start; i < stop; ++i) {
|
||||
const auto memAmountToCopy = inArrs[i]->lengthOf();
|
||||
const auto inputPtr = inArrs[i]->bufferAsT<T>();
|
||||
|
||||
auto zPtr = zPtrList[i];
|
||||
for (int j = 0; j < memAmountToCopy; j++) {
|
||||
zPtr[j] = inputPtr[j];
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_tad(func, 0, numOfInArrs, 1);
|
||||
return;
|
||||
}
|
||||
|
||||
// for one Array
|
||||
if (numOfInArrs < 2) {
|
||||
output.assign(inArrs[0]);
|
||||
return;
|
||||
}
|
||||
bool copyCase2 = copyCaseEws1 && output.ordering() == 'c';
|
||||
if (copyCase2) {
|
||||
sd::LongType times = 1;
|
||||
auto shapes = shape::shapeOf(output.shapeInfo());
|
||||
|
||||
T *z = output.bufferAsT<T>();
|
||||
for (int i = 0; i < axis; i++) {
|
||||
times = times * shapes[i];
|
||||
}
|
||||
|
||||
sd::LongType totalCopySize = output.lengthOf() / times;
|
||||
|
||||
std::vector<InputArgsCase2<T>> inputArgs;
|
||||
for (sd::LongType i = 0; i < numOfInArrs; i++) {
|
||||
InputArgsCase2<T> input = {inArrs[i]->bufferAsT<T>(),
|
||||
static_cast<int>(inArrs[i]->lengthOf()) / static_cast<int>(times)};
|
||||
inputArgs.push_back(input);
|
||||
}
|
||||
|
||||
auto func = [&inputArgs, z, totalCopySize](uint64_t thread_id, int64_t start, int64_t stop,
|
||||
int64_t increment) -> void {
|
||||
auto outPtr = &(z[start * totalCopySize]);
|
||||
auto numOfInArrs = inputArgs.size();
|
||||
for (int i = start; i < stop; i++) {
|
||||
for (size_t j = 0; j < numOfInArrs; j++) {
|
||||
auto inputCopySize = inputArgs[j].size;
|
||||
const T *inputBasePtr = inputArgs[j].ptr;
|
||||
auto inputPtr = &(inputBasePtr[i * inputCopySize]);
|
||||
// copy
|
||||
PRAGMA_OMP_SIMD
|
||||
for (int k = 0; k < inputCopySize; k++) {
|
||||
outPtr[k] = inputPtr[k];
|
||||
}
|
||||
outPtr += inputCopySize;
|
||||
}
|
||||
}
|
||||
};
|
||||
samediff::Threads::parallel_tad(func, 0, times, 1);
|
||||
return;
|
||||
}
|
||||
|
||||
// Cache shape and stride information for output
|
||||
const sd::LongType zRank = shape::rank(output.shapeInfo());
|
||||
const sd::LongType* zShape = shape::shapeOf(output.shapeInfo());
|
||||
const sd::LongType* zStride = shape::stride(output.shapeInfo());
|
||||
|
||||
// Pre-cache input arrays' shape information
|
||||
std::vector<const sd::LongType*> inShapes(numOfInArrs);
|
||||
std::vector<const sd::LongType*> inStrides(numOfInArrs);
|
||||
std::vector<sd::LongType> inRanks(numOfInArrs);
|
||||
|
||||
for (sd::LongType i = 0; i < numOfInArrs; i++) {
|
||||
inRanks[i] = shape::rank(inArrs[i]->shapeInfo());
|
||||
inShapes[i] = shape::shapeOf(inArrs[i]->shapeInfo());
|
||||
inStrides[i] = shape::stride(inArrs[i]->shapeInfo());
|
||||
}
|
||||
|
||||
// general case
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
sd::LongType coords[SD_MAX_RANK], temp;
|
||||
|
||||
for (sd::LongType i = start; i < stop; i += increment) {
|
||||
INDEX2COORDS(i, zRank, zShape, coords);
|
||||
|
||||
sd::LongType zOffset;
|
||||
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
||||
|
||||
sd::LongType inArrIdx = 0;
|
||||
sd::LongType xDim = inArrs[inArrIdx]->sizeAt(axis);
|
||||
|
||||
temp = coords[axis];
|
||||
while (coords[axis] >= xDim) {
|
||||
coords[axis] -= xDim;
|
||||
xDim = inArrs[++inArrIdx]->sizeAt(axis);
|
||||
}
|
||||
|
||||
const T *x = inArrs[inArrIdx]->bufferAsT<T>();
|
||||
sd::LongType xOffset;
|
||||
COORDS2INDEX(inRanks[inArrIdx], inStrides[inArrIdx], coords, xOffset);
|
||||
|
||||
zBuff[zOffset] = x[xOffset];
|
||||
|
||||
coords[axis] = temp;
|
||||
}
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_for(func, 0, output.lengthOf());
|
||||
}
|
||||
/**
|
||||
* Concatneate multi array of the same shape together
|
||||
* along a particular dimension
|
||||
*/
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::concatCpuGeneric(LongType dimension, int numArrays,NDArray **data,
|
||||
NDArray *vresult) {
|
||||
auto result = reinterpret_cast<T *>(vresult);
|
||||
std::vector<NDArray *> inputs(numArrays);
|
||||
|
||||
|
||||
for (sd::LongType i = 0; i < numArrays; ++i)
|
||||
inputs[i] =
|
||||
new NDArray(static_cast<void *>(data[i]), data[i]->shapeInfo(), nullptr, false, 0);
|
||||
|
||||
sd::SpecialMethods<T>::concatCpuGeneric(inputs, *vresult, dimension);
|
||||
|
||||
for (sd::LongType i = 0; i < numArrays; ++i) {
|
||||
delete inputs[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::splitCpuGeneric(NDArray& input, const std::vector<NDArray*>& outArrs, const LongType axis) {
|
||||
int numSplits = outArrs.size();
|
||||
const auto sizeofT = input.sizeOfT();
|
||||
auto xBuff = input.bufferAsT<T>();
|
||||
|
||||
bool luckCase1 = ((axis == 0 && input.ordering() == 'c') || (axis == input.rankOf() - 1 && input.ordering() == 'f'));
|
||||
|
||||
if (luckCase1) {
|
||||
T* x = const_cast<T*>(xBuff);
|
||||
for (sd::LongType i = 0; i < numSplits; ++i) {
|
||||
const auto memAmountToCopy = outArrs[i]->lengthOf();
|
||||
ops::safe_copy(x, static_cast<const T*>(outArrs[i]->buffer()), static_cast<size_t>(memAmountToCopy));
|
||||
x += memAmountToCopy;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Cache shape and stride information
|
||||
const sd::LongType xRank = shape::rank(input.shapeInfo());
|
||||
const sd::LongType* xShape = shape::shapeOf(input.shapeInfo());
|
||||
const sd::LongType* xStride = shape::stride(input.shapeInfo());
|
||||
|
||||
// Pre-cache output array ranks, shapes, and strides
|
||||
std::vector<const sd::LongType*> outShapes(numSplits);
|
||||
std::vector<const sd::LongType*> outStrides(numSplits);
|
||||
std::vector<sd::LongType> outRanks(numSplits);
|
||||
|
||||
for (int i = 0; i < numSplits; i++) {
|
||||
outRanks[i] = shape::rank(outArrs[i]->shapeInfo());
|
||||
outShapes[i] = shape::shapeOf(outArrs[i]->shapeInfo());
|
||||
outStrides[i] = shape::stride(outArrs[i]->shapeInfo());
|
||||
}
|
||||
|
||||
sd::LongType zDim = outArrs[0]->sizeAt(axis);
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
sd::LongType coords[SD_MAX_RANK], temp;
|
||||
|
||||
for (sd::LongType i = start; i < stop; i += increment) {
|
||||
INDEX2COORDS(i, xRank, xShape, coords);
|
||||
sd::LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, coords, xOffset);
|
||||
|
||||
sd::LongType outArrIdx = 0;
|
||||
temp = coords[axis];
|
||||
|
||||
while (coords[axis] >= zDim) {
|
||||
coords[axis] -= zDim;
|
||||
++outArrIdx;
|
||||
}
|
||||
|
||||
T* z = outArrs[outArrIdx]->bufferAsT<T>();
|
||||
sd::LongType zOffset;
|
||||
COORDS2INDEX(outRanks[outArrIdx], outStrides[outArrIdx], coords, zOffset);
|
||||
z[zOffset] = xBuff[xOffset];
|
||||
|
||||
coords[axis] = temp;
|
||||
}
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_for(func, 0, input.lengthOf());
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::sortGeneric(NDArray *input, bool descending) {
|
||||
auto x = input->bufferAsT<T>();
|
||||
auto xShapeInfo = input->shapeInfo();
|
||||
quickSort_parallel(input, Environment::getInstance().maxMasterThreads(), descending);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::quickSort_parallel_internal(NDArray *x, int left, int right, int cutoff, bool descending) {
|
||||
if (right - left <= cutoff) {
|
||||
// Use insertion sort for small arrays
|
||||
auto xBuff = x->bufferAsT<T>();
|
||||
for (int i = left + 1; i <= right; i++) {
|
||||
T key = xBuff[i];
|
||||
int j = i - 1;
|
||||
if (descending) {
|
||||
while (j >= left && xBuff[j] < key) {
|
||||
xBuff[j + 1] = xBuff[j];
|
||||
j--;
|
||||
}
|
||||
} else {
|
||||
while (j >= left && xBuff[j] > key) {
|
||||
xBuff[j + 1] = xBuff[j];
|
||||
j--;
|
||||
}
|
||||
}
|
||||
xBuff[j + 1] = key;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Choose pivot as median of three
|
||||
auto xBuff = x->bufferAsT<T>();
|
||||
int mid = (left + right) / 2;
|
||||
if (descending) {
|
||||
if (xBuff[right] > xBuff[left]) std::swap(xBuff[right], xBuff[left]);
|
||||
if (xBuff[mid] > xBuff[left]) std::swap(xBuff[mid], xBuff[left]);
|
||||
if (xBuff[right] > xBuff[mid]) std::swap(xBuff[right], xBuff[mid]);
|
||||
} else {
|
||||
if (xBuff[right] < xBuff[left]) std::swap(xBuff[right], xBuff[left]);
|
||||
if (xBuff[mid] < xBuff[left]) std::swap(xBuff[mid], xBuff[left]);
|
||||
if (xBuff[right] < xBuff[mid]) std::swap(xBuff[right], xBuff[mid]);
|
||||
}
|
||||
|
||||
// Partition
|
||||
T pivot = xBuff[mid];
|
||||
int i = left;
|
||||
int j = right;
|
||||
|
||||
while (i <= j) {
|
||||
if (descending) {
|
||||
while (xBuff[i] > pivot) i++;
|
||||
while (xBuff[j] < pivot) j--;
|
||||
} else {
|
||||
while (xBuff[i] < pivot) i++;
|
||||
while (xBuff[j] > pivot) j--;
|
||||
}
|
||||
|
||||
if (i <= j) {
|
||||
std::swap(xBuff[i], xBuff[j]);
|
||||
i++;
|
||||
j--;
|
||||
}
|
||||
}
|
||||
|
||||
// Recursively sort sub-arrays
|
||||
if (left < j) quickSort_parallel_internal(x, left, j, cutoff, descending);
|
||||
if (i < right) quickSort_parallel_internal(x, i, right, cutoff, descending);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::quickSort_parallel(NDArray *x, int numThreads, bool descending) {
|
||||
const int CUTOFF = 32; // Threshold for switching to insertion sort
|
||||
auto length = x->lengthOf();
|
||||
|
||||
if (length <= 1) return;
|
||||
|
||||
// For very small arrays, just use the internal sort
|
||||
if (length <= CUTOFF || numThreads <= 1) {
|
||||
quickSort_parallel_internal(x, 0, length - 1, CUTOFF, descending);
|
||||
return;
|
||||
}
|
||||
|
||||
// For larger arrays, partition into segments and sort in parallel
|
||||
int segmentSize = length / numThreads;
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
int threadLeft = start * segmentSize;
|
||||
int threadRight = (start == numThreads - 1) ? length - 1 : (start + 1) * segmentSize - 1;
|
||||
quickSort_parallel_internal(x, threadLeft, threadRight, CUTOFF, descending);
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_for(func, 0, numThreads);
|
||||
|
||||
// Merge sorted segments if we used multiple threads
|
||||
if (numThreads > 1) {
|
||||
auto xBuff = x->bufferAsT<T>();
|
||||
std::vector<T> temp(length);
|
||||
for (int size = segmentSize; size < length; size *= 2) {
|
||||
for (int left = 0; left < length; left += 2 * size) {
|
||||
int mid = std::min(left + size, (int)length);
|
||||
int right = std::min(left + 2 * size, (int)length);
|
||||
int i = left, j = mid, k = left;
|
||||
|
||||
// Merge two segments
|
||||
while (i < mid && j < right) {
|
||||
if (descending) {
|
||||
temp[k++] = (xBuff[i] >= xBuff[j]) ? xBuff[i++] : xBuff[j++];
|
||||
} else {
|
||||
temp[k++] = (xBuff[i] <= xBuff[j]) ? xBuff[i++] : xBuff[j++];
|
||||
}
|
||||
}
|
||||
while (i < mid) temp[k++] = xBuff[i++];
|
||||
while (j < right) temp[k++] = xBuff[j++];
|
||||
|
||||
// Copy back
|
||||
for (i = left; i < right; i++) {
|
||||
xBuff[i] = temp[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SpecialMethods<T>::sortTadGeneric(NDArray *input, sd::LongType *dimension, int dimensionLength, bool descending) {
|
||||
auto x = input->bufferAsT<T>();
|
||||
sd::LongType xLength = input->lengthOf();
|
||||
sd::LongType xTadLength = shape::tadLength(input->shapeInfo(), dimension, dimensionLength);
|
||||
int numTads = xLength / xTadLength;
|
||||
|
||||
const std::vector<sd::LongType> dimVector(dimension, dimension + dimensionLength);
|
||||
auto pack = sd::ConstantTadHelper::getInstance().tadForDimensions(
|
||||
const_cast<sd::LongType *>(input->shapeInfo()), const_cast<sd::LongType *>(dimVector.data()), false);
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto r = start; r < stop; r++) {
|
||||
NDArray *dx = pack->extractTadView(input, r);
|
||||
quickSort_parallel(dx, xTadLength, descending);
|
||||
delete dx;
|
||||
}
|
||||
};
|
||||
samediff::Threads::parallel_tad(func, 0, numTads);
|
||||
}
|
||||
|
||||
} // namespace sd
|
||||
Reference in New Issue
Block a user