/* ******************************************************************************
*
*
* 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), created on 20.04.2018
//
#include<ops/declarable/helpers/transforms.h>
#include <array/ResultSet.h>
#include <helpers/ShapeUtils.h>
#include <numeric>
#include <array/NDArrayFactory.h>
#include <helpers/TAD.h>
#include <exceptions/cuda_exception.h>
#include <helpers/PointersManager.h>
#include <helpers/ConstantTadHelper.h>
namespace sd {
namespace ops {
namespace helpers {
///////////////////////////////////////////////////////////////////
// x - input, y - paddings, z - output
template<typename X, typename Y>
__global__ static void padCuda(const int mode,
const void *vx, const Nd4jLong *xShapeInfo,
const void *vy, const Nd4jLong *yShapeInfo,
void *vz, const Nd4jLong *zShapeInfo,
const void *vPadVal) {
const X padVal = *reinterpret_cast<const X*>(vPadVal);
const auto x = reinterpret_cast<const X*>(vx);
const auto y = reinterpret_cast<const Y*>(vy);
auto z = reinterpret_cast<X*>(vz);
__shared__ int rank, rankMinusOne;
__shared__ Nd4jLong zLen, totalThreads, *coords, *xShape, *zShape, shift1, shift2, yStride0;
if (threadIdx.x == 0) {
extern __shared__ unsigned char shmem[];
coords = reinterpret_cast<Nd4jLong*>(shmem);
zLen = shape::length(zShapeInfo);
xShape = shape::shapeOf(const_cast<Nd4jLong*>(xShapeInfo));
zShape = shape::shapeOf(const_cast<Nd4jLong*>(zShapeInfo));
yStride0 = shape::stride(const_cast<Nd4jLong*>(yShapeInfo))[0];
rank = shape::rank(xShapeInfo);
zLen = shape::length(zShapeInfo);
rankMinusOne = rank - 1;
totalThreads = gridDim.x * blockDim.x;
shift1 = mode == 1 ? 0 : 1; // REFLECT : SYMMETRIC
shift2 = mode == 1 ? 2 : 1; // REFLECT : SYMMETRIC
}
__syncthreads();
auto xzCoord = coords + threadIdx.x * rank; // we use xzCoord storage both for x and z arrays
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
if(mode == 0) { // CONSTANT case
for (Nd4jLong i = tid; i < zLen; i += totalThreads) {
shape::index2coords(i, zShapeInfo, xzCoord);
const auto zOffset = shape::getOffset(zShapeInfo, xzCoord);
bool within = true;
for(int j = rankMinusOne; j >= 0; --j) {
if(xShape[j] == zShape[j]) continue;
const auto left = y[shape::getIndexOffset(yStride0 * j, yShapeInfo)];
if(xzCoord[j] < left || xzCoord[j] >= left + xShape[j]) {within = false; break;}
else {xzCoord[j] = xzCoord[j] - left;}
}
if(within)
z[zOffset] = x[shape::getOffset(xShapeInfo, xzCoord)];
else
z[zOffset] = padVal;
}
}
else { // REFLECT and SYMMETRIC cases
for (Nd4jLong i = tid; i < zLen; i += totalThreads) {
shape::index2coords(i, zShapeInfo, xzCoord);
const auto zOffset = shape::getOffset(zShapeInfo, xzCoord);
for(int j = rankMinusOne; j >= 0; --j) {
if(xShape[j] == zShape[j]) continue;
xzCoord[j] = xzCoord[j] - y[shape::getIndexOffset(yStride0 * j, yShapeInfo)]; // are ready to fill middle (within input dimension range)
if(xzCoord[j] < 0) xzCoord[j] = -xzCoord[j] - shift1; // means fill from left
else if(xzCoord[j] >= xShape[j]) xzCoord[j] = 2 * xShape[j] - xzCoord[j] - shift2; // means fill from right
}
const auto xOffset = shape::getOffset(xShapeInfo, xzCoord);
z[zOffset] = x[xOffset];
}
}
}
///////////////////////////////////////////////////////////////////
template<typename X, typename Y>
static void padCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem, const cudaStream_t *stream,
const int mode,
const void *vx, const Nd4jLong *xShapeInfo,
const void *vy, const Nd4jLong *yShapeInfo,
void *vz, const Nd4jLong *zShapeInfo,
const void* padVal) {
padCuda<X,Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(mode, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, padVal);
}
///////////////////////////////////////////////////////////////////
void pad(sd::LaunchContext * context, const int mode, const NDArray& input, const NDArray& paddings, NDArray& output, const NDArray& padValue) {
PointersManager manager(context, "pad");
NDArray::prepareSpecialUse({&output}, {&input, &paddings, &padValue});
const int threadsPerBlock = MAX_NUM_THREADS / 4;
const int blocksPerGrid = (output.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
const int sharedMem = 8 * threadsPerBlock * output.rankOf() + 128;
const auto xType = input.dataType();
const auto yType = paddings.dataType();
BUILD_DOUBLE_SELECTOR(xType, yType, padCudaLauncher, (blocksPerGrid, threadsPerBlock, sharedMem, context->getCudaStream(), mode, input.specialBuffer(), input.specialShapeInfo(), paddings.specialBuffer(), paddings.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), padValue.specialBuffer()), LIBND4J_TYPES, INDEXING_TYPES);
NDArray::registerSpecialUse({&output}, {&input, &paddings, &padValue});
manager.synchronize();
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
static __global__ void mirrorPadLinearKernel(void const* vx, const Nd4jLong* xShape, void* vz, const Nd4jLong* zShape, Nd4jLong leftSide, Nd4jLong leftSideCorrected, Nd4jLong xLen, Nd4jLong len, Nd4jLong zLen) {
__shared__ T const* x;
__shared__ T* z;
if (threadIdx.x == 0) {
x = reinterpret_cast<T const*>(vx);
z = reinterpret_cast<T*>(vz);
}
__syncthreads();
auto start = blockIdx.x * blockDim.x + threadIdx.x;
auto step = blockDim.x * gridDim.x;
for(int i = start; i < zLen; i+= step) {
auto zIndex = shape::getIndexOffset(i, zShape);
auto xIndex = shape::getIndexOffset(len - i, xShape);
if (i < leftSide) // left side
xIndex = shape::getIndexOffset(leftSideCorrected - i, xShape);
else if(i >= leftSide && i < leftSide + xLen) // middle
xIndex = shape::getIndexOffset(i - leftSide, xShape);
// else // right side
// z[i] = x[len - i];
z[zIndex] = x[xIndex];
}
}
template <typename F, typename I>
static __global__ void mirrorPadKernel(void const* vx, const Nd4jLong* xShape, void* vz, const Nd4jLong* zShape, Nd4jLong outLen, void const* paddings, const Nd4jLong* paddingShape, int reflBorder) {
__shared__ F const* x;
__shared__ I const* pads;
__shared__ F* z;
__shared__ Nd4jLong zRank, rank;
__shared__ Nd4jLong* xIdx;
if (threadIdx.x == 0) {
extern __shared__ unsigned char shmem[];
xIdx = reinterpret_cast<Nd4jLong*>(shmem);
rank = shape::rank(xShape);
x = reinterpret_cast<F const*>(vx);//
pads = reinterpret_cast<I const*>(paddings);
z = reinterpret_cast<F*>(vz);
}
__syncthreads();
auto start = threadIdx.x + blockIdx.x * blockDim.x;
auto step = blockDim.x * gridDim.x;
for(Nd4jLong i = start; i < outLen; i+= step) {
auto xzCoord = xIdx + threadIdx.x * rank;
//auto zxCoord = xIdx + (threadIdx.x + threadIdx.x % 2 + 1) * rank;
shape::index2coords(i, zShape, xzCoord);
auto outOffset = shape::getOffset(zShape, xzCoord);
// auto intStep = blockDim.y * gridDim.y;
for(int j = 0; j < rank; j++) {
const Nd4jLong inLen = shape::sizeAt(xShape, j);
Nd4jLong coords[2] = {j, 0};
auto padOffset = shape::getOffset(paddingShape, coords); // padding already has rank 2
const auto leftSide = pads[padOffset];
const auto leftSideCorrected = leftSide - reflBorder;
const Nd4jLong len = 2 * (inLen - 1) + leftSide + reflBorder;
if(xzCoord[j] < leftSide) // left side
xzCoord[j] = leftSideCorrected - xzCoord[j];
else if(xzCoord[j] >= leftSide && xzCoord[j] < leftSide + inLen) // middle
xzCoord[j] = xzCoord[j] - leftSide;
else if (len > xzCoord[j]) // right side
xzCoord[j] = len - xzCoord[j];
else
xzCoord[j] = xzCoord[j] - len;
}
auto inOffset = shape::getOffset(xShape, xzCoord);
z[outOffset] = x[inOffset];
}
}
template<typename F, typename I>
static void mirrorPad_(sd::LaunchContext * context, const NDArray& input, const NDArray& paddings, NDArray& output, const int mode) {
// mode: 0 - REFLECT, else - SYMMETRIC
const int reflBorder = (bool)mode ? 1 : 0;
const int rank = input.rankOf();
const Nd4jLong outLen = output.lengthOf();
auto stream = context->getCudaStream();
NDArray::prepareSpecialUse({&output}, {&input, &paddings});
if(rank <= 1) {
const Nd4jLong inLen = input.lengthOf();
const auto leftSide = paddings.e<Nd4jLong>(0);
const auto leftSideCorrected = leftSide - reflBorder;
const Nd4jLong len = 2*(inLen-1) + leftSide + reflBorder;
mirrorPadLinearKernel<F><<<256, 512, 256, *stream>>>(input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), leftSide, leftSideCorrected, inLen, len, outLen);
sd::DebugHelper::checkErrorCode(stream, "helpers::mirrorPadLinearKernel(...) failed");
}
else {
const int threadsPerBlock = MAX_NUM_THREADS / 2;
const int blocksPerGrid = (outLen + threadsPerBlock - 1) / threadsPerBlock;
const int sharedMem = threadsPerBlock * sizeof(Nd4jLong) * input.rankOf() + 256;
mirrorPadKernel<F, I><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), outLen, paddings.specialBuffer(), paddings.specialShapeInfo(), reflBorder);
sd::DebugHelper::checkErrorCode(stream, "helpers::mirrorPadKernel(...) failed");
}
NDArray::registerSpecialUse({&output}, {&input, &paddings});
}
void mirrorPad(sd::LaunchContext * context, const NDArray& input, const NDArray& paddings, NDArray& output, const int mode) {
BUILD_DOUBLE_SELECTOR(input.dataType(), paddings.dataType(), mirrorPad_, (context, input, paddings, output, mode), LIBND4J_TYPES, INDEXING_TYPES);
}
}
}
}