wehub-resource-sync
408 lines
16 KiB
C++
408 lines
16 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
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author Yurii Shyrma (iuriish@yahoo.com)
|
|
// @author raver119@gmail.com
|
|
//
|
|
#include <execution/Threads.h>
|
|
#include <ops/declarable/helpers/s_t_b.h>
|
|
#if NOT_EXCLUDED(OP_space_to_batch)
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void batchToSpace_(NDArray& input, NDArray& output, const sd::LongType cropBottom,
|
|
const sd::LongType cropTop, const sd::LongType cropLeft, const sd::LongType cropRight) {
|
|
// input [bS, H * blockSize, W * blockSize, iC]
|
|
// output [bS, H * blockSize - cropBottom - cropTop, W * blockSize - cropLeft - cropRight, iC]
|
|
|
|
// if (cropTop = cropBottom = cropRight = cropLeft = 0) shapes are the same
|
|
// else:
|
|
// oH -> [cropBottom, iH - cropTop]
|
|
// oW -> [cropLeft, iH - cropRight]
|
|
// xLen > zLen
|
|
|
|
const T* x = input.bufferAsT<T>();
|
|
T* z = output.bufferAsT<T>();
|
|
|
|
const int rank = 4;
|
|
|
|
const sd::LongType* xShapeInfo = input.shapeInfo();
|
|
const sd::LongType* zShapeInfo = output.shapeInfo();
|
|
|
|
const sd::LongType bS = xShapeInfo[1];
|
|
const sd::LongType iH = xShapeInfo[2];
|
|
const sd::LongType iW = xShapeInfo[3];
|
|
const sd::LongType iC = xShapeInfo[4];
|
|
|
|
// loop through output array
|
|
auto func = PRAGMA_THREADS_FOR_3D {
|
|
for (auto b = start_x; b < stop_x; b += inc_x) {
|
|
for (auto h = start_y; h < stop_y; h += inc_y) {
|
|
for (auto w = start_z; w < stop_z; w += inc_z) {
|
|
for (sd::LongType c = 0; c < iC; ++c) {
|
|
const sd::LongType xOffset = b * xShapeInfo[5] + h * xShapeInfo[6] + w * xShapeInfo[7] + c * xShapeInfo[8];
|
|
const sd::LongType zOffset = b * zShapeInfo[5] + (h - cropBottom) * zShapeInfo[6] +
|
|
(w - cropLeft) * zShapeInfo[7] + c * zShapeInfo[8];
|
|
|
|
z[zOffset] = x[xOffset];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, bS, 1, cropBottom, iH - cropTop, 1, cropLeft, iW - cropRight, 1);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void batchToSpace_,
|
|
(NDArray& input, NDArray& output, const sd::LongType cropBottom, const sd::LongType cropTop,
|
|
const sd::LongType cropLeft, const sd::LongType cropRight),
|
|
SD_COMMON_TYPES);
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void batchToSpace(sd::LaunchContext* context, NDArray input, NDArray& output, const sd::LongType cropBottom,
|
|
const sd::LongType cropTop, const sd::LongType cropLeft, const sd::LongType cropRight,
|
|
const sd::LongType blockSize) {
|
|
// [bS*blockSize*blockSize, H/blockSize, W/blockSize, iC] is rearranged/permuted to [bS, oH, oW, iC]
|
|
// oH = H - cropTop - cropBottom
|
|
// oW = W - cropLeft - cropRight
|
|
|
|
std::vector<sd::LongType> shape = {blockSize, blockSize, output.sizeAt(0), input.sizeAt(1), input.sizeAt(2), input.sizeAt(3)};
|
|
NDArray *inputRearranged0 = input.reshape(
|
|
input.ordering(),shape);
|
|
inputRearranged0->permutei({2, 3, 0, 4, 1, 5}, false, false);
|
|
|
|
if (input.lengthOf() == output.lengthOf())
|
|
output.assign(inputRearranged0);
|
|
else {
|
|
std::vector<sd::LongType> temp = {output.sizeAt(0), input.sizeAt(1) * blockSize, input.sizeAt(2) * blockSize, input.sizeAt(3)};
|
|
NDArray *inputRearranged1 = inputRearranged0->reshape(
|
|
input.ordering(),
|
|
temp);
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), batchToSpace_,
|
|
(*inputRearranged1, output, cropBottom, cropTop, cropLeft, cropRight), SD_COMMON_TYPES);
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void batchToSpaceND_(NDArray* input, NDArray* crop, NDArray* output,
|
|
const LongType numOfSpatialDims) {
|
|
// input [bS, H * blockShape[0], W * blockShape[1], iC]
|
|
// output [bS, H * blockShape[0] - cropBottom - cropTop, W * blockShape[1] - cropLeft - cropRight, iC]
|
|
|
|
// if (cropTop = cropBottom = cropRight = cropLeft = 0) shapes are the same
|
|
// else:
|
|
// oH -> [cropBottom, iH - cropTop]
|
|
// oW -> [cropLeft, iH - cropRight]
|
|
// xLen >= zLen
|
|
|
|
const T* x = input->bufferAsT<T>();
|
|
T* z = output->bufferAsT<T>();
|
|
|
|
const sd::LongType rank = input->rankOf();
|
|
const sd::LongType zLen = output->lengthOf();
|
|
|
|
// loop through input array
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
sd::LongType zCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
|
|
|
for (auto i = start; i < stop; i++) {
|
|
INDEX2COORDS(i, rank, shape::shapeOf(output->shapeInfo()), zCoords);
|
|
|
|
memcpy(xCoords, zCoords, rank * sizeof(sd::LongType));
|
|
|
|
// evaluate spatial coordinates for x
|
|
for (sd::LongType j = 1; j <= numOfSpatialDims; ++j)
|
|
xCoords[j] += crop->e<sd::LongType>(j - 1, 0); // add crop left
|
|
|
|
sd::LongType zOffset, xOffset;
|
|
COORDS2INDEX(rank, shape::stride(output->shapeInfo()), zCoords, zOffset);
|
|
COORDS2INDEX(rank, shape::stride(input->shapeInfo()), xCoords, xOffset);
|
|
|
|
z[zOffset] = x[xOffset];
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, zLen);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void batchToSpaceND_,
|
|
(NDArray* input, NDArray* crop, NDArray* output, const sd::LongType numOfSpatialDims),
|
|
SD_COMMON_TYPES);
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void batchToSpaceND(sd::LaunchContext* context, NDArray& input, NDArray& blockShape, NDArray& crop,
|
|
NDArray& output){
|
|
// 4D example, numOfSpatialDims = 2 - two spatial dimensions
|
|
// [bS*blockShape[0]*blockShape[1], iH, iW, iC] is rearranged/permuted to [bS, iH*blockShape[0] - cropTop -
|
|
// cropBottom, iW*blockShape[1] - cropLeft - cropRight, iC]
|
|
|
|
const sd::LongType rank = input.rankOf();
|
|
const sd::LongType numOfSpatialDims = blockShape.sizeAt(0);
|
|
|
|
//*** construct reshaping std::vector for first reshape of input array ***//
|
|
|
|
std::vector<sd::LongType> temp(numOfSpatialDims + rank);
|
|
|
|
sd::LongType i;
|
|
for (i = 0; i < numOfSpatialDims; ++i) temp[i] = blockShape.e<sd::LongType>(i);
|
|
temp[i++] = output.sizeAt(0);
|
|
for (sd::LongType j = 1; j < rank; ++i, ++j) temp[i] = input.sizeAt(j);
|
|
|
|
NDArray *inputRearranged0 = input.reshape(input.ordering(), temp);
|
|
|
|
//*** construct permuting std::vector for permutation of input array ***//
|
|
|
|
temp[0] = numOfSpatialDims;
|
|
|
|
for (i = 1; i <= numOfSpatialDims; ++i) {
|
|
temp[2 * i - 1] = numOfSpatialDims + i;
|
|
temp[2 * i] = i - 1;
|
|
}
|
|
for (i = 2 * numOfSpatialDims + 1; i < static_cast<sd::LongType>(temp.size()); ++i) temp[i] = i;
|
|
|
|
inputRearranged0->permutei(temp, false, false);
|
|
|
|
if (input.lengthOf() == output.lengthOf()) {
|
|
output.assign(inputRearranged0);
|
|
} else {
|
|
//*** construct reshaping std::vector for second reshape of input array ***//
|
|
|
|
temp.resize(rank);
|
|
|
|
temp[0] = output.sizeAt(0);
|
|
|
|
for (i = 1; i < rank; ++i)
|
|
temp[i] = (i <= numOfSpatialDims) ? input.sizeAt(i) * blockShape.e<sd::LongType>(i - 1) : input.sizeAt(i);
|
|
|
|
NDArray *inputRearranged1 = inputRearranged0->reshape(input.ordering(), temp);
|
|
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), batchToSpaceND_, (inputRearranged1, &crop, &output, numOfSpatialDims),
|
|
SD_COMMON_TYPES);
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void spaceToBatch_(NDArray& input, NDArray& output, const sd::LongType padBottom,
|
|
const sd::LongType padTop, const sd::LongType padLeft, const sd::LongType padRight) {
|
|
// input [bS, H * blockSize - padBottom - padTop, W * blockSize - padLeft - padRight, iC]
|
|
// output [bS, H * blockSize, W * blockSize, iC]
|
|
|
|
// if (padTop = padBottom = padRight = padLeft = 0) shapes are the same
|
|
// else:
|
|
// iH -> [padBottom, oH - padTop]
|
|
// iW -> [padLeft, oW - padRight]
|
|
// zLen > xLen
|
|
|
|
const T* x = input.bufferAsT<T>();
|
|
T* z = output.bufferAsT<T>();
|
|
|
|
const int rank = 4;
|
|
|
|
const sd::LongType* xShapeInfo = input.shapeInfo();
|
|
const sd::LongType* zShapeInfo = output.shapeInfo();
|
|
|
|
const sd::LongType bS = zShapeInfo[1];
|
|
const sd::LongType oH = zShapeInfo[2];
|
|
const sd::LongType oW = zShapeInfo[3];
|
|
const sd::LongType iC = zShapeInfo[4];
|
|
|
|
// loop through output array
|
|
auto func = PRAGMA_THREADS_FOR_2D {
|
|
for (auto b = start_x; b < stop_x; b += inc_x) {
|
|
for (auto h = start_y; h < stop_y; h += inc_y) {
|
|
for (sd::LongType w = 0; w < oW; ++w) {
|
|
for (sd::LongType c = 0; c < iC; ++c) {
|
|
const sd::LongType zOffset = b * zShapeInfo[5] + h * zShapeInfo[6] + w * zShapeInfo[7] + c * zShapeInfo[8];
|
|
|
|
if (h >= padBottom && h < oH - padTop && w >= padLeft && w < oW - padRight) {
|
|
const sd::LongType xOffset = b * xShapeInfo[5] + (h - padBottom) * xShapeInfo[6] +
|
|
(w - padLeft) * xShapeInfo[7] + c * xShapeInfo[8];
|
|
z[zOffset] = x[xOffset];
|
|
} else
|
|
z[zOffset] = 0.f;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_for(func, 0, bS, 1, 0, oH, 1);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void spaceToBatch_,
|
|
(NDArray& input, NDArray& output, const sd::LongType padBottom, const sd::LongType padTop,
|
|
const sd::LongType padLeft, const sd::LongType padRight),
|
|
SD_COMMON_TYPES);
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void spaceToBatch(sd::LaunchContext* context, NDArray& input, NDArray& output, const sd::LongType padBottom,
|
|
const sd::LongType padTop, const sd::LongType padLeft, const sd::LongType padRight,
|
|
const sd::LongType blockSize) {
|
|
// [bS, iH, iW, iC] is rearranged/permuted to [bS*blockSize*blockSize, (iH + padBottom + padTop)/blockSize, (iW +
|
|
// padLeft + padRight)/blockSize, iC]
|
|
|
|
std::vector<sd::LongType> shape1 = {blockSize, blockSize, input.sizeAt(0), output.sizeAt(1), output.sizeAt(2), output.sizeAt(3)};
|
|
NDArray *outputRearranged0 = output.reshape(
|
|
output.ordering(), shape1,
|
|
false);
|
|
outputRearranged0->permutei({2, 3, 0, 4, 1, 5}, false, false);
|
|
|
|
if (input.lengthOf() == output.lengthOf()) {
|
|
outputRearranged0->assign(&input);
|
|
} else {
|
|
std::vector<sd::LongType> shape2 = {input.sizeAt(0), output.sizeAt(1) * blockSize, output.sizeAt(2) * blockSize, output.sizeAt(3)};
|
|
NDArray *outputRearranged1 = outputRearranged0->reshape(
|
|
output.ordering(),
|
|
shape2, false);
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), spaceToBatch_,
|
|
(input, *outputRearranged1, padBottom, padTop, padLeft, padRight), SD_COMMON_TYPES);
|
|
|
|
if (output.buffer() != outputRearranged1->buffer()) outputRearranged0->assign(outputRearranged1);
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void spaceToBatchND_(NDArray& input, NDArray& padding, NDArray& output,
|
|
const LongType numOfSpatialDims) {
|
|
// 4D example
|
|
// input [bS, H * blockShape[0] - padBottom - padTop, W * blockShape[1] - padLeft - padRight, iC]
|
|
// output [bS, H * blockShape[0], W * blockShape[1], iC]
|
|
|
|
// if (padTop = padBottom = padRight = padLeft = 0) shapes are the same
|
|
// else:
|
|
// iH -> [padBottom, oH - padTop]
|
|
// iW -> [padLeft, oW - padRight]
|
|
// zLen > xLen
|
|
|
|
const T* x = input.bufferAsT<T>();
|
|
T* z = output.bufferAsT<T>();
|
|
|
|
const int rank = input.rankOf();
|
|
const sd::LongType zLen = output.lengthOf();
|
|
|
|
// loop through output array
|
|
auto func = PRAGMA_THREADS_FOR {
|
|
sd::LongType zCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
|
|
|
for (sd::LongType i = start; i < stop; i++) {
|
|
INDEX2COORDS(i, rank, shape::shapeOf(output.shapeInfo()), zCoords);
|
|
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(rank, shape::stride(output.shapeInfo()), zCoords, zOffset);
|
|
|
|
memcpy(xCoords, zCoords, rank * sizeof(LongType));
|
|
|
|
bool within = true;
|
|
|
|
for (sd::LongType j = 1; j <= numOfSpatialDims; ++j) {
|
|
const auto padLeft = padding.e<sd::LongType>(j - 1, 0);
|
|
const auto padRight = padding.e<sd::LongType>(j - 1, 1);
|
|
|
|
within &= zCoords[j] >= padLeft && zCoords[j] < output.sizeAt(j) - padRight;
|
|
|
|
if (!within) break;
|
|
|
|
xCoords[j] = zCoords[j] - padLeft; // get coordinates for x
|
|
}
|
|
|
|
if (within) {
|
|
sd::LongType xOffset;
|
|
COORDS2INDEX(rank, shape::stride(input.shapeInfo()), xCoords, xOffset);
|
|
z[zOffset] = x[xOffset];
|
|
} else {
|
|
z[zOffset] = 0.f;
|
|
}
|
|
}
|
|
};
|
|
|
|
samediff::Threads::parallel_tad(func, 0, zLen);
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE( void spaceToBatchND_,
|
|
(NDArray& input, NDArray& padding, NDArray& output,
|
|
const sd::LongType numOfSpatialDims),
|
|
SD_COMMON_TYPES);
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void spaceToBatchND(sd::LaunchContext* context, NDArray& input, NDArray& blockShape, NDArray& padding,
|
|
NDArray& output) {
|
|
// 4D example with two spatial dimensions
|
|
// [bS, iH, iW, iC] is rearranged/permuted to [bS*blockShape[0]*blockShape[1], (iH + padBottom +
|
|
// padTop)/blockShape[0], (iW + padLeft + padRight)/blockShape[1], iC]
|
|
|
|
const sd::LongType rank = input.rankOf();
|
|
|
|
const sd::LongType numOfSpatialDims = blockShape.sizeAt(0);
|
|
|
|
//*** construct reshaping std::vector for first reshape of output array ***//
|
|
std::vector<sd::LongType> temp(numOfSpatialDims + rank);
|
|
|
|
int i;
|
|
for (i = 0; i < numOfSpatialDims; ++i) temp[i] = blockShape.e<sd::LongType>(i);
|
|
temp[i++] = input.sizeAt(0);
|
|
for (int j = 1; j < rank; ++i, ++j) temp[i] = output.sizeAt(j);
|
|
|
|
NDArray *outputRearranged0 = output.reshape(output.ordering(), temp, false);
|
|
|
|
//*** construct permuting std::vector for permutation of output array ***//
|
|
|
|
temp[0] = numOfSpatialDims;
|
|
|
|
for (i = 1; i <= numOfSpatialDims; ++i) {
|
|
temp[2 * i - 1] = numOfSpatialDims + i;
|
|
temp[2 * i] = i - 1;
|
|
}
|
|
for (i = 2 * numOfSpatialDims + 1; i < static_cast<int>(temp.size()); ++i) temp[i] = i;
|
|
|
|
outputRearranged0->permutei(temp, false, false);
|
|
|
|
// ****** //
|
|
|
|
if (input.lengthOf() == output.lengthOf()) {
|
|
outputRearranged0->assign(&input);
|
|
} else {
|
|
//*** construct reshaping std::vector for second reshape of output array ***//
|
|
temp.resize(rank);
|
|
|
|
temp[0] = input.sizeAt(0);
|
|
|
|
for (i = 1; i < rank; ++i)
|
|
temp[i] = (i <= numOfSpatialDims) ? output.sizeAt(i) * blockShape.e<sd::LongType>(i - 1) : output.sizeAt(i);
|
|
|
|
NDArray *outputRearranged1 = outputRearranged0->reshape(output.ordering(), temp, false);
|
|
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), spaceToBatchND_, (input, padding, *outputRearranged1, numOfSpatialDims),
|
|
SD_COMMON_TYPES);
|
|
|
|
if (output.buffer() != outputRearranged1->buffer()) outputRearranged0->assign(outputRearranged1);
|
|
}
|
|
}
|
|
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|
|
#endif |