wehub-resource-sync
678 lines
31 KiB
Plaintext
678 lines
31 KiB
Plaintext
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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 GS <sgazeos@gmail.com>
|
|
//
|
|
#include <array/NDArrayFactory.h>
|
|
#include <exceptions/cuda_exception.h>
|
|
#include <execution/cuda/LaunchDims.h>
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
|
|
#include <ops/declarable/helpers/segment.h>
|
|
#include <ops/declarable/helpers/segment_common.h>
|
|
#include <system/selective_rendering.h>
|
|
|
|
#include "helpers/DebugHelper.h"
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// Segment ops linear kernels
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void segmentMaxLinearKernel(void* input, LongType const* inputShape, LongType* starts,
|
|
LongType* lengths, LongType numOfClasses, void* output,
|
|
LongType const* outputShape) {
|
|
__shared__ T* val;
|
|
__shared__ LongType xLen, zLen, zIndex;
|
|
__shared__ T* x;
|
|
__shared__ T* z;
|
|
__shared__ LongType threadsPerSegment, start, finish;
|
|
|
|
// Cache shape information
|
|
__shared__ sd::LongType inputRank, outputRank;
|
|
__shared__ const sd::LongType* inputStridePtr;
|
|
__shared__ const sd::LongType* outputStridePtr;
|
|
|
|
auto segment = blockIdx.x;
|
|
if (threadIdx.x == 0) {
|
|
x = reinterpret_cast<T*>(input);
|
|
z = reinterpret_cast<T*>(output);
|
|
extern __shared__ unsigned char shmem[];
|
|
val = reinterpret_cast<T*>(shmem);
|
|
xLen = shape::length(inputShape);
|
|
zLen = shape::length(outputShape);
|
|
|
|
// Cache shape information
|
|
inputRank = shape::rank(inputShape);
|
|
outputRank = shape::rank(outputShape);
|
|
inputStridePtr = shape::stride(inputShape);
|
|
outputStridePtr = shape::stride(outputShape);
|
|
|
|
if (segment < numOfClasses) {
|
|
LongType segmentCoords[] = {segment};
|
|
COORDS2INDEX(1, outputStridePtr, segmentCoords, zIndex);
|
|
start = starts[segment];
|
|
finish = start + lengths[segment];
|
|
LongType startCoords[] = {start};
|
|
LongType xOffset;
|
|
COORDS2INDEX(1, inputStridePtr, startCoords, xOffset);
|
|
z[zIndex] = x[xOffset];
|
|
val[segment] = z[zIndex];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for (auto e = start + threadIdx.x + 1; e < finish; e += blockDim.x) {
|
|
LongType eCoords[] = {e};
|
|
LongType xIndex;
|
|
COORDS2INDEX(1, inputStridePtr, eCoords, xIndex);
|
|
math::atomics::sd_atomicMax<T>(&z[zIndex], x[xIndex]);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void unsortedSegmentMaxLinearKernel(void* input, LongType const* inputShape, void* indices,
|
|
LongType const* indicesShape, LongType* starts,
|
|
LongType* lengths, LongType numOfClasses, void* output,
|
|
LongType const* outputShape) {
|
|
__shared__ LongType xLen, zLen, zIndex;
|
|
__shared__ T* x;
|
|
__shared__ T* z;
|
|
__shared__ I* y;
|
|
|
|
// Cache shape information
|
|
__shared__ sd::LongType inputRank, outputRank, indicesRank;
|
|
__shared__ const sd::LongType* inputStridePtr;
|
|
__shared__ const sd::LongType* outputStridePtr;
|
|
__shared__ const sd::LongType* indicesStridePtr;
|
|
|
|
auto segment = blockIdx.x;
|
|
|
|
if (threadIdx.x == 0) {
|
|
x = reinterpret_cast<T*>(input);
|
|
z = reinterpret_cast<T*>(output);
|
|
y = reinterpret_cast<I*>(indices);
|
|
xLen = shape::length(inputShape);
|
|
zLen = shape::length(outputShape);
|
|
|
|
// Cache shape information
|
|
inputRank = shape::rank(inputShape);
|
|
outputRank = shape::rank(outputShape);
|
|
indicesRank = shape::rank(indicesShape);
|
|
inputStridePtr = shape::stride(inputShape);
|
|
outputStridePtr = shape::stride(outputShape);
|
|
indicesStridePtr = shape::stride(indicesShape);
|
|
|
|
LongType segmentCoords[] = {segment};
|
|
COORDS2INDEX(1, outputStridePtr, segmentCoords, zIndex);
|
|
if (lengths[segment] > 0) {
|
|
LongType startCoords[] = {starts[segment]};
|
|
LongType xOffset;
|
|
COORDS2INDEX(1, inputStridePtr, startCoords, xOffset);
|
|
z[zIndex] = x[xOffset];
|
|
} else {
|
|
z[zIndex] = -DataTypeUtils::max<T>();
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
if (lengths[segment] > 0) {
|
|
for (auto e = threadIdx.x + 1; e < xLen; e += blockDim.x) {
|
|
LongType eCoords[] = {e};
|
|
LongType xIndex, yIndex;
|
|
COORDS2INDEX(1, inputStridePtr, eCoords, xIndex);
|
|
COORDS2INDEX(1, indicesStridePtr, eCoords, yIndex);
|
|
if (y[yIndex] == segment) {
|
|
math::atomics::sd_atomicMax<T>(&z[zIndex], x[xIndex]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void segmentMaxTadKernel(void* inputBuf, LongType const* inputShape, LongType const* inputTads,
|
|
LongType const* inputTadOffsets, I* indices, LongType* starts,
|
|
LongType* lengths, LongType numOfClasses, void* outputBuf,
|
|
LongType const* outputShape, LongType const* outputTads,
|
|
LongType const* outputTadOffsets, T filler, LongType indicesLength,
|
|
LongType numInputTads, LongType numOutputTads) {
|
|
__shared__ T* val;
|
|
__shared__ LongType len, zIndex, total, zLen;
|
|
__shared__ T* z;
|
|
__shared__ int start, finish;
|
|
__shared__ I segment;
|
|
|
|
// Cache shape information
|
|
__shared__ sd::LongType inputTadRank, outputTadRank;
|
|
__shared__ const sd::LongType* inputTadShapePtr;
|
|
__shared__ const sd::LongType* outputTadShapePtr;
|
|
__shared__ const sd::LongType* inputTadStridePtr;
|
|
__shared__ const sd::LongType* outputTadStridePtr;
|
|
|
|
if (threadIdx.x == 0 && blockIdx.x < indicesLength) {
|
|
segment = indices[blockIdx.x];
|
|
zLen = shape::length(outputShape);
|
|
auto zOffset = outputTadOffsets[segment];
|
|
z = reinterpret_cast<T*>(outputBuf) + outputTadOffsets[segment];
|
|
len = shape::length(inputTads);
|
|
|
|
// Cache shape information
|
|
inputTadRank = shape::rank(inputTads);
|
|
outputTadRank = shape::rank(outputTads);
|
|
inputTadShapePtr = shape::shapeOf(inputTads);
|
|
outputTadShapePtr = shape::shapeOf(outputTads);
|
|
inputTadStridePtr = shape::stride(inputTads);
|
|
outputTadStridePtr = shape::stride(outputTads);
|
|
|
|
start = starts[segment];
|
|
finish = start + lengths[segment];
|
|
total = shape::sizeAt(inputShape, 0);
|
|
}
|
|
__syncthreads();
|
|
|
|
auto idx = blockIdx.x;
|
|
if (idx < numInputTads) {
|
|
auto x = reinterpret_cast<T*>(inputBuf) + inputTadOffsets[idx];
|
|
if (blockIdx.x == start) {
|
|
for (auto e = threadIdx.x; e < len; e += blockDim.x) {
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xIndex;
|
|
LongType zIndex;
|
|
|
|
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, xCoords);
|
|
COORDS2INDEX(inputTadRank, inputTadStridePtr, xCoords, xIndex);
|
|
INDEX2COORDS(e, outputTadRank, outputTadShapePtr, zCoords);
|
|
COORDS2INDEX(outputTadRank, outputTadStridePtr, zCoords, zIndex);
|
|
|
|
math::atomics::sd_atomicMax<T>(&z[zIndex], x[xIndex]);
|
|
}
|
|
} else {
|
|
for (auto e = threadIdx.x; e < len; e += blockDim.x) {
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xIndex;
|
|
LongType zIndex;
|
|
|
|
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, xCoords);
|
|
COORDS2INDEX(inputTadRank, inputTadStridePtr, xCoords, xIndex);
|
|
INDEX2COORDS(e, outputTadRank, outputTadShapePtr, zCoords);
|
|
COORDS2INDEX(outputTadRank, outputTadStridePtr, zCoords, zIndex);
|
|
|
|
if (lengths[segment]) math::atomics::sd_atomicMax<T>(&z[zIndex], x[xIndex]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
|
|
template <typename T, typename I>
|
|
static void segmentMaxFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
T val = -DataTypeUtils::max<T>();
|
|
output->assign(val);
|
|
auto stream = context->getCudaStream();
|
|
indices->syncToHost();
|
|
LongType numOfClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
|
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
|
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
|
sd::LongType len = indices->lengthOf();
|
|
classesRangesBegs.assign(len);
|
|
int zero2 = 0;
|
|
classesRangesLens.assign(zero2);
|
|
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
|
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
|
fillUpSegments(indices, numOfClasses, classesRangesBegs, classesRangesLens);
|
|
|
|
NDArray::prepareSpecialUse({output}, {input, indices, &classesRangesBegs, &classesRangesLens});
|
|
|
|
if (input->isVector() || input->isScalar()) {
|
|
dim3 launchDims = segmentDims(numOfClasses,input->lengthOf());
|
|
segmentMaxLinearKernel<T, I><<<launchDims.y,launchDims.x,launchDims.z, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), begins, lengths, numOfClasses, output->specialBuffer(),
|
|
output->specialShapeInfo());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxLinearKernel failed");
|
|
|
|
} else {
|
|
LongType zero = 0;
|
|
std::vector<LongType> *dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,&zero);
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
|
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);
|
|
auto inputTads = packX->specialShapeInfo();
|
|
auto inputTadOffsets = packX->specialOffsets();
|
|
auto outputTads = packZ->specialShapeInfo();
|
|
auto outputTadOffsets = packZ->specialOffsets();
|
|
dim3 launchDims = segmentTad(packX->numberOfTads());
|
|
segmentMaxTadKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
|
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numOfClasses, output->specialBuffer(),
|
|
output->specialShapeInfo(), outputTads, outputTadOffsets,static_cast<T>(0),
|
|
indices->lengthOf(),packX->numberOfTads(),packZ->numberOfTads());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxTadKernel failed");
|
|
|
|
delete dimensions;
|
|
}
|
|
NDArray::registerSpecialUse({output}, {input, indices, &classesRangesBegs, &classesRangesLens});
|
|
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
void segmentMaxFunctor(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
|
NDArray::prepareSpecialUse({output}, {input, indices});
|
|
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), segmentMaxFunctor_, (context, input, indices, output),
|
|
SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, indices});
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
|
|
template <typename T, typename I>
|
|
static void unsortedSegmentMaxFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
|
auto stream = context->getCudaStream();
|
|
T val = DataTypeUtils::infOrMax<T>();
|
|
output->assign(val);
|
|
|
|
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
|
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
|
int zero2 = 0;
|
|
sd::LongType len = indices->lengthOf();
|
|
classesRangesBegs.assign(len);
|
|
classesRangesLens.assign(zero2);
|
|
|
|
dim3 dims = getFillUpSegmentsDims(numOfClasses, indices->lengthOf());
|
|
fillUpSegments(indices, numOfClasses, classesRangesBegs, classesRangesLens);
|
|
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
|
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
|
|
|
if (input->isVector() || input->isScalar()) {
|
|
unsortedSegmentMaxLinearKernel<T, I><<<dims.x, dims.y, dims.z, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
|
begins, lengths, numOfClasses, output->specialBuffer(), output->specialShapeInfo());
|
|
sd::DebugHelper::checkErrorCode(stream, "unsortedSegmentMaxLinearKernel failed");
|
|
|
|
} else {
|
|
LongType zero = 0;
|
|
std::vector<LongType> *dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,&zero);
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
|
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);
|
|
auto inputTads = packX->specialShapeInfo();
|
|
auto inputTadOffsets = packX->specialOffsets();
|
|
auto outputTads = packZ->specialShapeInfo();
|
|
auto outputTadOffsets = packZ->specialOffsets();
|
|
dims.x = input->sizeAt(0);
|
|
T val = -DataTypeUtils::max<T>();
|
|
output->assign(val);
|
|
segmentMaxTadKernel<T, I><<<dims.x, dims.y, dims.z, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
|
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numOfClasses, output->specialBuffer(),
|
|
output->specialShapeInfo(), outputTads, outputTadOffsets,static_cast<T>(0),indices->lengthOf(),packX->numberOfTads(),packZ->numberOfTads());
|
|
|
|
delete dimensions;
|
|
}
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
void unsortedSegmentMaxFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses,
|
|
NDArray* output) {
|
|
NDArray::prepareSpecialUse({output}, {input, indices});
|
|
output->nullify();
|
|
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentMaxFunctor_,
|
|
(context, input, indices, numOfClasses, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, indices});
|
|
}
|
|
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
// segment max
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void segmentMaxBPLinearKernel(void* inputBuf, LongType const* inputShape, void* forwardOutput,
|
|
LongType const* forwardShape, void* eps, LongType const* epsShape, void* indicesBuf, LongType const* indicesShape, void* outputBuf,
|
|
LongType const* outputShape, LongType indicesLen) {
|
|
__shared__ T* x;
|
|
__shared__ T* gradIn;
|
|
__shared__ T* gradOut;
|
|
__shared__ I* y;
|
|
__shared__ T* z;
|
|
__shared__ LongType xLen, gradLen;
|
|
|
|
// Cache shape/stride/rank information in shared memory
|
|
__shared__ sd::LongType xRank, yRank, zRank, gradInRank, gradOutRank;
|
|
__shared__ const sd::LongType *xShapePtr, *yShapePtr, *zShapePtr, *gradInShapePtr, *gradOutShapePtr;
|
|
__shared__ const sd::LongType *xStridePtr, *yStridePtr, *zStridePtr, *gradInStridePtr, *gradOutStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
xLen = shape::length(inputShape);
|
|
x = reinterpret_cast<T*>(inputBuf);
|
|
y = reinterpret_cast<I*>(indicesBuf);
|
|
z = reinterpret_cast<T*>(outputBuf);
|
|
gradIn = reinterpret_cast<T*>(forwardOutput);
|
|
gradOut = reinterpret_cast<T*>(eps);
|
|
gradLen = shape::length(epsShape);
|
|
|
|
// Cache all shape information
|
|
xRank = shape::rank(inputShape);
|
|
yRank = shape::rank(indicesShape);
|
|
zRank = shape::rank(outputShape);
|
|
gradInRank = shape::rank(forwardShape);
|
|
gradOutRank = shape::rank(epsShape);
|
|
|
|
xShapePtr = shape::shapeOf(inputShape);
|
|
yShapePtr = shape::shapeOf(indicesShape);
|
|
zShapePtr = shape::shapeOf(outputShape);
|
|
gradInShapePtr = shape::shapeOf(forwardShape);
|
|
gradOutShapePtr = shape::shapeOf(epsShape);
|
|
|
|
xStridePtr = shape::stride(inputShape);
|
|
yStridePtr = shape::stride(indicesShape);
|
|
zStridePtr = shape::stride(outputShape);
|
|
gradInStridePtr = shape::stride(forwardShape);
|
|
gradOutStridePtr = shape::stride(epsShape);
|
|
}
|
|
__syncthreads();
|
|
|
|
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
|
auto step = gridDim.x * blockDim.x;
|
|
|
|
for (auto e = start; e < indicesLen; e += step) {
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType yCoords[SD_MAX_RANK];
|
|
LongType gradICoords[SD_MAX_RANK];
|
|
LongType gradOCoords[SD_MAX_RANK];
|
|
LongType zOffset;
|
|
LongType xOffset;
|
|
LongType yOffset;
|
|
LongType gradOffsetI;
|
|
LongType gradOffsetO;
|
|
|
|
INDEX2COORDS(e, zRank, zShapePtr, zCoords);
|
|
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
|
INDEX2COORDS(e, xRank, xShapePtr, xCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
|
INDEX2COORDS(e, yRank, yShapePtr, yCoords);
|
|
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
|
|
|
|
auto classIndex = y[yOffset];
|
|
INDEX2COORDS(classIndex, gradInRank, gradInShapePtr, gradICoords);
|
|
COORDS2INDEX(gradInRank, gradInStridePtr, gradICoords, gradOffsetI);
|
|
INDEX2COORDS(classIndex, gradOutRank, gradOutShapePtr, gradOCoords);
|
|
COORDS2INDEX(gradOutRank, gradOutStridePtr, gradOCoords, gradOffsetO);
|
|
|
|
if (math::sd_abs<T,T>(gradIn[gradOffsetI] - x[xOffset]) <= T(1.e-6)) {
|
|
z[zOffset] = gradOut[gradOffsetO];
|
|
}
|
|
}
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void segmentMaxBPTadKernel(void* inputBuf, LongType const* inputShape,
|
|
void* forwardOutput,
|
|
LongType const* forwardShape,
|
|
void* eps, LongType const* epsShape,
|
|
void* indicesBuf, LongType const* indicesShape,
|
|
void* outputBuf,
|
|
LongType const* outputShape, LongType const* inputTadShapeInfo,
|
|
LongType const* inputOffsets, LongType const* gradInTadShapeInfo,
|
|
LongType const* gradInOffsets, LongType const* gradOutTadShapeInfo,
|
|
LongType const* gradOutOffsets, LongType const* outTadShapeInfo,
|
|
LongType const* outOffsets, LongType indicesLen) {
|
|
__shared__ T* x;
|
|
__shared__ I *indices;
|
|
__shared__ T* gradIn;
|
|
__shared__ T* gradOut;
|
|
__shared__ I* y;
|
|
__shared__ T* z;
|
|
__shared__ LongType xLen, yLen, gradLen, currentLen, gradOutLen, inLen;
|
|
|
|
// Cache shape information for all TADs
|
|
__shared__ sd::LongType inputTadRank;
|
|
__shared__ const sd::LongType* inputTadShapePtr;
|
|
__shared__ const sd::LongType* inputTadStridePtr;
|
|
|
|
__shared__ sd::LongType gradInTadRank;
|
|
__shared__ const sd::LongType* gradInTadShapePtr;
|
|
__shared__ const sd::LongType* gradInTadStridePtr;
|
|
|
|
__shared__ sd::LongType gradOutTadRank;
|
|
__shared__ const sd::LongType* gradOutTadShapePtr;
|
|
__shared__ const sd::LongType* gradOutTadStridePtr;
|
|
|
|
__shared__ sd::LongType outTadRank;
|
|
__shared__ const sd::LongType* outTadShapePtr;
|
|
__shared__ const sd::LongType* outTadStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
xLen = shape::length(inputShape);
|
|
indices = reinterpret_cast<I*>(indicesBuf);
|
|
x = reinterpret_cast<T*>(inputBuf);
|
|
y = reinterpret_cast<I*>(indicesBuf);
|
|
z = reinterpret_cast<T*>(outputBuf);
|
|
yLen = shape::length(indicesShape);
|
|
gradOut = reinterpret_cast<T*>(eps);
|
|
gradIn = reinterpret_cast<T*>(forwardOutput);
|
|
gradLen = shape::length(epsShape);
|
|
inLen = shape::length(gradInTadShapeInfo);
|
|
gradOutLen = shape::length(gradOutTadShapeInfo);
|
|
currentLen = shape::length(inputTadShapeInfo);
|
|
|
|
// Cache all TAD shape information
|
|
inputTadRank = shape::rank(inputTadShapeInfo);
|
|
inputTadShapePtr = shape::shapeOf(inputTadShapeInfo);
|
|
inputTadStridePtr = shape::stride(inputTadShapeInfo);
|
|
|
|
gradInTadRank = shape::rank(gradInTadShapeInfo);
|
|
gradInTadShapePtr = shape::shapeOf(gradInTadShapeInfo);
|
|
gradInTadStridePtr = shape::stride(gradInTadShapeInfo);
|
|
|
|
gradOutTadRank = shape::rank(gradOutTadShapeInfo);
|
|
gradOutTadShapePtr = shape::shapeOf(gradOutTadShapeInfo);
|
|
gradOutTadStridePtr = shape::stride(gradOutTadShapeInfo);
|
|
|
|
outTadRank = shape::rank(outTadShapeInfo);
|
|
outTadShapePtr = shape::shapeOf(outTadShapeInfo);
|
|
outTadStridePtr = shape::stride(outTadShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
for (auto i = blockIdx.x; i < indicesLen; i += gridDim.x) {
|
|
I segment = indices[i];
|
|
T* current = x;
|
|
T* currentOut = z;
|
|
auto classNum = segment;
|
|
auto currentOffset = inputOffsets[i];
|
|
auto currentOutOffset = outOffsets[i];
|
|
auto currentGradOutOffset = gradOutOffsets[classNum];
|
|
auto bPTensorOffset = gradInOffsets[classNum];
|
|
|
|
auto gradIn2 = gradIn + bPTensorOffset;
|
|
auto current2 = current + currentOffset;
|
|
auto currentGradOut2 = gradOut + currentGradOutOffset;
|
|
auto currentOut2 = currentOut + currentOutOffset;
|
|
|
|
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
|
sd::LongType xCoords[SD_MAX_RANK];
|
|
sd::LongType gradInCoords[SD_MAX_RANK];
|
|
sd::LongType gradOutCoords[SD_MAX_RANK];
|
|
sd::LongType outCoords[SD_MAX_RANK];
|
|
sd::LongType xIndex;
|
|
sd::LongType gradInIndex;
|
|
sd::LongType gradOutIndex;
|
|
sd::LongType outIndex;
|
|
|
|
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, xCoords);
|
|
COORDS2INDEX(inputTadRank, inputTadStridePtr, xCoords, xIndex);
|
|
INDEX2COORDS(e, gradInTadRank, gradInTadShapePtr, gradInCoords);
|
|
COORDS2INDEX(gradInTadRank, gradInTadStridePtr, gradInCoords, gradInIndex);
|
|
INDEX2COORDS(e, gradOutTadRank, gradOutTadShapePtr, gradOutCoords);
|
|
COORDS2INDEX(gradOutTadRank, gradOutTadStridePtr, gradOutCoords, gradOutIndex);
|
|
INDEX2COORDS(e, outTadRank, outTadShapePtr, outCoords);
|
|
COORDS2INDEX(outTadRank, outTadStridePtr, outCoords, outIndex);
|
|
|
|
if (math::sd_abs<T, T>(gradIn2[gradInIndex] - current2[xIndex]) <= T(1.e-6)) {
|
|
currentOut2[outIndex] = currentGradOut2[gradOutIndex];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
template <typename T, typename I>
|
|
Status segmentMaxFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
auto stream = context->getCudaStream();
|
|
/* NDArray tempRes(gradOut->ordering(), gradOut->getShapeAsVector(), DataTypeUtils::fromT<T>(),
|
|
context); */
|
|
auto outShape = gradOut->getShapeAsVector();
|
|
NDArray tempRes(gradOut->ordering(), outShape, DataTypeUtils::fromT<T>(), context);
|
|
segmentMaxFunctor_<T, I>(context, input, indices, &tempRes);
|
|
NDArray::prepareSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
|
if (input->isVector() || input->isScalar()) {
|
|
LongType loop_size = input->lengthOf();
|
|
auto numOfClasses = gradOut->lengthOf();
|
|
dim3 segmentBpDims2 = segmentBpDims(1 + gradOut->lengthOf(),input->lengthOf());
|
|
segmentMaxBPLinearKernel<T, I><<<segmentBpDims2.y, segmentBpDims2.x, segmentBpDims2.z, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), tempRes.specialBuffer(), tempRes.specialShapeInfo(),
|
|
gradOut->specialBuffer(), gradOut->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
|
output->specialBuffer(), output->specialShapeInfo(), indices->lengthOf());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxBPLinearKernel failed");
|
|
|
|
} else {
|
|
LongType zero = 0;
|
|
std::vector<LongType> *dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,&zero);
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
|
|
|
NDArray::preparePrimaryUse({&tempRes}, {&tempRes});
|
|
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);
|
|
auto packGradIn = ConstantTadHelper::getInstance().tadForDimensions(tempRes.shapeInfo(), dimensions);
|
|
auto packGradOut = ConstantTadHelper::getInstance().tadForDimensions(gradOut->shapeInfo(), dimensions);
|
|
LongType const* inputTadShapeInfo = packX->specialShapeInfo();
|
|
LongType const* inputTadOffsets = packX->specialOffsets();
|
|
LongType const* outputTadShapeInfo = packZ->specialShapeInfo();
|
|
LongType const* outputTadOffsets = packZ->specialOffsets();
|
|
LongType const* gradInTadShapeInfo = packGradIn->specialShapeInfo();
|
|
LongType const* gradInTadOffsets = packGradIn->specialOffsets();
|
|
LongType const* gradOutTadShapeInfo = packGradOut->specialShapeInfo();
|
|
LongType const* gradOutTadOffsets = packGradOut->specialOffsets();
|
|
dim3 segmentBpTad2 = segmentBpTad(gradOut->lengthOf(),input->lengthOf());
|
|
segmentMaxBPTadKernel<T, I><<<segmentBpTad2.x, segmentBpTad2.y, segmentBpTad2.z, *stream>>>(
|
|
input->specialBuffer(),
|
|
input->specialShapeInfo(),
|
|
tempRes.specialBuffer(),
|
|
tempRes.specialShapeInfo(),
|
|
|
|
gradOut->specialBuffer(),
|
|
gradOut->specialShapeInfo(),
|
|
|
|
indices->specialBuffer(),
|
|
indices->specialShapeInfo(),
|
|
output->specialBuffer(),
|
|
output->specialShapeInfo(),
|
|
inputTadShapeInfo,
|
|
inputTadOffsets, gradInTadShapeInfo,
|
|
gradInTadOffsets, gradOutTadShapeInfo,
|
|
gradOutTadOffsets, outputTadShapeInfo,
|
|
outputTadOffsets,
|
|
indices->lengthOf());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxBPTadKernel failed");
|
|
|
|
delete dimensions;
|
|
}
|
|
NDArray::registerSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
|
return Status::OK;
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
Status segmentMaxFunctorBP(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
NDArray* output) {
|
|
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
|
auto indicesDType = indices->dataType();
|
|
auto outputDType = output->dataType();
|
|
BUILD_DOUBLE_SELECTOR(output->dataType(), indices->dataType(), return segmentMaxFunctorBP_,
|
|
(context, input, indices, gradOut, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
|
}
|
|
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
template <typename T, typename I>
|
|
static Status unsortedSegmentMaxFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices,
|
|
NDArray* gradOut,
|
|
LongType numOfClasses, NDArray* output) {
|
|
// if input is a vector: (as if in doc sample)
|
|
auto stream = context->getCudaStream();
|
|
auto outShape = gradOut->getShapeAsVector();
|
|
NDArray tempRes(gradOut->ordering(), outShape, DataTypeUtils::fromT<T>(),
|
|
context);
|
|
unsortedSegmentMaxFunctor_<T, I>(context, input, indices, numOfClasses, &tempRes);
|
|
NDArray::prepareSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
|
if (input->isVector() || input->isScalar()) {
|
|
LongType loop_size = input->lengthOf();
|
|
auto numOfClasses = gradOut->lengthOf(); // indices->e<sd::LongType>(loop_size - 1);
|
|
segmentMaxBPLinearKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
|
input->specialBuffer(), input->specialShapeInfo(), tempRes.specialBuffer(), tempRes.specialShapeInfo(),
|
|
gradOut->specialBuffer(), gradOut->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
|
output->specialBuffer(), output->specialShapeInfo(),indices->lengthOf());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxBPLinearKernel failed");
|
|
} else {
|
|
LongType zero = 0;
|
|
std::vector<LongType> *dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,&zero);
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimensions);
|
|
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimensions);
|
|
auto packGradIn = ConstantTadHelper::getInstance().tadForDimensions(tempRes.shapeInfo(), dimensions);
|
|
auto packGradOut = ConstantTadHelper::getInstance().tadForDimensions(gradOut->shapeInfo(), dimensions);
|
|
LongType const* inputTads = packX->specialShapeInfo();
|
|
LongType const* inputTadOffsets = packX->specialOffsets();
|
|
LongType const* outputTads = packZ->specialShapeInfo();
|
|
LongType const* outputTadOffsets = packZ->specialOffsets();
|
|
LongType const* gradInTads = packGradIn->specialShapeInfo();
|
|
LongType const* gradInTadOffsets = packGradIn->specialOffsets();
|
|
LongType const* gradOutTads = packGradOut->specialShapeInfo();
|
|
LongType const* gradOutTadOffsets = packGradOut->specialOffsets();
|
|
|
|
|
|
segmentMaxBPTadKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
|
input->specialBuffer(),
|
|
input->specialShapeInfo(),
|
|
tempRes.specialBuffer(),
|
|
tempRes.specialShapeInfo(),
|
|
gradOut->specialBuffer(),
|
|
gradOut->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
|
output->specialBuffer(), output->specialShapeInfo(), inputTads, inputTadOffsets, gradInTads, gradInTadOffsets,
|
|
gradOutTads, gradOutTadOffsets, outputTads, outputTadOffsets, indices->lengthOf());
|
|
sd::DebugHelper::checkErrorCode(stream, "segmentMaxBPTadKernel failed");
|
|
|
|
|
|
delete dimensions;
|
|
}
|
|
NDArray::registerSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
|
return Status::OK;
|
|
}
|
|
// -------------------------------------------------------------------------------------------------------------- //
|
|
Status unsortedSegmentMaxFunctorBP(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
|
LongType numOfClasses, NDArray* output) {
|
|
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
|
auto indicesDType = indices->dataType();
|
|
auto outputDType = output->dataType();
|
|
BUILD_DOUBLE_SELECTOR(output->dataType(), indices->dataType(), return unsortedSegmentMaxFunctorBP_,
|
|
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
|
}
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|