chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,291 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 George A. Shulinok <sgazeos@gmail.com>, created on 4/18/2019
|
||||
//
|
||||
#include <ops/declarable/helpers/BarnesHutTsne.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// count rows kernel - count input pRows and pCols and put result onto pRowCounts
|
||||
// pRowCounts - array of ints, with length N
|
||||
// pRows - array of ints with length N, vals from 0 to N-1
|
||||
// pCols - array of ints with length < N and vals between 0 and max(pRows)
|
||||
//
|
||||
static SD_KERNEL void countRowsKernel(int* pRowCounts, int const* pRows, int const* pCols, LongType N) {
|
||||
auto start = blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
for (int n = threadIdx.x + start; n < N; n += step) {
|
||||
int begin = pRows[n]; //->e<int>(n);
|
||||
int end = pRows[n + 1]; // rowP->e<int>(n + 1);
|
||||
for (int i = begin; i < end; i++) {
|
||||
bool present = false;
|
||||
// loop between near pRows
|
||||
for (int m = pRows[pCols[i]]; m < pRows[pCols[i] + 1]; m++)
|
||||
if (pCols[m] == n) { // mark index as existed with columns array
|
||||
present = true;
|
||||
break;
|
||||
}
|
||||
|
||||
atomicAdd(&pRowCounts[n], 1);
|
||||
|
||||
if (!present) // increment row counter for given index
|
||||
atomicAdd(&pRowCounts[pCols[i]], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// row counter caller
|
||||
LongType barnes_row_count(NDArray* rowP, NDArray* colP, LongType N, NDArray& rowCounts) {
|
||||
int* pRowCounts = reinterpret_cast<int*>(rowCounts.specialBuffer());
|
||||
int const* pRows = reinterpret_cast<int const*>(rowP->specialBuffer());
|
||||
int const* pCols = reinterpret_cast<int const*>(colP->specialBuffer());
|
||||
auto stream = rowCounts.getContext()->getCudaStream();
|
||||
countRowsKernel<<<1, 1, 128, *stream>>>(pRowCounts, pRows, pCols, N);
|
||||
sd::DebugHelper::checkErrorCode(stream, "countRows failed");
|
||||
|
||||
NDArray numElementsArr = rowCounts.sumNumber();
|
||||
auto numElements = numElementsArr.e<LongType>(0);
|
||||
return numElements;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// extend symRowP with pRowCounts array vals
|
||||
// pRowCounts - int array with length N
|
||||
// symRowP - int array with length N+1
|
||||
// N - given array length
|
||||
//
|
||||
static SD_KERNEL void fillUpsymRow(int const* pRowCounts, int* symRowP, int N) {
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (int n = start; n < N + 1; n += step) { // to avoid race condition use shift only for given index
|
||||
symRowP[n] = 0;
|
||||
for (int i = 0; i < n; i++) atomicAdd(&symRowP[n], pRowCounts[i]);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// symmetrize routine kernel
|
||||
// pRows - rows buffer (ints)
|
||||
// pCols - column buffer (ints) with vals between 0 and max(pRows)
|
||||
// pVals - values vector (floats)
|
||||
// symRowP - ints, shifted pRows
|
||||
// symColP - ints, shifted pCols,
|
||||
// offset - ints, shitfs
|
||||
// pOutput - result matrix (floats)
|
||||
// N - pRows length
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void symmetrizeKernel(int const* pRows, int const* pCols, T const* pVals, int* symRowP, int* symColP,
|
||||
int* offset, T* pOutput, int N) {
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (int n = start; n < N; n += step) {
|
||||
int begin = pRows[n];
|
||||
int bound = pRows[n + 1];
|
||||
|
||||
for (int i = begin; i < bound; i++) {
|
||||
bool present = false;
|
||||
int colPI = pCols[i];
|
||||
int start = pRows[colPI];
|
||||
int end = pRows[colPI + 1];
|
||||
|
||||
for (int m = start; m < end; m++) {
|
||||
if (pCols[m] == n) {
|
||||
present = true;
|
||||
if (n <= colPI) {
|
||||
symColP[symRowP[n] + offset[n]] = colPI;
|
||||
symColP[symRowP[colPI] + offset[colPI]] = n;
|
||||
pOutput[symRowP[n] + offset[n]] = pVals[i] + pVals[m];
|
||||
pOutput[symRowP[colPI] + offset[colPI]] = pVals[i] + pVals[m];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// If (colP[i], n) is not present, there is no addition involved
|
||||
if (!present) {
|
||||
symColP[symRowP[n] + offset[n]] = colPI;
|
||||
symColP[symRowP[pCols[i]] + offset[colPI]] = n;
|
||||
pOutput[symRowP[n] + offset[n]] = pVals[i];
|
||||
pOutput[symRowP[colPI] + offset[colPI]] = pVals[i];
|
||||
}
|
||||
// Update offsets
|
||||
if (!present || (present && n <= colPI)) {
|
||||
atomicAdd(&offset[n], 1);
|
||||
|
||||
if (colPI != n) atomicAdd(&offset[colPI], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// symmetrize algorithm itself
|
||||
//
|
||||
template <typename T>
|
||||
static void barnes_symmetrize_(NDArray* rowP, NDArray* colP, NDArray* valP, LongType N,
|
||||
NDArray* outputRows, NDArray* outputCols, NDArray* outputVals, NDArray* rowCounts) {
|
||||
int const* pRows = reinterpret_cast<int const*>(rowP->specialBuffer());
|
||||
int* symRowP = reinterpret_cast<int*>(outputRows->specialBuffer());
|
||||
int* pRowCounts = reinterpret_cast<int*>(rowCounts->specialBuffer());
|
||||
auto stream = outputCols->getContext()->getCudaStream();
|
||||
// fill up syRowP array
|
||||
fillUpsymRow<<<1, N, 128, *stream>>>(pRowCounts, symRowP, N);
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillUpsymRow failed");
|
||||
|
||||
outputRows->syncToHost();
|
||||
int* symColP = reinterpret_cast<int*>(outputCols->specialBuffer());
|
||||
int const* pCols = reinterpret_cast<int const*>(colP->specialBuffer());
|
||||
T const* pVals = reinterpret_cast<T const*>(valP->specialBuffer());
|
||||
T* pOutput = reinterpret_cast<T*>(outputVals->specialBuffer());
|
||||
auto offsetArr = NDArrayFactory::create<int>('c', {N});
|
||||
int* offset = reinterpret_cast<int*>(offsetArr->specialBuffer());
|
||||
// symmetrize itself
|
||||
symmetrizeKernel<T><<<1, 1, 1024, *stream>>>(pRows, pCols, pVals, symRowP, symColP, offset, pOutput, N);
|
||||
sd::DebugHelper::checkErrorCode(stream, "symmetrizeKernel failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// symmetrize caller and adoption
|
||||
//
|
||||
void barnes_symmetrize(NDArray* rowP, NDArray* colP, NDArray* valP, LongType N,
|
||||
NDArray* outputRows, NDArray* outputCols, NDArray* outputVals, NDArray* rowCounts) {
|
||||
BUILD_SINGLE_SELECTOR(valP->dataType(), barnes_symmetrize_,
|
||||
(rowP, colP, valP, N, outputRows, outputCols, outputVals, rowCounts), SD_NUMERIC_TYPES);
|
||||
|
||||
*outputVals /= 2.0;
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void barnes_symmetrize_,
|
||||
(NDArray* rowP, NDArray* colP, NDArray* valP, sd::LongType N,
|
||||
NDArray* outputRows, NDArray* outputCols, NDArray* outputVals, NDArray* rowCounts),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// edge forces implementation
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void edgeForcesKernel(int const* pRows, int const* pCols, T const* dataP, T const* vals, T* outputP,
|
||||
int N, int colCount, int rowSize) {
|
||||
// std::vector<T> buffer(colCount);
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (int n = start; n < N; n += step) {
|
||||
int start = pRows[n];
|
||||
int end = pRows[n + 1];
|
||||
int shift = n * colCount;
|
||||
for (int i = start; i < end; i++) {
|
||||
T const* thisSlice = dataP + pCols[i] * colCount;
|
||||
T res = static_cast<T>(1);
|
||||
|
||||
for (int k = 0; k < colCount; k++) {
|
||||
auto valTemp = dataP[shift + k] - thisSlice[k]; // thisSlice[k];
|
||||
res += valTemp * valTemp; // (dataP[shift + k] * dataP[shift + k] - 2 * dataP[shift + k] * thisSlice[k] +
|
||||
// thisSlice[k] * thisSlice[k])
|
||||
}
|
||||
res = vals[i] / res;
|
||||
for (int k = 0; k < colCount; k++)
|
||||
math::atomics::sd_atomicAdd(&outputP[shift + k], T((dataP[shift + k] - thisSlice[k]) * res));
|
||||
}
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// edge forces algorithm
|
||||
//
|
||||
|
||||
template <typename T>
|
||||
static void barnes_edge_forces_(NDArray* rowP, NDArray * colP, NDArray * valP, int N,
|
||||
NDArray * data, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {data, rowP, colP, valP, valP});
|
||||
T const* dataP = reinterpret_cast<T const*>(data->specialBuffer());
|
||||
T const* vals = reinterpret_cast<T const*>(valP->specialBuffer());
|
||||
T* outputP = reinterpret_cast<T*>(output->specialBuffer());
|
||||
int const* pRows = reinterpret_cast<int const*>(rowP->specialBuffer());
|
||||
int const* pCols = reinterpret_cast<int const*>(colP->specialBuffer());
|
||||
int colCount = data->columns();
|
||||
// auto shift = 0;
|
||||
auto rowSize = sizeof(T) * colCount;
|
||||
auto stream = output->getContext()->getCudaStream();
|
||||
edgeForcesKernel<T><<<1, 128, 1024, *stream>>>(pRows, pCols, dataP, vals, outputP, N, colCount, rowSize);
|
||||
sd::DebugHelper::checkErrorCode(stream, "edgeForces failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {rowP, colP, valP, data});
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// edge forces caller
|
||||
//
|
||||
void barnes_edge_forces(NDArray* rowP, NDArray * colP, NDArray * valP, int N, NDArray* output,
|
||||
NDArray& data) {
|
||||
// Loop over all edges in the graph
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), barnes_edge_forces_, (rowP, colP, valP, N, &data, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void barnes_edge_forces_,
|
||||
(NDArray* rowP, NDArray * colP, NDArray * valP, int N, NDArray * data,
|
||||
NDArray* output),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// gains - run a function T((x + 2.) * sd::math::sd_sign<T,T>(grad) != sd::math::sd_sign<T,T>(eps)) + T(x * 0.8 *
|
||||
// sd::math::sd_sign<T,T>(grad) != sd::math::sd_sign<T,T>(eps)); for all members in input and put all in output
|
||||
//
|
||||
template <typename T>
|
||||
void barnes_gains_(NDArray* input, NDArray* gradX, NDArray* epsilon, NDArray* output) {
|
||||
auto gainsInternal = LAMBDA_TTT(x, grad, eps) {
|
||||
T res = math::sd_sign<T, T>(grad) != math::sd_sign<T, T>(eps) ? x + T(.2) : x * T(.8);
|
||||
if (res < .01) res = .01;
|
||||
return res;
|
||||
});
|
||||
|
||||
input->applyTriplewiseLambda(gradX, epsilon, gainsInternal, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// gains caller
|
||||
void barnes_gains(NDArray* input, NDArray* gradX, NDArray* epsilon, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), barnes_gains_, (input, gradX, epsilon, output), SD_NUMERIC_TYPES);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void barnes_gains_, (NDArray * input, NDArray* gradX, NDArray* epsilon, NDArray* output),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// cell contains - check cells for given point
|
||||
//
|
||||
bool cell_contains(NDArray* corner, NDArray* width, NDArray* point, LongType dimension) {
|
||||
auto cornerMinusWidth = *corner - *width;
|
||||
auto cornerPlusWidth = *corner + *width;
|
||||
// executes on host side, so sync all to host memory
|
||||
cornerMinusWidth.syncToHost();
|
||||
cornerPlusWidth.syncToHost();
|
||||
for (LongType i = 0; i < dimension; i++) {
|
||||
if (cornerMinusWidth.e<double>(i) > point->e<double>(i)) return false;
|
||||
if (cornerPlusWidth.e<double>(i) < point->e<double>(i)) return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,740 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 19.04.2018
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/activations.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
void SD_KERNEL preluCuda(const void *vx, const LongType *xShapeInfo, const void *vy, const LongType *yShapeInfo,
|
||||
void *vz) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
const auto y = reinterpret_cast<const Y *>(vy);
|
||||
auto z = reinterpret_cast<X *>(vz);
|
||||
|
||||
__shared__ LongType xzLen;
|
||||
__shared__ int xzRank, yRank;
|
||||
__shared__ const LongType *xzShape;
|
||||
__shared__ const LongType *xzStride;
|
||||
__shared__ const LongType *yShape;
|
||||
__shared__ const LongType *yStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xzLen = shape::length(xShapeInfo);
|
||||
xzRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
xzShape = shape::shapeOf(xShapeInfo);
|
||||
xzStride = shape::stride(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (int i = tid; i < xzLen; i += blockDim.x * gridDim.x) {
|
||||
INDEX2COORDS(i, xzRank, xzShape, coords);
|
||||
|
||||
LongType xzOffset;
|
||||
COORDS2INDEX(xzRank, xzStride, coords, xzOffset);
|
||||
const auto xVal = x[xzOffset];
|
||||
|
||||
if (xVal < 0) {
|
||||
for (LongType j = 0; j < yRank; ++j)
|
||||
if (yShapeInfo[j + 1] == 1) coords[j + 1] = 0;
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(yRank, yStride, coords + 1, yOffset);
|
||||
z[xzOffset] = xVal * y[yOffset];
|
||||
} else {
|
||||
z[xzOffset] = xVal;
|
||||
}
|
||||
}
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
void preluCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz) {
|
||||
preluCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz);
|
||||
sd::DebugHelper::checkGlobalErrorCode("prelu failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void prelu(LaunchContext *context, NDArray *input, NDArray *alpha, NDArray *output) {
|
||||
PointersManager manager(context, "prelu");
|
||||
|
||||
dim3 launchDims = getLaunchDims("prelu");
|
||||
|
||||
const auto xType = input->dataType();
|
||||
const auto yType = alpha->dataType();
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {&input, &alpha});
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
xType, preluCudaLauncher,
|
||||
(launchDims.x, launchDims.y, launchDims.z, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), alpha->specialBuffer(), alpha->specialShapeInfo(), output->specialBuffer()),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {&input, &alpha});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
void SD_KERNEL preluBPCuda(const void *vIn, const LongType *inShapeInfo, const void *vAlpha,
|
||||
const LongType *alphaShapeInfo, const void *vdLdO, const LongType *dLdOShapeInfo,
|
||||
void *vdLdI, const LongType *dLdIShapeInfo, void *vdLdA,
|
||||
const LongType *dLdAShapeInfo) {
|
||||
const auto in = reinterpret_cast<const X *>(vIn);
|
||||
const auto alpha = reinterpret_cast<const Y *>(vAlpha);
|
||||
const auto dLdO = reinterpret_cast<const Y *>(vdLdO);
|
||||
auto dLdI = reinterpret_cast<Y *>(vdLdI);
|
||||
auto dLdA = reinterpret_cast<Y *>(vdLdA);
|
||||
|
||||
__shared__ LongType inLen, totalThreads;
|
||||
__shared__ int inRank, alphaRank;
|
||||
__shared__ const LongType *inShape;
|
||||
__shared__ const LongType *inStride;
|
||||
__shared__ const LongType *dLdOStride;
|
||||
__shared__ const LongType *dLdIStride;
|
||||
__shared__ const LongType *alphaStride;
|
||||
__shared__ const LongType *dLdAStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inLen = shape::length(inShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
inRank = shape::rank(inShapeInfo);
|
||||
alphaRank = shape::rank(alphaShapeInfo);
|
||||
|
||||
// Cache shapes and strides
|
||||
inShape = shape::shapeOf(inShapeInfo);
|
||||
inStride = shape::stride(inShapeInfo);
|
||||
dLdOStride = shape::stride(dLdOShapeInfo);
|
||||
dLdIStride = shape::stride(dLdIShapeInfo);
|
||||
alphaStride = shape::stride(alphaShapeInfo);
|
||||
dLdAStride = shape::stride(dLdAShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (int i = tid; i < inLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, inRank, inShape, coords);
|
||||
|
||||
LongType inOffset, dLdOOffset, dLdIOffset;
|
||||
COORDS2INDEX(inRank, inStride, coords, inOffset);
|
||||
COORDS2INDEX(inRank, dLdOStride, coords, dLdOOffset);
|
||||
COORDS2INDEX(inRank, dLdIStride, coords, dLdIOffset);
|
||||
|
||||
const auto xVal = in[inOffset];
|
||||
const auto grO = dLdO[dLdOOffset];
|
||||
|
||||
if (xVal < 0) {
|
||||
for (LongType j = 0; j < alphaRank; ++j)
|
||||
if (alphaShapeInfo[j + 1] == 1) coords[j + 1] = 0;
|
||||
|
||||
LongType alphaOffset, dLdAOffset;
|
||||
COORDS2INDEX(alphaRank, alphaStride, coords + 1, alphaOffset);
|
||||
COORDS2INDEX(alphaRank, dLdAStride, coords + 1, dLdAOffset);
|
||||
|
||||
dLdI[dLdIOffset] = grO * alpha[alphaOffset];
|
||||
|
||||
math::atomics::sd_atomicAdd<Y>(&dLdA[dLdAOffset], static_cast<Y>(grO * xVal));
|
||||
} else {
|
||||
dLdI[dLdIOffset] = grO;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
void SD_HOST preluBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vIn, const LongType *inShapeInfo,
|
||||
const void *vAlpha, const LongType *alphaShapeInfo, const void *vdLdO,
|
||||
const LongType *dLdOShapeInfo, void *vdLdI, const LongType *dLdIShapeInfo,
|
||||
void *vdLdA, const LongType *dLdAShapeInfo) {
|
||||
preluBPCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
|
||||
vIn, inShapeInfo, vAlpha, alphaShapeInfo, vdLdO, dLdOShapeInfo, vdLdI, dLdIShapeInfo, vdLdA, dLdAShapeInfo);
|
||||
sd::DebugHelper::checkGlobalErrorCode("prelu bp failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void preluBP(LaunchContext *context, NDArray *input, NDArray *alpha, NDArray *dLdO, NDArray *dLdI,
|
||||
NDArray *dLdA) {
|
||||
dLdA->nullify();
|
||||
|
||||
PointersManager manager(context, "preluBP");
|
||||
|
||||
dim3 launchDims = getLaunchDims("prelu");
|
||||
|
||||
const auto xType = input->dataType();
|
||||
const auto zType = alpha->dataType();
|
||||
|
||||
NDArray::prepareSpecialUse({dLdI, dLdA}, {input, alpha, dLdO});
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
xType, preluBPCudaLauncher,
|
||||
(launchDims.x, launchDims.y, launchDims.z, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), alpha->specialBuffer(), alpha->specialShapeInfo(), dLdO->specialBuffer(),
|
||||
dLdO->specialShapeInfo(), dLdI->specialBuffer(), dLdI->specialShapeInfo(), dLdA->specialBuffer(),
|
||||
dLdA->specialShapeInfo()),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&dLdI, &dLdA}, {input, alpha, dLdO});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_DEVICE void softMaxForVectorCuda(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo) {
|
||||
auto inBuff = reinterpret_cast<const T *>(vx);
|
||||
auto outBuff = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ T shmemMax;
|
||||
__shared__ T shmemSum;
|
||||
__shared__ LongType tadLen;
|
||||
__shared__ int xRank;
|
||||
__shared__ int zRank;
|
||||
__shared__ const LongType *xShape;
|
||||
__shared__ const LongType *xStride;
|
||||
__shared__ const LongType *zShape;
|
||||
__shared__ const LongType *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLen = shape::length(xShapeInfo);
|
||||
shmemMax = -DataTypeUtils::max<T>();
|
||||
shmemSum = 0.f;
|
||||
|
||||
// Cache ranks
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
// Cache shapes and strides
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
T max = -DataTypeUtils::max<T>();
|
||||
T sum = static_cast<T>(0.f);
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Calculate max using cached values
|
||||
for (LongType j = 0; j < tadLen; ++j) {
|
||||
INDEX2COORDS(j, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
max = math::sd_max<T>(max, inBuff[xOffset]);
|
||||
}
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
|
||||
// Calculate exp(x - max) and sum using cached values
|
||||
for (LongType j = 0; j < tadLen; ++j) {
|
||||
INDEX2COORDS(j, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
T temp = math::sd_exp<T, T>(inBuff[xOffset] - max);
|
||||
INDEX2COORDS(j, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
outBuff[zOffset] = temp;
|
||||
sum += temp;
|
||||
}
|
||||
|
||||
// Final division step using cached values
|
||||
for (LongType j = 0; j < tadLen; ++j) {
|
||||
INDEX2COORDS(j, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
outBuff[zOffset] /= sum;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void SD_KERNEL softMaxForVectorCudaGlobal(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType numOfSubArrs) {
|
||||
softMaxForVectorCuda<T>(vx, xShapeInfo, vz, zShapeInfo);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void softMaxForVectorCudaLauncher(const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType numTads) {
|
||||
|
||||
softMaxForVectorCudaGlobal<T><<<1, SD_CUDA_BLOCK_SIZE, 1024, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, numTads);
|
||||
sd::DebugHelper::checkGlobalErrorCode("softmax failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void softmaxEws1Kernel(const T *input, const LongType *inputOffsets, T *output,
|
||||
const LongType *outputOffsets,
|
||||
LongType numOfSubArrs, LongType tadLen) {
|
||||
int i = blockIdx.x; // Each block handles one TAD
|
||||
|
||||
if (i >= numOfSubArrs) return; // Out-of-bounds check for TADs
|
||||
|
||||
auto inBuff = input + inputOffsets[i];
|
||||
auto outBuff = output + outputOffsets[i];
|
||||
|
||||
__shared__ T shmemMax;
|
||||
__shared__ T shmemSum;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
shmemMax = -DataTypeUtils::max<T>();
|
||||
shmemSum = 0.f;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
|
||||
// Calculate max
|
||||
for (LongType j = threadIdx.x; j < tadLen; j+= gridDim.x) {
|
||||
math::atomics::sd_atomicMax(&shmemMax, inBuff[j]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate exp(x - max) and sum
|
||||
for (LongType j = threadIdx.x; j < tadLen; j += gridDim.x) {
|
||||
T temp = math::sd_exp<T, T>(inBuff[j] - shmemMax);
|
||||
outBuff[j] = temp;
|
||||
math::atomics::sd_atomicAdd(&shmemSum, temp);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Final division step
|
||||
for (LongType j = threadIdx.x; j < tadLen; j += blockDim.x) {
|
||||
outBuff[j] /= shmemSum;
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
template <typename T>
|
||||
SD_KERNEL static void softMaxCuda(const void *vx, const LongType *xTadShapeInfo, const LongType *xOffsets,
|
||||
void *vz, const LongType *zTadShapeInfo, const LongType *zOffsets, LongType numTads) {
|
||||
int i = blockIdx.x;
|
||||
if(i >= numTads) return;
|
||||
|
||||
const auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
const auto *xTad = x + xOffsets[blockIdx.x];
|
||||
auto *zTad = z + zOffsets[blockIdx.x];
|
||||
softMaxForVectorCuda<T>(xTad, xTadShapeInfo, zTad, zTadShapeInfo);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
static void softMaxEws1CudaLauncher(const int blocksPerGrid,
|
||||
const int threadsPerBlock,
|
||||
const int sharedMem,
|
||||
const cudaStream_t *stream,
|
||||
const void *vx, const LongType *xOffsets, void *vz,
|
||||
const LongType *zOffsets, LongType numTads, LongType tadLength) {
|
||||
|
||||
|
||||
|
||||
auto reCastInputs = reinterpret_cast<const T *>(vx);
|
||||
auto reCastOutputs = reinterpret_cast<T *>(vz);
|
||||
softmaxEws1Kernel<T>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(reCastInputs,
|
||||
xOffsets,
|
||||
reCastOutputs,
|
||||
zOffsets,
|
||||
numTads,
|
||||
tadLength);
|
||||
sd::DebugHelper::checkGlobalErrorCode("softmaxews failed");
|
||||
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void softMaxCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xTadShapeInfo,
|
||||
const LongType *xOffsets, void *vz, const LongType *zTadShapeInfo,
|
||||
const LongType *zOffsets, LongType numTads) {
|
||||
|
||||
|
||||
softMaxCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xTadShapeInfo, xOffsets, vz, zTadShapeInfo,
|
||||
zOffsets ,numTads);
|
||||
sd::DebugHelper::checkGlobalErrorCode("softmax failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void softmax(LaunchContext *context, NDArray *input, NDArray *output, const int dimension) {
|
||||
const int rank = input->rankOf();
|
||||
|
||||
PointersManager manager(context, "helpers::softmax");
|
||||
|
||||
if (input->isVector()) {
|
||||
if (rank == 1 || input->sizeAt(dimension) != 1) {
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), softMaxForVectorCudaLauncher,
|
||||
(context->getCudaStream(), input->specialBuffer(), input->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo(),1),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
} else
|
||||
*output = 1.;
|
||||
} else {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {dimension});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {dimension});
|
||||
|
||||
dim3 softmaxDims = getSoftmaxDims(packZ->numberOfTads());
|
||||
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), softMaxCudaLauncher,
|
||||
(softmaxDims.x, softmaxDims.y,
|
||||
softmaxDims.z,
|
||||
context->getCudaStream(),
|
||||
input->specialBuffer(),
|
||||
packX->specialShapeInfo(),
|
||||
packX->specialOffsets(), output->specialBuffer(),
|
||||
packZ->specialShapeInfo(),
|
||||
packZ->specialOffsets(),packX->numberOfTads()),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
|
||||
}
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
output->tickWriteDevice();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void SD_KERNEL logSoftMaxForVectorCuda(const void *vx, const LongType *xzShapeInfo, void *vz) {
|
||||
// logic of this kernel is based on assumption gridDim = 1
|
||||
|
||||
const auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ LongType len;
|
||||
__shared__ int numOfIters;
|
||||
__shared__ int xzRank;
|
||||
__shared__ const LongType *xzShape;
|
||||
__shared__ const LongType *xzStride;
|
||||
__shared__ T shmem[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(xzShapeInfo);
|
||||
numOfIters = (len + blockDim.x - 1) / blockDim.x; // ceil (len / blockDim.x)
|
||||
|
||||
// Cache rank, shape and stride information
|
||||
xzRank = shape::rank(xzShapeInfo);
|
||||
xzShape = shape::shapeOf(xzShapeInfo);
|
||||
xzStride = shape::stride(xzShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
T temp = -DataTypeUtils::max<T>(); // set start value to compare with at first iteration, FIXME: what if T is unsigned ??
|
||||
|
||||
// ************ evaluate max element in input array x ************ //
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx < len) {
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
shmem[threadIdx.x] = (threadIdx.x != 0) ? x[offset] : math::sd_max<T>(x[offset], temp); // take into account max element evaluated on previous iteration and stored in temp
|
||||
} else {
|
||||
shmem[threadIdx.x] = -DataTypeUtils::max<T>(); // FIXME: what if T is unsigned ??
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s /= 2) {
|
||||
if (threadIdx.x < s) shmem[threadIdx.x] = math::sd_max<T>(shmem[threadIdx.x], shmem[threadIdx.x + s]);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
temp = shmem[0]; // save max value calculated at current iteration
|
||||
}
|
||||
|
||||
const T max = temp;
|
||||
temp = 0;
|
||||
|
||||
// ************ evaluate value of exp(x[offset] - max) per each element, store it to shared memory shmem ************
|
||||
// at the same time evaluate sum of exponents, sum will be stored in shmem[0]
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx < len) {
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
z[offset] = math::sd_exp<T, T>(x[offset] - max);
|
||||
shmem[threadIdx.x] = (threadIdx.x != 0) ? z[offset] : (z[offset] + temp); // take into account sum element evaluated on previous iteration and stored in temp
|
||||
} else {
|
||||
shmem[threadIdx.x] = 0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s /= 2) {
|
||||
if (threadIdx.x < s) shmem[threadIdx.x] += shmem[threadIdx.x + s];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
temp = shmem[0]; // save sum calculated at current iteration
|
||||
}
|
||||
|
||||
// ************ evaluate log(z[offset] / sum) ************ //
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx < len) { // Added bounds check that was missing in original
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
z[offset] = math::sd_log<T, T>(z[offset] / shmem[0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void logSoftMaxForVectorCudaLauncher(const cudaStream_t *stream, const void *vx, const LongType *xzShapeInfo,
|
||||
void *vz) {
|
||||
dim3 launchDims = getLaunchDims("softmax");
|
||||
logSoftMaxForVectorCuda<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vx, xzShapeInfo, vz);
|
||||
sd::DebugHelper::checkGlobalErrorCode("logsoftmax failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void logSoftmax(LaunchContext *context, NDArray *input, NDArray *output, const int dimension) {
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
const int rank = input->rankOf();
|
||||
|
||||
if (input->isVector()) {
|
||||
if (rank == 1 || input->sizeAt(dimension) != 1) {
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input->dataType(), logSoftMaxForVectorCudaLauncher,
|
||||
(context->getCudaStream(), input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer()),
|
||||
SD_FLOAT_TYPES);
|
||||
input->tickReadDevice();
|
||||
} else
|
||||
*output = 0.;
|
||||
} else {
|
||||
std::vector<LongType> dim = {static_cast<LongType>(dimension)};
|
||||
auto maxAlongDim = const_cast<NDArray *>(input)->reduceAlongDimension(reduce::Max, &dim, true);
|
||||
auto inputMinusMax = *input - maxAlongDim;
|
||||
inputMinusMax.applyTransform(transform::Exp, output); // output contains exponents temporarily
|
||||
auto sumAlongDim = output->reduceAlongDimension(reduce::Sum, &dim, true);
|
||||
*output /= sumAlongDim;
|
||||
output->applyTransform(transform::Log, output);
|
||||
input->tickReadDevice();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "helpers::logSoftmax");
|
||||
manager.synchronize();
|
||||
|
||||
output->tickWriteDevice();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void SD_KERNEL softMaxDerivForVectorCuda(const void *vx, const LongType *xzShapeInfo, void *vz) {
|
||||
// logic of this kernel is based on assumption gridDim = 1
|
||||
|
||||
const auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ LongType len;
|
||||
__shared__ int numOfIters;
|
||||
__shared__ int xzRank;
|
||||
__shared__ const LongType *xzShape;
|
||||
__shared__ const LongType *xzStride;
|
||||
__shared__ T shmem[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(xzShapeInfo);
|
||||
numOfIters = (len + blockDim.x - 1) / blockDim.x; // ceil (len / blockDim.x)
|
||||
|
||||
// Cache rank, shape and stride information
|
||||
xzRank = shape::rank(xzShapeInfo);
|
||||
xzShape = shape::shapeOf(xzShapeInfo);
|
||||
xzStride = shape::stride(xzShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
T temp = -DataTypeUtils::max<T>(); // set start value to compare with at first iteration, FIXME: what if T is unsigned ??
|
||||
|
||||
// ************ evaluate max element in input array x ************ //
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx < len) {
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
shmem[threadIdx.x] = (threadIdx.x != 0) ? x[offset] : math::sd_max<T>(x[offset], temp); // take into account max element evaluated on previous iteration and stored in temp
|
||||
} else {
|
||||
shmem[threadIdx.x] = -DataTypeUtils::max<T>(); // FIXME: what if T is unsigned ??
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s /= 2) {
|
||||
if (threadIdx.x < s) shmem[threadIdx.x] = math::sd_max<T>(shmem[threadIdx.x], shmem[threadIdx.x + s]);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
temp = shmem[0]; // save max value calculated at current iteration
|
||||
}
|
||||
|
||||
const T max = temp;
|
||||
temp = 0;
|
||||
|
||||
// ************ evaluate value of exp(x[offset] - max) per each element, store it to shared memory shmem ************
|
||||
// at the same evaluate sum of exponents, sum will be stored in shmem[0]
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx < len) {
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
z[offset] = math::sd_exp<T, T>(x[offset] - max);
|
||||
shmem[threadIdx.x] = (threadIdx.x != 0) ? z[offset] : (z[offset] + temp); // take into account sum element evaluated on previous iteration and stored in temp
|
||||
} else {
|
||||
shmem[threadIdx.x] = 0;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (int s = blockDim.x / 2; s > 0; s /= 2) {
|
||||
if (threadIdx.x < s) shmem[threadIdx.x] += shmem[threadIdx.x + s];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
temp = shmem[0]; // save sum calculated at current iteration
|
||||
}
|
||||
|
||||
// ************ evaluate (z[offset] / sum) and derivative z[offset] = z[offset] * (1 - z[offset]) ************ //
|
||||
for (int i = 0; i < numOfIters; ++i) {
|
||||
const LongType elemIdx = i * blockDim.x + threadIdx.x;
|
||||
if (elemIdx >= len) continue;
|
||||
|
||||
LongType offset;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(elemIdx, xzRank, xzShape, coords);
|
||||
COORDS2INDEX(xzRank, xzStride, coords, offset);
|
||||
z[offset] /= shmem[0];
|
||||
z[offset] *= (1.f - z[offset]); // derivative
|
||||
}
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void softMaxDerivForVectorCudaLauncher(const cudaStream_t *stream, const void *vx, const LongType *xzShapeInfo,
|
||||
void *vz) {
|
||||
dim3 launchDims = getLaunchDims("softmax");
|
||||
|
||||
softMaxDerivForVectorCuda<T><<<launchDims.x,launchDims.y, launchDims.z, *stream>>>(vx, xzShapeInfo, vz);
|
||||
sd::DebugHelper::checkGlobalErrorCode("softmax derivative failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void softmaxDerivative(LaunchContext *context, NDArray *input, NDArray *output, const int dimension) {
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
const int rank = input->rankOf();
|
||||
LongType temp;
|
||||
|
||||
if (shape::isCommonVector(input->shapeInfo(), temp)) {
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input->dataType(), softMaxDerivForVectorCudaLauncher,
|
||||
(context->getCudaStream(), input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer()),
|
||||
SD_FLOAT_TYPES);
|
||||
input->tickReadDevice();
|
||||
} else {
|
||||
std::vector<LongType> dim = {static_cast<LongType>(dimension)};
|
||||
auto maxAlongDim = const_cast<NDArray *>(input)->reduceAlongDimension(reduce::Max, &dim, true);
|
||||
auto inputMinusMax = *input - maxAlongDim;
|
||||
inputMinusMax.applyTransform(transform::Exp, output); // output contains exponents temporarily
|
||||
auto sumAlongDim = output->reduceAlongDimension(reduce::Sum, &dim, true);
|
||||
*output /= sumAlongDim;
|
||||
*output *= (1.f - *output); // derivative
|
||||
input->tickReadDevice();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "helpers::softmaxDerivative");
|
||||
manager.synchronize();
|
||||
|
||||
output->tickWriteDevice();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void thresholdRelu_(NDArray *input, double threshold, NDArray *output) {
|
||||
auto routine = LAMBDA_T(_x, threshold) { return _x > (T)threshold ? _x : (T)0.f; });
|
||||
input->applyLambda(routine, output);
|
||||
}
|
||||
|
||||
void thresholdRelu(LaunchContext *context, NDArray *input, double threshold, NDArray *output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), thresholdRelu_, (input, threshold, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void thresholdReluDerivative_(NDArray *input, double theta, NDArray *dLdO, NDArray *output) {
|
||||
auto derivative = LAMBDA_TT(_x, grO, theta) {
|
||||
if (_x > theta)
|
||||
return grO;
|
||||
else
|
||||
return static_cast<T>(0);
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(dLdO, derivative, output);
|
||||
}
|
||||
|
||||
void thresholdReluDerivative(LaunchContext *context, NDArray *input, double threshold, NDArray *dLdO,
|
||||
NDArray *output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), thresholdReluDerivative_, (input, threshold, dLdO, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,157 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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)
|
||||
//
|
||||
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/addBias.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void addBiasCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const bool isNCHW) {
|
||||
// bias [oC]
|
||||
|
||||
// if(input_rank == 4)
|
||||
// input and output have same shapes: [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
|
||||
// if(input_rank == 5)
|
||||
// input and output have same shapes: [bS, oD, oH, oW, oC] (NHWC) or [bS, oD, oC, oH, oW] (NCHW)
|
||||
|
||||
const X* x = reinterpret_cast<const X*>(vx);
|
||||
const Y* y = reinterpret_cast<const Y*>(vy);
|
||||
X* z = reinterpret_cast<X*>(vz);
|
||||
|
||||
__shared__ LongType rank, channelPosition, posOfNonUnityDim;
|
||||
__shared__ LongType len, *sharedMem;
|
||||
__shared__ bool xzSameOffsets, xzAreSame;
|
||||
__shared__ const LongType *xShape;
|
||||
__shared__ const LongType *xStride;
|
||||
__shared__ const LongType *zStride;
|
||||
__shared__ const LongType *yStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
rank = shape::rank(xShapeInfo); // xRank == zRank
|
||||
xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
len = shape::length(xShapeInfo);
|
||||
channelPosition = isNCHW ? 1 : rank - 1; // second or last
|
||||
xzAreSame = x == z;
|
||||
|
||||
// Cache shapes and strides
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
|
||||
shape::isCommonVector(yShapeInfo, posOfNonUnityDim);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += blockDim.x * gridDim.x) {
|
||||
INDEX2COORDS(i, rank, xShape, coords);
|
||||
|
||||
LongType xOffsets;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffsets);
|
||||
LongType zOffsets;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffsets);
|
||||
LongType yOffsets = coords[channelPosition] * yStride[posOfNonUnityDim];
|
||||
|
||||
if (xzAreSame)
|
||||
z[zOffsets] += static_cast<X>(y[yOffsets]);
|
||||
else
|
||||
z[zOffsets] = static_cast<X>(x[xOffsets]) + static_cast<X>(y[yOffsets]);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void addBiasCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const bool isNCHW) {
|
||||
addBiasCuda<X, Y>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, isNCHW);
|
||||
sd::DebugHelper::checkGlobalErrorCode("addbias failed");
|
||||
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void addBias2DCuda(const void* vx, const void* vy, void* vz, uint32_t blocks, uint32_t length) {
|
||||
auto y = reinterpret_cast<const Y*>(vy);
|
||||
|
||||
for (uint32_t b = blockIdx.x; b < blocks; b += gridDim.x) {
|
||||
auto x = reinterpret_cast<const X*>(vx) + length * b;
|
||||
auto z = reinterpret_cast<X*>(vz) + length * b;
|
||||
|
||||
for (uint32_t e = threadIdx.x; e < length; e += blockDim.x) {
|
||||
z[e] = x[e] + y[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
static void addBias2DCudaLauncher(const cudaStream_t* stream, const void* vx, const void* vy, void* vz, uint32_t blocks,
|
||||
uint32_t length) {
|
||||
dim3 dims = getAddBiasDims(2, 2);
|
||||
|
||||
addBias2DCuda<X, Y><<<dims.x, dims.y, dims.z, *stream>>>(vx, vy, vz, blocks, length);
|
||||
sd::DebugHelper::checkGlobalErrorCode("addbias 2d failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void addBias(graph::Context& block, NDArray& input, NDArray& bias, NDArray& output, const bool isNCHW) {
|
||||
PointersManager manager(block.launchContext(), "addBias");
|
||||
NDArray::prepareSpecialUse({&output}, {&input, &bias});
|
||||
|
||||
if (input.rankOf() == 2 && bias.rankOf() == 1 && input.ordering() == 'c' && output.ordering() == 'c' &&
|
||||
input.sizeAt(1) == bias.sizeAt(0)) {
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), bias.dataType(), addBias2DCudaLauncher,
|
||||
(block.launchContext()->getCudaStream(), input.specialBuffer(), bias.specialBuffer(),
|
||||
output.specialBuffer(), input.sizeAt(0), bias.sizeAt(0)),
|
||||
SD_FLOAT_TYPES, SD_FLOAT_TYPES);
|
||||
} else {
|
||||
// default case
|
||||
dim3 dims = getAddBiasDims(input.rankOf(), input.rankOf());
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), bias.dataType(), addBiasCudaLauncher,
|
||||
(dims.x, dims.y, dims.z, block.launchContext()->getCudaStream(),
|
||||
input.specialBuffer(), input.specialShapeInfo(), bias.specialBuffer(),
|
||||
bias.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), isNCHW),
|
||||
SD_FLOAT_TYPES, SD_FLOAT_TYPES);
|
||||
}
|
||||
NDArray::registerSpecialUse({&output}, {&input, &bias});
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,113 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/adjust_hue.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void SD_KERNEL adjustHueCuda(const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets,
|
||||
void* vz, const LongType* zShapeInfo, const LongType* zTadOffsets,
|
||||
const LongType numOfTads, const T delta, const LongType dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
T h, s, v;
|
||||
|
||||
rgbToHsv<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], h, s, v);
|
||||
|
||||
h += delta;
|
||||
if (h > 1)
|
||||
h -= 1;
|
||||
else if (h < 0)
|
||||
h += 1;
|
||||
|
||||
hsvToRgb<T>(h, s, v, zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_HOST void adjustHueCudaLauncher(const int blocksPerGrid, const int threadsPerBlock,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xTadOffsets, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zTadOffsets, const LongType numOfTads,
|
||||
NDArray* deltaScalarArr, const LongType dimC) {
|
||||
adjustHueCuda<T><<<blocksPerGrid, threadsPerBlock, 512, *stream>>>(
|
||||
vx, xShapeInfo, xTadOffsets, vz, zShapeInfo, zTadOffsets, numOfTads, deltaScalarArr->e<T>(0), dimC);
|
||||
sd::DebugHelper::checkGlobalErrorCode("sadjustHue failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
void adjustHue(LaunchContext* context, NDArray* input, NDArray* deltaScalarArr, NDArray* output,
|
||||
const LongType dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (numOfTads + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
dim3 adjustDims = getAdjustDims(numOfTads);
|
||||
PointersManager manager(context, "adjustHue");
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, deltaScalarArr});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), adjustHueCudaLauncher,
|
||||
(adjustDims.x, adjustDims.y, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), packX->platformOffsets(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), packZ->platformOffsets(), numOfTads, deltaScalarArr, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, deltaScalarArr});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,116 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/adjust_hue.h>
|
||||
#include <ops/declarable/helpers/adjust_saturation.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void SD_KERNEL adjustSaturationCuda(const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xTadOffsets, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zTadOffsets, const LongType numOfTads,
|
||||
const T factor, const LongType dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
T h, s, v;
|
||||
|
||||
rgbToHsv<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], h, s, v);
|
||||
|
||||
s *= factor;
|
||||
if (s > 1.f)
|
||||
s = 1.f;
|
||||
else if (s < 0.f)
|
||||
s = 0.f;
|
||||
|
||||
hsvToRgb<T>(h, s, v, zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_HOST void adjustSaturationCudaLauncher(const int blocksPerGrid, const int threadsPerBlock,
|
||||
const cudaStream_t* stream, const void* vx,
|
||||
const LongType* xShapeInfo, const LongType* xTadOffsets,
|
||||
void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zTadOffsets, const LongType numOfTads,
|
||||
NDArray* factorScalarArr, const LongType dimC) {
|
||||
adjustSaturationCuda<T><<<blocksPerGrid, threadsPerBlock, 256, *stream>>>(
|
||||
vx, xShapeInfo, xTadOffsets, vz, zShapeInfo, zTadOffsets, numOfTads, factorScalarArr->e<T>(0), dimC);
|
||||
sd::DebugHelper::checkGlobalErrorCode("adjustSaturation failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
void adjustSaturation(LaunchContext* context, NDArray* input, NDArray* factorScalarArr, NDArray* output,
|
||||
const LongType dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
dim3 adjustDims = getAdjustDims(numOfTads);
|
||||
|
||||
PointersManager manager(context, "adjustSaturation");
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, factorScalarArr});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), adjustSaturationCudaLauncher,
|
||||
(adjustDims.x,adjustDims.y, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), packX->platformOffsets(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), packZ->platformOffsets(), numOfTads, factorScalarArr, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, factorScalarArr});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,112 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/assign.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
#include "helpers/ShapeUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL static void assignKernel(const void* vx, const LongType* xShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType xOffset, const LongType zOffset) {
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ LongType len, totalThreads;
|
||||
__shared__ int rank;
|
||||
__shared__ const LongType *xShape;
|
||||
__shared__ const LongType *zShape;
|
||||
__shared__ const LongType *xStride;
|
||||
__shared__ const LongType *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(zShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
rank = shape::rank(zShapeInfo);
|
||||
|
||||
// Cache shapes and strides
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
LongType xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = tid; i < len; i += totalThreads) {
|
||||
INDEX2COORDS(i, rank, zShape, zCoords);
|
||||
INDEX2COORDS(i, rank, xShape, xCoords);
|
||||
|
||||
LongType xIndex, zIndex;
|
||||
COORDS2INDEX(rank, xStride, xCoords, xIndex);
|
||||
COORDS2INDEX(rank, zStride, zCoords, zIndex);
|
||||
|
||||
z[zIndex] = static_cast<Z>(x[xIndex]);
|
||||
}
|
||||
}
|
||||
template <typename X, typename Z>
|
||||
SD_HOST static void assignCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType xOffset, const LongType zOffset) {
|
||||
assignKernel<X, Z><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, xOffset, zOffset);
|
||||
DebugHelper::checkGlobalErrorCode("assignKernel(...) failed");
|
||||
}
|
||||
|
||||
void assign(sd::LaunchContext* context, sd::NDArray* target, sd::NDArray* source) {
|
||||
if (target->lengthOf() != source->lengthOf()) {
|
||||
std::string errorMsg = "assign helper: Source and target arrays must have the same length. ";
|
||||
errorMsg += "Source shape: " + ShapeUtils::shapeAsString(source) + ", ";
|
||||
errorMsg += "Target shape: " + ShapeUtils::shapeAsString(target) + ", ";
|
||||
errorMsg += "Source datatype: " + DataTypeUtils::asString(source->dataType()) + ", ";
|
||||
errorMsg += "Target datatype: " + DataTypeUtils::asString(target->dataType());
|
||||
THROW_EXCEPTION(errorMsg.c_str());
|
||||
}
|
||||
|
||||
NDArray::prepareSpecialUse({target}, {source});
|
||||
|
||||
auto xType = source->dataType();
|
||||
auto zType = target->dataType();
|
||||
|
||||
dim3 launchDims = traceDims(target->lengthOf());
|
||||
|
||||
PointersManager manager(context, "helpers::assign");
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(xType, zType, assignCudaLauncher,
|
||||
(launchDims.x, launchDims.y, launchDims.z, context->getCudaStream(),
|
||||
source->specialBuffer(), source->specialShapeInfo(),
|
||||
target->specialBuffer(), target->specialShapeInfo(),
|
||||
source->offset(), target->offset()),
|
||||
SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
|
||||
manager.synchronize();
|
||||
NDArray::registerSpecialUse({target}, {source});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,60 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <ops/declarable/helpers/axis.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
void adjustAxis(LongType rank, NDArray* axisVector, std::vector<LongType>& output) {
|
||||
if(axisVector->isScalar()) {
|
||||
output.resize(1);
|
||||
auto ca = axisVector->e<LongType>(0);
|
||||
if (ca < 0) // shift values on rank for negative vals
|
||||
ca += rank;
|
||||
output[0] = ca;
|
||||
return;
|
||||
}
|
||||
output.resize(axisVector->lengthOf());
|
||||
axisVector->tickReadDevice(); // mark input as read on device
|
||||
axisVector->syncToHost(); // sync to host
|
||||
for (int e = 0; e < axisVector->lengthOf(); e++) {
|
||||
auto ca = axisVector->e<LongType>(e);
|
||||
if (ca < 0) // shift values on rank for negative vals
|
||||
ca += rank;
|
||||
|
||||
output[e] = ca;
|
||||
}
|
||||
}
|
||||
|
||||
void adjustAxis(LongType rank, std::vector<LongType>& axisVector) {
|
||||
for (int e = 0; e < axisVector.size(); e++) {
|
||||
auto a = axisVector[e];
|
||||
if (a < 0) // shift vals on rank for negative vals
|
||||
axisVector[e] = a + rank;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,208 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <cublas_v2.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/batched_gemm.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/float16.h>
|
||||
|
||||
#include <indexing/NDIndexUtils.h>
|
||||
#include <ops/declarable/CustomOperations.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
|
||||
void bgemm(NDArray *a, NDArray *b, NDArray *c, NDArray *alphas, NDArray *betas,
|
||||
int transA, int transB, int M, int N, int K, int lda, int ldb, int ldc, NDArray *all) {
|
||||
NDArray *allIndex = nullptr;
|
||||
if(all != nullptr)
|
||||
allIndex = all;
|
||||
else {
|
||||
NDArray allLocal = NDIndexUtils::createAll();
|
||||
all = &allLocal;
|
||||
}
|
||||
|
||||
|
||||
int batchSize = a->sizeAt(0) / 2;
|
||||
std::vector<NDArray *>inputs;
|
||||
std::vector<NDArray *> keyInputs;
|
||||
std::vector<NDArray *> outputs;
|
||||
|
||||
create_view createView;
|
||||
|
||||
//add alpha and beta before the batch gemm, this just needs to be broadcasted
|
||||
inputs.push_back(alphas);
|
||||
inputs.push_back(betas);
|
||||
|
||||
//divide by 2: queries and keys
|
||||
for(int i = 0; i < batchSize; i++) {
|
||||
auto point = NDIndexUtils::createPoint(i);
|
||||
auto aSlice = createView.evaluate({a,&point,all,all},{},{});
|
||||
auto bSlice = createView.evaluate({b,&point,all,all},{},{});
|
||||
auto outSlice = createView.evaluate({c,&point,all,all},{},{});
|
||||
inputs.push_back(aSlice.at(0));
|
||||
keyInputs.push_back(bSlice.at(0));
|
||||
outputs.push_back(outSlice.at(0));
|
||||
}
|
||||
|
||||
bgemm(inputs,keyInputs,outputs,alphas,betas,transA,transB,M,N,K,lda,ldb,ldc);
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// bsxMXK x bSxKxN = bSxMxN
|
||||
void bgemm( std::vector<NDArray *> &vA, std::vector<NDArray *> &vB, std::vector<NDArray *> &vC,
|
||||
NDArray *alphas, NDArray *betas, int transA, int transB, int M, int N, int K, int lda,
|
||||
int ldb, int ldc) {
|
||||
const auto bS = vA.size(); // batch size
|
||||
|
||||
std::vector<NDArray*> pA(bS), pB(bS), pC(bS);
|
||||
|
||||
std::vector<NDArray*> toDelete;
|
||||
|
||||
for (int i = 0; i < bS; ++i) {
|
||||
pA[i] = vA[i]->dup('f'); // dup() already returns NDArray*
|
||||
toDelete.emplace_back(pA[i]);
|
||||
|
||||
|
||||
pB[i] = vB[i]->dup('f'); // dup() already returns NDArray*
|
||||
toDelete.emplace_back(pB[i]);
|
||||
|
||||
pC[i] = vC[i]->dup('f'); // dup() already returns NDArray*
|
||||
toDelete.emplace_back(pC[i]);
|
||||
|
||||
|
||||
if (pC[i]->ordering() != 'f') {
|
||||
auto temp = pA[i];
|
||||
std::vector<sd::LongType> permute = {1,0};
|
||||
pA[i] = pB[i]->permute(permute, false, false); // permute() already returns NDArray*
|
||||
pB[i] = temp->permute(permute, false, false);
|
||||
pC[i] = pC[i]->permute(permute, false, false);
|
||||
toDelete.push_back(pA[i]);
|
||||
toDelete.push_back(pB[i]);
|
||||
toDelete.push_back(pC[i]);
|
||||
M = pA[i]->sizeAt(0);
|
||||
K = pA[i]->sizeAt(1);
|
||||
N = pB[i]->sizeAt(1);
|
||||
}
|
||||
|
||||
NDArray::prepareSpecialUse({pC[i]}, {pA[i], pB[i]});
|
||||
NDArray::registerSpecialUse({pC[i]}, {pA[i], pB[i]});
|
||||
}
|
||||
|
||||
NDArray::prepareSpecialUse({}, {alphas, betas});
|
||||
NDArray::registerSpecialUse({}, {alphas, betas});
|
||||
|
||||
std::vector<void*> pAbuffs(bS), pBbuffs(bS), pCbuffs(bS);
|
||||
for (int i = 0; i < bS; ++i) {
|
||||
pAbuffs[i] = pA[i]->specialBuffer();
|
||||
pBbuffs[i] = pB[i]->specialBuffer();
|
||||
pCbuffs[i] = pC[i]->specialBuffer();
|
||||
}
|
||||
|
||||
LaunchContext * context = vA[0]->getContext();
|
||||
PointersManager manager(context, "helpers::bgemm cuda");
|
||||
|
||||
const void** aBuffers = reinterpret_cast<const void**>(manager.replicatePointer(pAbuffs.data(), bS * sizeof(void*)));
|
||||
const void** bBuffers = reinterpret_cast<const void**>(manager.replicatePointer(pBbuffs.data(), bS * sizeof(void*)));
|
||||
void** cBuffers = reinterpret_cast<void**>(manager.replicatePointer(pCbuffs.data(), bS * sizeof(void*)));
|
||||
|
||||
|
||||
|
||||
const cublasOperation_t transAblas = transA == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
||||
const cublasOperation_t transBblas = transB == 112 ? CUBLAS_OP_T : CUBLAS_OP_N;
|
||||
|
||||
if(M < 0) THROW_EXCEPTION("M < 0");
|
||||
if(N < 0) THROW_EXCEPTION("N < 0");
|
||||
if(K < 0) THROW_EXCEPTION("K < 0");
|
||||
|
||||
|
||||
const auto aType = pA[0]->dataType();
|
||||
const auto bType = pB[0]->dataType();
|
||||
const auto cType = pC[0]->dataType();
|
||||
|
||||
std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());
|
||||
|
||||
auto handle = reinterpret_cast<cublasHandle_t*>(context->getCublasHandle());
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
auto status = cublasSetStream_v2(*handle, *stream);
|
||||
|
||||
if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("MmulHelper::mmulMxM cuda set stream failed ! Please double check the passed in handle.", status);
|
||||
|
||||
const bool AB(aType == bType), AC(aType == cType), ABC(AB && AC);
|
||||
|
||||
// choose appropriate cuda gemm api depending on data types
|
||||
// Handle empty arrays by using default values (alpha=1.0, beta=0.0)
|
||||
if (ABC && aType == DOUBLE) {
|
||||
double alpha = alphas->lengthOf() > 0 ? alphas->e<double>(0) : 1.0;
|
||||
double beta = betas->lengthOf() > 0 ? betas->e<double>(0) : 0.0;
|
||||
status = cublasDgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const double**)aBuffers, lda,
|
||||
(const double**)bBuffers, ldb, &beta, (double**)cBuffers, ldc, bS);
|
||||
} else if (ABC && aType == FLOAT32) {
|
||||
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
||||
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
||||
status = cublasSgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const float**)aBuffers, lda,
|
||||
(const float**)bBuffers, ldb, &beta, (float**)cBuffers, ldc, bS);
|
||||
} else if (ABC && aType == HALF) {
|
||||
__half alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
||||
__half beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
||||
status = cublasHgemmBatched(*handle, transAblas, transBblas, M, N, K, &alpha, (const __half**)aBuffers, lda,
|
||||
(const __half**)bBuffers, ldb, &beta, (__half**)cBuffers, ldc, bS);
|
||||
} else if (AB && aType == INT8 && cType == FLOAT32) {
|
||||
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
||||
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
||||
status = cublasGemmBatchedEx(*handle, transAblas, transBblas, M, N, K, &alpha, aBuffers, CUDA_R_8I, lda, bBuffers,
|
||||
CUDA_R_8I, ldb, &beta, cBuffers, CUDA_R_32F, ldc, bS, CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
|
||||
} else if (AB && aType == HALF && cType == FLOAT32) {
|
||||
float alpha = alphas->lengthOf() > 0 ? alphas->e<float>(0) : 1.0f;
|
||||
float beta = betas->lengthOf() > 0 ? betas->e<float>(0) : 0.0f;
|
||||
status =
|
||||
cublasGemmBatchedEx(*handle, transAblas, transBblas, M, N, K, &alpha, aBuffers, CUDA_R_16F, lda, bBuffers,
|
||||
CUDA_R_16F, ldb, &beta, cBuffers, CUDA_R_32F, ldc, bS, CUDA_R_32F, CUBLAS_GEMM_DEFAULT);
|
||||
} else
|
||||
THROW_EXCEPTION("batched gemm cuda: this mode is not implemented yet !");
|
||||
|
||||
if (status != CUBLAS_STATUS_SUCCESS) {
|
||||
throw cuda_exception::build("MmulHelper::mmulMxM cuda execution failed !", status);
|
||||
}
|
||||
|
||||
auto cudaResult = cudaStreamSynchronize(*stream);
|
||||
if (cudaResult != 0) {
|
||||
throw cuda_exception::build("MmulHelper::mmulMxM cuda stream synchronize failed !", cudaResult);
|
||||
}
|
||||
|
||||
for (int i = 0; i < bS; ++i)
|
||||
vC[i]->assign(pC[i]);
|
||||
|
||||
for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,151 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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, created on 25.02.2018
|
||||
//
|
||||
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/OmpLaunchHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/batchnorm.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void batchnormCuda2(const void* vx, const LongType* xShapeInfo, const void* vMean,
|
||||
const LongType* meanShapeInfo, const void* vVariance,
|
||||
const LongType* varianceShapeInfo, const void* vGamma,
|
||||
const LongType* gammaShapeInfo, const void* vBeta,
|
||||
const LongType* betaShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const int numDims, const LongType* dims, const T epsilon) {
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
const auto mean = reinterpret_cast<const T*>(vMean);
|
||||
const auto variance = reinterpret_cast<const T*>(vVariance);
|
||||
const auto gamma = reinterpret_cast<const T*>(vGamma);
|
||||
const auto beta = reinterpret_cast<const T*>(vBeta);
|
||||
|
||||
__shared__ int xRank, minRank; // xRank == zRank, minRank = meanRank = varianceRank = gammaRank = betaRank
|
||||
__shared__ LongType xLen, totalThreads; // xLen = zLen
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
xLen = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
minRank = shape::rank(meanShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < xLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, xRank, shape::shapeOf(xShapeInfo), coords);
|
||||
|
||||
LongType xOffset, zOffset;
|
||||
COORDS2INDEX(xRank, shape::stride(xShapeInfo), coords, xOffset);
|
||||
COORDS2INDEX(xRank, shape::stride(zShapeInfo), coords, zOffset);
|
||||
|
||||
if (minRank == xRank) {
|
||||
for (LongType i = 0, j = 0; i < xRank; ++i) {
|
||||
if (j < numDims && i != dims[j])
|
||||
coords[i] = 0;
|
||||
else
|
||||
++j;
|
||||
}
|
||||
} else // minRank = numDims = 1 in this case
|
||||
coords[0] = coords[dims[0]];
|
||||
|
||||
LongType meanOffset, varianceOffset;
|
||||
COORDS2INDEX(minRank, shape::stride(meanShapeInfo), coords, meanOffset);
|
||||
COORDS2INDEX(minRank, shape::stride(varianceShapeInfo), coords, varianceOffset);
|
||||
|
||||
T sigmaInvGam = 1. / math::sd_sqrt<T, T>(variance[varianceOffset] + epsilon);
|
||||
|
||||
if (gamma != nullptr) {
|
||||
LongType gammaOffset;
|
||||
COORDS2INDEX(minRank, shape::stride(gammaShapeInfo), coords, gammaOffset);
|
||||
sigmaInvGam *= gamma[gammaOffset];
|
||||
}
|
||||
|
||||
z[zOffset] = (x[xOffset] - mean[meanOffset]) * sigmaInvGam;
|
||||
|
||||
if (beta != nullptr) {
|
||||
LongType betaOffset;
|
||||
COORDS2INDEX(minRank, shape::stride(betaShapeInfo), coords, betaOffset);
|
||||
z[zOffset] += beta[betaOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST static void batchnormCudaLauncher2(const int blocksPerGrid, const int threadsPerBlock,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vMean, const LongType* meanShapeInfo, const void* vVariance,
|
||||
const LongType* varianceShapeInfo, const void* vGamma,
|
||||
const LongType* gammaShapeInfo, const void* vBeta,
|
||||
const LongType* betaShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const int numDims, const LongType* dims, const double epsilon) {
|
||||
batchnormCuda2<T><<<blocksPerGrid, threadsPerBlock, 512, *stream>>>(
|
||||
vx, xShapeInfo, vMean, meanShapeInfo, vVariance, varianceShapeInfo, vGamma, gammaShapeInfo, vBeta, betaShapeInfo,
|
||||
vz, zShapeInfo, numDims, dims, static_cast<T>(epsilon));
|
||||
sd::DebugHelper::checkGlobalErrorCode("batchNormCuda2 failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void batchnorm(NDArray* input, NDArray* mean, NDArray* variance, NDArray* gamma,
|
||||
NDArray* beta, NDArray* output, const std::vector<LongType>& axes, const double epsilon) {
|
||||
|
||||
dim3 batchNormDims = getBatchNormDims(input->lengthOf());
|
||||
PointersManager manager(input->getContext(), "batchnorm");
|
||||
|
||||
const LongType* dims = reinterpret_cast<LongType*>(manager.replicatePointer(axes.data(), axes.size() * sizeof(LongType)));
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), batchnormCudaLauncher2,
|
||||
(batchNormDims.x, batchNormDims.y, input->getContext()->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), mean->specialBuffer(), mean->specialShapeInfo(),
|
||||
variance->specialBuffer(), variance->specialShapeInfo(),
|
||||
gamma ? gamma->specialBuffer() : nullptr, gamma ? gamma->specialShapeInfo() : nullptr,
|
||||
beta ? beta->specialBuffer() : nullptr, beta ? beta->specialShapeInfo() : nullptr,
|
||||
output->specialBuffer(), output->specialShapeInfo(), axes.size(), dims, epsilon),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,233 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <array/DataTypeUtils.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/betaInc.h>
|
||||
|
||||
#include <cmath>
|
||||
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// modified Lentz’s algorithm for continued fractions,
|
||||
// reference: Lentz, W.J. 1976, “Generating Bessel Functions in Mie Scattering Calculations Using Continued Fractions,”
|
||||
template <typename T>
|
||||
SD_DEVICE T continuedFractionCuda(const T a, const T b, const T x) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
T* coeffs = reinterpret_cast<T*>(shmem);
|
||||
|
||||
const T min = DataTypeUtils::min<T>() / DataTypeUtils::eps<T>();
|
||||
const T aPlusb = a + b;
|
||||
T val, aPlus2i;
|
||||
|
||||
T t2 = coeffs[1];
|
||||
T t1 = coeffs[0];
|
||||
if (math::sd_abs<T,T>(t1) < min) t1 = min;
|
||||
t1 = static_cast<T>(1) / t1;
|
||||
T result = t1;
|
||||
|
||||
for (LongType i = 1; i <= maxIter; ++i) {
|
||||
const LongType i2 = 2 * i;
|
||||
aPlus2i = a + static_cast<T>(i2);
|
||||
|
||||
// t1
|
||||
t1 = static_cast<T>(1) + coeffs[i2] * t1;
|
||||
if (math::sd_abs<T,T>(t1) < min) t1 = min;
|
||||
t1 = static_cast<T>(1) / t1;
|
||||
// t2
|
||||
t2 = static_cast<T>(1) + coeffs[i2] / t2;
|
||||
if (math::sd_abs<T,T>(t2) < min) t2 = min;
|
||||
// result
|
||||
result *= t2 * t1;
|
||||
// t1
|
||||
t1 = static_cast<T>(1) + coeffs[i2 + 1] * t1;
|
||||
if (math::sd_abs<T,T>(t1) < min) t1 = min;
|
||||
t1 = static_cast<T>(1) / t1;
|
||||
// t2
|
||||
t2 = static_cast<T>(1) + coeffs[i2 + 1] / t2;
|
||||
if (math::sd_abs<T,T>(t2) < min) t2 = min;
|
||||
// result
|
||||
val = t2 * t1;
|
||||
result *= val;
|
||||
|
||||
// condition to stop loop
|
||||
if (math::sd_abs<T,T>(val - static_cast<T>(1)) <= DataTypeUtils::eps<T>()) return result;
|
||||
}
|
||||
|
||||
return DataTypeUtils::infOrMax<T>(); // no convergence, more iterations is required, return infinity
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void betaIncForArrayCuda(const void* va, const LongType* aShapeInfo, const void* vb,
|
||||
const LongType* bShapeInfo, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
T* sharedMem = reinterpret_cast<T*>(shmem);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType aLen, bLen, xLen, zLen, aOffset, bOffset, xOffset, zOffset;
|
||||
__shared__ int aRank, bRank, xRank, zRank;
|
||||
__shared__ const LongType *aShape, *bShape, *xShape, *zShape;
|
||||
__shared__ const LongType *aStride, *bStride, *xStride, *zStride;
|
||||
__shared__ T a, b, x;
|
||||
__shared__ bool symmCond;
|
||||
|
||||
const LongType j = blockIdx.x; // one block per each element
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache lengths
|
||||
aLen = shape::length(aShapeInfo);
|
||||
bLen = shape::length(bShapeInfo);
|
||||
xLen = shape::length(xShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
// Cache ranks
|
||||
aRank = shape::rank(aShapeInfo);
|
||||
bRank = shape::rank(bShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
// Cache shapes
|
||||
aShape = shape::shapeOf(aShapeInfo);
|
||||
bShape = shape::shapeOf(bShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
|
||||
// Cache strides
|
||||
aStride = shape::stride(aShapeInfo);
|
||||
bStride = shape::stride(bShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
LongType aCoords[SD_MAX_RANK];
|
||||
LongType bCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
|
||||
INDEX2COORDS(j, aRank, aShape, aCoords);
|
||||
COORDS2INDEX(aRank, aStride, aCoords, aOffset);
|
||||
INDEX2COORDS(j, bRank, bShape, bCoords);
|
||||
COORDS2INDEX(bRank, bStride, bCoords, bOffset);
|
||||
INDEX2COORDS(j, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
INDEX2COORDS(j, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
|
||||
if (aOffset >= aLen || bOffset >= bLen || xOffset >= xLen || zOffset >= zLen)
|
||||
return;
|
||||
|
||||
a = *(reinterpret_cast<const T*>(va) + aOffset);
|
||||
b = *(reinterpret_cast<const T*>(vb) + bOffset);
|
||||
x = *(reinterpret_cast<const T*>(vx) + xOffset);
|
||||
symmCond = x > (a + T(1)) / (a + b + T(2));
|
||||
|
||||
if (symmCond) { // swap a and b, x = 1 - x
|
||||
T temp = a;
|
||||
a = b;
|
||||
b = temp;
|
||||
x = T(1) - x;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// t^{n-1} * (1 - t)^{n-1} is symmetric function with respect to x = 0.5
|
||||
if (zOffset < zLen && a == b && x == T(0.5)) {
|
||||
z[zOffset] = T(0.5);
|
||||
return;
|
||||
}
|
||||
|
||||
if (zOffset < zLen && (x == T(0) || x == T(1))) {
|
||||
if (symmCond) {
|
||||
z[zOffset] = T(1) - x;
|
||||
} else {
|
||||
z[zOffset] = x;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// calculate two coefficients per thread
|
||||
if (threadIdx.x != 0) {
|
||||
const int i = threadIdx.x;
|
||||
const T aPlus2i = a + T(2) * T(i);
|
||||
sharedMem[2 * i] = T(i) * (b - T(i)) * x / ((aPlus2i - T(1)) * aPlus2i);
|
||||
sharedMem[2 * i + 1] = -(a + T(i)) * (a + b + T(i)) * x / ((aPlus2i + T(1)) * aPlus2i);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
const T gammaPart = static_cast<T>(lgamma(a)) + static_cast<T>(lgamma(b)) - static_cast<T>(lgamma(a + b));
|
||||
const T front = math::sd_exp<T, T>(math::sd_log<T, T>(x) * a + math::sd_log<T, T>(T(1) - x) * b - gammaPart);
|
||||
|
||||
sharedMem[0] = T(1) - (a + b) * x / (a + T(1));
|
||||
sharedMem[1] = T(1);
|
||||
|
||||
z[zOffset] = front * continuedFractionCuda(a, b, x) / a;
|
||||
|
||||
if (symmCond) { // symmetry relation
|
||||
z[zOffset] = T(1) - z[zOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void betaIncForArrayCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* va, const LongType* aShapeInfo,
|
||||
const void* vb, const LongType* bShapeInfo, const void* vx,
|
||||
const LongType* xShapeInfo, void* vz, const LongType* zShapeInfo) {
|
||||
betaIncForArrayCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(va, aShapeInfo, vb, bShapeInfo, vx,
|
||||
xShapeInfo, vz, zShapeInfo);
|
||||
sd::DebugHelper::checkGlobalErrorCode("betaInc failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// overload betaInc for arrays, shapes of a, b and x must be the same !!!
|
||||
void betaInc(LaunchContext* context, NDArray& a, NDArray& b, NDArray& x, NDArray& output) {
|
||||
dim3 launchDims = getBetaInc(maxIter,output.lengthOf(),output.sizeOfT());
|
||||
|
||||
const auto xType = x.dataType();
|
||||
|
||||
PointersManager manager(context, "betaInc");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&a, &b, &x});
|
||||
BUILD_SINGLE_SELECTOR(xType, betaIncForArrayCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), a.specialBuffer(),
|
||||
a.specialShapeInfo(), b.specialBuffer(), b.specialShapeInfo(), x.specialBuffer(),
|
||||
x.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo()),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&a, &b, &x});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,428 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void clipByNormCuda(const void* vClipNorm, const void* vNorm, const LongType* normShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType* dimensions,
|
||||
const LongType dimsLen,
|
||||
const bool useAverage) {
|
||||
const T clipNorm = *reinterpret_cast<const T*>(vClipNorm);
|
||||
const T* norm = reinterpret_cast<const T*>(vNorm);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType zLen, tadLen, totalThreads;
|
||||
__shared__ int zRank, normRank;
|
||||
__shared__ const LongType *zShape;
|
||||
__shared__ const LongType *zStride;
|
||||
__shared__ const LongType *normStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
tadLen = zLen / shape::length(normShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Cache ranks
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
normRank = shape::rank(normShapeInfo);
|
||||
|
||||
// Cache shapes and strides
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
normStride = shape::stride(normShapeInfo);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
LongType zCoords[SD_MAX_RANK], normCoords[SD_MAX_RANK];
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < zLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
|
||||
// deduce norm coords
|
||||
for (int j = 0; j < dimsLen; ++j) normCoords[j] = zCoords[dimensions[j]];
|
||||
|
||||
LongType normOffset, zOffset;
|
||||
COORDS2INDEX(normRank, normStride, normCoords, normOffset);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
|
||||
const T actualNorm = useAverage ? static_cast<T>(norm[normOffset]) / static_cast<T>(tadLen) : static_cast<T>(norm[normOffset]);
|
||||
|
||||
if (actualNorm > clipNorm) z[zOffset] *= static_cast<T>(clipNorm) / static_cast<T>(actualNorm);
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST static void clipByNormCudaLauncher(const int blocksPerGrid, const int threadsPerBlock,
|
||||
const cudaStream_t* stream, const void* vClipNorm, const void* vNorm,
|
||||
const LongType* normShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* dimensions, const LongType dimsLen, const bool useAverage) {
|
||||
clipByNormCuda<T><<<blocksPerGrid, threadsPerBlock, 512, *stream>>>(vClipNorm, vNorm, normShapeInfo, vz, zShapeInfo,
|
||||
dimensions, dimsLen, useAverage);
|
||||
sd::DebugHelper::checkGlobalErrorCode("clipByNorm failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void clipByNorm(LaunchContext* context, NDArray* input, NDArray* output, const std::vector<LongType>& dims,
|
||||
NDArray* clipNorm, const bool isInplace, const bool useAverage) {
|
||||
NDArray* z = nullptr;
|
||||
|
||||
if (isInplace) {
|
||||
z = input;
|
||||
} else {
|
||||
output->assign(input);
|
||||
z = output;
|
||||
}
|
||||
|
||||
if (dims.empty()) {
|
||||
std::vector<LongType> empty;
|
||||
NDArray actualNorm = useAverage ? z->reduceAlongDimension(reduce::Norm2, &empty) / z->lengthOf()
|
||||
: z->reduceAlongDimension(reduce::Norm2, &empty);
|
||||
|
||||
if (actualNorm.e<float>(0) > clipNorm->e<float>(0)) *z *= *clipNorm / actualNorm;
|
||||
} else {
|
||||
NDArray actualNorms = z->reduceAlongDimension(reduce::Norm2, &dims);
|
||||
|
||||
std::vector<LongType> *dimsToExclude = ShapeUtils::evalDimsToExclude(z->rankOf(), dims.size(),dims.data());
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (z->lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
PointersManager manager(context, "clipByNorm");
|
||||
|
||||
const LongType* dimensions = reinterpret_cast<const LongType*>(
|
||||
manager.replicatePointer(dimsToExclude->data(), dimsToExclude->size() * sizeof(LongType)));
|
||||
|
||||
NDArray::prepareSpecialUse({z}, {z, &actualNorms, clipNorm});
|
||||
|
||||
|
||||
BUILD_SINGLE_SELECTOR(z->dataType(), clipByNormCudaLauncher,
|
||||
(blocksPerGrid,
|
||||
threadsPerBlock,
|
||||
context->getCudaStream(),
|
||||
clipNorm->specialBuffer(),
|
||||
actualNorms.specialBuffer(),
|
||||
actualNorms.specialShapeInfo(),
|
||||
z->specialBuffer(),
|
||||
z->specialShapeInfo(),
|
||||
dimensions,
|
||||
dimsToExclude->size(),
|
||||
useAverage),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({z}, {z, &actualNorms, clipNorm});
|
||||
|
||||
manager.synchronize();
|
||||
delete dimsToExclude;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void clipByNormBpCuda(const void* vClipNorm, const void* vx, const LongType* xShapeInfo, // input
|
||||
const void* vy, const LongType* yShapeInfo, // gradO
|
||||
const void* vNorm, const LongType* normShapeInfo, const void* vSum,
|
||||
const LongType* sumShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, // gradI
|
||||
const LongType* dimensions, const LongType dimsLen, const bool useAverage) {
|
||||
const T clipNorm = *reinterpret_cast<const T*>(vClipNorm);
|
||||
const T* norm = reinterpret_cast<const T*>(vNorm);
|
||||
const T* sum = reinterpret_cast<const T*>(vSum);
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
const T* y = reinterpret_cast<const T*>(vy);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType zLen, tadLen, totalThreads;
|
||||
__shared__ bool sameOffsets;
|
||||
__shared__ int zRank, yRank, normRank, sumRank, xRank;
|
||||
__shared__ const LongType *zShape;
|
||||
__shared__ const LongType *zStride;
|
||||
__shared__ const LongType *yStride;
|
||||
__shared__ const LongType *normStride;
|
||||
__shared__ const LongType *sumStride;
|
||||
__shared__ const LongType *xStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
tadLen = zLen / shape::length(normShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, yShapeInfo, zShapeInfo);
|
||||
|
||||
// Cache ranks
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
normRank = shape::rank(normShapeInfo);
|
||||
sumRank = shape::rank(sumShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
|
||||
// Cache shapes and strides
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
normStride = shape::stride(normShapeInfo);
|
||||
sumStride = shape::stride(sumShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
LongType zCoords[SD_MAX_RANK], normCoords[SD_MAX_RANK];
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < zLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
|
||||
LongType zOffset, yOffset;
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
if(sameOffsets) {
|
||||
yOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(yRank, yStride, zCoords, yOffset);
|
||||
}
|
||||
|
||||
// deduce norm coords
|
||||
for (int j = 0; j < dimsLen; ++j) normCoords[j] = zCoords[dimensions[j]];
|
||||
|
||||
LongType normOffset;
|
||||
COORDS2INDEX(normRank, normStride, normCoords, normOffset);
|
||||
|
||||
const T actualNorm = useAverage ? norm[normOffset] / tadLen : norm[normOffset];
|
||||
|
||||
if (actualNorm > clipNorm) {
|
||||
LongType sumOffset, xOffset;
|
||||
COORDS2INDEX(sumRank, sumStride, normCoords, sumOffset);
|
||||
if(sameOffsets) {
|
||||
xOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(xRank, xStride, zCoords, xOffset);
|
||||
}
|
||||
|
||||
const T sumVal = sum[sumOffset];
|
||||
z[zOffset] = (clipNorm / actualNorm) * y[yOffset] *
|
||||
(static_cast<T>(1.f) - (x[xOffset] * sumVal) / (actualNorm * actualNorm));
|
||||
} else {
|
||||
z[zOffset] = y[yOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void clipByNormBp_(LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
|
||||
const std::vector<LongType>& dims, NDArray* clipNorm, const bool useAverage) {
|
||||
const int rank = input->rankOf();
|
||||
|
||||
auto actualNorms = input->reduceAlongDimension(reduce::Norm2, &dims);
|
||||
|
||||
if (actualNorms.lengthOf() == 1) {
|
||||
const T norm = useAverage ? actualNorms.e<T>(0) / static_cast<T>(input->lengthOf()) : actualNorms.e<T>(0);
|
||||
|
||||
auto clipVal = clipNorm->e<T>(0);
|
||||
|
||||
if (norm > clipVal) {
|
||||
const T sum = input->reduceNumber(reduce::Sum).e<T>(0); // reduce to scalar
|
||||
const T factor1 = clipVal / norm;
|
||||
const T factor2 = static_cast<T>(1.f) / (norm * norm); // 1 / (norm*norm*norm)
|
||||
|
||||
auto lambda = LAMBDA_TT(x, y, sum, factor1, factor2) {
|
||||
return factor1 * y * (static_cast<T>(1.f) - factor2 * x * sum);
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(gradO, lambda, gradI);
|
||||
const_cast<NDArray*>(input)->applyPairwiseLambda(const_cast<NDArray*>(gradO), lambda, gradI);
|
||||
} else
|
||||
gradI->assign(gradO);
|
||||
} else {
|
||||
NDArray actualNorms = input->reduceAlongDimension(reduce::Norm2, &dims);
|
||||
NDArray sums = input->reduceAlongDimension(reduce::Sum, &dims);
|
||||
|
||||
std::vector<LongType> *dimsToExclude = ShapeUtils::evalDimsToExclude(gradI->rankOf(), dims.size(),dims.data());
|
||||
|
||||
|
||||
dim3 launchDims = clipDims(gradI->lengthOf());
|
||||
PointersManager manager(context, "clipByNormBp");
|
||||
|
||||
const LongType* dimensions = reinterpret_cast<const LongType*>(
|
||||
manager.replicatePointer(dimsToExclude->data(), dimsToExclude->size() * sizeof(LongType)));
|
||||
|
||||
NDArray::prepareSpecialUse({gradI}, {&actualNorms, &sums, clipNorm, input, gradO});
|
||||
clipByNormBpCuda<T><<<launchDims.y, launchDims.x,launchDims.z, *context->getCudaStream()>>>(
|
||||
clipNorm->specialBuffer(), input->specialBuffer(), input->specialShapeInfo(), gradO->specialBuffer(),
|
||||
gradO->specialShapeInfo(), actualNorms.specialBuffer(), actualNorms.specialShapeInfo(), sums.specialBuffer(),
|
||||
sums.specialShapeInfo(), gradI->specialBuffer(), gradI->specialShapeInfo(), dimensions, (LongType)dimsToExclude->size(),
|
||||
useAverage);
|
||||
sd::DebugHelper::checkGlobalErrorCode("clipByNorm failed");
|
||||
|
||||
NDArray::registerSpecialUse({gradI}, {&actualNorms, &sums, clipNorm, input, gradO});
|
||||
|
||||
manager.synchronize();
|
||||
delete dimsToExclude;
|
||||
}
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void clipByNormBp_,
|
||||
(sd::LaunchContext * context, NDArray* input, NDArray* gradO, NDArray* gradI,
|
||||
const std::vector<sd::LongType>& dimensions, NDArray* clipNorm, const bool useAverage),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void clipByNormBp(LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
|
||||
const std::vector<LongType>& dimensions, NDArray* clipNorm, const bool useAverage) {
|
||||
NDArray casted = clipNorm->cast(input->dataType());
|
||||
BUILD_SINGLE_SELECTOR(gradI->dataType(), clipByNormBp_,
|
||||
(context, &casted, gradO, gradI, dimensions, clipNorm, useAverage), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void clipByGlobalNorm_(LaunchContext* context, std::vector<NDArray*>& inputs, double clipNorm,
|
||||
memory::Workspace* workspace, std::vector<NDArray*>& outputs, bool isInplace) {
|
||||
T globalNorm = static_cast<T>(0.f);
|
||||
|
||||
for (auto i = 0; i < inputs.size(); i++) {
|
||||
auto input = inputs[i];
|
||||
auto l2norm = input->reduceNumber(reduce::Norm2);
|
||||
globalNorm += l2norm.e<T>(0) * l2norm.e<T>(0);
|
||||
}
|
||||
|
||||
globalNorm = math::sd_sqrt<T,T>(globalNorm);
|
||||
outputs[inputs.size()]->p(0, globalNorm);
|
||||
const T factor = static_cast<T>(clipNorm) / globalNorm;
|
||||
|
||||
for (size_t e = 0; e < inputs.size(); e++) {
|
||||
// all-reduce
|
||||
auto input = inputs[e];
|
||||
auto output = outputs[e];
|
||||
|
||||
if (static_cast<double>(globalNorm) <= clipNorm) {
|
||||
output->assign(input);
|
||||
} else {
|
||||
auto lambda = LAMBDA_T(_x, factor) { return _x * factor; });
|
||||
input->applyLambda(lambda, output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void clipByGlobalNorm(LaunchContext* context, std::vector<NDArray*>& inputs, double clipNorm,
|
||||
memory::Workspace* workspace, std::vector<NDArray*>& outputs, bool isInplace) {
|
||||
BUILD_SINGLE_SELECTOR(outputs[0]->dataType(), clipByGlobalNorm_,
|
||||
(context, inputs, clipNorm, workspace, outputs, isInplace), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void clipByGlobalNorm_,
|
||||
(sd::LaunchContext * context, std::vector<NDArray*> & inputs, double clipNorm,
|
||||
sd::memory::Workspace* workspace, std::vector<NDArray*>& outputs, bool isInplace),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
template <typename T>
|
||||
static void SD_KERNEL clipByValueKernel(void* input, const LongType* inputShape, void* output,
|
||||
const LongType* outputShape, double leftBound, double rightBound) {
|
||||
__shared__ T* outputBuf;
|
||||
__shared__ T* inputBuf;
|
||||
__shared__ LongType length;
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ LongType* inputShapePtr;
|
||||
__shared__ LongType* outputShapePtr;
|
||||
__shared__ LongType* inputStridePtr;
|
||||
__shared__ LongType* outputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outputBuf = reinterpret_cast<T*>(output);
|
||||
inputBuf = reinterpret_cast<T*>(input);
|
||||
length = shape::length(inputShape);
|
||||
|
||||
// Cache shape information
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputOffset;
|
||||
LongType outputOffset;
|
||||
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, inputOffset);
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, outputOffset);
|
||||
|
||||
if (inputBuf[inputOffset] > rightBound)
|
||||
outputBuf[outputOffset] = (T)rightBound;
|
||||
else if (inputBuf[inputOffset] < leftBound)
|
||||
outputBuf[outputOffset] = (T)leftBound;
|
||||
else
|
||||
outputBuf[outputOffset] = inputBuf[inputOffset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void clipByValue_(LaunchContext* context, NDArray* input, double leftBound, double rightBound,
|
||||
NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
dim3 launchDims = getLaunchDims("clip");
|
||||
clipByValueKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(input->specialBuffer(), input->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), leftBound,
|
||||
rightBound);
|
||||
sd::DebugHelper::checkGlobalErrorCode("clipByValue failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
void clipByValue(LaunchContext* context, NDArray* input, double leftBound, double rightBound, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), clipByValue_, (context, input, leftBound, rightBound, output),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void clipByValue_, (sd::LaunchContext * context, NDArray* input, double leftBound,
|
||||
double rightBound, NDArray* output);
|
||||
, SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,141 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 30.11.17.
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/col2im.h>
|
||||
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// columns [bS, iC, kH, kW, oH, oW] to be de-convoluted to image [bS, iC, iH, iW]
|
||||
template <typename T>
|
||||
static SD_KERNEL void col2imCuda(const void* columns, const LongType* colShapeInfo, void* image,
|
||||
const LongType* imShapeInfo, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW) {
|
||||
const T* col = reinterpret_cast<const T*>(columns);
|
||||
T* im = reinterpret_cast<T*>(image);
|
||||
|
||||
__shared__ LongType kH, kW, oH, oW;
|
||||
__shared__ LongType imLen;
|
||||
__shared__ LongType imRank;
|
||||
__shared__ LongType colRank;
|
||||
__shared__ LongType* imShape;
|
||||
__shared__ LongType* colShape;
|
||||
__shared__ LongType* imStride;
|
||||
__shared__ LongType* colStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
kH = dH * (colShapeInfo[3] - 1) + 1;
|
||||
kW = dW * (colShapeInfo[4] - 1) + 1;
|
||||
oH = colShapeInfo[5];
|
||||
oW = colShapeInfo[6];
|
||||
imLen = shape::length(imShapeInfo);
|
||||
|
||||
// Cache shape information
|
||||
imRank = shape::rank(imShapeInfo);
|
||||
colRank = shape::rank(colShapeInfo);
|
||||
imShape = shape::shapeOf(imShapeInfo);
|
||||
colShape = shape::shapeOf(colShapeInfo);
|
||||
imStride = shape::stride(imShapeInfo);
|
||||
colStride = shape::stride(colShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < imLen; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, imRank, imShape, coords);
|
||||
|
||||
LongType imOffset;
|
||||
COORDS2INDEX(imRank, imStride, coords, imOffset);
|
||||
|
||||
const auto bSiCoffset = coords[0] * colShape[7] + coords[1] * colShape[8];
|
||||
|
||||
const LongType imH = coords[2] + pH;
|
||||
const LongType imW = coords[3] + pW;
|
||||
|
||||
const LongType colHstart = (imH < kH) ? 0 : (imH - kH) / sH + 1;
|
||||
const LongType colWstart = (imW < kW) ? 0 : (imW - kW) / sW + 1;
|
||||
|
||||
const LongType colHend = sd::math::sd_min<LongType>(imH / sH + 1, oH);
|
||||
const LongType colWend = sd::math::sd_min<LongType>(imW / sW + 1, oW);
|
||||
|
||||
T val = static_cast<T>(0);
|
||||
|
||||
for (coords[4] = colHstart; coords[4] < colHend; ++coords[4]) {
|
||||
coords[2] = imH - coords[4] * sH;
|
||||
if (coords[2] % dH != 0) continue;
|
||||
|
||||
for (coords[5] = colWstart; coords[5] < colWend; ++coords[5]) {
|
||||
coords[3] = imW - coords[5] * sW;
|
||||
if (coords[3] % dW != 0) continue;
|
||||
|
||||
LongType colOffset;
|
||||
COORDS2INDEX(colRank, colStride, coords, colOffset);
|
||||
|
||||
val += col[bSiCoffset + colOffset];
|
||||
}
|
||||
}
|
||||
im[imOffset] = val;
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void col2imCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* columns, const LongType* colShapeInfo,
|
||||
void* image, const LongType* imShapeInfo, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW) {
|
||||
col2imCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(columns, colShapeInfo, image, imShapeInfo, sH,
|
||||
sW, pH, pW, dH, dW);
|
||||
DebugHelper::checkGlobalErrorCode( "col2im(...) failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void col2im(LaunchContext& context, NDArray* input, NDArray* output, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType iH, const LongType iW, const LongType dH, const LongType dW){
|
||||
PointersManager manager(&context, "col2im");
|
||||
dim3 dims = getCol2imLaunchParams(*input,*output);
|
||||
|
||||
NDArray::prepareSpecialUse({input}, {output});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), col2imCudaLauncher,
|
||||
(dims.x, dims.y, dims.z, context.getCudaStream(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), input->specialBuffer(), input->specialShapeInfo(), sH, sW, pH, pW, dH, dW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({input}, {output});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,200 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 AbdelRauf
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/LoopsCoordsHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/adjust_hue.h>
|
||||
#include <ops/declarable/helpers/imagesHelpers.h>
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename X>
|
||||
SD_HOST_DEVICE static uint8_t pack(const X* buff, const X& threshold) {
|
||||
uint8_t res;
|
||||
res = (buff[0] > threshold) << 7;
|
||||
res = res | ((buff[1] > threshold) << 6);
|
||||
res = res | ((buff[2] > threshold) << 5);
|
||||
res = res | ((buff[3] > threshold) << 4);
|
||||
res = res | ((buff[4] > threshold) << 3);
|
||||
res = res | ((buff[5] > threshold) << 2);
|
||||
res = res | ((buff[6] > threshold) << 1);
|
||||
res = res | (buff[7] > threshold);
|
||||
return res;
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST_DEVICE uint8_t pack<bool>(const bool* buff, const bool& threshold) {
|
||||
// ignore threshold
|
||||
uint8_t res;
|
||||
res = buff[0] << 7;
|
||||
res = res | (buff[1] << 6);
|
||||
res = res | (buff[2] << 5);
|
||||
res = res | (buff[3] << 4);
|
||||
res = res | (buff[4] << 3);
|
||||
res = res | (buff[5] << 2);
|
||||
res = res | (buff[6] << 1);
|
||||
res = res | buff[7];
|
||||
return res;
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
SD_HOST_DEVICE static uint8_t pack(const X* buff, int stride, const X& threshold) {
|
||||
uint8_t res;
|
||||
res = (buff[0] > threshold) << 7;
|
||||
res = res | ((buff[1 * stride] > threshold) << 6);
|
||||
res = res | ((buff[2 * stride] > threshold) << 5);
|
||||
res = res | ((buff[3 * stride] > threshold) << 4);
|
||||
res = res | ((buff[4 * stride] > threshold) << 3);
|
||||
res = res | ((buff[5 * stride] > threshold) << 2);
|
||||
res = res | ((buff[6 * stride] > threshold) << 1);
|
||||
res = res | (buff[7 * stride] > threshold);
|
||||
return res;
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST_DEVICE uint8_t pack<bool>(const bool* buff, int stride, const bool& threshold) {
|
||||
// ignore threshold
|
||||
uint8_t res;
|
||||
res = buff[0] << 7;
|
||||
res = res | (buff[1 * stride] << 6);
|
||||
res = res | (buff[2 * stride] << 5);
|
||||
res = res | (buff[3 * stride] << 4);
|
||||
res = res | (buff[4 * stride] << 3);
|
||||
res = res | (buff[5 * stride] << 2);
|
||||
res = res | (buff[6 * stride] << 1);
|
||||
res = res | buff[7 * stride];
|
||||
return res;
|
||||
}
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void SD_KERNEL cmpBitpack(const void* vx, void* vz, int rank, int len, const LongType* xStridesExtended,
|
||||
const LongType* outPutShapeInfo, T threshold) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
uint8_t* z = reinterpret_cast<uint8_t*>(vz);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto shapes = shape::shapeOf(outPutShapeInfo);
|
||||
auto zStrides = shape::stride(outPutShapeInfo);
|
||||
LongType coords[SD_MAX_RANK] = {};
|
||||
LongType* ptr_coords = (LongType*)&coords;
|
||||
// its extended as {rank+1} so xStridesExtended[rank] is valid
|
||||
auto inLastStride = xStridesExtended[rank];
|
||||
|
||||
for (auto k = tid; k < len; k += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(k, rank, shapes, ptr_coords);
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStridesExtended, ptr_coords, xOffset);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStrides, ptr_coords, zOffset);
|
||||
|
||||
auto buffPart = &(x[xOffset]);
|
||||
auto outBuffPart = &(z[zOffset]);
|
||||
*outBuffPart = pack<T>(buffPart, inLastStride, threshold);
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void SD_KERNEL cmpBitpackEws(const void* vx, void* vz, int len, const LongType xStride,
|
||||
const LongType yStride, T threshold) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
uint8_t* z = reinterpret_cast<uint8_t*>(vz);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (xStride == 1) {
|
||||
for (auto k = tid; k < len; k += gridDim.x * blockDim.x) {
|
||||
auto buffPart = &(x[k * 8]);
|
||||
auto outBuffPart = &(z[k * yStride]);
|
||||
*outBuffPart = pack<T>(buffPart, threshold);
|
||||
}
|
||||
} else {
|
||||
for (auto k = tid; k < len; k += gridDim.x * blockDim.x) {
|
||||
auto buffPart = &(x[k * 8 * xStride]);
|
||||
auto outBuffPart = &(z[k * yStride]);
|
||||
*outBuffPart = pack<T>(buffPart, xStride, threshold);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_HOST void cmpBitpackCudaLauncher(graph::Context& block, NDArray& input,
|
||||
|
||||
NDArray& thresholdScalar, NDArray& output) {
|
||||
T threshold = thresholdScalar.e<T>(0);
|
||||
|
||||
auto inStrides = input.stridesOf();
|
||||
auto rank = output.rankOf();
|
||||
|
||||
// threadblock size
|
||||
// grid size
|
||||
auto stream = block.launchContext()->getCudaStream();
|
||||
dim3 compareAndBitpackDims = getCompareAndBitpackDims(output.lengthOf());
|
||||
PointersManager manager(block.launchContext(), "compare_and_bitpack");
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
if (input.ordering() == 'c' && output.ordering() == 'c') {
|
||||
cmpBitpackEws<T><<<compareAndBitpackDims.y, compareAndBitpackDims.x,compareAndBitpackDims.z>>>(input.specialBuffer(), output.specialBuffer(),
|
||||
output.lengthOf(), inStrides[rank - 1],
|
||||
output.stridesOf()[rank - 1], threshold);
|
||||
sd::DebugHelper::checkGlobalErrorCode("cmpBitpackEws failed");
|
||||
|
||||
} else {
|
||||
// if output shape is {n1, n2, n3} then input shape is { n1. n2, n3 * 8}
|
||||
// therefore we can split input shape {n1, n2, n3 , 8} and correct its stride
|
||||
// as we do not need last shape info. lets just extend and correct its stride
|
||||
LongType extendedStrides[SD_MAX_RANK];
|
||||
for (int i = 0; i < rank; i++) {
|
||||
extendedStrides[i] = inStrides[i];
|
||||
}
|
||||
// lets correct new stride
|
||||
extendedStrides[rank - 1] = 8 * inStrides[rank - 1];
|
||||
extendedStrides[rank] = inStrides[rank - 1];
|
||||
|
||||
auto strideSize = (rank + 1) * sizeof(LongType);
|
||||
LongType* extendedStridesDevPtr =
|
||||
reinterpret_cast<LongType*>(manager.replicatePointer(extendedStrides, strideSize));
|
||||
cmpBitpack<T><<<compareAndBitpackDims.y, compareAndBitpackDims.x,compareAndBitpackDims.z>>>(input.specialBuffer(), output.specialBuffer(), rank,
|
||||
output.lengthOf(), extendedStridesDevPtr,
|
||||
output.specialShapeInfo(), threshold);
|
||||
sd::DebugHelper::checkGlobalErrorCode("compareAndBitpackDims failed");
|
||||
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void compareAndBitpack(graph::Context& block, NDArray& input, NDArray& threshold, NDArray& output) {
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), cmpBitpackCudaLauncher, (block, input, threshold, output), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,154 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
#include <ops/declarable/helpers/compare_elem.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
static SD_KERNEL void comparator(void *vx, const LongType *xShapeInfo, LongType length, const bool isStrict,
|
||||
void *reductionBuffer, bool *z) {
|
||||
auto x = reinterpret_cast<T *>(vx);
|
||||
auto reduction = reinterpret_cast<uint32_t *>(reductionBuffer);
|
||||
|
||||
extern __shared__ uint32_t shared[];
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// Cache shape information in shared memory
|
||||
__shared__ LongType xRank;
|
||||
__shared__ LongType *xShape;
|
||||
__shared__ LongType *xStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
shared[threadIdx.x] = 0;
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset0;
|
||||
LongType xOffset1;
|
||||
|
||||
// each thread will compare 2 elements: E and E+1
|
||||
for (int e = tid; e < length - 1; e += blockDim.x * gridDim.x) {
|
||||
INDEX2COORDS(e, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset0);
|
||||
INDEX2COORDS(e + 1, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset1);
|
||||
|
||||
auto val0 = x[xOffset0];
|
||||
auto val1 = x[xOffset1];
|
||||
|
||||
bool v = false;
|
||||
if (isStrict)
|
||||
v = val1 > val0;
|
||||
else
|
||||
v = val1 >= val0;
|
||||
|
||||
// store comparison result in shared memory
|
||||
shared[threadIdx.x] += v ? 0 : 1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// aggregate sums in shared memory
|
||||
for (LongType activeThreads = blockDim.x / 2; activeThreads > 0; activeThreads /= 2) {
|
||||
if (threadIdx.x < activeThreads) shared[threadIdx.x] += shared[threadIdx.x + activeThreads];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// store over the grid if we have more than 1 block
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int *>(reductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
|
||||
tid = threadIdx.x;
|
||||
if (threadIdx.x == 0) reduction[blockIdx.x] = shared[0];
|
||||
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
tc[16384] = 0;
|
||||
shared[threadIdx.x] = 0;
|
||||
|
||||
for (int i = threadIdx.x; i < gridDim.x; i += blockDim.x) shared[threadIdx.x] += reduction[i];
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (LongType activeThreads = blockDim.x / 2; activeThreads > 0; activeThreads /= 2) {
|
||||
if (threadIdx.x < activeThreads) shared[threadIdx.x] += shared[threadIdx.x + activeThreads];
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = shared[0] == 0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// if we have only 1 block, we just store results right away
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int *>(reductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = shared[0] == 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void _compare_elem(LaunchContext *context, NDArray *input, bool isStrictlyIncreasing, bool &output) {
|
||||
auto z = NDArrayFactory::create<bool>(false, context);
|
||||
|
||||
dim3 compareElemDims = getCompareElem(input->lengthOf());
|
||||
comparator<T><<<compareElemDims.x,compareElemDims.y,compareElemDims.z, *context->getCudaStream()>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), input->lengthOf(), isStrictlyIncreasing,
|
||||
context->getReductionPointer(), reinterpret_cast<bool *>(z.specialBuffer()));
|
||||
|
||||
z->tickWriteDevice();
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(), "is_strictly_increasing");
|
||||
|
||||
output = z->e<bool>(0);
|
||||
}
|
||||
|
||||
void compare_elem(LaunchContext *context, NDArray *input, bool isStrictlyIncreasing, bool &output) {
|
||||
auto xType = input->dataType();
|
||||
input->syncToDevice();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, _compare_elem, (context, input, isStrictlyIncreasing, output), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void _compare_elem,
|
||||
(sd::LaunchContext * context, NDArray *A, bool isStrictlyIncreasing, bool &output);
|
||||
, SD_COMMON_TYPES);
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,169 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
///
|
||||
///
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL static void concatCuda(void* pVx, void* pxShapeInfo, void* vz, const sd::LongType* zShapeInfo,
|
||||
const int axis) {
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType zLen, totalThreads;
|
||||
__shared__ LongType zRank;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Cache shape information
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = tid; i < zLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, zRank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(zRank, zStride, coords, zOffset);
|
||||
|
||||
int inArrIdx = 0;
|
||||
LongType* xShapeInfo = reinterpret_cast<sd::LongType**>(pxShapeInfo)[inArrIdx];
|
||||
|
||||
// Cache the input array's shape information for the current iteration
|
||||
LongType xRank = shape::rank(xShapeInfo);
|
||||
LongType* xStride = shape::stride(xShapeInfo);
|
||||
|
||||
while (coords[axis] >= xShapeInfo[axis + 1]) {
|
||||
coords[axis] -= xShapeInfo[axis + 1];
|
||||
xShapeInfo = reinterpret_cast<sd::LongType**>(pxShapeInfo)[++inArrIdx];
|
||||
// Update shape information for new input array
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
}
|
||||
|
||||
const auto* x = reinterpret_cast<T*>(reinterpret_cast<void**>(pVx)[inArrIdx]);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, coords, xOffset);
|
||||
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST static void concatCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, void* pVx, void* pxShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int axis) {
|
||||
concatCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(pVx, pxShapeInfo, vz, zShapeInfo, axis);
|
||||
DebugHelper::checkGlobalErrorCode("concat general case failed(...) failed");
|
||||
}
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void concat(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output, const int axis) {
|
||||
const int numInArrs = inArrs.size();
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, inArrs);
|
||||
|
||||
bool luckCase1 = false;
|
||||
|
||||
// prepare arrays of pointers on buffers and shapes
|
||||
std::vector<const void*> hInBuffers(numInArrs);
|
||||
std::vector<const LongType*> hInShapeInfo(numInArrs);
|
||||
std::vector <int> lenPerArray(numInArrs);
|
||||
for (int i = 0; i < numInArrs; i++) {
|
||||
hInBuffers[i] = inArrs[i]->specialBuffer();
|
||||
hInShapeInfo[i] = inArrs[i]->specialShapeInfo();
|
||||
lenPerArray[i] = inArrs[i]->isEmpty() ? 0 : inArrs[i]->isScalar() ? 1 : inArrs[i]->lengthOf();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "helpers::concat");
|
||||
|
||||
void* dInBuffers = manager.replicatePointer(hInBuffers.data(), hInBuffers.size() * sizeof(void*));
|
||||
|
||||
dim3 dims = getConcat(output.lengthOf());
|
||||
|
||||
if (luckCase1) { // for example {1,10} + {2,10} + {3,10} = {6, 10} order c; or {10,1} + {10,2} + {10,3} = {10, 6}
|
||||
void* z = static_cast<int8_t*>(output.specialBuffer());
|
||||
|
||||
for (sd::LongType i = 0; i < numInArrs; ++i) {
|
||||
const auto sizeofT = output.sizeOfT();
|
||||
const auto memAmountToCopy = inArrs[i]->lengthOf() * sizeofT;
|
||||
cudaMemcpyAsync(z, reinterpret_cast<const int8_t*>(inArrs[i]->specialBuffer()), memAmountToCopy,
|
||||
cudaMemcpyDeviceToDevice, *context->getCudaStream());
|
||||
z = static_cast<int8_t*>(z) + memAmountToCopy;
|
||||
}
|
||||
|
||||
if (cudaStreamSynchronize(*context->getCudaStream()) != 0)
|
||||
THROW_EXCEPTION("concat cuda: luckCase1 failed!");
|
||||
|
||||
for (int i = 0; i < numInArrs; ++i) inArrs[i]->tickReadDevice();
|
||||
output.tickWriteDevice();
|
||||
manager.synchronize();
|
||||
output.syncToHost();
|
||||
return;
|
||||
}
|
||||
|
||||
void* dInShapeInfo = manager.replicatePointer(hInShapeInfo.data(), hInShapeInfo.size() * sizeof(LongType*));
|
||||
|
||||
BUILD_SINGLE_SELECTOR(inArrs[0]->dataType(), concatCudaLauncher,
|
||||
(dims.x, dims.y, dims.z, context->getCudaStream(), dInBuffers, dInShapeInfo,
|
||||
output.specialBuffer(), output.specialShapeInfo(), axis),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
manager.synchronize();
|
||||
manager.synchronize();
|
||||
output.syncToHost();
|
||||
NDArray::registerSpecialUse({&output}, inArrs);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,159 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
|
||||
#include <ops/declarable/helpers/confusion.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL static void copyBuffers(LongType* destination, void const* source, LongType bufferLength) {
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
const T * sourceCast = reinterpret_cast<T const*>(source);
|
||||
for (int t = tid; t < bufferLength; t += step) {
|
||||
destination[t] = static_cast<LongType>(sourceCast[t]);
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL static void confusionFunctorKernel(LongType* labelsBuffer, LongType* predictionBuffer, LongType bufferLength, void const* weightsBuffer, void* outputBuffer,
|
||||
const LongType* tadShape, const LongType* tadOffsets) {
|
||||
__shared__ int arrIdx, blocksPerArr;
|
||||
__shared__ T* z;
|
||||
__shared__ T const* w;
|
||||
__shared__ LongType *zShapeInfo, *xShapeInfo, arrLen;
|
||||
__shared__ LongType tadRank;
|
||||
__shared__ LongType* tadShapePtr;
|
||||
__shared__ LongType* tadStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z = reinterpret_cast<T*>(outputBuffer);
|
||||
w = reinterpret_cast<T const*>(weightsBuffer);
|
||||
arrLen = shape::length(tadShape);
|
||||
|
||||
// Cache shape information
|
||||
tadRank = shape::rank(tadShape);
|
||||
tadShapePtr = shape::shapeOf(tadShape);
|
||||
tadStridePtr = shape::stride(tadShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
LongType predCoords[SD_MAX_RANK];
|
||||
LongType predOffset;
|
||||
|
||||
for (int t = tid; t < bufferLength; t += step) {
|
||||
auto label = labelsBuffer[t];
|
||||
auto pred = predictionBuffer[t];
|
||||
auto tZ = z + tadOffsets[label];
|
||||
T val = (weightsBuffer == nullptr ? (T)1.0f : w[t]);
|
||||
|
||||
INDEX2COORDS(pred, tadRank, tadShapePtr, predCoords);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, predCoords, predOffset);
|
||||
tZ[predOffset] = val;
|
||||
}
|
||||
}
|
||||
template <typename X, typename Z>
|
||||
void _confusionFunctor(LaunchContext* context, NDArray* labels, NDArray* predictions, NDArray* weights,
|
||||
NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
auto pack = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), 1);
|
||||
PointersManager manager(context, "helpers::confusion");
|
||||
predictions->syncToDevice();
|
||||
LongType* labelsLongBuffer = labels->dataType() == INT64 ? (LongType*)labels->specialBuffer() : nullptr;
|
||||
LongType* predictionLongBuffer =
|
||||
predictions->dataType() == INT64 ? (LongType*)predictions->specialBuffer() : nullptr;
|
||||
|
||||
|
||||
|
||||
dim3 conf = getLaunchDims("confusion_matrix");
|
||||
if (labelsLongBuffer == nullptr) {
|
||||
auto err = cudaMalloc(&labelsLongBuffer, labels->lengthOf() * sizeof(LongType));
|
||||
if (err != 0) throw cuda_exception::build("Cannot allocate memory for labels long buffer", err);
|
||||
// copy with type conversion
|
||||
copyBuffers<X><<<conf.x, conf.y, conf.z, *stream>>>(labelsLongBuffer, labels->specialBuffer(), labels->lengthOf());
|
||||
sd::DebugHelper::checkGlobalErrorCode("copyBuffers failed");
|
||||
|
||||
}
|
||||
|
||||
if (predictionLongBuffer == nullptr) {
|
||||
auto err = cudaMalloc(&predictionLongBuffer, predictions->lengthOf() * sizeof(LongType));
|
||||
if (err != 0) throw cuda_exception::build("Cannot allocate memory for predictions long buffer", err);
|
||||
// copy with type conversion
|
||||
copyBuffers<X>
|
||||
<<<256, 512, 1024, *stream>>>(predictionLongBuffer, predictions->specialBuffer(), predictions->lengthOf());
|
||||
sd::DebugHelper::checkGlobalErrorCode("copyBuffers failed");
|
||||
|
||||
}
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
|
||||
|
||||
auto bufferLength = labels->lengthOf();
|
||||
dim3 launchDims = getLaunchDims("confusionMatrix");
|
||||
confusionFunctorKernel<Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
labelsLongBuffer, predictionLongBuffer, bufferLength, weights != nullptr ? weights->specialBuffer() : nullptr,
|
||||
output->specialBuffer(), pack->specialShapeInfo(), pack->specialOffsets());
|
||||
sd::DebugHelper::checkGlobalErrorCode("confusionFunctorKernel failed");
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
if (predictionLongBuffer != predictions->specialBuffer()) {
|
||||
cudaError_t err = cudaFree(predictionLongBuffer);
|
||||
if (err != 0) throw cuda_exception::build("Cannot deallocate memory for predictions long buffer", err);
|
||||
}
|
||||
|
||||
if (labelsLongBuffer != labels->specialBuffer()) {
|
||||
cudaError_t err = cudaFree(labelsLongBuffer);
|
||||
if (err != 0) throw cuda_exception::build("Cannot deallocate memory for labels long buffer", err);
|
||||
}
|
||||
}
|
||||
|
||||
void confusionFunctor(LaunchContext* context, NDArray* labels, NDArray* predictions, NDArray* weights,
|
||||
NDArray* output) {
|
||||
auto xType = predictions->dataType();
|
||||
auto zType = output->dataType(); // weights can be null
|
||||
NDArray::prepareSpecialUse({output}, {labels, predictions, weights});
|
||||
BUILD_DOUBLE_SELECTOR(xType, zType, _confusionFunctor, (context, labels, predictions, weights, output),
|
||||
SD_INDEXING_TYPES, SD_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {labels, predictions, weights});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,155 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <math/templatemath.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include <helpers/DebugHelper.h>
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// columns [bS, iC, kD, kH, kW, oD, oH, oW] to be de-convoluted to volume [bS, iC, iD, iH, iW]
|
||||
template <typename T>
|
||||
static SD_KERNEL void col2volCuda(const void* columns, const LongType* colShapeInfo, void* volume,
|
||||
const LongType* volShapeInfo, const int sD, const int sH, const int sW,
|
||||
const int pD, const int pH, const int pW, const int dD, const int dH, const int dW) {
|
||||
const T* col = reinterpret_cast<const T*>(columns);
|
||||
T* vol = reinterpret_cast<T*>(volume);
|
||||
|
||||
__shared__ LongType kD, kH, kW, oD, oH, oW, *sharedMem;
|
||||
__shared__ LongType volLen;
|
||||
__shared__ LongType volRank;
|
||||
__shared__ LongType* volShape;
|
||||
__shared__ LongType* volStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
oD = colShapeInfo[6];
|
||||
oH = colShapeInfo[7];
|
||||
oW = colShapeInfo[8];
|
||||
|
||||
kD = dD * (colShapeInfo[3] - 1) + 1;
|
||||
kH = dH * (colShapeInfo[4] - 1) + 1;
|
||||
kW = dW * (colShapeInfo[5] - 1) + 1;
|
||||
|
||||
volLen = shape::length(volShapeInfo);
|
||||
|
||||
// Cache shape information
|
||||
volRank = shape::rank(volShapeInfo);
|
||||
volShape = shape::shapeOf(volShapeInfo);
|
||||
volStride = shape::stride(volShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * 8;
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < volLen; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, volRank, volShape, coords);
|
||||
|
||||
sd::LongType volOffset;
|
||||
COORDS2INDEX(volRank, volStride, coords, volOffset);
|
||||
|
||||
const auto bSiCoffset = coords[0] * colShapeInfo[9] + coords[1] * colShapeInfo[10];
|
||||
|
||||
const LongType imD = coords[2] + pD;
|
||||
const LongType imH = coords[3] + pH;
|
||||
const LongType imW = coords[4] + pW;
|
||||
|
||||
const LongType colDstart = (imD < kD) ? 0 : (imD - kD) / sD + 1;
|
||||
const LongType colHstart = (imH < kH) ? 0 : (imH - kH) / sH + 1;
|
||||
const LongType colWstart = (imW < kW) ? 0 : (imW - kW) / sW + 1;
|
||||
|
||||
const LongType colDend = sd::math::sd_min<LongType>(imD / sD + 1, oD);
|
||||
const LongType colHend = sd::math::sd_min<LongType>(imH / sH + 1, oH);
|
||||
const LongType colWend = sd::math::sd_min<LongType>(imW / sW + 1, oW);
|
||||
|
||||
T val = static_cast<T>(0);
|
||||
|
||||
for (LongType colD = colDstart; colD < colDend; ++colD) {
|
||||
coords[2] = imD - colD * sD;
|
||||
if (coords[2] % dD != 0) continue;
|
||||
|
||||
for (LongType colH = colHstart; colH < colHend; ++colH) {
|
||||
coords[3] = imH - colH * sH;
|
||||
if (coords[3] % dH != 0) continue;
|
||||
|
||||
for (LongType colW = colWstart; colW < colWend; ++colW) {
|
||||
coords[4] = imW - colW * sW;
|
||||
if (coords[4] % dW != 0) continue;
|
||||
|
||||
val += col[bSiCoffset + (coords[2] / dD) * colShapeInfo[11] + (coords[3] / dH) * colShapeInfo[12] +
|
||||
(coords[4] / dW) * colShapeInfo[13] + colD * colShapeInfo[14] + colH * colShapeInfo[15] +
|
||||
colW * colShapeInfo[16]];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
vol[volOffset] = val;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void col2volCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* columns, const LongType* colShapeInfo,
|
||||
void* volume, const LongType* volShapeInfo, const LongType sD, const LongType sH,
|
||||
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
|
||||
const LongType dW) {
|
||||
col2volCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(columns, colShapeInfo, volume, volShapeInfo,
|
||||
sD, sH, sW, pD, pH, pW, dD, dH, dW);
|
||||
DebugHelper::checkGlobalErrorCode( "col2vol(...) failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::col2vol(graph::Context& block, NDArray& col, NDArray& vol, const LongType sD, const LongType sH,
|
||||
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
|
||||
const LongType dW) {
|
||||
PointersManager manager(block.launchContext(), "col2vol");
|
||||
|
||||
|
||||
dim3 col2VolDims = getCol2VolDims(vol.lengthOf(), col.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&vol}, {&col});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
vol.dataType(), col2volCudaLauncher,
|
||||
(col2VolDims.x, col2VolDims.y, col2VolDims.z, block.launchContext()->getCudaStream(), col.specialBuffer(),
|
||||
col.specialShapeInfo(), vol.specialBuffer(), vol.specialShapeInfo(), sD, sH, sW, pD, pH, pW, dD, dH, dW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&vol}, {&col});
|
||||
|
||||
manager.synchronize();
|
||||
DebugHelper::checkGlobalErrorCode( "col2vol(...) failed");
|
||||
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,137 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/addBias.h>
|
||||
#include <ops/declarable/helpers/col2im.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <ops/declarable/helpers/im2col.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void conv2d_(sd::graph::Context& block, NDArray* input, NDArray* weights, NDArray* bias,
|
||||
NDArray* output, const LongType kH, const LongType kW, const LongType sH, const LongType sW, LongType pH, LongType pW,
|
||||
const LongType dH, const LongType dW, const int paddingMode, const int isNCHW, const int wFormat) {
|
||||
|
||||
// input [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
|
||||
// weights [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
// bias [oC]
|
||||
// output [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
|
||||
|
||||
LongType bS = input->sizeAt(0);
|
||||
LongType iC = ConvolutionUtils::inChannels(weights->shapeInfo(), wFormat);
|
||||
LongType oC = ConvolutionUtils::outChannels(weights->shapeInfo(), wFormat);
|
||||
LongType iH = ConvolutionUtils::inputHeight(input->shapeInfo(), isNCHW);
|
||||
LongType iW = ConvolutionUtils::inputWidth(input->shapeInfo(), isNCHW);
|
||||
LongType oH = ConvolutionUtils::calcOutDimConv(iH, kH, sH, pH, dH, paddingMode);
|
||||
LongType oW = ConvolutionUtils::calcOutDimConv(iW, kW, sW, pW, dW, paddingMode);
|
||||
|
||||
std::vector<LongType> wAxes;
|
||||
if (0 == wFormat)
|
||||
wAxes = {0, 1, 2};
|
||||
else if (1 == wFormat)
|
||||
wAxes = {2, 3, 1};
|
||||
else
|
||||
wAxes = {1, 2, 3};
|
||||
|
||||
|
||||
std::vector<sd::LongType> colShape = {bS, iC, kH, kW, oH, oW};
|
||||
NDArray *col = new NDArray('c', colShape, input->dataType(), input->getContext());
|
||||
std::vector<LongType> colPermute = {0, 3, 4, 5, 1, 2}; // {bS, iC, kH, kW, oH, oW}
|
||||
NDArray *colP = col->permute(colPermute, false, false); // permute() already returns NDArray*
|
||||
std::vector<sd::LongType> mmulResShape = {bS * oH * oW, oC};
|
||||
NDArray mmulResult('f', mmulResShape, output->dataType(), output->getContext());
|
||||
|
||||
std::vector<LongType> permuteForOutput = {0, 3, 1, 2};
|
||||
|
||||
//----- calculation of output -----//
|
||||
auto ctx = block.launchContext();
|
||||
|
||||
|
||||
NDArray zero = NDArrayFactory::create(0.f, input->getContext());
|
||||
if (isNCHW) {
|
||||
helpers::im2col(*ctx, *input, *colP, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
zero);
|
||||
} else {
|
||||
std::vector<sd::LongType> permute = {0, 3, 1, 2};
|
||||
// For NHWC, we need to permute the input to NCHW before im2col
|
||||
NDArray* inputNchw = input->permute(permute, 0, false); // permute() already returns NDArray*
|
||||
helpers::im2col(*ctx, *inputNchw, *colP, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
zero);
|
||||
delete inputNchw; // Clean up permuted array
|
||||
}
|
||||
|
||||
|
||||
|
||||
std::vector<sd::LongType> permute = {0, 3, 4, 5, 1, 2};
|
||||
block.pushIntermediateResult(col);
|
||||
|
||||
std::vector<sd::LongType> shape = {bS * oH * oW, kW * kH * iC};
|
||||
auto im2colReshape = col->reshape('c', shape, true);
|
||||
|
||||
auto weightsPermuted = weights->permute(permuteForOutput, 0, false);
|
||||
std::vector<LongType> weightShape = {iC * kH * kW, oC};
|
||||
auto reshapedW = weightsPermuted.reshape('f', weightShape, false);
|
||||
MmulHelper::matmul(&im2colReshape, &reshapedW, &mmulResult, false, false, 1.0, 0.0);
|
||||
|
||||
|
||||
std::vector<LongType> mmulResultShape = {oH, oW, bS, oC};
|
||||
auto reshaped = mmulResult.reshape('f', mmulResultShape, false);
|
||||
std::vector<sd::LongType> permutedShape = {2, 3, 1,0};
|
||||
auto permuted = reshaped.permute(permutedShape, 0, false);
|
||||
|
||||
// Reshape and copy result to output
|
||||
if (isNCHW) {
|
||||
output->assign(&permuted);
|
||||
} else {
|
||||
std::vector<sd::LongType> otherPermute = {0,2,3,1};
|
||||
permuted = permuted.permute(otherPermute, 0, false);
|
||||
output->assign(&permuted);
|
||||
}
|
||||
|
||||
//----- add biases if required -----//
|
||||
if (bias) {
|
||||
helpers::addBias(block, *output, *bias, *output, isNCHW);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::conv2d(sd::graph::Context& block, NDArray* input, NDArray* weights,
|
||||
NDArray* bias, NDArray* output, const LongType kH, const LongType kW, const LongType sH,
|
||||
const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW, const int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input->dataType(), conv2d_,
|
||||
(block, input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,176 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <execution/Threads.h>
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <ops/declarable/helpers/addBias.h>
|
||||
#include <ops/declarable/helpers/col2im.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <ops/declarable/helpers/im2col.h>
|
||||
|
||||
#include "helpers/ShapeUtils.h"
|
||||
#if NOT_EXCLUDED(OP_col2im) && NOT_EXCLUDED(OP_im2col)
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
static void conv2dBP_(sd::graph::Context& block, NDArray* input, NDArray* weights, NDArray* bias,
|
||||
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW,
|
||||
const int paddingMode, const int isNCHW, const int wFormat) {
|
||||
|
||||
// input [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
|
||||
// weights [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
// bias [oC]
|
||||
// gradO [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
|
||||
// gradI [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
|
||||
// gradW [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
// gradB [oC]
|
||||
|
||||
const LongType bS = input->sizeAt(0); // batch size
|
||||
const LongType iC = isNCHW ? input->sizeAt(1) : input->sizeAt(3); // input channels
|
||||
const LongType iH = isNCHW ? input->sizeAt(2) : input->sizeAt(1); // input height
|
||||
const LongType iW = isNCHW ? input->sizeAt(3) : input->sizeAt(2); // input width
|
||||
|
||||
const LongType oC = isNCHW ? gradO->sizeAt(1) : gradO->sizeAt(3); // output channels
|
||||
const LongType oH = isNCHW ? gradO->sizeAt(2) : gradO->sizeAt(1); // output height
|
||||
const LongType oW = isNCHW ? gradO->sizeAt(3) : gradO->sizeAt(2); // output width
|
||||
NDArray *inputPermuted, *gradOPermuted, *gradIPermuted;
|
||||
if (!isNCHW) {
|
||||
std::vector<sd::LongType> permute = {0, 3, 1, 2};
|
||||
inputPermuted = input->permute(permute,false,false); // permute() already returns NDArray*
|
||||
gradOPermuted = gradO->permute(permute,false,false);
|
||||
gradIPermuted = gradI->permute(permute,false,false);
|
||||
} else {
|
||||
inputPermuted = input;
|
||||
gradOPermuted = gradO;
|
||||
gradIPermuted = gradI;
|
||||
}
|
||||
|
||||
std::vector<sd::LongType> reshape = {oC,bS * oH * oW};
|
||||
// Reshape gradO to 2D: [oC, bS * oH * oW]
|
||||
NDArray gradO2d = gradOPermuted->reshape(gradOPermuted->ordering(), reshape,false);
|
||||
|
||||
// Perform im2col
|
||||
NDArray* columns;
|
||||
if (block.hasIntermediateResults()) {
|
||||
columns = block.intermediateResult(0);
|
||||
if (columns->rankOf() < 6) {
|
||||
columns->reshapei({bS, iC, kH, kW, oH, oW});
|
||||
}
|
||||
} else {
|
||||
std::vector<sd::LongType> colShape = {bS, iC, kH, kW, oH, oW};
|
||||
columns = new NDArray(inputPermuted->ordering(), colShape, inputPermuted->dataType(), inputPermuted->getContext());
|
||||
auto ctx = block.launchContext();
|
||||
NDArray zero = NDArrayFactory::create<double>(0., inputPermuted->getContext());
|
||||
helpers::im2col(*ctx, *inputPermuted, *columns, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
zero);
|
||||
}
|
||||
|
||||
// Calculate gradW
|
||||
if (gradW) {
|
||||
std::vector<sd::LongType> columnsShape = {bS * oH * oW, iC * kH * kW};
|
||||
std::vector<sd::LongType> wShape = {oC, iC * kH * kW};
|
||||
NDArray columns2d = columns->reshape('c', columnsShape,false);
|
||||
std::vector<sd::LongType > permu = {1,0};
|
||||
NDArray &gradW2d = gradW->reshape('f', wShape,false).permute(permu,false,false);
|
||||
|
||||
MmulHelper::matmul( &columns2d,&gradO2d, &gradW2d, true, true, 1.0, 0.0, &gradW2d);
|
||||
gradW->assign(&gradW2d);
|
||||
|
||||
}
|
||||
|
||||
// Calculate gradB
|
||||
if (gradB) {
|
||||
std::vector<LongType> axes = {1}; // Sum over bS, oH, oW
|
||||
gradO2d.reduceAlongDimension(reduce::Sum, gradB, &axes);
|
||||
}
|
||||
|
||||
// Calculate gradI
|
||||
NDArray weights2d;
|
||||
if (wFormat == 0) {
|
||||
std::vector<sd::LongType> permute = {3,2,1,0};
|
||||
std::vector<sd::LongType> shape2 = {iC * kH * kW, oC};
|
||||
weights2d = weights->permute(permute,false,false).reshape('f', shape2);
|
||||
} else if (wFormat == 1) {
|
||||
std::vector<sd::LongType> shape2 = {iC * kH * kW, oC};
|
||||
weights2d = weights->reshape('f', shape2);
|
||||
} else {
|
||||
std::vector<sd::LongType> shape2 = {iC * kH * kW, oC};
|
||||
std::vector<sd::LongType> permute3 = {0, 2, 3, 1};
|
||||
weights2d = weights->permute(permute3,false,false).reshape('f', shape2);
|
||||
}
|
||||
|
||||
std::vector<sd::LongType> colShape2 = {iC * kH * kW, bS * oH * oW};
|
||||
NDArray columns2d('c', colShape2, columns->dataType(), columns->getContext());
|
||||
|
||||
|
||||
MmulHelper::matmul(&weights2d, &gradO2d, &columns2d, false, false, 1.0, 0.0);
|
||||
//Calculate epsilonNext by doing im2col reduction.
|
||||
//Current col2im implementation expects input with order: [miniBatch,channels,kH,kW,outH,outW]
|
||||
//currently have [kH,kW,inDepth,outW,outH,miniBatch] -> permute first
|
||||
auto eps6d = columns2d.newShapeNoCopy({kH, kW,iC, oW, oH, bS }, 'f');
|
||||
std::vector<sd::LongType> epsPerm = {5,2,1,0,4,3};
|
||||
auto permuted = eps6d->permute(epsPerm,false,false);
|
||||
|
||||
// Perform col2im
|
||||
auto ctx = block.launchContext();
|
||||
helpers::col2im(*ctx, &permuted, gradIPermuted, sH, sW, pH, pW, iH, iW, dH, dW);
|
||||
// Handle NHWC format if necessary
|
||||
if (!isNCHW) {
|
||||
std::vector<sd::LongType> permute = {0, 2,3,1};
|
||||
auto permuted = gradIPermuted->permute(permute,false,false); // [bS, iC, iH, iW] -> [bS, iH, iW, iC]
|
||||
gradI->assign(&permuted); // [bS, iC, iH, iW] -> [bS, iH, iW, iC]
|
||||
}
|
||||
|
||||
// Clean up
|
||||
if (!isNCHW) {
|
||||
delete inputPermuted;
|
||||
delete gradOPermuted;
|
||||
delete gradIPermuted;
|
||||
}
|
||||
if (!block.hasIntermediateResults()) {
|
||||
delete columns;
|
||||
}
|
||||
}
|
||||
|
||||
void ConvolutionUtils::conv2dBP(sd::graph::Context& block, NDArray* input, NDArray* weights,
|
||||
NDArray* bias, NDArray* gradO, NDArray* gradI, NDArray* gradW,
|
||||
NDArray* gradB, const LongType kH, const LongType kW, const LongType sH, const LongType sW, LongType pH, LongType pW,
|
||||
const LongType dH, const LongType dW, const int paddingMode, const int isNCHW,
|
||||
const int wFormat) {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(input->dataType(), conv2dBP_,
|
||||
(block, input, weights, bias, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
paddingMode, isNCHW, wFormat),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,122 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/addBias.h>
|
||||
#include <ops/declarable/helpers/col2im.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <ops/declarable/helpers/im2col.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void depthwiseConv2d_(sd::graph::Context& block, NDArray* input, NDArray* weights,
|
||||
NDArray* bias, NDArray* output, const LongType kH, const LongType kW, const LongType sH,
|
||||
const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW, const int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
// input [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
|
||||
// weights [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
|
||||
// bias [oC] = iC*mC
|
||||
// output [bS, oH, oW, iC*mC] (NHWC) or [bS, iC*mC, oH, oW] (NCHW)
|
||||
|
||||
// kH filter(kernel) height
|
||||
// kW filter(kernel) width
|
||||
// sH strides height
|
||||
// sW strides width
|
||||
// pH paddings height
|
||||
// pW paddings width
|
||||
// dH dilations height
|
||||
// dW dilations width
|
||||
// paddingMode 0-VALID, 1-SAME
|
||||
// isNCHW 0-NCHW, 1-NHWC
|
||||
|
||||
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
|
||||
// iC*mC), output channels, output height/width
|
||||
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWiC, indWmC, indWkH, indOoH);
|
||||
mC = weights->sizeAt(indWmC); // channels multiplier
|
||||
|
||||
std::vector<std::vector<sd::LongType>> modifColumns = {
|
||||
{1, 0, 4, 5, 2, 3},
|
||||
{iC, bS * oH * oW, kH * kW}}; // [bS,iC,kH,kW,oH,oW] -> [iC,bS,oH,oW,kH,kW] -> [iC,bS*oH*oW,kH*kW]
|
||||
std::vector<std::vector<sd::LongType>> modifOutput, modifWeights;
|
||||
std::vector<sd::LongType> outReShape;
|
||||
|
||||
if (!isNCHW) {
|
||||
outReShape = {bS, oH, oW, iC, mC}; // [bS,oH,oW,iC*mC] -> [bS,oH,oW,iC,mC]
|
||||
modifOutput = {{3, 0, 1, 2, 4},
|
||||
{iC, bS * oH * oW, mC}}; // [bS,oH,oW,iC,mC] -> [iC,bS,oH,oW,mC] -> [iC,bS*oH*oW,mC]
|
||||
std::vector<sd::LongType> permuteVec = {0, 3, 1, 2};
|
||||
input = input->permute(permuteVec, false, false); // permute() already returns NDArray*
|
||||
} else {
|
||||
outReShape = {bS, iC, mC, oH, oW}; // [bS,iC*mC,oH,oW] -> [bS,iC,mC,oH,oW]
|
||||
modifOutput = {{1, 0, 3, 4, 2},
|
||||
{iC, bS * oH * oW, mC}}; // [bS,iC,mC,oH,oW] -> [iC,bS,oH,oW,mC] -> [iC,bS*oH*oW,mC]
|
||||
}
|
||||
|
||||
if (0 == wFormat)
|
||||
modifWeights = {{2, 0, 1, 3}, {iC, kH * kW, mC}};
|
||||
else if (1 == wFormat)
|
||||
modifWeights = {{1, 2, 3, 0}, {iC, kH * kW, mC}};
|
||||
else
|
||||
modifWeights = {{3, 1, 2, 0}, {iC, kH * kW, mC}};
|
||||
|
||||
if (paddingMode == 1) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
std::vector<sd::LongType> colShape = {bS, iC, kH, kW, oH, oW};
|
||||
|
||||
NDArray columns(input->ordering(),colShape, input->dataType(), input->getContext());
|
||||
NDArray outputReshaped = output->reshape(output->ordering(), outReShape, false);
|
||||
|
||||
NDArray zero = NDArrayFactory::create(0.f, input->getContext());
|
||||
helpers::im2col(
|
||||
*output->getContext(), *input, columns, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
zero); // [bS, iC, iH, iW] is convoluted to [bS, iC, kH, kW, oH, oW]
|
||||
MmulHelper::tensorDot(&columns, weights, &outputReshaped, modifColumns, modifWeights,
|
||||
modifOutput); // [iC, bS*oH*oW, kW*kH] x [iC, kH*kW, mC] = [iC, bS*oH*oW, mC]
|
||||
|
||||
if (bias)
|
||||
helpers::addBias(block, *output, *bias, *output, isNCHW);
|
||||
|
||||
if (!isNCHW) delete input;
|
||||
}
|
||||
|
||||
void ConvolutionUtils::depthwiseConv2d(sd::graph::Context& block, NDArray* input, NDArray* weights,
|
||||
NDArray* bias, NDArray* output, const LongType kH, const LongType kW, const LongType sH,
|
||||
const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW, const int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input->dataType(), depthwiseConv2d_,
|
||||
(block, input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,146 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/col2im.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <ops/declarable/helpers/im2col.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void depthwiseConv2dBP_(NDArray* input, NDArray* weights, NDArray* bias, NDArray* gradO,
|
||||
NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH, const LongType kW, const LongType sH,
|
||||
const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW, const int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
// input [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
|
||||
// weights [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
|
||||
// bias [oC] = [iC*mC]
|
||||
// gradO [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
|
||||
// gradI [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW), epsilon
|
||||
// gradW [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
|
||||
// gradB [oC]
|
||||
|
||||
// kH filter(kernel) height
|
||||
// kW filter(kernel) width
|
||||
// sH strides height
|
||||
// sW strides width
|
||||
// pH paddings height
|
||||
// pW paddings width
|
||||
// dH dilations height
|
||||
// dW dilations width
|
||||
// paddingMode 0-VALID, 1-SAME
|
||||
// isNCHW 0-NHWC, 1-NCHW
|
||||
|
||||
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
|
||||
// iC*mC), output channels, output height/width
|
||||
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWiC, indWmC, indWkH, indOoH);
|
||||
mC = weights->sizeAt(indWmC); // channels multiplier
|
||||
|
||||
std::vector<std::vector<LongType>> modifColumns = {
|
||||
{1, 2, 3, 0, 4, 5}, {iC, kH * kW, bS * oH * oW}}; // [bS,iC,kH,kW,oH,oW] -> [iC, kH*kW, bS*oH*oW]
|
||||
std::vector<std::vector<LongType>> modifGradO1, modifGradO2, modifWeights;
|
||||
std::vector<LongType> gradOreShape;
|
||||
|
||||
if (!isNCHW) {
|
||||
gradOreShape = {bS, oH, oW, iC, mC}; // [bS,oH,oW,iC*mC] -> [bS,oH,oW,iC,mC]
|
||||
modifGradO1 = {{3, 0, 1, 2, 4},
|
||||
{iC, bS * oH * oW, mC}}; // [bS,oH,oW,iC,mC] -> [iC,bS,oH,oW,mC] -> [iC,bS*oH*oW,mC]
|
||||
modifGradO2 = {{3, 0, 1, 2}, {iC, mC, bS * oH * oW}}; // [bS,oH,oW,iC*mC] -> [iC*mC,bS,oH,oW] -> [iC,mC,bS*oH*oW]
|
||||
std::vector<sd::LongType> permuteVec = {0, 3, 1, 2};
|
||||
input = input->permute(permuteVec, false, false); // permute() already returns NDArray*
|
||||
gradI = gradI->permute(permuteVec, false, false);
|
||||
} else {
|
||||
gradOreShape = {bS, iC, mC, oH, oW}; // [bS,iC*mC,oH,oW] -> [bS,iC,mC,oH,oW]
|
||||
modifGradO1 = {{1, 0, 3, 4, 2},
|
||||
{iC, bS * oH * oW, mC}}; // [bS,iC,mC,oH,oW] -> [iC,bS,oH,oW,mC] -> [iC,bS*oH*oW,mC]
|
||||
modifGradO2 = {{1, 0, 2, 3}, {iC, mC, bS * oH * oW}}; // [bS,iC*mC,oH,oW] -> [iC*mC,bS,oH,oW] -> [iC,mC,bS*oH*oW]
|
||||
}
|
||||
|
||||
if (0 == wFormat)
|
||||
modifWeights = {{2, 0, 1, 3}, {iC, kH * kW, mC}};
|
||||
else if (1 == wFormat)
|
||||
modifWeights = {{1, 2, 3, 0}, {iC, kH * kW, mC}};
|
||||
else
|
||||
modifWeights = {{3, 1, 2, 0}, {iC, kH * kW, mC}};
|
||||
|
||||
if (paddingMode == 1) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
std::vector<sd::LongType> colShape = {bS, iC, kH, kW, oH, oW};
|
||||
NDArray columns(input->ordering(), colShape, input->dataType(), input->getContext());
|
||||
NDArray gradOreshaped = gradO->reshape(gradO->ordering(), gradOreShape,false);
|
||||
|
||||
// ----- calculation of gradW and gradB ----- //
|
||||
|
||||
NDArray zero = NDArrayFactory::create(0.f, input->getContext());
|
||||
helpers::im2col(
|
||||
*input->getContext(), *input, columns, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
zero); // [bS, iC, iH, iW] is convoluted to [bS, iC, kH, kW, oH, oW]
|
||||
MmulHelper::tensorDot(&columns, &gradOreshaped, gradW, modifColumns, modifGradO1,
|
||||
modifWeights); // [iC, kW*kH, bS*oH*oW] x [iC, bS*oH*oW, mC] = [iC, kH*kW, mC]
|
||||
|
||||
// ----- calculation of gradB ----- //
|
||||
if (gradB) {
|
||||
NDArray* gradBR = gradB;
|
||||
if (gradB->rankOf() == 2) {
|
||||
std::vector<sd::LongType> lenShape = {gradB->lengthOf()};
|
||||
gradBR = gradB->reshape(gradB->ordering(), lenShape, false);
|
||||
}
|
||||
std::vector<LongType> dims = {0, indOoH, indOoH + 1};
|
||||
gradO->reduceAlongDimension(reduce::Sum, gradBR,&dims, false); // sum over bS, oH, oW
|
||||
if (gradBR != gradB) delete gradBR;
|
||||
}
|
||||
|
||||
//----- calculation of gradI -----//
|
||||
MmulHelper::tensorDot(weights, gradO, &columns, modifWeights, modifGradO2,
|
||||
modifColumns); // [iC, kH*kW, mC] x [iC, mC, bS*oH*oW] = [iC, kW*kH, bS*oH*oW]
|
||||
helpers::col2im(*input->getContext(), &columns, gradI, sH, sW, pH, pW, iH, iW, dH,
|
||||
dW); // [bS, iC, kH, kW, oH, oW] is de-convoluted to [bS, iC, iH, iW]
|
||||
|
||||
if (!isNCHW) {
|
||||
delete input;
|
||||
delete gradI;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::depthwiseConv2dBP(graph::Context& block, NDArray* input, NDArray* weights,
|
||||
NDArray* bias, NDArray* gradO, NDArray* gradI, NDArray* gradW,
|
||||
NDArray* gradB, const LongType kH, const LongType kW, const LongType sH, const LongType sW, LongType pH,
|
||||
LongType pW, const LongType dH, const LongType dW, const int paddingMode, const int isNCHW,
|
||||
const int wFormat) {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input->dataType(), depthwiseConv2dBP_,
|
||||
(input, weights, bias, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,381 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <math/templatemath.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void avgPooling2dCuda(const void *vx, const LongType *xShapeInfo, void *vz, const LongType *zShapeInfo,
|
||||
const LongType kH, const LongType kW, const LongType sH, const LongType sW,
|
||||
const LongType pH, const LongType pW, const LongType dH, const LongType dW,
|
||||
const int extraParam0) {
|
||||
// input is [bS, iC, iH, iW]
|
||||
// output is [bS, iC, oH, oW]
|
||||
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
|
||||
__shared__ LongType bS, iC, oH, oW, iH, iW, strideB, strideC, strideY, strideX, strideOB, strideOC, strideOY,
|
||||
strideOX, length, kHEff, kWEff;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
bS = shape::sizeAt(xShapeInfo, 0);
|
||||
iC = shape::sizeAt(xShapeInfo, 1);
|
||||
oH = shape::sizeAt(zShapeInfo, 2);
|
||||
oW = shape::sizeAt(zShapeInfo, 3);
|
||||
iH = shape::sizeAt(xShapeInfo, 2);
|
||||
iW = shape::sizeAt(xShapeInfo, 3);
|
||||
|
||||
strideB = shape::stride(xShapeInfo)[0];
|
||||
strideC = shape::stride(xShapeInfo)[1];
|
||||
strideY = shape::stride(xShapeInfo)[2];
|
||||
strideX = shape::stride(xShapeInfo)[3];
|
||||
|
||||
strideOB = shape::stride(zShapeInfo)[0];
|
||||
strideOC = shape::stride(zShapeInfo)[1];
|
||||
strideOY = shape::stride(zShapeInfo)[2];
|
||||
strideOX = shape::stride(zShapeInfo)[3];
|
||||
|
||||
length = shape::length(zShapeInfo);
|
||||
|
||||
// Replace kernel H/W with *effective* kernel H/W accounting for dilation
|
||||
kHEff = kH + (kH - 1) * (dH - 1);
|
||||
kWEff = kW + (kW - 1) * (dW - 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int index = tid; index < length; index += blockDim.x * gridDim.x) {
|
||||
const LongType pw = index % oW;
|
||||
const LongType ph = (index / oW) % oH;
|
||||
const LongType c = (index / oW / oH) % iC;
|
||||
const LongType n = index / oW / oH / iC;
|
||||
|
||||
LongType hstart = sH * ph - pH;
|
||||
LongType wstart = sW * pw - pW;
|
||||
LongType hend = hstart + kHEff;
|
||||
LongType wend = wstart + kWEff;
|
||||
|
||||
if (hstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-hstart) / static_cast<double>(dH));
|
||||
hstart += f * dH;
|
||||
}
|
||||
if (wstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-wstart) / static_cast<double>(dW));
|
||||
wstart += f * dW;
|
||||
}
|
||||
if (hend > iH) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(hend - iH) / static_cast<double>(dH));
|
||||
hend -= f * dH;
|
||||
}
|
||||
if (wend > iW) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(wend - iW) / static_cast<double>(dW));
|
||||
wend -= f * dW;
|
||||
}
|
||||
|
||||
// Accounts for dilation
|
||||
int pool_size = sd::math::sd_ceil<double, int>(static_cast<double>(hend - hstart) / static_cast<double>(dH)) *
|
||||
sd::math::sd_ceil<double, int>(static_cast<double>(wend - wstart) / static_cast<double>(dW));
|
||||
|
||||
Z sum = static_cast<Z>(0.0f);
|
||||
|
||||
const X *inSlice = x + (n * strideB + c * strideC);
|
||||
|
||||
for (int h = hstart; h < hend; h += dH)
|
||||
for (int w = wstart; w < wend; w += dW) sum += static_cast<Z>(inSlice[h * strideY + w * strideX]);
|
||||
|
||||
int divide_factor = pool_size; // Case 0: exclude padding
|
||||
if (extraParam0 == 1) // Case 1: include padding
|
||||
divide_factor = kH * kW;
|
||||
|
||||
z[n * strideOB + c * strideOC + pw * strideOX + ph * strideOY] = sum / static_cast<Z>(divide_factor);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static void avgPooling2dCudaLauncher(LaunchContext &block, const void *vx, const LongType *vxShapeInfo, void *vz,
|
||||
const LongType *vzShapeInfo, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, const LongType pH, const LongType pW,
|
||||
const LongType dH, const LongType dW, const int extraParam0) {
|
||||
dim3 launchDims = getLaunchDims("avg_pooling");
|
||||
avgPooling2dCuda<X, Z><<<launchDims.y, launchDims.x, launchDims.z, *block.getCudaStream()>>>(
|
||||
vx, vxShapeInfo, vz, vzShapeInfo, kH, kW, sH, sW, pH, pW, dH, dW, extraParam0);
|
||||
DebugHelper::checkErrorCode(block.getCudaStream(), "avgb pooling 2d failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void pnormPooling2dCuda(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, const LongType pH, const LongType pW,
|
||||
const LongType dH, const LongType dW, const int extraParam0) {
|
||||
// input is [bS, iC, iH, iW]
|
||||
// output is [bS, iC, oH, oW]
|
||||
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
|
||||
__shared__ LongType bS, iC, oH, oW, iH, iW, strideB, strideC, strideY, strideX, strideOB, strideOC, strideOY,
|
||||
strideOX, length, kHEff, kWEff;
|
||||
__shared__ bool fOrder;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
bS = shape::sizeAt(xShapeInfo, 0);
|
||||
iC = shape::sizeAt(xShapeInfo, 1);
|
||||
oH = shape::sizeAt(zShapeInfo, 2);
|
||||
oW = shape::sizeAt(zShapeInfo, 3);
|
||||
iH = shape::sizeAt(xShapeInfo, 2);
|
||||
iW = shape::sizeAt(xShapeInfo, 3);
|
||||
|
||||
strideB = shape::stride(xShapeInfo)[0];
|
||||
strideC = shape::stride(xShapeInfo)[1];
|
||||
strideY = shape::stride(xShapeInfo)[2];
|
||||
strideX = shape::stride(xShapeInfo)[3];
|
||||
|
||||
strideOB = shape::stride(zShapeInfo)[0];
|
||||
strideOC = shape::stride(zShapeInfo)[1];
|
||||
strideOY = shape::stride(zShapeInfo)[2];
|
||||
strideOX = shape::stride(zShapeInfo)[3];
|
||||
|
||||
length = shape::length(zShapeInfo);
|
||||
|
||||
// Replace kernel H/W with *effective* kernel H/W accounting for dilation
|
||||
kHEff = kH + (kH - 1) * (dH - 1);
|
||||
kWEff = kW + (kW - 1) * (dW - 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int index = tid; index < length; index += blockDim.x * gridDim.x) {
|
||||
const LongType pw = index % oW;
|
||||
const LongType ph = (index / oW) % oH;
|
||||
const LongType c = (index / oW / oH) % iC;
|
||||
const LongType n = index / oW / oH / iC;
|
||||
|
||||
LongType hstart = sH * ph - pH;
|
||||
LongType wstart = sW * pw - pW;
|
||||
LongType hend = hstart + kHEff;
|
||||
LongType wend = wstart + kWEff;
|
||||
|
||||
if (hstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-hstart) / static_cast<double>(dH));
|
||||
hstart += f * dH;
|
||||
}
|
||||
if (wstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-wstart) / static_cast<double>(dW));
|
||||
wstart += f * dW;
|
||||
}
|
||||
if (hend > iH) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(hend - iH) / static_cast<double>(dH));
|
||||
hend -= f * dH;
|
||||
}
|
||||
if (wend > iW) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(wend - iW) / static_cast<double>(dW));
|
||||
wend -= f * dW;
|
||||
}
|
||||
|
||||
Z sum = static_cast<Z>(0.0f);
|
||||
|
||||
const X *inSlice = x + (n * strideB + c * strideC);
|
||||
|
||||
for (int h = hstart; h < hend; h += dH)
|
||||
for (int w = wstart; w < wend; w += dW)
|
||||
sum += math::sd_pow<Z, Z, Z>(static_cast<Z>(math::sd_abs<X,X>(inSlice[h * strideY + w * strideX])), static_cast<Z>(extraParam0));
|
||||
|
||||
z[n * strideOB + c * strideOC + pw * strideOX + ph * strideOY] = math::sd_pow<Z, Z, Z>(sum, static_cast<Z>(1.0f) / static_cast<Z>(extraParam0));
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static void pnormPooling2dCudaLauncher(LaunchContext &block, const void *vx, const LongType *vxShapeInfo, void *vz,
|
||||
const LongType *vzShapeInfo, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, const LongType pH, const LongType pW,
|
||||
const LongType dH, const LongType dW, const int extraParam0) {
|
||||
dim3 launchDims = getLaunchDims("avg_pooling");
|
||||
pnormPooling2dCuda<X, Z><<<launchDims.y, launchDims.x, launchDims.z, *block.getCudaStream()>>>(
|
||||
vx, vxShapeInfo, vz, vzShapeInfo, kH, kW, sH, sW, pH, pW, dH, dW, extraParam0);
|
||||
DebugHelper::checkErrorCode(block.getCudaStream(), "pnorm pooling 2d failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void maxPooling2dCuda(const void *vx, const LongType *xShapeInfo, void *vz, const LongType *zShapeInfo,
|
||||
const int kH, const LongType kW, const LongType sH, const LongType sW,
|
||||
const LongType pH, const LongType pW, const LongType dH, const LongType dW,
|
||||
const int extraParam0) {
|
||||
// input is [bS, iC, iH, iW]
|
||||
// output is [bS, iC, oH, oW]
|
||||
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
|
||||
__shared__ LongType bS, iC, oH, oW, iH, iW, strideB, strideC, strideY, strideX, strideOB, strideOC, strideOY,
|
||||
strideOX, length, kHEff, kWEff;
|
||||
__shared__ bool fOrder;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
bS = shape::sizeAt(xShapeInfo, 0);
|
||||
iC = shape::sizeAt(xShapeInfo, 1);
|
||||
oH = shape::sizeAt(zShapeInfo, 2);
|
||||
oW = shape::sizeAt(zShapeInfo, 3);
|
||||
iH = shape::sizeAt(xShapeInfo, 2);
|
||||
iW = shape::sizeAt(xShapeInfo, 3);
|
||||
|
||||
strideB = shape::stride(xShapeInfo)[0];
|
||||
strideC = shape::stride(xShapeInfo)[1];
|
||||
strideY = shape::stride(xShapeInfo)[2];
|
||||
strideX = shape::stride(xShapeInfo)[3];
|
||||
|
||||
strideOB = shape::stride(zShapeInfo)[0];
|
||||
strideOC = shape::stride(zShapeInfo)[1];
|
||||
strideOY = shape::stride(zShapeInfo)[2];
|
||||
strideOX = shape::stride(zShapeInfo)[3];
|
||||
|
||||
length = shape::length(zShapeInfo);
|
||||
|
||||
// Replace kernel H/W with *effective* kernel H/W accounting for dilation
|
||||
kHEff = kH + (kH - 1) * (dH - 1);
|
||||
kWEff = kW + (kW - 1) * (dW - 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (int index = tid; index < length; index += blockDim.x * gridDim.x) {
|
||||
const LongType pw = index % oW;
|
||||
const LongType ph = (index / oW) % oH;
|
||||
const LongType c = (index / oW / oH) % iC;
|
||||
const LongType n = index / oW / oH / iC;
|
||||
|
||||
LongType hstart = sH * ph - pH;
|
||||
LongType wstart = sW * pw - pW;
|
||||
LongType hend = hstart + kHEff;
|
||||
LongType wend = wstart + kWEff;
|
||||
|
||||
if (hstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-hstart) / static_cast<double>(dH));
|
||||
hstart += f * dH;
|
||||
}
|
||||
if (wstart < 0) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(-wstart) / static_cast<double>(dW));
|
||||
wstart += f * dW;
|
||||
}
|
||||
if (hend > iH) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(hend - iH) / static_cast<double>(dH));
|
||||
hend -= f * dH;
|
||||
}
|
||||
if (wend > iW) {
|
||||
int f = math::sd_ceil<double, int>(static_cast<double>(wend - iW) / static_cast<double>(dW));
|
||||
wend -= f * dW;
|
||||
}
|
||||
// Accounts for dilation
|
||||
int pool_size = sd::math::sd_ceil<double, int>(static_cast<double>(hend - hstart) / static_cast<double>(dH)) *
|
||||
sd::math::sd_ceil<double, int>(static_cast<double>(wend - wstart) / static_cast<double>(dW));
|
||||
|
||||
Z max = -DataTypeUtils::max<Z>();
|
||||
|
||||
const X *inSlice = x + (n * strideB + c * strideC);
|
||||
|
||||
for (int h = hstart; h < hend; h += dH) {
|
||||
for (int w = wstart; w < wend; w += dW) {
|
||||
Z v = static_cast<Z>(inSlice[h * strideY + w * strideX]);
|
||||
if (v > max) max = v;
|
||||
}
|
||||
}
|
||||
|
||||
z[n * strideOB + c * strideOC + pw * strideOX + ph * strideOY] = max;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static void maxPooling2dCudaLauncher(LaunchContext &block, const void *vx, const LongType *vxShapeInfo, void *vz,
|
||||
const LongType *vzShapeInfo, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, const LongType pH, const LongType pW,
|
||||
const LongType dH, const LongType dW, const int extraParam0, const int rank,
|
||||
const int len) {
|
||||
dim3 poolingDims = getPoolingDims(len, rank);
|
||||
maxPooling2dCuda<X, Z><<<poolingDims.y, poolingDims.x, poolingDims.z, *block.getCudaStream()>>>(
|
||||
vx, vxShapeInfo, vz, vzShapeInfo, kH, kW, sH, sW, pH, pW, dH, dW, extraParam0);
|
||||
DebugHelper::checkErrorCode(block.getCudaStream(), "max pooling 2d failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::pooling2d(graph::Context &block, NDArray&input, NDArray &output, const LongType kH,
|
||||
const LongType kW, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW, const PoolingType poolingMode,
|
||||
const int extraParam0) {
|
||||
if (!input.isActualOnDeviceSide()) input.syncToDevice();
|
||||
|
||||
switch (poolingMode) {
|
||||
case MAX_POOL: {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input.dataType(), maxPooling2dCudaLauncher,
|
||||
(*block.launchContext(), input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), kH, kW, sH, sW, pH, pW, dH, dW, extraParam0, output.rankOf(), output.lengthOf()),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
} break;
|
||||
case AVG_POOL: {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input.dataType(), avgPooling2dCudaLauncher,
|
||||
(*block.launchContext(), input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), kH, kW, sH, sW, pH, pW, dH, dW, extraParam0),
|
||||
SD_NUMERIC_TYPES);
|
||||
} break;
|
||||
case PNORM_POOL: {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input.dataType(), pnormPooling2dCudaLauncher,
|
||||
(*block.launchContext(), input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), kH, kW, sH, sW, pH, pW, dH, dW, extraParam0),
|
||||
SD_FLOAT_TYPES);
|
||||
} break;
|
||||
default:
|
||||
THROW_EXCEPTION("Pooling2D: Unknown PoolingType used");
|
||||
}
|
||||
|
||||
output.tickWriteDevice();
|
||||
input.tickReadDevice();
|
||||
|
||||
auto result = cudaStreamSynchronize(*block.launchContext()->getCudaStream());
|
||||
if (result != 0) throw cuda_exception::build("Pooling2D failed", result);
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,214 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <math/templatemath.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void pooling2dBPCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType kH, const LongType kW, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW, const int poolingMode,
|
||||
const int extraParam0) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
const T* y = reinterpret_cast<const T*>(vy);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
LongType coord2, coord3;
|
||||
__shared__ LongType rank, kHeff, kWeff, iH, iW, kProd;
|
||||
__shared__ LongType xLen, yLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* yShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* yStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
yLen = shape::length(yShapeInfo);
|
||||
xLen = shape::length(xShapeInfo);
|
||||
rank = 4;
|
||||
|
||||
kHeff = kH + (kH - 1) * (dH - 1);
|
||||
kWeff = kW + (kW - 1) * (dW - 1);
|
||||
|
||||
iH = xShapeInfo[3];
|
||||
iW = xShapeInfo[4];
|
||||
|
||||
kProd = kH * kW;
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto yInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (yInd >= yLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(yInd, rank, yShape, coords);
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(rank, yStride, coords, yOffset);
|
||||
|
||||
LongType hstart = coords[2] * sH - pH;
|
||||
LongType wstart = coords[3] * sW - pW;
|
||||
LongType hend = hstart + kHeff;
|
||||
LongType wend = wstart + kWeff;
|
||||
if (hstart < 0) hstart += dH * ((-hstart + dH - 1) / dH);
|
||||
if (wstart < 0) wstart += dW * ((-wstart + dW - 1) / dW);
|
||||
if (hend > iH) hend -= dH * ((hend - iH + dH - 1) / dH);
|
||||
if (wend > iW) wend -= dW * ((wend - iW + dW - 1) / dW);
|
||||
|
||||
switch (poolingMode) {
|
||||
/*** max ***/
|
||||
case 0: {
|
||||
coord2 = hstart;
|
||||
coord3 = wstart;
|
||||
bool out_of_range = false;
|
||||
|
||||
T max = -DataTypeUtils::max<T>();
|
||||
for (coords[2] = hstart; coords[2] < hend; coords[2] += dH) {
|
||||
for (coords[3] = wstart; coords[3] < wend; coords[3] += dW) {
|
||||
LongType offset;
|
||||
COORDS2INDEX(rank, xStride, coords, offset);
|
||||
T val = x[offset];
|
||||
if (val > max) {
|
||||
max = val;
|
||||
coord2 = coords[2];
|
||||
coord3 = coords[3];
|
||||
}
|
||||
}
|
||||
}
|
||||
coords[2] = coord2;
|
||||
coords[3] = coord3;
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(&z[zOffset], y[yOffset]);
|
||||
} break;
|
||||
|
||||
/*** avg ***/
|
||||
case 1: {
|
||||
T val = y[yOffset];
|
||||
|
||||
if (extraParam0 == 0) // Exclude padding
|
||||
val /= static_cast<T>(math::sd_ceil<double, int>(static_cast<double>(hend - hstart) / static_cast<double>(dH)) *
|
||||
math::sd_ceil<double, int>(static_cast<double>(wend - wstart) / static_cast<double>(dW)));
|
||||
else if (extraParam0 == 1) // Include padding
|
||||
val /= static_cast<T>(kProd);
|
||||
|
||||
for (coords[2] = hstart; coords[2] < hend; coords[2] += dH)
|
||||
for (coords[3] = wstart; coords[3] < wend; coords[3] += dW) {
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(&z[zOffset], val);
|
||||
}
|
||||
} break;
|
||||
|
||||
/*** pnorm ***/
|
||||
case 2: {
|
||||
T sum = static_cast<T>(0.);
|
||||
T val = y[yOffset];
|
||||
|
||||
for (coords[2] = hstart; coords[2] < hend; coords[2] += dH)
|
||||
for (coords[3] = wstart; coords[3] < wend; coords[3] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
sum += math::sd_pow<T, T, T>(math::sd_abs<T,T>(x[xOffset]), static_cast<T>(extraParam0));
|
||||
}
|
||||
|
||||
val *= math::sd_pow<T, T, T>(sum, (static_cast<T>(1.0f) - static_cast<T>(extraParam0)) / static_cast<T>(extraParam0));
|
||||
|
||||
for (coords[2] = hstart; coords[2] < hend; coords[2] += dH) {
|
||||
for (coords[3] = wstart; coords[3] < wend; coords[3] += dW) {
|
||||
LongType xOffset, zOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(
|
||||
&z[zOffset], val * math::sd_pow<T, T, T>(math::sd_abs<T,T>(x[xOffset]), static_cast<T>(extraParam0) - static_cast<T>(1.0f)) *
|
||||
math::sd_sgn<T, T>(x[xOffset]));
|
||||
}
|
||||
}
|
||||
} break;
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void pooling2dBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType kH, const LongType kW, const LongType sH,
|
||||
const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW,
|
||||
const int poolingMode, const int extraParam0) {
|
||||
pooling2dBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
|
||||
vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, kH, kW, sH, sW, pH, pW, dH, dW, poolingMode, extraParam0);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"pooling2dBPCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::pooling2dBP(graph::Context& block, NDArray& input, NDArray& gradO,
|
||||
NDArray& gradI, const LongType kH, const LongType kW, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW, const int poolingMode,
|
||||
const int extraParam0) {
|
||||
// initial zeroing of gradI
|
||||
gradI.nullify();
|
||||
|
||||
PointersManager manager(block.launchContext(), "pooling2dBP");
|
||||
|
||||
auto inputBuff = input.specialBuffer();
|
||||
dim3 poolingDims = getPoolingDims(gradO.lengthOf(),gradO.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&gradI}, {&input, &gradO});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), pooling2dBPCudaLauncher,
|
||||
(poolingDims.x, poolingDims.y, poolingDims.z, block.launchContext()->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), gradO.specialBuffer(), gradO.specialShapeInfo(), gradI.specialBuffer(),
|
||||
gradI.specialShapeInfo(), kH, kW, sH, sW, pH, pW, dH, dW, poolingMode, extraParam0),
|
||||
SD_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({&gradI}, {&input, &gradO});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,195 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <math/templatemath.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void pooling3dCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int kD, const int kH, const int kW,
|
||||
const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
|
||||
const int dD, const int dH, const int dW, const int poolingMode,
|
||||
const int extraParam0) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank, kDeff, kHeff, kWeff, iD, iH, iW, kProd;
|
||||
__shared__ LongType zLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = 5;
|
||||
|
||||
kDeff = kD + (kD - 1) * (dD - 1);
|
||||
kHeff = kH + (kH - 1) * (dH - 1);
|
||||
kWeff = kW + (kW - 1) * (dW - 1);
|
||||
|
||||
iD = xShapeInfo[3];
|
||||
iH = xShapeInfo[4];
|
||||
iW = xShapeInfo[5];
|
||||
|
||||
kProd = kD * kH * kW;
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(zInd, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
int dstart = coords[2] * sD - pD;
|
||||
int hstart = coords[3] * sH - pH;
|
||||
int wstart = coords[4] * sW - pW;
|
||||
int dend = dstart + kDeff;
|
||||
int hend = hstart + kHeff;
|
||||
int wend = wstart + kWeff;
|
||||
|
||||
if (dstart < 0) dstart += dD * ((-dstart + dD - 1) / dD);
|
||||
if (hstart < 0) hstart += dH * ((-hstart + dH - 1) / dH);
|
||||
if (wstart < 0) wstart += dW * ((-wstart + dW - 1) / dW);
|
||||
if (dend > iD) dend -= dD * ((dend - iD + dD - 1) / dD);
|
||||
if (hend > iH) hend -= dH * ((hend - iH + dH - 1) / dH);
|
||||
if (wend > iW) wend -= dW * ((wend - iW + dW - 1) / dW);
|
||||
|
||||
switch (poolingMode) {
|
||||
/*** max ***/
|
||||
case 0: {
|
||||
T max = -DataTypeUtils::max<T>();
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD) {
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH) {
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
T val = x[xOffset];
|
||||
if (val > max) max = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
z[zOffset] = max;
|
||||
} break;
|
||||
|
||||
/*** avg ***/
|
||||
case 1: {
|
||||
T sum = static_cast<T>(0.);
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD)
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH)
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
sum += x[xOffset];
|
||||
}
|
||||
|
||||
if (extraParam0 == 0) { // Exclude padding
|
||||
LongType a = (dend - dstart) / dD + ((dend - dstart) % dD == 0 ? 0 : 1);
|
||||
LongType b = (hend - hstart) / dH + ((hend - hstart) % dH == 0 ? 0 : 1);
|
||||
LongType c = (wend - wstart) / dW + ((wend - wstart) % dW == 0 ? 0 : 1);
|
||||
sum /= static_cast<T>(a * b * c); // Accounts for dilation
|
||||
} else if (extraParam0 == 1) // Include padding
|
||||
sum /= kProd;
|
||||
|
||||
z[zOffset] = sum;
|
||||
} break;
|
||||
|
||||
/*** pnorm ***/
|
||||
case 2: {
|
||||
T sum = static_cast<T>(0.);
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD)
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH)
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
sum += math::sd_pow<T, T, T>(math::sd_abs<T, T>(x[xOffset]), static_cast<T>(extraParam0));
|
||||
}
|
||||
|
||||
sum = math::sd_pow<T, T, T>(sum, static_cast<T>(1.0f) / static_cast<T>(extraParam0));
|
||||
|
||||
z[zOffset] = sum;
|
||||
} break;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void pooling3dCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int kD, const int kH, const int kW,
|
||||
const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
|
||||
const int dD, const int dH, const int dW, const int poolingMode,
|
||||
const int extraParam0) {
|
||||
pooling3dCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
|
||||
vx, xShapeInfo, vz, zShapeInfo, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, poolingMode, extraParam0);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"pooling3dBPCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::pooling3d(graph::Context& block, NDArray& input, NDArray& output, const LongType kD,
|
||||
const LongType kH, const LongType kW, const LongType sD, const LongType sH, const LongType sW, const LongType pD,
|
||||
const LongType pH, const LongType pW, const LongType dD, const LongType dH, const LongType dW,
|
||||
const int poolingMode, const int extraParam0) {
|
||||
PointersManager manager(block.launchContext(), "pooling3d");
|
||||
|
||||
dim3 poolingDims = getPoolingDims(output.lengthOf(),output.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), pooling3dCudaLauncher,
|
||||
(poolingDims.x, poolingDims.y, poolingDims.z, block.launchContext()->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), kD, kH, kW, sD, sH, sW, pD, pH, pW,
|
||||
dD, dH, dW, poolingMode, extraParam0),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,222 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <math/templatemath.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void pooling3dBPCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const int kD, const int kH, const int kW, const int sD, const int sH,
|
||||
const int sW, const int pD, const int pH, const int pW, const int dD,
|
||||
const int dH, const int dW, const int poolingMode, const int extraParam0) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
const T* y = reinterpret_cast<const T*>(vy);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
LongType coord2, coord3, coord4;
|
||||
__shared__ int rank, kDeff, kHeff, kWeff, iD, iH, iW, kProd;
|
||||
__shared__ LongType yLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* yShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* yStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
yLen = shape::length(yShapeInfo);
|
||||
rank = 5;
|
||||
|
||||
kDeff = kD + (kD - 1) * (dD - 1);
|
||||
kHeff = kH + (kH - 1) * (dH - 1);
|
||||
kWeff = kW + (kW - 1) * (dW - 1);
|
||||
|
||||
iD = xShapeInfo[3];
|
||||
iH = xShapeInfo[4];
|
||||
iW = xShapeInfo[5];
|
||||
|
||||
kProd = kD * kH * kW;
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto yInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (yInd >= yLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(yInd, rank, yShape, coords);
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(rank, yStride, coords, yOffset);
|
||||
|
||||
int dstart = coords[2] * sD - pD;
|
||||
int hstart = coords[3] * sH - pH;
|
||||
int wstart = coords[4] * sW - pW;
|
||||
int dend = dstart + kDeff;
|
||||
int hend = hstart + kHeff;
|
||||
int wend = wstart + kWeff;
|
||||
|
||||
if (dstart < 0) dstart += dD * ((-dstart + dD - 1) / dD);
|
||||
if (hstart < 0) hstart += dH * ((-hstart + dH - 1) / dH);
|
||||
if (wstart < 0) wstart += dW * ((-wstart + dW - 1) / dW);
|
||||
if (dend > iD) dend -= dD * ((dend - iD + dD - 1) / dD);
|
||||
if (hend > iH) hend -= dH * ((hend - iH + dH - 1) / dH);
|
||||
if (wend > iW) wend -= dW * ((wend - iW + dW - 1) / dW);
|
||||
|
||||
switch (poolingMode) {
|
||||
/*** max ***/
|
||||
case 0: {
|
||||
T max = -DataTypeUtils::max<T>();
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD) {
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH) {
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
T val = x[xOffset];
|
||||
if (val > max) {
|
||||
max = val;
|
||||
coord2 = coords[2];
|
||||
coord3 = coords[3];
|
||||
coord4 = coords[4];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
coords[2] = coord2;
|
||||
coords[3] = coord3;
|
||||
coords[4] = coord4;
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(&z[zOffset], y[yOffset]);
|
||||
} break;
|
||||
|
||||
/*** avg ***/
|
||||
case 1: {
|
||||
T val = y[yOffset];
|
||||
|
||||
if (extraParam0 == 0) // Exclude padding
|
||||
val /= static_cast<T>(math::sd_ceil<double, int>(static_cast<double>(dend - dstart) / static_cast<double>(dD)) *
|
||||
math::sd_ceil<double, int>(static_cast<double>(hend - hstart) / static_cast<double>(dH)) *
|
||||
math::sd_ceil<double, int>(static_cast<double>(wend - wstart) / static_cast<double>(dW)));
|
||||
else if (extraParam0 == 1) // Include padding
|
||||
val /= static_cast<T>(kProd);
|
||||
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD)
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH)
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(&z[zOffset], val);
|
||||
}
|
||||
} break;
|
||||
|
||||
/*** pnorm ***/
|
||||
case 2: {
|
||||
T sum = static_cast<T>(0.);
|
||||
T val = y[yOffset];
|
||||
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD)
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH)
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
sum += math::sd_pow<T, T, T>(math::sd_abs<T,T>(x[xOffset]), static_cast<T>(extraParam0));
|
||||
}
|
||||
|
||||
val *= math::sd_pow<T, T, T>(sum, (static_cast<T>(1.0f) - static_cast<T>(extraParam0)) / static_cast<T>(extraParam0));
|
||||
|
||||
for (coords[2] = dstart; coords[2] < dend; coords[2] += dD)
|
||||
for (coords[3] = hstart; coords[3] < hend; coords[3] += dH)
|
||||
for (coords[4] = wstart; coords[4] < wend; coords[4] += dW) {
|
||||
LongType xOffset, zOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
math::atomics::sd_atomicAdd<T>(
|
||||
&z[zOffset], val * math::sd_pow<T, T, T>(math::sd_abs<T,T>(x[xOffset]), static_cast<T>(extraParam0) - static_cast<T>(1.0f)) *
|
||||
math::sd_sgn<T, T>(x[xOffset]));
|
||||
}
|
||||
} break;
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void pooling3dBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int kD, const int kH, const int kW,
|
||||
const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
|
||||
const int dD, const int dH, const int dW, const int poolingMode,
|
||||
const int extraParam0) {
|
||||
pooling3dBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz,
|
||||
zShapeInfo, kD, kH, kW, sD, sH, sW, pD, pH,
|
||||
pW, dD, dH, dW, poolingMode, extraParam0);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"pooling3dBPCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::pooling3dBP(graph::Context& block, NDArray& input, NDArray& gradO,
|
||||
NDArray& gradI, const LongType kD, const LongType kH, const LongType kW, const LongType sD, const LongType sH,
|
||||
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
|
||||
const LongType dW, const int poolingMode, const int extraParam0) {
|
||||
// initial zeroing of gradI
|
||||
gradI.nullify();
|
||||
|
||||
PointersManager manager(block.launchContext(), "pooling3dBP");
|
||||
dim3 poolingDims = getPoolingDims(gradO.lengthOf(),gradO.rankOf());
|
||||
NDArray::prepareSpecialUse({&gradI}, {&input, &gradO});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), pooling3dBPCudaLauncher,
|
||||
(poolingDims.y, poolingDims.x, poolingDims.z, block.launchContext()->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), gradO.specialBuffer(), gradO.specialShapeInfo(), gradI.specialBuffer(),
|
||||
gradI.specialShapeInfo(), kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, poolingMode, extraParam0),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&gradI}, {&input, &gradO});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace s
|
||||
@@ -0,0 +1,90 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void sconv2d_(graph::Context& block, NDArray* input, NDArray* weightsDepth,
|
||||
NDArray* weightsPoint, NDArray* bias, NDArray* output, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW, const int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
// input [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
|
||||
// weightsDepth [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
|
||||
// weightsPoint [1, 1, iC*mC, oC], [oC, iC*mC, 1, 1], [oC, 1, 1, iC*mC]
|
||||
// bias [oC], oC = iC*mC if weightsPoint=nullptr
|
||||
// output is [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
|
||||
|
||||
// kH filter(kernel) height
|
||||
// kW filter(kernel) width
|
||||
// sH strides height
|
||||
// sW strides width
|
||||
// pH paddings height
|
||||
// pW paddings width
|
||||
// dH dilations height
|
||||
// dW dilations width
|
||||
// paddingMode 0-VALID, 1-SAME
|
||||
// isNCHW 1-NCHW, 0-NHWC
|
||||
|
||||
LongType bS, iC, iH, iW, mC, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, channels multiplier, output channels, output height/width
|
||||
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWiC, indWmC, indWkH, indOoH);
|
||||
mC = weightsDepth->sizeAt(indWmC); // channels multiplier
|
||||
|
||||
NDArray* outputDepth = output;
|
||||
if (weightsPoint) { // if pointwise convolution is expected
|
||||
std::vector<sd::LongType> shape1 = {bS, oH, oW, iC * mC};
|
||||
std::vector<sd::LongType> shape2 = {bS, iC * mC, oH, oW};
|
||||
outputDepth = new NDArray(output->ordering(), !isNCHW ? shape1 : shape2, input->dataType(), input->getContext());
|
||||
}
|
||||
// ----- perform depthwise convolution (if weightsPoint is absent then oC = iC*mC) ----- //
|
||||
ConvolutionUtils::depthwiseConv2d(block, input, weightsDepth, weightsPoint ? nullptr : bias, outputDepth, kH, kW, sH,
|
||||
sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat);
|
||||
|
||||
// ----- perform pointwise convolution (oH = iH, oW = iW) ----- //
|
||||
if (weightsPoint) {
|
||||
ConvolutionUtils::conv2d(block, outputDepth, weightsPoint, bias, output, 1, 1, 1, 1, 0, 0, 1, 1, paddingMode,
|
||||
isNCHW, wFormat); // in this case oH=iH, oW=iW
|
||||
delete outputDepth;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::sconv2d(graph::Context& block, NDArray* input, NDArray* weightsDepth,
|
||||
NDArray* weightsPoint, NDArray* bias, NDArray* output, const LongType kH,
|
||||
const LongType kW, const LongType sH, const LongType sW, LongType pH, LongType pW, const LongType dH, const LongType dW,
|
||||
const int paddingMode, const int isNCHW, const int wFormat) {
|
||||
BUILD_SINGLE_SELECTOR_TWICE(input->dataType(), sconv2d_,
|
||||
(block, input, weightsDepth, weightsPoint, bias, output, kH, kW, sH, sW, pH, pW, dH, dW,
|
||||
paddingMode, isNCHW, wFormat),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,114 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void upsampling2dCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType factorH, const LongType factorW,
|
||||
const bool isNCHW) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType rank, dimIH;
|
||||
__shared__ LongType zLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
dimIH = isNCHW ? 2 : 1;
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = 4;
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(zInd, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
coords[dimIH] /= factorH;
|
||||
coords[dimIH + 1] /= factorW;
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void upsampling2dCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType factorH, const LongType factorW,
|
||||
const bool isNCHW) {
|
||||
upsampling2dCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, factorH,
|
||||
factorW, isNCHW);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"upsampling2dCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::upsampling2d(graph::Context& block, NDArray& input, NDArray& output, const LongType factorH,
|
||||
const LongType factorW, const bool isNCHW) {
|
||||
PointersManager manager(block.launchContext(), "upsampling2d");
|
||||
|
||||
dim3 getUpSampling = getUpsamplingDims(output.lengthOf(),output.rankOf());
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), upsampling2dCudaLauncher,
|
||||
(getUpSampling.x, getUpSampling.y, getUpSampling.z, block.launchContext()->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), factorH, factorW, isNCHW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,120 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void upsampling2dBPCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const bool isNCHW) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType rank, dimIH;
|
||||
__shared__ LongType factorH, factorW;
|
||||
__shared__ LongType zLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
dimIH = isNCHW ? 2 : 1;
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = 4;
|
||||
|
||||
factorH = xShapeInfo[dimIH + 1] / zShapeInfo[dimIH + 1];
|
||||
factorW = xShapeInfo[dimIH + 2] / zShapeInfo[dimIH + 2];
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(zInd, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
z[zOffset] = 0;
|
||||
|
||||
const LongType zCoord2 = coords[dimIH] * factorH;
|
||||
const LongType zCoord3 = coords[dimIH + 1] * factorW;
|
||||
|
||||
for (coords[dimIH] = zCoord2; coords[dimIH] < zCoord2 + factorH; ++coords[dimIH])
|
||||
for (coords[dimIH + 1] = zCoord3; coords[dimIH + 1] < zCoord3 + factorW; ++coords[dimIH + 1]) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
z[zOffset] += x[xOffset];
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void upsampling2dBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const bool isNCHW) {
|
||||
upsampling2dBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, isNCHW);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"upsampling2dBPCuda failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::upsampling2dBP(graph::Context& block, NDArray& gradO, NDArray& gradI,
|
||||
const bool isNCHW) {
|
||||
PointersManager manager(block.launchContext(), "upsampling2d_bp");
|
||||
|
||||
dim3 getUpSampling = getUpsamplingDims(gradI.lengthOf(),gradI.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&gradI}, {&gradO});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
gradI.dataType(), upsampling2dBPCudaLauncher,
|
||||
(getUpSampling.x, getUpSampling.y, getUpSampling.z, block.launchContext()->getCudaStream(), gradO.specialBuffer(),
|
||||
gradO.specialShapeInfo(), gradI.specialBuffer(), gradI.specialShapeInfo(), isNCHW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&gradI}, {&gradO});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,118 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void upsampling3dCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int factorD, const int factorH,
|
||||
const int factorW, const bool isNCDHW) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank, dimID;
|
||||
__shared__ LongType zLen, *sharedMem;
|
||||
__shared__ LongType* xShape;
|
||||
__shared__ LongType* zShape;
|
||||
__shared__ LongType* xStride;
|
||||
__shared__ LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
dimID = isNCDHW ? 2 : 1;
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = 5;
|
||||
|
||||
// Cache shape information
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(zInd, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
coords[dimID] /= factorD;
|
||||
coords[dimID + 1] /= factorH;
|
||||
coords[dimID + 2] /= factorW;
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void upsampling3dCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const int factorD, const int factorH,
|
||||
const int factorW, const bool isNCDHW) {
|
||||
upsampling3dCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, factorD,
|
||||
factorH, factorW, isNCDHW);
|
||||
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"upsampling3dCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::upsampling3d(graph::Context& block, NDArray& input, NDArray& output, const LongType factorD,
|
||||
const LongType factorH, const LongType factorW, const bool isNCDHW) {
|
||||
PointersManager manager(block.launchContext(), "upsampling3d");
|
||||
|
||||
|
||||
dim3 getUpSampling = getUpsamplingDims(output.lengthOf(),output.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), upsampling3dCudaLauncher,
|
||||
(getUpSampling.x, getUpSampling.y, getUpSampling.z, block.launchContext()->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), factorD, factorH, factorW, isNCDHW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,128 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void upsampling3dBPCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const bool isNCDHW) {
|
||||
// x (gradO) has shape [bS, iC, iD, iH, iW] (NCDHW) or [bS, iD, iH, iW, iC] (NDHWC)
|
||||
// z (gradI) has shape [bS, iC, factorD*iD, factorH*iH, factorW*iW ] (NCDHW) or [bS, factorD*iD, factorH*iH,
|
||||
// factorW*iW, iC] (NDHWC)
|
||||
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank, dimID;
|
||||
__shared__ LongType factorD, factorH, factorW;
|
||||
__shared__ LongType zLen;
|
||||
__shared__ LongType *sharedMem;
|
||||
__shared__ LongType *xShape, *zShape;
|
||||
__shared__ LongType *xStride, *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
// Cache shape and stride pointers
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
dimID = isNCDHW ? 2 : 1;
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = 5;
|
||||
|
||||
factorD = xShape[dimID + 1] / zShape[dimID + 1];
|
||||
factorH = xShape[dimID + 2] / zShape[dimID + 2];
|
||||
factorW = xShape[dimID + 3] / zShape[dimID + 3];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
INDEX2COORDS(zInd, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
z[zOffset] = 0;
|
||||
|
||||
const LongType zCoord2 = coords[dimID] * factorD;
|
||||
const LongType zCoord3 = coords[dimID + 1] * factorH;
|
||||
const LongType zCoord4 = coords[dimID + 2] * factorW;
|
||||
|
||||
for (coords[dimID] = zCoord2; coords[dimID] < zCoord2 + factorD; ++coords[dimID])
|
||||
for (coords[dimID + 1] = zCoord3; coords[dimID + 1] < zCoord3 + factorH; ++coords[dimID + 1])
|
||||
for (coords[dimID + 2] = zCoord4; coords[dimID + 2] < zCoord4 + factorW; ++coords[dimID + 2]) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
z[zOffset] += x[xOffset];
|
||||
}
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void upsampling3dBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const bool isNCDHW) {
|
||||
upsampling3dBPCuda<T>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, isNCDHW);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"upsampling3dBPCudaLauncher failed");
|
||||
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::upsampling3dBP(graph::Context& block, NDArray& gradO, NDArray& gradI,
|
||||
const bool isNCDHW) {
|
||||
PointersManager manager(block.launchContext(), "upsampling3d_bp");
|
||||
|
||||
dim3 getUpSampling = getUpsamplingDims(gradI.lengthOf(),gradI.rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({&gradI}, {&gradO});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
gradI.dataType(), upsampling3dBPCudaLauncher,
|
||||
(getUpSampling.x, getUpSampling.y, getUpSampling.z, block.launchContext()->getCudaStream(), gradO.specialBuffer(),
|
||||
gradO.specialShapeInfo(), gradI.specialBuffer(), gradI.specialShapeInfo(), isNCDHW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&gradI}, {&gradO});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,128 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// vol [bS, iC, iD, iH, iW] is convoluted to col [bS, iC, kD, kH, kW, oD, oH, oW]
|
||||
template <typename T>
|
||||
static SD_KERNEL void vol2colCuda(const void* volume, const LongType* volShapeInfo, void* columns,
|
||||
const LongType* colShapeInfo, const LongType sD, const LongType sH, const LongType sW,
|
||||
const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH, const LongType dW) {
|
||||
const T* vol = reinterpret_cast<const T*>(volume);
|
||||
T* col = reinterpret_cast<T*>(columns);
|
||||
|
||||
__shared__ LongType colRank, volRank, colLen;
|
||||
__shared__ LongType iD, iH, iW;
|
||||
__shared__ const LongType *volShape, *volStride, *colShape, *colStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
volRank = 5;
|
||||
colRank = 8;
|
||||
|
||||
volShape = shape::shapeOf(volShapeInfo);
|
||||
volStride = shape::stride(volShapeInfo);
|
||||
|
||||
colShape = shape::shapeOf(colShapeInfo);
|
||||
colStride = shape::stride(colShapeInfo);
|
||||
|
||||
colLen = shape::length(colShapeInfo);
|
||||
|
||||
iD = volShape[2];
|
||||
iH = volShape[3];
|
||||
iW = volShape[4];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto colInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (colInd >= colLen) return;
|
||||
|
||||
extern __shared__ LongType sharedMem[];
|
||||
auto coords = sharedMem + threadIdx.x * colRank;
|
||||
|
||||
INDEX2COORDS(colInd, colRank, colShape, coords);
|
||||
|
||||
LongType colOffset;
|
||||
COORDS2INDEX(colRank, colStride, coords, colOffset);
|
||||
|
||||
// Compute volumetric indices
|
||||
const auto volDep = -pD + coords[2] * dD + coords[5] * sD;
|
||||
const auto volRow = -pH + coords[3] * dH + coords[6] * sH;
|
||||
const auto volCol = -pW + coords[4] * dW + coords[7] * sW;
|
||||
|
||||
if (volDep < 0 || volDep >= iD || volRow < 0 || volRow >= iH || volCol < 0 || volCol >= iW) {
|
||||
col[colOffset] = static_cast<T>(0.);
|
||||
} else {
|
||||
coords[2] = volDep;
|
||||
coords[3] = volRow;
|
||||
coords[4] = volCol;
|
||||
|
||||
LongType volOffset;
|
||||
COORDS2INDEX(volRank, volStride, coords, volOffset);
|
||||
col[colOffset] = vol[volOffset];
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void vol2colCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* volume, const LongType* volShapeInfo,
|
||||
void* columns, const LongType* colShapeInfo, const int sD, const LongType sH,
|
||||
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
|
||||
const LongType dW) {
|
||||
vol2colCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(volume, volShapeInfo, columns, colShapeInfo,
|
||||
sD, sH, sW, pD, pH, pW, dD, dH, dW);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"vol2colCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void ConvolutionUtils::vol2col(graph::Context& block, NDArray* vol, NDArray* col, const LongType sD, const LongType sH,
|
||||
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
|
||||
const LongType dW) {
|
||||
PointersManager manager(block.launchContext(), "vol2col");
|
||||
|
||||
dim3 vol2ColDims = getVol2ColDims(col->lengthOf(),col->rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({col}, {vol});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
vol->dataType(), vol2colCudaLauncher,
|
||||
(vol2ColDims.x, vol2ColDims.y, vol2ColDims.z, block.launchContext()->getCudaStream(), vol->specialBuffer(),
|
||||
vol->specialShapeInfo(), col->specialBuffer(), col->specialShapeInfo(), sD, sH, sW, pD, pH, pW, dD, dH, dW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({col}, {vol});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,135 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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, created on 10.06.2019
|
||||
//
|
||||
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/cross.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void crossCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo) {
|
||||
__shared__ const T* x;
|
||||
__shared__ const T* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType rank;
|
||||
__shared__ LongType lenWithoutLastDim, totalThreads;
|
||||
__shared__ const LongType *xShape, *xStride, *yShape, *yStride, *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
x = reinterpret_cast<const T*>(vx);
|
||||
y = reinterpret_cast<const T*>(vy);
|
||||
z = reinterpret_cast<T*>(vz);
|
||||
|
||||
rank = shape::rank(xShapeInfo);
|
||||
lenWithoutLastDim = shape::length(xShapeInfo) / xShapeInfo[rank];
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
extern __shared__ LongType sharedMem[];
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < lenWithoutLastDim; i += totalThreads) {
|
||||
// Compute coordinates without the last dimension
|
||||
INDEX2COORDS(i, rank - 1, xShape, coords);
|
||||
coords[rank - 1] = 0;
|
||||
|
||||
LongType xOffset, yOffset, zOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
COORDS2INDEX(rank, yStride, coords, yOffset);
|
||||
|
||||
// Fetch elements for cross product
|
||||
const auto x0 = x[xOffset];
|
||||
const auto y0 = y[yOffset];
|
||||
|
||||
xOffset += xStride[rank - 1];
|
||||
yOffset += yStride[rank - 1];
|
||||
|
||||
const auto x1 = x[xOffset];
|
||||
const auto y1 = y[yOffset];
|
||||
|
||||
xOffset += xStride[rank - 1];
|
||||
yOffset += yStride[rank - 1];
|
||||
|
||||
const auto x2 = x[xOffset];
|
||||
const auto y2 = y[yOffset];
|
||||
|
||||
// Compute offsets for output
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
z[zOffset] = x1 * y2 - x2 * y1;
|
||||
|
||||
zOffset += zStride[rank - 1];
|
||||
z[zOffset] = x2 * y0 - x0 * y2;
|
||||
|
||||
zOffset += zStride[rank - 1];
|
||||
z[zOffset] = x0 * y1 - x1 * y0;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
SD_HOST static void crossCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo) {
|
||||
crossCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t*>(stream),"crossCuda failed");
|
||||
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void crossCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, const void* vy,
|
||||
const sd::LongType* yShapeInfo, void* vz, const sd::LongType* zShapeInfo),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
void crossBatched(LaunchContext* context, NDArray* x, NDArray* y, NDArray* z) {
|
||||
dim3 launchDims = getCross(x->lengthOf(),x->rankOf(),x->sizeAt(-1));
|
||||
PointersManager manager(context, "cross");
|
||||
|
||||
NDArray::prepareSpecialUse({z}, {x, y});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
x->dataType(), crossCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), x->specialBuffer(), x->specialShapeInfo(),
|
||||
y->specialBuffer(), y->specialShapeInfo(), z->specialBuffer(), z->specialShapeInfo()),
|
||||
SD_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({z}, {x, y});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,43 @@
|
||||
/*******************************************************************************
|
||||
* Copyright (c) 2021 Deeplearning4j Contributors
|
||||
*
|
||||
* 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.
|
||||
*
|
||||
* 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 AbdelRauf
|
||||
//
|
||||
#include <execution/ThreadPool.h>
|
||||
#include <execution/Threads.h>
|
||||
#include <helpers/LoopsCoordsHelper.h>
|
||||
#include <ops/declarable/helpers/ctc.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <type_traits>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
void ctcLoss(graph::Context &block, NDArray&logitsInput, NDArray&targetLabels,
|
||||
NDArray&logitsLengths, NDArray&targetLabelLengths, NDArray &logLosses,
|
||||
NDArray &gradients, int blankIndex) {
|
||||
// not imeplemented
|
||||
THROW_EXCEPTION("ctcLoss:: Not implemented yet");
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,116 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
//
|
||||
//
|
||||
#include <ops/declarable/helpers/d_t_s.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void depthToSpaceKernel(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const int block_size, const bool isNHWC) {
|
||||
auto input_ptr = reinterpret_cast<const T *>(vx);
|
||||
auto output_ptr = reinterpret_cast<T *>(vz);
|
||||
|
||||
const int batch_size = shape::sizeAt(xShapeInfo, 0);
|
||||
const int input_depth = isNHWC ? shape::sizeAt(xShapeInfo, 3) : shape::sizeAt(xShapeInfo, 1);
|
||||
const int input_height = isNHWC ? shape::sizeAt(xShapeInfo, 1) : shape::sizeAt(xShapeInfo, 2);
|
||||
const int input_width = isNHWC ? shape::sizeAt(xShapeInfo, 2) : shape::sizeAt(xShapeInfo, 3);
|
||||
|
||||
const int output_depth = isNHWC ? shape::sizeAt(zShapeInfo, 3) : shape::sizeAt(zShapeInfo, 1);
|
||||
const int output_height = isNHWC ? shape::sizeAt(zShapeInfo, 1) : shape::sizeAt(zShapeInfo, 2);
|
||||
const int output_width = isNHWC ? shape::sizeAt(zShapeInfo, 2) : shape::sizeAt(zShapeInfo, 3);
|
||||
|
||||
const int input_area = input_width * input_height;
|
||||
const int input_depth_by_input_area = input_depth * input_area;
|
||||
const int output_depth_by_input_height = output_depth * input_height;
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (isNHWC) {
|
||||
const int total_count = batch_size * output_height * output_width * output_depth;
|
||||
for (int out_idx = tid; out_idx < total_count; out_idx += blockDim.x * gridDim.x) {
|
||||
const int d = out_idx % output_depth;
|
||||
const int out_idx2 = out_idx / output_depth;
|
||||
const int w = out_idx2 % output_width;
|
||||
const int out_idx3 = out_idx2 / output_width;
|
||||
const int h = out_idx3 % output_height;
|
||||
const int b = out_idx3 / output_height;
|
||||
|
||||
const int in_h = h / block_size;
|
||||
const int offset_h = h % block_size;
|
||||
const int in_w = w / block_size;
|
||||
const int offset_w = w % block_size;
|
||||
const int offset_d = (offset_h * block_size + offset_w) * output_depth;
|
||||
const int in_d = d + offset_d;
|
||||
const int inp_idx = in_d + input_depth * (in_w + input_width * (in_h + input_height * b));
|
||||
(output_ptr + out_idx)[0] = (input_ptr + inp_idx)[0];
|
||||
}
|
||||
} else {
|
||||
const int total_count = batch_size * input_depth_by_input_area;
|
||||
|
||||
for (int input_idx = tid; input_idx < total_count; input_idx += blockDim.x * gridDim.x) {
|
||||
const int n_bY_bX_oC_iY = input_idx / input_width;
|
||||
const int iX = input_idx - n_bY_bX_oC_iY * input_width;
|
||||
|
||||
const int n_bY_bX = n_bY_bX_oC_iY / output_depth_by_input_height;
|
||||
const int oC_iY = n_bY_bX_oC_iY - n_bY_bX * output_depth_by_input_height;
|
||||
|
||||
const int n_bY = n_bY_bX / block_size;
|
||||
const int bX = n_bY_bX - n_bY * block_size;
|
||||
|
||||
const int n = n_bY / block_size;
|
||||
const int bY = n_bY - n * block_size;
|
||||
|
||||
const int output_idx =
|
||||
bX + block_size * (iX + input_width * (bY + block_size * (oC_iY + n * output_depth_by_input_height)));
|
||||
|
||||
(output_ptr + output_idx)[0] = (input_ptr + input_idx)[0];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void __depthToSpace(LaunchContext *context, NDArray&input, NDArray *output, int block_size,
|
||||
bool isNHWC) {
|
||||
dim3 launchDims = getLaunchDims("depth_to_space");
|
||||
depthToSpaceKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *context->getCudaStream()>>>(input.specialBuffer(), input.specialShapeInfo(),
|
||||
output->specialBuffer(),
|
||||
output->specialShapeInfo(), block_size, isNHWC);
|
||||
DebugHelper::checkGlobalErrorCode("depthToSpaceKernel failed");
|
||||
|
||||
}
|
||||
|
||||
void _depthToSpace(LaunchContext *context, NDArray&input, NDArray *output, int block_size, bool isNHWC) {
|
||||
auto xType = input.dataType();
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {&input});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, __depthToSpace, (context, input, output, block_size, isNHWC), SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {&input});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,108 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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)
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <ops/declarable/helpers/gammaMathFunc.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void diGammaCuda(const void *vx, const LongType *xShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ LongType len;
|
||||
__shared__ bool sameOffset;
|
||||
__shared__ LongType xRank, zRank;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape, *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(xShapeInfo);
|
||||
sameOffset = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType zOffset;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < len; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
|
||||
if (sameOffset) {
|
||||
zOffset = xOffset;
|
||||
} else {
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
}
|
||||
|
||||
z[zOffset] = diGammaScalar<T>(x[xOffset]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void diGammaCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo) {
|
||||
diGammaCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory, *stream>>>(vx, xShapeInfo, vz, zShapeInfo);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "crossCuda failed");
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void diGamma(LaunchContext *context, NDArray&x, NDArray &z) {
|
||||
dim3 digammaDims2 = digammaDims(z.lengthOf());
|
||||
NDArray::prepareSpecialUse({&z}, {&x});
|
||||
BUILD_SINGLE_SELECTOR(x.dataType(), diGammaCudaLauncher,
|
||||
(digammaDims2.y, digammaDims2.x, digammaDims2.z, context->getCudaStream(), x.specialBuffer(),
|
||||
x.specialShapeInfo(), z.specialBuffer(), z.specialShapeInfo()),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&z}, {&x});
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void diGammaCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
||||
const cudaStream_t *stream, const void *vx, const sd::LongType *xShapeInfo, void *vz,
|
||||
const sd::LongType *zShapeInfo),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,200 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// Created by GS <sgazeos@gmail.com> on 4/6/2018.
|
||||
//
|
||||
#include <array/ResultSet.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <ops/declarable/helpers/diag.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// diag functor cuda kernel
|
||||
// outputBuffer - output tensor buffer
|
||||
// outputShape - output tensor shape
|
||||
// inputBuffer - input tensor buffer - this tensor should be placed on diagonal position of output
|
||||
// inputShape - input tensor shape
|
||||
// inputLength - length for input tensor
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void diagFunctorKernel(void* outputBuffer, const LongType* outputShape, void const* inputBuffer,
|
||||
const LongType* inputShape, LongType inputLength) {
|
||||
__shared__ T* z;
|
||||
__shared__ T const* x;
|
||||
__shared__ LongType outputRank, inputRank, outputLength;
|
||||
__shared__ const LongType *outputShapePtr, *outputStridePtr;
|
||||
__shared__ const LongType *inputShapePtr, *inputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z = reinterpret_cast<T*>(outputBuffer);
|
||||
x = reinterpret_cast<T const*>(inputBuffer);
|
||||
|
||||
outputRank = shape::rank(outputShape);
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputLength = shape::length(outputShape);
|
||||
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
LongType xOffset;
|
||||
|
||||
for (LongType t = tid; t < inputLength; t += step) {
|
||||
// Compute coordinates and offsets for output
|
||||
INDEX2COORDS(t * (inputLength + 1), outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
|
||||
// Compute coordinates and offsets for input
|
||||
INDEX2COORDS(t, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
|
||||
// Assign the value to the diagonal position
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// diag part functor cuda kernel
|
||||
// outputBuffer - output tensor buffer - linear sequence of diagonal values
|
||||
// outputShape - output tensor shape
|
||||
// inputBuffer - input tensor buffer - this tensor should be placed on diagonal position of output
|
||||
// inputShape - input tensor shape
|
||||
// outputLength - given length of output
|
||||
// inputLength - given length for input tensor
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void diagPartFunctorKernel(void* outputBuffer, const LongType* outputShape,
|
||||
void const* inputBuffer, const LongType* inputShape, LongType outputLength, LongType inputLength) {
|
||||
__shared__ T* z;
|
||||
__shared__ T const* x;
|
||||
__shared__ LongType outputRank, inputRank;
|
||||
__shared__ const LongType *outputShapePtr, *outputStridePtr;
|
||||
__shared__ const LongType *inputShapePtr, *inputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z = reinterpret_cast<T*>(outputBuffer);
|
||||
x = reinterpret_cast<T const*>(inputBuffer);
|
||||
|
||||
outputRank = shape::rank(outputShape);
|
||||
inputRank = shape::rank(inputShape);
|
||||
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
LongType xOffset;
|
||||
LongType i = tid * (outputLength + 1); // position of the diagonal value in the input
|
||||
|
||||
for (LongType t = tid; t < outputLength && i < inputLength; t += step) {
|
||||
// Compute coordinates and offsets for output
|
||||
INDEX2COORDS(t, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
|
||||
// Compute coordinates and offsets for input
|
||||
INDEX2COORDS(i, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
|
||||
// Assign diagonal value
|
||||
z[zOffset] = x[xOffset];
|
||||
|
||||
// Move to the next diagonal value
|
||||
i += outputLength + 1;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Returns a batched matrix tensor with new batched diagonal values.
|
||||
// for detailed explanations please take a look on web page:
|
||||
// https://www.tensorflow.org/api_docs/python/tf/matrix_set_diag
|
||||
template <typename T>
|
||||
static void _diagFunctor(LaunchContext* context, NDArray* input, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
auto inputLength = input->isScalar() ? 1 : input->lengthOf();
|
||||
dim3 launchDims = getLaunchDims("diagPart");
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
diagFunctorKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), input->specialBuffer(), input->specialShapeInfo(),
|
||||
inputLength);
|
||||
DebugHelper::checkErrorCode(stream,"diagFunctorKernel failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// diagFunctor - caller for diag functor processor
|
||||
void diagFunctor(LaunchContext* context, NDArray* input, NDArray* output) {
|
||||
auto xType = input->dataType();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, _diagFunctor, (context, input, output), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void _diagFunctor, (sd::LaunchContext * context, NDArray* input, NDArray* output);
|
||||
, SD_COMMON_TYPES);
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// diagPartFunctor - caller for diag part functor kernel
|
||||
template <typename T>
|
||||
void _diagPartFunctor(LaunchContext* context, NDArray * input, NDArray* output) {
|
||||
const int outLen = output->lengthOf();
|
||||
const int inLen = input->isScalar() ? 1 : input->lengthOf();
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
dim3 launchDims = getLaunchDims("diagPart");
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
|
||||
diagPartFunctorKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), input->specialBuffer(), input->specialShapeInfo(), outLen,
|
||||
inLen);
|
||||
DebugHelper::checkErrorCode(stream,"diagFunctorKernel failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// diagPartFunctor - caller for diag part functor processor
|
||||
void diagPartFunctor(LaunchContext* context, NDArray * input, NDArray* output) {
|
||||
auto zType = output->dataType();
|
||||
BUILD_SINGLE_SELECTOR(zType, _diagPartFunctor, (context, input, output), SD_NUMERIC_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,143 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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)
|
||||
//
|
||||
#include <array/DataTypeUtils.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/dilation2d.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL static void dilation2dCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const int sH, const int sW, const int pH, const int pW, const int dH,
|
||||
const int dW) {
|
||||
// x [bS, iH, iW, iC]
|
||||
// y [kH, kW, iC]
|
||||
// z [bS, oH, oW, iC]
|
||||
|
||||
const X* x = reinterpret_cast<const X*>(vx);
|
||||
const X* y = reinterpret_cast<const X*>(vy);
|
||||
Z* z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ LongType xRank, yRank, zRank;
|
||||
__shared__ const LongType *xShape, *xStride, *yShape, *yStride, *zShape, *zStride;
|
||||
__shared__ LongType iH, iW, kH, kW, zLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
iH = xShape[1];
|
||||
iW = xShape[2];
|
||||
|
||||
kH = yShape[0];
|
||||
kW = yShape[1];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto zInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (zInd >= zLen) return;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
|
||||
INDEX2COORDS(zInd, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
|
||||
yCoords[2] = zCoords[3]; // iC coordinate is the same for x, y, and z
|
||||
const auto oh = zCoords[1];
|
||||
const auto ow = zCoords[2];
|
||||
|
||||
X max = -DataTypeUtils::max<X>();
|
||||
|
||||
for (yCoords[0] = 0; yCoords[0] < kH; ++yCoords[0]) {
|
||||
xCoords[1] = oh * sH - pH + yCoords[0] * dH;
|
||||
if (xCoords[1] < 0 || xCoords[1] >= iH) continue;
|
||||
|
||||
for (yCoords[1] = 0; yCoords[1] < kW; ++yCoords[1]) {
|
||||
xCoords[2] = ow * sW - pW + yCoords[1] * dW;
|
||||
if (xCoords[2] < 0 || xCoords[2] >= iW) continue;
|
||||
|
||||
LongType xOffset, yOffset;
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
|
||||
const X val = x[xOffset] + y[yOffset];
|
||||
if (val > max) max = val;
|
||||
}
|
||||
}
|
||||
|
||||
z[zOffset] = static_cast<Z>(max);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static void dilation2dCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW) {
|
||||
dilation2dCuda<X, Z><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz,
|
||||
zShapeInfo, sH, sW, pH, pW, dH, dW);
|
||||
DebugHelper::checkGlobalErrorCode( "dilation2d(...) failed");
|
||||
|
||||
}
|
||||
|
||||
void dilation2d(LaunchContext* context, NDArray* input, NDArray* weights, NDArray* output, const LongType sH,
|
||||
const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW) {
|
||||
PointersManager manager(context, "dilation2d");
|
||||
dim3 dilation = getDilation(output->lengthOf(),weights->rankOf(),output->rankOf());
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, weights});
|
||||
BUILD_SINGLE_SELECTOR_TWICE(
|
||||
input->dataType(), dilation2dCudaLauncher,
|
||||
(dilation.y, dilation.x, dilation.z, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), weights->specialBuffer(), weights->specialShapeInfo(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), sH, sW, pH, pW, dH, dW),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, weights});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,393 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <legacy/NativeOps.h>
|
||||
#include <ops/declarable/helpers/dropout.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void dropoutSimpleKernel(void const* inputBuf, LongType const* inputShape, void* outputBuf,
|
||||
LongType const* outputShape, double probVal, int inLen,
|
||||
RandomGenerator* nodeRng) {
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
T const* input = reinterpret_cast<T const*>(inputBuf);
|
||||
T* output = reinterpret_cast<T*>(outputBuf);
|
||||
|
||||
__shared__ LongType inputRank, outputRank;
|
||||
__shared__ const LongType *inputShapePtr, *inputStridePtr;
|
||||
__shared__ const LongType *outputShapePtr, *outputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputOffset;
|
||||
LongType outputOffset;
|
||||
|
||||
// Loop through all elements and nullify based on probability
|
||||
for (LongType e = tid; e < inLen; e += step) {
|
||||
T val = nodeRng->relativeT(e, T(0.f), T(1.f));
|
||||
|
||||
// If probability is acceptable, save the scaled value
|
||||
if (double(val) < probVal) {
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, outputOffset);
|
||||
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, inputOffset);
|
||||
|
||||
output[outputOffset] = T(input[inputOffset] / probVal);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void dropoutSimple(LaunchContext* context, NDArray * input, NDArray* output, double probValue,
|
||||
int seed) {
|
||||
RandomGenerator nodeRng(3019L, seed);
|
||||
int inLen = input->lengthOf();
|
||||
RandomGenerator* dRandom;
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
|
||||
auto err = cudaMalloc(&dRandom, sizeof(RandomGenerator));
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::dropoutSimple: Cannot allocate device memory for random generator.", err);
|
||||
}
|
||||
err = cudaMemcpy(dRandom, &nodeRng, sizeof(RandomGenerator), cudaMemcpyHostToDevice);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::dropoutSimple: Cannot set up device memory for random generator.", err);
|
||||
}
|
||||
|
||||
dim3 getDims = getLaunchDims("dropout");
|
||||
dropoutSimpleKernel<T><<<getDims.x, getDims.y, getDims.z, *stream>>>(input->specialBuffer(), input->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), probValue,
|
||||
inLen, dRandom);
|
||||
err = cudaFree(dRandom);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::dropoutSimple: Cannot deallocate device memory for random generator.", err);
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status _dropOutFunctor(sd::graph::Context& context, NDArray* input, NDArray* output, NDArray* reduceShape, int seed,
|
||||
double probValue) {
|
||||
if (reduceShape == nullptr) {
|
||||
dropoutSimple<T>(context.launchContext(), input, output, probValue, seed);
|
||||
} else {
|
||||
REQUIRE_TRUE(reduceShape->lengthOf() <= input->rankOf(), 0, "dropout: Noise shape should be fittable to input");
|
||||
|
||||
std::vector<LongType> dims(reduceShape->lengthOf());
|
||||
reduceShape->syncToHost(); // to ensure that follows are actual
|
||||
bool fit = true;
|
||||
|
||||
for (int i = 0; i < dims.size(); i++) {
|
||||
if (fit) {
|
||||
dims[i] = reduceShape->e<LongType>(i);
|
||||
for (int e = 0; e < input->rankOf(); ++e)
|
||||
if (fit)
|
||||
if (input->sizeAt(e) % dims[i]) {
|
||||
fit = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// check dims to fit input
|
||||
REQUIRE_TRUE(fit, 0, "dropout: Noise shape should fit to input rank.");
|
||||
std::unique_ptr<NDArray> chunk(new NDArray('c', dims, output->dataType(), context.launchContext()));
|
||||
float one = 1.f;
|
||||
chunk->assign(one);
|
||||
|
||||
dropoutSimple<T>(context.launchContext(), chunk.get(), chunk.get(), probValue, seed);
|
||||
// broadcast chunk to full matrix
|
||||
std::unique_ptr<NDArray> dropOutMultiplier(new NDArray(*input));
|
||||
dropOutMultiplier->assign(one);
|
||||
|
||||
*dropOutMultiplier += *chunk;
|
||||
|
||||
// FIXME: we could do this in one step, aren't we?
|
||||
NDArray ret = *input * *dropOutMultiplier;
|
||||
output->assign(&ret);
|
||||
}
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
Status dropOutFunctor(sd::graph::Context& context, NDArray* input, NDArray* output, NDArray* reduceShape, int seed, double probValue, NDArray* mask) {
|
||||
auto xType = input->dataType();
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, return _dropOutFunctor, (context, input, output, reduceShape, seed, probValue),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
/////////////////////////////////// backpropagations ///////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void dropoutBPKernel(void* outputBuf, LongType const* outputShape, void* gradOutBuf,
|
||||
LongType const* gradOutShape, double probValue) {
|
||||
__shared__ T* output;
|
||||
__shared__ T* input;
|
||||
__shared__ LongType len;
|
||||
__shared__ LongType outputRank, gradOutRank;
|
||||
__shared__ const LongType *outputShapePtr, *outputStridePtr;
|
||||
__shared__ const LongType *gradOutShapePtr, *gradOutStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(outputShape);
|
||||
|
||||
output = reinterpret_cast<T*>(outputBuf);
|
||||
input = reinterpret_cast<T*>(gradOutBuf);
|
||||
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
gradOutRank = shape::rank(gradOutShape);
|
||||
gradOutShapePtr = shape::shapeOf(gradOutShape);
|
||||
gradOutStridePtr = shape::stride(gradOutShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType gradOutCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
LongType gradOutOffset;
|
||||
|
||||
for (LongType e = tid; e < len; e += step) {
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, zOffset);
|
||||
|
||||
INDEX2COORDS(e, gradOutRank, gradOutShapePtr, gradOutCoords);
|
||||
COORDS2INDEX(gradOutRank, gradOutStridePtr, gradOutCoords, gradOutOffset);
|
||||
|
||||
// Scale gradients back if the output wasn't zero
|
||||
if (output[zOffset] != T(0.)) {
|
||||
output[zOffset] = T(input[gradOutOffset] / probValue);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static Status dropOutFunctorBP_(sd::graph::Context& context, NDArray* input, NDArray* gradOut, NDArray* output,
|
||||
NDArray* reduceShape, int seed, double probValue, NDArray* mask) {
|
||||
// we're making additional FF run to see how probabilities played out with given seeds
|
||||
auto res = dropOutFunctor(context, input, output, reduceShape, seed, probValue,mask);
|
||||
auto stream = context.launchContext()->getCudaStream();
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, gradOut});
|
||||
|
||||
|
||||
if (Status::OK == res) {
|
||||
dim3 launchDims = getLaunchDims("dropout");
|
||||
dropoutBPKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
probValue);
|
||||
|
||||
DebugHelper::checkGlobalErrorCode( "dropout_bp(...) failed");
|
||||
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, gradOut});
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void alphaDropoutSimpleKernel(void const* inputBuf, LongType const* inputShape, void* outputBuf,
|
||||
LongType const* outputShape, double probValue, double alpha,
|
||||
double alpha1, double beta, int inLen, RandomGenerator* nodeRng) {
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
T const* input = reinterpret_cast<T const*>(inputBuf);
|
||||
T* output = reinterpret_cast<T*>(outputBuf);
|
||||
|
||||
__shared__ LongType inputRank, outputRank;
|
||||
__shared__ const LongType *inputShapePtr, *inputStridePtr;
|
||||
__shared__ const LongType *outputShapePtr, *outputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputOffset;
|
||||
LongType outputOffset;
|
||||
|
||||
for (auto e = tid; e < inLen; e += step) {
|
||||
T val = nodeRng->relativeT(e, T(0.f), T(1.f));
|
||||
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, inputOffset);
|
||||
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, outputOffset);
|
||||
|
||||
output[outputOffset] = (val >= T(probValue)
|
||||
? T(alpha * beta + alpha1)
|
||||
: T(alpha * static_cast<double>(input[inputOffset]) + alpha1));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void alphaDropoutSimple(LaunchContext* context, NDArray * input, NDArray* output, int seed,
|
||||
double probValue, double alpha, double alpha1, double beta) {
|
||||
RandomGenerator nodeRng(3019L, seed), *dRandom;
|
||||
auto stream = context->getCudaStream();
|
||||
auto err = cudaMalloc(&dRandom, sizeof(RandomGenerator));
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::alphaDropoutSimple: Cannot allocate device memory for random generator.",
|
||||
err);
|
||||
}
|
||||
err = cudaMemcpy(dRandom, &nodeRng, sizeof(RandomGenerator), cudaMemcpyHostToDevice);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::alphaDropoutSimple: Cannot set up device memory for random generator.", err);
|
||||
}
|
||||
|
||||
dim3 launchDims = getLaunchDims("dropout");
|
||||
alphaDropoutSimpleKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(), probValue,
|
||||
alpha, alpha1, beta, output->lengthOf(), dRandom);
|
||||
|
||||
DebugHelper::checkGlobalErrorCode( "alphaDropoutSimpleKernel(...) failed");
|
||||
|
||||
err = cudaFree(dRandom);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::alphaDropoutSimple: Cannot deallocate device memory for random generator.",
|
||||
err);
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static Status alphaDropOutFunctor_(sd::graph::Context& context, NDArray* input, NDArray* output, NDArray* reduceShape, int seed, double probValue, double alpha,
|
||||
double alpha1, double beta, NDArray* mask) {
|
||||
if (reduceShape == nullptr) {
|
||||
alphaDropoutSimple<T>(context.launchContext(), input, output, seed, probValue, alpha, alpha1, beta);
|
||||
} else {
|
||||
REQUIRE_TRUE(reduceShape->lengthOf() <= input->rankOf(), 0, "dropout: Noise shape should be fittable to input");
|
||||
|
||||
std::vector<LongType> dims(reduceShape->lengthOf());
|
||||
reduceShape->syncToHost(); // to ensure that follows are actual
|
||||
bool fit = true;
|
||||
|
||||
for (int i = 0; i < dims.size(); i++) {
|
||||
if (fit) {
|
||||
dims[i] = reduceShape->e<LongType>(i);
|
||||
for (int e = 0; e < input->rankOf(); ++e)
|
||||
if (fit)
|
||||
if (input->sizeAt(e) % dims[i]) {
|
||||
fit = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// check dims to fit input
|
||||
REQUIRE_TRUE(fit, 0, "alpha_dropout: Noise shape should fit to input rank.");
|
||||
std::unique_ptr<NDArray> chunk(new NDArray('c', dims, output->dataType(), context.launchContext()));
|
||||
float one = 1.f;
|
||||
|
||||
chunk->assign(one);
|
||||
|
||||
alphaDropoutSimple<T>(context.launchContext(), chunk.get(), chunk.get(), seed, probValue, alpha, alpha1, beta);
|
||||
|
||||
// broadcast chunk to full matrix
|
||||
std::unique_ptr<NDArray> dropOutMultiplier(new NDArray(*input));
|
||||
dropOutMultiplier->assign(one);
|
||||
|
||||
*dropOutMultiplier += *chunk;
|
||||
NDArray ret = *input * *dropOutMultiplier;
|
||||
output->assign(&ret);
|
||||
|
||||
}
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status alphaDropOutFunctorBP_(sd::graph::Context& context, NDArray* input, NDArray* gradOut, NDArray* output, NDArray* reduceShape,
|
||||
int seed, double probValue, double alpha, double alpha1, double beta, NDArray* mask) {
|
||||
auto res = alphaDropOutFunctor(context, input, output, reduceShape, seed, probValue, alpha, alpha1, beta, mask);
|
||||
if (res == Status::OK) {
|
||||
// FIXME: can we make it single-loop?
|
||||
(*output) *= alpha;
|
||||
(*output) *= (*gradOut);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
Status dropOutFunctorBP(sd::graph::Context& context, NDArray* input, NDArray* gradOut, NDArray* output, NDArray* reduceShape,
|
||||
int seed, double probValue, NDArray* mask) {
|
||||
BUILD_SINGLE_SELECTOR(context.dataType(), return dropOutFunctorBP_,
|
||||
(context, input, gradOut, output, reduceShape, seed, probValue, mask), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
Status alphaDropOutFunctor(sd::graph::Context& context, NDArray* input, NDArray* output, NDArray* reduceShape, int seed,
|
||||
double probValue, double alpha, double alpha1, double beta, NDArray* mask) {
|
||||
BUILD_SINGLE_SELECTOR(context.dataType(), return alphaDropOutFunctor_,
|
||||
(context, input, output, reduceShape, seed, probValue, alpha, alpha1, beta, mask),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
Status alphaDropOutFunctorBP(sd::graph::Context& context, NDArray* input, NDArray* gradOut, NDArray* output, NDArray* reduceShape, int seed, double probValue,
|
||||
double alpha, double alpha1, double beta, NDArray* mask) {
|
||||
BUILD_SINGLE_SELECTOR(context.dataType(), return alphaDropOutFunctorBP_,
|
||||
(context, input, gradOut, output, reduceShape, seed, probValue, alpha, alpha1, beta,mask),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,529 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/dynamic.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename X, typename Y>
|
||||
static SD_KERNEL void dynamicPartitionScalarKernel(const void *vx, const LongType *xShapeInfo, const void *vi,
|
||||
const LongType *iShapeInfo, void **vz, LongType **zShapeInfos, const LongType numOutputs) {
|
||||
auto x = reinterpret_cast<const X *>(vx);
|
||||
auto i = reinterpret_cast<const Y *>(vi);
|
||||
|
||||
__shared__ LongType xRank, iRank;
|
||||
__shared__ const LongType *xShape, *xStride;
|
||||
__shared__ const LongType *iShape, *iStride;
|
||||
|
||||
// Shared variables for CUDA kernel
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
iRank = shape::rank(iShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
iShape = shape::shapeOf(iShapeInfo);
|
||||
iStride = shape::stride(iShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto xLength = shape::length(xShapeInfo);
|
||||
auto iLength = shape::length(iShapeInfo);
|
||||
|
||||
extern __shared__ char shmem[];
|
||||
__shared__ Y *rawIndices;
|
||||
__shared__ Y *trueIndices;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rawIndices = reinterpret_cast<Y *>(shmem);
|
||||
trueIndices = rawIndices + blockDim.x;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Process partitions
|
||||
for (LongType o = blockIdx.x; o < numOutputs; o += gridDim.x) {
|
||||
auto z = reinterpret_cast<X *>(vz[o]);
|
||||
auto zShapeInfo = zShapeInfos[o];
|
||||
|
||||
__shared__ LongType zLength, zRank;
|
||||
__shared__ const LongType *zShape, *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLength = shape::length(zShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Ensure iLimit is a multiple of blockDim.x
|
||||
auto iLimit = (iLength <= blockDim.x) ? blockDim.x : (iLength + (blockDim.x - (iLength % blockDim.x)));
|
||||
int cnt = 0;
|
||||
|
||||
for (LongType e = threadIdx.x; e < iLimit; e += blockDim.x) {
|
||||
if (e < iLength) {
|
||||
LongType iOffset, iCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(e, iRank, iShape, iCoords);
|
||||
COORDS2INDEX(iRank, iStride, iCoords, iOffset);
|
||||
rawIndices[threadIdx.x] = i[iOffset];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Map updates using prefix-like approach
|
||||
if (threadIdx.x == 0) {
|
||||
for (int f = 0; f < blockDim.x; f++) {
|
||||
if (rawIndices[f] == static_cast<Y>(o))
|
||||
trueIndices[f] = cnt++;
|
||||
else
|
||||
trueIndices[f] = -1;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Perform actual update
|
||||
if (e < iLength && trueIndices[threadIdx.x] >= 0) {
|
||||
LongType xOffset, xCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(e, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
z[trueIndices[threadIdx.x]] = x[xOffset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
static SD_KERNEL void dynamicPartitionTadKernel(const void *vx, const LongType *xTadShapeInfo,
|
||||
const LongType *xTadOffsets, LongType xLength,
|
||||
const void *vindices, const LongType *iShapeInfo, LongType iLength, void **vz,
|
||||
LongType **zTadShapeInfos, LongType **zTadOffsets,
|
||||
LongType numOutputs) {
|
||||
auto x = reinterpret_cast<const X *>(vx);
|
||||
auto indices = reinterpret_cast<const Y *>(vindices);
|
||||
|
||||
// we run things in blocks, 1 partition per block of threads
|
||||
for (int i = blockIdx.x; i < numOutputs; i += gridDim.x) {
|
||||
auto z = reinterpret_cast<X *>(vz[i]);
|
||||
|
||||
// each thread has own counter for partitions
|
||||
int outCnt = 0;
|
||||
|
||||
for (LongType e = 0; e < iLength; e++) {
|
||||
LongType iCoords[SD_MAX_RANK];
|
||||
LongType iOffset;
|
||||
INDEX2COORDS(e, shape::rank(iShapeInfo), shape::shapeOf(iShapeInfo), iCoords);
|
||||
COORDS2INDEX(shape::rank(iShapeInfo), shape::stride(iShapeInfo), iCoords, iOffset);
|
||||
|
||||
if (indices[iOffset] == i) {
|
||||
auto dx = x + xTadOffsets[e];
|
||||
auto dz = z + zTadOffsets[i][outCnt++];
|
||||
|
||||
for (int f = threadIdx.x; f < xLength; f += blockDim.x) {
|
||||
LongType fCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType zOffset;
|
||||
INDEX2COORDS(f, shape::rank(xTadShapeInfo), shape::shapeOf(xTadShapeInfo), fCoords);
|
||||
COORDS2INDEX(shape::rank(xTadShapeInfo), shape::stride(xTadShapeInfo), fCoords, xOffset);
|
||||
INDEX2COORDS(f, shape::rank(zTadShapeInfos[i]), shape::shapeOf(zTadShapeInfos[i]), fCoords);
|
||||
COORDS2INDEX(shape::rank(zTadShapeInfos[i]), shape::stride(zTadShapeInfos[i]), fCoords, zOffset);
|
||||
|
||||
dz[zOffset] = dx[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
static void _dynamicPartitionFunctor(LaunchContext *context, NDArray *input, NDArray *indices,
|
||||
std::vector<NDArray *> &outputList) {
|
||||
std::vector<std::pair<NDArray *, int>> outputs(outputList.size());
|
||||
int sourceDimsLen = input->rankOf() - indices->rankOf();
|
||||
|
||||
unsigned int outSize = outputList.size();
|
||||
|
||||
PointersManager pm(context, "dynamicPartition");
|
||||
|
||||
if (sourceDimsLen) { // non-linear case
|
||||
std::vector<LongType> sourceDims(sourceDimsLen);
|
||||
|
||||
for (int i = sourceDimsLen; i > 0; i--) sourceDims[sourceDimsLen - i] = input->rankOf() - i;
|
||||
// compute tad array for given dimensions
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), &sourceDims);
|
||||
|
||||
std::vector<void *> outBuffers(outSize);
|
||||
std::vector<const LongType *> tadShapes(outSize);
|
||||
std::vector<const LongType *> tadOffsets(outSize);
|
||||
std::vector<LongType> numTads(outSize);
|
||||
// fill up dimensions array for before kernel
|
||||
for (unsigned int i = 0; i < outSize; i++) {
|
||||
outputs[i].first = outputList[i];
|
||||
std::vector<LongType> outDims(outputs[i].first->rankOf() - 1);
|
||||
|
||||
int r = outputs[i].first->rankOf();
|
||||
|
||||
for (int k = 1; k < r; k++) outDims[k - 1] = k;
|
||||
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(outputList.at(i)->shapeInfo(), &outDims);
|
||||
|
||||
outBuffers[i] = outputList.at(i)->specialBuffer();
|
||||
tadShapes[i] = packZ->platformShapeInfo();
|
||||
tadOffsets[i] = packZ->platformOffsets();
|
||||
}
|
||||
|
||||
// we copy pointers to device
|
||||
auto dOutBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(outBuffers.data(), outBuffers.size() * sizeof(void *)));
|
||||
auto dOutTadShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(tadShapes.data(), tadShapes.size() * sizeof(LongType *)));
|
||||
auto dOutTadOffsets = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(tadOffsets.data(), tadOffsets.size() * sizeof(LongType *)));
|
||||
// run kernel on device
|
||||
dim3 launchDims = getDynamicPartitionDims(256,sizeof(Y));
|
||||
|
||||
dynamicPartitionTadKernel<X, Y><<<launchDims.y,launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
||||
input->specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(),
|
||||
shape::length(packX->primaryShapeInfo()), indices->specialBuffer(), indices->specialShapeInfo(),
|
||||
indices->lengthOf(), dOutBuffers, dOutTadShapes, dOutTadOffsets, outSize);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"dynamicPartitionTadKernel failed");
|
||||
|
||||
|
||||
} else { // linear case
|
||||
dim3 launchDims = getDynamicPartitionDims(256,sizeof(Y));
|
||||
std::vector<void *> outBuffers;
|
||||
std::vector<const LongType *> outShapes;
|
||||
|
||||
for (auto v : outputList) {
|
||||
outBuffers.emplace_back(v->specialBuffer());
|
||||
outShapes.emplace_back(v->specialShapeInfo());
|
||||
}
|
||||
|
||||
auto dOutBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(outBuffers.data(), outBuffers.size() * sizeof(void *)));
|
||||
auto dOutShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(outShapes.data(), outShapes.size() * sizeof(LongType *)));
|
||||
|
||||
dynamicPartitionScalarKernel<X, Y><<<launchDims.y,launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
||||
dOutBuffers, dOutShapes, outSize);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"dynamicPartitionScalarKernel failed");
|
||||
|
||||
}
|
||||
|
||||
pm.synchronize();
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
static SD_KERNEL void dynamicStitchScalarKernel(void **vx, LongType **xShapeInfos, void **vindices,
|
||||
LongType **iShapeInfos, int inputSize, void *vz,
|
||||
const LongType *zShapeInfo, LongType zLength) {
|
||||
__shared__ LongType zRank;
|
||||
__shared__ const LongType *zShapePtr, *zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto z = reinterpret_cast<X *>(vz);
|
||||
|
||||
// Process each input array
|
||||
for (int e = blockIdx.x; e < inputSize; e += gridDim.x) {
|
||||
auto x = reinterpret_cast<X *>(vx[e]);
|
||||
auto indices = reinterpret_cast<Y *>(vindices[e]);
|
||||
|
||||
auto xShapeInfo = xShapeInfos[e];
|
||||
auto iShapeInfo = iShapeInfos[e];
|
||||
|
||||
auto iLength = shape::length(iShapeInfo);
|
||||
|
||||
// Loop over indices in parallel
|
||||
for (int i = threadIdx.x; i < iLength; i += blockDim.x) {
|
||||
LongType iCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType iOffset, xOffset, zOffset;
|
||||
|
||||
// Compute index for indices array
|
||||
INDEX2COORDS(i, shape::rank(iShapeInfo), shape::shapeOf(iShapeInfo), iCoords);
|
||||
COORDS2INDEX(shape::rank(iShapeInfo), shape::stride(iShapeInfo), iCoords, iOffset);
|
||||
|
||||
auto idx = indices[iOffset];
|
||||
if (idx >= 0 && idx < zLength) {
|
||||
// Compute z offset
|
||||
INDEX2COORDS(idx, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
// Compute x offset
|
||||
INDEX2COORDS(i, shape::rank(xShapeInfo), shape::shapeOf(xShapeInfo), xCoords);
|
||||
COORDS2INDEX(shape::rank(xShapeInfo), shape::stride(xShapeInfo), xCoords, xOffset);
|
||||
|
||||
// Assign value to z
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
static SD_KERNEL void dynamicStitchTadKernel(void **vx, LongType **xTadShapeInfos, LongType **xTadOffsets,
|
||||
void **vindices, LongType **iShapeInfos, int inputSize, void *vz,
|
||||
const LongType *zTadShapeInfo, const LongType *zTadOffsets,
|
||||
LongType *numTadsPerInput, LongType numOutputsTad) {
|
||||
__shared__ LongType zRank, zLength, zTadLength;
|
||||
__shared__ const LongType *zShapePtr, *zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zRank = shape::rank(zTadShapeInfo);
|
||||
zLength = shape::length(zTadShapeInfo);
|
||||
zTadLength = shape::length(zTadShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zTadShapeInfo);
|
||||
zStridePtr = shape::stride(zTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto bz = reinterpret_cast<X *>(vz);
|
||||
|
||||
// Process each input array
|
||||
for (int e = threadIdx.x; e < inputSize; e += blockDim.x) {
|
||||
auto indices = reinterpret_cast<Y *>(vindices[e]);
|
||||
auto iShapeInfo = iShapeInfos[e];
|
||||
auto numTads = numTadsPerInput[e];
|
||||
|
||||
if (shape::isEmptyConst(iShapeInfo)) continue;
|
||||
|
||||
auto iLength = shape::length(iShapeInfo);
|
||||
auto xTadShapeInfo = xTadShapeInfos[e];
|
||||
auto xTadLength = shape::length(xTadShapeInfo);
|
||||
auto xShapePtr = shape::shapeOf(xTadShapeInfo);
|
||||
auto xStridePtr = shape::stride(xTadShapeInfo);
|
||||
|
||||
// Process each index in the input
|
||||
for (int i = 0; i < iLength; i++) {
|
||||
LongType iCoords[SD_MAX_RANK], iOffset;
|
||||
INDEX2COORDS(i, shape::rank(iShapeInfo), shape::shapeOf(iShapeInfo), iCoords);
|
||||
COORDS2INDEX(shape::rank(iShapeInfo), shape::stride(iShapeInfo), iCoords, iOffset);
|
||||
|
||||
auto idx = indices[iOffset];
|
||||
|
||||
// Input array offset for current TAD
|
||||
auto x = reinterpret_cast<X *>(vx[e]) + xTadOffsets[e][i];
|
||||
auto zTad = bz + zTadOffsets[idx];
|
||||
|
||||
// Copy data from input to output
|
||||
for (int j = threadIdx.x; j < xTadLength; j += blockDim.x) {
|
||||
LongType xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType xIdx, zIdx;
|
||||
|
||||
INDEX2COORDS(j, shape::rank(xTadShapeInfo), xShapePtr, xCoords);
|
||||
COORDS2INDEX(shape::rank(xTadShapeInfo), xStridePtr, xCoords, xIdx);
|
||||
|
||||
INDEX2COORDS(j, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zIdx);
|
||||
|
||||
if (xIdx < xTadLength && zIdx < zLength) {
|
||||
zTad[zIdx] = x[xIdx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
static Status _dynamicStitchFunctor(LaunchContext *context, std::vector<NDArray *> const &inputs,
|
||||
std::vector<NDArray *> const &indices, NDArray *output) {
|
||||
LongType inputSize = inputs.size();
|
||||
|
||||
PointersManager pm(context, "dynamicStitch");
|
||||
|
||||
if (output->isVector()) {
|
||||
std::vector<const void *> inputBuffers(inputSize);
|
||||
std::vector<const LongType *> inputShapes(inputSize);
|
||||
std::vector<const void *> indicesBuffers(inputSize);
|
||||
std::vector<const LongType *> indicesShapes(inputSize);
|
||||
|
||||
for (LongType e = 0; e < inputSize; e++) {
|
||||
inputBuffers[e] = inputs.at(e)->specialBuffer();
|
||||
indicesBuffers[e] = indices.at(e)->specialBuffer();
|
||||
|
||||
inputShapes[e] = inputs.at(e)->specialShapeInfo();
|
||||
indicesShapes[e] = indices.at(e)->specialShapeInfo();
|
||||
}
|
||||
|
||||
// copying pointers to buffers to device
|
||||
auto dInputBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(inputBuffers.data(), inputSize * sizeof(void *)));
|
||||
auto dIndicesBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(indicesBuffers.data(), inputSize * sizeof(void *)));
|
||||
auto dInputShapes =
|
||||
reinterpret_cast<LongType **>(pm.replicatePointer(inputShapes.data(), inputSize * sizeof(LongType *)));
|
||||
auto dIndicesShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(indicesShapes.data(), inputSize * sizeof(LongType *)));
|
||||
dim3 launchDims = getLaunchDims("dynamic_stitch_tad");
|
||||
|
||||
dynamicStitchScalarKernel<X, Y><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
||||
dInputBuffers, dInputShapes, dIndicesBuffers, dIndicesShapes, inputSize, output->specialBuffer(),
|
||||
output->specialShapeInfo(), output->lengthOf());
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"dynamicStitchScalarKernel failed");
|
||||
|
||||
} else {
|
||||
std::vector<LongType> restDims(output->rankOf() - 1);
|
||||
for (int i = restDims.size(); i > 0; i--) restDims[restDims.size() - i] = output->rankOf() - i;
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), &restDims);
|
||||
|
||||
std::vector<const void *> inputBuffers(inputSize);
|
||||
std::vector<const LongType *> inputTadShapes(inputSize);
|
||||
std::vector<const LongType *> inputTadOffsets(inputSize);
|
||||
|
||||
std::vector<const void *> indicesBuffers(inputSize);
|
||||
std::vector<const LongType *> indicesShapes(inputSize);
|
||||
std::vector<LongType> inputsNumTads(inputSize);
|
||||
|
||||
for (LongType e = 0; e < inputSize; e++) {
|
||||
std::vector<LongType> sourceDims(inputs[e]->rankOf() - indices[e]->rankOf());
|
||||
for (LongType i = sourceDims.size(); i > 0; i--) sourceDims[sourceDims.size() - i] = inputs[e]->rankOf() - i;
|
||||
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(inputs[e]->shapeInfo(), &sourceDims);
|
||||
indicesBuffers[e] = indices[e]->specialBuffer();
|
||||
indicesShapes[e] = indices[e]->specialShapeInfo();
|
||||
inputsNumTads[e] = packX->numberOfTads();
|
||||
inputBuffers[e] = inputs[e]->specialBuffer();
|
||||
inputTadShapes[e] = packX->platformShapeInfo();
|
||||
inputTadOffsets[e] = packX->platformOffsets();
|
||||
}
|
||||
|
||||
// copying pointers to buffers to device
|
||||
auto dInputBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(inputBuffers.data(), inputSize * sizeof(void *)));
|
||||
auto dInputTadShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(inputTadShapes.data(), inputSize * sizeof(LongType *)));
|
||||
auto dInputTadOffsets = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(inputTadOffsets.data(), inputSize * sizeof(LongType *)));
|
||||
|
||||
auto dIndicesBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(indicesBuffers.data(), inputSize * sizeof(void *)));
|
||||
auto dIndicesShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(indicesShapes.data(), inputSize * sizeof(LongType *)));
|
||||
|
||||
auto dNumTadsInputs = reinterpret_cast<LongType *>(
|
||||
pm.replicatePointer(inputsNumTads.data(), inputSize * sizeof(LongType *)));
|
||||
|
||||
|
||||
dim3 launchDims = getLaunchDims("dynamic_stitch_tad");
|
||||
dynamicStitchTadKernel<X, Y><<<launchDims.x, launchDims.y, launchDims.z, *context->getCudaStream()>>>(
|
||||
dInputBuffers, dInputTadShapes, dInputTadOffsets, dIndicesBuffers, dIndicesShapes, inputSize,
|
||||
output->specialBuffer(), packZ->platformShapeInfo(), packZ->platformOffsets(),dNumTadsInputs, packZ->numberOfTads());
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"dynamicStitchTadKernel failed");
|
||||
|
||||
}
|
||||
|
||||
pm.synchronize();
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void _dynamicPartitionFunctorBP(NDArray *input, NDArray *indices,
|
||||
std::vector<NDArray *> const &inputGradientList,
|
||||
std::vector<NDArray *> &outputList) {}
|
||||
|
||||
void dynamicPartitionFunctor(LaunchContext *context, NDArray *input, NDArray *indices,
|
||||
std::vector<NDArray *> &outputList) {
|
||||
auto xType = input->dataType();
|
||||
auto yType = indices->dataType();
|
||||
|
||||
NDArray::prepareSpecialUse({}, {indices, input});
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, _dynamicPartitionFunctor, (context, input, indices, outputList), SD_NUMERIC_TYPES,
|
||||
SD_INDEXING_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({}, {indices, input});
|
||||
|
||||
// TODO: it would be nice to have NDArray::registerSpecialUse signature that accepts something else beyond
|
||||
// initializer_list
|
||||
for (auto v : outputList) {
|
||||
v->tickWriteDevice();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static Status _dynamicStitchFunctorBP(std::vector<NDArray *> const &inputs, std::vector<NDArray *> const &indices,
|
||||
NDArray *gradInput, std::vector<NDArray *> &outputList) {
|
||||
THROW_EXCEPTION("Not implemented yet");
|
||||
}
|
||||
|
||||
Status dynamicStitchFunctor(LaunchContext *context, std::vector<NDArray *> const &inputs,
|
||||
std::vector<NDArray *> const &indices, NDArray *output) {
|
||||
auto xType = inputs.at(0)->dataType();
|
||||
auto yType = indices.at(0)->dataType();
|
||||
|
||||
for (auto v : indices) {
|
||||
v->syncToDevice();
|
||||
v->tickReadDevice();
|
||||
}
|
||||
|
||||
for (auto v : inputs) {
|
||||
v->syncToDevice();
|
||||
v->tickReadDevice();
|
||||
}
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {});
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, _dynamicStitchFunctor, (context, inputs, indices, output), SD_NUMERIC_TYPES,
|
||||
SD_INDEXING_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({output}, {});
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
Status dynamicStitchFunctorBP(LaunchContext *context, std::vector<NDArray *> const &inputs,
|
||||
std::vector<NDArray *> const &indices, NDArray *gradInput,
|
||||
std::vector<NDArray *> &outputList) {
|
||||
auto xType = inputs.at(0)->dataType();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, return _dynamicStitchFunctorBP, (inputs, indices, gradInput, outputList),
|
||||
SD_NUMERIC_TYPES);
|
||||
}
|
||||
|
||||
void dynamicPartitionFunctorBP(LaunchContext *context, NDArray *input, NDArray *indices,
|
||||
std::vector<NDArray *> const &inputGradientList, std::vector<NDArray *> &outputList) {
|
||||
auto xType = input->dataType();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, _dynamicPartitionFunctorBP, (input, indices, inputGradientList, outputList),
|
||||
SD_NUMERIC_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,184 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
// @author sgazeos@gmail.com
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
|
||||
#include <ops/declarable/helpers/axis.h>
|
||||
#include <execution/Threads.h>
|
||||
|
||||
#include <array>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// extract patches kernel
|
||||
// - theSame - SAME or VALID - output format
|
||||
// - batchCount - batches - the first dimension of input
|
||||
// - sizeRow, sizeCol - rows and cols sizes for batch
|
||||
// - rowDim, colDim - rows and cols dimensions for input patches
|
||||
// - outRowDim, outColDim - rows and cols dimensions for output patches
|
||||
// - strideRow, strideCol - step between input elements with patches
|
||||
// - rateRow, rateCol - counts for input patches
|
||||
// - rowCast, colCast - shifts for output placement (1 or 0)
|
||||
// - lastDim - last dimension of input/output
|
||||
// - input - input tensor buffer
|
||||
// - patchShape - input patch TAD shape
|
||||
// - inputOffsets - input TAD offsets
|
||||
// - output - output tensor buffer
|
||||
// - outTadShape - output TAD shape
|
||||
// - outputOffsets - output TAD offsets
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void globalExtractPatchesKernel(bool theSame, int batchCount,
|
||||
int sizeRow, int sizeCol, int rowDim,
|
||||
int colDim, int outRowDim,
|
||||
int outColDim, int strideRow,
|
||||
int strideCol,
|
||||
int rateRow, int rateCol,
|
||||
int rowCast,
|
||||
int colCast, int lastDim,
|
||||
const T* input,
|
||||
const LongType* patchShape,
|
||||
const LongType* inputOffsets,
|
||||
T* output,
|
||||
const LongType* outTadShape,
|
||||
const LongType* outputOffsets) {
|
||||
|
||||
for (auto batch = threadIdx.x; batch < batchCount; batch+= gridDim.x) {
|
||||
auto patch = input + inputOffsets[batch];
|
||||
auto outMatrix = output + outputOffsets[batch];
|
||||
|
||||
for (LongType i = 0; i < outRowDim; i++) {
|
||||
for (LongType j = 0; j < outColDim; j++) {
|
||||
LongType pos = 0;
|
||||
auto rowStart = i * strideRow - (theSame ? rowCast : 0);
|
||||
auto colStart = j * strideCol - (theSame ? colCast : 0);
|
||||
auto rowEnd = rowStart + sizeRow * rateRow;
|
||||
auto colEnd = colStart + sizeCol * rateCol;
|
||||
if (!theSame) {
|
||||
rowEnd = math::sd_min<T>(rowStart + sizeRow * rateRow, rowDim);
|
||||
colEnd = math::sd_min<T>(colStart + sizeCol * rateCol, colDim);
|
||||
}
|
||||
for (auto row = rowStart; row < rowEnd; row += rateRow)
|
||||
for (auto col = colStart; col < colEnd; col += rateCol)
|
||||
for (auto pixel = 0; pixel < lastDim; pixel++) {
|
||||
LongType zPos[] = {i, j, pos};
|
||||
LongType xPos[] = {row, col, pixel};
|
||||
bool setUp = (theSame && row >= 0 && col >= 0 && row < rowDim && col < colDim) || (!theSame);
|
||||
|
||||
if (setUp) { // VALID or SAME cases
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(shape::rank(outTadShape), shape::stride(outTadShape), zPos, zOffset);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(shape::rank(patchShape), shape::stride(patchShape), xPos, xOffset);
|
||||
outMatrix[zOffset] = patch[xOffset];
|
||||
}
|
||||
pos++;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void _extractPatches(LaunchContext* context, NDArray* images, NDArray* output, int sizeRow, int sizeCol,
|
||||
int strideRow, int strideCol, int rateRow, int rateCol, bool theSame) {
|
||||
std::vector<LongType> restDims({1, 2, 3}); // the first and the last dims
|
||||
ResultSet listOfMatricies = images->allTensorsAlongDimension(restDims);
|
||||
ResultSet listOfOutputs = output->allTensorsAlongDimension(restDims);
|
||||
// 3D matrices - 2D matrices of vectors (if last dim is greater than 1)
|
||||
|
||||
int batchCount = listOfMatricies.size();
|
||||
LongType lastDim = images->sizeAt(3);
|
||||
LongType rowDim = images->sizeAt(1);
|
||||
LongType colDim = images->sizeAt(2);
|
||||
LongType outRowDim = output->sizeAt(1);
|
||||
LongType outColDim = output->sizeAt(2);
|
||||
auto rowCast = 1;
|
||||
auto colCast = 1;
|
||||
if (sizeRow * rateRow < 3) rowCast = 0;
|
||||
if (sizeCol * rateCol < 3) colCast = 0;
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto batch = 0; batch < stop; batch++) {
|
||||
auto patch = listOfMatricies.at(batch);
|
||||
std::vector<sd::LongType> inShape = {patch->sizeAt(1), patch->sizeAt(2), patch->sizeAt(3)};
|
||||
auto inPatch = patch->rankOf() > 3 && patch->sizeAt(0) == 1 ? patch->reshape('c', inShape, false) : patch;
|
||||
auto outMatrix = listOfOutputs.at(batch);
|
||||
std::vector<sd::LongType> outShape = {outMatrix->sizeAt(1), outMatrix->sizeAt(2), outMatrix->sizeAt(3)};
|
||||
auto outReshape = outMatrix->rankOf() > 3 && outMatrix->sizeAt(0) == 1 ? outMatrix->reshape('c', outShape, false) : outMatrix;
|
||||
for (LongType i = 0; i < outRowDim; i++) {
|
||||
for (LongType j = 0; j < outColDim; j++) {
|
||||
LongType pos = 0;
|
||||
auto rowStart = i * strideRow - (theSame ? rowCast : 0);
|
||||
auto colStart = j * strideCol - (theSame ? colCast : 0);
|
||||
auto rowEnd = rowStart + sizeRow * rateRow;
|
||||
auto colEnd = colStart + sizeCol * rateCol;
|
||||
if (!theSame) {
|
||||
rowEnd = math::sd_min(rowStart + sizeRow * rateRow, rowDim);
|
||||
colEnd = math::sd_min(colStart + sizeCol * rateCol, colDim);
|
||||
}
|
||||
for (auto row = rowStart; row < rowEnd; row += rateRow)
|
||||
for (auto col = colStart; col < colEnd; col += rateCol)
|
||||
for (auto pixel = 0; pixel < lastDim; pixel++) {
|
||||
bool setUp = (theSame && row >= 0
|
||||
&& col >= 0 && row < rowDim
|
||||
&& col < colDim)
|
||||
|| (!theSame);
|
||||
if (setUp) {
|
||||
outReshape->p<T>(i,j,pos,inPatch->e<T>(row, col, pixel));
|
||||
}
|
||||
pos++;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Cleanup reshaped arrays if they were created
|
||||
if (inPatch != patch) delete inPatch;
|
||||
if (outReshape != outMatrix) delete outReshape;
|
||||
}
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_tad(func, 0, batchCount);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void _extractPatches,
|
||||
(sd::LaunchContext * context, NDArray* input, NDArray* output, int sizeRow, int sizeCol,
|
||||
int stradeRow, int stradeCol, int rateRow, int rateCol, bool theSame),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
void extractPatches(LaunchContext* context, NDArray* images, NDArray* output, int sizeRow, int sizeCol,
|
||||
int stradeRow, int stradeCol, int rateRow, int rateCol, bool theSame) {
|
||||
auto xType = images->dataType();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(xType, _extractPatches,
|
||||
(context, images, output, sizeRow, sizeCol, stradeRow, stradeCol, rateRow, rateCol, theSame),
|
||||
SD_COMMON_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,179 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <ops/declarable/helpers/fake_quantization.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// fakeQuantWithMinMaxVars_
|
||||
// input - input tensor
|
||||
// min - min scalar tensor
|
||||
// max - max scalar tensor
|
||||
// numBits - (default 16bit)
|
||||
// narrowed - shrink is true
|
||||
// output - output tensor
|
||||
//
|
||||
template <typename T>
|
||||
static SD_HOST_DEVICE void nudge(T min, T max, int quantMin, int quantMax, T* scale, T* nudgedMin, T* nudgedMax) {
|
||||
T quantMaxF = static_cast<T>(quantMax);
|
||||
T quantMinF = static_cast<T>(quantMin);
|
||||
*scale = (max - min) / (quantMaxF - quantMinF);
|
||||
auto zeroPointFromMin = quantMinF - min / *scale;
|
||||
uint16_t const nudgedZeroPoint = [zeroPointFromMin, quantMin, quantMax, quantMaxF, quantMinF] {
|
||||
if (zeroPointFromMin < quantMinF) {
|
||||
return static_cast<uint16_t>(quantMin);
|
||||
}
|
||||
if (zeroPointFromMin > quantMaxF) {
|
||||
return static_cast<uint16_t>(quantMax);
|
||||
}
|
||||
return math::sd_round<T, uint16_t>(zeroPointFromMin);
|
||||
}();
|
||||
*nudgedMax = (quantMaxF - static_cast<T>(nudgedZeroPoint)) * (*scale);
|
||||
*nudgedMin = (quantMinF - static_cast<T>(nudgedZeroPoint)) * (*scale);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void fakeQuantWithMinMaxVars_(NDArray* input, NDArray* min, NDArray* max, int numBits, bool narrowed, NDArray* output) {
|
||||
int lowIntBound = narrowed ? 1 : 0;
|
||||
int upperIntBound = (1 << numBits) - 1;
|
||||
min->syncToHost(); // these are scalars, so nothing much happened
|
||||
max->syncToHost();
|
||||
T scale, nudgedMin, nudgedMax;
|
||||
nudge(min->t<T>(0), max->t<T>(0), lowIntBound, upperIntBound, &scale, &nudgedMin, &nudgedMax);
|
||||
|
||||
auto wiseMinMaxAndSoOn = LAMBDA_T(x, nudgedMin, nudgedMax, scale) {
|
||||
T val = x;
|
||||
if (x < nudgedMin) {
|
||||
val = nudgedMin;
|
||||
} else if (x > nudgedMax) {
|
||||
val = nudgedMax;
|
||||
} else
|
||||
val = x;
|
||||
return (math::sd_floor<T, T>((val - nudgedMin) / scale + T(0.5)) * scale + nudgedMin);
|
||||
});
|
||||
|
||||
input->applyLambda(wiseMinMaxAndSoOn, output);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void fakeQuantWithMinMaxKernel(const T* input, const LongType* inputShape, T* min, T* max,
|
||||
int lowIntBound, int upperIntBound, LongType channels, T* output,
|
||||
const LongType* outputShape, LongType length) {
|
||||
__shared__ LongType inputRank, outputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
__shared__ LongType blockSize;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
blockSize = length / channels; // Calculate block size based on the last dimension
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputOffset;
|
||||
LongType outputOffset;
|
||||
|
||||
// Loop over channels
|
||||
for (auto i = blockIdx.x; i < (int)channels; i += gridDim.x) {
|
||||
T scale, nudgedMin, nudgedMax;
|
||||
|
||||
// Nudge values for quantization
|
||||
nudge(min[i], max[i], lowIntBound, upperIntBound, &scale, &nudgedMin, &nudgedMax);
|
||||
|
||||
// Loop over blocks for quantization
|
||||
for (auto b = threadIdx.x; b < blockSize; b += blockDim.x) {
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(b * channels + i, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, inputOffset);
|
||||
|
||||
T val = input[inputOffset];
|
||||
|
||||
// Clamp value within nudged min and max
|
||||
if (val < nudgedMin) {
|
||||
val = nudgedMin;
|
||||
} else if (val > nudgedMax) {
|
||||
val = nudgedMax;
|
||||
}
|
||||
|
||||
// Compute output coordinates and offset
|
||||
INDEX2COORDS(b * channels + i, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, outputOffset);
|
||||
|
||||
// Quantize and assign the value to output
|
||||
output[outputOffset] = math::sd_floor<T, T>((val - nudgedMin) / scale + T(0.5f)) * scale + nudgedMin;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
void fakeQuantWithMinMaxVarsPerChannel_(LaunchContext* context, NDArray* input, NDArray* min, NDArray* max, int numBits,
|
||||
bool narrowed, NDArray* output) {
|
||||
int lowIntBound = narrowed ? 1 : 0;
|
||||
int upperIntBound = (1 << numBits) - 1;
|
||||
auto channels = min->lengthOf();
|
||||
auto length = input->lengthOf();
|
||||
NDArray::prepareSpecialUse({output}, {min, max, input});
|
||||
auto stream = context->getCudaStream();
|
||||
T* inputBuf = input->dataBuffer()->specialAsT<T>();
|
||||
T* outputBuf = output->dataBuffer()->specialAsT<T>();
|
||||
T* minBuf = min->dataBuffer()->specialAsT<T>();
|
||||
T* maxBuf = max->dataBuffer()->specialAsT<T>();
|
||||
dim3 launchDims = getLaunchDims("fake_quantization");
|
||||
fakeQuantWithMinMaxKernel<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(inputBuf, input->specialShapeInfo(), minBuf, maxBuf,
|
||||
lowIntBound, upperIntBound, channels, outputBuf,
|
||||
output->specialShapeInfo(), length);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"fakeQuantWithMinMaxKernel failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {min, max, input});
|
||||
}
|
||||
|
||||
void fakeQuantWithMinMaxVars(NDArray* input, NDArray* min, NDArray* max, int numBits, bool narrowed, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), fakeQuantWithMinMaxVars_, (input, min, max, numBits, narrowed, output),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
void fakeQuantWithMinMaxVarsPerChannel(LaunchContext* context, NDArray* input, NDArray* min, NDArray* max, int numBits,
|
||||
bool narrowed, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), fakeQuantWithMinMaxVarsPerChannel_,
|
||||
(context, input, min, max, numBits, narrowed, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,109 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/flatten.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
static void SD_KERNEL flattenKernel(void **xBuffers, LongType **xShapeInfos, LongType *offsets, LongType numInputs, void *zBuffer, const LongType *zShapeInfo, char order) {
|
||||
__shared__ LongType xRank, xLength;
|
||||
__shared__ const LongType *xShapePtr, *xStridePtr;
|
||||
|
||||
int xCoord[SD_MAX_RANK];
|
||||
|
||||
// Each block of threads works on one input array
|
||||
for (LongType e = blockIdx.x; e < numInputs; e += gridDim.x) {
|
||||
auto z = reinterpret_cast<T *>(zBuffer) + offsets[e];
|
||||
|
||||
auto xBuffer = reinterpret_cast<T *>(xBuffers[e]);
|
||||
auto xShapeInfo = xShapeInfos[e];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xLength = shape::length(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each element of this input array has its own place within the common output array
|
||||
for (LongType i = threadIdx.x; i < xLength; i += blockDim.x) {
|
||||
LongType xOffset;
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
|
||||
// Compute x coordinates and offset
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
// Write the value from xBuffer to the flattened zBuffer
|
||||
z[i] = xBuffer[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void flatten_(LaunchContext *context, std::vector<NDArray *> &inputs, NDArray *output, char order) {
|
||||
PointersManager pm(context, "flatten");
|
||||
|
||||
std::vector<const void *> hdBuffers(inputs.size());
|
||||
std::vector<LongType> hOffsets(inputs.size());
|
||||
std::vector<const LongType *> hdShapes(inputs.size());
|
||||
LongType cOffset = 0;
|
||||
|
||||
// calculating offsets in output
|
||||
for (int e = 0; e < inputs.size(); e++) {
|
||||
hOffsets[e] = cOffset;
|
||||
cOffset += inputs[e]->lengthOf();
|
||||
|
||||
hdBuffers[e] = inputs[e]->specialBuffer();
|
||||
hdShapes[e] = inputs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
// copying pointers to device
|
||||
auto dBuffers = (void **)pm.replicatePointer(hdBuffers.data(), inputs.size() * sizeof(void *));
|
||||
auto dShapes = (LongType **)pm.replicatePointer(hdShapes.data(), inputs.size() * sizeof(LongType *));
|
||||
auto dOffsets = (LongType *)pm.replicatePointer(hOffsets.data(), inputs.size() * sizeof(LongType));
|
||||
dim3 launchDims = getLaunchDims("flatten");
|
||||
flattenKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *context->getCudaStream()>>>(
|
||||
dBuffers, dShapes, dOffsets, inputs.size(), output->specialBuffer(), output->specialShapeInfo(), order);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"flattenKernel failed");
|
||||
|
||||
pm.synchronize();
|
||||
}
|
||||
|
||||
void flatten(LaunchContext *context, std::vector<NDArray *> &inputs, NDArray *output, char order) {
|
||||
// FIXME: we want NDArrayFactory::prepareSpecialUse here eventually
|
||||
const std::vector<NDArray *> v(inputs.begin(), inputs.end());
|
||||
//prepareSpecialUse requires const
|
||||
NDArray::prepareSpecialUse({output}, v, {});
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), flatten_, (context, inputs, output, order), SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,266 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 07.03.2019
|
||||
//
|
||||
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/gather.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void gatherCudaLinearKernel(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo) {
|
||||
__shared__ const X* x;
|
||||
__shared__ const Y* y;
|
||||
__shared__ X* z;
|
||||
__shared__ LongType xLen, yLen, zLen;
|
||||
__shared__ LongType xRank, yRank, zRank;
|
||||
__shared__ const LongType *xShapePtr, *xStridePtr;
|
||||
__shared__ const LongType *yShapePtr, *yStridePtr;
|
||||
__shared__ const LongType *zShapePtr, *zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
x = reinterpret_cast<const X*>(vx);
|
||||
z = reinterpret_cast<X*>(vz);
|
||||
y = reinterpret_cast<const Y*>(vy);
|
||||
|
||||
xLen = shape::length(xShapeInfo);
|
||||
yLen = shape::length(yShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (LongType j = start; j < zLen; j += step) {
|
||||
LongType zIndex, yIndex, xIndex;
|
||||
LongType zCoords[SD_MAX_RANK], yCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
||||
|
||||
// Compute z coordinates and offset
|
||||
INDEX2COORDS(j, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zIndex);
|
||||
|
||||
// Compute y coordinates and offset
|
||||
INDEX2COORDS(j, yRank, yShapePtr, yCoords);
|
||||
COORDS2INDEX(yRank, yStridePtr, yCoords, yIndex);
|
||||
|
||||
// Compute x coordinates and offset
|
||||
INDEX2COORDS(y[yIndex], xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xIndex);
|
||||
|
||||
// Assign value to z
|
||||
z[zIndex] = x[xIndex];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void gatherCuda(const int numOfSubArrs, const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xOffsets, const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zOffsets) {
|
||||
const Y* y = reinterpret_cast<const Y*>(vy);
|
||||
|
||||
__shared__ const X* x;
|
||||
__shared__ X* z;
|
||||
__shared__ LongType xLen, yRank, xRank, zRank;
|
||||
__shared__ const LongType *xShapePtr, *xStridePtr, *yShapePtr, *yStridePtr, *zShapePtr, *zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType i = blockIdx.x; i < numOfSubArrs; i += gridDim.x) {
|
||||
if (threadIdx.x == 0) {
|
||||
LongType yIndex, xOffset, zOffset;
|
||||
LongType yCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
|
||||
// Calculate y index
|
||||
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
|
||||
COORDS2INDEX(yRank, yStridePtr, yCoords, yIndex);
|
||||
|
||||
// Calculate x offset
|
||||
INDEX2COORDS(y[yIndex], xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
// Calculate z offset
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
x = reinterpret_cast<const X*>(vx) + xOffsets[xOffset];
|
||||
z = reinterpret_cast<X*>(vz) + zOffsets[zOffset];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Copy data from x to z
|
||||
for (LongType j = threadIdx.x; j < xLen; j += blockDim.x) {
|
||||
LongType zIndex, xIndex;
|
||||
LongType zCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
||||
|
||||
// Calculate z index
|
||||
INDEX2COORDS(j, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zIndex);
|
||||
|
||||
// Calculate x index
|
||||
INDEX2COORDS(j, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xIndex);
|
||||
|
||||
// Copy value
|
||||
z[zIndex] = x[xIndex];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_HOST static void gatherCudaLinear(const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo) {
|
||||
//note gather linear and gather are different kernels
|
||||
dim3 gatherLinear = getLaunchDims("gather_linear");
|
||||
gatherCudaLinearKernel<X, Y><<<gatherLinear.x, gatherLinear.y, gatherLinear.z, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"gatherCudaLinearKernel failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST static void gatherCudaLauncher(const cudaStream_t* stream, const int numOfSubArrs, const void* vx,
|
||||
const LongType* xShapeInfo, const LongType* xOffsets, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zOffsets) {
|
||||
dim3 gatherLinear = getGatherLinear(numOfSubArrs);
|
||||
gatherCuda<X, Y><<<gatherLinear.y, gatherLinear.x, gatherLinear.z, *stream>>>(numOfSubArrs, vx, xShapeInfo, xOffsets, vy,
|
||||
yShapeInfo, vz, zShapeInfo, zOffsets);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"gatherCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
void gather(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output,
|
||||
const std::vector<LongType>& intArgs) {
|
||||
const LongType inputRank = input->rankOf();
|
||||
const LongType numOfIntArgs = intArgs.size();
|
||||
|
||||
LongType axis = numOfIntArgs > 0 ? intArgs[0] : 0;
|
||||
if (axis < 0) axis += inputRank;
|
||||
|
||||
if (indices == nullptr && numOfIntArgs == 2) { // scalar case
|
||||
NDArray scalar = (*input)(intArgs[1], {axis});
|
||||
output->assign(&scalar);
|
||||
} else if (indices != nullptr && indices->isScalar()) {
|
||||
if (input->rankOf() <= 1) { // For scalar indices, rank 0 or 1 input: can't do tensor along dimension 0 as this is
|
||||
// whole array... instead, we want to get a scalar
|
||||
auto idx = indices->e<LongType>(0);
|
||||
auto scalarNDArray = input->e(idx);
|
||||
output->assign(&scalarNDArray);
|
||||
} else {
|
||||
NDArray inSubArr = (*input)(indices->e<LongType>(0), {axis});
|
||||
output->assign(&inSubArr);
|
||||
}
|
||||
} else {
|
||||
NDArray* pIndices = const_cast<NDArray*>(indices);
|
||||
if (indices == nullptr) {
|
||||
std::vector<LongType> firstShape = {numOfIntArgs - 1};
|
||||
std::vector<double> data = std::vector<double>(intArgs.begin() + 1, intArgs.end());
|
||||
pIndices = new NDArray(input->ordering(),firstShape,
|
||||
data, INT64, input->getContext());
|
||||
}
|
||||
std::vector<LongType> dimsOut(pIndices->rankOf());
|
||||
std::iota(dimsOut.begin(), dimsOut.end(), axis); // fill with axis, axis+1, ... axis+pIndices->rankOf()-1
|
||||
|
||||
const LongType numOfSubArrs = pIndices->lengthOf();
|
||||
|
||||
LongType *outSubArrShapeInfo(nullptr), *inSubArrShapeInfo(nullptr), *outSubArrOffsets(nullptr),
|
||||
*inSubArrOffsets(nullptr);
|
||||
input->getSubArrShapeAndOffsets({axis}, inSubArrShapeInfo, inSubArrOffsets);
|
||||
output->getSubArrShapeAndOffsets(dimsOut, outSubArrShapeInfo, outSubArrOffsets);
|
||||
if (output->rankOf() > 1) {
|
||||
PointersManager manager(context, "gather");
|
||||
auto xShapeInfo = reinterpret_cast<LongType*>(
|
||||
manager.replicatePointer(inSubArrShapeInfo, shape::shapeInfoByteLength(inSubArrShapeInfo)));
|
||||
auto zShapeInfo = reinterpret_cast<LongType*>(
|
||||
manager.replicatePointer(outSubArrShapeInfo, shape::shapeInfoByteLength(outSubArrShapeInfo)));
|
||||
auto xOffsets = reinterpret_cast<LongType*>(manager.replicatePointer(
|
||||
inSubArrOffsets, (input->lengthOf() / shape::length(inSubArrShapeInfo)) * sizeof(LongType)));
|
||||
auto zOffsets = reinterpret_cast<LongType*>(manager.replicatePointer(
|
||||
outSubArrOffsets, (output->lengthOf() / shape::length(outSubArrShapeInfo)) * sizeof(LongType)));
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input, pIndices});
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
input->dataType(), pIndices->dataType(), gatherCudaLauncher,
|
||||
(context->getCudaStream(), numOfSubArrs, input->specialBuffer(), xShapeInfo, xOffsets,
|
||||
pIndices->specialBuffer(), pIndices->specialShapeInfo(), output->specialBuffer(), zShapeInfo, zOffsets),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, pIndices});
|
||||
manager.synchronize();
|
||||
} else {
|
||||
NDArray::prepareSpecialUse({output}, {input, pIndices});
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
input->dataType(), pIndices->dataType(), gatherCudaLinear,
|
||||
(context->getCudaStream(), input->specialBuffer(), input->specialShapeInfo(), pIndices->specialBuffer(),
|
||||
pIndices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo()),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, pIndices});
|
||||
}
|
||||
|
||||
if (indices == nullptr) delete pIndices;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,159 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - input, y - indices, z - output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void gatherNDCuda(const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
const auto y = reinterpret_cast<const Y *>(vy);
|
||||
auto z = reinterpret_cast<X *>(vz);
|
||||
|
||||
__shared__ int xRank, yRank, zRank, maxRank, yLastDim;
|
||||
__shared__ LongType zLen, totalThreads;
|
||||
__shared__ const LongType *xShapePtr, *xStridePtr;
|
||||
__shared__ const LongType *yShapePtr, *yStridePtr;
|
||||
__shared__ const LongType *zShapePtr, *zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
maxRank = sd::math::sd_max<int>(yRank, sd::math::sd_max<int>(xRank, zRank));
|
||||
|
||||
zLen = shape::length(zShapeInfo);
|
||||
yLastDim = shape::shapeOf(yShapeInfo)[yRank - 1];
|
||||
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
extern __shared__ unsigned char shmem[];
|
||||
auto coord = reinterpret_cast<LongType *>(shmem) + threadIdx.x * maxRank;
|
||||
|
||||
LongType *zCoordStart, *xCoordStart;
|
||||
|
||||
if (yLastDim == xRank) {
|
||||
zCoordStart = coord;
|
||||
xCoordStart = coord;
|
||||
} else if (zRank >= xRank) {
|
||||
zCoordStart = coord;
|
||||
xCoordStart = coord + zRank - xRank;
|
||||
} else {
|
||||
zCoordStart = coord + xRank - zRank;
|
||||
xCoordStart = coord;
|
||||
}
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < zLen; i += totalThreads) {
|
||||
// Compute z coordinates and offset
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoordStart);
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoordStart, zOffset);
|
||||
|
||||
// Save and modify last y coordinate
|
||||
int coordToRestore = (yLastDim != xRank) ? static_cast<int>(zCoordStart[yRank - 1]) : 0;
|
||||
zCoordStart[yRank - 1] = 0;
|
||||
|
||||
// Compute y offset
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(yRank, yStridePtr, zCoordStart, yOffset);
|
||||
|
||||
// Restore z coordinate
|
||||
if (yLastDim != xRank) zCoordStart[yRank - 1] = coordToRestore;
|
||||
|
||||
// Compute x coordinates
|
||||
for (LongType j = 0; j < yLastDim; ++j) {
|
||||
xCoordStart[j] = y[yOffset + j * yStridePtr[yRank - 1]];
|
||||
}
|
||||
|
||||
// Compute x offset
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoordStart, xOffset);
|
||||
|
||||
// Assign value to z
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void gatherNDCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo,
|
||||
const void *vy, const LongType *yShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo) {
|
||||
gatherNDCuda<X, Y>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"gatherNDCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void gatherND(LaunchContext *context, NDArray &input, NDArray &indices, NDArray &output) {
|
||||
const int maxRank = sd::math::sd_max<int>(indices.rankOf(), sd::math::sd_max<int>(input.rankOf(), output.rankOf()));
|
||||
|
||||
|
||||
dim3 gatherNdDims = getGatherNd(output.lengthOf(),maxRank);
|
||||
const auto xType = input.dataType();
|
||||
const auto yType = indices.dataType();
|
||||
|
||||
PointersManager manager(context, "gatherND");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input, &indices});
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, gatherNDCudaLauncher,
|
||||
(gatherNdDims.y, gatherNdDims.x, gatherNdDims.z, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), indices.specialBuffer(), indices.specialShapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo()),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input, &indices});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,47 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <ops/declarable/helpers/axis.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
void applyGradientDescent_(LaunchContext* context, NDArray* input, NDArray* step, double weight,
|
||||
NDArray* output) {
|
||||
// classic one
|
||||
auto lambda = LAMBDA_TT(_x, _y, weight) { return _x - (_y * weight); });
|
||||
|
||||
input->applyPairwiseLambda(step, lambda, output);
|
||||
}
|
||||
|
||||
void applyGradientDescent(LaunchContext* context, NDArray* input, NDArray* step, double weight, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), applyGradientDescent_, (context, input, step, weight, output),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void applyGradientDescent_,
|
||||
(LaunchContext * context, NDArray* input, NDArray* step, double weight, NDArray* output),
|
||||
SD_FLOAT_TYPES);
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,120 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <ops/declarable/helpers/hamming.h>
|
||||
#include <ops/declarable/helpers/helpers.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void _hammingKernel(const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz, void *reductionBuffer, LongType length) {
|
||||
auto x = reinterpret_cast<const X *>(vx);
|
||||
auto y = reinterpret_cast<const X *>(vy);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
|
||||
__shared__ LongType shared[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ LongType xRank, yRank;
|
||||
__shared__ const LongType *xShapePtr, *xStridePtr;
|
||||
__shared__ const LongType *yShapePtr, *yStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Initialize shared memory for intermediate results
|
||||
shared[threadIdx.x] = 0;
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
for (LongType e = tid; e < length; e += blockDim.x * gridDim.x) {
|
||||
LongType xCoords[SD_MAX_RANK], yCoords[SD_MAX_RANK];
|
||||
LongType xOffset, yOffset;
|
||||
|
||||
// Calculate coordinates and offsets using cached values
|
||||
INDEX2COORDS(e, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(e, yRank, yShapePtr, yCoords);
|
||||
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
|
||||
|
||||
auto _x = static_cast<unsigned long long>(x[xOffset]);
|
||||
auto _y = static_cast<unsigned long long>(y[yOffset]);
|
||||
|
||||
// Save intermediate results into shared memory
|
||||
shared[threadIdx.x] += __popcll(_x ^ _y);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Reduction within a block
|
||||
auto numItems = sd::math::sd_min<LongType>(blockDim.x, length);
|
||||
auto floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) floorPow2 &= floorPow2 - 1;
|
||||
|
||||
if (threadIdx.x >= floorPow2)
|
||||
shared[threadIdx.x - floorPow2] += shared[threadIdx.x];
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (threadIdx.x < activeThreads && threadIdx.x + activeThreads < numItems)
|
||||
shared[threadIdx.x] += shared[threadIdx.x + activeThreads];
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Write the final result to global memory
|
||||
if (threadIdx.x == 0 && shared[0] > 0) {
|
||||
math::atomics::sd_atomicAdd<Z>(&z[0], static_cast<Z>(shared[0]));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Z>
|
||||
static void _hamming(LaunchContext *context, NDArray &x, NDArray &y, NDArray &z) {
|
||||
dim3 launchDims = getLaunchDims("hamming");
|
||||
_hammingKernel<X, Z><<<launchDims.x,launchDims.y, launchDims.z, *context->getCudaStream()>>>(
|
||||
x.specialBuffer(), x.specialShapeInfo(), y.specialBuffer(), y.specialShapeInfo(), z.specialBuffer(), nullptr,
|
||||
x.lengthOf());
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"_hammingKernel failed");
|
||||
|
||||
}
|
||||
|
||||
void hamming(LaunchContext *context, NDArray &x, NDArray &y, NDArray &output) {
|
||||
NDArray::prepareSpecialUse({&output}, {&x, &y});
|
||||
BUILD_DOUBLE_SELECTOR(x.dataType(), output.dataType(), _hamming, (context, x, y, output), SD_INTEGER_TYPES,
|
||||
SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&x, &y});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,137 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <ops/declarable/helpers/hashcode.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
static SD_KERNEL void splitBufferToChuncks(T* buffer, LongType* tempBuffer, LongType numBlocks, LongType blockSize,
|
||||
LongType length) {
|
||||
for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < numBlocks; b += gridDim.x * blockDim.x) {
|
||||
auto blockBuffer = buffer + b * numBlocks;
|
||||
|
||||
LongType r = 1LL;
|
||||
for (int e = 0; e < blockSize && e + (b * numBlocks) < length; e++) {
|
||||
auto v = longBytes<T>(blockBuffer[e]);
|
||||
r = 31LL * r + v;
|
||||
}
|
||||
|
||||
tempBuffer[b] = r;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void internalHash(LongType* tempBuffer, LongType* tempResult, LongType numBlocks, LongType blockSize,
|
||||
LongType lastLength) {
|
||||
for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < numBlocks; b += gridDim.x * blockDim.x) {
|
||||
auto blockBuffer = tempBuffer + b * numBlocks;
|
||||
LongType r = 1LL;
|
||||
|
||||
for (LongType e = 0; e < blockSize && e + (b * numBlocks) < lastLength; e++) {
|
||||
auto v = longBytes<T>(blockBuffer[e]);
|
||||
r = 31LL * r + v;
|
||||
}
|
||||
|
||||
tempResult[b] = r;
|
||||
}
|
||||
}
|
||||
|
||||
static SD_KERNEL void lastStep(LongType* resultBuf, LongType* tempBufferA, LongType* tempResult, LongType length,
|
||||
LongType blockSize) {
|
||||
if (threadIdx.x == 0) {
|
||||
if (length <= blockSize)
|
||||
*resultBuf = *tempBufferA;
|
||||
else
|
||||
*resultBuf = *tempResult;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void hashCode_(LaunchContext* context, NDArray& array, NDArray& result) {
|
||||
auto blockSize = 32;
|
||||
auto stream = context->getCudaStream();
|
||||
array.syncToDevice();
|
||||
|
||||
NDArray::prepareSpecialUse({&result}, {&array});
|
||||
auto length = array.lengthOf();
|
||||
int numBlocks = length / blockSize + ((length % blockSize == 0) ? 0 : 1);
|
||||
auto tempA = NDArrayFactory::create<LongType>('c', {numBlocks}, context);
|
||||
auto tempB = NDArrayFactory::create<LongType>('c', {numBlocks / blockSize + 1}, context);
|
||||
|
||||
auto buffer = reinterpret_cast<T*>(array.specialBuffer()); // bufferAsT<T>();
|
||||
auto tempBufferA = reinterpret_cast<LongType*>(tempA.specialBuffer()); // bufferAsT<sd::LongType>();
|
||||
auto tempBufferB = reinterpret_cast<LongType*>(tempB.specialBuffer()); // bufferAsT<sd::LongType>();
|
||||
|
||||
dim3 launchDims = getHashCodeSplit(length,numBlocks);
|
||||
// default buffer is the first one, because it might be the last one in case of small arrays (< blockSize)
|
||||
auto tempBuffer = tempBufferA;
|
||||
auto tempResult = tempBufferB;
|
||||
|
||||
// we divide array into 32 element chunks, and store intermediate results once
|
||||
splitBufferToChuncks<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(buffer, tempBuffer, numBlocks, blockSize, length);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"splitBufferToChuncks failed");
|
||||
|
||||
// we replace pointer with intermediate one, and repeat only one chunk left
|
||||
int iterationCount = 0;
|
||||
while (numBlocks > 1) {
|
||||
int lastLength = numBlocks;
|
||||
numBlocks = lastLength / blockSize + ((lastLength % blockSize == 0) ? 0 : 1);
|
||||
|
||||
dim3 internalLaunchDims = getHashCodeInternal(numBlocks);
|
||||
internalHash<LongType>
|
||||
<<<internalLaunchDims.y,internalLaunchDims.x, internalLaunchDims.z, *stream>>>(tempBuffer, tempResult, numBlocks, blockSize, lastLength);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"internalHash failed");
|
||||
|
||||
iterationCount++;
|
||||
// swapping buffers
|
||||
if (iterationCount % 2 == 0) {
|
||||
tempBuffer = tempBufferA;
|
||||
tempResult = tempBufferB;
|
||||
} else {
|
||||
tempBuffer = tempBufferB;
|
||||
tempResult = tempBufferA;
|
||||
}
|
||||
}
|
||||
|
||||
dim3 lastDims = getLaunchDims("hashcode_last");
|
||||
lastStep<<<lastDims.x, lastDims.y, lastDims.z, *stream>>>(reinterpret_cast<LongType*>(result.specialBuffer()), tempBufferA, tempResult,
|
||||
length, blockSize);
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"lastStep failed");
|
||||
|
||||
delete tempA;
|
||||
delete tempB;
|
||||
NDArray::registerSpecialUse({&result}, {&array});
|
||||
}
|
||||
|
||||
void hashCode(LaunchContext* context, NDArray& array, NDArray& result) {
|
||||
BUILD_SINGLE_SELECTOR(array.dataType(), hashCode_, (context, array, result), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void hashCode_, (LaunchContext * context, NDArray& array, NDArray& result),
|
||||
SD_COMMON_TYPES);
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,147 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <ops/declarable/helpers/histogram.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename X, typename Z>
|
||||
static void SD_KERNEL histogramKernel(void *xBuffer, const LongType *xShapeInfo, void *zBuffer,
|
||||
const LongType *zShapeInfo, void *allocationPointer, void *reductionPointer,
|
||||
LongType numBins, X *min_val, X *max_val) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto dx = reinterpret_cast<X *>(xBuffer);
|
||||
auto result = reinterpret_cast<Z *>(zBuffer);
|
||||
|
||||
__shared__ Z *bins;
|
||||
__shared__ int length;
|
||||
__shared__ Z *reductor;
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
bins = (Z *)shmem;
|
||||
reductor = ((Z *)allocationPointer) + (numBins * blockIdx.x);
|
||||
|
||||
length = shape::length(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
X binSize = X((*max_val - *min_val) / numBins);
|
||||
|
||||
// nullify bins
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
bins[e] = (Z)0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int e = tid; e < length; e += blockDim.x * gridDim.x) {
|
||||
int idx = int((dx[e] - *min_val) / binSize);
|
||||
idx = math::sd_max(idx, 0); // atomicMax(&idx, 0);//atomicMax(&idx, 0);
|
||||
idx = math::sd_min(idx, int(numBins - 1)); // atomicMin(&idx, int(numBins - 1));
|
||||
math::atomics::sd_atomicAdd<Z>(&bins[idx], (Z)1);
|
||||
}
|
||||
__syncthreads();
|
||||
// at this point all bins in shared memory are calculated, so we aggregate them now via threadfence trick
|
||||
|
||||
// transfer shared memory to reduction memory
|
||||
if (gridDim.x > 1) {
|
||||
unsigned int *tc = (unsigned int *)reductionPointer;
|
||||
__shared__ bool amLast;
|
||||
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
reductor[e] = bins[e];
|
||||
}
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
tc[16384] = 0;
|
||||
|
||||
// nullify shared memory for future accumulation
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
bins[e] = (Z)0;
|
||||
}
|
||||
|
||||
// accumulate reduced bins
|
||||
for (int r = 0; r < gridDim.x; r++) {
|
||||
Z *ptrBuf = ((Z *)allocationPointer) + (r * numBins);
|
||||
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
math::atomics::sd_atomicAdd(&bins[e], ptrBuf[e]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// write them out to Z
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
result[e] = bins[e];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// if there's only 1 block - just write away data
|
||||
for (int e = threadIdx.x; e < numBins; e += blockDim.x) {
|
||||
result[e] = bins[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
static void histogram_(LaunchContext *context, void *xBuffer, const LongType *xShapeInfo,
|
||||
const LongType *dxShapeInfo, void *zBuffer, const LongType *zShapeInfo, LongType numBins, void *min_val, void *max_val) {
|
||||
dim3 histogramDims = getHistogramDims(shape::length(xShapeInfo),numBins);
|
||||
int workspaceSize = histogramDims.x * numBins;
|
||||
auto tmp = NDArrayFactory::create<Z>('c', {workspaceSize}, context);
|
||||
|
||||
histogramKernel<X, Z><<<histogramDims.x, histogramDims.y, histogramDims.z, *context->getCudaStream()>>>(
|
||||
xBuffer, dxShapeInfo, zBuffer, zShapeInfo, tmp.specialBuffer(), context->getReductionPointer(), numBins,
|
||||
reinterpret_cast<X *>(min_val), reinterpret_cast<X *>(max_val));
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(),"histogramKernel failed");
|
||||
|
||||
cudaStreamSynchronize(*context->getCudaStream());
|
||||
}
|
||||
|
||||
void histogramHelper(LaunchContext *context, NDArray &input, NDArray &output) {
|
||||
LongType numBins = output.lengthOf();
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
auto min_val = input.reduceNumber(reduce::SameOps::Min);
|
||||
auto max_val = input.reduceNumber(reduce::SameOps::Max);
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
input.dataType(), output.dataType(), histogram_,
|
||||
(context, input.specialBuffer(), input.shapeInfo(), input.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), numBins, min_val.specialBuffer(), max_val.specialBuffer()),
|
||||
SD_COMMON_TYPES, SD_INTEGER_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,136 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 31.08.2018
|
||||
//
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/histogramFixedWidth.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL static void histogramFixedWidthCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const X leftEdge, const X rightEdge) {
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
// Shared memory caching
|
||||
__shared__ LongType xLen, zLen, totalThreads, nbins;
|
||||
__shared__ X binWidth, secondEdge, lastButOneEdge;
|
||||
__shared__ LongType xRank, zRank;
|
||||
__shared__ const LongType* xShapePtr;
|
||||
__shared__ const LongType* xStridePtr;
|
||||
__shared__ const LongType* zShapePtr;
|
||||
__shared__ const LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(xShapeInfo);
|
||||
nbins = shape::length(zShapeInfo); // nbins = zLen
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
binWidth = (rightEdge - leftEdge) / nbins;
|
||||
secondEdge = leftEdge + binWidth;
|
||||
lastButOneEdge = rightEdge - binWidth;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < xLen; i += totalThreads) {
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Use cached rank and shape to compute coordinates and offset for x
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
const X value = x[xOffset];
|
||||
|
||||
// Determine the histogram bin index
|
||||
LongType zIndex;
|
||||
if (value < secondEdge)
|
||||
zIndex = 0;
|
||||
else if (value >= lastButOneEdge)
|
||||
zIndex = nbins - 1;
|
||||
else
|
||||
zIndex = static_cast<LongType>((value - leftEdge) / binWidth);
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
|
||||
// Use cached rank and shape to compute coordinates and offset for z
|
||||
INDEX2COORDS(zIndex, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
// Atomic addition to the histogram bin
|
||||
math::atomics::sd_atomicAdd<Z>(&z[zOffset], 1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_HOST static void histogramFixedWidthCudaLauncher(const cudaStream_t* stream, NDArray& input,
|
||||
NDArray& range, NDArray& output) {
|
||||
const X leftEdge = range.e<X>(0);
|
||||
const X rightEdge = range.e<X>(1);
|
||||
|
||||
dim3 launchDims = getLaunchDims("histogram_fixed_width");
|
||||
histogramFixedWidthCuda<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(input.specialBuffer(), input.specialShapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo(),
|
||||
leftEdge, rightEdge);
|
||||
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"histogramKernel failed");
|
||||
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
void histogramFixedWidth(LaunchContext* context, NDArray& input, NDArray& range, NDArray& output) {
|
||||
// firstly initialize output with zeros
|
||||
output.nullify();
|
||||
|
||||
PointersManager manager(context, "histogramFixedWidth");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), output.dataType(), histogramFixedWidthCudaLauncher,
|
||||
(context->getCudaStream(), input, range, output), SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,125 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// Created by raver119 on 30.11.17.
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/im2col.h>
|
||||
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// input [bS, iC, iH, iW] is convoluted to output [bS, iC, kH, kW, oH, oW]
|
||||
template <typename T>
|
||||
SD_KERNEL static void im2colCuda(const void *image, void *columns, const LongType *imShapeInfo,
|
||||
const LongType *colShapeInfo, const LongType sH, const LongType sW, const LongType pH,
|
||||
const LongType pW, const LongType dH, const LongType dW, const double zeroPadValD) {
|
||||
T zeroPadVal = static_cast<T>(zeroPadValD); // Value to use when value is padding
|
||||
const auto im = reinterpret_cast<const T *>(image);
|
||||
auto col = reinterpret_cast<T *>(columns);
|
||||
|
||||
// Shared memory caching
|
||||
__shared__ LongType colLen, imLen, iH, iW;
|
||||
__shared__ LongType imRank, colRank;
|
||||
__shared__ const LongType *imShapePtr, *imStridePtr;
|
||||
__shared__ const LongType *colShapePtr, *colStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
colRank = 6;
|
||||
imRank = 4;
|
||||
|
||||
colLen = shape::length(colShapeInfo);
|
||||
imLen = shape::length(imShapeInfo);
|
||||
|
||||
iH = shape::shapeOf(imShapeInfo)[2];
|
||||
iW = shape::shapeOf(imShapeInfo)[3];
|
||||
|
||||
imShapePtr = shape::shapeOf(imShapeInfo);
|
||||
imStridePtr = shape::stride(imShapeInfo);
|
||||
|
||||
colShapePtr = shape::shapeOf(colShapeInfo);
|
||||
colStridePtr = shape::stride(colShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto colInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (colInd >= colLen) return; // Boundary check for threads
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
// Calculate coordinates and offsets
|
||||
INDEX2COORDS(colInd, colRank, colShapePtr, coords);
|
||||
|
||||
LongType colOffset;
|
||||
COORDS2INDEX(colRank, colStridePtr, coords, colOffset);
|
||||
|
||||
coords[2] = (-pH + coords[2] * dH) + coords[4] * sH; // imH
|
||||
coords[3] = (-pW + coords[3] * dW) + coords[5] * sW; // imW
|
||||
|
||||
// Check bounds and assign appropriate values
|
||||
if (coords[2] >= iH || coords[3] >= iW || coords[2] < 0 || coords[3] < 0) {
|
||||
if (colOffset < colLen)
|
||||
col[colOffset] = zeroPadVal;
|
||||
} else {
|
||||
LongType imOffset;
|
||||
COORDS2INDEX(imRank, imStridePtr, coords, imOffset);
|
||||
if (imOffset < imLen && colOffset < colLen)
|
||||
col[colOffset] = im[imOffset];
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void im2colCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
||||
LaunchContext &context, const void *image, void *columns,
|
||||
const LongType *imShapeInfo, const LongType *colShapeInfo, LongType sH,
|
||||
LongType sW, LongType pH, LongType pW, LongType dH, LongType dW, double zeroPadVal) {
|
||||
im2colCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory /* rank of columns = 6 */, *context.getCudaStream()>>>(
|
||||
image, columns, imShapeInfo, colShapeInfo, sH, sW, pH, pW, dH, dW, zeroPadVal);
|
||||
DebugHelper::checkErrorCode(context.getCudaStream(), "im2colCuda(...) failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void im2col(LaunchContext &context, NDArray&image, NDArray &columns, const LongType kH, const LongType kW,
|
||||
const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW,
|
||||
NDArray&arrZeroPadVal) {
|
||||
PointersManager manager(&context, "im2col");
|
||||
|
||||
dim3 im2colDevs = getim2ColLaunchParams(columns);
|
||||
NDArray::prepareSpecialUse({&columns}, {&image});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
columns.dataType(), im2colCudaLauncher,
|
||||
(im2colDevs.x, im2colDevs.y,im2colDevs.z, context, image.specialBuffer(), columns.specialBuffer(),
|
||||
image.specialShapeInfo(), columns.specialShapeInfo(), sH, sW, pH, pW, dH, dW, arrZeroPadVal.e<double>(0)),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&columns}, {&image});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,194 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <array/NDArray.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
typedef NDArray ColorTable_t;
|
||||
static NDArray DefaultColorTable(int depth, LaunchContext* context) {
|
||||
// std::vector<std::vector<float>> colorTable;
|
||||
const LongType kDefaultTableLength = 10;
|
||||
const LongType kDefaultChannelLength = 4;
|
||||
std::vector<sd::LongType> shape = {kDefaultTableLength, kDefaultChannelLength};
|
||||
std::vector<double> table = {
|
||||
1, 1, 0, 1, // yellow
|
||||
0, 0, 1, 1, // 1: blue
|
||||
1, 0, 0, 1, // 2: red
|
||||
0, 1, 0, 1, // 3: lime
|
||||
0.5, 0, 0.5, 1, // 4: purple
|
||||
0.5, 0.5, 0, 1, // 5: olive
|
||||
0.5, 0, 0, 1, // 6: maroon
|
||||
0, 0, 0.5, 1, // 7: navy blue
|
||||
0, 1, 1, 1, // 8: aqua
|
||||
1, 0, 1, 1 // 9: fuchsia
|
||||
};
|
||||
NDArray colorTable('c', shape,
|
||||
table,
|
||||
FLOAT32, context);
|
||||
|
||||
if (depth == 1) {
|
||||
float one = 1.f;
|
||||
colorTable.assign(one); // all to white when black and white colors
|
||||
}
|
||||
return colorTable;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void drawBoundingBoxesKernel(T const* images, const LongType* imagesShape, float const* boxes,
|
||||
const LongType* boxesShape, float const* colorTable,
|
||||
const LongType* colorTableShape, T* output,
|
||||
const LongType* outputShape,
|
||||
LongType batchSize, LongType width, LongType height, LongType channels,
|
||||
LongType boxSize, LongType colorTableLen) {
|
||||
for (auto batch = blockIdx.x; batch < (int)batchSize; batch += gridDim.x) { // loop by batch
|
||||
for (auto boxIndex = 0; boxIndex < boxSize; ++boxIndex) {
|
||||
auto colorIndex = boxIndex % colorTableLen; // colorSet->at(c);
|
||||
LongType indices0[] = {batch, boxIndex, 0};
|
||||
LongType indices1[] = {batch, boxIndex, 1};
|
||||
LongType indices2[] = {batch, boxIndex, 2};
|
||||
LongType indices3[] = {batch, boxIndex, 3};
|
||||
|
||||
LongType rowStartOffset, rowEndOffset, colStartOffset, colEndOffset;
|
||||
COORDS2INDEX(3, shape::stride(boxesShape), indices0, rowStartOffset);
|
||||
COORDS2INDEX(3, shape::stride(boxesShape), indices2, rowEndOffset);
|
||||
COORDS2INDEX(3,shape::stride(boxesShape), indices1, colStartOffset);
|
||||
COORDS2INDEX(3, shape::stride(boxesShape), indices3, colEndOffset);
|
||||
|
||||
auto rowStart = LongType((height - 1) * boxes[rowStartOffset]);
|
||||
auto rowStartBound = math::sd_max(LongType(0), rowStart);
|
||||
auto rowEnd = LongType((height - 1) * boxes[rowEndOffset]);
|
||||
auto rowEndBound = math::sd_min(LongType(height - 1), rowEnd);
|
||||
auto colStart = LongType((width - 1) * boxes[colStartOffset]);
|
||||
auto colStartBound = math::sd_max(LongType(0), colStart);
|
||||
auto colEnd = LongType((width - 1) * boxes[colEndOffset]);
|
||||
auto colEndBound = math::sd_min(LongType(width - 1), colEnd);
|
||||
|
||||
if (rowStart > rowEnd || colStart > colEnd) {
|
||||
continue;
|
||||
}
|
||||
if (rowStart >= height || rowEnd < 0 || colStart >= width || colEnd < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Draw upper line
|
||||
if (rowStart >= 0) {
|
||||
for (auto j = colStartBound + threadIdx.x; j <= colEndBound; j += blockDim.x)
|
||||
for (auto c = 0; c < channels; c++) {
|
||||
LongType zPos[] = {batch, rowStart, j, c};
|
||||
LongType cPos[] = {colorIndex, c};
|
||||
LongType cIndex, zIndex;
|
||||
COORDS2INDEX(2, shape::stride(colorTableShape), cPos, cIndex);
|
||||
COORDS2INDEX(4, shape::stride(outputShape), zPos, zIndex);
|
||||
output[zIndex] = (T)colorTable[cIndex];
|
||||
}
|
||||
}
|
||||
// Draw bottom line.
|
||||
if (rowEnd < height) {
|
||||
for (auto j = colStartBound + threadIdx.x; j <= colEndBound; j += blockDim.x)
|
||||
for (auto c = 0; c < channels; c++) {
|
||||
LongType zPos[] = {batch, rowEnd, j, c};
|
||||
LongType cPos[] = {colorIndex, c};
|
||||
LongType cIndex, zIndex;
|
||||
COORDS2INDEX(2, shape::stride(colorTableShape), cPos, cIndex);
|
||||
COORDS2INDEX(4, shape::stride(outputShape), zPos, zIndex);
|
||||
output[zIndex] = (T)colorTable[cIndex];
|
||||
}
|
||||
}
|
||||
|
||||
// Draw left line.
|
||||
if (colStart >= 0) {
|
||||
for (auto i = rowStartBound + threadIdx.x; i <= rowEndBound; i += blockDim.x)
|
||||
for (auto c = 0; c < channels; c++) {
|
||||
LongType zPos[] = {batch, i, colStart, c};
|
||||
LongType cPos[] = {colorIndex, c};
|
||||
LongType cIndex, zIndex;
|
||||
COORDS2INDEX(2, shape::stride(colorTableShape), cPos, cIndex);
|
||||
COORDS2INDEX(4, shape::stride(outputShape), zPos, zIndex);
|
||||
output[zIndex] = (T)colorTable[cIndex];
|
||||
}
|
||||
}
|
||||
// Draw right line.
|
||||
if (colEnd < width) {
|
||||
for (auto i = rowStartBound + threadIdx.x; i <= rowEndBound; i += blockDim.x)
|
||||
for (auto c = 0; c < channels; c++) {
|
||||
LongType zPos[] = {batch, i, colEnd, c};
|
||||
LongType cPos[] = {colorIndex, c};
|
||||
LongType cIndex, zIndex;
|
||||
COORDS2INDEX(2, shape::stride(colorTableShape), cPos, cIndex);
|
||||
COORDS2INDEX(4, shape::stride(outputShape), zPos, zIndex);
|
||||
output[zIndex] = (T)colorTable[cIndex];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void drawBoundingBoxesH(LaunchContext* context, NDArray * images, NDArray * boxes, NDArray * colors,
|
||||
NDArray* output) {
|
||||
auto batchSize = images->sizeAt(0);
|
||||
auto height = images->sizeAt(1);
|
||||
auto width = images->sizeAt(2);
|
||||
auto channels = images->sizeAt(3);
|
||||
auto stream = context->getCudaStream();
|
||||
auto boxSize = boxes->sizeAt(1);
|
||||
NDArray colorsTable = DefaultColorTable(channels, context);
|
||||
if ((colors != nullptr && colors->lengthOf() > 0)) {
|
||||
colorsTable = *colors;
|
||||
}
|
||||
|
||||
auto imagesBuf = images->getDataBuffer()->specialAsT<T>();
|
||||
auto boxesBuf = boxes->getDataBuffer()->specialAsT<float>(); // boxes should be float32
|
||||
auto colorsTableBuf = colorsTable.getDataBuffer()->specialAsT<float>(); // color table is float32
|
||||
auto outputBuf = output->dataBuffer()->specialAsT<T>();
|
||||
dim3 launchDims = getLaunchDims("draw_bounding_boxes");
|
||||
drawBoundingBoxesKernel<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
imagesBuf, images->specialShapeInfo(), boxesBuf, boxes->specialShapeInfo(), colorsTableBuf,
|
||||
colorsTable.specialShapeInfo(), outputBuf, output->specialShapeInfo(), batchSize, width, height, channels,
|
||||
boxSize, colorsTable.lengthOf());
|
||||
}
|
||||
|
||||
void drawBoundingBoxesFunctor(LaunchContext* context, NDArray* images, NDArray* boxes, NDArray* colors,
|
||||
NDArray* output) {
|
||||
// images - batch of 3D images with BW (last dim = 1), RGB (last dim = 3) or RGBA (last dim = 4) channel set
|
||||
// boxes - batch of 2D bounds with last dim (y_start, x_start, y_end, x_end) to compute i and j as
|
||||
// floor((height - 1 ) * y_start) => rowStart, floor((height - 1) * y_end) => rowEnd
|
||||
// floor((width - 1 ) * x_start) => colStart, floor((width - 1) * x_end) => colEnd
|
||||
// height = images->sizeAt(1), width = images->sizeAt(2)
|
||||
// colors - colors for each box given
|
||||
// set up color for each box as frame
|
||||
NDArray::prepareSpecialUse({output}, {images, boxes, colors});
|
||||
output->assign(images);
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), drawBoundingBoxesH, (context, images, boxes, colors, output),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {images, boxes, colors});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,396 @@
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/image_resize.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
static SD_INLINE SD_HOST_DEVICE LongType boundsAmp(LongType const low, LongType const high, LongType const value) {
|
||||
if (high < value) return high;
|
||||
if (value < low) return low;
|
||||
return value;
|
||||
}
|
||||
|
||||
template <typename TKernelFunc>
|
||||
static SD_KERNEL void computeSpansKernel(TKernelFunc* kernel, int* startsVec, float* weightsVector, LongType outSize,
|
||||
LongType inSize, float kernelScale, int spanSize,
|
||||
float const invScale, float const invTranslate, float invKernelScale,
|
||||
float* tempWeightsBuf) {
|
||||
// return value if within bounds or bounds otherwise
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
__shared__ int maxSpanSize;
|
||||
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
maxSpanSize = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto x = tid; x < outSize; x += step) {
|
||||
const float columnFloat = x + 0.5f;
|
||||
const float sampleFloat = columnFloat * invScale + invTranslate;
|
||||
|
||||
// Don't sample when the sampling location is outside the source image.
|
||||
if (sampleFloat < 0 || sampleFloat > inSize) {
|
||||
// Add an empty span.
|
||||
startsVec[x] = 0;
|
||||
continue;
|
||||
}
|
||||
LongType spanStart = math::sd_ceil<float, float>(sampleFloat - kernel->radius() * kernelScale - 0.5f);
|
||||
LongType spanEnd = math::sd_floor<float, float>(sampleFloat + kernel->radius() * kernelScale - 0.5f);
|
||||
spanStart = boundsAmp(0LL, inSize - 1, spanStart);
|
||||
spanEnd = boundsAmp(0LL, inSize - 1, spanEnd) + 1;
|
||||
int const spanSize = spanEnd - spanStart;
|
||||
if (spanSize > spanSize) {
|
||||
return;
|
||||
}
|
||||
float totalWeightSum = 0.f;
|
||||
auto tempWeights = &tempWeightsBuf[x];
|
||||
auto actualWeights = 0;
|
||||
for (int source = spanStart; source < spanEnd; ++source) {
|
||||
float kernelPos = static_cast<float>(source) + 0.5f - sampleFloat;
|
||||
float weight = (*kernel)(kernelPos * invKernelScale);
|
||||
totalWeightSum += weight;
|
||||
tempWeights[actualWeights++] = weight;
|
||||
}
|
||||
maxSpanSize = math::sd_max(maxSpanSize, spanSize);
|
||||
if (math::sd_abs<float,float>(totalWeightSum) >= 1000.f * DataTypeUtils::min<float>()) { //
|
||||
auto totalWeightSumInverted = 1.0f / totalWeightSum;
|
||||
auto outIndex = spanSize * x;
|
||||
for (auto weightIndex = 0; weightIndex < actualWeights; ++weightIndex) {
|
||||
weightsVector[outIndex] = tempWeights[weightIndex] * totalWeightSumInverted;
|
||||
++outIndex;
|
||||
}
|
||||
}
|
||||
startsVec[x] = spanStart;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename TKernelFunc>
|
||||
static Status computeSpans(LaunchContext* context, TKernelFunc& kernel, LongType const outSize, LongType const inSize, float const scale, float const translate,
|
||||
bool const antialias, Spans& spans) {
|
||||
// When sampling, we need the inverse scale and translation, to map from an
|
||||
// output to an input pixel.
|
||||
float const invScale = 1.f / scale;
|
||||
float const invTranslate = -invScale * translate;
|
||||
// When downsampling the kernel should be scaled since we want to low pass
|
||||
// filter and interpolate, but when upsampling it should not be since we only
|
||||
// want to interpolate.
|
||||
float const kernelScale = antialias ? math::sd_max(invScale, 1.f) : 1.f;
|
||||
spans._spanSize =
|
||||
math::sd_min(2 * static_cast<int>(std::ceil(kernel.radius() * kernelScale)) + 1, static_cast<int>(inSize));
|
||||
spans._starts = NDArrayFactory::create<int>('c', {outSize});
|
||||
spans._starts.syncToHost();
|
||||
spans._weights = NDArrayFactory::create<float>('c', {outSize, spans._spanSize});
|
||||
spans._weights.syncToHost();
|
||||
|
||||
auto startsVec = reinterpret_cast<int*>(spans._starts.buffer());
|
||||
auto weightsVector = reinterpret_cast<float*>(spans._weights.buffer());
|
||||
spans._weights.nullify();
|
||||
|
||||
const float invKernelScale = 1.f / kernelScale;
|
||||
auto stream = context->getCudaStream();
|
||||
auto maxSpanSize = 0;
|
||||
std::vector<float> tempWeights;
|
||||
for (auto x = 0; x < outSize; x++) {
|
||||
const float columnFloat = x + 0.5f;
|
||||
const float sampleFloat = columnFloat * invScale + invTranslate;
|
||||
|
||||
// Don't sample when the sampling location is outside the source image.
|
||||
if (sampleFloat < 0 || sampleFloat > inSize) {
|
||||
// Add an empty span.
|
||||
startsVec[x] = 0;
|
||||
continue;
|
||||
}
|
||||
LongType spanStart = math::sd_ceil<float, float>(sampleFloat - kernel.radius() * kernelScale - 0.5f);
|
||||
LongType spanEnd = math::sd_floor<float, float>(sampleFloat + kernel.radius() * kernelScale - 0.5f);
|
||||
spanStart = boundsAmp(0LL, inSize - 1, spanStart);
|
||||
spanEnd = boundsAmp(0LL, inSize - 1, spanEnd) + 1;
|
||||
int const spanSize = spanEnd - spanStart;
|
||||
if (spanSize > spans._spanSize) {
|
||||
return Logger::logStatusMsg(
|
||||
Status::BAD_INPUT,
|
||||
"Span is too large: "); // + spanSize + " vs " + spans._spanSize);//, spanSize, spans._spanSize));
|
||||
}
|
||||
float totalWeightSum = 0.f;
|
||||
tempWeights.clear();
|
||||
|
||||
for (int source = spanStart; source < spanEnd; ++source) {
|
||||
float kernelPos = static_cast<float>(source) + 0.5f - sampleFloat;
|
||||
float weight = kernel(kernelPos * invKernelScale);
|
||||
totalWeightSum += weight;
|
||||
tempWeights.push_back(weight);
|
||||
}
|
||||
maxSpanSize = math::sd_max(maxSpanSize, spanSize);
|
||||
if (math::sd_abs<float,float>(totalWeightSum) >= 1000.f * DataTypeUtils::min<float>()) { //
|
||||
auto totalWeightSumInverted = 1.0f / totalWeightSum;
|
||||
auto outIndex = spans._spanSize * x;
|
||||
for (auto weightIndex = 0; weightIndex < tempWeights.size(); ++weightIndex) {
|
||||
weightsVector[outIndex++] = tempWeights[weightIndex] * totalWeightSumInverted;
|
||||
}
|
||||
}
|
||||
startsVec[x] = spanStart;
|
||||
}
|
||||
spans._starts.tickWriteHost();
|
||||
spans._weights.tickWriteHost();
|
||||
spans._starts.syncToDevice();
|
||||
spans._weights.syncToDevice();
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void batchedGatherSpan(LongType outputWidth, LongType outputHeight, int rowSpanSize,
|
||||
int const* rowStartsBuf, Z const* rowWeightBuf, int columnSpanSize,
|
||||
int const* columnStartsBuf, Z const* columnWeightBuf, X const* pImages,
|
||||
const LongType* imageSpecialShapeInfo, Z* pIntermediate, Z* pOutput,
|
||||
LongType outputPixPerBatch) {
|
||||
auto batchSize = shape::sizeAt(imageSpecialShapeInfo, 0);
|
||||
auto inputHeight = shape::sizeAt(imageSpecialShapeInfo, 1);
|
||||
auto inputWidth = shape::sizeAt(imageSpecialShapeInfo, 2);
|
||||
auto channels = shape::sizeAt(imageSpecialShapeInfo, 3);
|
||||
bool inputEws1 = shape::elementWiseStride(imageSpecialShapeInfo) == 1;
|
||||
auto inputPixPerBatch = shape::strideAt(imageSpecialShapeInfo, 0);
|
||||
auto inRowStride = shape::strideAt(imageSpecialShapeInfo, 1);
|
||||
auto wStride = shape::strideAt(imageSpecialShapeInfo, 2);
|
||||
auto cStride = shape::strideAt(imageSpecialShapeInfo, 3);
|
||||
auto intermediatePixPerBatch = inputWidth * outputHeight * channels;
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (int b = tid; b < batchSize; b += step) {
|
||||
auto imagePtr = pImages + b * inputPixPerBatch;
|
||||
auto intermediatePtr = pIntermediate + b * intermediatePixPerBatch;
|
||||
auto outputPtr = pOutput + b * outputPixPerBatch;
|
||||
gatherRows<X, Z>(rowSpanSize, rowStartsBuf, rowWeightBuf, imagePtr, inputHeight, inputWidth, outputHeight,
|
||||
inputWidth, channels, intermediatePtr, inputEws1, inRowStride, wStride, cStride);
|
||||
gatherColumns<Z>(columnSpanSize, columnStartsBuf, columnWeightBuf, intermediatePtr, outputHeight, inputWidth,
|
||||
outputHeight, outputWidth, channels, outputPtr);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
static void gatherSpans(LaunchContext* context, int const rowSpanSize, NDArray& rowStarts,
|
||||
NDArray& rowWeights, int const colSpanSize, NDArray& columnStarts,
|
||||
NDArray& columnWeights, NDArray * images, NDArray& intermediate, NDArray* output) {
|
||||
const auto imageSpecialShapeInfo = images->specialShapeInfo();
|
||||
auto outputHeight = output->sizeAt(1);
|
||||
auto outputWidth = output->sizeAt(2);
|
||||
auto channels = images->sizeAt(3);
|
||||
auto outputPixPerBatch = outputWidth * outputHeight * channels;
|
||||
auto intermediatePtr = reinterpret_cast<Z*>(intermediate.specialBuffer());
|
||||
|
||||
auto imagePtr = reinterpret_cast<X const*>(images->specialBuffer());
|
||||
auto outputPtr = reinterpret_cast<Z*>(output->specialBuffer());
|
||||
auto stream = context->getCudaStream();
|
||||
auto rowStartsBuf = reinterpret_cast<int const*>(rowStarts.specialBuffer());
|
||||
auto rowWeightBuf = reinterpret_cast<Z const*>(rowWeights.specialBuffer());
|
||||
auto columnStartsBuf = reinterpret_cast<int const*>(columnStarts.specialBuffer());
|
||||
auto columnWeightBuf = reinterpret_cast<Z const*>(columnWeights.specialBuffer());
|
||||
dim3 launchDims = getLaunchDims("image_resize_v2_gather");
|
||||
batchedGatherSpan<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
outputWidth, outputHeight, rowSpanSize, rowStartsBuf, rowWeightBuf, colSpanSize, columnStartsBuf, columnWeightBuf,
|
||||
imagePtr, imageSpecialShapeInfo, intermediatePtr, outputPtr, outputPixPerBatch);
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
static Status resizeKernel(LaunchContext* context, ImageResizeMethods method, NDArray * input, LongType outWidth,
|
||||
LongType outHeight, bool antialias, double coefficient,
|
||||
NDArray* output) {
|
||||
LongType const batchSize = input->sizeAt(0);
|
||||
LongType const inputHeight = input->sizeAt(1);
|
||||
LongType const inputWidth = input->sizeAt(2);
|
||||
LongType const channels = input->sizeAt(3);
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
Z rowScale = Z(outHeight) / Z(inputHeight);
|
||||
Z columnScale = Z(outWidth) / Z(inputWidth);
|
||||
|
||||
// Return if the output is empty.
|
||||
if (output->lengthOf() == 0) return Status::OK;
|
||||
|
||||
Spans colSpans;
|
||||
Spans rowSpans;
|
||||
auto res = Status::OK;
|
||||
switch (method) {
|
||||
case kResizeBilinear: {
|
||||
TriangleKernelFunc kernel;
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
|
||||
} break;
|
||||
case kResizeBicubic: {
|
||||
KeysCubicKernelFunc<float> kernel(static_cast<float>(coefficient));
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
} break;
|
||||
case kResizeLanczos3: {
|
||||
LanczosKernelFunc kernel(3.f);
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
|
||||
} break;
|
||||
|
||||
case kResizeLanczos5: {
|
||||
LanczosKernelFunc kernel(5.f);
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
|
||||
} break;
|
||||
case kResizeGaussian: {
|
||||
GaussianKernelFunc kernel;
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
|
||||
} break;
|
||||
case kResizeMitchellcubic: {
|
||||
MitchellCubicKernelFunc kernel;
|
||||
res = computeSpans(context, kernel, outWidth, inputWidth, columnScale, 0.f, antialias, colSpans);
|
||||
if (res != Status::OK) return res;
|
||||
res = computeSpans(context, kernel, outHeight, inputHeight, rowScale, 0.f, antialias, rowSpans);
|
||||
|
||||
} break;
|
||||
};
|
||||
|
||||
NDArray intermediate = NDArrayFactory::create<Z>('c', {batchSize, outHeight, inputWidth, channels});
|
||||
|
||||
// const functor::Spans& const_row_spans = row_spans;
|
||||
// typename TTypes<int32, 1>::ConstTensor row_starts(
|
||||
// const_row_spans.starts.tensor<int32, 1>());
|
||||
auto& rowStarts = rowSpans._starts; // shape {outWidth}
|
||||
auto& rowWeights = rowSpans._weights; // shape {outWidth, numSpans}
|
||||
auto& columnStarts = colSpans._starts; // shape {outHeights}
|
||||
auto& columnWeights = colSpans._weights; // shape {outHeights, numSpans}
|
||||
|
||||
gatherSpans<X, Z>(context, rowSpans._spanSize, rowStarts, rowWeights, colSpans._spanSize, columnStarts, columnWeights,
|
||||
input, intermediate, output);
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
return res;
|
||||
}
|
||||
|
||||
#if defined(HAS_FLOAT32)
|
||||
#define SD_FLOAT_TYPES_FLOAT32 SKIP_FIRST_COMMA(TTYPE_FLOAT32)
|
||||
|
||||
static Status resizeTriangle(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
bool const antialias, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeBilinear, image, width, height, antialias, 0, output), SD_NUMERIC_TYPES,
|
||||
SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::resizeTriangle: This resize method is avaliable in future versions");
|
||||
}
|
||||
|
||||
static Status resizeLanczos3(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
bool const antialias, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeLanczos3, image, width, height, antialias, 0, output), SD_NUMERIC_TYPES,
|
||||
SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::resizeLanczos3: This resize method is avaliable in future versions");
|
||||
}
|
||||
|
||||
static Status resizeLanczos5(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
bool const antialias, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeLanczos5, image, width, height, antialias, 0, output), SD_NUMERIC_TYPES,
|
||||
SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::resizeLanczos5: This resize method is avaliable in future versions");
|
||||
}
|
||||
|
||||
static Status resizeGaussian(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
bool const antialias, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeGaussian, image, width, height, antialias, 0, output), SD_NUMERIC_TYPES,
|
||||
SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::resizeGaussian: This resize method is avaliable in future versions");
|
||||
}
|
||||
static Status resizeMitchellcubic(LaunchContext* context, NDArray * image, int const width,
|
||||
int const height, bool const antialias, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeMitchellcubic, image, width, height, antialias, 0, output), SD_NUMERIC_TYPES,
|
||||
SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::ResizeMitchellcubic: This resize method is avaliable in future versions");
|
||||
}
|
||||
|
||||
static Status resizeBicubicA(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
CoordinateTransformationMode coorMode, bool exclude_outside, double coefficient,
|
||||
NDArray* output) {
|
||||
constexpr bool alignCorners = false;
|
||||
return resizeBicubicFunctorA(context, image, width, height, alignCorners, coorMode, exclude_outside, coefficient,
|
||||
output);
|
||||
}
|
||||
|
||||
static Status resizeBicubicAntialias(LaunchContext* context, NDArray * image, int const width,
|
||||
int const height, bool const antialias, double coefficient, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(image->dataType(), output->dataType(), return resizeKernel,
|
||||
(context, kResizeBicubic, image, width, height, antialias, coefficient, output),
|
||||
SD_NUMERIC_TYPES, SD_FLOAT_TYPES_FLOAT32);
|
||||
return Logger::logStatusMsg(Status::VALIDATION,
|
||||
"helpers::ResizeMitchellcubic: This resize method is avaliable in future versions");
|
||||
}
|
||||
|
||||
#endif
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
Status resizeFunctor(LaunchContext* context, NDArray * image, int const width, int const height,
|
||||
ImageResizeMethods method, CoordinateTransformationMode coorMode, bool exclude_outside,
|
||||
NearestMode nearestMode, double coefficient, bool antialias, NDArray* output) {
|
||||
switch (method) {
|
||||
case kResizeNearest:
|
||||
return resizeNeighborFunctor(context, image, width, height, coorMode, nearestMode, false, output);
|
||||
case kResizeArea:
|
||||
return resizeAreaFunctor(context, image, width, height, false, output);
|
||||
|
||||
#if defined(HAS_FLOAT32)
|
||||
case kResizeBilinear:
|
||||
return resizeTriangle(context, image, width, height, antialias, output);
|
||||
case kResizeLanczos3:
|
||||
return resizeLanczos3(context, image, width, height, antialias, output);
|
||||
case kResizeLanczos5:
|
||||
return resizeLanczos5(context, image, width, height, antialias, output);
|
||||
case kResizeGaussian:
|
||||
return resizeGaussian(context, image, width, height, antialias, output);
|
||||
case kResizeMitchellcubic:
|
||||
return resizeMitchellcubic(context, image, width, height, antialias, output);
|
||||
case kResizeBicubic: {
|
||||
// if antialias then coorMode is HALF_PIXEL and exlude_outside is true
|
||||
if (antialias) {
|
||||
return resizeBicubicAntialias(context, image, width, height, antialias, coefficient, output);
|
||||
} else {
|
||||
// use modified v1
|
||||
return resizeBicubicA(context, image, width, height, coorMode, exclude_outside, coefficient, output);
|
||||
}
|
||||
}
|
||||
#else
|
||||
case kResizeBilinear:
|
||||
case kResizeLanczos3:
|
||||
case kResizeLanczos5:
|
||||
case kResizeGaussian:
|
||||
case kResizeMitchellcubic:
|
||||
case kResizeBicubic: {
|
||||
sd_printf("helper::resizeFunctor: only float type is supported by this resize method %i\n", (int)method);
|
||||
return Logger::logStatusMsg(Status::BAD_INPUT, "helper::resizeFunctor: only float type supported");
|
||||
}
|
||||
#endif
|
||||
default:
|
||||
sd_printf("helper::resizeFunctor: Wrong resize method %i\n", (int)method);
|
||||
THROW_EXCEPTION("helper::resizeFunctor: Wrong resize method.");
|
||||
}
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,551 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <legacy/NativeOps.h>
|
||||
#include <ops/declarable/helpers/image_suppression.h>
|
||||
|
||||
#include <queue>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// needToSuppressWithThreshold - predicate for suppression
|
||||
// boxes - boxes tensor buffer
|
||||
// boxesShape boxes tensor shape
|
||||
// previousIndex - index for current pos value
|
||||
// nextIndex - index for neighbor pos value
|
||||
// threshold - threashold value to suppress
|
||||
//
|
||||
// return value: true, if threshold is overcome, false otherwise
|
||||
//
|
||||
template <typename T>
|
||||
static SD_DEVICE bool needToSuppressWithThreshold(T* boxes, LongType const* boxesShape, int previousIndex,
|
||||
int nextIndex, T threshold) {
|
||||
LongType previous0[] = {previousIndex, 0};
|
||||
LongType previous1[] = {previousIndex, 1};
|
||||
LongType previous2[] = {previousIndex, 2};
|
||||
LongType previous3[] = {previousIndex, 3};
|
||||
LongType next0[] = {nextIndex, 0};
|
||||
LongType next1[] = {nextIndex, 1};
|
||||
LongType next2[] = {nextIndex, 2};
|
||||
LongType next3[] = {nextIndex, 3};
|
||||
|
||||
LongType prevOffset0, prevOffset1, prevOffset2, prevOffset3;
|
||||
LongType nextOffset0, nextOffset1, nextOffset2, nextOffset3;
|
||||
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous0, prevOffset0);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous1, prevOffset1);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous2, prevOffset2);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous3, prevOffset3);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next0, nextOffset0);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next1, nextOffset1);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next2, nextOffset2);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next3, nextOffset3);
|
||||
|
||||
// we have rectangle with given max values. Compute vexes of rectangle first
|
||||
|
||||
T minYPrev = math::sd_min(boxes[prevOffset0], boxes[prevOffset2]);
|
||||
T minXPrev = math::sd_min(boxes[prevOffset1], boxes[prevOffset3]);
|
||||
T maxYPrev = math::sd_max(boxes[prevOffset0], boxes[prevOffset2]);
|
||||
T maxXPrev = math::sd_max(boxes[prevOffset1], boxes[prevOffset3]);
|
||||
T minYNext = math::sd_min(boxes[nextOffset0], boxes[nextOffset2]);
|
||||
T minXNext = math::sd_min(boxes[nextOffset1], boxes[nextOffset3]);
|
||||
T maxYNext = math::sd_max(boxes[nextOffset0], boxes[nextOffset2]);
|
||||
T maxXNext = math::sd_max(boxes[nextOffset1], boxes[nextOffset3]);
|
||||
|
||||
// compute areas for comparation
|
||||
T areaPrev = (maxYPrev - minYPrev) * (maxXPrev - minXPrev);
|
||||
T areaNext = (maxYNext - minYNext) * (maxXNext - minXNext);
|
||||
|
||||
// of course, areas should be positive
|
||||
if (areaNext <= T(0.f) || areaPrev <= T(0.f)) return false;
|
||||
|
||||
// compute intersection of rectangles
|
||||
T minIntersectionY = math::sd_max(minYPrev, minYNext);
|
||||
T minIntersectionX = math::sd_max(minXPrev, minXNext);
|
||||
T maxIntersectionY = math::sd_min(maxYPrev, maxYNext);
|
||||
T maxIntersectionX = math::sd_min(maxXPrev, maxXNext);
|
||||
T intersectionArea = math::sd_max(T(maxIntersectionY - minIntersectionY), T(0.0f)) *
|
||||
math::sd_max(T(maxIntersectionX - minIntersectionX), T(0.0f));
|
||||
T intersectionValue = intersectionArea / (areaPrev + areaNext - intersectionArea);
|
||||
// final check
|
||||
return intersectionValue > threshold;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static inline T similirityV3_(NDArray& boxes, LongType i, LongType j) {
|
||||
const T zero = static_cast<T>(0.f);
|
||||
const T yminI = math::sd_min(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
|
||||
const T xminI = math::sd_min(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
|
||||
const T ymaxI = math::sd_max(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
|
||||
const T xmaxI = math::sd_max(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
|
||||
const T yminJ = math::sd_min(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
|
||||
const T xminJ = math::sd_min(boxes.t<T>(j, 1), boxes.t<T>(j, 3));
|
||||
const T ymaxJ = math::sd_max(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
|
||||
const T xmaxJ = math::sd_max(boxes.t<T>(j, 1), boxes.t<T>(j, 3));
|
||||
const T areaI = (ymaxI - yminI) * (xmaxI - xminI);
|
||||
const T areaJ = (ymaxJ - yminJ) * (xmaxJ - xminJ);
|
||||
if (areaI <= zero || areaJ <= zero) {
|
||||
return zero;
|
||||
}
|
||||
const T intersectionYmin = math::sd_max(yminI, yminJ);
|
||||
const T intersectionXmin = math::sd_max(xminI, xminJ);
|
||||
const T intersectionYmax = math::sd_min(ymaxI, ymaxJ);
|
||||
const T intersectionXmax = math::sd_min(xmaxI, xmaxJ);
|
||||
const T intersectionY = intersectionYmax - intersectionYmin;
|
||||
const T intersectionX = intersectionXmax - intersectionXmin;
|
||||
const T intersectionArea = math::sd_max(intersectionY, zero) * math::sd_max(intersectionX, zero);
|
||||
return intersectionArea / (areaI + areaJ - intersectionArea);
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static SD_DEVICE T similirityV3(T* boxes, LongType const* boxesShape, int previousIndex, int nextIndex) {
|
||||
LongType previous0[] = {previousIndex, 0};
|
||||
LongType previous1[] = {previousIndex, 1};
|
||||
LongType previous2[] = {previousIndex, 2};
|
||||
LongType previous3[] = {previousIndex, 3};
|
||||
LongType next0[] = {nextIndex, 0};
|
||||
LongType next1[] = {nextIndex, 1};
|
||||
LongType next2[] = {nextIndex, 2};
|
||||
LongType next3[] = {nextIndex, 3};
|
||||
|
||||
LongType prevOffset0, prevOffset1, prevOffset2, prevOffset3;
|
||||
LongType nextOffset0, nextOffset1, nextOffset2, nextOffset3;
|
||||
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous0, prevOffset0);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous1, prevOffset1);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous2, prevOffset2);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), previous3, prevOffset3);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next0, nextOffset0);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next1, nextOffset1);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next2, nextOffset2);
|
||||
COORDS2INDEX(2, shape::stride(boxesShape), next3, nextOffset3);
|
||||
|
||||
// we have rectangle with given max values. Compute vexes of rectangle first
|
||||
|
||||
T minYPrev = math::sd_min(boxes[prevOffset0], boxes[prevOffset2]);
|
||||
T minXPrev = math::sd_min(boxes[prevOffset1], boxes[prevOffset3]);
|
||||
T maxYPrev = math::sd_max(boxes[prevOffset0], boxes[prevOffset2]);
|
||||
T maxXPrev = math::sd_max(boxes[prevOffset1], boxes[prevOffset3]);
|
||||
T minYNext = math::sd_min(boxes[nextOffset0], boxes[nextOffset2]);
|
||||
T minXNext = math::sd_min(boxes[nextOffset1], boxes[nextOffset3]);
|
||||
T maxYNext = math::sd_max(boxes[nextOffset0], boxes[nextOffset2]);
|
||||
T maxXNext = math::sd_max(boxes[nextOffset1], boxes[nextOffset3]);
|
||||
|
||||
// compute areas for comparator
|
||||
T areaPrev = (maxYPrev - minYPrev) * (maxXPrev - minXPrev);
|
||||
T areaNext = (maxYNext - minYNext) * (maxXNext - minXNext);
|
||||
|
||||
// of course, areas should be positive
|
||||
if (areaNext <= T(0.f) || areaPrev <= T(0.f)) return false;
|
||||
|
||||
// compute intersection of rectangles
|
||||
T minIntersectionY = math::sd_max(minYPrev, minYNext);
|
||||
T minIntersectionX = math::sd_max(minXPrev, minXNext);
|
||||
T maxIntersectionY = math::sd_min(maxYPrev, maxYNext);
|
||||
T maxIntersectionX = math::sd_min(maxXPrev, maxXNext);
|
||||
T intersectionArea = math::sd_max(T(maxIntersectionY - minIntersectionY), T(0.0f)) *
|
||||
math::sd_max(T(maxIntersectionX - minIntersectionX), T(0.0f));
|
||||
T intersectionValue = intersectionArea / (areaPrev + areaNext - intersectionArea);
|
||||
// final check
|
||||
return intersectionValue;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// shouldSelectKernel - compute status for all selected rectangles (boxes)
|
||||
//
|
||||
// we compute boolean flag as shared uint32 and return it on final only for the first thread
|
||||
//
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void shouldSelectKernel(T* boxesBuf, LongType const* boxesShape, I* indexBuf,
|
||||
I* selectedIndicesData, double threshold, int numSelected, int i,
|
||||
bool* shouldSelect) {
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
__shared__ unsigned int shouldSelectShared;
|
||||
if (threadIdx.x == 0) {
|
||||
shouldSelectShared = (unsigned int)shouldSelect[0];
|
||||
}
|
||||
__syncthreads();
|
||||
for (int j = numSelected - 1 - tid; j >= 0; j -= step) {
|
||||
if (shouldSelectShared) {
|
||||
if (needToSuppressWithThreshold(boxesBuf, boxesShape, indexBuf[i], indexBuf[selectedIndicesData[j]],
|
||||
T(threshold)))
|
||||
atomicCAS(&shouldSelectShared, 1, 0); // exchange only when need to suppress
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// final move: collect result
|
||||
if (threadIdx.x == 0) {
|
||||
*shouldSelect = shouldSelectShared > 0;
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// indices - type depended, indicesLong - type defined (only 64bit integers)
|
||||
//
|
||||
template <typename I>
|
||||
static SD_KERNEL void copyIndices(void* indices, void* indicesLong, LongType len) {
|
||||
I* indexBuf = reinterpret_cast<I*>(indices);
|
||||
LongType* srcBuf = reinterpret_cast<LongType*>(indicesLong);
|
||||
;
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (auto i = tid; i < len; i += step) indexBuf[i] = (I)srcBuf[i];
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void suppressScores(T* scores, I* indices, LongType length, T scoreThreshold) {
|
||||
auto start = blockIdx.x * blockDim.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (auto e = start + threadIdx.x; e < (int)length; e += step) {
|
||||
if (scores[e] < scoreThreshold) {
|
||||
scores[e] = scoreThreshold;
|
||||
indices[e] = -1;
|
||||
} else {
|
||||
indices[e] = I(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// nonMaxSuppressionV2 algorithm - given from TF NonMaxSuppressionV2 implementation
|
||||
//
|
||||
template <typename T, typename I>
|
||||
static void nonMaxSuppressionV2_(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize,
|
||||
double threshold, double scoreThreshold, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {boxes, scales});
|
||||
std::vector<sd::LongType> shape = {scales->lengthOf()};
|
||||
NDArray indices (NDArrayFactory::create_<I>(
|
||||
'c', shape, context)); // - 1, scales->lengthOf()); //, scales->getContext());
|
||||
|
||||
NDArray scores(*scales);
|
||||
Pointer extras[2] = {nullptr, stream};
|
||||
auto indexBuf = indices.dataBuffer()->specialAsT<I>();
|
||||
auto scoreBuf = scores.dataBuffer()->specialAsT<T>();
|
||||
dim3 launchDims = getLaunchDims("image_suppress_scores");
|
||||
suppressScores<T, I><<<launchDims.x, launchDims.y,launchDims.z, *stream>>>(scoreBuf, indexBuf, scores.lengthOf(), T(scoreThreshold));
|
||||
indices.tickWriteDevice();
|
||||
sortByValue(extras, &indices,
|
||||
&scores,true);
|
||||
indices.tickWriteDevice();
|
||||
NDArray selectedIndices = NDArrayFactory::create<I>('c', {output->lengthOf()}, context);
|
||||
int numSelected = 0;
|
||||
int numBoxes = boxes->sizeAt(0), tt(0);
|
||||
auto boxesBuf = reinterpret_cast<T*>(boxes->specialBuffer());
|
||||
|
||||
auto selectedIndicesData = reinterpret_cast<I*>(selectedIndices.specialBuffer());
|
||||
auto outputBuf = reinterpret_cast<I*>(output->specialBuffer());
|
||||
|
||||
bool* shouldSelectD;
|
||||
auto err = cudaMalloc(&shouldSelectD, sizeof(bool));
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot allocate memory for bool flag", err);
|
||||
}
|
||||
for (I i = 0; i < boxes->sizeAt(0); ++i) {
|
||||
bool shouldSelect = numSelected < output->lengthOf();
|
||||
if (shouldSelect) {
|
||||
err = cudaMemcpy(shouldSelectD, &shouldSelect, sizeof(bool), cudaMemcpyHostToDevice);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot set up bool flag to device", err);
|
||||
}
|
||||
|
||||
dim3 selectDims = getLaunchDims("image_suppress_select");
|
||||
shouldSelectKernel<T, I><<<selectDims.y,selectDims.x,selectDims.z, *stream>>>(
|
||||
boxesBuf, boxes->specialShapeInfo(), indexBuf, selectedIndicesData, threshold, numSelected, i, shouldSelectD);
|
||||
err = cudaMemcpy(&shouldSelect, shouldSelectD, sizeof(bool), cudaMemcpyDeviceToHost);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot set up bool flag to host", err);
|
||||
}
|
||||
}
|
||||
|
||||
if (shouldSelect) {
|
||||
cudaMemcpy(reinterpret_cast<I*>(output->specialBuffer()) + numSelected, indexBuf + i, sizeof(I),
|
||||
cudaMemcpyDeviceToDevice);
|
||||
cudaMemcpy(selectedIndicesData + numSelected, &i, sizeof(I), cudaMemcpyHostToDevice);
|
||||
numSelected++;
|
||||
}
|
||||
}
|
||||
|
||||
err = cudaFree(shouldSelectD);
|
||||
if (err) {
|
||||
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot deallocate memory for bool flag", err);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T, typename I>
|
||||
static SD_DEVICE bool checkOverlapBoxes(T* boxes, LongType const* shape, T* scores, I* indices, I* selectedIndices,
|
||||
I* startIndices, I selectedSize, I nextCandidateIndex, T overlapThreshold,
|
||||
T scoreThreshold, bool simple) {
|
||||
bool shouldHardSuppress = false;
|
||||
T& nextCandidateScore = scores[nextCandidateIndex];
|
||||
I selectedIndex = indices[nextCandidateIndex];
|
||||
I finish = startIndices[nextCandidateIndex];
|
||||
|
||||
for (int j = selectedSize; j > finish; --j) {
|
||||
T boxVal;
|
||||
if (simple) {
|
||||
LongType xPos[] = {selectedIndex, selectedIndices[j - 1]};
|
||||
LongType xShift;
|
||||
COORDS2INDEX(shape::rank(shape), shape::stride(shape), xPos, xShift);
|
||||
boxVal = boxes[xShift];
|
||||
} else {
|
||||
boxVal = similirityV3(boxes, shape, selectedIndex, selectedIndices[j - 1]);
|
||||
}
|
||||
if (boxVal > static_cast<T>(overlapThreshold)) nextCandidateScore = static_cast<T>(0.f);
|
||||
|
||||
// First decide whether to perform hard suppression
|
||||
if (boxVal >= overlapThreshold) {
|
||||
shouldHardSuppress = true;
|
||||
break;
|
||||
}
|
||||
|
||||
// If nextCandidate survives hard suppression, apply soft suppression
|
||||
if (nextCandidateScore <= static_cast<T>(scoreThreshold)) break;
|
||||
}
|
||||
|
||||
return shouldHardSuppress;
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void suppressNonMaxOverlapKernel(T* boxes, LongType const* boxesShape, T* scoresData, I* indices,
|
||||
I* startIndices, LongType length, I maxOutputLen,
|
||||
T overlapThreshold, T scoreThreshold, I* output, LongType const* outputShape, I* outputLength, bool simple) {
|
||||
__shared__ I selectedSize;
|
||||
__shared__ I* tempOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
selectedSize = outputLength ? *outputLength : maxOutputLen;
|
||||
extern __shared__ unsigned char shmem[];
|
||||
tempOutput = (I*)shmem;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (I nextCandidateIndex = start + threadIdx.x; selectedSize < maxOutputLen && nextCandidateIndex < (I)length;) {
|
||||
auto originalScore = scoresData[nextCandidateIndex];
|
||||
I nextCandidateBoxIndex = indices[nextCandidateIndex];
|
||||
auto selectedSizeMark = selectedSize;
|
||||
|
||||
// skip for cases when index is less than 0 (under score threshold)
|
||||
if (nextCandidateBoxIndex < 0) {
|
||||
nextCandidateIndex += step;
|
||||
continue;
|
||||
}
|
||||
// check for overlaps
|
||||
bool shouldHardSuppress =
|
||||
checkOverlapBoxes(boxes, boxesShape, scoresData, indices, tempOutput, startIndices, selectedSize,
|
||||
nextCandidateIndex, overlapThreshold, scoreThreshold, simple); // false;
|
||||
T nextCandidateScore = scoresData[nextCandidateIndex];
|
||||
|
||||
startIndices[nextCandidateIndex] = selectedSize;
|
||||
if (!shouldHardSuppress) {
|
||||
if (nextCandidateScore == originalScore) {
|
||||
// Suppression has not occurred, so select nextCandidate
|
||||
I currSize = math::atomics::sd_atomicAdd(&selectedSize, (I)1);
|
||||
if (output) {
|
||||
output[currSize] = nextCandidateBoxIndex;
|
||||
}
|
||||
tempOutput[currSize] = nextCandidateBoxIndex;
|
||||
|
||||
}
|
||||
|
||||
if ((float) nextCandidateScore > (float) scoreThreshold) {
|
||||
// Soft suppression has occurred and current score is still greater than
|
||||
// scoreThreshold; add nextCandidate back onto priority queue.
|
||||
continue; // in some cases, this index not 0
|
||||
}
|
||||
}
|
||||
nextCandidateIndex += step;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
printf("selectedSize: %i\n", selectedSize);
|
||||
if (outputLength) *outputLength = selectedSize;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
typedef NDArray (*SimilarityFunc)(NDArray& boxes, LongType i, LongType j);
|
||||
template <typename T>
|
||||
static inline T similarityOverlaps_(NDArray& boxes, LongType i, LongType j) {
|
||||
return boxes.t<T>(i, j);
|
||||
}
|
||||
|
||||
static NDArray similiratyOverlaps(NDArray& boxes, LongType i, LongType j) {
|
||||
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
|
||||
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similarityOverlaps_, (boxes, i, j), SD_FLOAT_TYPES);
|
||||
return res;
|
||||
}
|
||||
|
||||
static NDArray similarityV3(NDArray& boxes, LongType i, LongType j) {
|
||||
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
|
||||
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similirityV3_, (boxes, i, j), SD_FLOAT_TYPES);
|
||||
return res;
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T, typename I>
|
||||
static LongType nonMaxSuppressionGeneric_(LaunchContext* context, NDArray* boxes, NDArray* scores,
|
||||
int outputSize, float overlapThreshold, float scoreThreshold,
|
||||
NDArray* output, SimilarityFunc f) {
|
||||
auto stream = context->getCudaStream();
|
||||
if (output)
|
||||
NDArray::preparePrimaryUse({output}, {boxes, scores});
|
||||
else {
|
||||
if (!boxes->isActualOnHostSide()) boxes->syncToHost();
|
||||
if (!scores->isActualOnHostSide()) scores->syncToHost();
|
||||
}
|
||||
|
||||
auto numBoxes = boxes->sizeAt(0);
|
||||
T* scoresData = scores->dataBuffer()->primaryAsT<T>();
|
||||
|
||||
// Data structure for a selection candidate in NMS.
|
||||
struct Candidate {
|
||||
int _boxIndex;
|
||||
T _score;
|
||||
int _suppressBeginIndex;
|
||||
};
|
||||
|
||||
auto cmp = [](const Candidate& bsI, const Candidate& bsJ) -> bool {
|
||||
return ((bsI._score == bsJ._score) && (bsI._boxIndex > bsJ._boxIndex)) || (bsI._score < bsJ._score);
|
||||
};
|
||||
|
||||
std::priority_queue<Candidate, std::deque<Candidate>, decltype(cmp)> candidatePriorityQueue(cmp);
|
||||
for (auto i = 0; i < scores->lengthOf(); ++i) {
|
||||
if ((float)scoresData[i] > (float)scoreThreshold) {
|
||||
candidatePriorityQueue.emplace(Candidate({i, scoresData[i], 0}));
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<I> selected;
|
||||
T similarity, originalScore;
|
||||
Candidate nextCandidate;
|
||||
|
||||
while (selected.size() < outputSize && !candidatePriorityQueue.empty()) {
|
||||
nextCandidate = candidatePriorityQueue.top();
|
||||
originalScore = nextCandidate._score;
|
||||
candidatePriorityQueue.pop();
|
||||
|
||||
// Overlapping boxes are likely to have similar scores, therefore we
|
||||
// iterate through the previously selected boxes backwards in order to
|
||||
// see if `nextCandidate` should be suppressed. We also enforce a property
|
||||
// that a candidate can be suppressed by another candidate no more than
|
||||
// once via `suppress_begin_index` which tracks which previously selected
|
||||
// boxes have already been compared against next_candidate prior to a given
|
||||
// iteration. These previous selected boxes are then skipped over in the
|
||||
// following loop.
|
||||
bool shouldHardSuppress = false;
|
||||
for (int j = static_cast<int>(selected.size()) - 1; j >= nextCandidate._suppressBeginIndex; --j) {
|
||||
auto similarityA =
|
||||
f(*boxes, nextCandidate._boxIndex, selected[j]); // boxes->t<T>(nextCandidate._boxIndex, selected[j]);
|
||||
similarity = similarityA.template t<T>(0);
|
||||
nextCandidate._score *= T(similarity <= overlapThreshold ? 1.0 : 0.); // suppressWeightFunc(similarity);
|
||||
|
||||
// First decide whether to perform hard suppression
|
||||
if ((float)similarity >= static_cast<float>(overlapThreshold)) {
|
||||
shouldHardSuppress = true;
|
||||
break;
|
||||
}
|
||||
|
||||
// If next_candidate survives hard suppression, apply soft suppression
|
||||
if ((float)nextCandidate._score <= (float)scoreThreshold) break;
|
||||
}
|
||||
// If `nextCandidate._score` has not dropped below `scoreThreshold`
|
||||
// by this point, then we know that we went through all of the previous
|
||||
// selections and can safely update `suppress_begin_index` to
|
||||
// `selected.size()`. If on the other hand `next_candidate.score`
|
||||
// *has* dropped below the score threshold, then since `suppressWeight`
|
||||
// always returns values in [0, 1], further suppression by items that were
|
||||
// not covered in the above for loop would not have caused the algorithm
|
||||
// to select this item. We thus do the same update to
|
||||
// `suppressBeginIndex`, but really, this element will not be added back
|
||||
// into the priority queue in the following.
|
||||
nextCandidate._suppressBeginIndex = selected.size();
|
||||
|
||||
if (!shouldHardSuppress) {
|
||||
if (nextCandidate._score == originalScore) {
|
||||
// Suppression has not occurred, so select next_candidate
|
||||
selected.push_back(nextCandidate._boxIndex);
|
||||
}
|
||||
if ((float)nextCandidate._score > (float)scoreThreshold) {
|
||||
// Soft suppression has occurred and current score is still greater than
|
||||
// score_threshold; add next_candidate back onto priority queue.
|
||||
candidatePriorityQueue.push(nextCandidate);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (output) {
|
||||
DataBuffer buf(selected.data(), selected.size() * sizeof(I), DataTypeUtils::fromT<I>());
|
||||
output->dataBuffer()->copyBufferFrom(buf, buf.getLenInBytes());
|
||||
}
|
||||
|
||||
return (LongType)selected.size();
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void nonMaxSuppression(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize, double threshold,
|
||||
double scoreThreshold, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output->dataType(), nonMaxSuppressionV2_,
|
||||
(context, boxes, scales, maxSize, threshold, scoreThreshold, output), SD_FLOAT_TYPES,
|
||||
SD_INDEXING_TYPES);
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
LongType nonMaxSuppressionGeneric(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize,
|
||||
double threshold, double scoreThreshold, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output ? output->dataType() : DataType::INT32,
|
||||
return nonMaxSuppressionGeneric_,
|
||||
(context, boxes, scales, maxSize, threshold, scoreThreshold, output, similiratyOverlaps),
|
||||
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
return boxes->sizeAt(0);
|
||||
}
|
||||
|
||||
LongType nonMaxSuppressionV3(LaunchContext* context, NDArray* boxes, NDArray* scores, int maxSize,
|
||||
double overlapThreshold, double scoreThreshold, NDArray* output) {
|
||||
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output ? output->dataType() : DataType::INT32,
|
||||
return nonMaxSuppressionGeneric_,
|
||||
(context, boxes, scores, maxSize, overlapThreshold, scoreThreshold, output, similarityV3),
|
||||
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
return boxes->sizeAt(0);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,463 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 Oleh Semeniv (oleg.semeniv@gmail.com)
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/adjust_hue.h>
|
||||
#include <ops/declarable/helpers/imagesHelpers.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void rgbToYuvCuda(const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zTadOffsets,
|
||||
const LongType numOfTads, const int dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
rgbYuv<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void rgbToYuvCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t* stream,
|
||||
const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zTadOffsets, const LongType numOfTads,
|
||||
const int dimC) {
|
||||
rgbToYuvCuda<T><<<blocksPerGrid, threadsPerBlock, 256, *stream>>>(vx, xShapeInfo, xTadOffsets, vz, zShapeInfo,
|
||||
zTadOffsets, numOfTads, dimC);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "rgbToYuvCudaLauncher failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void transformRgbYuv(LaunchContext* context, NDArray& input, NDArray& output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output.shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (numOfTads + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
PointersManager manager(context, "yuv_to_rgb");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), rgbToYuvCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), packX->platformOffsets(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), packZ->platformOffsets(), numOfTads, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void yuvToRgbCuda(const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zTadOffsets,
|
||||
const LongType numOfTads, const int dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
yuvRgb<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void yuvToRgbCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t* stream,
|
||||
const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zTadOffsets, const LongType numOfTads,
|
||||
const int dimC) {
|
||||
yuvToRgbCuda<T><<<blocksPerGrid, threadsPerBlock, 256, *stream>>>(vx, xShapeInfo, xTadOffsets, vz, zShapeInfo,
|
||||
zTadOffsets, numOfTads, dimC);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "yuvToRgbCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void transformYuvRgb(LaunchContext* context, NDArray& input, NDArray& output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output.shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (numOfTads + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
PointersManager manager(context, "yuv_to_rgb");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), yuvToRgbCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), packX->platformOffsets(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), packZ->platformOffsets(), numOfTads, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// for example xShapeInfo = {2,3,4}, zShapeInfo = {2,1,4}
|
||||
template <typename T>
|
||||
SD_KERNEL void rgbToGrsCuda(const void* vx, const LongType* xShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const int dimC) {
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType zLen;
|
||||
__shared__ LongType rank;
|
||||
__shared__ const LongType* xShapePtr;
|
||||
__shared__ const LongType* zShapePtr;
|
||||
__shared__ const LongType* xStridePtr;
|
||||
__shared__ const LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = shape::rank(zShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
extern __shared__ unsigned char shmem[];
|
||||
auto coords = reinterpret_cast<LongType*>(shmem) + threadIdx.x * rank;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < zLen; i += gridDim.x * blockDim.x) {
|
||||
// Compute coordinates for the current index
|
||||
INDEX2COORDS(i, rank, zShapePtr, coords);
|
||||
|
||||
// Compute z offset
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStridePtr, coords, zOffset);
|
||||
|
||||
// Compute x offsets for R, G, B channels
|
||||
LongType xOffset0;
|
||||
COORDS2INDEX(rank, xStridePtr, coords, xOffset0);
|
||||
const auto xOffset1 = xOffset0 + xStridePtr[dimC];
|
||||
const auto xOffset2 = xOffset1 + xStridePtr[dimC];
|
||||
|
||||
// Convert RGB to grayscale
|
||||
z[zOffset] = 0.2989f * x[xOffset0] + 0.5870f * x[xOffset1] + 0.1140f * x[xOffset2];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void rgbToGrsCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const int dimC) {
|
||||
rgbToGrsCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, dimC);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "rgbToGrsCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void transformRgbGrs(LaunchContext* context, NDArray& input, NDArray& output, const int dimC) {
|
||||
PointersManager manager(context, "rgbToGrs");
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 4;
|
||||
const int blocksPerGrid = (input.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
const int sharedMem = input.rankOf() * sizeof(LongType) * threadsPerBlock + 128;
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), rgbToGrsCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, sharedMem, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), dimC),
|
||||
SD_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void SD_KERNEL rgbToHsvCuda(const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets,
|
||||
void* vz, const LongType* zShapeInfo, const LongType* zTadOffsets,
|
||||
const LongType numOfTads, const int dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
rgbToHsv<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void SD_KERNEL hsvToRgbCuda(const void* vx, const LongType* xShapeInfo, const LongType* xTadOffsets,
|
||||
void* vz, const LongType* zShapeInfo, const LongType* zTadOffsets,
|
||||
const LongType numOfTads, const int dimC) {
|
||||
const T* x = reinterpret_cast<const T*>(vx);
|
||||
T* z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ LongType xDimCstride, zDimCstride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < numOfTads; i += gridDim.x * blockDim.x) {
|
||||
const T* xTad = x + xTadOffsets[i];
|
||||
T* zTad = z + zTadOffsets[i];
|
||||
|
||||
hsvToRgb<T>(xTad[0], xTad[xDimCstride], xTad[2 * xDimCstride], zTad[0], zTad[zDimCstride], zTad[2 * zDimCstride]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_HOST void hsvToRgbCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xTadOffsets, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zTadOffsets, const LongType numOfTads,
|
||||
const int dimC) {
|
||||
hsvToRgbCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, xTadOffsets, vz, zShapeInfo,
|
||||
zTadOffsets, numOfTads, dimC);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "hsvToRgbCuda failed");
|
||||
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_HOST void rgbToHsvCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xTadOffsets, void* vz, const LongType* zShapeInfo,
|
||||
const LongType* zTadOffsets, const LongType numOfTads,
|
||||
const int dimC) {
|
||||
rgbToHsvCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory, *stream>>>(vx, xShapeInfo, xTadOffsets, vz, zShapeInfo,
|
||||
zTadOffsets, numOfTads, dimC);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "rgbToHsvCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void transformHsvRgb(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
dim3 launchDims = imageHelper(numOfTads);
|
||||
PointersManager manager(context, "hsv_to_rgb");
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), hsvToRgbCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z,context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), packX->platformOffsets(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), packZ->platformOffsets(), numOfTads, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void transformRgbHsv(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {dimC});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {dimC});
|
||||
|
||||
const LongType numOfTads = packX->numberOfTads();
|
||||
|
||||
dim3 launchDims = imageHelper(numOfTads);
|
||||
|
||||
PointersManager manager(context, "rgb_to_hsv");
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), rgbToHsvCudaLauncher,
|
||||
(launchDims.y, launchDims.x,launchDims.z, context->getCudaStream(), input->specialBuffer(),
|
||||
input->specialShapeInfo(), packX->platformOffsets(), output->specialBuffer(),
|
||||
output->specialShapeInfo(), packZ->platformOffsets(), numOfTads, dimC),
|
||||
SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void tripleTransformerCuda(const void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xTadShapeInfo, const LongType* xOffsets, void* vz,
|
||||
const LongType* zShapeInfo, const LongType* zTadShapeInfo,
|
||||
const LongType* zOffsets, const int dimC, int mode, uint64_t numTads) {
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType zLen, *sharedMem;
|
||||
__shared__ int rank; // xRank == zRank
|
||||
|
||||
float yiqarr[3][3] = {
|
||||
{0.299f, 0.59590059f, 0.2115f}, {0.587f, -0.27455667f, -0.52273617f}, {0.114f, -0.32134392f, 0.31119955f}};
|
||||
|
||||
float rgbarr[3][3] = {
|
||||
{1.f, 1.f, 1.f}, {0.95598634f, -0.27201283f, -1.10674021f}, {0.6208248f, -0.64720424f, 1.70423049f}};
|
||||
|
||||
auto tr = mode == 1 ? yiqarr : rgbarr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = shape::rank(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType* coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
if (dimC == (rank - 1) && 'c' == shape::order(xShapeInfo) && 1 == shape::elementWiseStride(xShapeInfo) &&
|
||||
'c' == shape::order(zShapeInfo) && 1 == shape::elementWiseStride(zShapeInfo)) {
|
||||
for (uint64_t f = blockIdx.x * blockDim.x + threadIdx.x; f < zLen / 3; f += gridDim.x * blockDim.x) {
|
||||
auto i = f * 3;
|
||||
|
||||
auto xi0 = x[i];
|
||||
auto xi1 = x[i + 1];
|
||||
auto xi2 = x[i + 2];
|
||||
|
||||
for (int e = 0; e < 3; e++) z[i + e] = xi0 * tr[0][e] + xi1 * tr[1][e] + xi2 * tr[2][e];
|
||||
}
|
||||
} else {
|
||||
// TAD based case
|
||||
const LongType xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
const LongType zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
|
||||
for (uint64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < numTads; i += blockDim.x * gridDim.x) {
|
||||
const T* xTad = x + xOffsets[i];
|
||||
T* zTad = z + zOffsets[i];
|
||||
|
||||
auto xi0 = xTad[0];
|
||||
auto xi1 = xTad[xDimCstride];
|
||||
auto xi2 = xTad[xDimCstride * 2];
|
||||
|
||||
for (int e = 0; e < 3; e++) zTad[zDimCstride * e] = xi0 * tr[0][e] + xi1 * tr[1][e] + xi2 * tr[2][e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void rgbYiq(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimC);
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimC);
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
dim3 launchDims = getLaunchDims("image_helpers_triple");
|
||||
tripleTransformerCuda<T><<<launchDims.x,launchDims.y, launchDims.z, *context->getCudaStream()>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), packX->platformShapeInfo(), packX->platformOffsets(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), packZ->platformShapeInfo(), packZ->platformOffsets(), dimC, 1,
|
||||
packZ->numberOfTads());
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "tripleTransformerCuda failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_INLINE static void yiqRgb(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), dimC);
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), dimC);
|
||||
|
||||
dim3 launchDims = getLaunchDims("image_helpers_triple");
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
tripleTransformerCuda<T><<<launchDims.x, launchDims.y,launchDims.z, *context->getCudaStream()>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), packX->platformShapeInfo(), packX->platformOffsets(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), packZ->platformShapeInfo(), packZ->platformOffsets(), dimC, 2,
|
||||
packZ->numberOfTads());
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "tripleTransformerCuda failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
void transformYiqRgb(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), yiqRgb, (context, input, output, dimC), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
void transformRgbYiq(LaunchContext* context, NDArray* input, NDArray* output, const int dimC) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), rgbYiq, (context, input, output, dimC), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,112 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <legacy/NativeOpExecutioner.h>
|
||||
#include <ops/declarable/helpers/reductions.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void argMax(NDArray& input, NDArray& output, const std::vector<LongType>& dimensions) {
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
if (output.isScalar()) {
|
||||
NativeOpExecutioner::execIndexReduceScalar(LaunchContext::defaultContext(), indexreduce::Ops::IndexMax,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo());
|
||||
} else {
|
||||
auto tadPack = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), (LongType*)&dimensions,dimensions.size());
|
||||
|
||||
NativeOpExecutioner::execIndexReduce(LaunchContext::defaultContext(), indexreduce::Ops::IndexMax, input.buffer(),
|
||||
input.shapeInfo(), input.specialBuffer(), input.specialShapeInfo(), nullptr,
|
||||
output.buffer(), output.shapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), (LongType*)nullptr, dimensions.size(),
|
||||
tadPack->specialShapeInfo(), tadPack->specialOffsets());
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
|
||||
void argMin(NDArray& input, NDArray& output, const std::vector<LongType>& dimensions) {
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
if (output.isScalar()) {
|
||||
NativeOpExecutioner::execIndexReduceScalar(LaunchContext::defaultContext(), indexreduce::Ops::IndexMin,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo());
|
||||
} else {
|
||||
auto tadPack = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), (LongType)&dimensions);
|
||||
|
||||
NativeOpExecutioner::execIndexReduce(LaunchContext::defaultContext(), indexreduce::Ops::IndexMin, input.buffer(),
|
||||
input.shapeInfo(), input.specialBuffer(), input.specialShapeInfo(), nullptr,
|
||||
output.buffer(), output.shapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), (LongType*)nullptr, dimensions.size(),
|
||||
tadPack->specialShapeInfo(), tadPack->specialOffsets());
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
|
||||
void argAbsMax(NDArray& input, NDArray& output, const std::vector<LongType>& dimensions) {
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
if (output.isScalar()) {
|
||||
NativeOpExecutioner::execIndexReduceScalar(LaunchContext::defaultContext(), indexreduce::Ops::IndexAbsoluteMax,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo());
|
||||
} else {
|
||||
auto tadPack = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), (LongType)&dimensions);
|
||||
|
||||
NativeOpExecutioner::execIndexReduce(LaunchContext::defaultContext(), indexreduce::Ops::IndexAbsoluteMax,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo(), (LongType*)nullptr,
|
||||
dimensions.size(), tadPack->specialShapeInfo(), tadPack->specialOffsets());
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
|
||||
void argAbsMin(NDArray& input, NDArray& output, const std::vector<LongType>& dimensions) {
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
if (output.isScalar()) {
|
||||
NativeOpExecutioner::execIndexReduceScalar(LaunchContext::defaultContext(), indexreduce::Ops::IndexAbsoluteMin,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo());
|
||||
} else {
|
||||
auto tadPack = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), (LongType)&dimensions);
|
||||
|
||||
NativeOpExecutioner::execIndexReduce(LaunchContext::defaultContext(), indexreduce::Ops::IndexAbsoluteMin,
|
||||
input.buffer(), input.shapeInfo(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), nullptr, output.buffer(), output.shapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo(), (LongType*)nullptr,
|
||||
dimensions.size(), tadPack->specialShapeInfo(), tadPack->specialOffsets());
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,112 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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, created on 21.09.2018
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
|
||||
#include <loops/special_kernels.h>
|
||||
#include <ops/declarable/helpers/ismax.h>
|
||||
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static void ismax_(LaunchContext* context, NDArray* input, NDArray* output,
|
||||
const std::vector<LongType>& dimensions) {
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
auto xRank = input->rankOf();
|
||||
auto zRank = output->rankOf();
|
||||
auto xType = input->dataType();
|
||||
auto zType = output->dataType();
|
||||
input->syncToDevice();
|
||||
LongType* special = nullptr;
|
||||
PointersManager manager(context, "IsMaxHelper");
|
||||
if (dimensions.size() == 0) {
|
||||
/**
|
||||
* In case of vector-input for IsMax, it just turns into IndexReduce call + subsequent filler call
|
||||
*/
|
||||
auto indexMax = input->applyIndexReduce(indexreduce::IndexMax, &dimensions);
|
||||
auto targetIdx = indexMax->e<LongType>(0);
|
||||
|
||||
dim3 launchDims = getLaunchDims("ismaxFill");
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
zType, fillIsMaxGeneric,
|
||||
(launchDims, stream, output->specialBuffer(),
|
||||
const_cast<sd::LongType *>(output->specialShapeInfo()), output->lengthOf(), targetIdx),
|
||||
SD_COMMON_TYPES);
|
||||
manager.synchronize();
|
||||
|
||||
delete indexMax;
|
||||
} else {
|
||||
LongType* hostYShapeInfo = nullptr;
|
||||
LongType* hostTShapeInfo = nullptr;
|
||||
LongType* dimension = nullptr;
|
||||
|
||||
LongType dimensionLength = dimensions.size();
|
||||
std::vector<LongType> copy(dimensions);
|
||||
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), copy.data(), copy.size());
|
||||
|
||||
// we launch legacy IndexMax op, to get indices of max values along dimension
|
||||
auto indexMaxArr = input->applyIndexReduce(indexreduce::IndexMax, &dimensions);
|
||||
|
||||
dim3 launchDims = getLaunchDims("ismax");
|
||||
dimension = (LongType*)manager.replicatePointer(dimensions.data(), dimensions.size() * sizeof(LongType));
|
||||
|
||||
// at this point, all IMax indexes are gathered, and we execute filler
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
zType, fillDimensionalIsMaxGeneric,
|
||||
(launchDims, stream, indexMaxArr->specialBuffer(), output->specialBuffer(),
|
||||
const_cast<sd::LongType *>(output->specialShapeInfo()),
|
||||
const_cast<sd::LongType *>(packZ->specialShapeInfo()),
|
||||
dimension, dimensionLength, const_cast<sd::LongType *>(packZ->specialOffsets())),
|
||||
SD_COMMON_TYPES);
|
||||
manager.synchronize();
|
||||
|
||||
delete indexMaxArr;
|
||||
}
|
||||
}
|
||||
|
||||
void ismax(LaunchContext* context, NDArray* input, NDArray* output, const std::vector<LongType>& dimensions) {
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), ismax_, (context, input, output, dimensions), SD_COMMON_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void ismax_,
|
||||
(sd::LaunchContext * context, NDArray* input, NDArray* output,
|
||||
const std::vector<sd::LongType>& dimensions),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,106 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <ops/declarable/helpers/legacy_helpers.h>
|
||||
#include <ops/ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
void reluDerivative__(NDArray* theFirst, NDArray* theSecond) {
|
||||
auto functor = LAMBDA_TT(x, y) { return x > (T)0.f ? y : T(0.f); });
|
||||
|
||||
theFirst->applyPairwiseLambda(theSecond, functor,theFirst);
|
||||
}
|
||||
|
||||
void reluDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), reluDerivative__, (theFirst, theSecond), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void reluDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return x > (T)0.f ? y : T(0.f); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void reluDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), reluDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void relu6Derivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return x > (T)0.f && x < (T)6.f ? y : T(0.f); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void relu6Derivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), relu6Derivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void leakyReluDerivative_(NDArray* input, NDArray* epsilon, NDArray* output, const float alpha) {
|
||||
const T alphaT = static_cast<T>(alpha);
|
||||
|
||||
auto functor = LAMBDA_TT(x, y, alphaT) { return x < 0 ? alphaT * y : y; });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void leakyReluDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput,
|
||||
const float alpha) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), leakyReluDerivative_, (theFirst, theSecond, theOutput, alpha),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void eluDerivative_(NDArray* input, NDArray* epsilon, NDArray* output, const float alpha) {
|
||||
const T alphaT = static_cast<T>(alpha);
|
||||
|
||||
auto functor = LAMBDA_TT(x, y, alphaT) { return y * math::sd_eluderivative<T, T>(x, alphaT); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void eluDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput,
|
||||
const float alpha) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), eluDerivative_, (theFirst, theSecond, theOutput, alpha), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void seluDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return y * simdOps::SELUDerivative<T>::op(x, nullptr); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void seluDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), seluDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,85 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <ops/declarable/helpers/legacy_helpers.h>
|
||||
#include <ops/ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void tanhDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
T th = math::sd_tanh<T, T>(x);
|
||||
return y * ((T)1.0f - (th * th));
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void tanhDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), tanhDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void hardTanhDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
T th = math::sd_tanh<T, T>(x);
|
||||
return y * simdOps::HardTanhDerivative<T>::op(x, nullptr);
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void hardTanhDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), hardTanhDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void rationalTanhDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return y * simdOps::RationalTanhDerivative<T>::op(x, nullptr); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void rationalTanhDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), rationalTanhDerivative_, (theFirst, theSecond, theOutput),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void rectifiedTanhDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return x > (T)0.0f ? y * (math::sd_tanhderivative<T, T>(x)) : (T)0.0f; });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void rectifiedTanhDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), rectifiedTanhDerivative_, (theFirst, theSecond, theOutput),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,236 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <ops/declarable/helpers/legacy_helpers.h>
|
||||
#include <ops/ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void cubeDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return y * (3 * x * x); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void cubeDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), cubeDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// return (x >= X(0.f) ? y: -y);
|
||||
template <typename T>
|
||||
void reduceNorm1_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return x > T(0.f) ? y : -y; });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void reduceNorm1(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), reduceNorm1_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void sigmCrossEntropy_(NDArray* logits, NDArray* labels, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
return math::sd_max<T>(x, (T)0.f) - x * y + math::sd_log<T, T>((T)1.f + math::sd_exp<T, T>(-math::sd_abs<T,T>(x)));
|
||||
});
|
||||
|
||||
logits->applyPairwiseLambda(labels, functor, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void sigmCrossEntropy(LaunchContext* context, NDArray* logits, NDArray* labels, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(logits->dataType(), sigmCrossEntropy_, (logits, labels, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void sigmCrossEntropyGrad_(NDArray* logits, NDArray* labels, NDArray* output) {
|
||||
// 1 - labels - 1 / (1 + exp(logits))
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
if (x <= 0) return static_cast<T>(1.) - y - static_cast<T>(1.) / (static_cast<T>(1.) + math::sd_exp<T, T>(x));
|
||||
auto e = math::sd_exp<T, T>(-x);
|
||||
return static_cast<T>(1.) - y - e / (static_cast<T>(1.) + e);
|
||||
});
|
||||
|
||||
logits->applyPairwiseLambda(labels, functor, output);
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void sigmCrossEntropyGrad(LaunchContext* context, NDArray* logits, NDArray* labels, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(logits->dataType(), sigmCrossEntropyGrad_, (logits, labels, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void softSignDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
T ss = (T)1.f + math::sd_abs<T,T>(x);
|
||||
return y * ((T)1.0f / (ss * ss));
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void softSignDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), softSignDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void softPlusDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
T p = math::sd_pow<T, T, T>(static_cast<T>(M_E), x);
|
||||
return y * (p / (p + 1.));
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void softPlusDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), softPlusDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
///
|
||||
/// \param input
|
||||
/// \param epsilon
|
||||
/// \param output
|
||||
template <typename T>
|
||||
void sigmoidDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) {
|
||||
T s = math::sd_sigmoid<T, T>(x);
|
||||
return y * (s * ((T)1.0f - s));
|
||||
});
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void sigmoidDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), sigmoidDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void hardSigmoidDerivative_(NDArray* input, NDArray* epsilon, NDArray* output) {
|
||||
auto functor = LAMBDA_TT(x, y) { return y * simdOps::HardSigmoidDerivative<T>::op(x, nullptr); });
|
||||
|
||||
input->applyPairwiseLambda(epsilon, functor, output);
|
||||
}
|
||||
|
||||
void hardSigmoidDerivative(LaunchContext* context, NDArray* theFirst, NDArray* theSecond, NDArray* theOutput) {
|
||||
BUILD_SINGLE_SELECTOR(theFirst->dataType(), hardSigmoidDerivative_, (theFirst, theSecond, theOutput), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void logSumExp_(NDArray* input, NDArray* axis, NDArray* output) {
|
||||
// reduce along axis with
|
||||
NDArray *tempInput = input->dup();
|
||||
input->applyTransform(transform::Exp, tempInput);
|
||||
std::vector<LongType> axisVector;
|
||||
if (axis != nullptr) {
|
||||
axisVector.resize(axis->lengthOf());
|
||||
for (size_t i = 0; i < axisVector.size(); ++i) axisVector[i] = axis->e<int>(i);
|
||||
}
|
||||
tempInput.reduceAlongDimension(reduce::Sum, output, &axisVector);
|
||||
output->applyTransform(transform::Log, output);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void logSumExp_(NDArray* input, NDArray* subtrah, NDArray* axis, NDArray* output) {
|
||||
// reduce along axis with
|
||||
NDArray *tempInput = input->dup();
|
||||
input->applyPairwiseTransform(pairwise::Subtract, subtrah, &tempInput);
|
||||
tempInput.applyTransform(transform::Exp, &tempInput);
|
||||
|
||||
std::vector<LongType> axisVector;
|
||||
if (axis != nullptr) {
|
||||
axisVector.resize(axis->lengthOf());
|
||||
for (size_t i = 0; i < axisVector.size(); ++i) axisVector[i] = axis->e<int>(i);
|
||||
}
|
||||
tempInput.reduceAlongDimension(reduce::Sum, output, &axisVector);
|
||||
output->applyTransform(transform::Log, output);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void logSumExp(LaunchContext* context, NDArray* input, NDArray* axis, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), logSumExp_, (input, axis, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void logSumExp(LaunchContext* context, NDArray* input, NDArray* subtrah, NDArray* axis, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), logSumExp_, (input, subtrah, axis, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
void weightedCrossEntropyWithLogitsFunctor_(NDArray * targets, NDArray * input, NDArray * weights,
|
||||
NDArray* output) {
|
||||
T posWeight = weights->e<T>(0);
|
||||
|
||||
auto mainRoutineT1 = LAMBDA_TT(_x, _z, posWeight) {
|
||||
T targetWeight = (1. + (posWeight - (T)1.f) * _z);
|
||||
return (1. - _z) * _x +
|
||||
targetWeight * (math::sd_log<T, T>((T)1.f + math::sd_exp<T, T>(-math::sd_abs<T,T>(_x))) +
|
||||
math::sd_max(-_x, T(0.f)));
|
||||
});
|
||||
|
||||
auto mainRoutineT2 = LAMBDA_TTT(_x, _z, _w) {
|
||||
return (((T)1.0 - _z) * _x) + _w * (math::sd_log<T, T>(T(1.) + math::sd_exp<T, T>(-math::sd_abs<T,T>(_x))) + math::sd_max(-_x, T(0.f)));
|
||||
});
|
||||
|
||||
if (weights->isScalar()) {
|
||||
input->applyPairwiseLambda(targets, mainRoutineT1, output);
|
||||
} else {
|
||||
std::unique_ptr<NDArray> targetVector(new NDArray(*weights));
|
||||
targetVector->applyScalar(scalar::Add, -1.f, targetVector.get());
|
||||
|
||||
*targets = (*targetVector * *targets) + T(1.f);
|
||||
input->applyPairwiseLambda(targets, mainRoutineT1, output);
|
||||
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void weightedCrossEntropyWithLogitsFunctor(LaunchContext* context, NDArray * targets, NDArray * input,
|
||||
NDArray * weights, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {targets, input, weights});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(targets->dataType(), weightedCrossEntropyWithLogitsFunctor_, (targets, input, weights, output),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({output}, {targets, input, weights});
|
||||
}
|
||||
|
||||
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,52 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 George A. Shulinok <sgazeos@gmail.com>
|
||||
//
|
||||
#include <ops/declarable/helpers/lgamma.h>
|
||||
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// calculate digamma function for array elements
|
||||
template <typename T>
|
||||
void lgamma_(NDArray* x, NDArray* z) {
|
||||
auto lgammaProc = LAMBDA_T(x_) {
|
||||
return T(DataTypeUtils::fromT<T>() == DOUBLE ? ::lgamma(x_)
|
||||
: ::lgammaf(x_));
|
||||
});
|
||||
|
||||
x->applyLambda(lgammaProc, z);
|
||||
}
|
||||
|
||||
void lgamma(LaunchContext* context, NDArray* x, NDArray* z) {
|
||||
BUILD_SINGLE_SELECTOR(x->dataType(), lgamma_, (x, z), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void lgamma_, (NDArray * x, NDArray* z), SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,188 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <ops/declarable/helpers/lrn.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void lrnKernel(void* vx, LongType const* xTadShapeInfo, LongType const* xTadOffsets, void* vz,
|
||||
LongType const* zTadShapeInfo, LongType const* zTadOffsets, LongType numTads,
|
||||
LongType tadLength, int depth, double bias, double alpha,
|
||||
double beta) {
|
||||
extern __shared__ char sharedChar[];
|
||||
T* shared = reinterpret_cast<T*>(sharedChar);
|
||||
|
||||
auto xEws = shape::elementWiseStride(xTadShapeInfo);
|
||||
auto zEws = shape::elementWiseStride(zTadShapeInfo);
|
||||
|
||||
auto xOrder = shape::order(xTadShapeInfo);
|
||||
auto zOrder = shape::order(zTadShapeInfo);
|
||||
|
||||
const T tbias = static_cast<T>(bias);
|
||||
const T tbeta = static_cast<T>(beta);
|
||||
const T talpha = static_cast<T>(alpha);
|
||||
|
||||
// one block of threads processes 1 example within batch
|
||||
for (LongType i = blockIdx.x; i < numTads; i += gridDim.x) {
|
||||
auto x = reinterpret_cast<T*>(vx) + xTadOffsets[i];
|
||||
auto z = reinterpret_cast<T*>(vz) + zTadOffsets[i];
|
||||
|
||||
// load everything into shared memory, so we'll operate on shared memory from now on
|
||||
shared[threadIdx.x] = x[threadIdx.x * xEws];
|
||||
__syncthreads();
|
||||
|
||||
const LongType begin = sd::math::sd_max<int>(0, threadIdx.x - depth);
|
||||
const LongType last = depth + threadIdx.x + 1;
|
||||
const LongType end = sd::math::sd_min<int>(last, tadLength);
|
||||
|
||||
T prev = static_cast<T>(0.);
|
||||
for (int s = begin; s < end; s++) prev = prev + shared[s] * shared[s];
|
||||
|
||||
z[threadIdx.x * zEws] = shared[threadIdx.x] / math::sd_pow<T, T, T>(tbias + alpha * prev, tbeta);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void lrnBPKernel(void const* vx, LongType const* xTadShapeInfo, LongType const* xTadOffsets,
|
||||
void* vz,
|
||||
LongType const* zTadShapeInfo, LongType const* zTadOffsets, LongType numTads,
|
||||
LongType tadLength, int depth, double bias, double alpha,
|
||||
double beta) {
|
||||
extern __shared__ char sharedChar[];
|
||||
X* sharedX = reinterpret_cast<X*>(sharedChar);
|
||||
Z* sharedY = reinterpret_cast<Z*>(sharedX + blockDim.x);
|
||||
|
||||
auto xEws = shape::elementWiseStride(xTadShapeInfo);
|
||||
auto zEws = shape::elementWiseStride(zTadShapeInfo);
|
||||
|
||||
auto xOrder = shape::order(xTadShapeInfo);
|
||||
auto zOrder = shape::order(zTadShapeInfo);
|
||||
|
||||
const Z tbias = static_cast<Z>(bias);
|
||||
const Z tbeta = static_cast<Z>(beta);
|
||||
const Z talpha = static_cast<Z>(alpha);
|
||||
const Z coeff = talpha * tbeta;
|
||||
|
||||
for (LongType i = blockIdx.x; i < numTads; i += gridDim.x) {
|
||||
auto x = reinterpret_cast<X const*>(vx) + xTadOffsets[i];
|
||||
auto z = reinterpret_cast<Z*>(vz) + zTadOffsets[i];
|
||||
|
||||
const LongType begin = sd::math::sd_max<int>(0, threadIdx.x - depth);
|
||||
const LongType last = depth + threadIdx.x + 1;
|
||||
const LongType end = sd::math::sd_min<int>(last, tadLength);
|
||||
|
||||
// load everything into shared memory
|
||||
sharedX[threadIdx.x] = x[threadIdx.x * xEws];
|
||||
sharedY[threadIdx.x] = static_cast<Z>(0.f);
|
||||
__syncthreads();
|
||||
|
||||
// we're operating in shared memory
|
||||
for (int s = begin; s < end; s++) sharedY[threadIdx.x] = sharedY[threadIdx.x] + sharedX[s] * sharedX[s];
|
||||
__syncthreads();
|
||||
|
||||
Z factor[1024];
|
||||
Z init = tbias + talpha * sharedY[threadIdx.x];
|
||||
|
||||
Z prev = static_cast<Z>(0.f);
|
||||
for (LongType s = begin; s < end; ++s) {
|
||||
factor[s] = math::sd_pow<Z, Z, Z>(tbias + talpha * sharedY[s], -tbeta - 1);
|
||||
prev = prev + sharedX[s] * factor[s];
|
||||
}
|
||||
|
||||
z[threadIdx.x * zEws] = factor[threadIdx.x] * init - 2 * sharedX[threadIdx.x] * coeff * prev;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
static void lrnBP_(graph::Context& block, NDArray& input, NDArray& gradO, NDArray& gradI,
|
||||
const int depth, const float bias, const float alpha, const float beta) {
|
||||
auto rank = input.rankOf();
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), {rank - 1});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(gradI.shapeInfo(), {rank - 1});
|
||||
|
||||
const auto tadLength = shape::length(packX->primaryShapeInfo());
|
||||
const int numThreads = tadLength;
|
||||
|
||||
if (tadLength > 1024 || tadLength < 1) THROW_EXCEPTION("LRN: tadLength > 1024 isn't implemented yet");
|
||||
|
||||
dim3 launchDims = lrnDims(tadLength,packX->numberOfTads(),DataTypeUtils::sizeOf(input.dataType()),DataTypeUtils::sizeOf(gradI.dataType()));
|
||||
|
||||
lrnBPKernel<X, Z><<<launchDims.y, launchDims.x, launchDims.z,
|
||||
*block.launchContext()->getCudaStream()>>>(
|
||||
input.specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), gradI.specialBuffer(),
|
||||
packZ->platformShapeInfo(), packZ->platformOffsets(), packX->numberOfTads(), tadLength, depth, bias, alpha, beta);
|
||||
|
||||
gradI.tickWriteDevice();
|
||||
gradI *= gradO;
|
||||
}
|
||||
|
||||
void lrnBP(graph::Context& block, NDArray& input, NDArray& gradO, NDArray& gradI, const int depth,
|
||||
const float bias, const float alpha, const float beta) {
|
||||
input.syncToDevice();
|
||||
gradO.syncToDevice();
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), gradO.dataType(), lrnBP_,
|
||||
(block, input, gradO, gradI, depth, bias, alpha, beta), SD_FLOAT_TYPES, SD_FLOAT_TYPES);
|
||||
|
||||
gradI.tickWriteDevice();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void lrnFunctor_(graph::Context& block, NDArray* input, NDArray* output, int depth, double bias, double alpha,
|
||||
double beta) {
|
||||
auto rank = input->rankOf();
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {rank - 1});
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), {rank - 1});
|
||||
|
||||
const auto tadLength = shape::length(packX->primaryShapeInfo());
|
||||
const int numBlocks = sd::math::sd_min<LongType>(1024, packX->numberOfTads());
|
||||
const int numThreads = tadLength;
|
||||
dim3 launchDims = lrnDims(tadLength, packX->numberOfTads(), DataTypeUtils::sizeOf(input->dataType()),
|
||||
DataTypeUtils::sizeOf(input->dataType()));
|
||||
|
||||
if (tadLength > 1024 || tadLength < 1) THROW_EXCEPTION("LRN: tadLength > 1024 isn't implemented yet");
|
||||
|
||||
lrnKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *block.launchContext()->getCudaStream()>>>(
|
||||
input->specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), output->specialBuffer(),
|
||||
packZ->platformShapeInfo(), packZ->platformOffsets(), packX->numberOfTads(), tadLength, depth, bias, alpha, beta);
|
||||
}
|
||||
|
||||
Status lrnFunctor(graph::Context& block, NDArray* input, NDArray* output, int depth, double bias, double alpha,
|
||||
double beta) {
|
||||
input->syncToDevice();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), lrnFunctor_, (block, input, output, depth, bias, alpha, beta),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
output->tickWriteDevice();
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,199 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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, created on 14.02.2018
|
||||
//
|
||||
|
||||
// implementation of operation for LSTM cell with peep hole connections:
|
||||
// http://www.bioinf.jku.at/publications/older/2604.pdf
|
||||
// S. Hochreiter and J. Schmidhuber. "Long Short-Term Memory". Neural Computation, 9(8):1735-1780, 1997.
|
||||
// and
|
||||
// https://research.google.com/pubs/archive/43905.pdf
|
||||
// Hasim Sak, Andrew Senior, and Francoise Beaufays. "Long short-term memory recurrent neural network architectures for
|
||||
// large scale acoustic modeling." INTERSPEECH, 2014.
|
||||
|
||||
#include <array/NDArrayList.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/CustomOperations.h>
|
||||
#include <ops/declarable/helpers/lstm.h>
|
||||
#include <ops/declarable/helpers/lstmBlock.h>
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <iterator>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void lstmCell(LaunchContext* context, NDArray* xt, NDArray* ht_1, NDArray* ct_1,
|
||||
NDArray* Wx, NDArray* Wh, NDArray* Wc, NDArray* Wp, NDArray* b, NDArray* ht,
|
||||
NDArray* ct, const std::vector<double>& params) {
|
||||
// xt input [bS x nIn]
|
||||
// ht_1 previous cell output [bS x numProj], that is at previous time step t-1, in case of projection=false ->
|
||||
// numProj=nOut!!! ct_1 previous cell state [bS x nOut], that is at previous time step t-1
|
||||
|
||||
// Wx input-to-hidden weights, [nIn x 4*nOut]
|
||||
// Wh hidden-to-hidden weights, [numProj x 4*nOut]
|
||||
// Wc diagonal weights for peephole connections [3*nOut]
|
||||
// Wp projection weights [nOut x numProj]
|
||||
// b biases, [4*nOut]
|
||||
|
||||
// ht current cell output [bS x numProj], that is at current time step t
|
||||
// ct current cell state [bS x nOut], that is at current time step t
|
||||
|
||||
const bool peephole = (bool)params[0]; // if true, provide peephole connections
|
||||
const bool projection =
|
||||
(bool)params[1]; // if true, then projection is performed, if false then numProj==nOut is mandatory!!!!
|
||||
double clippingCellValue =
|
||||
params[2]; // clipping value for ct, if it is not equal to zero, then cell state is clipped
|
||||
double clippingProjValue =
|
||||
params[3]; // clipping value for projected ht, if it is not equal to zero, then projected cell output is clipped
|
||||
const double forgetBias = params[4];
|
||||
|
||||
const int bS = xt->sizeAt(0);
|
||||
const int nIn = xt->sizeAt(1);
|
||||
const int numProj = ht_1->sizeAt(1);
|
||||
const int nOut = ct_1->sizeAt(1);
|
||||
|
||||
auto z = mmul(*xt, *Wx) + mmul(*ht_1, *Wh) + *b; // [bS x 4*nOut] + [bS x 4*nOut] + [1 x 4*nOut] = [bS x 4*nOut]
|
||||
|
||||
auto zit = z({0, 0, 0, nOut}); // z for input gate, = mmul(Wxi,xt) + mmul(Whi,ht_1) + bi = [bS x nOut]
|
||||
auto zft = z({0, 0, nOut, 2 * nOut}); // z for forget gate, = mmul(Wxf,xt) + mmul(Whf,ht_1) + bf = [bS x nOut]
|
||||
auto zct = z({0, 0, 2 * nOut, 3 * nOut}); // z for cell state, = mmul(Wxc,xt) + mmul(Whc,ht_1) + bc = [bS x nOut]
|
||||
auto zot = z({0, 0, 3 * nOut, 4 * nOut}); // z for output gate, = mmul(Wxo,xt) + mmul(Who,ht_1) + bo = [bS x nOut]
|
||||
|
||||
if (peephole) { // add peephole connections: z + ct_1*Wc
|
||||
zit += (*ct_1) * (*Wc)({0, nOut}); // add peephole connections to input gate
|
||||
zft += (*ct_1) * (*Wc)({nOut, 2 * nOut}); // add peephole connections to forget gate
|
||||
}
|
||||
|
||||
// current sell state = ft*ct_1 + it*tanh(mmul(Wxc,xt) + mmul(Whc,ht_1) + bc
|
||||
NDArray zftPlusForgetBias = zft + forgetBias;
|
||||
NDArray toAssign = sigmoid(zftPlusForgetBias) * (*ct_1) + sigmoid(zit) * tanh(zct);
|
||||
ct->assign(&toAssign);
|
||||
|
||||
// if clipping value is provided then cell state is clipped by this value prior to the cell output activation
|
||||
if (clippingCellValue > 0.0) ct->applyScalar(scalar::LstmClip, clippingCellValue, ct);
|
||||
|
||||
if (peephole) zot += (*ct) * (*Wc)({{2 * nOut, 3 * nOut}}); // add peephole connections to output gate zot + ct*Wc
|
||||
|
||||
// current cell output = ot*tanh(ct)
|
||||
auto htNoPeepHole = sigmoid(zot) * tanh(*ct); // = [bS x nOut]
|
||||
|
||||
// apply projection
|
||||
if (projection) {
|
||||
NDArray restultOne = mmul(htNoPeepHole, *Wp);
|
||||
ht->assign(&restultOne); // [bS x nOut] * [ nOut x numProj] = [bS x numProj]
|
||||
// if clipping projection is provided then projected cell output state is clipped by this value
|
||||
if (clippingProjValue != 0.) ht->applyScalar(scalar::LstmClip, clippingProjValue, ht);
|
||||
} else
|
||||
ht->assign(&htNoPeepHole);
|
||||
}
|
||||
|
||||
void lstmBlockCell(NDArray* xt, NDArray* cLast, NDArray* yLast, NDArray* W, NDArray* Wci,
|
||||
NDArray* Wcf, NDArray* Wco, NDArray* b, NDArray* i, NDArray* c, NDArray* f,
|
||||
NDArray* o, NDArray* z, NDArray* h, NDArray* y, const std::vector<double>& params) {
|
||||
/* Input arrays:
|
||||
* 0: xt - input [bS, nIn] at time t
|
||||
* 1: cLast (cs_prev) - previous cell state [bS, nOut], time t-1
|
||||
* 2: yLast (h_prev) - previous output [bS, nOut], time t-1
|
||||
* 3: W - Weights - concatenated (input-to-hidden, hidden-to-hidden weights) weights, [(nIn+nOut),
|
||||
* 4*nOut] 4: Wci - weights - cell peephole (t-1) connections to input modulation gate, [nOut] 5: Wcf -
|
||||
* weights - cell peephole (t-1) connections to forget gate, [nOut] 6: Wco - weights - cell peephole (t)
|
||||
* connections to output gate, [nOut] 7: b - biases, [4*nOut]
|
||||
*
|
||||
* Input integer arguments:
|
||||
* 0: if not zero, provide peephole connections
|
||||
*
|
||||
* Input float arguments:
|
||||
* 0: the bias added to forget gates in order to reduce the scale of forgetting in the beginning of the training
|
||||
* 1: clipping value for cell state, if it is not equal to zero, then cell state is clipped
|
||||
*
|
||||
* Output arrays:
|
||||
* 0: i - Input modulation gate activations [bS, nOut]
|
||||
* 1: c (cs) - Cell state (pre tanh) [bs, nOut] (cs)
|
||||
* 2: f - Output - forget gate activations [bs, nOut]
|
||||
* 3: o - Output - output gate activations [bs, nOut]
|
||||
* 4: z (ci) - Output - block input [bs, nOut]
|
||||
* 5: h (co) - Cell state, post tanh [bs, nOut]
|
||||
* 6: y (h) - Current cell output [bS, nOut], time t
|
||||
*/
|
||||
const bool peephole = (bool)params[0]; // if true, provide peephole connections
|
||||
const double forgetBias = params[1];
|
||||
const double clippingCellValue =
|
||||
params[2]; // clipping value for ct, if it is not equal to zero, then cell state is clipped
|
||||
|
||||
const int bS = xt->sizeAt(0);
|
||||
const int nIn = xt->sizeAt(1);
|
||||
const int nOut = cLast->sizeAt(1);
|
||||
|
||||
std::vector<sd::LongType> shape = {xt->sizeAt(0), xt->sizeAt(1) + yLast->sizeAt(1)};
|
||||
// Concat inputs: [xt, yt-1]: concat([bs,nIn],[bs,nOut]) -> [bs, (nIn+nOut)]
|
||||
NDArray concatOut(xt->ordering(), shape, xt->dataType(),
|
||||
xt->getContext());
|
||||
concat(xt->getContext(), {const_cast<NDArray*>(xt), const_cast<NDArray*>(yLast)}, concatOut, {1});
|
||||
|
||||
auto m = mmul(concatOut, *W); // mmul: [bs, (nIn+nOut)] * [(nIn+nOut), 4*nOut] = [bs, 4*nOut]
|
||||
m += (*b); // addiRowVector
|
||||
|
||||
// Note: weights are ordered [inputGate, blockInput, forgetGate, outputGate] to match TF (TF code comments state
|
||||
// [i,f,z/ci,o] but behaviour is [i,z,f,o])
|
||||
auto zi = m({0, 0, 0, nOut}); // z for input modulation gate, [bS, nOut]
|
||||
auto zz = m({0, 0, nOut, 2 * nOut}); // z for block input, [bS, nOut]
|
||||
auto zf = m({0, 0, 2 * nOut, 3 * nOut}); // z for forget gate, [bS, nOut]
|
||||
auto zo = m({0, 0, 3 * nOut, 4 * nOut}); // z for output gate, [bS, nOut]
|
||||
|
||||
if (peephole) { // add peephole connections: z + ct_1*Wc
|
||||
zi += (*cLast) * (*Wci); // add peephole connections to input gate
|
||||
zf += (*cLast) * (*Wcf); // add peephole connections to forget gate
|
||||
}
|
||||
|
||||
// current sell state = ft*cLast + it*tanh(mmul(Wxc,xt) + mmul(Whc,ht_1) + bc
|
||||
if (forgetBias != 0.0) zf += forgetBias;
|
||||
|
||||
zz.applyTransform(transform::Tanh, z); // z = tanh(zz)
|
||||
zi.applyTransform(transform::Sigmoid, i); // i = sigmoid(zi)
|
||||
zf.applyTransform(transform::Sigmoid, f); // f = sigmoid(zf);
|
||||
|
||||
// cell state = blockInput .* inputGate + prevCellState .* forgetGate
|
||||
z->applyPairwiseTransform(pairwise::Multiply, i, c); // c = z * i
|
||||
auto temp = (*f) * (*cLast);
|
||||
*c += temp; // c = (i * z) + (zf * (*cLast))
|
||||
c->applyTransform(transform::Tanh, h); // h = tanh(c)
|
||||
|
||||
// if clipping value is provided then cell state is clipped by this value prior to the cell output activation
|
||||
if (clippingCellValue > 0.0) c->applyScalar(scalar::LstmClip, clippingCellValue, c);
|
||||
|
||||
if (peephole) {
|
||||
// add peephole connections to output gate zot + ct*Wc
|
||||
auto prod = *c * (*Wco);
|
||||
zo += prod;
|
||||
}
|
||||
zo.applyTransform(transform::Sigmoid, o); // o = sigmoid(zo)
|
||||
|
||||
// current cell output = ot*tanh(ct)
|
||||
c->applyTransform(transform::Tanh, h); // h = tanh(c)
|
||||
o->applyPairwiseTransform(pairwise::Multiply, h, y); // y = o * h
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,132 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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/NDArray.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/lstsq.h>
|
||||
#include <ops/declarable/helpers/lup.h>
|
||||
#include <ops/declarable/helpers/qr.h>
|
||||
#include <ops/declarable/helpers/triangular_solve.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void fillRegularizerKernel(T* ioMatrixData, const LongType* ioMatrixShape,
|
||||
const LongType* ioMatrixTads, const LongType* ioMatrixOffsets,
|
||||
LongType batchSize, LongType rows, T const value) {
|
||||
for (auto x = blockIdx.x; x < batchSize; x += gridDim.x) {
|
||||
auto z = ioMatrixData + ioMatrixOffsets[x];
|
||||
for (auto r = threadIdx.x; r < rows; r += blockDim.x) {
|
||||
LongType pos[] = {r, r};
|
||||
LongType zIndex;
|
||||
COORDS2INDEX(2, shape::stride(ioMatrixTads), pos, zIndex);
|
||||
z[zIndex] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void fillRegularizer(LaunchContext* context, NDArray* ioMatrix, double const value) {
|
||||
std::vector<LongType> dims = {-2, -1};
|
||||
auto lastDimsTads = ConstantTadHelper::getInstance().tadForDimensions(ioMatrix->shapeInfo(), &dims);
|
||||
auto stream = context->getCudaStream();
|
||||
auto rows = ioMatrix->sizeAt(-2);
|
||||
dim3 launchDims = getLaunchDims("lstsq_reg");
|
||||
fillRegularizerKernel<T><<<launchDims.y,launchDims.x,launchDims.z, *stream>>>(
|
||||
ioMatrix->dataBuffer()->specialAsT<T>(), ioMatrix->specialShapeInfo(), lastDimsTads->specialShapeInfo(),
|
||||
lastDimsTads->specialOffsets(), lastDimsTads->numberOfTads(), rows, (T)value);
|
||||
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status leastSquaresSolveFunctor_(LaunchContext* context, NDArray* leftInput, NDArray* rightInput,
|
||||
double const l2Regularizer, bool const fast, NDArray* output) {
|
||||
if (fast) { // Cholesky decomposition approach
|
||||
// Equation for solve A^T * Ax = A^T * b, so
|
||||
// 1. Computing A2:
|
||||
auto tAtShape = ShapeUtils::evalShapeForMatmul(leftInput->shapeInfo(), leftInput->shapeInfo(), true, false);
|
||||
// tAtShape[tAtShape.size() - 2] = output->sizeAt(-2);
|
||||
NDArray leftOutput(leftInput->ordering(), tAtShape, output->dataType(), context);
|
||||
MmulHelper::matmul(leftInput, leftInput, &leftOutput, true, false,1.0,0.0,&leftOutput); // Computing A2 = A^T * A
|
||||
// 2. Computing B' = A^T * b
|
||||
auto rightOutput = output->ulike();
|
||||
|
||||
MmulHelper::matmul(leftInput, rightInput, rightOutput, true, false,1.0,0.0,rightOutput); // Computing B' = A^T * b
|
||||
// 3. Regularization ( indeed A' = A2 - l2Regularizer * I)
|
||||
if (l2Regularizer != 0.0) {
|
||||
auto regularizer = leftOutput.ulike();
|
||||
regularizer->nullify();
|
||||
fillRegularizer<T>(context, regularizer, (T)l2Regularizer);
|
||||
leftOutput += *regularizer;
|
||||
}
|
||||
|
||||
// 4. Cholesky decomposition -- output matrix is square and lower triangular
|
||||
cholesky(context, &leftOutput, &leftOutput, true); // inplace decomposition
|
||||
// 5. Solve two triangular systems:
|
||||
auto rightB = rightOutput->ulike();
|
||||
rightB->nullify();
|
||||
|
||||
triangularSolveFunctor(context, &leftOutput, rightOutput, true, false, rightB);
|
||||
|
||||
adjointMatrix(context, &leftOutput, true, &leftOutput);
|
||||
triangularSolveFunctor(context, &leftOutput, rightB, false, false, output);
|
||||
// All done
|
||||
} else { // QR decomposition approach
|
||||
// Equation for solve Rx = Q^T * b, where A = Q * R, where Q - orthogonal matrix, and R - upper triangular
|
||||
// 1. QR decomposition
|
||||
auto* qShapePtr = leftInput->getShapeAsVector();
|
||||
std::vector<LongType> qShape = *qShapePtr;
|
||||
delete qShapePtr;
|
||||
auto* rShapePtr = leftInput->getShapeAsVector();
|
||||
std::vector<LongType> rShape = *rShapePtr;
|
||||
delete rShapePtr;
|
||||
qShape[leftInput->rankOf() - 1] = leftInput->sizeAt(-2);
|
||||
|
||||
NDArray Q(leftInput->ordering(), qShape, leftInput->dataType(), context);
|
||||
NDArray R(leftInput->ordering(), rShape, leftInput->dataType(), context);
|
||||
qr(context, leftInput, &Q, &R, true);
|
||||
// 2. b` = Q^t * b:
|
||||
auto rightOutput = rightInput->ulike();
|
||||
MmulHelper::matmul(&Q, rightInput, rightOutput, true, false,1.0,0.0,rightOutput);
|
||||
// 3. Solve triangular system
|
||||
triangularSolveFunctor(context, &R, rightOutput, false, false, output);
|
||||
}
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
Status leastSquaresSolveFunctor(LaunchContext* context, NDArray* leftInput, NDArray* rightInput,
|
||||
double const l2Regularizer, bool const fast, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(leftInput->dataType(), return leastSquaresSolveFunctor_,
|
||||
(context, leftInput, rightInput, l2Regularizer, fast, output), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,122 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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)
|
||||
//
|
||||
#include <array/ResultSet.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/matrixSetDiag.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void matrixSetDiagCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const bool zeroPad) {
|
||||
// x - input, shape [A,B,C]
|
||||
// y - diagonal, shape [A,B]
|
||||
// z - output, shape [A,B,C]
|
||||
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
const auto y = reinterpret_cast<const T*>(vy);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int xRank;
|
||||
__shared__ LongType xLen;
|
||||
__shared__ const LongType *shapeX, *strideX, *strideY, *strideZ;
|
||||
__shared__ bool areSameOffsets;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xLen = shape::length(xShapeInfo);
|
||||
shapeX = shape::shapeOf(xShapeInfo);
|
||||
strideX = shape::stride(xShapeInfo);
|
||||
strideY = shape::stride(yShapeInfo);
|
||||
strideZ = shape::stride(zShapeInfo);
|
||||
areSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (LongType i = tid; i < xLen; i += step) {
|
||||
// Compute coordinates and offsets
|
||||
INDEX2COORDS(i, xRank, shapeX, coords);
|
||||
LongType xOffset, zOffset, yOffset;
|
||||
|
||||
COORDS2INDEX(xRank, strideX, coords, xOffset);
|
||||
zOffset = areSameOffsets ? xOffset : 0;
|
||||
if (!areSameOffsets) {
|
||||
COORDS2INDEX(xRank, strideZ, coords, zOffset);
|
||||
}
|
||||
|
||||
// Check if on the diagonal
|
||||
if (coords[xRank - 2] == coords[xRank - 1]) {
|
||||
COORDS2INDEX(xRank - 1, strideY, coords, yOffset);
|
||||
z[zOffset] = y[yOffset];
|
||||
} else {
|
||||
z[zOffset] = zeroPad ? static_cast<T>(0) : x[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void matrixSetDiagCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const bool zeroPad) {
|
||||
matrixSetDiagCuda<T>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, zeroPad);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "matrixSetDiagCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void matrixSetDiag(LaunchContext* context, NDArray& input, NDArray& diagonal, NDArray& output,
|
||||
const bool zeroPad) {
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (input.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
const int sharedMem = threadsPerBlock * sizeof(LongType) * input.rankOf() + 128;
|
||||
|
||||
dim3 launchDims = matrixSetDiagDims(input.lengthOf(),input.rankOf());
|
||||
PointersManager manager(context, "matrixSetDiag");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input, &diagonal});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), matrixSetDiagCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), diagonal.specialBuffer(), diagonal.specialShapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo(), zeroPad),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input, &diagonal});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,126 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 George A. Shulinok <sgazeos@gmail.com>
|
||||
//
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/matrix_band.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// matrix band kernel
|
||||
//
|
||||
// inputBuffer - buffer of input tensor
|
||||
// inputShape - shape of input tensor
|
||||
// outputBuffer - buffer of output tensor
|
||||
// outputShape - shape of output tensor
|
||||
// lowerBand - lower band of matrix
|
||||
// upperBand - upper band of matrix
|
||||
// tadOnlyInputShapeInfo - TAD shape for input
|
||||
// tadInputOffsets - TAD offsets for input
|
||||
// tadOnlyOutputShapeInfo - TAD output shape
|
||||
// tadOutputOffsets - TAD output offsets
|
||||
// numTads - number of subarrays
|
||||
// inputLength - input subarray length
|
||||
//
|
||||
template <typename T>
|
||||
static SD_KERNEL void matrixBandKernel(const void* inputBuffer, const LongType* inputShape, void* outputBuffer,
|
||||
const LongType* outputShape, LongType lowerBand, LongType upperBand,
|
||||
const LongType* tadOnlyInputShapeInfo, const LongType* tadInputOffsets,
|
||||
const LongType* tadOnlyOutputShapeInfo, const LongType* tadOutputOffsets,
|
||||
LongType numTads, LongType inputLength) {
|
||||
int totalThreads = blockDim.x;
|
||||
LongType rows = shape::sizeAt(inputShape, -2);
|
||||
LongType cols = shape::sizeAt(inputShape, -1);
|
||||
auto resetBuffer = reinterpret_cast<T *>(outputBuffer);
|
||||
auto input = reinterpret_cast<T const *>(inputBuffer);
|
||||
|
||||
for (LongType e = blockIdx.x; e < numTads; e += gridDim.x) {
|
||||
auto yOffset = tadInputOffsets[e];
|
||||
auto xOffset = tadOutputOffsets[e];
|
||||
if (outputBuffer != inputBuffer) // if not inplace
|
||||
for(int i = 0; i < inputLength; i++) {
|
||||
resetBuffer[i] = input[i];
|
||||
}
|
||||
for (LongType i = blockIdx.y; i < rows; i += gridDim.y) {
|
||||
for (LongType j = threadIdx.x; j < cols; j += totalThreads) {
|
||||
LongType coords[2] = {i, j};
|
||||
LongType tadOffsetOut, tadOffsetIn;
|
||||
COORDS2INDEX(shape::rank(tadOnlyOutputShapeInfo), shape::stride(tadOnlyOutputShapeInfo), coords, tadOffsetOut);
|
||||
COORDS2INDEX(shape::rank(tadOnlyInputShapeInfo), shape::stride(tadOnlyInputShapeInfo), coords, tadOffsetIn);
|
||||
|
||||
// If not inplace, copy the input to the output
|
||||
*(resetBuffer + xOffset + tadOffsetOut) = *(input + yOffset + tadOffsetIn);
|
||||
|
||||
// Check the lower diagonals
|
||||
if (lowerBand >= 0 && (i - j) > lowerBand)
|
||||
*(resetBuffer + xOffset + tadOffsetOut) = T(0);
|
||||
|
||||
// Check the upper diagonals
|
||||
if (upperBand >= 0 && (j - i) > upperBand)
|
||||
*(resetBuffer + xOffset + tadOffsetOut) = T(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// matrixBandPart_ - main algorithm caller
|
||||
//
|
||||
template <typename T>
|
||||
void matrixBandPart_(LaunchContext* context, NDArray* input, NDArray* output, LongType lowerBand, LongType upperBand) {
|
||||
dim3 launchDims = getLaunchDims("matrixBand");
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
std::vector<LongType> lastDims({input->rankOf() - 2, input->rankOf() - 1});
|
||||
std::vector<LongType> *dimsToExclude = ShapeUtils::evalDimsToExclude(input->rankOf(), lastDims.size(),lastDims.data());
|
||||
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), &lastDims);
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), &lastDims);
|
||||
|
||||
const LongType numTads = packX->numberOfTads();
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
matrixBandKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(), lowerBand,
|
||||
upperBand, packX->specialShapeInfo(), packX->specialOffsets(), packZ->specialShapeInfo(), packZ->specialOffsets(),
|
||||
numTads, input->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "matrixBandKernel failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
|
||||
delete dimsToExclude;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
void matrixBandPart(LaunchContext* context, NDArray* input, NDArray* output, LongType lowerBand, LongType upperBand) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), matrixBandPart_, (context, input, output, lowerBand, upperBand),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,124 @@
|
||||
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// Created by GS <sgazeos@gmail.com> on 3/21/2018.
|
||||
//
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/matrix_diag_part.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// put diagonals from input batched matrices to output batched vectors
|
||||
template <typename T>
|
||||
static SD_KERNEL void matrixDiagPartKernel(void* inputBuffer, void* outputBuffer, LongType numTads,
|
||||
LongType inputLength, LongType* tadOnlyInputShapeInfo,
|
||||
LongType* tadInputOffsets,
|
||||
LongType* tadOnlyOutputShapeInfo,
|
||||
LongType* tadOutputOffsets) {
|
||||
|
||||
if(blockIdx.x >= numTads)
|
||||
return;
|
||||
auto outputBuffer2 = reinterpret_cast<T*>(outputBuffer);
|
||||
auto inputBuffer2 = reinterpret_cast<T const*>(inputBuffer);
|
||||
|
||||
int totalThreads = blockDim.x;
|
||||
for (LongType i = blockIdx.x; i < numTads; i += gridDim.x) {
|
||||
auto yOffset = tadInputOffsets[i];
|
||||
auto xOffset = tadOutputOffsets[i];
|
||||
for (LongType j = threadIdx.x; j < inputLength; j += totalThreads) {
|
||||
LongType coords[2] = {j, j};
|
||||
LongType tadOffset, indexOffset;
|
||||
COORDS2INDEX(shape::rank(tadOnlyInputShapeInfo), shape::stride(tadOnlyInputShapeInfo), coords, tadOffset);
|
||||
COORDS2INDEX(shape::rank(tadOnlyOutputShapeInfo), shape::stride(tadOnlyOutputShapeInfo), coords, indexOffset);
|
||||
*(reinterpret_cast<T*>(outputBuffer) + xOffset + indexOffset) =
|
||||
*(reinterpret_cast<T const*>(inputBuffer) + yOffset + tadOffset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Returns a batched matrix tensor with new batched diagonal values.
|
||||
// for detailed explanations please take a look on web page:
|
||||
// https://www.tensorflow.org/api_docs/python/tf/matrix_set_diag
|
||||
//
|
||||
template <typename T>
|
||||
static Status _matrixDiagPart(LaunchContext* context, NDArray* input, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
auto listOut = output->allTensorsAlongDimension({output->rankOf() - 1});
|
||||
auto listDiag = input->allTensorsAlongDimension({input->rankOf() - 2, input->rankOf() - 1});
|
||||
|
||||
if (listOut.size() != listDiag.size()) {
|
||||
sd_printf("matrix_diag_part: Input matrix has wrong shape.", "");
|
||||
return Status::VALIDATION;
|
||||
}
|
||||
LongType lastDimension = math::sd_min(input->sizeAt(-2), input->sizeAt(-1));
|
||||
|
||||
LongType dims = output->rankOf() - 1;
|
||||
std::vector<LongType> *dimsToExclude = ShapeUtils::evalDimsToExclude(output->rankOf(), 1,&dims);
|
||||
const LongType numTads =
|
||||
ShapeUtils::getNumOfSubArrs(input->shapeInfo(),*dimsToExclude);
|
||||
std::vector<LongType> outputDims({output->rankOf() - 1});
|
||||
std::vector<LongType> inputDims({input->rankOf() - 2, input->rankOf() - 1});
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), &inputDims);
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), &outputDims);
|
||||
|
||||
if (!output->isActualOnDeviceSide()) input->syncToDevice();
|
||||
|
||||
if (!input->isActualOnDeviceSide()) input->syncToDevice();
|
||||
|
||||
dim3 launchDims = getLaunchDims("matrixDiag");
|
||||
matrixDiagPartKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(),
|
||||
output->specialBuffer(),numTads, lastDimension, const_cast<sd::LongType *>(packX->specialShapeInfo()),
|
||||
const_cast<sd::LongType *>(packX->specialOffsets()),
|
||||
const_cast<sd::LongType *>(packZ->specialShapeInfo()), const_cast<sd::LongType *>(packZ->specialOffsets()));
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "matrixDiagPartKernel failed");
|
||||
|
||||
delete dimsToExclude;
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// caller for _matrixDiagPart
|
||||
//
|
||||
Status matrixDiagPart(LaunchContext* context, NDArray* input, NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), return _matrixDiagPart, (context, input, output), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( sd::Status _matrixDiagPart,
|
||||
(sd::LaunchContext * context, NDArray* input, NDArray* output), SD_COMMON_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,119 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <ops/declarable/helpers/max_pooling.h>
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename Z>
|
||||
static SD_KERNEL void indicesFiller(void* vz, const LongType* zShapeInfo, LongType part, LongType bSize) {
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ const LongType *shape, *stride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(zShapeInfo);
|
||||
shape = shape::shapeOf(zShapeInfo);
|
||||
stride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType b = blockIdx.x; b < bSize; b += gridDim.x) {
|
||||
for (LongType e = threadIdx.x; e < part; e += blockDim.x) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
|
||||
// Compute coordinates and offset
|
||||
INDEX2COORDS(e + b * part, rank, shape, zCoords);
|
||||
COORDS2INDEX(rank, stride, zCoords, zOffset);
|
||||
|
||||
// Assign the index value
|
||||
z[zOffset] = static_cast<Z>(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Y>
|
||||
static void maxPoolingFunctor_(graph::Context& block, NDArray* input, NDArray* values,
|
||||
std::vector<LongType> const& params, NDArray* indices) {
|
||||
LongType kY = params[0];
|
||||
LongType kX = params[1];
|
||||
|
||||
LongType sY = params[2];
|
||||
LongType sX = params[3];
|
||||
|
||||
LongType pY = params[4];
|
||||
LongType pX = params[5];
|
||||
|
||||
LongType dY = params[6];
|
||||
LongType dX = params[7];
|
||||
|
||||
LongType oY = 0;
|
||||
LongType oX = 0;
|
||||
|
||||
const LongType bSize = input->sizeAt(0);
|
||||
const LongType inD = input->sizeAt(1);
|
||||
const LongType inY = input->sizeAt(2);
|
||||
const LongType inX = input->sizeAt(3);
|
||||
|
||||
const bool isSameMode = params[8] != 0;
|
||||
|
||||
ConvolutionUtils::calcOutSizePool2D(oY, oX, kY, kX, sY, sX, pY, pX, dY, dX, inY, inX, isSameMode);
|
||||
|
||||
if (isSameMode)
|
||||
ConvolutionUtils::calcPadding2D(pY, pX, oY, oX, inY, inX, params[0], params[1], params[2], params[3], params[6],
|
||||
params[7]);
|
||||
|
||||
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
|
||||
// poolingMode; 9 - divisor;
|
||||
ConvolutionUtils::pooling2d(block, *input, *values, kY, kX, sY, sX, pY, pX, dY, dX, MAX_POOL, 1);
|
||||
|
||||
if (nullptr != indices) {
|
||||
// for max_pool_with_argmax
|
||||
auto total = input->lengthOf();
|
||||
auto part = total / bSize;
|
||||
|
||||
indicesFiller<Y><<<256, 256, 1024, *block.launchContext()->getCudaStream()>>>(
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), part, bSize);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(block.launchContext()->getCudaStream(), "indicesFiller failed");
|
||||
}
|
||||
}
|
||||
|
||||
void maxPoolingFunctor(LaunchContext* context, graph::Context& block, NDArray* input, NDArray* values,
|
||||
std::vector<LongType> const& params, NDArray* indices) {
|
||||
NDArray::prepareSpecialUse({values, indices}, {input});
|
||||
auto yType = indices == nullptr ? INT64 : indices->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), yType, maxPoolingFunctor_, (block, input, values, params, indices),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({values, indices}, {input});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,105 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <array/NDArray.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
void maximumBPFunctor_(NDArray* x, NDArray* y, NDArray* epsNext, NDArray* gradX, NDArray* gradY) {
|
||||
auto lambdaX = LAMBDA_TTT(_e, _x, _y) { return _x >= _y ? _e : (T)0.; });
|
||||
|
||||
auto lambdaY = LAMBDA_TTT(_e, _x, _y) { return _x <= _y ? _e : (T)0.; });
|
||||
|
||||
if (x->isSameShape(y)) {
|
||||
// PWT case case
|
||||
|
||||
// X gradient
|
||||
epsNext->applyTriplewiseLambda(x, y, lambdaX, gradX);
|
||||
|
||||
// Y gradient
|
||||
epsNext->applyTriplewiseLambda(x, y, lambdaY, gradY);
|
||||
|
||||
} else if (y->isScalar()) {
|
||||
T s = y->e<T>(0);
|
||||
auto lambdaS = LAMBDA_TT(_e, _x, s) { return _x >= s ? _e : (T)0.; });
|
||||
|
||||
float zero = 0.f;
|
||||
// scalar case
|
||||
auto tmp = epsNext->reduceNumber(reduce::Sum);
|
||||
if (x <= y)
|
||||
gradY->assign(&tmp);
|
||||
else
|
||||
gradY->assign(zero);
|
||||
|
||||
epsNext->applyPairwiseLambda(x, lambdaS, gradX);
|
||||
} else {
|
||||
// broadcast case
|
||||
|
||||
// in this case we want to boost our X and Y shapes to the size of FF pass output (or epsNext, which has the same
|
||||
// shape)
|
||||
auto preX = x->dup();
|
||||
auto preY = y->dup();
|
||||
|
||||
auto* targetShapePtr = epsNext->getShapeAsVector();
|
||||
std::vector<LongType> targetShape = *targetShapePtr;
|
||||
delete targetShapePtr;
|
||||
|
||||
preX.tileToShape(targetShape, preX);
|
||||
preY.tileToShape(targetShape, preY);
|
||||
|
||||
epsNext->applyTriplewiseLambda(&preX, &preY, lambdaX, &preX);
|
||||
epsNext->applyTriplewiseLambda(&preX, &preY, lambdaY, &preY);
|
||||
|
||||
auto axisX = ShapeUtils::evalBroadcastBackwardAxis(x->shapeInfo(), epsNext->shapeInfo());
|
||||
auto axisY = ShapeUtils::evalBroadcastBackwardAxis(y->shapeInfo(), epsNext->shapeInfo());
|
||||
|
||||
if (axisX.size() > 0) {
|
||||
auto sum = preX.reduceAlongDimension(reduce::Sum, &axisX);
|
||||
gradX->assign(&sum);
|
||||
} else
|
||||
gradX->assign(&preX);
|
||||
|
||||
if (axisY.size() > 0) {
|
||||
auto sum = preY.reduceAlongDimension(reduce::Sum, &axisY);
|
||||
gradY->assign(&sum);
|
||||
} else
|
||||
gradY->assign(&preY);
|
||||
}
|
||||
}
|
||||
|
||||
void maximumBPFunctor(LaunchContext* context, NDArray* x, NDArray* y, NDArray* epsNext, NDArray* gradX,
|
||||
NDArray* gradY) {
|
||||
NDArray::prepareSpecialUse({gradX, gradY}, {x, y, epsNext});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(x->dataType(), maximumBPFunctor_, (x, y, epsNext, gradX, gradY), SD_NUMERIC_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({gradX, gradY}, {x, y, epsNext});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,735 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T, typename Z>
|
||||
static SD_KERNEL void mergeMaxIndexCudaLauncher(void** inArrs, void** inShapes, const int numArrays, void* voutput,
|
||||
const LongType* outputShape, LongType length) {
|
||||
auto output = reinterpret_cast<Z*>(voutput);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int rankOutput;
|
||||
__shared__ const LongType *shapeOutput, *strideOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankOutput = shape::rank(outputShape);
|
||||
shapeOutput = shape::shapeOf(outputShape);
|
||||
strideOutput = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
T mVal = -DataTypeUtils::max<T>();
|
||||
Z mIdx(0);
|
||||
|
||||
// Iterate through all input arrays to find the maximum value and its index
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto x = reinterpret_cast<const T*>(inArrs[i]);
|
||||
auto xShape = reinterpret_cast<const LongType*>(inShapes[i]);
|
||||
|
||||
__shared__ int rankInput;
|
||||
__shared__ const LongType *shapeInput, *strideInput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankInput = shape::rank(xShape);
|
||||
shapeInput = shape::shapeOf(xShape);
|
||||
strideInput = shape::stride(xShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(e, rankInput, shapeInput, xCoords);
|
||||
COORDS2INDEX(rankInput, strideInput, xCoords, xOffset);
|
||||
|
||||
// Update maximum value and index
|
||||
const auto val = x[xOffset];
|
||||
if (mVal < val) {
|
||||
mIdx = static_cast<Z>(i);
|
||||
mVal = val;
|
||||
}
|
||||
}
|
||||
|
||||
// Compute output coordinates and offset
|
||||
LongType outputOffset;
|
||||
INDEX2COORDS(e, rankOutput, shapeOutput, outputCoords);
|
||||
COORDS2INDEX(rankOutput, strideOutput, outputCoords, outputOffset);
|
||||
|
||||
// Store the index of the maximum value in the output
|
||||
output[outputOffset] = mIdx;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T, typename Z>
|
||||
static void mergeMaxIndex_(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
int nArrSize = static_cast<int>(inArrs.size());
|
||||
std::vector<const void*> inBuffers(nArrSize), inShapes(nArrSize);
|
||||
|
||||
for (int e = 0; e < nArrSize; e++) {
|
||||
inBuffers[e] = inArrs[e]->specialBuffer();
|
||||
inShapes[e] = inArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeMaxIndex");
|
||||
|
||||
auto pInBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(inBuffers.data(), inBuffers.size() * sizeof(void*)));
|
||||
auto pInShapes = reinterpret_cast<void**>(manager.replicatePointer(inShapes.data(), inShapes.size() * sizeof(void*)));
|
||||
auto length = output.lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
mergeMaxIndexCudaLauncher<T, Z><<<mergeLaunchDims.y, mergeLaunchDims.x, mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
pInBuffers, pInShapes, nArrSize, output.specialBuffer(), output.specialShapeInfo(), length);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeMaxIndexCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeMaxIndex(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
NDArray::prepareSpecialUse({&output}, inArrs);
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(inArrs[0]->dataType(), output.dataType(), mergeMaxIndex_, (context, inArrs, output),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, inArrs);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeMaxCudaLauncher(void** inArrs, void** inShapes, const int numArrays, void* voutput,
|
||||
const LongType* outputShape, LongType length) {
|
||||
auto output = reinterpret_cast<T*>(voutput);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int rankOutput;
|
||||
__shared__ const LongType *shapeOutput, *strideOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankOutput = shape::rank(outputShape);
|
||||
shapeOutput = shape::shapeOf(outputShape);
|
||||
strideOutput = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
T mVal = -DataTypeUtils::max<T>();
|
||||
|
||||
// Iterate through all input arrays to find the maximum value
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto x = reinterpret_cast<const T*>(inArrs[i]);
|
||||
auto xShape = reinterpret_cast<const LongType*>(inShapes[i]);
|
||||
|
||||
__shared__ int rankInput;
|
||||
__shared__ const LongType *shapeInput, *strideInput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankInput = shape::rank(xShape);
|
||||
shapeInput = shape::shapeOf(xShape);
|
||||
strideInput = shape::stride(xShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(e, rankInput, shapeInput, xCoords);
|
||||
COORDS2INDEX(rankInput, strideInput, xCoords, xOffset);
|
||||
|
||||
// Update maximum value
|
||||
const auto val = x[xOffset];
|
||||
if (mVal < val) {
|
||||
mVal = val;
|
||||
}
|
||||
}
|
||||
|
||||
// Compute output coordinates and offset
|
||||
LongType outputOffset;
|
||||
INDEX2COORDS(e, rankOutput, shapeOutput, outputCoords);
|
||||
COORDS2INDEX(rankOutput, strideOutput, outputCoords, outputOffset);
|
||||
|
||||
// Store the maximum value in the output
|
||||
output[outputOffset] = mVal;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void mergeMax_(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
int nArrsSize = static_cast<int>(inArrs.size());
|
||||
|
||||
std::vector<const void*> inBuffers(nArrsSize), inShapes(nArrsSize);
|
||||
|
||||
for (int e = 0; e < nArrsSize; e++) {
|
||||
inBuffers[e] = inArrs[e]->specialBuffer();
|
||||
inShapes[e] = inArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeMax");
|
||||
|
||||
auto pInBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(inBuffers.data(), inBuffers.size() * sizeof(void*)));
|
||||
auto pInShapes = reinterpret_cast<void**>(manager.replicatePointer(inShapes.data(), inShapes.size() * sizeof(void*)));
|
||||
auto length = output.lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
mergeMaxCudaLauncher<T><<<mergeLaunchDims.y, mergeLaunchDims.x, mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
pInBuffers, pInShapes, nArrsSize, output.specialBuffer(), output.specialShapeInfo(), length);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeMaxCudaLauncher failed");
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeMax(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
NDArray::prepareSpecialUse({&output}, inArrs);
|
||||
|
||||
BUILD_SINGLE_SELECTOR(output.dataType(), mergeMax_, (context, inArrs, output), SD_COMMON_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, inArrs);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeMaxBpCudaLauncher(void** inArrs, void** inShapes, const void* vgradient,
|
||||
const LongType* gradientShape, const int numArrays, void** outArrs,
|
||||
void** outShapes, LongType length, bool bSameOrderAndEws1) {
|
||||
const auto grad = reinterpret_cast<const T*>(vgradient);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int gradRank;
|
||||
__shared__ const LongType *gradShape, *gradStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
gradRank = shape::rank(gradientShape);
|
||||
gradShape = shape::shapeOf(gradientShape);
|
||||
gradStride = shape::stride(gradientShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
T mVal = -DataTypeUtils::max<T>();
|
||||
int nMaxIndex = 0;
|
||||
LongType gradOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
// Compute gradient offset if not same order and EWS=1
|
||||
if (!bSameOrderAndEws1) {
|
||||
INDEX2COORDS(e, gradRank, gradShape, coords);
|
||||
COORDS2INDEX(gradRank, gradStride, coords, gradOffset);
|
||||
}
|
||||
|
||||
// Find the maximum value and its index across all input arrays
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto x = reinterpret_cast<T*>(inArrs[i]);
|
||||
LongType xOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
if (!bSameOrderAndEws1) {
|
||||
auto xShape = reinterpret_cast<const LongType*>(inShapes[i]);
|
||||
COORDS2INDEX(shape::rank(xShape), shape::stride(xShape), coords, xOffset);
|
||||
}
|
||||
|
||||
const auto val = x[xOffset];
|
||||
if (mVal < val) {
|
||||
mVal = val;
|
||||
nMaxIndex = i;
|
||||
}
|
||||
}
|
||||
|
||||
// Assign gradient to the corresponding output array at the max index
|
||||
auto output = reinterpret_cast<T*>(outArrs[nMaxIndex]);
|
||||
LongType zOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
if (!bSameOrderAndEws1) {
|
||||
auto outShape = reinterpret_cast<const LongType*>(outShapes[nMaxIndex]);
|
||||
COORDS2INDEX(shape::rank(outShape), shape::stride(outShape), coords, zOffset);
|
||||
}
|
||||
|
||||
output[zOffset] = grad[gradOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void mergeMaxBp_(LaunchContext* context, const std::vector<NDArray*>& inArrs,
|
||||
std::vector<NDArray*>& outArrs, int nArrSize, bool bSameOrderAndEws1) {
|
||||
std::vector<const void*> inBuffers(nArrSize), inShapes(nArrSize), outBuffers(nArrSize), outShapes(nArrSize);
|
||||
|
||||
for (int e = 0; e < nArrSize; e++) {
|
||||
inBuffers[e] = inArrs[e]->specialBuffer();
|
||||
inShapes[e] = inArrs[e]->specialShapeInfo();
|
||||
outBuffers[e] = outArrs[e]->specialBuffer();
|
||||
outShapes[e] = outArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeMaxBp");
|
||||
|
||||
auto pInBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(inBuffers.data(), inBuffers.size() * sizeof(void*)));
|
||||
auto pInShapes = reinterpret_cast<void**>(manager.replicatePointer(inShapes.data(), inShapes.size() * sizeof(void*)));
|
||||
|
||||
auto pOutBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outBuffers.data(), outBuffers.size() * sizeof(void*)));
|
||||
auto pOutShapes =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outShapes.data(), outShapes.size() * sizeof(void*)));
|
||||
|
||||
auto length = inArrs[nArrSize]->lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
|
||||
mergeMaxBpCudaLauncher<T><<<mergeLaunchDims.y, mergeLaunchDims.x, mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
pInBuffers, pInShapes, inArrs[nArrSize]->specialBuffer(), inArrs[nArrSize]->specialShapeInfo(), nArrSize,
|
||||
pOutBuffers, pOutShapes, length, bSameOrderAndEws1);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeMaxBpCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeMaxBp(LaunchContext* context, const std::vector<NDArray*>& inArrs, std::vector<NDArray*>& outArrs) {
|
||||
// not use gradient
|
||||
int nArrSize = static_cast<int>(inArrs.size() - 1);
|
||||
|
||||
const std::vector<NDArray*>& out = reinterpret_cast<const std::vector<NDArray*>&>(outArrs);
|
||||
|
||||
NDArray::prepareSpecialUse(out, inArrs);
|
||||
|
||||
bool bSameOrderAndEws1 = false;
|
||||
auto ordering = inArrs[nArrSize]->ordering();
|
||||
|
||||
|
||||
BUILD_SINGLE_SELECTOR(inArrs[nArrSize]->dataType(), mergeMaxBp_,
|
||||
(context, inArrs, outArrs, nArrSize, bSameOrderAndEws1), SD_COMMON_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse(out, inArrs);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeAvgCudaLauncher(void** inArrs, void** inShapes, const int numArrays, void* voutput,
|
||||
const LongType* outputShape, LongType length) {
|
||||
auto output = reinterpret_cast<T*>(voutput);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int rankOutput;
|
||||
__shared__ const LongType *shapeOutput, *strideOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankOutput = shape::rank(outputShape);
|
||||
shapeOutput = shape::shapeOf(outputShape);
|
||||
strideOutput = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
T sum = static_cast<T>(0.0);
|
||||
|
||||
// Sum values from all input arrays
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto x = reinterpret_cast<T*>(inArrs[i]);
|
||||
auto xShape = reinterpret_cast<const LongType*>(inShapes[i]);
|
||||
|
||||
__shared__ int rankInput;
|
||||
__shared__ const LongType *shapeInput, *strideInput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankInput = shape::rank(xShape);
|
||||
shapeInput = shape::shapeOf(xShape);
|
||||
strideInput = shape::stride(xShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(e, rankInput, shapeInput, xCoords);
|
||||
COORDS2INDEX(rankInput, strideInput, xCoords, xOffset);
|
||||
|
||||
sum += x[xOffset];
|
||||
}
|
||||
|
||||
// Compute output coordinates and offset
|
||||
LongType outputOffset;
|
||||
INDEX2COORDS(e, rankOutput, shapeOutput, outputCoords);
|
||||
COORDS2INDEX(rankOutput, strideOutput, outputCoords, outputOffset);
|
||||
|
||||
// Store the averaged value in the output
|
||||
output[outputOffset] = sum / static_cast<T>(numArrays);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void mergeAvg_(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
std::vector<const void*> inBuffers(inArrs.size()), inShapes(inArrs.size());
|
||||
|
||||
for (int e = 0; e < inArrs.size(); e++) {
|
||||
inBuffers[e] = inArrs[e]->specialBuffer();
|
||||
inShapes[e] = inArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeAvg");
|
||||
|
||||
auto pInBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(inBuffers.data(), inBuffers.size() * sizeof(void*)));
|
||||
auto pInShapes = reinterpret_cast<void**>(manager.replicatePointer(inShapes.data(), inShapes.size() * sizeof(void*)));
|
||||
auto length = output.lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
|
||||
mergeAvgCudaLauncher<T><<<mergeLaunchDims.y, mergeLaunchDims.x, mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
pInBuffers, pInShapes, (int)inArrs.size(), output.specialBuffer(), output.specialShapeInfo(), length);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeAvgCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeAvg(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
NDArray::prepareSpecialUse({&output}, inArrs);
|
||||
|
||||
BUILD_SINGLE_SELECTOR(output.dataType(), mergeAvg_, (context, inArrs, output), SD_FLOAT_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, inArrs);
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeAvgBpCudaLauncher(const void* vgradient, const LongType* gradientShape, void** outArrs,
|
||||
void** outShapes, const int numArrays, LongType length,
|
||||
bool bSameOrderAndEws1) {
|
||||
const auto grad = reinterpret_cast<const T*>(vgradient);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int gradRank;
|
||||
__shared__ const LongType *gradShape, *gradStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
gradRank = shape::rank(gradientShape);
|
||||
gradShape = shape::shapeOf(gradientShape);
|
||||
gradStride = shape::stride(gradientShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
LongType gradOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
// Compute gradient offset if not using the same order and EWS=1
|
||||
if (!bSameOrderAndEws1) {
|
||||
INDEX2COORDS(e, gradRank, gradShape, coords);
|
||||
COORDS2INDEX(gradRank, gradStride, coords, gradOffset);
|
||||
}
|
||||
|
||||
// Iterate through each output array and compute the average gradient
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto output = reinterpret_cast<T*>(outArrs[i]);
|
||||
LongType zOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
if (!bSameOrderAndEws1) {
|
||||
auto outShape = reinterpret_cast<const LongType*>(outShapes[i]);
|
||||
COORDS2INDEX(shape::rank(outShape), shape::stride(outShape), coords, zOffset);
|
||||
}
|
||||
|
||||
// Assign averaged gradient value to output
|
||||
output[zOffset] = grad[gradOffset] / static_cast<T>(numArrays);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void mergeAvgBp_(LaunchContext* context, NDArray& gradient, std::vector<NDArray*>& outArrs,
|
||||
bool bSameOrderAndEws1) {
|
||||
int nArrSize = static_cast<int>(outArrs.size());
|
||||
|
||||
std::vector<const void*> outBuffers(nArrSize), outShapes(nArrSize);
|
||||
|
||||
for (int e = 0; e < nArrSize; e++) {
|
||||
outBuffers[e] = outArrs[e]->specialBuffer();
|
||||
outShapes[e] = outArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeAvgBp");
|
||||
|
||||
auto pOutBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outBuffers.data(), outBuffers.size() * sizeof(void*)));
|
||||
auto pOutShapes =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outShapes.data(), outShapes.size() * sizeof(void*)));
|
||||
|
||||
auto length = gradient.lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
|
||||
mergeAvgBpCudaLauncher<T><<<mergeLaunchDims.y, mergeLaunchDims.x,mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
gradient.specialBuffer(), gradient.specialShapeInfo(), pOutBuffers, pOutShapes, nArrSize, length,
|
||||
bSameOrderAndEws1);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeAvgBpCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeAvgBp(LaunchContext* context, NDArray& gradient, std::vector<NDArray*>& outArrs) {
|
||||
const std::vector<NDArray*>& out = reinterpret_cast<const std::vector<NDArray*>&>(outArrs);
|
||||
|
||||
NDArray::prepareSpecialUse(out, {&gradient});
|
||||
|
||||
bool bSameOrderAndEws1 = false;
|
||||
auto ordering = gradient.ordering();
|
||||
|
||||
for (const auto& v : outArrs) {
|
||||
bSameOrderAndEws1 &= (ordering == v->ordering());
|
||||
}
|
||||
|
||||
BUILD_SINGLE_SELECTOR(gradient.dataType(), mergeAvgBp_, (context, gradient, outArrs, bSameOrderAndEws1),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
NDArray::prepareSpecialUse(out, {&gradient});
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeAddCudaLauncher(void** inArrs, void** inShapes, const int numArrays, void* voutput,
|
||||
const LongType* outputShape, LongType length) {
|
||||
auto output = reinterpret_cast<T*>(voutput);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int rankOutput;
|
||||
__shared__ const LongType *shapeOutput, *strideOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankOutput = shape::rank(outputShape);
|
||||
shapeOutput = shape::shapeOf(outputShape);
|
||||
strideOutput = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
T sum(0.0f);
|
||||
|
||||
// Compute the sum across all input arrays
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto x = reinterpret_cast<T*>(inArrs[i]);
|
||||
auto xShape = reinterpret_cast<const LongType*>(inShapes[i]);
|
||||
|
||||
__shared__ int rankInput;
|
||||
__shared__ const LongType *shapeInput, *strideInput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rankInput = shape::rank(xShape);
|
||||
shapeInput = shape::shapeOf(xShape);
|
||||
strideInput = shape::stride(xShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(e, rankInput, shapeInput, xCoords);
|
||||
COORDS2INDEX(rankInput, strideInput, xCoords, xOffset);
|
||||
|
||||
sum += x[xOffset];
|
||||
}
|
||||
|
||||
// Compute output coordinates and offset
|
||||
LongType outputOffset;
|
||||
INDEX2COORDS(e, rankOutput, shapeOutput, outputCoords);
|
||||
COORDS2INDEX(rankOutput, strideOutput, outputCoords, outputOffset);
|
||||
|
||||
// Store the computed sum in the output
|
||||
output[outputOffset] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void mergeAdd_(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
int nArrSize = static_cast<int>(inArrs.size());
|
||||
std::vector<const void*> inBuffers(nArrSize), inShapes(nArrSize);
|
||||
|
||||
for (int e = 0; e < nArrSize; e++) {
|
||||
inBuffers[e] = inArrs[e]->specialBuffer();
|
||||
inShapes[e] = inArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeAdd");
|
||||
|
||||
auto pInBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(inBuffers.data(), inBuffers.size() * sizeof(void*)));
|
||||
auto pInShapes = reinterpret_cast<void**>(manager.replicatePointer(inShapes.data(), inShapes.size() * sizeof(void*)));
|
||||
auto length = output.lengthOf();
|
||||
|
||||
dim3 mergeLaunchDims = mergeDims(length);
|
||||
|
||||
mergeAddCudaLauncher<T><<<mergeLaunchDims.x, mergeLaunchDims.y, mergeLaunchDims.z, *context->getCudaStream()>>>(
|
||||
pInBuffers, pInShapes, nArrSize, output.specialBuffer(), output.specialShapeInfo(), length);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeAddCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void mergeAdd_,
|
||||
(sd::LaunchContext * context, const std::vector<NDArray*>& inArrs, NDArray& output),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
void mergeAdd(LaunchContext* context, const std::vector<NDArray*>& inArrs, NDArray& output) {
|
||||
NDArray::prepareSpecialUse({&output}, inArrs);
|
||||
|
||||
BUILD_SINGLE_SELECTOR(output.dataType(), mergeAdd_, (context, inArrs, output), SD_NUMERIC_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, inArrs);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeAddBpCudaLauncher(const void* vgradient, const LongType* gradientShape, void** outArrs,
|
||||
void** outShapes, const int numArrays, LongType length,
|
||||
bool bSameOrderAndEws1) {
|
||||
const auto grad = reinterpret_cast<const T*>(vgradient);
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ int gradRank;
|
||||
__shared__ const LongType *gradShape, *gradStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
gradRank = shape::rank(gradientShape);
|
||||
gradShape = shape::shapeOf(gradientShape);
|
||||
gradStride = shape::stride(gradientShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (LongType e = tid; e < length; e += step) {
|
||||
LongType gradOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
// Compute gradient offset if not using same order and EWS=1
|
||||
if (!bSameOrderAndEws1) {
|
||||
INDEX2COORDS(e, gradRank, gradShape, coords);
|
||||
COORDS2INDEX(gradRank, gradStride, coords, gradOffset);
|
||||
}
|
||||
|
||||
for (int i = 0; i < numArrays; ++i) {
|
||||
auto output = reinterpret_cast<T*>(outArrs[i]);
|
||||
LongType zOffset = bSameOrderAndEws1 ? e : 0;
|
||||
|
||||
// Compute output offset if not using same order and EWS=1
|
||||
if (!bSameOrderAndEws1) {
|
||||
auto outShape = reinterpret_cast<const LongType*>(outShapes[i]);
|
||||
COORDS2INDEX(shape::rank(outShape), shape::stride(outShape), coords, zOffset);
|
||||
}
|
||||
|
||||
// Assign gradient value to output
|
||||
output[zOffset] = grad[gradOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void mergeAddBp_(LaunchContext* context, NDArray& gradient, std::vector<NDArray*>& outArrs,
|
||||
bool bSameOrderAndEws1) {
|
||||
int nArrSize = static_cast<int>(outArrs.size());
|
||||
|
||||
std::vector<const void*> outBuffers(nArrSize), outShapes(nArrSize);
|
||||
|
||||
for (int e = 0; e < nArrSize; e++) {
|
||||
outBuffers[e] = outArrs[e]->specialBuffer();
|
||||
outShapes[e] = outArrs[e]->specialShapeInfo();
|
||||
}
|
||||
|
||||
PointersManager manager(context, "mergeAddBp");
|
||||
|
||||
auto pOutBuffers =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outBuffers.data(), outBuffers.size() * sizeof(void*)));
|
||||
auto pOutShapes =
|
||||
reinterpret_cast<void**>(manager.replicatePointer(outShapes.data(), outShapes.size() * sizeof(void*)));
|
||||
|
||||
auto length = gradient.lengthOf();
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (length + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
mergeAddBpCudaLauncher<T><<<blocksPerGrid, threadsPerBlock, 512, *context->getCudaStream()>>>(
|
||||
gradient.specialBuffer(), gradient.specialShapeInfo(), pOutBuffers, pOutShapes, nArrSize, length,
|
||||
bSameOrderAndEws1);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeAddBpCudaLauncher failed");
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void mergeAddBp(LaunchContext* context, NDArray& gradient, std::vector<NDArray*>& outArrs) {
|
||||
const std::vector<NDArray*>& out = reinterpret_cast<const std::vector<NDArray*>&>(outArrs);
|
||||
NDArray::prepareSpecialUse(out, {&gradient});
|
||||
|
||||
bool bSameOrderAndEws1 = false;
|
||||
auto ordering = gradient.ordering();
|
||||
|
||||
for (const auto& v : outArrs) {
|
||||
bSameOrderAndEws1 &= (ordering == v->ordering());
|
||||
}
|
||||
|
||||
BUILD_SINGLE_SELECTOR(gradient.dataType(), mergeAddBp_, (context, gradient, outArrs, bSameOrderAndEws1),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
NDArray::prepareSpecialUse(out, {&gradient});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,160 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
|
||||
#include <array/ResultSet.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/meshgrid.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_DEVICE void assign_(void *vx, LongType *xShapeInfo, void *vz, LongType *zShapeInfo) {
|
||||
auto x = reinterpret_cast<T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const auto step = blockDim.x * gridDim.x;
|
||||
|
||||
__shared__ LongType length, rankX, rankZ;
|
||||
__shared__ const LongType *shapeX, *strideX, *shapeZ, *strideZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
rankX = shape::rank(xShapeInfo);
|
||||
rankZ = shape::rank(zShapeInfo);
|
||||
shapeX = shape::shapeOf(xShapeInfo);
|
||||
strideX = shape::stride(xShapeInfo);
|
||||
shapeZ = shape::shapeOf(zShapeInfo);
|
||||
strideZ = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = tid; i < length; i += step) {
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(i, rankX, shapeX, xCoords);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rankX, strideX, xCoords, xOffset);
|
||||
|
||||
// Compute output coordinates and offset
|
||||
INDEX2COORDS(i, rankZ, shapeZ, zCoords);
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rankZ, strideZ, zCoords, zOffset);
|
||||
|
||||
// Assign value from input to output
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void meshgridKernel(int rank, void **outBuffers, LongType **tadShapes, LongType **tadOffsets,
|
||||
LongType *numTads, void **inBuffers, LongType **inShapes) {
|
||||
// for all arrays
|
||||
for (int i = blockIdx.x; i < rank; i += gridDim.x) {
|
||||
// for all tads in this array
|
||||
for (LongType j = 0; j < numTads[i]; j++) {
|
||||
assign_<T>(inBuffers[i], inShapes[i], reinterpret_cast<T *>(outBuffers[i]) + tadOffsets[i][j], tadShapes[i]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void meshgrid_(LaunchContext *context, const std::vector<NDArray *> &inArrs,
|
||||
const std::vector<NDArray *> &outArrs, const bool swapFirst2Dims) {
|
||||
const int rank = inArrs.size();
|
||||
int inIndices[SD_MAX_RANK];
|
||||
std::iota(inIndices, inIndices + rank, 0);
|
||||
if (swapFirst2Dims && rank > 1) {
|
||||
inIndices[0] = 1;
|
||||
inIndices[1] = 0;
|
||||
}
|
||||
|
||||
PointersManager pm(context, "meshgrid");
|
||||
std::vector<const void *> hInBuffers(rank);
|
||||
std::vector<void *> hOutBuffers(rank);
|
||||
std::vector<const LongType *> hInShapes(rank);
|
||||
|
||||
std::vector<const LongType *> hOutTadShapes(rank);
|
||||
std::vector<const LongType *> hOutTadOffsets(rank);
|
||||
|
||||
std::vector<LongType> hNumTads(rank);
|
||||
|
||||
for (int i = 0; i < rank; ++i) {
|
||||
hInBuffers[i] = inArrs[i]->specialBuffer();
|
||||
hInShapes[i] = inArrs[i]->specialShapeInfo();
|
||||
|
||||
hOutBuffers[i] = outArrs[i]->specialBuffer();
|
||||
|
||||
auto pack = ConstantTadHelper::getInstance().tadForDimensions(outArrs[i]->shapeInfo(), {inIndices[i]});
|
||||
hOutTadShapes[i] = pack->specialShapeInfo();
|
||||
hOutTadOffsets[i] = pack->specialOffsets();
|
||||
hNumTads[i] = pack->numberOfTads();
|
||||
|
||||
}
|
||||
|
||||
auto dInBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(hInBuffers.data(), hInBuffers.size() * sizeof(void *)));
|
||||
auto dOutBuffers =
|
||||
reinterpret_cast<void **>(pm.replicatePointer(hOutBuffers.data(), hOutBuffers.size() * sizeof(void *)));
|
||||
|
||||
auto dInShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(hInShapes.data(), hInShapes.size() * sizeof(LongType *)));
|
||||
auto dOutTadShapes = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(hOutTadShapes.data(), hOutTadShapes.size() * sizeof(LongType *)));
|
||||
auto dOutTadOffsets = reinterpret_cast<LongType **>(
|
||||
pm.replicatePointer(hOutTadOffsets.data(), hOutTadOffsets.size() * sizeof(LongType *)));
|
||||
|
||||
auto dNumTads =
|
||||
reinterpret_cast<LongType *>(pm.replicatePointer(hNumTads.data(), hNumTads.size() * sizeof(LongType)));
|
||||
|
||||
dim3 launchDims = getLaunchDims("meshgrid");
|
||||
meshgridKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(rank, dOutBuffers, dOutTadShapes, dOutTadOffsets,
|
||||
dNumTads, dInBuffers, dInShapes);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "meshgridKernel failed");
|
||||
|
||||
pm.synchronize();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void meshgrid(LaunchContext *context, const std::vector<NDArray *> &inArrs, const std::vector<NDArray *> &outArrs,
|
||||
const bool swapFirst2Dims) {
|
||||
BUILD_SINGLE_SELECTOR(inArrs.at(0)->dataType(), meshgrid_, (context, inArrs, outArrs, swapFirst2Dims),
|
||||
SD_NUMERIC_TYPES);
|
||||
|
||||
for (auto v : outArrs) v->tickWriteDevice();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,108 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#ifndef __MIN_I_MAX_H_HELPERS__
|
||||
#define __MIN_I_MAX_H_HELPERS__
|
||||
#include <array/NDArray.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
void minimumBPFunctor_(NDArray* x, NDArray* y, NDArray* epsNext, NDArray* gradX, NDArray* gradY) {
|
||||
auto lambdaX = LAMBDA_TTT(_e, _x, _y) { return _x <= _y ? _e : (T)0.; });
|
||||
|
||||
auto lambdaY = LAMBDA_TTT(_e, _x, _y) { return _x >= _y ? _e : (T)0.; });
|
||||
|
||||
if (x->isSameShape(y)) {
|
||||
// PWT case case
|
||||
|
||||
// X gradient
|
||||
epsNext->applyTriplewiseLambda(x, y, lambdaX, gradX);
|
||||
|
||||
// Y gradient
|
||||
epsNext->applyTriplewiseLambda(x, y, lambdaY, gradY);
|
||||
|
||||
} else if (y->isScalar()) {
|
||||
T s = y->e<T>(0);
|
||||
auto lambdaS = LAMBDA_TT(_e, _x, s) { return _x <= s ? _e : (T)0.; });
|
||||
|
||||
float zero = 0.f;
|
||||
// scalar case
|
||||
auto tmp = epsNext->reduceNumber(reduce::Sum);
|
||||
if (x <= y)
|
||||
gradY->assign(&tmp);
|
||||
else
|
||||
gradY->assign(zero);
|
||||
|
||||
epsNext->applyPairwiseLambda(x, lambdaS, gradX);
|
||||
} else {
|
||||
// broadcast case
|
||||
|
||||
// in this case we want to boost our X and Y shapes to the size of FF pass output (or epsNext, which has the same
|
||||
// shape)
|
||||
auto preX = x->dup();
|
||||
auto preY = y->dup();
|
||||
|
||||
auto* targetShapePtr = epsNext->getShapeAsVector();
|
||||
std::vector<LongType> targetShape = *targetShapePtr;
|
||||
delete targetShapePtr;
|
||||
|
||||
preX.tileToShape(targetShape, preX);
|
||||
preY.tileToShape(targetShape, preY);
|
||||
|
||||
epsNext->applyTriplewiseLambda(&preX, &preY, lambdaX, &preX);
|
||||
epsNext->applyTriplewiseLambda(&preX, &preY, lambdaY, &preY);
|
||||
|
||||
auto axisX = ShapeUtils::evalBroadcastBackwardAxis(x->shapeInfo(), epsNext->shapeInfo());
|
||||
auto axisY = ShapeUtils::evalBroadcastBackwardAxis(y->shapeInfo(), epsNext->shapeInfo());
|
||||
|
||||
if (axisX.size() > 0) {
|
||||
auto sum = preX.reduceAlongDimension(reduce::Sum, &axisX);
|
||||
gradX->assign(&sum);
|
||||
} else
|
||||
gradX->assign(&preX);
|
||||
|
||||
if (axisY.size() > 0) {
|
||||
auto sum = preY.reduceAlongDimension(reduce::Sum, &axisY);
|
||||
gradY->assign(&sum);
|
||||
} else
|
||||
gradY->assign(&preY);
|
||||
}
|
||||
}
|
||||
|
||||
void minimumBPFunctor(LaunchContext* context, NDArray* x, NDArray* y, NDArray* epsNext, NDArray* gradX,
|
||||
NDArray* gradY) {
|
||||
NDArray::prepareSpecialUse({gradX, gradY}, {x, y, epsNext});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(x->dataType(), minimumBPFunctor_, (x, y, epsNext, gradX, gradY), SD_NUMERIC_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({gradX, gradY}, {x, y, epsNext});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
#endif
|
||||
@@ -0,0 +1,129 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <legacy/NativeOps.h>
|
||||
#include <ops/declarable/helpers/nth_element.h>
|
||||
|
||||
#include "array/NDArrayFactory.h"
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void fillUpElementKernel(void* outputBuffer, const LongType* outputShapeInfo, void* inputBuffer,
|
||||
const LongType* inputShapeInfo, const LongType* pTadShape,
|
||||
const LongType* pTadOffsets, LongType n) {
|
||||
__shared__ LongType bufferLength;
|
||||
__shared__ int rankOutput, rankTad;
|
||||
__shared__ const LongType *shapeOutput, *strideOutput, *shapeTad, *strideTad;
|
||||
|
||||
auto z = reinterpret_cast<T*>(outputBuffer);
|
||||
auto x = reinterpret_cast<T*>(inputBuffer);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
bufferLength = shape::length(outputShapeInfo);
|
||||
rankOutput = shape::rank(outputShapeInfo);
|
||||
rankTad = shape::rank(pTadShape);
|
||||
shapeOutput = shape::shapeOf(outputShapeInfo);
|
||||
strideOutput = shape::stride(outputShapeInfo);
|
||||
shapeTad = shape::shapeOf(pTadShape);
|
||||
strideTad = shape::stride(pTadShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType t = tid; t < bufferLength; t += step) {
|
||||
// Compute output coordinates and offset
|
||||
INDEX2COORDS(t, rankOutput, shapeOutput, zCoords);
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rankOutput, strideOutput, zCoords, zOffset);
|
||||
|
||||
// Compute input coordinates and offset
|
||||
INDEX2COORDS(n, rankTad, shapeTad, xCoords);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rankTad, strideTad, xCoords, xOffset);
|
||||
|
||||
// Access and assign the value
|
||||
z[zOffset] = x[pTadOffsets[t] + xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void nthElementFunctor_(LaunchContext* context, NDArray* input, LongType n, NDArray* output, bool reverse) {
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
NDArray sortedVals(*input);
|
||||
Pointer params[2];
|
||||
params[0] = context;
|
||||
params[1] = context->getCudaStream();
|
||||
// Nth element in sorted sequence : basic algorithm sort and retrieve nth element in sorted
|
||||
if (input->isVector()) {
|
||||
sort(params, &sortedVals, reverse);
|
||||
|
||||
cudaMemcpy(reinterpret_cast<T*>(output->specialBuffer()), reinterpret_cast<T*>(sortedVals.specialBuffer()) + n,
|
||||
sizeof(T), cudaMemcpyDeviceToDevice);
|
||||
} else { // rank greater than 1
|
||||
std::vector<LongType> lastDims(
|
||||
{input->rankOf() - 1});
|
||||
NDArray *dimData = NDArrayFactory::create_<LongType>('c',{2},lastDims, context);
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(sortedVals.shapeInfo(), &lastDims);
|
||||
|
||||
auto pTadShape = packX->specialShapeInfo();
|
||||
auto pTadShapeH = packX->primaryShapeInfo();
|
||||
auto pTadOffsets = packX->specialOffsets();
|
||||
sortTad(params, &sortedVals,
|
||||
reinterpret_cast<sd::LongType *>(lastDims.data()),
|
||||
lastDims.size(),
|
||||
const_cast<sd::LongType *>(pTadShape),
|
||||
const_cast<sd::LongType *>(pTadOffsets),
|
||||
reverse);
|
||||
sortedVals.tickWriteDevice();
|
||||
sortedVals.syncToHost();
|
||||
auto stream = context->getCudaStream();
|
||||
dim3 launchDims = getLaunchDims("nth_element_fill");
|
||||
fillUpElementKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(output->specialBuffer(), output->specialShapeInfo(),
|
||||
sortedVals.specialBuffer(), sortedVals.specialShapeInfo(),
|
||||
pTadShape, pTadOffsets, n);
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillUpElementKernel failed");
|
||||
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
void nthElementFunctor(LaunchContext* context, NDArray* input, LongType n, NDArray* output, bool reverse) {
|
||||
auto inputDType = input->dataType();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), nthElementFunctor_, (context, input, n, output, reverse), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,128 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 30.05.2019
|
||||
//
|
||||
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/one_hot.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, z - output
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL static void onehotCuda(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const LongType axis, const LongType depth,
|
||||
const Z on, const Z off) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
|
||||
__shared__ int xRank, zRank;
|
||||
__shared__ LongType zLen, totalThreads;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape, *zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
LongType coord[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = tid; i < zLen; i += totalThreads) {
|
||||
// Compute output coordinate and offset
|
||||
INDEX2COORDS(i, zRank, zShape, coord);
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(zRank, zStride, coord, zOffset);
|
||||
|
||||
// Extract depth coordinate and shift axis
|
||||
const auto depthCoord = coord[axis];
|
||||
for (LongType j = axis; j < zRank - 1; ++j) {
|
||||
coord[j] = coord[j + 1];
|
||||
}
|
||||
|
||||
// Compute input offset
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, coord, xOffset);
|
||||
|
||||
// Check if the depth matches the index
|
||||
const LongType idx = static_cast<LongType>(x[xOffset]);
|
||||
z[zOffset] = (depthCoord == idx) ? on : off;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void onehotCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const LongType axis, const LongType depth,
|
||||
const double on, const double off) {
|
||||
onehotCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, axis, depth,
|
||||
static_cast<Y>(on), static_cast<Y>(off));
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "onehotCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void onehot(const LaunchContext *context, NDArray *indices, NDArray *output, const LongType axis,
|
||||
const LongType depth, const double on, const double off) {
|
||||
const auto xType = indices->dataType();
|
||||
const auto zType = output->dataType();
|
||||
|
||||
dim3 oneHotLaunch = oneHotDims(output->lengthOf(),output->rankOf(), sizeof(decltype(*output->shapeInfo())));
|
||||
PointersManager manager(context, "onehot");
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {indices});
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
xType, zType, onehotCudaLauncher,
|
||||
(oneHotLaunch.y, oneHotLaunch.x, oneHotLaunch.z, context->getCudaStream(), indices->specialBuffer(),
|
||||
indices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(), axis, depth, on, off),
|
||||
SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {indices});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,322 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - input, y - paddings, z - output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void padCuda(const int mode, const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* 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__ LongType zLen, totalThreads;
|
||||
__shared__ const LongType *xShape, *zShape, *xStride, *zStride;
|
||||
__shared__ LongType yStride0, shift1, shift2;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
rankMinusOne = rank - 1;
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
yStride0 = shape::stride(yShapeInfo)[0];
|
||||
zLen = shape::length(zShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
shift1 = (mode == 1) ? 0 : 1; // REFLECT : SYMMETRIC
|
||||
shift2 = (mode == 1) ? 2 : 1; // REFLECT : SYMMETRIC
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = totalThreads;
|
||||
|
||||
LongType xzCoord[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = start; i < zLen; i += step) {
|
||||
// Compute output coordinate and offset
|
||||
INDEX2COORDS(i, rank, zShape, xzCoord);
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, xzCoord, zOffset);
|
||||
|
||||
bool within = true;
|
||||
|
||||
for (int j = rankMinusOne; j >= 0; --j) {
|
||||
if (xShape[j] == zShape[j]) continue;
|
||||
|
||||
LongType leftOffset;
|
||||
LongType leftCoords[] = {yStride0 * j};
|
||||
COORDS2INDEX(1, shape::stride(yShapeInfo), leftCoords, leftOffset);
|
||||
const auto left = y[leftOffset];
|
||||
|
||||
if (xzCoord[j] < left || xzCoord[j] >= left + xShape[j]) {
|
||||
within = false;
|
||||
|
||||
if (mode != 0) { // REFLECT or SYMMETRIC
|
||||
xzCoord[j] = xzCoord[j] - left;
|
||||
|
||||
if (xzCoord[j] < 0) { // Left boundary
|
||||
xzCoord[j] = -xzCoord[j] - shift1;
|
||||
} else if (xzCoord[j] >= xShape[j]) { // Right boundary
|
||||
xzCoord[j] = 2 * xShape[j] - xzCoord[j] - shift2;
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
} else {
|
||||
xzCoord[j] -= left;
|
||||
}
|
||||
}
|
||||
|
||||
if (within || mode != 0) {
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, xzCoord, xOffset);
|
||||
z[zOffset] = within ? x[xOffset] : x[xOffset]; // Handles REFLECT or SYMMETRIC
|
||||
} else {
|
||||
z[zOffset] = padVal; // CONSTANT padding
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
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 LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const void* padVal) {
|
||||
padCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(mode, vx, xShapeInfo, vy, yShapeInfo, vz,
|
||||
zShapeInfo, padVal);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "padCuda failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void pad(LaunchContext* context, const int mode, NDArray& input, NDArray& paddings, NDArray& output,
|
||||
NDArray& padValue) {
|
||||
PointersManager manager(context, "pad");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input, &paddings, &padValue});
|
||||
|
||||
dim3 padLaunch = padDims(output.lengthOf(),output.rankOf());
|
||||
const auto xType = input.dataType();
|
||||
const auto yType = paddings.dataType();
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
xType, yType, padCudaLauncher,
|
||||
(padLaunch.y, padLaunch.x, padLaunch.z, context->getCudaStream(), mode, input.specialBuffer(),
|
||||
input.specialShapeInfo(), paddings.specialBuffer(), paddings.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), padValue.specialBuffer()),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input, &paddings, &padValue});
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void mirrorPadLinearKernel(void const* vx, const LongType* xShape, void* vz,
|
||||
const LongType* zShape,
|
||||
LongType leftSide, LongType leftSideCorrected, LongType xLen, LongType len,
|
||||
LongType zLen) {
|
||||
__shared__ T const* x;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType rankX, rankZ;
|
||||
__shared__ const LongType* shapeX;
|
||||
__shared__ const LongType* strideX;
|
||||
__shared__ const LongType* shapeZ;
|
||||
__shared__ const LongType* strideZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
x = reinterpret_cast<T const*>(vx);
|
||||
z = reinterpret_cast<T*>(vz);
|
||||
|
||||
rankX = shape::rank(xShape);
|
||||
rankZ = shape::rank(zShape);
|
||||
shapeX = shape::shapeOf(xShape);
|
||||
strideX = shape::stride(xShape);
|
||||
shapeZ = shape::shapeOf(zShape);
|
||||
strideZ = shape::stride(zShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = blockDim.x * gridDim.x;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xOffset, zOffset;
|
||||
|
||||
for (LongType i = start; i < zLen; i += step) {
|
||||
// Compute coordinates and offset for the output
|
||||
INDEX2COORDS(i, rankZ, shapeZ, zCoords);
|
||||
COORDS2INDEX(rankZ, strideZ, zCoords, zOffset);
|
||||
|
||||
// Adjust input offset based on the mirror padding logic
|
||||
if (i < leftSide) { // Left side
|
||||
const LongType mirrorIndex = leftSideCorrected - i;
|
||||
COORDS2INDEX(rankX, strideX, &mirrorIndex, xOffset);
|
||||
} else if (i < leftSide + xLen) { // Middle section
|
||||
const LongType middleIndex = i - leftSide;
|
||||
COORDS2INDEX(rankX, strideX, &middleIndex, xOffset);
|
||||
} else { // Right side
|
||||
const LongType mirrorIndex = len - i;
|
||||
COORDS2INDEX(rankX, strideX, &mirrorIndex, xOffset);
|
||||
}
|
||||
|
||||
// Assign value from input to output
|
||||
if (zOffset < zLen && xOffset < xLen) {
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename F, typename I>
|
||||
static SD_KERNEL void mirrorPadKernel(void const* vx, const LongType* xShape, void* vz, const LongType* zShape,
|
||||
LongType outLen, void const* paddings, const LongType* paddingShape,
|
||||
int reflBorder) {
|
||||
__shared__ F const* x;
|
||||
__shared__ I const* pads;
|
||||
__shared__ F* z;
|
||||
__shared__ LongType rank;
|
||||
__shared__ sd::LongType *zStride;
|
||||
__shared__ sd::LongType *xStride;
|
||||
__shared__ LongType* zShapeArr;
|
||||
__shared__ LongType* xShapeArr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShape);
|
||||
zShapeArr = shape::shapeOf(zShape);
|
||||
zStride = shape::stride(zShape);
|
||||
xShapeArr = shape::shapeOf(xShape);
|
||||
xStride = shape::stride(xShape);
|
||||
|
||||
x = reinterpret_cast<F const*>(vx);
|
||||
pads = reinterpret_cast<I const*>(paddings);
|
||||
z = reinterpret_cast<F*>(vz);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto start = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const auto step = blockDim.x * gridDim.x;
|
||||
|
||||
LongType xzCoord[SD_MAX_RANK];
|
||||
LongType coords[2];
|
||||
|
||||
for (LongType i = start; i < outLen; i += step) {
|
||||
// Calculate output coordinate and offset
|
||||
INDEX2COORDS(i, rank, zShapeArr, xzCoord);
|
||||
LongType outOffset;
|
||||
COORDS2INDEX(rank, zStride, xzCoord, outOffset);
|
||||
|
||||
// Adjust input coordinates based on mirror padding
|
||||
for (LongType j = 0; j < rank; ++j) {
|
||||
const auto inLen = shape::sizeAt(xShape, j);
|
||||
|
||||
coords[0] = j;
|
||||
coords[1] = 0;
|
||||
|
||||
LongType padOffset;
|
||||
COORDS2INDEX(2, shape::stride(paddingShape), coords, padOffset);
|
||||
const auto leftSide = pads[padOffset];
|
||||
const auto leftSideCorrected = leftSide - reflBorder;
|
||||
const auto len = 2 * (inLen - 1) + leftSide + reflBorder;
|
||||
|
||||
if (xzCoord[j] < leftSide) { // Left side
|
||||
xzCoord[j] = leftSideCorrected - xzCoord[j];
|
||||
} else if (xzCoord[j] < leftSide + inLen) { // Middle
|
||||
xzCoord[j] = xzCoord[j] - leftSide;
|
||||
} else if (xzCoord[j] < len) { // Right side
|
||||
xzCoord[j] = len - xzCoord[j];
|
||||
} else { // Beyond the mirrored region
|
||||
xzCoord[j] = xzCoord[j] - len;
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate input offset and assign value
|
||||
LongType inOffset;
|
||||
COORDS2INDEX(rank, xStride, xzCoord, inOffset);
|
||||
z[outOffset] = x[inOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename F, typename I>
|
||||
static void mirrorPad_(LaunchContext* context, NDArray& input, NDArray& paddings, NDArray& output,
|
||||
const int mode) {
|
||||
// mode: 0 - REFLECT, else - SYMMETRIC
|
||||
const int reflBorder = (bool)mode ? 1 : 0;
|
||||
const LongType rank = input.rankOf();
|
||||
const LongType outLen = output.lengthOf();
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({&output}, {&input, &paddings});
|
||||
|
||||
if (rank <= 1) {
|
||||
const LongType inLen = input.isScalar() ? 1 : input.lengthOf();
|
||||
const auto leftSide = paddings.e<LongType>(0);
|
||||
const auto leftSideCorrected = leftSide - reflBorder;
|
||||
const LongType len = 2 * (inLen - 1) + leftSide + reflBorder;
|
||||
|
||||
dim3 mirrorPadLinearDims2 = mirrorPadLinearDims(len);
|
||||
mirrorPadLinearKernel<F><<<mirrorPadLinearDims2.y, mirrorPadLinearDims2.x, mirrorPadLinearDims2.z, *stream>>>(
|
||||
input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), leftSide,
|
||||
leftSideCorrected, inLen, len, outLen);
|
||||
DebugHelper::checkErrorCode(stream, "helpers::mirrorPadLinearKernel(...) failed");
|
||||
} else {
|
||||
dim3 mirrorPadDims = mirrorPadTad(output.lengthOf(),input.rankOf());
|
||||
mirrorPadKernel<F, I><<<mirrorPadDims.y, mirrorPadDims.x, mirrorPadDims.z, *stream>>>(
|
||||
input.specialBuffer(), input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), outLen,
|
||||
paddings.specialBuffer(), paddings.specialShapeInfo(), reflBorder);
|
||||
DebugHelper::checkErrorCode(stream, "helpers::mirrorPadKernel(...) failed");
|
||||
}
|
||||
NDArray::registerSpecialUse({&output}, {&input, &paddings});
|
||||
}
|
||||
|
||||
void mirrorPad(LaunchContext* context, NDArray& input, NDArray& paddings, NDArray& output,
|
||||
const int mode) {
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), paddings.dataType(), mirrorPad_, (context, input, paddings, output, mode),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,148 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 17.05.2018
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <ops/declarable/helpers/percentile.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename X>
|
||||
static SD_KERNEL void percentileKernel(void* vx, const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
||||
const LongType numTads, const LongType tadLength, void* vz,
|
||||
const LongType* zShapeInfo, const LongType zLength,
|
||||
const LongType position) {
|
||||
const auto x = reinterpret_cast<X*>(vx);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
|
||||
__shared__ LongType xRank, zRank;
|
||||
__shared__ const LongType* xShape;
|
||||
__shared__ const LongType* xStride;
|
||||
__shared__ const LongType* zShape;
|
||||
__shared__ const LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xTadShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
xShape = shape::shapeOf(xTadShapeInfo);
|
||||
xStride = shape::stride(xTadShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType t = blockIdx.x; t < numTads; t += gridDim.x) {
|
||||
auto tad = x + xTadOffsets[t];
|
||||
|
||||
// Sort TAD using odd-even transposition sort
|
||||
for (LongType m = 0; m < tadLength; ++m) {
|
||||
for (LongType tid = threadIdx.x; tid < tadLength; tid += blockDim.x) {
|
||||
const auto top = (m % 2 == 0) ? 2 * tid + 1 : 2 * tid + 2;
|
||||
if (top < tadLength) {
|
||||
if (tad[top - 1] > tad[top]) {
|
||||
// Swap values
|
||||
X temp = tad[top - 1];
|
||||
tad[top - 1] = tad[top];
|
||||
tad[top] = temp;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Save the final value to the output
|
||||
if (threadIdx.x == 0) {
|
||||
const auto value = tad[position];
|
||||
LongType zOffset;
|
||||
|
||||
COORDS2INDEX(zRank, zStride, &t, zOffset);
|
||||
z[zOffset] = value;
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void _percentile(LaunchContext* context, NDArray& input, NDArray& output, std::vector<LongType>& axis,
|
||||
const float q, const int interpolation) {
|
||||
const int inputRank = input.rankOf();
|
||||
|
||||
if (axis.empty())
|
||||
for (int i = 0; i < inputRank; ++i) axis.push_back(i);
|
||||
else
|
||||
shape::checkDimensions(inputRank, &axis);
|
||||
|
||||
auto tempArray = input.dup();
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(tempArray.shapeInfo(), &axis);
|
||||
|
||||
auto tadLength = shape::length(packX->primaryShapeInfo());
|
||||
|
||||
const float fraction = 1.f - q / 100.;
|
||||
LongType position = 0;
|
||||
|
||||
switch (interpolation) {
|
||||
case 0: // lower
|
||||
position = static_cast<LongType>(math::sd_ceil<float, T>((tadLength - 1) * fraction));
|
||||
break;
|
||||
case 1: // higher
|
||||
position = static_cast<LongType>(math::sd_floor<float, T>((tadLength - 1) * fraction));
|
||||
break;
|
||||
case 2: // nearest
|
||||
position = static_cast<LongType>(math::sd_round<float, T>((tadLength - 1) * fraction));
|
||||
break;
|
||||
}
|
||||
position = tadLength - position - 1;
|
||||
|
||||
dim3 launchDims = getLaunchDims("percentile");
|
||||
percentileKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
||||
tempArray.specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), packX->numberOfTads(), tadLength,
|
||||
output.specialBuffer(), output.specialShapeInfo(), output.lengthOf(), position);
|
||||
|
||||
DebugHelper::checkErrorCode(context->getCudaStream(), "percentile");
|
||||
}
|
||||
|
||||
void percentile(LaunchContext* context, NDArray& input, NDArray& output, std::vector<LongType>& axises,
|
||||
const float q, const int interpolation) {
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), _percentile, (context, input, output, axises, q, interpolation),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void _percentile,
|
||||
(sd::LaunchContext * context, NDArray& input, NDArray& output, std::vector<sd::LongType>& axises,
|
||||
const float q, const int interpolation),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,125 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 26.04.2019
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <ops/declarable/helpers/gammaMathFunc.h>
|
||||
#include <ops/declarable/helpers/zeta.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void polyGammaCuda(const void *vn, const LongType *nShapeInfo, const void *vx,
|
||||
const LongType *xShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto n = reinterpret_cast<const T *>(vn);
|
||||
const auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ LongType len;
|
||||
__shared__ bool sameOffsetNX, sameOffsetNZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(nShapeInfo);
|
||||
sameOffsetNX = shape::haveSameShapeAndStrides(xShapeInfo, nShapeInfo);
|
||||
sameOffsetNZ = shape::haveSameShapeAndStrides(zShapeInfo, nShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (LongType i = tid; i < len; i += totalThreads) {
|
||||
LongType nOffset, xOffset, zOffset;
|
||||
|
||||
// Compute offsets for n
|
||||
nOffset = i * (sameOffsetNX ? 0 : shape::stride(nShapeInfo)[0]);
|
||||
|
||||
// Compute offsets for x
|
||||
if (sameOffsetNX) {
|
||||
xOffset = nOffset;
|
||||
} else {
|
||||
xOffset = i * shape::stride(xShapeInfo)[0];
|
||||
}
|
||||
|
||||
// Compute offsets for z
|
||||
if (sameOffsetNZ) {
|
||||
zOffset = nOffset;
|
||||
} else {
|
||||
zOffset = i * shape::stride(zShapeInfo)[0];
|
||||
}
|
||||
|
||||
const T order = n[nOffset];
|
||||
const int sign = ((static_cast<int>(order) + 1) % 2 == 0) ? 1 : -1;
|
||||
|
||||
if (order != static_cast<int>(order)) {
|
||||
z[zOffset] = DataTypeUtils::nanOrZero<T>();
|
||||
} else if (order == 0) {
|
||||
z[zOffset] = diGammaScalar<T>(x[xOffset]);
|
||||
} else {
|
||||
T factorial = static_cast<T>(1);
|
||||
for (int j = 2; j <= order; ++j) {
|
||||
factorial *= j;
|
||||
}
|
||||
z[zOffset] = sign * factorial * zetaScalar<T>(order + 1, x[xOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void polyGammaCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
||||
const cudaStream_t *stream, const void *vn, const LongType *nShapeInfo,
|
||||
const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo) {
|
||||
polyGammaCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory, *stream>>>(vn, nShapeInfo, vx, xShapeInfo, vz, zShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "print_device failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void polyGamma(LaunchContext *context, NDArray&n, NDArray&x, NDArray &z) {
|
||||
NDArray::prepareSpecialUse({&z}, {&n, &x});
|
||||
|
||||
dim3 launchDims = polygammaDims(z.lengthOf());
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
n.dataType(), polyGammaCudaLauncher,
|
||||
(launchDims.y,launchDims.x,launchDims.z, context->getCudaStream(), n.specialBuffer(), n.specialShapeInfo(),
|
||||
x.specialBuffer(), x.specialShapeInfo(), z.specialBuffer(), z.specialShapeInfo()),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&z}, {&n, &x});
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void polyGammaCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,const cudaStream_t *stream, const void *vn,
|
||||
const sd::LongType *nShapeInfo, const void *vx, const sd::LongType *xShapeInfo, void *vz,
|
||||
const sd::LongType *zShapeInfo),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,140 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 12.06.2019
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/prefix.h>
|
||||
#include <ops/ops.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void prefix_(scalar::Ops op, const void* vx, const LongType* xShapeInfo, void* vz, const LongType* zShapeInfo, bool exclusive, bool reverse) {
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
const auto length = shape::length(xShapeInfo);
|
||||
const auto rankX = shape::rank(xShapeInfo);
|
||||
const auto rankZ = shape::rank(zShapeInfo);
|
||||
const auto shapeX = shape::shapeOf(xShapeInfo);
|
||||
const auto strideX = shape::stride(xShapeInfo);
|
||||
const auto shapeZ = shape::shapeOf(zShapeInfo);
|
||||
const auto strideZ = shape::stride(zShapeInfo);
|
||||
|
||||
T prevSum = (op == scalar::Add) ? static_cast<T>(0) : static_cast<T>(1);
|
||||
T sum = prevSum;
|
||||
|
||||
LongType coordsX[SD_MAX_RANK];
|
||||
LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
if (reverse) {
|
||||
for (LongType e = length - 1; e >= 0; --e) {
|
||||
LongType offsetX, offsetZ;
|
||||
|
||||
// Compute input and output offsets
|
||||
INDEX2COORDS(e, rankX, shapeX, coordsX);
|
||||
COORDS2INDEX(rankX, strideX, coordsX, offsetX);
|
||||
INDEX2COORDS(e, rankZ, shapeZ, coordsZ);
|
||||
COORDS2INDEX(rankZ, strideZ, coordsZ, offsetZ);
|
||||
|
||||
// Perform operation
|
||||
sum = (op == scalar::Add) ? simdOps::Add<T, T, T>::op(sum, x[offsetX]) : simdOps::Multiply<T, T, T>::op(sum, x[offsetX]);
|
||||
|
||||
if (!exclusive) prevSum = sum;
|
||||
|
||||
z[offsetZ] = prevSum;
|
||||
prevSum = sum;
|
||||
}
|
||||
} else {
|
||||
for (LongType e = 0; e < length; ++e) {
|
||||
LongType offsetX, offsetZ;
|
||||
|
||||
// Compute input and output offsets
|
||||
INDEX2COORDS(e, rankX, shapeX, coordsX);
|
||||
COORDS2INDEX(rankX, strideX, coordsX, offsetX);
|
||||
INDEX2COORDS(e, rankZ, shapeZ, coordsZ);
|
||||
COORDS2INDEX(rankZ, strideZ, coordsZ, offsetZ);
|
||||
|
||||
// Perform operation
|
||||
sum = (op == scalar::Add) ? simdOps::Add<T, T, T>::op(sum, x[offsetX]) : simdOps::Multiply<T, T, T>::op(sum, x[offsetX]);
|
||||
|
||||
if (!exclusive) prevSum = sum;
|
||||
|
||||
z[offsetZ] = prevSum;
|
||||
prevSum = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void prefix_(scalar::Ops op, NDArray* x, NDArray* z, const std::vector<LongType>& dims, bool exclusive,
|
||||
bool reverse) {
|
||||
NDArray::preparePrimaryUse({z}, {x});
|
||||
auto xTads = x->allTensorsAlongDimension(dims);
|
||||
auto zTads = z->allTensorsAlongDimension(dims);
|
||||
auto t = xTads.size();
|
||||
|
||||
for (int e = 0; e < t; e++) {
|
||||
auto tx = xTads.at(e);
|
||||
auto tz = zTads.at(e);
|
||||
|
||||
prefix_<T>(op, tx->buffer(), tx->shapeInfo(), tz->buffer(), tz->shapeInfo(), exclusive, reverse);
|
||||
}
|
||||
|
||||
NDArray::registerPrimaryUse({z}, {x});
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
|
||||
template <typename T>
|
||||
static void prefix_(scalar::Ops op, NDArray* x, NDArray* z, bool exclusive, bool reverse) {
|
||||
prefix_<T>(op, x->buffer(), x->shapeInfo(), z->buffer(), z->shapeInfo(), exclusive, reverse);
|
||||
};
|
||||
|
||||
void prefix(LaunchContext* context, scalar::Ops op, NDArray* x, NDArray* z, bool exclusive, bool reverse) {
|
||||
BUILD_SINGLE_SELECTOR(x->dataType(), prefix_, (op, x, z, exclusive, reverse), SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
void prefix(LaunchContext* context, scalar::Ops op, NDArray* x, NDArray* z, const std::vector<LongType>& dims,
|
||||
bool exclusive, bool reverse) {
|
||||
BUILD_SINGLE_SELECTOR(x->dataType(), prefix_, (op, x, z, dims, exclusive, reverse), SD_COMMON_TYPES);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void prefix_,
|
||||
(scalar::Ops op, const void* vx, sd::LongType const* xShapeInfo, void* vz,
|
||||
sd::LongType const* zShapeInfo, bool exclusive, bool reverse),
|
||||
SD_COMMON_TYPES);
|
||||
BUILD_SINGLE_TEMPLATE( void prefix_,
|
||||
(scalar::Ops op, NDArray* x, NDArray* z, const std::vector<sd::LongType>& dims, bool exclusive,
|
||||
bool reverse),
|
||||
SD_COMMON_TYPES);
|
||||
BUILD_SINGLE_TEMPLATE( void prefix_,
|
||||
(scalar::Ops op, NDArray* x, NDArray* z, bool exclusive, bool reverse), SD_COMMON_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,85 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/print_variable.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
static SD_KERNEL void print_device(const void *special, const LongType *shapeInfo) {
|
||||
__shared__ LongType length;
|
||||
__shared__ LongType rank;
|
||||
__shared__ const LongType *shape, *stride;
|
||||
const auto x = reinterpret_cast<const T *>(special);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(shapeInfo);
|
||||
rank = shape::rank(shapeInfo);
|
||||
shape = shape::shapeOf(shapeInfo);
|
||||
stride = shape::stride(shapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0 && blockIdx.x == 0) {
|
||||
printf("[");
|
||||
|
||||
LongType coords[SD_MAX_RANK];
|
||||
LongType offset;
|
||||
|
||||
for (LongType e = 0; e < length; e++) {
|
||||
INDEX2COORDS(e, rank, shape, coords);
|
||||
COORDS2INDEX(rank, stride, coords, offset);
|
||||
|
||||
printf("%f", static_cast<float>(x[offset]));
|
||||
if (e < length - 1) printf(", ");
|
||||
}
|
||||
|
||||
printf("]\n");
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static SD_HOST void exec_print_device(LaunchContext &ctx, const void *special, const LongType *shapeInfo) {
|
||||
dim3 launchDims = getLaunchDims("print");
|
||||
print_device<T><<<launchDims.x, launchDims.y, launchDims.z, *ctx.getCudaStream()>>>(special, shapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(ctx.getCudaStream(), "print_device failed");
|
||||
|
||||
}
|
||||
|
||||
void print_special(LaunchContext &ctx, NDArray&array, const std::string &message) {
|
||||
NDArray::prepareSpecialUse({}, {&array});
|
||||
|
||||
PointersManager pm(&ctx, "print_device");
|
||||
BUILD_SINGLE_SELECTOR(array.dataType(), exec_print_device, (ctx, array.specialBuffer(), array.specialShapeInfo()),
|
||||
SD_COMMON_TYPES)
|
||||
pm.synchronize();
|
||||
|
||||
NDArray::registerSpecialUse({}, {&array});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,196 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 George A. Shulinok <sgazeos@gmail.com>
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <helpers/MmulHelper.h>
|
||||
#include <ops/declarable/helpers/qr.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void matrixMinorKernel(T* outBuffer, LongType* outShape, T* inBuffer, LongType* inShape,
|
||||
LongType column, LongType rows, LongType columns) {
|
||||
|
||||
for (auto i = blockIdx.x; i < rows; i += gridDim.x)
|
||||
for (auto j = threadIdx.x; j < columns; j += blockDim.x) {
|
||||
LongType pos[] = {i, j};
|
||||
|
||||
LongType zIndex;
|
||||
COORDS2INDEX(shape::rank(outShape), shape::stride(outShape), pos, zIndex);
|
||||
|
||||
LongType xIndex;
|
||||
COORDS2INDEX(shape::rank(inShape), shape::stride(inShape), pos, xIndex);
|
||||
|
||||
if (i < column || j < column) {
|
||||
outBuffer[zIndex] = i != j ? T(0.f) : T(1.f);
|
||||
} else {
|
||||
outBuffer[zIndex] = inBuffer[xIndex]; // m.t<T>(i,j) = in.t<T>(i,j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
NDArray matrixMinor(LaunchContext* context, NDArray& in, LongType col) {
|
||||
NDArray *m = in.ulike();
|
||||
m->setIdentity();
|
||||
NDArray view = *m;
|
||||
NDArray assign = in({col, m->rows(), col, m->columns()});
|
||||
view({col, m->rows(), col, m->columns()}).assign(&assign);
|
||||
|
||||
m->tickWriteDevice();
|
||||
return *m;
|
||||
}
|
||||
|
||||
/* m = I - v v^T */
|
||||
template <typename T>
|
||||
static SD_KERNEL void vmulKernel(T* resBuf, const LongType* resShape, T const* vBuff, LongType const* vShape,
|
||||
LongType n) {
|
||||
for (auto i = blockIdx.x; i < n; i += gridDim.x)
|
||||
for (auto j = threadIdx.x; j < n; j += blockDim.x) {
|
||||
LongType posR[] = {i, j};
|
||||
LongType indexR, indexX, indexY;
|
||||
COORDS2INDEX(shape::rank(resShape), shape::stride(resShape), posR, indexR);
|
||||
COORDS2INDEX(1, shape::stride(vShape), &i, indexX);
|
||||
COORDS2INDEX(1, shape::stride(vShape), &j, indexY);
|
||||
|
||||
resBuf[indexR] = T(-2.f) * vBuff[indexX] * vBuff[indexY] + (i != j ? T(0.f) : T(1.f));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
NDArray vmul(LaunchContext* context, NDArray& v, int n) {
|
||||
std::vector<LongType> shape = {n, n};
|
||||
NDArray res('c', shape, v.dataType(), context); // x = matrix_new(n, n);
|
||||
|
||||
auto stream = context->getCudaStream();
|
||||
dim3 launchDims = getLaunchDims("qr");
|
||||
vmulKernel<T><<<launchDims.x,launchDims.y, launchDims.z, *stream>>>(res.dataBuffer()->specialAsT<T>(), res.specialShapeInfo(),
|
||||
reinterpret_cast<T const*>(v.specialBuffer()), v.specialShapeInfo(), n);
|
||||
sd::DebugHelper::checkErrorCode(stream, "vmulKernel failed");
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static bool diagonalIsPositive(NDArray* matrix, LongType k) {
|
||||
T hVal;
|
||||
LongType pos[] = {k, k};
|
||||
LongType shift;
|
||||
COORDS2INDEX(shape::rank(matrix->shapeInfo()), shape::stride(matrix->shapeInfo()), pos, shift);
|
||||
cudaMemcpy(&hVal, matrix->specialBuffer(), sizeof(T), cudaMemcpyDeviceToHost);
|
||||
return hVal > T(0.f);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void qrSingle(LaunchContext* context, NDArray* matrix, NDArray* Q, NDArray* R, bool const fullMatrices) {
|
||||
LongType M = matrix->sizeAt(0);
|
||||
LongType N = matrix->sizeAt(1);
|
||||
auto resQ = fullMatrices ? *Q->ulike() : NDArrayFactory::create<T>(matrix->ordering(), {M, M}, Q->getContext());
|
||||
auto resR = fullMatrices ? R->ulike() : matrix->ulike();
|
||||
std::vector<NDArray*> q(M, nullptr);
|
||||
NDArray z = *matrix;
|
||||
std::vector<LongType> shape = {M};
|
||||
|
||||
NDArray e('c', shape, DataTypeUtils::fromT<T>(), context); // two internal buffers and scalar for squared norm
|
||||
for (auto k = 0; k < N && k < M - 1; k++) { // loop for columns, but not further then row number
|
||||
e.nullify();
|
||||
z = matrixMinor<T>(context, z,
|
||||
k); // minor computing for current column with given matrix z (initally is a input matrix)
|
||||
|
||||
auto currentColumn = z({0, 0, k, k + 1}); // retrieve k column from z to x buffer
|
||||
std::vector<LongType> zero = {0};
|
||||
auto norm = currentColumn.reduceAlongDimension(reduce::Norm2, &zero);
|
||||
if (diagonalIsPositive<T>(matrix, k)) // matrix->t<T>(k,k) > T(0.f)) // negate on positive matrix diagonal element
|
||||
norm.applyTransform(transform::Neg, &norm); // *= -1.f;//-norm.t<T>(0);
|
||||
|
||||
e.p(k, &norm); // e - is filled by 0 vector except diagonal element (filled by 1)
|
||||
e += currentColumn; // e[i] = x[i] + a * e[i] for each i from 0 to n - 1
|
||||
auto normE = e.reduceAlongDimension(reduce::Norm2, &zero);
|
||||
e /= normE;
|
||||
q[k] = new NDArray(vmul<T>(context, e, M));
|
||||
auto qQ = z.ulike();
|
||||
MmulHelper::matmul(q[k], &z, qQ, false, false,1.0,0.0,qQ);
|
||||
z = std::move(*qQ);
|
||||
}
|
||||
resQ.assign(q[0]);
|
||||
|
||||
for (int i = 1; i < N && i < M - 1; i++) {
|
||||
auto tempResQ = resQ;
|
||||
MmulHelper::matmul(q[i],&resQ, &tempResQ, false, false,1.0,0.0,&tempResQ);
|
||||
resQ = std::move(tempResQ);
|
||||
}
|
||||
MmulHelper::matmul(&resQ, matrix, resR, false, false,1.0,0.0,resR);
|
||||
// resR *= -1.f;
|
||||
resQ.transposei();
|
||||
|
||||
if (fullMatrices) {
|
||||
Q->assign(&resQ);
|
||||
R->assign(resR);
|
||||
} else {
|
||||
NDArray resRRef = *resR;
|
||||
NDArray qAssign = resQ({0, 0, 0, N});
|
||||
Q->assign(&qAssign);
|
||||
NDArray rAssign = resRRef({0, N, 0, 0});
|
||||
R->assign(&rAssign);
|
||||
}
|
||||
|
||||
// Clean up allocated NDArrays in q vector
|
||||
for (LongType i = 0; i < M; i++) {
|
||||
if (q[i] != nullptr) {
|
||||
delete q[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void qr_(LaunchContext* context, NDArray * input, NDArray* outputQ, NDArray* outputR, bool const fullMatricies) {
|
||||
LongType lastDim = input->rankOf() - 1;
|
||||
LongType preLastDim = input->rankOf() - 2;
|
||||
|
||||
NDArray::prepareSpecialUse({outputQ, outputR}, {input});
|
||||
ResultSet listOutQ(outputQ->allTensorsAlongDimension({(int)preLastDim, (int)lastDim}));
|
||||
ResultSet listOutR(outputR->allTensorsAlongDimension({(int)preLastDim, (int)lastDim}));
|
||||
ResultSet listInput(input->allTensorsAlongDimension({(int)preLastDim, (int)lastDim}));
|
||||
auto start = 0;
|
||||
auto stop = listInput.size();
|
||||
auto increment = 1;
|
||||
|
||||
for (auto batch = start; batch < stop; batch += increment) {
|
||||
// qr here
|
||||
qrSingle<T>(context, listInput.at(batch), listOutQ.at(batch), listOutR.at(batch), fullMatricies);
|
||||
}
|
||||
NDArray::registerSpecialUse({outputQ, outputR}, {input});
|
||||
}
|
||||
|
||||
void qr(LaunchContext* context, NDArray * input, NDArray* outputQ, NDArray* outputR,
|
||||
bool const fullMatricies) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), qr_, (context, input, outputQ, outputR, fullMatricies), SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,542 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <graph/Context.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/RandomLauncher.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/random.h>
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
/**
|
||||
* gammaLess - compute gamma distributed value for shapes (alpha) from 0 to 1
|
||||
* @tparam T - any float types are acceptable
|
||||
* @param U - uniform random generated vals
|
||||
* @param alpha - shape of distribution
|
||||
* @param beta - scale of distributed values
|
||||
* @return gamma distributed value
|
||||
*/
|
||||
template <typename T>
|
||||
T SD_DEVICE gammaLess(T const* U, LongType index, LongType maxLength, T const alpha, T const beta) {
|
||||
auto d = T(1.0334f) - T(0.0766f) * math::p_exp(T(2.2942f) * alpha);
|
||||
auto a = math::p_pow(T(2.f), alpha) * math::p_pow<T>(T(1.f) - math::p_exp(-d * T(0.5f)), alpha);
|
||||
auto b = alpha * math::p_pow(d, alpha - T(1.f)) * exp(-d);
|
||||
auto c = a + b;
|
||||
T rawX;
|
||||
auto indexV = index;
|
||||
auto underAlpha = T(1.f) / alpha;
|
||||
auto powerAlpha = math::p_pow<T>(T(2.f), alpha - T(1.f));
|
||||
|
||||
for (;;) {
|
||||
auto u = (indexV < maxLength) ? U[indexV++] : U[0];
|
||||
if (indexV >= maxLength) indexV = 0LL;
|
||||
if (u <= a / c)
|
||||
rawX = -T(2.f) * math::p_log(T(1.f) - T(0.5f) * math::p_pow(c * u, underAlpha));
|
||||
else
|
||||
rawX = -math::p_log(c * (T(1.f) - u) / (alpha * math::p_pow(d, alpha - T(1.f))));
|
||||
|
||||
T v = indexV < maxLength ? U[indexV++] : U[0];
|
||||
if (indexV >= maxLength) indexV = 0LL;
|
||||
// math::atomics::sd_atomicAdd(index, 1LL);
|
||||
|
||||
if (rawX <= d) {
|
||||
auto testVal = (math::p_pow<T>(rawX, alpha - 1.f) * math::p_exp(-T(0.5f) * rawX)) /
|
||||
(powerAlpha * math::p_pow<T>(T(1.f) - math::p_exp(-T(0.5f) * rawX), alpha - T(1.f)));
|
||||
if (testVal < v) continue;
|
||||
break;
|
||||
} else {
|
||||
if (v <= math::p_pow<T>(d / rawX, T(1.f) - alpha)) break;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
return rawX / beta;
|
||||
}
|
||||
|
||||
/**
|
||||
* gammaGreat - generate gamma distributed value for shape (alpha) greater then 1
|
||||
* @tparam T - given type (any float type is accepted.)
|
||||
* @param rng - random generator
|
||||
* @param alpha - shape of the gamma distribution (alpha)
|
||||
* @param beta - scale of the gamma distribution (beta)
|
||||
* @return - gamma distributed value with given params
|
||||
*/
|
||||
template <typename T>
|
||||
T SD_DEVICE gammaGreat(T const* U, LongType index, LongType maxLength, T const alpha, T const beta) {
|
||||
auto decreasedAlpha = alpha - T(1.f / 3.f);
|
||||
auto c = T(1.) / math::p_sqrt(T(9.f) * decreasedAlpha);
|
||||
auto indexV = index;
|
||||
T x;
|
||||
auto normalDistributed = [U, maxLength](LongType& index) {
|
||||
auto v1 = index < maxLength ? U[index++] : U[0];
|
||||
if (index >= maxLength) index = 0LL;
|
||||
auto v2 = index < maxLength ? U[index++] : U[0];
|
||||
if (index >= maxLength) index = 0LL;
|
||||
return math::p_cos(T(2.f * 3.141592f) * v2) * math::p_sqrt(T(-2.f) * math::p_log(v1));
|
||||
};
|
||||
|
||||
float normalizedVar;
|
||||
for (;;) {
|
||||
do {
|
||||
x = normalDistributed(indexV);
|
||||
normalizedVar = T(1.f) + c * x;
|
||||
} while (normalizedVar < T(0.f));
|
||||
normalizedVar = normalizedVar * normalizedVar * normalizedVar; // v * v * v;
|
||||
|
||||
auto u = U[indexV++];
|
||||
if (indexV >= maxLength) indexV = 0LL;
|
||||
if (u < T(1.f) - T(.0331f) * (x * x) * (x * x)) break; // return (d * v / b);
|
||||
if (log(u) < 0.5f * x * x + decreasedAlpha * (1. - normalizedVar + math::p_log(normalizedVar))) break;
|
||||
}
|
||||
return (decreasedAlpha * normalizedVar / beta);
|
||||
}
|
||||
|
||||
/*
|
||||
* fillGammaKernel - fill up output with gamma distributed values
|
||||
*
|
||||
* uList - uniformly distributed values set
|
||||
* uLength - length of uList
|
||||
* alpha - alpha param
|
||||
* beta - beta param
|
||||
* output - distributed output.
|
||||
* */
|
||||
template <typename T>
|
||||
static SD_KERNEL void fillGammaKernel(T const* uList, LongType uLength, T const* alpha,
|
||||
const LongType* alphaShape, T const* beta, const LongType* betaShape,
|
||||
T* output, const LongType* outputShape) {
|
||||
__shared__ LongType aLength;
|
||||
__shared__ LongType outLength;
|
||||
__shared__ LongType rankAlpha, rankBeta, rankOutput;
|
||||
__shared__ const LongType *alphaShapeArr, *alphaStride, *betaShapeArr, *betaStride, *outputShapeArr, *outputStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
aLength = shape::length(alphaShape);
|
||||
outLength = shape::length(outputShape) / aLength;
|
||||
|
||||
rankAlpha = shape::rank(alphaShape);
|
||||
alphaShapeArr = shape::shapeOf(alphaShape);
|
||||
alphaStride = shape::stride(alphaShape);
|
||||
|
||||
rankBeta = betaShape ? shape::rank(betaShape) : 0;
|
||||
betaShapeArr = betaShape ? shape::shapeOf(betaShape) : nullptr;
|
||||
betaStride = betaShape ? shape::stride(betaShape) : nullptr;
|
||||
|
||||
rankOutput = shape::rank(outputShape);
|
||||
outputShapeArr = shape::shapeOf(outputShape);
|
||||
outputStride = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType k = blockIdx.x; k < outLength; k += gridDim.x) {
|
||||
const auto pos = k * aLength;
|
||||
|
||||
for (LongType e = threadIdx.x; e < aLength; e += blockDim.x) {
|
||||
LongType aCoords[SD_MAX_RANK], bCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType aIndex, bIndex = -1LL, zIndex;
|
||||
|
||||
// Alpha coordinates and index
|
||||
INDEX2COORDS(e, rankAlpha, alphaShapeArr, aCoords);
|
||||
COORDS2INDEX(rankAlpha, alphaStride, aCoords, aIndex);
|
||||
|
||||
// Beta coordinates and index (if beta is provided)
|
||||
if (betaShape) {
|
||||
INDEX2COORDS(e, rankBeta, betaShapeArr, bCoords);
|
||||
COORDS2INDEX(rankBeta, betaStride, bCoords, bIndex);
|
||||
}
|
||||
|
||||
// Output coordinates and index
|
||||
INDEX2COORDS(e + pos, rankOutput, outputShapeArr, zCoords);
|
||||
COORDS2INDEX(rankOutput, outputStride, zCoords, zIndex);
|
||||
|
||||
// Get beta value or default to 1
|
||||
const auto betaV = beta ? beta[bIndex] : T(1.f);
|
||||
|
||||
// Fill the output with the computed gamma value
|
||||
output[zIndex] = alpha[aIndex] > T(1.f)
|
||||
? gammaGreat(uList, pos, uLength, alpha[aIndex], betaV)
|
||||
: gammaLess(uList, pos, uLength, alpha[aIndex], betaV);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void fillRandomGamma_(LaunchContext* context, graph::RandomGenerator& rng, NDArray* alpha, NDArray* beta,
|
||||
NDArray* output) {
|
||||
// To fill up output need to broadcast alpha and beta to the same shape and in
|
||||
LongType* broadcasted = nullptr;
|
||||
if (beta != nullptr)
|
||||
ShapeUtils::evalBroadcastShapeInfo(alpha->shapeInfo(), beta->shapeInfo(), true, broadcasted, context->getWorkspace());
|
||||
else
|
||||
broadcasted = alpha->shapeInfo();
|
||||
auto step = shape::length(broadcasted);
|
||||
auto shift = output->lengthOf() * 4LL; // 2-wise greater case for uniform vals
|
||||
|
||||
auto copyAlpha = alpha;
|
||||
auto copyBeta = beta;
|
||||
if (beta != nullptr) {
|
||||
NDArray alphaBroadcasted(broadcasted, alpha->dataType(), true, context);
|
||||
NDArray betaBroadcasted(broadcasted, beta->dataType(), true, context);
|
||||
|
||||
copyAlpha = new NDArray(alphaBroadcasted.applyTrueBroadcast(BroadcastOpsTuple::Assign(), alpha));
|
||||
copyBeta = new NDArray(betaBroadcasted.applyTrueBroadcast(BroadcastOpsTuple::Assign(), beta));
|
||||
}
|
||||
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray uniform = NDArrayFactory::create<T>('c', {shift}, context);
|
||||
uniform.syncToDevice();
|
||||
// fill up uniform with given length
|
||||
RandomLauncher::fillUniform(context, rng, &uniform, 0.0000000001, 0.9999999999);
|
||||
uniform.syncToDevice();
|
||||
dim3 launchDims = getLaunchDims("random_gamma");
|
||||
fillGammaKernel<T><<<launchDims.x, launchDims.y,launchDims.z, *stream>>>(
|
||||
uniform.dataBuffer()->specialAsT<T>(), shift, copyAlpha->dataBuffer()->specialAsT<T>(),
|
||||
copyAlpha->specialShapeInfo(), beta ? copyBeta->dataBuffer()->specialAsT<T>() : (T const*)nullptr,
|
||||
beta ? copyBeta->specialShapeInfo() : (LongType const*)nullptr, output->dataBuffer()->specialAsT<T>(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillGammaKernel failed");
|
||||
|
||||
if (beta != nullptr) {
|
||||
delete copyAlpha;
|
||||
delete copyBeta;
|
||||
}
|
||||
}
|
||||
|
||||
void fillRandomGamma(LaunchContext* context, graph::RandomGenerator& rng, NDArray* alpha, NDArray* beta,
|
||||
NDArray* output) {
|
||||
if (beta)
|
||||
NDArray::prepareSpecialUse({output}, {alpha, beta});
|
||||
else
|
||||
NDArray::prepareSpecialUse({output}, {alpha});
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), fillRandomGamma_, (context, rng, alpha, beta, output), SD_FLOAT_NATIVE);
|
||||
if (beta)
|
||||
NDArray::registerSpecialUse({output}, {alpha, beta});
|
||||
else
|
||||
NDArray::prepareSpecialUse({output}, {alpha});
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void fillRandomGamma_,
|
||||
(LaunchContext * context, graph::RandomGenerator& rng, NDArray* alpha, NDArray* beta,
|
||||
NDArray* output),
|
||||
SD_FLOAT_NATIVE);
|
||||
|
||||
/*
|
||||
* algorithm Poisson generator based upon the inversion by sequential search
|
||||
*
|
||||
init:
|
||||
Let x ← 0, p ← e−λ, s ← p.
|
||||
using uniformly random sequence U (u in U) distributed at [0, 1].
|
||||
while u > s do:
|
||||
x ← x + 1.
|
||||
p ← p * λ / x.
|
||||
s ← s + p.
|
||||
return x.
|
||||
* */
|
||||
template <typename T>
|
||||
static SD_KERNEL void fillPoissonKernel(T* uList, LongType uLength, T* lambda, const LongType* lambdaShape,
|
||||
T* output, const LongType* outputShape) {
|
||||
__shared__ LongType lambdaLen;
|
||||
__shared__ LongType rankLambda, rankOutput;
|
||||
__shared__ const LongType *lambdaShapeArr, *lambdaStride, *outputShapeArr, *outputStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
lambdaLen = shape::length(lambdaShape);
|
||||
|
||||
rankLambda = shape::rank(lambdaShape);
|
||||
rankOutput = shape::rank(outputShape);
|
||||
|
||||
lambdaShapeArr = shape::shapeOf(lambdaShape);
|
||||
lambdaStride = shape::stride(lambdaShape);
|
||||
|
||||
outputShapeArr = shape::shapeOf(outputShape);
|
||||
outputStride = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType k = blockIdx.x; k < uLength; k += gridDim.x) {
|
||||
const auto pos = k * lambdaLen;
|
||||
const auto u = uList[k];
|
||||
|
||||
for (LongType e = threadIdx.x; e < lambdaLen; e += blockDim.x) {
|
||||
auto p = math::sd_exp<T, T>(-lambda[e]);
|
||||
auto s = p;
|
||||
auto x = T(0.);
|
||||
|
||||
LongType lCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType lIndex, zIndex;
|
||||
|
||||
// Compute coordinates and indices for lambda
|
||||
INDEX2COORDS(e, rankLambda, lambdaShapeArr, lCoords);
|
||||
COORDS2INDEX(rankLambda, lambdaStride, lCoords, lIndex);
|
||||
|
||||
// Compute coordinates and indices for output
|
||||
INDEX2COORDS(e + pos, rankOutput, outputShapeArr, zCoords);
|
||||
COORDS2INDEX(rankOutput, outputStride, zCoords, zIndex);
|
||||
|
||||
// Compute Poisson distributed value
|
||||
while (u > s) {
|
||||
x += T(1.);
|
||||
p *= lambda[lIndex] / x;
|
||||
s += p;
|
||||
}
|
||||
|
||||
// Assign computed value to output
|
||||
output[zIndex] = x;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void fillRandomPoisson_(LaunchContext* context, graph::RandomGenerator& rng, NDArray* lambda, NDArray* output) {
|
||||
auto shift = output->lengthOf() / lambda->lengthOf();
|
||||
std::vector<LongType> shape = {shift};
|
||||
NDArray uniform('c',shape, DOUBLE);
|
||||
PointersManager manager(context, "fillRandomPoisson");
|
||||
auto stream = context->getCudaStream();
|
||||
// fill up uniform with given length
|
||||
NDArray tempOutput = output->cast(DOUBLE);
|
||||
RandomLauncher::fillUniform(context, rng, &uniform, 0., 1.);
|
||||
|
||||
NDArray tempLambda = lambda->cast(DOUBLE);
|
||||
NDArray::prepareSpecialUse({output,&tempOutput}, {lambda,&tempLambda});
|
||||
|
||||
dim3 launchDims = getLaunchDims("random_poisson");
|
||||
fillPoissonKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(uniform.dataBuffer()->specialAsT<T>(), uniform.lengthOf(),
|
||||
tempLambda.dataBuffer()->specialAsT<T>(), tempLambda.specialShapeInfo(),
|
||||
tempOutput.dataBuffer()->specialAsT<T>(), tempOutput.specialShapeInfo());
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillPoissonKernel failed");
|
||||
NDArray ret = tempOutput.cast(output->dataType());
|
||||
output->assign(&ret);
|
||||
NDArray::registerSpecialUse({output,&tempOutput}, {lambda,&tempLambda});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
void fillRandomPoisson(LaunchContext* context, graph::RandomGenerator& rng, NDArray* lambda, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {lambda});
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), fillRandomPoisson_, (context, rng, lambda, output), SD_FLOAT_NATIVE);
|
||||
NDArray::registerSpecialUse({output}, {lambda});
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void fillRandomPoisson_,
|
||||
(LaunchContext * context, graph::RandomGenerator& rng, NDArray* lambda, NDArray* output),
|
||||
SD_FLOAT_NATIVE);
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void fillUniformKernel(graph::RandomGenerator* devRng, T from, T to, T* output,
|
||||
const LongType* outputShape) {
|
||||
const auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = blockDim.x * gridDim.x;
|
||||
|
||||
__shared__ LongType outputLen;
|
||||
__shared__ LongType rank;
|
||||
__shared__ const LongType *shape, *stride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outputLen = shape::length(outputShape);
|
||||
rank = shape::rank(outputShape);
|
||||
shape = shape::shapeOf(outputShape);
|
||||
stride = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = start; i < outputLen; i += step) {
|
||||
LongType zIndex;
|
||||
|
||||
// Calculate output index
|
||||
INDEX2COORDS(i, rank, shape, zCoords);
|
||||
COORDS2INDEX(rank, stride, zCoords, zIndex);
|
||||
|
||||
// Fill output with a random value in the range [from, to]
|
||||
output[zIndex] = devRng->relativeT<T>(i, from, to);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void fillRandomUniform_(LaunchContext* context, graph::RandomGenerator& rng, NDArray* min, NDArray* max,
|
||||
NDArray* output) {
|
||||
T minVal = T(0);
|
||||
T maxVal = DataTypeUtils::infOrMax<T>();
|
||||
if (min) minVal = min->t<T>(0);
|
||||
if (max) maxVal = max->t<T>(0);
|
||||
|
||||
if (output->isR())
|
||||
RandomLauncher::fillUniform(context, rng, output, minVal, maxVal);
|
||||
else {
|
||||
auto stream = context->getCudaStream();
|
||||
graph::RandomGenerator* devRng;
|
||||
auto err = cudaMalloc(&devRng, sizeof(graph::RandomGenerator));
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomUniform_: Cannot allocate device memory for random generator due error", err);
|
||||
}
|
||||
|
||||
err = cudaMemcpy(devRng, &rng, sizeof(graph::RandomGenerator), cudaMemcpyHostToDevice);
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomUniform_: Cannot copy random generator to device", err);
|
||||
}
|
||||
auto outputBuf = output->dataBuffer()->specialAsT<T>();
|
||||
auto outputShape = output->specialShapeInfo();
|
||||
dim3 launchDims = getLaunchDims("random_uniform");
|
||||
fillUniformKernel<T><<<launchDims.x,launchDims.y, launchDims.z, *stream>>>(devRng, minVal, maxVal, outputBuf, outputShape);
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillUniformKernel failed");
|
||||
|
||||
err = cudaStreamSynchronize(*stream);
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomUniform_: Cannot successfully finish kernel call", err);
|
||||
}
|
||||
|
||||
err = cudaFree(devRng);
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomUniform_: Cannot deallocate device memory for random generator", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fillRandomUniform(LaunchContext* context, graph::RandomGenerator& rng, NDArray* min, NDArray* max,
|
||||
NDArray* output) {
|
||||
BUILD_SINGLE_SELECTOR(output->dataType(), fillRandomUniform_, (context, rng, min, max, output), SD_NUMERIC_TYPES);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// used https://en.wikipedia.org/wiki/Categorical_distribution
|
||||
// methods: gumbel trick + softmax + argmax
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL static void fillMultiNomialCuda_(graph::RandomGenerator* devRng, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType batchValue,
|
||||
const LongType numOfSamples, const LongType numOfClassX, const LongType dimA, const X minVal,
|
||||
const X maxVal) {
|
||||
const X* x = reinterpret_cast<const X*>(vx);
|
||||
Z* z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ LongType xDimAstride, zDimAstride, xDimCstride, zDimCstride, dimC;
|
||||
|
||||
if (0 == threadIdx.x) {
|
||||
dimC = (0 == dimA) ? 1 : 0;
|
||||
zDimAstride = shape::stride(zShapeInfo)[dimA];
|
||||
xDimAstride = shape::stride(xShapeInfo)[dimA];
|
||||
zDimCstride = shape::stride(zShapeInfo)[dimC];
|
||||
xDimCstride = shape::stride(xShapeInfo)[dimC];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType index = tid; index < batchValue * numOfSamples; index += gridDim.x * blockDim.x) {
|
||||
LongType nBatchIndex = index / numOfSamples;
|
||||
LongType nSampleIndexInBatch = index - (nBatchIndex * numOfSamples);
|
||||
|
||||
const X* xTad = x + (nBatchIndex * xDimCstride);
|
||||
Z* zTad = z + (nBatchIndex * zDimCstride);
|
||||
Z& arg = zTad[nSampleIndexInBatch * zDimAstride];
|
||||
|
||||
X Max = -minVal;
|
||||
LongType nSamplesPerBatch = nBatchIndex * numOfClassX * numOfSamples;
|
||||
LongType nClassPerSamples = nSampleIndexInBatch * numOfClassX;
|
||||
|
||||
for (LongType nClass = 0; nClass < numOfClassX; nClass++) {
|
||||
LongType nIndex = nSamplesPerBatch + nClassPerSamples + nClass;
|
||||
X tValue = (xTad[nClass * xDimAstride] -
|
||||
math::sd_log<X, X>(-math::sd_log<X, X>(devRng->relativeT<X>(nIndex, minVal, maxVal))));
|
||||
if (tValue > Max) {
|
||||
Max = tValue;
|
||||
arg = nClass;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_HOST static void fillMultiNomialCudaLauncher(const int blocksPerGrid, const int threadsPerBlock,
|
||||
const cudaStream_t* stream, graph::RandomGenerator* devRng,
|
||||
const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType batchValue,
|
||||
const LongType numOfSamples, const LongType numOfClassX,
|
||||
const LongType dimA) {
|
||||
const X minVal = DataTypeUtils::min<X>();
|
||||
const X maxVal = static_cast<X>(1.0);
|
||||
|
||||
fillMultiNomialCuda_<X, Z><<<blocksPerGrid, threadsPerBlock, 256, *stream>>>(
|
||||
devRng, vx, xShapeInfo, vz, zShapeInfo, batchValue, numOfSamples, numOfClassX, dimA, minVal, maxVal);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "fillMultiNomialCuda_ failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void fillRandomMultiNomial(LaunchContext* context, graph::RandomGenerator& rng, NDArray& input, NDArray& output,
|
||||
const LongType numOfSamples, const int dimC) {
|
||||
LongType dimA = (0 == dimC) ? 1 : 0;
|
||||
|
||||
const LongType batchValue = output.sizeAt(dimC);
|
||||
const LongType numOfClassX = input.sizeAt(dimA);
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (batchValue * numOfSamples + threadsPerBlock - 1) / threadsPerBlock;
|
||||
|
||||
PointersManager manager(context, "fillMultinomial");
|
||||
graph::RandomGenerator* devRng;
|
||||
|
||||
auto err = cudaMalloc(&devRng, sizeof(graph::RandomGenerator));
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomMultiNomial: Cannot allocate device memory for random generator due error", err);
|
||||
}
|
||||
err = cudaStreamSynchronize(*context->getCudaStream());
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomMultiNomial: Cannot synchronize stream for random generator due error", err);
|
||||
}
|
||||
err =
|
||||
cudaMemcpyAsync(devRng, &rng, sizeof(graph::RandomGenerator), cudaMemcpyHostToDevice, *context->getCudaStream());
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomMultiNomial: Cannot copy random generator to device", err);
|
||||
}
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&input});
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), output.dataType(), fillMultiNomialCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, context->getCudaStream(), devRng, input.specialBuffer(),
|
||||
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), batchValue,
|
||||
numOfSamples, numOfClassX, dimA),
|
||||
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&input});
|
||||
manager.synchronize();
|
||||
|
||||
err = cudaFree(devRng);
|
||||
if (err != 0) {
|
||||
cuda_exception::build("fillRandomMultiNomial: Cannot deallocate device memory for random generator", err);
|
||||
}
|
||||
rng.rewindH(output.lengthOf() * numOfClassX);
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,228 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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)
|
||||
// implemented algorithm is GPU adaptation of algorithm described in following article:
|
||||
// "MergeShuffle: A Very Fast, Parallel Random Permutation Algorithm", https://arxiv.org/abs/1508.03167
|
||||
//
|
||||
#include <array/ResultSet.h>
|
||||
#include <execution/Threads.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static SD_KERNEL void fisherYatesCuda(graph::RandomGenerator* rng, void* vx, const LongType ews,
|
||||
const LongType len, const int power) {
|
||||
T* x = reinterpret_cast<T*>(vx);
|
||||
|
||||
__shared__ T *shmem, temp;
|
||||
__shared__ LongType ind, blockOffset, lenPerBlock;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char sharedMemory[];
|
||||
shmem = reinterpret_cast<T*>(sharedMemory);
|
||||
|
||||
blockOffset = (len * blockIdx.x) >> power;
|
||||
lenPerBlock = ((len * (blockIdx.x + 1)) >> power) - blockOffset;
|
||||
ind = blockOffset;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// copy from global memory to shared memory
|
||||
if (threadIdx.x < lenPerBlock) shmem[threadIdx.x] = x[(blockOffset + threadIdx.x) * ews];
|
||||
__syncthreads();
|
||||
|
||||
// *** apply Fisher-Yates shuffle to lenPerBlock number of elements
|
||||
if (threadIdx.x == 0) {
|
||||
for (LongType i = lenPerBlock - 1; i > 0; --i) {
|
||||
const LongType j = rng->relativeLong(ind++) % (i + 1);
|
||||
if (i != j) {
|
||||
temp = shmem[i];
|
||||
shmem[i] = shmem[j];
|
||||
shmem[j] = temp;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// copy from shared memory to global memory
|
||||
if (threadIdx.x < lenPerBlock) x[(blockOffset + threadIdx.x) * ews] = shmem[threadIdx.x];
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void mergeShuffleCuda(graph::RandomGenerator* rng, void* vx, const LongType ews,
|
||||
const LongType len, const int power, const LongType iterNum) {
|
||||
T* x = reinterpret_cast<T*>(vx);
|
||||
|
||||
__shared__ LongType ind, blockOffset, factor, beg, mid, totLen, iterExp;
|
||||
|
||||
// *** apply mergeShuffle algorithm
|
||||
if (threadIdx.x == 0) {
|
||||
factor = blockIdx.x << iterNum;
|
||||
iterExp = 1 << (iterNum - 1);
|
||||
blockOffset = (len * factor) >> power;
|
||||
mid = ((len * (factor + iterExp)) >> power) - blockOffset; // middle
|
||||
totLen = ((len * (factor + 2 * iterExp)) >> power) - blockOffset;
|
||||
ind = iterNum * len + blockOffset;
|
||||
beg = 0; // beginning
|
||||
|
||||
while (true) {
|
||||
if (rng->relativeLong(ind++) % 2) {
|
||||
if (mid == totLen) break;
|
||||
int first = (blockOffset + beg) * ews;
|
||||
int second = blockOffset + mid * ews;
|
||||
if(first >= len || second >= len) {
|
||||
break;
|
||||
}
|
||||
math::sd_swap<T>(x[(blockOffset + beg) * ews], x[(blockOffset + mid++) * ews]);
|
||||
} else {
|
||||
if (beg == mid) break;
|
||||
}
|
||||
++beg;
|
||||
}
|
||||
|
||||
// Fisher-Yates
|
||||
while (beg < totLen) {
|
||||
const LongType e = rng->relativeLong(ind++) % (beg + 1);
|
||||
int first = (blockOffset + beg) * ews;
|
||||
int second = blockOffset + e * ews;
|
||||
if(first >= len || second >= len) {
|
||||
break;
|
||||
}
|
||||
if (beg != e) math::sd_swap<T>(x[(blockOffset + beg) * ews], x[(blockOffset + e) * ews]);
|
||||
++beg;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// Fisher-Yates shuffle
|
||||
template <typename T>
|
||||
static void fisherYates(graph::RandomGenerator& rng, T* buff, const LongType& len, const LongType& ews, LongType ind) {
|
||||
for (LongType i = len - 1; i > 0; --i) {
|
||||
const LongType j = rng.relativeLong(ind++) % (i + 1);
|
||||
if (i != j) math::sd_swap<T>(buff[i * ews], buff[j * ews]);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void randomShuffle_(LaunchContext* context, NDArray& input, NDArray& output, graph::RandomGenerator& rng,
|
||||
const bool isInplace) {
|
||||
const int firstDim = input.sizeAt(0);
|
||||
LongType temp;
|
||||
|
||||
if (input.lengthOf() == 1 || firstDim == 1) {
|
||||
if (!isInplace) output.assign(&input);
|
||||
} else if (shape::isCommonVector(input.shapeInfo(), temp)) {
|
||||
NDArray* arr = &input;
|
||||
|
||||
if (!isInplace) {
|
||||
output.assign(&input);
|
||||
arr = &output;
|
||||
}
|
||||
|
||||
const LongType len = arr->lengthOf();
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS;
|
||||
|
||||
int power = 0;
|
||||
while ((len >> power) > threadsPerBlock) ++power;
|
||||
|
||||
dim3 fisherDims = randomShuffleFisherDims(power,input.sizeOfT());
|
||||
const int blocksPerGrid = fisherDims.y;
|
||||
const int sharedMem = fisherDims.z;
|
||||
|
||||
PointersManager manager(context, "NDArray::randomShuffle cuda");
|
||||
|
||||
graph::RandomGenerator* pRng = reinterpret_cast<graph::RandomGenerator*>(
|
||||
manager.replicatePointer(&rng, sizeof(graph::RandomGenerator)));
|
||||
|
||||
NDArray::prepareSpecialUse({arr}, {arr});
|
||||
|
||||
fisherYatesCuda<T><<<fisherDims.y, fisherDims.x, fisherDims.z, *context->getCudaStream()>>>(
|
||||
pRng, arr->specialBuffer(),0, len, power);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "fisherYatesCuda failed");
|
||||
|
||||
for (LongType j = 1, i = 1; j < blocksPerGrid; j += j, ++i) {
|
||||
dim3 mergeShuffleDims = randomShuffleMergeDims(j, power);
|
||||
mergeShuffleCuda<T><<<mergeShuffleDims.x, mergeShuffleDims.y, mergeShuffleDims.z, *context->getCudaStream()>>>(
|
||||
pRng, arr->specialBuffer(), 0, len, power, i);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "mergeShuffleCuda failed");
|
||||
|
||||
NDArray::registerSpecialUse({arr}, {arr});
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
rng.rewindH((len + 1) * power);
|
||||
}
|
||||
} else {
|
||||
LongType dim = 0;
|
||||
auto dimsToExclude = ShapeUtils::evalDimsToExclude(input.rankOf(),1 ,&dim);
|
||||
|
||||
if (isInplace) {
|
||||
auto subArrsList = input.allTensorsAlongDimension(*dimsToExclude);
|
||||
|
||||
// Fisher-Yates shuffle
|
||||
for (int i = firstDim - 1; i > 0; --i) {
|
||||
const int j = rng.relativeInt(i) % (i + 1);
|
||||
if (i != j) subArrsList.at(i)->swapUnsafe(*subArrsList.at(j));
|
||||
}
|
||||
} else {
|
||||
auto subArrsListIn = input.allTensorsAlongDimension(*dimsToExclude);
|
||||
auto subArrsListOut = output.allTensorsAlongDimension(*dimsToExclude);
|
||||
|
||||
std::vector<int> indices(firstDim);
|
||||
std::iota(indices.begin(), indices.end(), 0); // 0,1,2,3, ... firstDim-1
|
||||
|
||||
// shuffle indices
|
||||
fisherYates<int>(rng, indices.data(), firstDim, 1, 0);
|
||||
|
||||
auto func = PRAGMA_THREADS_FOR {
|
||||
for (auto i = start; i < stop; ++i) subArrsListOut.at(i)->assign(subArrsListIn.at(indices[i]));
|
||||
};
|
||||
|
||||
samediff::Threads::parallel_for(func, 0, firstDim);
|
||||
}
|
||||
|
||||
rng.rewindH(firstDim - 1);
|
||||
|
||||
delete dimsToExclude;
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
void randomShuffle(LaunchContext* context, NDArray& input, NDArray& output, graph::RandomGenerator& rng,
|
||||
const bool isInplace) {
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), randomShuffle_, (context, input, output, rng, isInplace), SD_COMMON_TYPES);
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,50 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 sgazeos@gmail.com
|
||||
//
|
||||
#include <ops/declarable/helpers/random_crop.h>
|
||||
#include <graph/Context.h>
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static Status _randomCropFunctor(graph::Context& context, NDArray* input, NDArray* shape, NDArray* output, int seed) {
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
Status randomCropFunctor(graph::Context& context, NDArray* input, NDArray* shape, NDArray* output, int seed) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), return _randomCropFunctor, (context, input, shape, output, seed),
|
||||
SD_FLOAT_TYPES);
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( sd::Status _randomCropFunctor,
|
||||
(graph::Context & context, NDArray* input, NDArray* shape, NDArray* output, int seed),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,61 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 27.08.2018
|
||||
//
|
||||
|
||||
#include <ops/declarable/helpers/range.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void global_range(void* output, LongType length, T start, T delta) {
|
||||
auto buff = reinterpret_cast<T*>(output);
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
for (LongType i = tid; i < length; i += step) {
|
||||
buff[i] = static_cast<T>(start) + static_cast<T>(i) * static_cast<T>(delta);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
// be careful: outVector must have c-order and ews = 1 !!!
|
||||
template <typename T>
|
||||
static void _range(LaunchContext* context, NDArray& start, NDArray& delta, NDArray& outVector) {
|
||||
dim3 launchDims = getLaunchDims("range");
|
||||
global_range<T><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(outVector.specialBuffer(), outVector.lengthOf(),
|
||||
start.e<T>(0), delta.e<T>(0));
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "global_range failed");
|
||||
|
||||
}
|
||||
|
||||
void range(LaunchContext* context, NDArray& start, NDArray& delta, NDArray& outVector) {
|
||||
NDArray::prepareSpecialUse({&outVector}, {&start, &delta});
|
||||
BUILD_SINGLE_SELECTOR(outVector.dataType(), _range, (context, start, delta, outVector), SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({&outVector}, {&start, &delta});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,319 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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, created on 16.04.2018
|
||||
//
|
||||
#include <array/ResultSet.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/reverse.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void reverseTadKernel(const void* vinput, const LongType* inputShape, void* voutput,
|
||||
const LongType* outputShape, const LongType* inputTadShape,
|
||||
const LongType* inputTadOffsets, const LongType* outputTadShape,
|
||||
const LongType* outputTadOffsets, uint64_t limit,
|
||||
uint64_t numOfElemsToReverse, uint64_t numTads) {
|
||||
auto input = reinterpret_cast<const T*>(vinput);
|
||||
auto output = reinterpret_cast<T*>(voutput);
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ LongType tadRankInput, tadRankOutput;
|
||||
__shared__ const LongType *tadShapeInput, *strideInput, *tadShapeOutput, *strideOutput;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadRankInput = shape::rank(inputTadShape);
|
||||
tadShapeInput = shape::shapeOf(inputTadShape);
|
||||
strideInput = shape::stride(inputTadShape);
|
||||
|
||||
tadRankOutput = shape::rank(outputTadShape);
|
||||
tadShapeOutput = shape::shapeOf(outputTadShape);
|
||||
strideOutput = shape::stride(outputTadShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto div = numOfElemsToReverse / 2;
|
||||
const auto odd = numOfElemsToReverse % 2 != 0;
|
||||
const auto rlimit = odd ? limit / 2 + 1 : limit / 2;
|
||||
|
||||
// Main loop for element swaps
|
||||
for (uint64_t e = tid; e < rlimit; e += step) {
|
||||
const auto tadId = e / div;
|
||||
if (tadId >= numTads) continue;
|
||||
|
||||
const auto idx = e % div;
|
||||
|
||||
const auto tadInput = input + inputTadOffsets[tadId];
|
||||
const auto tadOutput = output + outputTadOffsets[tadId];
|
||||
|
||||
LongType fCoords[SD_MAX_RANK], lCoords[SD_MAX_RANK];
|
||||
LongType fOffset, lOffset;
|
||||
|
||||
// Input coordinates and offsets
|
||||
INDEX2COORDS(idx, tadRankInput, tadShapeInput, fCoords);
|
||||
COORDS2INDEX(tadRankInput, strideInput, fCoords, fOffset);
|
||||
|
||||
INDEX2COORDS(numOfElemsToReverse - idx - 1, tadRankInput, tadShapeInput, lCoords);
|
||||
COORDS2INDEX(tadRankInput, strideInput, lCoords, lOffset);
|
||||
|
||||
auto v1 = tadInput[fOffset];
|
||||
auto v2 = tadInput[lOffset];
|
||||
|
||||
LongType zfCoords[SD_MAX_RANK], zlCoords[SD_MAX_RANK];
|
||||
LongType zfOffset, zlOffset;
|
||||
|
||||
// Output coordinates and offsets
|
||||
INDEX2COORDS(idx, tadRankOutput, tadShapeOutput, zfCoords);
|
||||
COORDS2INDEX(tadRankOutput, strideOutput, zfCoords, zfOffset);
|
||||
|
||||
INDEX2COORDS(numOfElemsToReverse - idx - 1, tadRankOutput, tadShapeOutput, zlCoords);
|
||||
COORDS2INDEX(tadRankOutput, strideOutput, zlCoords, zlOffset);
|
||||
|
||||
// Store swapped values
|
||||
tadOutput[zfOffset] = v2;
|
||||
tadOutput[zlOffset] = v1;
|
||||
}
|
||||
|
||||
// Handle odd middle element
|
||||
if (odd && threadIdx.x == 0) {
|
||||
for (uint64_t e = blockIdx.x; e < numTads; e += gridDim.x) {
|
||||
const auto tadInput = input + inputTadOffsets[e];
|
||||
const auto tadOutput = output + outputTadOffsets[e];
|
||||
|
||||
LongType xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType xOffset, zOffset;
|
||||
|
||||
// Coordinates and offsets for the middle element
|
||||
INDEX2COORDS(numOfElemsToReverse / 2, tadRankInput, tadShapeInput, xCoords);
|
||||
COORDS2INDEX(tadRankInput, strideInput, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(numOfElemsToReverse / 2, tadRankOutput, tadShapeOutput, zCoords);
|
||||
COORDS2INDEX(tadRankOutput, strideOutput, zCoords, zOffset);
|
||||
|
||||
tadOutput[zOffset] = tadInput[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static SD_KERNEL void reverseArrayKernel(const void* input, const LongType* inputShape, void* output,
|
||||
const LongType* outputShape, LongType numOfElemsToReverse) {
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto step = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ const T* inputArr;
|
||||
__shared__ T* outputArr;
|
||||
__shared__ LongType rankInput, rankOutput;
|
||||
__shared__ const LongType *inputShapeArr, *inputStride, *outputShapeArr, *outputStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inputArr = reinterpret_cast<const T*>(input);
|
||||
outputArr = reinterpret_cast<T*>(output);
|
||||
|
||||
rankInput = shape::rank(inputShape);
|
||||
rankOutput = shape::rank(outputShape);
|
||||
|
||||
inputShapeArr = shape::shapeOf(inputShape);
|
||||
inputStride = shape::stride(inputShape);
|
||||
|
||||
outputShapeArr = shape::shapeOf(outputShape);
|
||||
outputStride = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto odd = numOfElemsToReverse % 2 != 0;
|
||||
const auto limit = numOfElemsToReverse / 2;
|
||||
|
||||
for (LongType e = tid; e < limit; e += step) {
|
||||
LongType fCoords[SD_MAX_RANK], lCoords[SD_MAX_RANK];
|
||||
LongType fOffset, lOffset;
|
||||
|
||||
// Input indices
|
||||
INDEX2COORDS(e, rankInput, inputShapeArr, fCoords);
|
||||
COORDS2INDEX(rankInput, inputStride, fCoords, fOffset);
|
||||
|
||||
INDEX2COORDS(numOfElemsToReverse - e - 1, rankInput, inputShapeArr, lCoords);
|
||||
COORDS2INDEX(rankInput, inputStride, lCoords, lOffset);
|
||||
|
||||
auto v1 = inputArr[fOffset];
|
||||
auto v2 = inputArr[lOffset];
|
||||
|
||||
LongType zfCoords[SD_MAX_RANK], zlCoords[SD_MAX_RANK];
|
||||
LongType zfOffset, zlOffset;
|
||||
|
||||
// Output indices
|
||||
INDEX2COORDS(e, rankOutput, outputShapeArr, zfCoords);
|
||||
COORDS2INDEX(rankOutput, outputStride, zfCoords, zfOffset);
|
||||
|
||||
INDEX2COORDS(numOfElemsToReverse - e - 1, rankOutput, outputShapeArr, zlCoords);
|
||||
COORDS2INDEX(rankOutput, outputStride, zlCoords, zlOffset);
|
||||
|
||||
outputArr[zfOffset] = v2;
|
||||
outputArr[zlOffset] = v1;
|
||||
}
|
||||
|
||||
// Handle the odd middle element if applicable
|
||||
if (odd && tid == 0) {
|
||||
LongType xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
||||
LongType xOffset, zOffset;
|
||||
|
||||
INDEX2COORDS(limit, rankInput, inputShapeArr, xCoords);
|
||||
COORDS2INDEX(rankInput, inputStride, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(limit, rankOutput, outputShapeArr, zCoords);
|
||||
COORDS2INDEX(rankOutput, outputStride, zCoords, zOffset);
|
||||
|
||||
outputArr[zOffset] = inputArr[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void reverseTad(LaunchContext* context, NDArray* input, NDArray* output,
|
||||
const LongType* inputTadShape, const LongType* inputTadOffsets,
|
||||
const LongType* outputTadShape, const LongType* outputTadOffsets, uint64_t tadLength) {
|
||||
auto stream = context->getCudaStream();
|
||||
dim3 launchDims = getLaunchDims("reverse");
|
||||
|
||||
reverseTadKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(input->specialBuffer(), input->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), inputTadShape,
|
||||
inputTadOffsets, outputTadShape, outputTadOffsets, input->lengthOf(),
|
||||
tadLength, input->lengthOf() / tadLength);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reverseTadKernel failed");
|
||||
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void reverseArray(LaunchContext* context, NDArray* input, NDArray* output, LongType numOfElemsToReverse) {
|
||||
auto stream = context->getCudaStream();
|
||||
LongType numOfReverse = numOfElemsToReverse;
|
||||
if (numOfElemsToReverse == 0) numOfReverse = input->lengthOf();
|
||||
dim3 launchDims = getLaunchDims("reverse");
|
||||
|
||||
reverseArrayKernel<T><<<launchDims.y,launchDims.x, launchDims.z, *stream>>>(input->specialBuffer(), input->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo(), numOfReverse);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reverseArrayKernel failed");
|
||||
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void reverseSequence_(LaunchContext* context, NDArray* input, NDArray* seqLengths,
|
||||
NDArray* output, int seqDim, const int batchDim) {
|
||||
int posOfNonUnityDim = -1;
|
||||
seqLengths->syncToHost();
|
||||
auto stream = context->getCudaStream();
|
||||
dim3 launchDims = getLaunchDims("reverse");
|
||||
if (input->isVector() || shape::isLikeVector(input->shapeInfo(), posOfNonUnityDim) || seqLengths->lengthOf() == 1) {
|
||||
LongType numOfElemsToReverse = seqLengths->e<LongType>(0);
|
||||
if ((seqDim == 0 && input->sizeAt(0) == 1) || (batchDim == posOfNonUnityDim))
|
||||
output->assign(input);
|
||||
else
|
||||
reverseArrayKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
numOfElemsToReverse);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reverseArrayKernel failed");
|
||||
|
||||
} else {
|
||||
if (seqDim > batchDim) --seqDim;
|
||||
|
||||
std::vector<LongType> dim = {batchDim};
|
||||
std::vector<LongType> *dimensions = ShapeUtils::evalDimsToExclude(input->rankOf(), 1,dim.data());
|
||||
|
||||
auto inSubArrsSet = input->allTensorsAlongDimension(*dimensions);
|
||||
auto outSubArrsSet = output->allTensorsAlongDimension(*dimensions);
|
||||
|
||||
for (int i = 0; i < inSubArrsSet.size(); ++i) {
|
||||
LongType numOfElemsToReverse = seqLengths->e<LongType>(i);
|
||||
|
||||
if (numOfElemsToReverse == 0 || numOfElemsToReverse == 1) {
|
||||
outSubArrsSet.at(i)->assign(inSubArrsSet.at(i));
|
||||
} else {
|
||||
auto inInnerSet = inSubArrsSet.at(i)->allTensorsAlongDimension({seqDim});
|
||||
auto outInnerSet = outSubArrsSet.at(i)->allTensorsAlongDimension({seqDim});
|
||||
for (int j = 0; j < inInnerSet.size(); ++j)
|
||||
reverseArray<T>(context, inInnerSet.at(j), outInnerSet.at(j), numOfElemsToReverse);
|
||||
}
|
||||
}
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
|
||||
void reverseSequence(LaunchContext* context, NDArray* input, NDArray* seqLengths, NDArray* output,
|
||||
int seqDim, const int batchDim) {
|
||||
NDArray::prepareSpecialUse({output}, {input, seqLengths});
|
||||
|
||||
// if op isn't inplace - copy original data into output array
|
||||
if (output->specialBuffer() != input->specialBuffer()) output->assign(input);
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), reverseSequence_, (context, input, seqLengths, output, seqDim, batchDim),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, seqLengths});
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void reverse(LaunchContext* context, NDArray* input, NDArray* output, const std::vector<LongType>* intArgs) {
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), reinterpret_cast<LongType*>(*intArgs->data()),static_cast<sd::LongType>(intArgs->size()));
|
||||
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(output->shapeInfo(), reinterpret_cast<LongType*>(*intArgs->data()),static_cast<sd::LongType>(intArgs->size()));
|
||||
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
|
||||
if (packX->numberOfTads() == 1) {
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), reverseArray, (context, input, output, 0), SD_COMMON_TYPES);
|
||||
} else {
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input->dataType(), reverseTad,
|
||||
(context, input, output, packX->platformShapeInfo(), packX->platformOffsets(), packZ->platformShapeInfo(),
|
||||
packZ->platformOffsets(), (uint64_t)(input->lengthOf() / packX->numberOfTads())),
|
||||
SD_COMMON_TYPES);
|
||||
}
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,373 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/roll.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename T>
|
||||
static void SD_DEVICE rollKernelLinearStage1Dev(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType fullLength,
|
||||
int actualShift) {
|
||||
auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
// Cache shape information for x buffer
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
// Cache shape information for z buffer
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache x shape information
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache z shape information
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xOffsetA;
|
||||
LongType xOffsetB;
|
||||
LongType zOffsetA;
|
||||
LongType zOffsetB;
|
||||
|
||||
for (LongType i = tid; i < actualShift; i += blockDim.x * gridDim.x) {
|
||||
int sourceIndex = fullLength - actualShift + i;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffsetA);
|
||||
INDEX2COORDS(sourceIndex, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffsetB);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffsetA);
|
||||
INDEX2COORDS(sourceIndex, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffsetB);
|
||||
|
||||
auto eA = x[xOffsetA];
|
||||
auto eB = x[xOffsetB];
|
||||
|
||||
z[zOffsetA] = eB;
|
||||
z[zOffsetB] = eA;
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void SD_KERNEL rollKernelLinearStage1(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType fullLength, int actualShift) {
|
||||
rollKernelLinearStage1Dev<T>(vx, xShapeInfo, vz, zShapeInfo, fullLength, actualShift);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void SD_KERNEL rollKernelLinearStage2(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType fullLength, int actualShift,
|
||||
int shiftCount) {
|
||||
auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
// Cache shape information for x buffer
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
// Cache shape information for z buffer
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache x shape information
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache z shape information
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xOffsetA;
|
||||
LongType xOffsetB;
|
||||
LongType zOffsetA;
|
||||
LongType zOffsetB;
|
||||
|
||||
for (int count = 1; count < shiftCount; ++count) {
|
||||
for (int i = tid; i < actualShift; i += blockDim.x * gridDim.x) {
|
||||
int destinationIndex = fullLength - (count + 1) * actualShift + i;
|
||||
int sourceIndex = fullLength - count * actualShift + i;
|
||||
|
||||
INDEX2COORDS(destinationIndex, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffsetA);
|
||||
INDEX2COORDS(sourceIndex, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffsetB);
|
||||
|
||||
INDEX2COORDS(destinationIndex, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffsetA);
|
||||
INDEX2COORDS(sourceIndex, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffsetB);
|
||||
|
||||
auto eA = x[xOffsetB];
|
||||
auto eB = x[xOffsetA];
|
||||
|
||||
z[zOffsetA] = eA;
|
||||
z[zOffsetB] = eB;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void SD_KERNEL rollKernelLinearStage3(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, LongType fullLength, int actualShift,
|
||||
int remainShift) {
|
||||
auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
// Cache shape information for x buffer
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
// Cache shape information for z buffer
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache x shape information
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache z shape information
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
for (int i = tid; i < actualShift; i += blockDim.x * gridDim.x) {
|
||||
int remainIdx = i + actualShift;
|
||||
int sourceIndex = remainIdx + remainShift;
|
||||
|
||||
LongType xCoordsA[SD_MAX_RANK];
|
||||
LongType xCoordsB[SD_MAX_RANK];
|
||||
LongType zCoordsA[SD_MAX_RANK];
|
||||
LongType zCoordsB[SD_MAX_RANK];
|
||||
LongType xOffsetA;
|
||||
LongType xOffsetB;
|
||||
LongType zOffsetA;
|
||||
LongType zOffsetB;
|
||||
|
||||
INDEX2COORDS(remainIdx, xRank, xShapePtr, xCoordsA);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoordsA, xOffsetA);
|
||||
INDEX2COORDS(sourceIndex, xRank, xShapePtr, xCoordsB);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoordsB, xOffsetB);
|
||||
|
||||
INDEX2COORDS(remainIdx, zRank, zShapePtr, zCoordsA);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoordsA, zOffsetA);
|
||||
INDEX2COORDS(sourceIndex, zRank, zShapePtr, zCoordsB);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoordsB, zOffsetB);
|
||||
|
||||
auto eA = x[xOffsetA];
|
||||
auto eB = x[xOffsetB];
|
||||
|
||||
z[zOffsetA] = eB;
|
||||
z[zOffsetB] = eA;
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void SD_DEVICE swapTadsKernel(void *vx, void *vz, const LongType *zShapeInfo, LongType tadLength) {
|
||||
auto x = reinterpret_cast<T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
// Cache shape information for z buffer
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache z shape information
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
for (int e = threadIdx.x; e < tadLength; e += blockDim.x) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
|
||||
INDEX2COORDS(e, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
auto eA = x[zOffset];
|
||||
auto eB = z[zOffset];
|
||||
|
||||
x[zOffset] = eB;
|
||||
z[zOffset] = eA;
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
static void SD_KERNEL rollKernelFullAnyDimensionStage1(const void *vx, const LongType *xTadShapeInfo,
|
||||
const LongType *xTadOffsets, void *vz,
|
||||
const LongType *zTadShapeInfo,
|
||||
const LongType *zTadOffsets, int numTads, LongType tadLength, int dim, LongType sizeAt,
|
||||
int theShift) {
|
||||
auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
for (int e = blockIdx.x + theShift; e < sizeAt - theShift; e += gridDim.x) {
|
||||
int sourceIndex = dim * sizeAt + e - theShift;
|
||||
int targetIndex = dim * sizeAt + e;
|
||||
|
||||
swapTadsKernel<T>(z + xTadOffsets[sourceIndex], z + xTadOffsets[targetIndex], zTadShapeInfo, tadLength);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void SD_KERNEL rollKernelFullAnyDimensionStage2(void *vx, const LongType *xTadShapeInfo,
|
||||
const LongType *xTadOffsets, void *vz,
|
||||
const LongType *zTadShapeInfo,
|
||||
const LongType *zTadOffsets, int numTads, LongType tadLength, int dim, LongType sizeAt,
|
||||
int theShift) {
|
||||
auto x = reinterpret_cast<const T *>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
for (int e = blockIdx.x; e < theShift; e += gridDim.x) {
|
||||
int sourceIndex = dim * sizeAt + sizeAt - theShift + e;
|
||||
int targetIndex = dim * sizeAt + e;
|
||||
|
||||
swapTadsKernel<T>(z + zTadOffsets[sourceIndex], z + zTadOffsets[targetIndex], zTadShapeInfo, tadLength);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void rollFunctorFull_(NDArray *input, NDArray *output, std::vector<LongType> const &shifts,
|
||||
std::vector<LongType> const &axes, bool inplace) {
|
||||
if (!inplace) output->assign(input);
|
||||
|
||||
for (size_t i = 0; i < axes.size(); i++) {
|
||||
int axe = axes[i];
|
||||
ResultSet listOfTensors = input->allTensorsAlongDimension({axe});
|
||||
ResultSet listOfOutTensors = output->allTensorsAlongDimension({axe});
|
||||
int fullLen = listOfTensors.size();
|
||||
int theShift = shifts[i];
|
||||
for (int k = 0; k < fullLen; k++) {
|
||||
rollFunctorLinear(output->getContext(), listOfTensors.at(k), listOfOutTensors.at(k), theShift, true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static void rollFunctorLinear_(NDArray *input, NDArray *output, int shift, bool inplace) {
|
||||
if (!inplace) output->assign(input);
|
||||
|
||||
dim3 launchDims = getLaunchDims("roll");
|
||||
auto fullLen = input->lengthOf();
|
||||
int actualShift = shift; // % fullLen; // shift already non-negative then
|
||||
if (actualShift < 0) {
|
||||
actualShift -= fullLen * (actualShift / fullLen - 1);
|
||||
} else
|
||||
actualShift %= fullLen;
|
||||
|
||||
if (actualShift) {
|
||||
int shiftCount = fullLen / actualShift - 1;
|
||||
int remainShift = fullLen % actualShift;
|
||||
|
||||
// stage 1) swap last actualShift elements with first ones.
|
||||
rollKernelLinearStage1<T><<<launchDims.y, launchDims.x, launchDims.z, *(output->getContext()->getCudaStream())>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
fullLen, actualShift);
|
||||
sd::DebugHelper::checkErrorCode(output->getContext()->getCudaStream(), "rollKernelLinearStage1 failed");
|
||||
|
||||
// stage 2) swap swapped actualShift elements with rest remainShiftCount times.
|
||||
rollKernelLinearStage2<T><<<launchDims.y, launchDims.x, launchDims.z, *(output->getContext()->getCudaStream())>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
fullLen, actualShift, shiftCount);
|
||||
sd::DebugHelper::checkErrorCode(output->getContext()->getCudaStream(), "rollKernelLinearStage2 failed");
|
||||
// FIXME: no parallelism here :(
|
||||
// stage 3) swap remainer of items.
|
||||
if (remainShift && shiftCount)
|
||||
rollKernelLinearStage3<T><<<launchDims.y,launchDims.x,launchDims.z, *(output->getContext()->getCudaStream())>>>(
|
||||
output->specialBuffer(), output->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
fullLen, actualShift, remainShift);
|
||||
sd::DebugHelper::checkErrorCode(output->getContext()->getCudaStream(), "rollKernelLinearStage3 failed");
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
void rollFunctorFull(LaunchContext *context, NDArray *input, NDArray *output, std::vector<LongType> const &shifts,
|
||||
std::vector<LongType> const &axes, bool inplace) {
|
||||
input->syncToDevice();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), rollFunctorFull_, (input, output, shifts, axes, inplace), SD_COMMON_TYPES);
|
||||
|
||||
output->tickWriteDevice();
|
||||
}
|
||||
|
||||
void rollFunctorLinear(LaunchContext *context, NDArray *input, NDArray *output, int shift, bool inplace) {
|
||||
input->syncToDevice();
|
||||
|
||||
BUILD_SINGLE_SELECTOR(input->dataType(), rollFunctorLinear_, (input, output, shift, inplace), SD_COMMON_TYPES);
|
||||
|
||||
output->tickWriteDevice();
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void rollFunctorLinear_, (NDArray * input, NDArray *output, int shift, bool inplace),
|
||||
SD_COMMON_TYPES);
|
||||
BUILD_SINGLE_TEMPLATE( void rollFunctorFull_,
|
||||
(NDArray * input, NDArray *output, std::vector<sd::LongType> const &shifts, std::vector<sd::LongType> const &axes,
|
||||
bool inplace),
|
||||
SD_COMMON_TYPES);
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,532 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// Created by raver119 on 19.01.18.
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <ops/declarable/helpers/s_t_b.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void batchToSpaceCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType cropBottom,
|
||||
const LongType cropLeft) {
|
||||
const auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType rank, zLen;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape, *zStride;
|
||||
__shared__ LongType* sharedMem;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
rank = shape::rank(zShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType* coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < zLen; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
// Adjust spatial coordinates for cropping
|
||||
coords[1] += cropBottom;
|
||||
coords[2] += cropLeft;
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
|
||||
// Assign the value from input to output
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void batchToSpaceCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType cropBottom,
|
||||
const LongType cropLeft) {
|
||||
batchToSpaceCuda<T>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, cropBottom, cropLeft);
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void batchToSpaceCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, void* vz,
|
||||
const sd::LongType* zShapeInfo, const sd::LongType cropBottom, const sd::LongType cropLeft),
|
||||
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> rearrShape = {blockSize, blockSize, output.sizeAt(0), input.sizeAt(1), input.sizeAt(2), input.sizeAt(3)};
|
||||
NDArray inputRearranged0 = input.reshape(
|
||||
input.ordering(), rearrShape,false);
|
||||
inputRearranged0.permutei({2, 3, 0, 4, 1, 5}, false, false);
|
||||
|
||||
if (input.lengthOf() == output.lengthOf()) {
|
||||
output.assign(&inputRearranged0);
|
||||
} else {
|
||||
std::vector<sd::LongType> outputShape = {output.sizeAt(0), input.sizeAt(1) * blockSize, input.sizeAt(2) * blockSize, input.sizeAt(3)};
|
||||
NDArray inputRearranged1 = inputRearranged0.reshape(
|
||||
input.ordering(),
|
||||
outputShape);
|
||||
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (output.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
const int sharedMem = threadsPerBlock * sizeof(LongType) * output.rankOf() + 128;
|
||||
|
||||
PointersManager manager(context, "batchToSpace");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&inputRearranged1});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
input.dataType(), batchToSpaceCudaLauncher,
|
||||
(blocksPerGrid, threadsPerBlock, sharedMem, context->getCudaStream(), inputRearranged1.specialBuffer(),
|
||||
inputRearranged1.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), cropBottom, cropLeft),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&inputRearranged1});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void batchToSpaceNDCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType numOfSpatialDims) {
|
||||
// x - input, y - crop, z - output
|
||||
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
const auto y = reinterpret_cast<const Y*>(vy);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
|
||||
__shared__ LongType rank, zLen;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape, *zStride;
|
||||
__shared__ LongType* sharedMem;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
rank = shape::rank(zShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType* coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < zLen; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
// Adjust spatial coordinates for cropping
|
||||
for (LongType j = 1; j <= numOfSpatialDims; ++j) {
|
||||
const LongType yOffset = (j - 1) * yShapeInfo[3]; // yRank = 2, calculate offset manually
|
||||
coords[j] += y[yOffset]; // Add crop offset (cropLeft for each spatial dimension)
|
||||
}
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
|
||||
z[zOffset] = x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void batchToSpaceNDCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType numOfSpatialDims) {
|
||||
batchToSpaceNDCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz,
|
||||
zShapeInfo, numOfSpatialDims);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "batchToSpaceNDCuda failed");
|
||||
|
||||
}
|
||||
BUILD_DOUBLE_TEMPLATE( void batchToSpaceNDCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, const void* vy,
|
||||
const sd::LongType* yShapeInfo, void* vz, const sd::LongType* zShapeInfo,
|
||||
const sd::LongType numOfSpatialDims),
|
||||
SD_COMMON_TYPES, SD_INTEGER_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 LongType rank = input.rankOf();
|
||||
const LongType numOfSpatialDims = blockShape.sizeAt(0);
|
||||
|
||||
//*** construct reshaping std::vector for first reshape of input array ***//
|
||||
|
||||
std::vector<LongType> temp(numOfSpatialDims + rank);
|
||||
|
||||
int i;
|
||||
for (i = 0; i < numOfSpatialDims; ++i) temp[i] = blockShape.e<LongType>(i);
|
||||
temp[i++] = output.sizeAt(0);
|
||||
for (int 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 < temp.size(); ++i) temp[i] = i;
|
||||
|
||||
inputRearranged0.permutei(temp, 0, 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<LongType>(i - 1) : input.sizeAt(i);
|
||||
|
||||
NDArray inputRearranged1 = inputRearranged0.reshape(input.ordering(), temp);
|
||||
|
||||
dim3 launchDims = batchToSpaceNdLaunch(output.lengthOf(),output.rankOf());
|
||||
|
||||
PointersManager manager(context, "batchToSpaceND");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&inputRearranged1, &crop});
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
input.dataType(), crop.dataType(), batchToSpaceNDCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), inputRearranged1.specialBuffer(),
|
||||
inputRearranged1.specialShapeInfo(), crop.specialBuffer(), crop.specialShapeInfo(), output.specialBuffer(),
|
||||
output.specialShapeInfo(), numOfSpatialDims),
|
||||
SD_COMMON_TYPES, SD_INTEGER_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&inputRearranged1, &crop});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void spaceToBatchCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType padBottom,
|
||||
const LongType padTop, const LongType padLeft,
|
||||
const 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 auto x = reinterpret_cast<const T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ LongType rank, *sharedMem;
|
||||
__shared__ LongType zLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
rank = shape::rank(zShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType* coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
const LongType i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (i >= zLen) return;
|
||||
|
||||
INDEX2COORDS(i, rank, shape::shapeOf(zShapeInfo), coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, shape::stride(zShapeInfo), coords, zOffset);
|
||||
|
||||
if (coords[1] >= padBottom && coords[1] < zShapeInfo[2] - padTop && coords[2] >= padLeft &&
|
||||
coords[2] < zShapeInfo[3] - padRight) {
|
||||
coords[1] -= padBottom;
|
||||
coords[2] -= padLeft;
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, shape::stride(xShapeInfo), coords, xOffset);
|
||||
|
||||
z[zOffset] = x[xOffset];
|
||||
} else
|
||||
z[zOffset] = static_cast<T>(0.f);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
static void spaceToBatchCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo, const LongType padBottom,
|
||||
const LongType padTop, const LongType padLeft, const LongType padRight) {
|
||||
spaceToBatchCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, padBottom,
|
||||
padTop, padLeft, padRight);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "spaceToBatchCudaLauncher failed");
|
||||
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void spaceToBatchCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, void* vz,
|
||||
const sd::LongType* zShapeInfo, const sd::LongType padBottom, const sd::LongType padTop,
|
||||
const sd::LongType padLeft, const sd::LongType padRight),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void spaceToBatch(LaunchContext* context, NDArray& input, NDArray& output, const LongType padBottom,
|
||||
const LongType padTop, const LongType padLeft, const LongType padRight,
|
||||
const 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> outputShape = {blockSize, blockSize, input.sizeAt(0), output.sizeAt(1), output.sizeAt(2), input.sizeAt(3)};
|
||||
NDArray outputRearranged0 = output.reshape(
|
||||
output.ordering(), outputShape,
|
||||
false);
|
||||
outputRearranged0.permutei({2, 3, 0, 4, 1, 5}, false, false);
|
||||
|
||||
if (input.lengthOf() == output.lengthOf()) {
|
||||
outputRearranged0.assign(&input);
|
||||
} else {
|
||||
std::vector<sd::LongType> outReArrShape = {input.sizeAt(0), output.sizeAt(1) * blockSize, output.sizeAt(2) * blockSize, input.sizeAt(3)};
|
||||
NDArray outputRearranged1 = outputRearranged0.reshape(
|
||||
output.ordering(),
|
||||
outReArrShape, false);
|
||||
|
||||
|
||||
dim3 launchDims = spaceToBatchLaunch(output.lengthOf(),output.rankOf());
|
||||
|
||||
PointersManager manager(context, "spaceToBatch");
|
||||
|
||||
NDArray::prepareSpecialUse({&outputRearranged1}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), spaceToBatchCudaLauncher,
|
||||
(launchDims.y,launchDims.x,launchDims.z, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), outputRearranged1.specialBuffer(),
|
||||
outputRearranged1.specialShapeInfo(), padBottom, padTop, padLeft, padRight),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({&outputRearranged1}, {&input});
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
if (output.specialBuffer() != outputRearranged1.specialBuffer()) outputRearranged0.assign(&outputRearranged1);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void spaceToBatchNDCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
||||
const LongType numOfSpatialDims) {
|
||||
// x - input, y - padding, z - output
|
||||
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
const auto y = reinterpret_cast<const Y*>(vy);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
|
||||
__shared__ LongType rank, zLen, totalThreads;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape, *zStride;
|
||||
__shared__ LongType *sharedMem;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sharedMem = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
rank = shape::rank(zShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
auto coords = sharedMem + threadIdx.x * rank;
|
||||
|
||||
for (LongType i = blockDim.x * blockIdx.x + threadIdx.x; i < zLen; i += totalThreads) {
|
||||
INDEX2COORDS(i, rank, zShape, coords);
|
||||
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(rank, zStride, coords, zOffset);
|
||||
|
||||
bool within = true;
|
||||
|
||||
for (LongType j = 1; j <= numOfSpatialDims; ++j) {
|
||||
// Manually calculate y offsets for padding
|
||||
const LongType yOffset = (j - 1) * yShapeInfo[3];
|
||||
const LongType padLeft = y[yOffset];
|
||||
const LongType padRight = y[yOffset + yShapeInfo[4]];
|
||||
|
||||
// Check if coordinates are within the valid range
|
||||
within &= (coords[j] >= padLeft && coords[j] < zShape[j] - padRight);
|
||||
|
||||
if (!within) break;
|
||||
|
||||
// Adjust coordinates for x
|
||||
coords[j] -= padLeft;
|
||||
}
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(rank, xStride, coords, xOffset);
|
||||
|
||||
// Assign values to z
|
||||
if (within)
|
||||
z[zOffset] = x[xOffset];
|
||||
else
|
||||
z[zOffset] = static_cast<X>(0.f);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void spaceToBatchNDCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
||||
const void* vy, const LongType* yShapeInfo, void* vz,
|
||||
const LongType* zShapeInfo, const LongType numOfSpatialDims) {
|
||||
spaceToBatchNDCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz,
|
||||
zShapeInfo, numOfSpatialDims);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "spaceToBatchNDCuda failed");
|
||||
|
||||
}
|
||||
BUILD_DOUBLE_TEMPLATE( void spaceToBatchNDCudaLauncher,
|
||||
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, const void* vy,
|
||||
const sd::LongType* yShapeInfo, void* vz, const sd::LongType* zShapeInfo,
|
||||
const sd::LongType numOfSpatialDims),
|
||||
SD_COMMON_TYPES, SD_INTEGER_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void spaceToBatchND(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 LongType rank = input.rankOf();
|
||||
|
||||
const LongType numOfSpatialDims = blockShape.sizeAt(0);
|
||||
|
||||
//*** construct reshaping std::vector for first reshape of output array ***//
|
||||
std::vector<LongType> temp(numOfSpatialDims + rank);
|
||||
|
||||
int i;
|
||||
for (i = 0; i < numOfSpatialDims; ++i) temp[i] = blockShape.e<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 < 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<LongType>(i - 1) : output.sizeAt(i);
|
||||
|
||||
NDArray outputRearranged1 = outputRearranged0.reshape(output.ordering(), temp, false);
|
||||
|
||||
dim3 launchDims = spaceToBatchNdLaunch(output.lengthOf(),output.rankOf());
|
||||
PointersManager manager(context, "spaceToBatchND");
|
||||
|
||||
NDArray::prepareSpecialUse({&outputRearranged1}, {&input, &padding});
|
||||
BUILD_DOUBLE_SELECTOR(input.dataType(), padding.dataType(), spaceToBatchNDCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), input.specialBuffer(),
|
||||
input.specialShapeInfo(), padding.specialBuffer(), padding.specialShapeInfo(),
|
||||
outputRearranged1.specialBuffer(), outputRearranged1.specialShapeInfo(), numOfSpatialDims),
|
||||
SD_COMMON_TYPES, SD_INTEGER_TYPES);
|
||||
NDArray::registerSpecialUse({&outputRearranged1}, {&input, &padding});
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
if (output.specialBuffer() != outputRearranged1.specialBuffer()) outputRearranged0.assign(&outputRearranged1);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,117 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
//
|
||||
//
|
||||
#include <ops/declarable/helpers/s_t_d.h>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename T>
|
||||
static SD_KERNEL void spaceToDepthKernel(const void *vx, const LongType *xShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const int block_size, const bool isNHWC) {
|
||||
auto input_ptr = reinterpret_cast<const T *>(vx);
|
||||
auto output_ptr = reinterpret_cast<T *>(vz);
|
||||
|
||||
const int batch_size = shape::sizeAt(xShapeInfo, 0);
|
||||
const int input_depth = isNHWC ? shape::sizeAt(xShapeInfo, 3) : shape::sizeAt(xShapeInfo, 1);
|
||||
const int input_height = isNHWC ? shape::sizeAt(xShapeInfo, 1) : shape::sizeAt(xShapeInfo, 2);
|
||||
const int input_width = isNHWC ? shape::sizeAt(xShapeInfo, 2) : shape::sizeAt(xShapeInfo, 3);
|
||||
|
||||
const int output_depth = isNHWC ? shape::sizeAt(zShapeInfo, 3) : shape::sizeAt(zShapeInfo, 1);
|
||||
const int output_height = isNHWC ? shape::sizeAt(zShapeInfo, 1) : shape::sizeAt(zShapeInfo, 2);
|
||||
const int output_width = isNHWC ? shape::sizeAt(zShapeInfo, 2) : shape::sizeAt(zShapeInfo, 3);
|
||||
|
||||
const int input_depth_by_output_height = input_depth * output_height;
|
||||
|
||||
const int output_area = output_width * output_height;
|
||||
const int output_depth_by_output_area = output_depth * output_area;
|
||||
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
if (isNHWC) {
|
||||
const int total_count = batch_size * input_height * input_width * input_depth;
|
||||
|
||||
for (int inp_idx = tid; inp_idx < total_count; inp_idx += blockDim.x * gridDim.x) {
|
||||
// inp_idx = d + input_depth * (w + input_width * (h + input_height * b))
|
||||
const int d = inp_idx % input_depth;
|
||||
const int inp_idx2 = inp_idx / input_depth;
|
||||
const int w = inp_idx2 % input_width;
|
||||
const int inp_idx3 = inp_idx2 / input_width;
|
||||
const int h = inp_idx3 % input_height;
|
||||
const int b = inp_idx3 / input_height;
|
||||
|
||||
const int out_h = h / block_size;
|
||||
const int offset_h = h % block_size;
|
||||
const int out_w = w / block_size;
|
||||
const int offset_w = w % block_size;
|
||||
const int offset_d = (offset_h * block_size + offset_w) * input_depth;
|
||||
const int out_d = d + offset_d;
|
||||
|
||||
const int out_idx = out_d + output_depth * (out_w + output_width * (out_h + output_height * b));
|
||||
*(output_ptr + out_idx) = *(input_ptr + inp_idx);
|
||||
}
|
||||
} else {
|
||||
const int total_count = batch_size * output_depth_by_output_area;
|
||||
|
||||
for (int inp_idx = tid; inp_idx < total_count; inp_idx += blockDim.x * gridDim.x) {
|
||||
const int n_iC_oY_bY_oX = inp_idx / block_size;
|
||||
const int bX = inp_idx - n_iC_oY_bY_oX * block_size;
|
||||
|
||||
const int n_iC_oY_bY = n_iC_oY_bY_oX / output_width;
|
||||
const int oX = n_iC_oY_bY_oX - n_iC_oY_bY * output_width;
|
||||
|
||||
const int n_iC_oY = n_iC_oY_bY / block_size;
|
||||
const int bY = n_iC_oY_bY - n_iC_oY * block_size;
|
||||
|
||||
const int n = n_iC_oY / input_depth_by_output_height;
|
||||
const int iC_oY = n_iC_oY - n * input_depth_by_output_height;
|
||||
|
||||
const int output_idx =
|
||||
oX + (((n * block_size + bY) * block_size + bX) * input_depth_by_output_height + iC_oY) * output_width;
|
||||
|
||||
*(output_ptr + output_idx) = *(input_ptr + inp_idx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static void _spaceTodepth_(LaunchContext *context, NDArray&input, NDArray *output, int block_size,
|
||||
bool isNHWC) {
|
||||
dim3 launchDims = getLaunchDims("space_to_depth");
|
||||
spaceToDepthKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(input.specialBuffer(), input.specialShapeInfo(),
|
||||
output->specialBuffer(),
|
||||
output->specialShapeInfo(), block_size, isNHWC);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "spaceToDepthKernel failed");
|
||||
|
||||
}
|
||||
|
||||
void _spaceTodepth(LaunchContext *context, NDArray&input, NDArray *output, int block_size, bool isNHWC) {
|
||||
NDArray::prepareSpecialUse({output}, {&input});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), _spaceTodepth_, (context, input, output, block_size, isNHWC),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {&input});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,751 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/ConstantShapeHelper.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/scatter.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
#include "helpers/DebugHelper.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, y - contains number of bad indices, z - input/output
|
||||
template <typename X>
|
||||
SD_KERNEL static void checkIndicesCuda(const void *vx, const LongType *xShapeInfo, LongType *y,
|
||||
const LongType *zShapeInfo, const int axis) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
|
||||
__shared__ LongType xRank, xLen, numOfBadIndxPerBlock;
|
||||
__shared__ const LongType *xShape, *xStride, *zShape;
|
||||
__shared__ LongType *coords;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
coords = reinterpret_cast<LongType *>(shmem);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xLen = shape::length(xShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
|
||||
numOfBadIndxPerBlock = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto xCoords = coords + threadIdx.x * xRank;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < xLen; i += gridDim.x * blockDim.x) {
|
||||
INDEX2COORDS(i, xRank, xShape, xCoords);
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
|
||||
const LongType currentInd = x[xOffset];
|
||||
|
||||
const LongType limit = shape::sizeAt(zShapeInfo, axis == -1 ? xCoords[xRank - 1] : axis);
|
||||
if (currentInd >= limit) {
|
||||
sd::math::atomics::sd_atomicAdd<LongType>(&numOfBadIndxPerBlock, 1);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0 && numOfBadIndxPerBlock != 0) {
|
||||
sd::math::atomics::sd_atomicAdd<LongType>(y, numOfBadIndxPerBlock);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
static void checkIndicesCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo,
|
||||
LongType *y, const LongType *zShapeInfo, const int axis) {
|
||||
checkIndicesCuda<X><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, y, zShapeInfo, axis);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "checkIndicesCuda failed");
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
LongType checkIndices(LaunchContext *context, NDArray&indices, NDArray&output, const int axis) {
|
||||
const int threadsPerBlock = SD_MAX_NUM_THREADS / 2;
|
||||
const int blocksPerGrid = (indices.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
||||
const int sharedMem = threadsPerBlock * sizeof(LongType) * indices.rankOf() + 256;
|
||||
dim3 scatterDimsIndices = scatterDimsCheckIndices(indices.lengthOf(), indices.rankOf());
|
||||
const auto xType = indices.dataType();
|
||||
|
||||
PointersManager manager(context, "scatterNDcheckIndices");
|
||||
|
||||
// scalar, initial value = 0
|
||||
NDArray numOfBadIndx(INT64, context, true);
|
||||
|
||||
NDArray::prepareSpecialUse({&numOfBadIndx}, {&indices});
|
||||
BUILD_SINGLE_SELECTOR(
|
||||
xType, checkIndicesCudaLauncher,
|
||||
(scatterDimsIndices.y, scatterDimsIndices.x, scatterDimsIndices.z, context->getCudaStream(),
|
||||
indices.specialBuffer(), indices.specialShapeInfo(),
|
||||
reinterpret_cast<sd::LongType *>(numOfBadIndx.specialBuffer()), output.specialShapeInfo(), axis),
|
||||
SD_INTEGER_TYPES);
|
||||
NDArray::registerSpecialUse({&numOfBadIndx}, {&indices});
|
||||
|
||||
manager.synchronize();
|
||||
|
||||
return numOfBadIndx.t<LongType>(0);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, y - updates, z - input/output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void scatterLockCuda(const int opCode, const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
const auto y = reinterpret_cast<const Y *>(vy);
|
||||
auto z = reinterpret_cast<Y *>(vz);
|
||||
|
||||
__shared__ LongType xRank, yRank, zRank, xNonUnitDim, yNonUnitDim, zNonUnitDim;
|
||||
__shared__ const LongType *xShape, *yShape, *zShape, *xStride, *yStride, *zStride;
|
||||
__shared__ LongType xLen, zLen;
|
||||
__shared__ bool is1Dcase, xySameStride;
|
||||
__shared__ LongType *coords;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
coords = reinterpret_cast<LongType *>(shmem);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
xLen = shape::length(xShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
xNonUnitDim = yNonUnitDim = zNonUnitDim = 0;
|
||||
|
||||
is1Dcase = (shape::isCommonVector(zShapeInfo, zNonUnitDim) || shape::isScalar(zShapeInfo)) &&
|
||||
(shape::isCommonVector(yShapeInfo, yNonUnitDim) || shape::isScalar(yShapeInfo)) &&
|
||||
(shape::isCommonVector(xShapeInfo, xNonUnitDim) || shape::isScalar(xShapeInfo));
|
||||
|
||||
if (is1Dcase) xySameStride = xStride[xNonUnitDim] == yStride[yNonUnitDim];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType yOffset, zOffset;
|
||||
LongType zFirstCoord, *yCoords, *zCoords;
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < zLen; i += gridDim.x * blockDim.x) {
|
||||
if (!is1Dcase) {
|
||||
yCoords = coords + threadIdx.x * (yRank + zRank);
|
||||
zCoords = yCoords + yRank;
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
}
|
||||
|
||||
for (LongType j = 0; j < xLen; ++j) {
|
||||
if (is1Dcase) {
|
||||
yOffset = j * yStride[yNonUnitDim];
|
||||
zFirstCoord = x[xySameStride ? yOffset : j];
|
||||
|
||||
if (i != zFirstCoord) continue;
|
||||
|
||||
zOffset = i * zStride[zNonUnitDim];
|
||||
} else {
|
||||
INDEX2COORDS(j, xRank, xShape, yCoords);
|
||||
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, yCoords, xOffset);
|
||||
zFirstCoord = x[xOffset];
|
||||
|
||||
if (zCoords[0] != zFirstCoord) continue;
|
||||
|
||||
for (LongType k = 0; k < yRank - xRank; ++k) yCoords[xRank + k] = zCoords[k + 1];
|
||||
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
}
|
||||
|
||||
switch (opCode) {
|
||||
case pairwise::Add:
|
||||
z[zOffset] += y[yOffset];
|
||||
break;
|
||||
case pairwise::Subtract:
|
||||
z[zOffset] -= y[yOffset];
|
||||
break;
|
||||
case pairwise::Multiply:
|
||||
z[zOffset] *= y[yOffset];
|
||||
break;
|
||||
case pairwise::Divide:
|
||||
z[zOffset] /= y[yOffset];
|
||||
break;
|
||||
case pairwise::ReverseSubtract:
|
||||
z[zOffset] = y[yOffset] - z[zOffset];
|
||||
break;
|
||||
case pairwise::ReverseDivide:
|
||||
z[zOffset] = y[yOffset] / z[zOffset];
|
||||
break;
|
||||
case pairwise::CopyPws:
|
||||
z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MaxPairwise:
|
||||
if (z[zOffset] < y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MinPairwise:
|
||||
if (z[zOffset] > y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
default:
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, y - updates, z - input/output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void scatterCuda(const int opCode, const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
const auto y = reinterpret_cast<const Y *>(vy);
|
||||
auto z = reinterpret_cast<Y *>(vz);
|
||||
|
||||
__shared__ LongType xRank, yRank, zRank, xNonUnitDim, yNonUnitDim, zNonUnitDim;
|
||||
__shared__ const LongType *xShape, *yShape, *zShape, *xStride, *yStride, *zStride;
|
||||
__shared__ LongType yLen;
|
||||
__shared__ bool is1Dcase, xySameStride;
|
||||
__shared__ LongType *coords;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
coords = reinterpret_cast<LongType *>(shmem);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
yLen = shape::length(yShapeInfo);
|
||||
|
||||
xNonUnitDim = yNonUnitDim = zNonUnitDim = 0;
|
||||
|
||||
is1Dcase = (shape::isCommonVector(zShapeInfo, zNonUnitDim) || shape::isScalar(zShapeInfo)) &&
|
||||
(shape::isCommonVector(yShapeInfo, yNonUnitDim) || shape::isScalar(yShapeInfo)) &&
|
||||
(shape::isCommonVector(xShapeInfo, xNonUnitDim) || shape::isScalar(xShapeInfo));
|
||||
|
||||
if (is1Dcase) xySameStride = xStride[xNonUnitDim] == yStride[yNonUnitDim];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType xOffset, yOffset, zOffset;
|
||||
LongType *yCoords, *zCoords;
|
||||
|
||||
if (!is1Dcase) {
|
||||
yCoords = coords + threadIdx.x * (yRank + zRank);
|
||||
zCoords = yCoords + yRank;
|
||||
}
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < yLen; i += gridDim.x * blockDim.x) {
|
||||
if (is1Dcase) {
|
||||
yOffset = i * yStride[yNonUnitDim];
|
||||
zOffset = x[xySameStride ? yOffset : i * xStride[xNonUnitDim]] * zStride[zNonUnitDim];
|
||||
} else {
|
||||
INDEX2COORDS(i, yRank, yShape, yCoords);
|
||||
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
COORDS2INDEX(xRank, xStride, yCoords, xOffset);
|
||||
|
||||
zCoords[0] = x[xOffset];
|
||||
|
||||
for (LongType j = 0; j < yRank - xRank; ++j) {
|
||||
zCoords[j + 1] = yCoords[xRank + j];
|
||||
}
|
||||
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
}
|
||||
|
||||
switch (opCode) {
|
||||
case pairwise::Add:
|
||||
z[zOffset] += y[yOffset];
|
||||
break;
|
||||
case pairwise::Subtract:
|
||||
z[zOffset] -= y[yOffset];
|
||||
break;
|
||||
case pairwise::Multiply:
|
||||
z[zOffset] *= y[yOffset];
|
||||
break;
|
||||
case pairwise::Divide:
|
||||
z[zOffset] /= y[yOffset];
|
||||
break;
|
||||
case pairwise::ReverseSubtract:
|
||||
z[zOffset] = y[yOffset] - z[zOffset];
|
||||
break;
|
||||
case pairwise::ReverseDivide:
|
||||
z[zOffset] = y[yOffset] / z[zOffset];
|
||||
break;
|
||||
case pairwise::CopyPws:
|
||||
z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MaxPairwise:
|
||||
if (z[zOffset] < y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MinPairwise:
|
||||
if (z[zOffset] > y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
default:
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void scatterCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const int opCode, const void *vx,
|
||||
const LongType *xShapeInfo, const void *vy, const LongType *yShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const bool lock) {
|
||||
if (lock)
|
||||
scatterLockCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(opCode, vx, xShapeInfo, vy,
|
||||
yShapeInfo, vz, zShapeInfo);
|
||||
else
|
||||
scatterCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(opCode, vx, xShapeInfo, vy, yShapeInfo,
|
||||
vz, zShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "scatterLockCuda failed");
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void scatter(LaunchContext *context, pairwise::Ops op, NDArray&indices, NDArray&updates, NDArray &output,
|
||||
const bool lock) {
|
||||
const auto xType = indices.dataType();
|
||||
const auto yType = updates.dataType();
|
||||
|
||||
dim3 launchDims = scatterDims(lock ? output.lengthOf() : updates.lengthOf(), updates.rankOf() + output.rankOf());
|
||||
PointersManager manager(context, "scatter");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&updates, &indices});
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, scatterCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), op,
|
||||
indices.specialBuffer(), indices.specialShapeInfo(), updates.specialBuffer(),
|
||||
updates.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), lock),
|
||||
SD_INDEXING_TYPES, SD_GENERIC_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&updates, &indices});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, y - updates, z - output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void scatterNDLockCuda(const int opCode, const void *vx, const LongType *xShapeInfo, const void *vy,
|
||||
const LongType *yShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
const auto x = reinterpret_cast<const X *>(vx);
|
||||
const auto y = reinterpret_cast<const Y *>(vy);
|
||||
auto z = reinterpret_cast<Y *>(vz);
|
||||
|
||||
__shared__ LongType xRank, yRank, zRank, biggerXYRank, xLastDim, xNonUnitDim, yNonUnitDim, zNonUnitDim;
|
||||
__shared__ const LongType *xShape, *yShape, *zShape, *xStride, *yStride, *zStride;
|
||||
__shared__ LongType zLen, len;
|
||||
__shared__ bool is1Dcase;
|
||||
__shared__ LongType *coords;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
coords = reinterpret_cast<LongType *>(shmem);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
xLastDim = shape::sizeAt(xShapeInfo, -1);
|
||||
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
|
||||
biggerXYRank = xRank > yRank ? xRank : yRank;
|
||||
|
||||
xNonUnitDim = yNonUnitDim = zNonUnitDim = 0;
|
||||
|
||||
is1Dcase = (shape::isCommonVector(zShapeInfo, zNonUnitDim) || shape::isScalar(zShapeInfo)) &&
|
||||
(shape::isCommonVector(yShapeInfo, yNonUnitDim) || shape::isScalar(yShapeInfo)) &&
|
||||
(shape::isCommonVector(xShapeInfo, xNonUnitDim) || shape::isScalar(xShapeInfo));
|
||||
|
||||
len = is1Dcase ? shape::length(xShapeInfo) : shape::length(xShapeInfo) / xLastDim;
|
||||
zLen = shape::length(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType yOffset, zOffset, xOffset;
|
||||
LongType *yCoords, *zCoords;
|
||||
|
||||
if (!is1Dcase) {
|
||||
yCoords = coords + threadIdx.x * (biggerXYRank + zRank);
|
||||
zCoords = yCoords + biggerXYRank;
|
||||
}
|
||||
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < zLen; i += gridDim.x * blockDim.x) {
|
||||
if (!is1Dcase) INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
|
||||
for (LongType j = 0; j < len; j++) {
|
||||
if (is1Dcase) {
|
||||
if (x[j * xStride[xNonUnitDim]] != i) continue;
|
||||
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
} else {
|
||||
INDEX2COORDS(j, xRank - 1, xShape, yCoords);
|
||||
|
||||
yCoords[xRank - 1] = 0;
|
||||
COORDS2INDEX(xRank, xStride, yCoords, xOffset);
|
||||
if (zCoords[0] != x[xOffset]) continue;
|
||||
|
||||
bool matched = true;
|
||||
for (LongType k = 1; k < xLastDim; k++) {
|
||||
yCoords[xRank - 1] = k;
|
||||
COORDS2INDEX(xRank, xStride, yCoords, xOffset);
|
||||
if (zCoords[k] != x[xOffset]) {
|
||||
matched = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!matched) continue;
|
||||
|
||||
for (LongType k = xLastDim; k < zRank; ++k) yCoords[yRank - zRank + k] = zCoords[k];
|
||||
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
}
|
||||
|
||||
switch (opCode) {
|
||||
case pairwise::Add:
|
||||
z[zOffset] += y[yOffset];
|
||||
break;
|
||||
case pairwise::Subtract:
|
||||
z[zOffset] -= y[yOffset];
|
||||
break;
|
||||
case pairwise::Multiply:
|
||||
z[zOffset] *= y[yOffset];
|
||||
break;
|
||||
case pairwise::Divide:
|
||||
z[zOffset] /= y[yOffset];
|
||||
break;
|
||||
case pairwise::ReverseSubtract:
|
||||
z[zOffset] = y[yOffset] - z[zOffset];
|
||||
break;
|
||||
case pairwise::ReverseDivide:
|
||||
z[zOffset] = y[yOffset] / z[zOffset];
|
||||
break;
|
||||
case pairwise::CopyPws:
|
||||
z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MaxPairwise:
|
||||
if (z[zOffset] < y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MinPairwise:
|
||||
if (z[zOffset] > y[yOffset]) z[zOffset] = y[yOffset];
|
||||
break;
|
||||
default:
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// x - indices, y - updates, z - output
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL static void scatterNDCuda(const int opCode, const void* vx, const LongType* xShapeInfo, const void* vy,
|
||||
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo) {
|
||||
// Cast input and output pointers
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
const auto y = reinterpret_cast<const Y*>(vy);
|
||||
auto z = reinterpret_cast<Y*>(vz);
|
||||
|
||||
// Shared memory for shape information and flags
|
||||
__shared__ LongType xRank, yRank, zRank, biggerXYRank, xLastDim, xNonUnitDim, yNonUnitDim, zNonUnitDim, yLen;
|
||||
__shared__ bool is1Dcase;
|
||||
|
||||
// Shared memory for coordinates
|
||||
__shared__ LongType* coords;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Dynamically allocated shared memory
|
||||
extern __shared__ unsigned char shmem[];
|
||||
coords = reinterpret_cast<LongType*>(shmem);
|
||||
|
||||
// Initialize shared values
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
xLastDim = shape::sizeAt(xShapeInfo, -1);
|
||||
yLen = shape::length(yShapeInfo);
|
||||
|
||||
biggerXYRank = max(xRank, yRank);
|
||||
|
||||
xNonUnitDim = yNonUnitDim = zNonUnitDim = 0;
|
||||
|
||||
// Check if the operation involves 1D cases
|
||||
is1Dcase = (shape::isCommonVector(zShapeInfo, zNonUnitDim) || shape::isScalar(zShapeInfo)) &&
|
||||
(shape::isCommonVector(yShapeInfo, yNonUnitDim) || shape::isScalar(yShapeInfo)) &&
|
||||
(shape::isCommonVector(xShapeInfo, xNonUnitDim) || shape::isScalar(xShapeInfo));
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Dynamically allocated memory for local coordinates
|
||||
LongType* yCoords = coords + threadIdx.x * (biggerXYRank + zRank);
|
||||
LongType* zCoords = yCoords + biggerXYRank;
|
||||
|
||||
// Process each element in y
|
||||
for (LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < yLen; i += gridDim.x * blockDim.x) {
|
||||
LongType yOffset, zOffset;
|
||||
|
||||
// Convert linear index to multi-dimensional coordinates for y
|
||||
INDEX2COORDS(i, yRank, shape::shapeOf(yShapeInfo), yCoords);
|
||||
COORDS2INDEX(yRank, shape::stride(yShapeInfo), yCoords, yOffset);
|
||||
|
||||
// Save the last coordinate of y if needed
|
||||
if (yRank >= xRank) {
|
||||
zCoords[xLastDim] = yCoords[xRank - 1];
|
||||
}
|
||||
|
||||
// Map y coordinates to x and z coordinates
|
||||
for (LongType j = 0; j < xLastDim; ++j) {
|
||||
yCoords[xRank - 1] = j;
|
||||
COORDS2INDEX(xRank, shape::stride(xShapeInfo), yCoords, zCoords[j]);
|
||||
}
|
||||
|
||||
// Adjust remaining coordinates for z
|
||||
for (LongType j = xLastDim + 1; j < zRank; ++j) {
|
||||
zCoords[j] = yCoords[yRank - zRank + j];
|
||||
}
|
||||
|
||||
// Compute linear index for z
|
||||
COORDS2INDEX(zRank, shape::stride(zShapeInfo), zCoords, zOffset);
|
||||
|
||||
// Perform the operation based on opCode
|
||||
switch (opCode) {
|
||||
case pairwise::Add:
|
||||
z[zOffset] += y[yOffset];
|
||||
break;
|
||||
case pairwise::Subtract:
|
||||
z[zOffset] -= y[yOffset];
|
||||
break;
|
||||
case pairwise::Multiply:
|
||||
z[zOffset] *= y[yOffset];
|
||||
break;
|
||||
case pairwise::Divide:
|
||||
z[zOffset] /= y[yOffset];
|
||||
break;
|
||||
case pairwise::ReverseSubtract:
|
||||
z[zOffset] = y[yOffset] - z[zOffset];
|
||||
break;
|
||||
case pairwise::ReverseDivide:
|
||||
z[zOffset] = y[yOffset] / z[zOffset];
|
||||
break;
|
||||
case pairwise::CopyPws:
|
||||
z[zOffset] = y[yOffset];
|
||||
break;
|
||||
case pairwise::MaxPairwise:
|
||||
z[zOffset] = max(z[zOffset], y[yOffset]);
|
||||
break;
|
||||
case pairwise::MinPairwise:
|
||||
z[zOffset] = min(z[zOffset], y[yOffset]);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
static void scatterNDCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const int opCode, const void *vx,
|
||||
const LongType *xShapeInfo, const void *vy, const LongType *yShapeInfo, void *vz,
|
||||
const LongType *zShapeInfo, const bool lock) {
|
||||
if (lock)
|
||||
scatterNDLockCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(opCode, vx, xShapeInfo, vy,
|
||||
yShapeInfo, vz, zShapeInfo);
|
||||
else
|
||||
scatterNDCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(opCode, vx, xShapeInfo, vy, yShapeInfo,
|
||||
vz, zShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "scatterNDCuda failed");
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void scatterND(LaunchContext *context, pairwise::Ops op, NDArray&indices, NDArray&updates,
|
||||
NDArray &output, const bool lock) {
|
||||
const int xRank = indices.rankOf();
|
||||
const int yRank = updates.rankOf();
|
||||
const int zRank = output.rankOf();
|
||||
|
||||
dim3 launchDims =
|
||||
scatterNdDims(lock ? output.lengthOf() : updates.lengthOf(), ((yRank > xRank ? yRank : xRank) + zRank));
|
||||
const auto xType = indices.dataType();
|
||||
const auto yType = updates.dataType();
|
||||
|
||||
PointersManager manager(context, "scatterND");
|
||||
|
||||
NDArray::prepareSpecialUse({&output}, {&updates, &indices});
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, scatterNDCudaLauncher,
|
||||
(launchDims.y, launchDims.x, launchDims.z, context->getCudaStream(), op,
|
||||
indices.specialBuffer(), indices.specialShapeInfo(), updates.specialBuffer(),
|
||||
updates.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), lock),
|
||||
SD_INDEXING_TYPES, SD_GENERIC_NUMERIC_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&updates, &indices});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL void scatterForLossCuda(const void* vx, const LongType* xShapeInfo, void* vy, const LongType* yShapeInfo,
|
||||
void* vz, const LongType* zShapeInfo) {
|
||||
// Cast input and output pointers
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<Z*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
// Shared memory for shape information and coordinates
|
||||
__shared__ LongType xLen;
|
||||
__shared__ LongType xRank;
|
||||
__shared__ const LongType* xShape;
|
||||
__shared__ const LongType* xStride;
|
||||
__shared__ const LongType* yStride;
|
||||
__shared__ const LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Initialize shared memory variables
|
||||
xLen = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
zStride = zShapeInfo ? shape::stride(zShapeInfo) : nullptr;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index
|
||||
const LongType xInd = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// Return if the thread index exceeds the length of x
|
||||
if (xInd >= xLen) return;
|
||||
|
||||
// Dynamically allocated shared memory for coordinates
|
||||
extern __shared__ unsigned char shmem[];
|
||||
auto coords = reinterpret_cast<LongType*>(shmem) + threadIdx.x * (xRank + 1);
|
||||
|
||||
// Convert linear index to coordinates for x
|
||||
INDEX2COORDS(xInd, xRank, xShape, coords);
|
||||
|
||||
// Calculate offset for x
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(xRank, xStride, coords, xOffset);
|
||||
|
||||
// Update the last coordinate with the value from x
|
||||
coords[xRank] = x[xOffset];
|
||||
|
||||
// Calculate offset for y
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(xRank + 1, yStride, coords, yOffset);
|
||||
|
||||
if (z == nullptr) {
|
||||
// Gradient calculation
|
||||
y[yOffset] -= 1.f;
|
||||
} else {
|
||||
// Calculate offset for z
|
||||
LongType zOffset;
|
||||
COORDS2INDEX(xRank + 1, zStride, coords, zOffset);
|
||||
|
||||
// Update z with the value from y
|
||||
z[zOffset] = y[yOffset];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
static void scatterForLossCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
||||
const cudaStream_t *stream, const void *vx, const LongType *xShapeInfo, void *vy,
|
||||
const LongType *yShapeInfo, void *vz, const LongType *zShapeInfo) {
|
||||
scatterForLossCuda<X, Z>
|
||||
<<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "scatterUpdateCuda failed");
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
void scatterForLoss(LaunchContext *context, NDArray&indices, NDArray &updates, NDArray &output,
|
||||
const bool calcGrad) {
|
||||
// shapes of indices and output must be the same
|
||||
// shape of indices should be the same as updates shape with last dimension excluded, for example if updates is
|
||||
// {a,b,c} then indices should be {a,b}
|
||||
|
||||
PointersManager manager(context, "scatterForLoss");
|
||||
|
||||
dim3 launchDIms = scatterDims(indices.lengthOf(), updates.rankOf());
|
||||
if (calcGrad) {
|
||||
NDArray::prepareSpecialUse({&updates}, {&indices});
|
||||
BUILD_DOUBLE_SELECTOR(
|
||||
indices.dataType(), updates.dataType(), scatterForLossCudaLauncher,
|
||||
(launchDIms.y, launchDIms.x, launchDIms.z, context->getCudaStream(), indices.specialBuffer(),
|
||||
indices.specialShapeInfo(), updates.specialBuffer(), updates.specialShapeInfo(), nullptr, nullptr),
|
||||
SD_INDEXING_TYPES, SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&updates}, {&indices});
|
||||
} else {
|
||||
NDArray::prepareSpecialUse({&output}, {&indices, &updates});
|
||||
BUILD_DOUBLE_SELECTOR(indices.dataType(), updates.dataType(), scatterForLossCudaLauncher,
|
||||
(launchDIms.y, launchDIms.x, launchDIms.z, context->getCudaStream(), indices.specialBuffer(),
|
||||
indices.specialShapeInfo(), updates.specialBuffer(), updates.specialShapeInfo(),
|
||||
output.specialBuffer(), output.specialShapeInfo()),
|
||||
SD_INDEXING_TYPES, SD_FLOAT_TYPES);
|
||||
NDArray::registerSpecialUse({&output}, {&indices, &updates});
|
||||
}
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,143 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
template <typename X, typename Y>
|
||||
static SD_KERNEL void scatterSimpleKernel(void* vx, const LongType* xTadShape, const LongType* xTadOffsets,
|
||||
LongType xLength, LongType numTads, const void* vi,
|
||||
const LongType* iShapeInfo, LongType iLength, const void* vu,
|
||||
const LongType* uShapeInfo, LongType uLength) {
|
||||
// Shared memory caching for shape and stride information
|
||||
__shared__ LongType iRank, xTadRank, uRank;
|
||||
__shared__ const LongType* iShape;
|
||||
__shared__ const LongType* xTadShapePtr;
|
||||
__shared__ const LongType* uShape;
|
||||
__shared__ const LongType* iStride;
|
||||
__shared__ const LongType* xTadStride;
|
||||
__shared__ const LongType* uStride;
|
||||
|
||||
// Initialize shared memory
|
||||
if (threadIdx.x == 0) {
|
||||
iRank = shape::rank(iShapeInfo);
|
||||
xTadRank = shape::rank(xTadShape);
|
||||
uRank = shape::rank(uShapeInfo);
|
||||
|
||||
iShape = shape::shapeOf(iShapeInfo);
|
||||
xTadShapePtr = shape::shapeOf(xTadShape);
|
||||
uShape = shape::shapeOf(uShapeInfo);
|
||||
|
||||
iStride = shape::stride(iShapeInfo);
|
||||
xTadStride = shape::stride(xTadShape);
|
||||
uStride = shape::stride(uShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Cast input pointers
|
||||
const X* u = reinterpret_cast<const X*>(vu);
|
||||
const Y* indices = reinterpret_cast<const Y*>(vi);
|
||||
|
||||
// Calculate thread ID
|
||||
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// Iterate over the indices
|
||||
for (int i = tid; i < iLength; i += blockDim.x * gridDim.x) {
|
||||
// Offset for `x`
|
||||
auto x = reinterpret_cast<X*>(vx) + xTadOffsets[i];
|
||||
|
||||
// Compute coordinates and offsets for index tensor
|
||||
LongType idxCoords[SD_MAX_RANK];
|
||||
LongType idxOffset;
|
||||
INDEX2COORDS(i, iRank, iShape, idxCoords);
|
||||
COORDS2INDEX(iRank, iStride, idxCoords, idxOffset);
|
||||
auto idx = indices[idxOffset];
|
||||
|
||||
// Compute coordinates and offsets for x
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
INDEX2COORDS(idx, xTadRank, xTadShapePtr, xCoords);
|
||||
COORDS2INDEX(xTadRank, xTadStride, xCoords, xOffset);
|
||||
|
||||
// Compute coordinates and offsets for u
|
||||
LongType uCoords[SD_MAX_RANK];
|
||||
LongType uOffset;
|
||||
INDEX2COORDS(i, uRank, uShape, uCoords);
|
||||
COORDS2INDEX(uRank, uStride, uCoords, uOffset);
|
||||
|
||||
// Perform the scatter update
|
||||
x[xOffset] = u[uOffset];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename X, typename Y>
|
||||
void scatterSimple_(LaunchContext* context, const int opId, NDArray& input, NDArray& updates,
|
||||
NDArray& indices, const std::vector<LongType>& dimensions) {
|
||||
auto dims = ShapeUtils::evalDimsToExclude(input.rankOf(),dimensions.size(),dimensions.data());
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), dims);
|
||||
|
||||
auto xLength = shape::length(packX->primaryShapeInfo());
|
||||
auto iLength = indices.lengthOf();
|
||||
auto uLength = updates.lengthOf();
|
||||
|
||||
dim3 launchDims = getLaunchDims("scatter_simple");
|
||||
scatterSimpleKernel<X, Y><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
||||
input.specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), xLength, packX->numberOfTads(),
|
||||
indices.specialBuffer(), indices.specialShapeInfo(), iLength, updates.specialBuffer(), updates.specialShapeInfo(),
|
||||
uLength);
|
||||
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "scatterUpdateCuda failed");
|
||||
|
||||
|
||||
}
|
||||
|
||||
void scatterSimple(LaunchContext* context, const int opId, NDArray& input, NDArray& updates,
|
||||
NDArray& indices, const std::vector<LongType>& dimensions) {
|
||||
auto xType = input.dataType();
|
||||
auto yType = indices.dataType();
|
||||
|
||||
if (opId != 6) THROW_EXCEPTION("scatterSimple: only copy op is supported");
|
||||
|
||||
NDArray::prepareSpecialUse({&input}, {&updates, &indices});
|
||||
|
||||
BUILD_DOUBLE_SELECTOR(xType, yType, scatterSimple_, (context, opId, input, updates, indices, dimensions),
|
||||
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
||||
|
||||
NDArray::registerSpecialUse({&input}, {&updates, &indices});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,169 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <array/NDArrayFactory.h>
|
||||
#include <array/ResultSet.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <helpers/ConstantTadHelper.h>
|
||||
#include <helpers/PointersManager.h>
|
||||
#include <helpers/ShapeUtils.h>
|
||||
|
||||
#include <ops/declarable/helpers/transforms.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
///////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void scatterUpdateCuda(const int opCode, const int numOfInd, void* vx, const LongType* xShapeInfo,
|
||||
const LongType* xOffsets, void* vy, const LongType* yShapeInfo,
|
||||
const LongType* yOffsets, const LongType* indexes) {
|
||||
// Shared memory caching for shape and pointers
|
||||
__shared__ T *x, *y;
|
||||
__shared__ LongType arrLenX, arrLenY;
|
||||
__shared__ LongType xRank, yRank;
|
||||
__shared__ const LongType* xShape;
|
||||
__shared__ const LongType* yShape;
|
||||
__shared__ const LongType* xStride;
|
||||
__shared__ const LongType* yStride;
|
||||
|
||||
// Initialize shared variables
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
arrLenX = shape::length(xShapeInfo);
|
||||
arrLenY = shape::length(yShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Iterate through the number of indices
|
||||
for (int e = 0; e < numOfInd; e++) {
|
||||
const auto xIndex = indexes[e];
|
||||
const bool isOwner = xIndex < gridDim.x ? blockIdx.x == xIndex : blockIdx.x == xIndex % gridDim.x;
|
||||
|
||||
if (!isOwner) continue;
|
||||
|
||||
// Initialize x and y pointers
|
||||
if (threadIdx.x == 0) {
|
||||
x = reinterpret_cast<T*>(vx) + xOffsets[xIndex];
|
||||
y = reinterpret_cast<T*>(vy) + yOffsets[e];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Validate array lengths
|
||||
if (arrLenX != arrLenY) return;
|
||||
|
||||
// Process the elements
|
||||
for (LongType i = threadIdx.x; i < arrLenX; i += blockDim.x) {
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType xOffset, yOffset;
|
||||
|
||||
// Compute coordinates and offsets for x and y
|
||||
INDEX2COORDS(i, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
INDEX2COORDS(i, yRank, yShape, yCoords);
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
|
||||
// Perform the specified operation
|
||||
switch (opCode) {
|
||||
case 0:
|
||||
x[xOffset] += y[yOffset];
|
||||
break;
|
||||
case 1:
|
||||
x[xOffset] -= y[yOffset];
|
||||
break;
|
||||
case 2:
|
||||
x[xOffset] *= y[yOffset];
|
||||
break;
|
||||
case 3:
|
||||
x[xOffset] /= y[yOffset];
|
||||
break;
|
||||
case 4:
|
||||
x[xOffset] = y[yOffset] - x[xOffset];
|
||||
break;
|
||||
case 5:
|
||||
x[xOffset] = y[yOffset] / x[xOffset];
|
||||
break;
|
||||
case 6:
|
||||
x[xOffset] = y[yOffset];
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
SD_HOST static void scatterUpdateCudaLauncher(const cudaStream_t* stream, const int opCode, const int numOfInd,
|
||||
void* vx, const LongType* xShapeInfo, const LongType* xOffsets,
|
||||
void* vy, const LongType* yShapeInfo, const LongType* yOffsets,
|
||||
const LongType* indexes) {
|
||||
dim3 launchDims = getLaunchDims("scatter_update");
|
||||
scatterUpdateCuda<T><<<launchDims.y, launchDims.x, SD_MAX_NUM_THREADS, *stream>>>(opCode, numOfInd, vx, xShapeInfo, xOffsets, vy,
|
||||
yShapeInfo, yOffsets, indexes);
|
||||
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "scatterUpdateCuda failed");
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void scatterUpdate(LaunchContext* context, NDArray& input, NDArray& updates, const std::vector<LongType>* intArgs) {
|
||||
const int opCode = (*intArgs)[0];
|
||||
const int numOfDims = (*intArgs)[1];
|
||||
const int numOfInd = (*intArgs)[2 + numOfDims];
|
||||
|
||||
std::vector<LongType> tadDimensions(numOfDims);
|
||||
for (int e = 2; e < 2 + numOfDims; e++) tadDimensions[e - 2] = (*intArgs)[e];
|
||||
|
||||
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input.shapeInfo(), &tadDimensions);
|
||||
auto packY = ConstantTadHelper::getInstance().tadForDimensions(updates.shapeInfo(), &tadDimensions);
|
||||
|
||||
std::vector<LongType> shape = {numOfInd};
|
||||
NDArray indices(const_cast<LongType*>(intArgs->data()) + numOfDims + 3, 'c', shape, INT32, context);
|
||||
|
||||
PointersManager manager(context, "scatterUpdate");
|
||||
|
||||
NDArray::prepareSpecialUse({&input}, {&input, &updates, &indices});
|
||||
BUILD_SINGLE_SELECTOR(input.dataType(), scatterUpdateCudaLauncher,
|
||||
(context->getCudaStream(), opCode, numOfInd, input.specialBuffer(), packX->platformShapeInfo(),
|
||||
packX->platformOffsets(), updates.specialBuffer(), packY->platformShapeInfo(),
|
||||
packY->platformOffsets(), reinterpret_cast<sd::LongType *>(indices.specialBuffer())),
|
||||
SD_COMMON_TYPES);
|
||||
NDArray::registerSpecialUse({&input}, {&input, &updates, &indices});
|
||||
|
||||
manager.synchronize();
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,154 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 {
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Sorted segments ops implementations
|
||||
|
||||
template <typename T, typename I>
|
||||
static bool segmentIndicesValidate_(NDArray* indices, NDArray& aexpected, NDArray& aoutput) {
|
||||
return true;
|
||||
}
|
||||
|
||||
bool segmentIndicesValidate(LaunchContext* context, NDArray* indices, NDArray& expected, NDArray& output) {
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output.dataType();
|
||||
BUILD_DOUBLE_SELECTOR(output.dataType(), indices->dataType(), return segmentIndicesValidate_,
|
||||
(indices, expected, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Unsorted segment ops functors implementation
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename I>
|
||||
static SD_KERNEL void unsortedSegmentIndexValidateKernel(const I* indices, const LongType* indicesShape, I expected,
|
||||
I* found) {
|
||||
__shared__ bool onlyTrue;
|
||||
__shared__ LongType len;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
onlyTrue = true;
|
||||
len = shape::length(indicesShape);
|
||||
}
|
||||
__syncthreads();
|
||||
auto start = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
for (LongType e = start; e < len && onlyTrue; e += step) {
|
||||
math::atomics::sd_atomicMax(found, indices[e]);
|
||||
if (expected < *found) onlyTrue = false;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename I>
|
||||
static bool unsortedSegmentIndicesValidate_(LaunchContext* context, NDArray* indices, LongType expected,
|
||||
LongType& output) {
|
||||
output = expected;
|
||||
I found = output;
|
||||
I exp = expected;
|
||||
auto stream = context->getCudaStream();
|
||||
I* devFound;
|
||||
cudaMalloc(&devFound, sizeof(I));
|
||||
cudaMemcpy(devFound, &found, sizeof(I), cudaMemcpyHostToDevice);
|
||||
|
||||
dim3 launchDims = segmentValidateIndices(indices->lengthOf());
|
||||
unsortedSegmentIndexValidateKernel<I><<<launchDims.y,launchDims.x, launchDims.z, *stream>>>(
|
||||
reinterpret_cast<I*>(indices->specialBuffer()), indices->specialShapeInfo(), exp, devFound);
|
||||
sd::DebugHelper::checkErrorCode(stream, "unsortedSegmentIndexValidateKernel failed");
|
||||
|
||||
cudaMemcpy(&found, devFound, sizeof(I), cudaMemcpyDeviceToHost);
|
||||
cudaFree(devFound);
|
||||
output = found;
|
||||
return expected == output;
|
||||
}
|
||||
|
||||
bool unsortedSegmentIndicesValidate(LaunchContext* context, NDArray* indices, LongType expected, LongType& output) {
|
||||
BUILD_SINGLE_SELECTOR(indices->dataType(), return unsortedSegmentIndicesValidate_,
|
||||
(context, indices, expected, output), SD_INDEXING_TYPES);
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// fill up segments starts and ends - splitted ordered case
|
||||
template <typename I>
|
||||
static SD_KERNEL void fillUpSegmentsKernel(const void* indices, const LongType* indexShape, LongType numClasses,
|
||||
LongType* classesRangesStart, LongType* classesRangesLengths) {
|
||||
__shared__ const I* idxBuf;
|
||||
__shared__ LongType idxLen;
|
||||
__shared__ LongType* result;
|
||||
if (threadIdx.x == 0) {
|
||||
idxBuf = reinterpret_cast<const I*>(indices);
|
||||
idxLen = shape::length(indexShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto tid = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
for (auto j = tid; j < idxLen; j += step) {
|
||||
auto pos = idxBuf[j];
|
||||
math::atomics::sd_atomicMin<LongType>(&classesRangesStart[pos], (LongType)j);
|
||||
math::atomics::sd_atomicAdd<LongType>(&classesRangesLengths[pos], 1);
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename I>
|
||||
static void fillUpSegments_(NDArray* indices, LongType numClasses, NDArray& classesRangesBegs,
|
||||
NDArray& classesRangesLens) {
|
||||
dim3 dims = getFillUpSegmentsDims(numClasses, indices->lengthOf());
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
auto stream = classesRangesBegs.getContext()->getCudaStream();
|
||||
fillUpSegmentsKernel<I><<<dims.x, dims.y, dims.z, *stream>>>(indices->specialBuffer(), indices->specialShapeInfo(),
|
||||
numClasses, begins, lengths);
|
||||
sd::DebugHelper::checkErrorCode(stream, "fillUpSegmentsKernel failed");
|
||||
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
void fillUpSegments(NDArray* indices, LongType numClasses, NDArray& classesRangesBegs, NDArray& classesRangesLens) {
|
||||
BUILD_SINGLE_SELECTOR(indices->dataType(), fillUpSegments_,
|
||||
(indices, numClasses, classesRangesBegs, classesRangesLens), SD_INDEXING_TYPES);
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
@@ -0,0 +1,677 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
@@ -0,0 +1,632 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 "helpers/DebugHelper.h"
|
||||
#include <system/selective_rendering.h>
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Segment ops linear kernels
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMeanLinearKernel(void* input, LongType const* inputShape, LongType* indices,
|
||||
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* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__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);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
if (segment < numOfClasses) {
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType zOffset;
|
||||
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, zIndex);
|
||||
start = indices[segment];
|
||||
finish = start + lengths[segment];
|
||||
INDEX2COORDS(start, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, xOffset);
|
||||
if (lengths[segment] > 0)
|
||||
z[zIndex] = T(x[xOffset] / T(lengths[segment]));
|
||||
else
|
||||
z[zIndex] = 0;
|
||||
val[segment] = z[zIndex];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto e = start + threadIdx.x + 1; e < finish; e += blockDim.x) {
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, xOffset);
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], T(x[xOffset] / static_cast<T>(lengths[segment])));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void unsortedSegmentMeanLinearKernel(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* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__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);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType zOffset;
|
||||
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, zIndex);
|
||||
INDEX2COORDS(starts[segment], inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, xOffset);
|
||||
|
||||
if (lengths[segment] > 0)
|
||||
z[zIndex] = T(x[xOffset] / T(lengths[segment]));
|
||||
else
|
||||
z[zIndex] = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (lengths[segment] > 0) {
|
||||
for (auto e = threadIdx.x; e < xLen; e += blockDim.x) {
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType yIndex;
|
||||
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, xOffset);
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, inputCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, inputCoords, yIndex);
|
||||
|
||||
if (y[yIndex] == segment && e != starts[segment]) {
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], T(x[xOffset] / T(lengths[segment])));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMeanTadKernel(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, LongType indicesLen) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType len, zIndex, total;
|
||||
__shared__ T* z;
|
||||
__shared__ int threadsPerSegment, start, finish;
|
||||
|
||||
// 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(blockIdx.x >= indicesLen)
|
||||
return;
|
||||
|
||||
auto segment = indices[blockIdx.x];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z = reinterpret_cast<T*>(outputBuf) + outputTadOffsets[segment];
|
||||
len = shape::length(inputTads);
|
||||
start = starts[segment];
|
||||
finish = start + lengths[segment];
|
||||
total = shape::sizeAt(inputShape, 0);
|
||||
|
||||
// Cache TAD 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);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto idx = blockIdx.x;
|
||||
if (blockIdx.x <= total) {
|
||||
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_atomicAdd(&z[zIndex], T(x[xIndex] / lengths[segment]));
|
||||
}
|
||||
} 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_atomicAdd(&z[zIndex], T(x[xIndex] / lengths[segment]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// segment mean
|
||||
template <typename T, typename I>
|
||||
static void segmentMeanFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
LongType numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
|
||||
int zero2 = 0;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero2);
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
dim3 launchDims = segmentDims(numClasses,input->lengthOf());
|
||||
segmentMeanLinearKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanLinearKernel 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(input->sizeAt(0));
|
||||
segmentMeanTadKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
||||
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo(), outputTads, outputTadOffsets,indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void segmentMeanFunctor(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
UILD_DOUBLE_SELECTOR(output->dataType(), indices->dataType(), segmentMeanFunctor_, (context, input, indices, output),
|
||||
SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static void unsortedSegmentMeanFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
|
||||
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()) {
|
||||
unsortedSegmentMeanLinearKernel<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, "unsortedSegmentMeanLinearKernel failed");
|
||||
|
||||
} else {
|
||||
LongType zero = 0;
|
||||
output->assign(zero);
|
||||
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);
|
||||
LongType const* inputTads = packX->specialShapeInfo();
|
||||
LongType const* inputTadOffsets = packX->specialOffsets();
|
||||
LongType const* outputTads = packZ->specialShapeInfo();
|
||||
LongType const* outputTadOffsets = packZ->specialOffsets();
|
||||
dims.x = input->sizeAt(0);
|
||||
segmentMeanTadKernel<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, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void unsortedSegmentMeanFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses,
|
||||
NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
auto indicesDType = indices->dataType();
|
||||
auto inputDType = input->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentMeanFunctor_,
|
||||
(context, input, indices, numOfClasses, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMeanBPLinearKernel(void* inputBuf, LongType const* inputShape, void* eps,
|
||||
LongType const* epsShape, void* indicesBuf,
|
||||
LongType const* indicesShape, LongType* lengths, void* outputBuf,
|
||||
LongType const* outputShape) {
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, gradLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank, epsRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* epsShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
__shared__ const sd::LongType* epsStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
y = reinterpret_cast<I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
gradOut = reinterpret_cast<T*>(eps);
|
||||
gradLen = shape::length(epsShape);
|
||||
|
||||
// Cache all shape information
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
epsRank = shape::rank(epsShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (auto e = start; e < xLen; e += step) {
|
||||
LongType zOffset, xOffset, yOffset, gradOffsetO;
|
||||
sd::LongType zCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK], yCoords[SD_MAX_RANK], gradCoords[SD_MAX_RANK];
|
||||
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yOffset);
|
||||
|
||||
auto classIndex = y[yOffset];
|
||||
|
||||
INDEX2COORDS(classIndex, epsRank, epsShapePtr, gradCoords);
|
||||
COORDS2INDEX(epsRank, epsStridePtr, gradCoords, gradOffsetO);
|
||||
|
||||
z[zOffset] = T(gradOut[gradOffsetO] / float(lengths[classIndex]));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMeanBPTadKernel(void* inputBuf, LongType const* inputShape, void* eps,
|
||||
LongType const* epsShape, void* indicesBuf, LongType const* indicesShape,
|
||||
LongType* lengths, void* outputBuf, LongType const* outputShape,
|
||||
LongType const* inputTad, LongType const* inputOffsets,
|
||||
LongType const* gradOutTad, LongType const* gradOutOffsets,
|
||||
LongType const* outTad, LongType const* outOffsets) {
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, yLen, gradLen, currentLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType outTadRank, gradOutTadRank;
|
||||
__shared__ const sd::LongType* outTadShapePtr;
|
||||
__shared__ const sd::LongType* gradOutTadShapePtr;
|
||||
__shared__ const sd::LongType* outTadStridePtr;
|
||||
__shared__ const sd::LongType* gradOutTadStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
y = reinterpret_cast<I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
yLen = shape::length(indicesShape);
|
||||
gradOut = reinterpret_cast<T*>(eps);
|
||||
gradLen = shape::length(epsShape);
|
||||
currentLen = shape::length(outTad);
|
||||
|
||||
// Cache TAD shape information
|
||||
outTadRank = shape::rank(outTad);
|
||||
gradOutTadRank = shape::rank(gradOutTad);
|
||||
outTadShapePtr = shape::shapeOf(outTad);
|
||||
gradOutTadShapePtr = shape::shapeOf(gradOutTad);
|
||||
outTadStridePtr = shape::stride(outTad);
|
||||
gradOutTadStridePtr = shape::stride(gradOutTad);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto i = blockIdx.x; i < yLen; i += gridDim.x) {
|
||||
auto segment = y[i];
|
||||
T* currentOut = z + outOffsets[i];
|
||||
T* outGrad = gradOut + gradOutOffsets[segment];
|
||||
|
||||
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType gradCoords[SD_MAX_RANK];
|
||||
sd::LongType zIndex;
|
||||
sd::LongType gradIndex;
|
||||
|
||||
INDEX2COORDS(e, outTadRank, outTadShapePtr, zCoords);
|
||||
COORDS2INDEX(outTadRank, outTadStridePtr, zCoords, zIndex);
|
||||
INDEX2COORDS(e, gradOutTadRank, gradOutTadShapePtr, gradCoords);
|
||||
COORDS2INDEX(gradOutTadRank, gradOutTadStridePtr, gradCoords, gradIndex);
|
||||
|
||||
if (lengths[segment] > 0) currentOut[zIndex] = T(outGrad[gradIndex] / float(lengths[segment]));
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// backrop for mean
|
||||
template <typename T, typename I>
|
||||
Status segmentMeanFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
||||
NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
auto numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
sd::LongType zero2 = 0;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(zero2);
|
||||
classesRangesLens.assign(len);
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
LongType loop_size = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf(); // indices->e<sd::LongType>(loop_size - 1);
|
||||
dim3 segmentBpDims2 = segmentBpDims(gradOut->lengthOf(),input->lengthOf());
|
||||
segmentMeanBPLinearKernel<T, I><<<segmentBpDims2.y, segmentBpDims2.x, segmentBpDims2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanBPLinearKernel 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 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* gradOutTads = packGradOut->specialShapeInfo();
|
||||
LongType const* gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
dim3 segmentBpTad2 = segmentBpTad(indices->lengthOf(),input->lengthOf());
|
||||
|
||||
segmentMeanBPTadKernel<T, I><<<segmentBpTad2.y, segmentBpTad2.x, segmentBpTad2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo(), inputTads, inputTadOffsets, gradOutTads, gradOutTadOffsets, outputTads,
|
||||
outputTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// segmen mean bp main
|
||||
Status segmentMeanFunctorBP(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 segmentMeanFunctorBP_,
|
||||
(context, input, indices, gradOut, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static Status unsortedSegmentMeanFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices,
|
||||
NDArray* gradOut,
|
||||
LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
auto numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
|
||||
sd::LongType zero2 = 0;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(zero2);
|
||||
classesRangesLens.assign(len);
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
LongType loop_size = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf();
|
||||
dim3 segmentBpDims2 = segmentBpDims(gradOut->lengthOf(),input->lengthOf());
|
||||
segmentMeanBPLinearKernel<T, I><<<segmentBpDims2.y,segmentBpDims2.x,segmentBpDims2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanBPLinearKernel 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 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* gradOutTads = packGradOut->specialShapeInfo();
|
||||
LongType const* gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
dim3 segmentBpTad2 = segmentBpTad(indices->lengthOf(),input->lengthOf());
|
||||
|
||||
segmentMeanBPTadKernel<T, I><<<segmentBpTad2.y,segmentBpTad2.x, segmentBpTad2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo(), inputTads, inputTadOffsets, gradOutTads, gradOutTadOffsets, outputTads,
|
||||
outputTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMeanBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
Status unsortedSegmentMeanFunctorBP(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 unsortedSegmentMeanFunctorBP_,
|
||||
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,620 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 "helpers/DebugHelper.h"
|
||||
#include <system/selective_rendering.h>
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Segment ops linear kernels
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMinLinearKernel(const void* input, const LongType* inputShape, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, void* output,
|
||||
const LongType* outputShape) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType xLen, zLen, zIndex;
|
||||
__shared__ const T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType threadsPerSegment, start, finish;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
|
||||
auto segment = blockIdx.x;
|
||||
if(blockIdx.x >= numOfClasses)
|
||||
return;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
x = reinterpret_cast<const 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);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
if (segment < numOfClasses) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zIndex);
|
||||
if(zIndex >= zLen)
|
||||
return;
|
||||
start = starts[segment];
|
||||
finish = start + lengths[segment];
|
||||
LongType startCoords[SD_MAX_RANK];
|
||||
LongType startIndex;
|
||||
INDEX2COORDS(start, inputRank, inputShapePtr, startCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, startCoords, startIndex);
|
||||
z[zIndex] = x[startIndex];
|
||||
val[segment] = z[zIndex];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto e = start + threadIdx.x + 1; e < finish; e += blockDim.x) {
|
||||
LongType eCoords[SD_MAX_RANK];
|
||||
LongType eIndex;
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, eCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, eCoords, eIndex);
|
||||
if (eIndex >= xLen) return;
|
||||
math::atomics::sd_atomicMin(&z[zIndex], x[eIndex]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void unsortedSegmentMinLinearKernel(const void* input, const LongType* inputShape,
|
||||
const void* indices, const LongType* indicesShape, LongType* starts, LongType* lengths,
|
||||
LongType numOfClasses, void* output,
|
||||
const LongType* outputShape) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType xLen, zLen, segment, zIndex;
|
||||
__shared__ const T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ const I* y;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
segment = blockIdx.x;
|
||||
x = reinterpret_cast<const T*>(input);
|
||||
z = reinterpret_cast<T*>(output);
|
||||
y = reinterpret_cast<const 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);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zIndex);
|
||||
if (lengths[segment] > 0) {
|
||||
LongType startCoords[SD_MAX_RANK];
|
||||
LongType startIndex;
|
||||
INDEX2COORDS(starts[segment], inputRank, inputShapePtr, startCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, startCoords, startIndex);
|
||||
z[zIndex] = x[startIndex];
|
||||
} else {
|
||||
z[zIndex] = DataTypeUtils::max<T>();
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (lengths[segment] > 0) {
|
||||
for (auto e = threadIdx.x + 1; e < xLen; e += blockDim.x) {
|
||||
LongType eCoords[SD_MAX_RANK];
|
||||
LongType eIndex;
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, eCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, eCoords, eIndex);
|
||||
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType yIndex;
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
|
||||
if (y[yIndex] == segment) {
|
||||
math::atomics::sd_atomicMin(&z[zIndex], x[eIndex]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMinTadKernel(const void* inputBuf, const LongType* inputShape,
|
||||
const LongType* inputTads, const LongType* inputTadOffsets,
|
||||
I* indices, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, void* outputBuf, const LongType* outputShape,
|
||||
const LongType* outputTads, const LongType* outputTadOffsets, LongType indicesLen) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType len, zIndex, total;
|
||||
__shared__ T* z;
|
||||
__shared__ int threadsPerSegment, start, finish;
|
||||
|
||||
// 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(blockIdx.x >= indicesLen)
|
||||
return;
|
||||
|
||||
auto segment = indices[blockIdx.x];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z = reinterpret_cast<T*>(outputBuf) + outputTadOffsets[segment];
|
||||
len = shape::length(inputTads);
|
||||
start = starts[segment];
|
||||
finish = start + lengths[segment];
|
||||
total = shape::sizeAt(inputShape, 0);
|
||||
|
||||
// Cache TAD 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);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto idx = blockIdx.x;
|
||||
if (blockIdx.x <= total) {
|
||||
auto x = reinterpret_cast<const T*>(inputBuf) + inputTadOffsets[idx];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
LongType zOffset;
|
||||
|
||||
for (auto e = threadIdx.x; e < len; e += blockDim.x) {
|
||||
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, xCoords);
|
||||
COORDS2INDEX(inputTadRank, inputTadStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(e, outputTadRank, outputTadShapePtr, zCoords);
|
||||
COORDS2INDEX(outputTadRank, outputTadStridePtr, zCoords, zOffset);
|
||||
math::atomics::sd_atomicMin(&z[zOffset], x[xOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// segmen min
|
||||
template <typename T, typename I>
|
||||
static void segmentMinFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
LongType numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
auto classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
auto classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
T val = DataTypeUtils::infOrMax<T>();
|
||||
output->assign(val);
|
||||
sd::LongType zero2 = 0;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(zero2);
|
||||
classesRangesLens.assign(len);
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, &classesRangesBegs, &classesRangesLens});
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
dim3 launchDims = segmentDims(numClasses,input->lengthOf());
|
||||
segmentMinLinearKernel<T, I><<<launchDims.y,launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinLinearKernel 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(input->sizeAt(0));
|
||||
segmentMinTadKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
||||
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo(), outputTads, outputTadOffsets, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, &classesRangesBegs, &classesRangesLens});
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void segmentMinFunctor(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
output->nullify();
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), segmentMinFunctor_, (context, input, indices, output),
|
||||
SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static void unsortedSegmentMinFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
T val = DataTypeUtils::infOrMax<T>();
|
||||
sd::LongType len = indices->lengthOf();
|
||||
output->assign(val);
|
||||
sd::LongType zero = 0;
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
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());
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
unsortedSegmentMinLinearKernel<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, "unsortedSegmentMinLinearKernel failed");
|
||||
|
||||
} else {
|
||||
T val = DataTypeUtils::max<T>();
|
||||
output->assign(val);
|
||||
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);
|
||||
segmentMinTadKernel<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, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void unsortedSegmentMinFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses,
|
||||
NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
output->nullify();
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentMinFunctor_,
|
||||
(context, input, indices, numOfClasses, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentMinBPLinearKernel(const void* inputBuf, const LongType* inputShape,
|
||||
void* forwardOutput, const LongType* forwardShape, void* eps,
|
||||
const LongType* epsShape, const void* indicesBuf,
|
||||
const LongType* indicesShape, void* outputBuf,
|
||||
const LongType* outputShape) {
|
||||
__shared__ const T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ const I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, gradLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank, forwardRank, epsRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* forwardShapePtr;
|
||||
__shared__ const sd::LongType* epsShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
__shared__ const sd::LongType* forwardStridePtr;
|
||||
__shared__ const sd::LongType* epsStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<const T*>(inputBuf);
|
||||
y = reinterpret_cast<const 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
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
forwardRank = shape::rank(forwardShape);
|
||||
epsRank = shape::rank(epsShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
forwardShapePtr = shape::shapeOf(forwardShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
forwardStridePtr = shape::stride(forwardShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (auto e = start; e < xLen; 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, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yOffset);
|
||||
auto classIndex = y[yOffset];
|
||||
INDEX2COORDS(classIndex, forwardRank, forwardShapePtr, gradICoords);
|
||||
COORDS2INDEX(forwardRank, forwardStridePtr, gradICoords, gradOffsetI);
|
||||
INDEX2COORDS(classIndex, epsRank, epsShapePtr, gradOCoords);
|
||||
COORDS2INDEX(epsRank, epsStridePtr, 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 segmentMinBPTadKernel(const void* inputBuf, const LongType* inputShape, void* forwardOutput,
|
||||
const LongType* forwardShape, void* eps, const LongType* epsShape,
|
||||
const void* indicesBuf, const LongType* indicesShape, void* outputBuf,
|
||||
const LongType* outputShape, const LongType* inputTad,
|
||||
const LongType* inputOffsets, const LongType* gradInTad,
|
||||
const LongType* gradInOffsets, const LongType* gradOutTad,
|
||||
const LongType* gradOutOffsets, const LongType* outTad,
|
||||
const LongType* outOffsets) {
|
||||
__shared__ const T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ const I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, yLen, gradLen, currentLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType indicesRank;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<const T*>(inputBuf);
|
||||
y = reinterpret_cast<const 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);
|
||||
currentLen = shape::length(outTad);
|
||||
|
||||
// Cache indices shape information (only needed for segment calculation)
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto i = blockIdx.x; i < yLen; i += gridDim.x) {
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType yIndex;
|
||||
INDEX2COORDS(i, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
auto segment = y[yIndex];
|
||||
auto current = x + inputOffsets[i];
|
||||
auto currentOut = z + outOffsets[i];
|
||||
auto in = gradIn + gradInOffsets[segment];
|
||||
auto outGrad = gradOut + gradOutOffsets[segment];
|
||||
|
||||
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
||||
if (math::sd_abs<T,T>(in[e] - current[e]) <= T(1.e-6)) currentOut[e] = outGrad[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
Status segmentMinFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
||||
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);
|
||||
segmentMinFunctor_<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();
|
||||
|
||||
segmentMinBPLinearKernel<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());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinBPLinearKernel 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);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradInTads = packGradIn->specialShapeInfo();
|
||||
auto gradInTadOffsets = packGradIn->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
|
||||
segmentMinBPTadKernel<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);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinBPTadKernel failed");
|
||||
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// segment min
|
||||
Status segmentMinFunctorBP(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 segmentMinFunctorBP_,
|
||||
(context, input, indices, gradOut, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static Status unsortedSegmentMinFunctorBP_(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);
|
||||
unsortedSegmentMinFunctor_<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();
|
||||
segmentMinBPLinearKernel<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());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinBPLinearKernel 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);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradInTads = packGradIn->specialShapeInfo();
|
||||
auto gradInTadOffsets = packGradIn->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
|
||||
segmentMinBPTadKernel<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);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentMinBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut, &tempRes});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
Status unsortedSegmentMinFunctorBP(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 unsortedSegmentMinFunctorBP_,
|
||||
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,862 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 "helpers/DebugHelper.h"
|
||||
#include <system/selective_rendering.h>
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Segment Prod ops linear kernels
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentProdLinearKernel(void* input, LongType const* inputShape, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, void* output,
|
||||
LongType const* outputShape) {
|
||||
|
||||
// Shared memory for caching shape, stride, and rank information
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
|
||||
// Shared memory for pointers and lengths initialized by thread 0
|
||||
__shared__ T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen;
|
||||
__shared__ LongType zLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache rank, shape, and stride for inputShape
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
// Cache rank, shape, and stride for outputShape
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
// Cache lengths
|
||||
xLen = shape::length(inputShape);
|
||||
zLen = shape::length(outputShape);
|
||||
|
||||
// Initialize pointers
|
||||
x = reinterpret_cast<T*>(input);
|
||||
z = reinterpret_cast<T*>(output);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index and step size
|
||||
LongType startIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
LongType step = blockDim.x * gridDim.x;
|
||||
|
||||
// Coordinate arrays
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
// Offset variables
|
||||
LongType xIndex;
|
||||
LongType zIndex;
|
||||
|
||||
// Iterate over each class segment assigned to this block
|
||||
for (LongType segment = blockIdx.x; segment < numOfClasses; segment += gridDim.x) {
|
||||
// Convert segment index to coordinates for outputShape
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, outputCoords);
|
||||
// Convert coordinates back to linear index for outputShape
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, zIndex);
|
||||
|
||||
// Skip processing if zIndex is out of bounds
|
||||
if (zIndex >= zLen)
|
||||
continue;
|
||||
|
||||
// Retrieve start and finish indices for the current segment
|
||||
auto start = starts[segment];
|
||||
auto finish = start + lengths[segment];
|
||||
|
||||
// Skip processing if the length for the segment is zero
|
||||
if (lengths[segment] == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Iterate over elements within the segment, distributing work among threads
|
||||
for (LongType e = startIdx; e < finish; e += step) {
|
||||
// Convert linear index to coordinates for inputShape
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, inputCoords);
|
||||
// Convert coordinates back to linear index for inputShape
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, xIndex);
|
||||
|
||||
// Skip processing if xIndex is out of bounds
|
||||
if (xIndex >= xLen)
|
||||
continue;
|
||||
|
||||
// Perform atomic multiplication on the output buffer
|
||||
math::atomics::sd_atomicMul(&z[zIndex], x[xIndex]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void unsortedSegmentProdLinearKernel(T* input, LongType const* inputShape, I* indices,
|
||||
LongType const* indicesShape, LongType* starts, LongType* lengths,
|
||||
LongType numOfClasses, T* output, LongType const* outputShape) {
|
||||
|
||||
// Shared memory for caching shape, stride, and rank information
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
|
||||
__shared__ LongType indicesRank;
|
||||
__shared__ const LongType* indicesShapePtr;
|
||||
__shared__ const LongType* indicesStridePtr;
|
||||
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
|
||||
// Shared memory for pointers and lengths initialized by thread 0
|
||||
__shared__ T* x;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen;
|
||||
__shared__ LongType zLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache rank, shape, and stride for inputShape
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
// Cache rank, shape, and stride for indicesShape
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
|
||||
// Cache rank, shape, and stride for outputShape
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
// Cache lengths
|
||||
xLen = shape::length(inputShape);
|
||||
zLen = shape::length(outputShape);
|
||||
|
||||
// Initialize pointers
|
||||
x = input;
|
||||
y = indices;
|
||||
z = output;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index and step size
|
||||
LongType startIdx = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
LongType step = blockDim.x * gridDim.x;
|
||||
|
||||
// Coordinate arrays
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
|
||||
// Offset variables
|
||||
LongType xIndex;
|
||||
LongType yIndex;
|
||||
LongType zIndex;
|
||||
|
||||
for (LongType idx = startIdx; idx < xLen; idx += step) {
|
||||
// Convert linear index to coordinates for inputShape
|
||||
INDEX2COORDS(idx, inputRank, inputShapePtr, xCoords);
|
||||
// Convert coordinates back to linear index for inputShape
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xIndex);
|
||||
|
||||
// Convert linear index to coordinates for indicesShape
|
||||
INDEX2COORDS(idx, indicesRank, indicesShapePtr, yCoords);
|
||||
// Convert coordinates back to linear index for indicesShape
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
|
||||
// Retrieve the segment index from indices
|
||||
auto segment = y[yIndex];
|
||||
|
||||
// Convert segment index to coordinates for outputShape
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, zCoords);
|
||||
// Convert coordinates back to linear index for outputShape
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zIndex);
|
||||
|
||||
// Skip processing if the length for the segment is zero
|
||||
if (lengths[segment] == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Perform atomic multiplication on the output buffer
|
||||
math::atomics::sd_atomicMul(&z[zIndex], x[xIndex]);
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// SegmentProd kernel
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentProdTadKernel(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, LongType indicesLen) {
|
||||
|
||||
// Early exit if block index is out of range
|
||||
if (blockIdx.x >= indicesLen)
|
||||
return;
|
||||
|
||||
// Shared memory for caching shape, stride, and rank information
|
||||
__shared__ LongType inputTadRank;
|
||||
__shared__ const LongType* inputTadShapePtr;
|
||||
__shared__ const LongType* inputTadStridePtr;
|
||||
|
||||
__shared__ LongType outputTadRank;
|
||||
__shared__ const LongType* outputTadShapePtr;
|
||||
__shared__ const LongType* outputTadStridePtr;
|
||||
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
|
||||
// Shared memory for pointers and lengths initialized by thread 0
|
||||
__shared__ T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType len;
|
||||
__shared__ LongType total;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache rank, shape, and stride for inputTads
|
||||
inputTadRank = shape::rank(inputTads);
|
||||
inputTadShapePtr = shape::shapeOf(inputTads);
|
||||
inputTadStridePtr = shape::stride(inputTads);
|
||||
|
||||
// Cache rank, shape, and stride for outputTads
|
||||
outputTadRank = shape::rank(outputTads);
|
||||
outputTadShapePtr = shape::shapeOf(outputTads);
|
||||
outputTadStridePtr = shape::stride(outputTads);
|
||||
|
||||
// Cache rank, shape, and stride for inputShape
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
// Cache rank, shape, and stride for outputShape
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
// Cache lengths and total size
|
||||
total = shape::sizeAt(inputShape, 0);
|
||||
len = shape::length(inputTads);
|
||||
|
||||
// Initialize pointers
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index and step size
|
||||
LongType startIdx = blockIdx.x;
|
||||
LongType step = gridDim.x;
|
||||
|
||||
// Coordinate arrays
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
|
||||
// Offset variables
|
||||
LongType xIndex;
|
||||
LongType zIndex;
|
||||
|
||||
for (auto idx = startIdx; idx < total; idx += step) {
|
||||
// Retrieve the segment index from indices
|
||||
auto segment = indices[idx];
|
||||
|
||||
// Pointers to the current input and output TADs
|
||||
T* current = x + inputTadOffsets[idx];
|
||||
T* currentOut = z + outputTadOffsets[segment];
|
||||
|
||||
// Retrieve start and finish indices for the current segment
|
||||
LongType start = starts[segment];
|
||||
LongType finish = start + lengths[segment];
|
||||
|
||||
// Skip processing if the length for the segment is zero
|
||||
if (lengths[segment] == 0) continue;
|
||||
|
||||
// Iterate over elements within the TAD
|
||||
for (auto e = threadIdx.x; e < len; e += blockDim.x) {
|
||||
// Convert linear index to coordinates for inputTads
|
||||
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, inputCoords);
|
||||
// Convert coordinates back to linear index for inputTads
|
||||
COORDS2INDEX(inputTadRank, inputTadStridePtr, inputCoords, xIndex);
|
||||
|
||||
// Convert linear index to coordinates for outputTads
|
||||
INDEX2COORDS(e, outputTadRank, outputTadShapePtr, outputCoords);
|
||||
// Convert coordinates back to linear index for outputTads
|
||||
COORDS2INDEX(outputTadRank, outputTadStridePtr, outputCoords, zIndex);
|
||||
|
||||
// Perform atomic multiplication on the output buffer
|
||||
math::atomics::sd_atomicMul(¤tOut[zIndex], current[xIndex]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static void segmentProdFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
LongType numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
output->assign(one);
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
dim3 launchDims = segmentDims(indices->lengthOf(),input->lengthOf());
|
||||
segmentProdLinearKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(input->specialBuffer(), input->specialShapeInfo(), begins,
|
||||
lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdLinearKernel 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(input->lengthOf());
|
||||
segmentProdTadKernel<T, I><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
||||
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo(), outputTads, outputTadOffsets, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void segmentProdFunctor(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(output->dataType(), indices->dataType(), segmentProdFunctor_, (context, input, indices, output),
|
||||
SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static void unsortedSegmentProdFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
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());
|
||||
output->assign(one);
|
||||
|
||||
dim3 launchDims = getLaunchDims("unsorted_segment_prod_2");
|
||||
if (input->isVector()) {
|
||||
unsortedSegmentProdLinearKernel<T, I><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->dataBuffer()->specialAsT<T>(), input->specialShapeInfo(), indices->dataBuffer()->specialAsT<I>(),
|
||||
indices->specialShapeInfo(), begins, lengths, numOfClasses, output->dataBuffer()->specialAsT<T>(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "unsortedSegmentProdLinearKernel 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);
|
||||
segmentProdTadKernel<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, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void unsortedSegmentProdFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses,
|
||||
NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentProdFunctor_,
|
||||
(context, input, indices, numOfClasses, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentProdBPLinearKernel(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) {
|
||||
|
||||
// Shared memory for caching shape, stride, and rank information
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
|
||||
__shared__ LongType forwardRank;
|
||||
__shared__ const LongType* forwardShapePtr;
|
||||
__shared__ const LongType* forwardStridePtr;
|
||||
|
||||
__shared__ LongType epsRank;
|
||||
__shared__ const LongType* epsShapePtr;
|
||||
__shared__ const LongType* epsStridePtr;
|
||||
|
||||
__shared__ LongType indicesRank;
|
||||
__shared__ const LongType* indicesShapePtr;
|
||||
__shared__ const LongType* indicesStridePtr;
|
||||
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
|
||||
// Shared memory for pointers and lengths initialized by thread 0
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen;
|
||||
__shared__ LongType gradLen;
|
||||
__shared__ LongType currentLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache rank, shape, and stride for inputShape
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
// Cache rank, shape, and stride for forwardShape
|
||||
forwardRank = shape::rank(forwardShape);
|
||||
forwardShapePtr = shape::shapeOf(forwardShape);
|
||||
forwardStridePtr = shape::stride(forwardShape);
|
||||
|
||||
// Cache rank, shape, and stride for epsShape
|
||||
epsRank = shape::rank(epsShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
|
||||
// Cache rank, shape, and stride for indicesShape
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
|
||||
// Cache rank, shape, and stride for outputShape
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
// Initialize pointers and lengths
|
||||
xLen = shape::length(inputShape);
|
||||
gradLen = shape::length(epsShape);
|
||||
currentLen = shape::length(outputShape); // Assuming 'currentLen' corresponds to outputShape
|
||||
|
||||
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);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index and step size
|
||||
LongType start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
LongType step = gridDim.x * blockDim.x;
|
||||
|
||||
// Coordinate arrays
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType gradICoords[SD_MAX_RANK];
|
||||
LongType gradOCoords[SD_MAX_RANK];
|
||||
|
||||
// Offset variables
|
||||
LongType xOffset;
|
||||
LongType yOffset;
|
||||
LongType zOffset;
|
||||
LongType gradOffsetI;
|
||||
LongType gradOffsetO;
|
||||
|
||||
for (LongType e = start; e < xLen; e += step) {
|
||||
// Convert linear index to coordinates for inputShape
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, xCoords);
|
||||
// Convert coordinates back to linear index for inputShape
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
|
||||
// Convert linear index to coordinates for indicesShape
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
// Convert coordinates back to linear index for indicesShape
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yOffset);
|
||||
|
||||
// Retrieve the class index from indices
|
||||
auto classIndex = y[yOffset];
|
||||
|
||||
// Convert class index to coordinates for forwardShape
|
||||
INDEX2COORDS(classIndex, forwardRank, forwardShapePtr, gradICoords);
|
||||
// Convert coordinates back to linear index for forwardShape
|
||||
COORDS2INDEX(forwardRank, forwardStridePtr, gradICoords, gradOffsetI);
|
||||
|
||||
// Convert class index to coordinates for epsShape
|
||||
INDEX2COORDS(classIndex, epsRank, epsShapePtr, gradOCoords);
|
||||
// Convert coordinates back to linear index for epsShape
|
||||
COORDS2INDEX(epsRank, epsStridePtr, gradOCoords, gradOffsetO);
|
||||
|
||||
// Convert linear index to coordinates for outputShape
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, zCoords);
|
||||
// Convert coordinates back to linear index for outputShape
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
|
||||
// Perform the computation: z[zOffset] = gradOut[gradOffsetO] * gradIn[gradOffsetI] / x[xOffset];
|
||||
z[zOffset] = gradOut[gradOffsetO] * gradIn[gradOffsetI] / x[xOffset];
|
||||
}
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentProdBPTadKernel(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* inputTad,
|
||||
LongType const* inputOffsets, LongType const* gradInTad,
|
||||
LongType const* gradInOffsets, LongType const* gradOutTad,
|
||||
LongType const* gradOutOffsets, LongType const* outTad,
|
||||
LongType const* outOffsets) {
|
||||
|
||||
// Shared memory for caching shape, stride, and rank information
|
||||
__shared__ LongType inputRank;
|
||||
__shared__ const LongType* inputShapePtr;
|
||||
__shared__ const LongType* inputStridePtr;
|
||||
|
||||
__shared__ LongType forwardRank;
|
||||
__shared__ const LongType* forwardShapePtr;
|
||||
__shared__ const LongType* forwardStridePtr;
|
||||
|
||||
__shared__ LongType epsRank;
|
||||
__shared__ const LongType* epsShapePtr;
|
||||
__shared__ const LongType* epsStridePtr;
|
||||
|
||||
__shared__ LongType indicesRank;
|
||||
__shared__ const LongType* indicesShapePtr;
|
||||
__shared__ const LongType* indicesStridePtr;
|
||||
|
||||
__shared__ LongType outputRank;
|
||||
__shared__ const LongType* outputShapePtr;
|
||||
__shared__ const LongType* outputStridePtr;
|
||||
|
||||
__shared__ LongType inputTadRank;
|
||||
__shared__ const LongType* inputTadShapePtr;
|
||||
__shared__ const LongType* inputTadStridePtr;
|
||||
|
||||
__shared__ LongType gradInTadRank;
|
||||
__shared__ const LongType* gradInTadShapePtr;
|
||||
__shared__ const LongType* gradInTadStridePtr;
|
||||
|
||||
__shared__ LongType gradOutTadRank;
|
||||
__shared__ const LongType* gradOutTadShapePtr;
|
||||
__shared__ const LongType* gradOutTadStridePtr;
|
||||
|
||||
__shared__ LongType outTadRank;
|
||||
__shared__ const LongType* outTadShapePtr;
|
||||
__shared__ const LongType* outTadStridePtr;
|
||||
|
||||
// Shared memory for pointers and lengths initialized by thread 0
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen;
|
||||
__shared__ LongType yLen;
|
||||
__shared__ LongType gradLen;
|
||||
__shared__ LongType currentLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache rank, shape, and stride for inputShape
|
||||
inputRank = shape::rank(inputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
|
||||
// Cache rank, shape, and stride for forwardShape
|
||||
forwardRank = shape::rank(forwardShape);
|
||||
forwardShapePtr = shape::shapeOf(forwardShape);
|
||||
forwardStridePtr = shape::stride(forwardShape);
|
||||
|
||||
// Cache rank, shape, and stride for epsShape
|
||||
epsRank = shape::rank(epsShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
|
||||
// Cache rank, shape, and stride for indicesShape
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
|
||||
// Cache rank, shape, and stride for outputShape
|
||||
outputRank = shape::rank(outputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
|
||||
// Cache rank, shape, and stride for inputTad
|
||||
inputTadRank = shape::rank(inputTad);
|
||||
inputTadShapePtr = shape::shapeOf(inputTad);
|
||||
inputTadStridePtr = shape::stride(inputTad);
|
||||
|
||||
// Cache rank, shape, and stride for gradInTad
|
||||
gradInTadRank = shape::rank(gradInTad);
|
||||
gradInTadShapePtr = shape::shapeOf(gradInTad);
|
||||
gradInTadStridePtr = shape::stride(gradInTad);
|
||||
|
||||
// Cache rank, shape, and stride for gradOutTad
|
||||
gradOutTadRank = shape::rank(gradOutTad);
|
||||
gradOutTadShapePtr = shape::shapeOf(gradOutTad);
|
||||
gradOutTadStridePtr = shape::stride(gradOutTad);
|
||||
|
||||
// Cache rank, shape, and stride for outTad
|
||||
outTadRank = shape::rank(outTad);
|
||||
outTadShapePtr = shape::shapeOf(outTad);
|
||||
outTadStridePtr = shape::stride(outTad);
|
||||
|
||||
// Initialize pointers and lengths
|
||||
xLen = shape::length(inputShape);
|
||||
yLen = shape::length(indicesShape);
|
||||
gradLen = shape::length(epsShape);
|
||||
currentLen = shape::length(outTad);
|
||||
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
y = reinterpret_cast<I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
gradOut = reinterpret_cast<T*>(eps);
|
||||
gradIn = reinterpret_cast<T*>(forwardOutput);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate global thread index and step size
|
||||
LongType startIdx = blockIdx.x;
|
||||
LongType step = gridDim.x;
|
||||
|
||||
// Coordinate arrays
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType yIndex;
|
||||
|
||||
// Iterate over all relevant indices
|
||||
for (auto i = startIdx; i < yLen; i += step) {
|
||||
// Convert linear index to coordinates for indicesShape
|
||||
INDEX2COORDS(i, indicesRank, indicesShapePtr, yCoords);
|
||||
// Convert coordinates back to linear index for indicesShape
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
|
||||
// Retrieve the segment index from indices
|
||||
auto segment = y[yIndex];
|
||||
|
||||
// Pointers to the current input and output TADs
|
||||
T* current = x + inputOffsets[i];
|
||||
T* currentOut = z + outOffsets[i];
|
||||
|
||||
// Pointers to the corresponding gradIn and gradOut TADs
|
||||
T* in = gradIn + gradInOffsets[segment];
|
||||
T* outGrad = gradOut + gradOutOffsets[segment];
|
||||
|
||||
// Perform element-wise computation within the current TAD
|
||||
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
||||
// Compute output: currentOut[e] = outGrad[e] * in[e] / current[e];
|
||||
currentOut[e] = outGrad[e] * in[e] / current[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
Status segmentProdFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
||||
NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
auto outShape = gradOut->getShapeAsVector();
|
||||
|
||||
NDArray tempRes(gradOut->ordering(), outShape, DataTypeUtils::fromT<T>(),
|
||||
context); //->shapeInfo(), context);
|
||||
segmentProdFunctor_<T, I>(context, input, indices, &tempRes);
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
if (input->isVector()) {
|
||||
LongType loopSize = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf(); // indices->e<sd::LongType>(loop_size - 1);
|
||||
segmentProdBPLinearKernel<T, I><<<gradOut->lengthOf(), loopSize, 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), tempRes.specialBuffer(), tempRes.specialShapeInfo(),
|
||||
gradOut->specialBuffer(), gradOut->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdBPLinearKernel 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);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradInTads = packGradIn->specialShapeInfo();
|
||||
auto gradInTadOffsets = packGradIn->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
dim3 segmentBpTad2 = segmentBpTad(gradOut->lengthOf(),input->lengthOf());
|
||||
|
||||
segmentProdBPTadKernel<T, I><<<segmentBpTad2.y,segmentBpTad2.x, segmentBpTad2.z, *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);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
Status segmentProdFunctorBP(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 segmentProdFunctorBP_,
|
||||
(context, input, indices, gradOut, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static Status unsortedSegmentProdFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices,
|
||||
NDArray* gradOut,
|
||||
LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
auto outShape = gradOut->getShapeAsVector();
|
||||
NDArray tempRes(gradOut->ordering(),outShape, DataTypeUtils::fromT<T>(),
|
||||
context);
|
||||
unsortedSegmentProdFunctor_<T, I>(context, input, indices, numOfClasses, &tempRes);
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
if (input->isVector()) {
|
||||
LongType loopSize = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf();
|
||||
dim3 segmentBpTad2 = segmentBpDims(gradOut->lengthOf(),input->lengthOf());
|
||||
segmentProdBPLinearKernel<T, I><<<segmentBpTad2.y, segmentBpTad2.x,segmentBpTad2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), tempRes.specialBuffer(), tempRes.specialShapeInfo(),
|
||||
gradOut->specialBuffer(), gradOut->specialShapeInfo(), indices->specialBuffer(), indices->specialShapeInfo(),
|
||||
output->specialBuffer(), output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdBPLinearKernel 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);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradInTads = packGradIn->specialShapeInfo();
|
||||
auto gradInTadOffsets = packGradIn->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
dim3 segmentBpTad2 = segmentBpTad(gradOut->lengthOf(),input->lengthOf());
|
||||
segmentProdBPTadKernel<T, I><<<indices->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);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentProdBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
Status unsortedSegmentProdFunctorBP(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 unsortedSegmentProdFunctorBP_,
|
||||
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,404 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 "helpers/DebugHelper.h"
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void unsortedSegmentSqrtNLinearKernel(T* input, LongType const* inputShape, I* indices,
|
||||
LongType const* indicesShape, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, T* output,
|
||||
LongType const* outputShape) {
|
||||
__shared__ LongType xLen, zLen;
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
zLen = shape::length(outputShape);
|
||||
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
auto step = blockDim.x * gridDim.x;
|
||||
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yIndex;
|
||||
LongType zIndex;
|
||||
LongType xIndex;
|
||||
|
||||
for (auto idx = start; idx < xLen; idx += step) {
|
||||
INDEX2COORDS(idx, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
auto segment = indices[yIndex];
|
||||
INDEX2COORDS(segment, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zIndex);
|
||||
if (lengths[segment] == 0) continue;
|
||||
INDEX2COORDS(idx, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xIndex);
|
||||
if (xIndex >= xLen) continue;
|
||||
math::atomics::sd_atomicAdd(&output[zIndex], input[xIndex] / math::sd_sqrt<LongType, T>(lengths[segment]));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSqrtNTadKernel(T* 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, LongType numIndices) {
|
||||
if(blockIdx.x >= numIndices)
|
||||
return;
|
||||
|
||||
__shared__ LongType len, total;
|
||||
__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) {
|
||||
total = shape::sizeAt(inputShape, 0);
|
||||
len = shape::length(inputTads);
|
||||
|
||||
inputTadRank = shape::rank(inputTads);
|
||||
outputTadRank = shape::rank(outputTads);
|
||||
|
||||
inputTadShapePtr = shape::shapeOf(inputTads);
|
||||
outputTadShapePtr = shape::shapeOf(outputTads);
|
||||
|
||||
inputTadStridePtr = shape::stride(inputTads);
|
||||
outputTadStridePtr = shape::stride(outputTads);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType xIndex;
|
||||
LongType zIndex;
|
||||
|
||||
for (auto idx = blockIdx.x; idx < total; idx += gridDim.x) {
|
||||
auto segment = indices[idx];
|
||||
auto x = inputBuf + inputTadOffsets[idx];
|
||||
auto z = reinterpret_cast<T*>(outputBuf) + outputTadOffsets[segment];
|
||||
auto start = starts[segment];
|
||||
auto finish = start + lengths[segment];
|
||||
|
||||
for (auto e = threadIdx.x; e < len; e += blockDim.x) {
|
||||
INDEX2COORDS(e, inputTadRank, inputTadShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputTadRank, inputTadStridePtr, inputCoords, xIndex);
|
||||
INDEX2COORDS(e, outputTadRank, outputTadShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputTadRank, outputTadStridePtr, outputCoords, zIndex);
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], x[xIndex] / math::sd_sqrt<LongType, T>(lengths[segment]));
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static void unsortedSegmentSqrtNFunctor_(LaunchContext* context, NDArray* input, NDArray* indices,
|
||||
LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
dim3 dims= getLaunchDims("segmentSqrtN");
|
||||
fillUpSegments(indices, numOfClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
output->nullify();
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
unsortedSegmentSqrtNLinearKernel<T, I><<<dims.x, dims.y, dims.z, *stream>>>(
|
||||
input->dataBuffer()->specialAsT<T>(), input->specialShapeInfo(), indices->dataBuffer()->specialAsT<I>(),
|
||||
indices->specialShapeInfo(), begins, lengths, numOfClasses, output->dataBuffer()->specialAsT<T>(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "unsortedSegmentSqrtNLinearKernel failed");
|
||||
|
||||
} else {
|
||||
output->nullify();
|
||||
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);
|
||||
segmentSqrtNTadKernel<T, I><<<dims.x, dims.y, dims.z, *stream>>>(
|
||||
input->dataBuffer()->specialAsT<T>(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
||||
indices->dataBuffer()->specialAsT<I>(), begins, lengths, numOfClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo(), outputTads, outputTadOffsets, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSqrtNTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void unsortedSegmentSqrtNFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = input->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentSqrtNFunctor_,
|
||||
(context, input, indices, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSqrtNBPLinearKernel(void* inputBuf, LongType const* inputShape, void* eps,
|
||||
LongType const* epsShape, void* indicesBuf,
|
||||
LongType const* indicesShape, LongType* lengths, void* outputBuf,
|
||||
LongType const* outputShape) {
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradIn;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, gradLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank, epsRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* epsShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
__shared__ const sd::LongType* epsStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
y = reinterpret_cast<I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
gradOut = reinterpret_cast<T*>(eps);
|
||||
gradLen = shape::length(epsShape);
|
||||
|
||||
// Cache all shape information
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
epsRank = shape::rank(epsShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
LongType xOffset;
|
||||
LongType yOffset;
|
||||
LongType gradOffsetO;
|
||||
|
||||
for (auto e = start; e < xLen; e += step) {
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yOffset);
|
||||
auto classIndex = y[yOffset];
|
||||
INDEX2COORDS(classIndex, epsRank, epsShapePtr, zCoords);
|
||||
COORDS2INDEX(epsRank, epsStridePtr, zCoords, gradOffsetO);
|
||||
|
||||
z[zOffset] = T(gradOut[gradOffsetO] / math::sd_sqrt<LongType, float>(lengths[classIndex]));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSqrtNBPTadKernel(void* inputBuf, LongType const* inputShape, void* eps,
|
||||
LongType const* epsShape, void* indicesBuf, LongType const* indicesShape,
|
||||
LongType* lengths, void* outputBuf, LongType const* outputShape,
|
||||
LongType const* inputTad, LongType const* inputOffsets,
|
||||
LongType const* gradOutTad, LongType const* gradOutOffsets,
|
||||
LongType const* outTad, LongType const* outOffsets) {
|
||||
__shared__ T* x;
|
||||
__shared__ T* gradOut;
|
||||
__shared__ I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, yLen, gradLen, currentLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType outTadRank, gradOutTadRank;
|
||||
__shared__ const sd::LongType* outTadShapePtr;
|
||||
__shared__ const sd::LongType* gradOutTadShapePtr;
|
||||
__shared__ const sd::LongType* outTadStridePtr;
|
||||
__shared__ const sd::LongType* gradOutTadStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<T*>(inputBuf);
|
||||
y = reinterpret_cast<I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
yLen = shape::length(indicesShape);
|
||||
gradOut = reinterpret_cast<T*>(eps);
|
||||
gradLen = shape::length(epsShape);
|
||||
currentLen = shape::length(outTad);
|
||||
|
||||
// Cache TAD shape information
|
||||
outTadRank = shape::rank(outTad);
|
||||
gradOutTadRank = shape::rank(gradOutTad);
|
||||
outTadShapePtr = shape::shapeOf(outTad);
|
||||
gradOutTadShapePtr = shape::shapeOf(gradOutTad);
|
||||
outTadStridePtr = shape::stride(outTad);
|
||||
gradOutTadStridePtr = shape::stride(gradOutTad);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto i = blockIdx.x; i < yLen; i += gridDim.x) {
|
||||
auto segment = y[i];
|
||||
T* currentOut = z + outOffsets[i];
|
||||
T* outGrad = gradOut + gradOutOffsets[segment];
|
||||
|
||||
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType gradCoords[SD_MAX_RANK];
|
||||
LongType zIndex;
|
||||
LongType gradIndex;
|
||||
|
||||
INDEX2COORDS(e, outTadRank, outTadShapePtr, zCoords);
|
||||
COORDS2INDEX(outTadRank, outTadStridePtr, zCoords, zIndex);
|
||||
INDEX2COORDS(e, gradOutTadRank, gradOutTadShapePtr, gradCoords);
|
||||
COORDS2INDEX(gradOutTadRank, gradOutTadStridePtr, gradCoords, gradIndex);
|
||||
|
||||
if (lengths[segment] > 0)
|
||||
currentOut[zIndex] = T(outGrad[gradIndex] / math::sd_sqrt<LongType, float>(lengths[segment]));
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static Status unsortedSegmentSqrtNFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices,
|
||||
NDArray* gradOut,
|
||||
LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
auto numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
LongType loop_size = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf();
|
||||
segmentSqrtNBPLinearKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSqrtNBPLinearKernel 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 packGradOut = ConstantTadHelper::getInstance().tadForDimensions(gradOut->shapeInfo(), dimensions);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
dim3 segmentBpTad2 = segmentBpTad(indices->lengthOf(),input->lengthOf());
|
||||
|
||||
segmentSqrtNBPTadKernel<T, I><<<segmentBpTad2.y, segmentBpTad2.x, segmentBpTad2.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), lengths, output->specialBuffer(),
|
||||
output->specialShapeInfo(), inputTads, inputTadOffsets, gradOutTads, gradOutTadOffsets, outputTads,
|
||||
outputTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSqrtNBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
Status unsortedSegmentSqrtNFunctorBP(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 unsortedSegmentSqrtNFunctorBP_,
|
||||
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,508 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 "helpers/DebugHelper.h"
|
||||
#include <system/selective_rendering.h>
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Segment ops linear kernels
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSumLinearKernel(const void* input, const LongType* inputShape, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, void* output,
|
||||
const LongType* outputShape) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType xLen, zLen, segment, zIndex;
|
||||
__shared__ const T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ int threadsPerSegment, start, finish;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
threadsPerSegment = (gridDim.x + numOfClasses - 1) / numOfClasses;
|
||||
segment = blockIdx.x / threadsPerSegment;
|
||||
x = reinterpret_cast<const T*>(input);
|
||||
z = reinterpret_cast<T*>(output);
|
||||
|
||||
xLen = shape::length(inputShape);
|
||||
zLen = shape::length(outputShape);
|
||||
|
||||
if (segment < numOfClasses) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(segment, shape::rank(outputShape), shape::shapeOf(outputShape), zCoords);
|
||||
COORDS2INDEX(shape::rank(outputShape), shape::stride(outputShape), zCoords, zIndex);
|
||||
if(zIndex >= zLen)
|
||||
return;
|
||||
start = starts[segment];
|
||||
finish = start + lengths[segment];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(start, shape::rank(inputShape), shape::shapeOf(inputShape), xCoords);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(shape::rank(inputShape), shape::stride(inputShape), xCoords, xOffset);
|
||||
z[zIndex] = x[xOffset];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto e = start + threadIdx.x + 1; e < finish; e += blockDim.x) {
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(e, shape::rank(inputShape), shape::shapeOf(inputShape), xCoords);
|
||||
LongType xOffset;
|
||||
COORDS2INDEX(shape::rank(inputShape), shape::stride(inputShape), xCoords, xOffset);
|
||||
if (xOffset >= xLen) return;
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], x[xOffset]);
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void unsortedSegmentSumLinearKernel(const void* input, const LongType* inputShape,
|
||||
const void* indices, const LongType* indicesShape, LongType* starts, LongType* lengths,
|
||||
LongType numOfClasses, void* output,
|
||||
const LongType* outputShape) {
|
||||
__shared__ T* val;
|
||||
__shared__ LongType xLen, zLen, segment, zIndex;
|
||||
__shared__ const T* x;
|
||||
__shared__ T* z;
|
||||
__shared__ const I* y;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
segment = blockIdx.x;
|
||||
x = reinterpret_cast<const T*>(input);
|
||||
z = reinterpret_cast<T*>(output);
|
||||
y = reinterpret_cast<const I*>(indices);
|
||||
xLen = shape::length(inputShape);
|
||||
zLen = shape::length(outputShape);
|
||||
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(segment, shape::rank(outputShape), shape::shapeOf(outputShape), zCoords);
|
||||
COORDS2INDEX(shape::rank(outputShape), shape::stride(outputShape), zCoords, zIndex);
|
||||
if (lengths[segment] > 0) {
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType xOffset;
|
||||
INDEX2COORDS(starts[segment], shape::rank(inputShape), shape::shapeOf(inputShape), xCoords);
|
||||
COORDS2INDEX(shape::rank(inputShape), shape::stride(inputShape), xCoords, xOffset);
|
||||
z[zIndex] = x[xOffset];
|
||||
} else {
|
||||
z[zIndex] = 0;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (lengths[segment] > 0) {
|
||||
for (auto e = threadIdx.x; e < xLen; e += blockDim.x) {
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType xIndex;
|
||||
LongType yIndex;
|
||||
|
||||
INDEX2COORDS(e, shape::rank(inputShape), shape::shapeOf(inputShape), xCoords);
|
||||
COORDS2INDEX(shape::rank(inputShape), shape::stride(inputShape), xCoords, xIndex);
|
||||
INDEX2COORDS(e, shape::rank(indicesShape), shape::shapeOf(indicesShape), yCoords);
|
||||
COORDS2INDEX(shape::rank(indicesShape), shape::stride(indicesShape), yCoords, yIndex);
|
||||
|
||||
if (y[yIndex] == segment && e != starts[segment]) {
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], x[xIndex]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// SegmentSum kernel
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSumTadKernel(void* inputBuf, const LongType* inputShape,
|
||||
const LongType* inputTads, const LongType* inputTadOffsets,
|
||||
const I* indices, LongType* starts,
|
||||
LongType* lengths, LongType numOfClasses, void* outputBuf, const LongType* outputShape,
|
||||
const LongType* outputTads, const LongType* outputTadOffsets, LongType numIndices) {
|
||||
|
||||
__shared__ LongType len, total;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
total = shape::sizeAt(inputShape, 0);
|
||||
len = shape::length(inputTads);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto idx = blockIdx.x; idx < total; idx += gridDim.x) {
|
||||
auto x = reinterpret_cast<T*>(inputBuf) + inputTadOffsets[idx];
|
||||
auto segment = indices[idx];
|
||||
auto z = reinterpret_cast<T*>(outputBuf) + outputTadOffsets[segment];
|
||||
auto start = starts[segment];
|
||||
auto finish = start + lengths[segment];
|
||||
if (lengths[segment] == 0) continue;
|
||||
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, shape::rank(inputTads), shape::shapeOf(inputTads), xCoords);
|
||||
COORDS2INDEX(shape::rank(inputTads), shape::stride(inputTads), xCoords, xIndex);
|
||||
INDEX2COORDS(e, shape::rank(outputTads), shape::shapeOf(outputTads), zCoords);
|
||||
COORDS2INDEX(shape::rank(outputTads), shape::stride(outputTads), zCoords, zIndex);
|
||||
|
||||
math::atomics::sd_atomicAdd(&z[zIndex], x[xIndex]);
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
template <typename T, typename I>
|
||||
static void segmentSumFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
LongType numClasses = indices->e<LongType>(indices->lengthOf() - 1) + 1;
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
|
||||
fillUpSegments(indices, numClasses, classesRangesBegs, classesRangesLens);
|
||||
LongType* begins = reinterpret_cast<LongType*>(classesRangesBegs.specialBuffer());
|
||||
LongType* lengths = reinterpret_cast<LongType*>(classesRangesLens.specialBuffer());
|
||||
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
segmentSumLinearKernel<T, I><<<numClasses, input->lengthOf(), numClasses * 32 + 32, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumLinearKernel 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 segmentTadDims = segmentTad(input->sizeAt(0));
|
||||
segmentSumTadKernel<T, I><<<segmentTadDims.y,segmentTadDims.x,segmentTadDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), inputTads, inputTadOffsets,
|
||||
reinterpret_cast<I*>(indices->specialBuffer()), begins, lengths, numClasses, output->specialBuffer(),
|
||||
output->specialShapeInfo(), outputTads, outputTadOffsets, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void segmentSumFunctor(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
output->nullify();
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = input->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), segmentSumFunctor_, (context, input, indices, output),
|
||||
SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
static void unsortedSegmentSumFunctor_(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray classesRangesBegs = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
NDArray classesRangesLens = NDArrayFactory::create<LongType>('c', {numOfClasses}, context);
|
||||
sd::LongType zero = 0;
|
||||
sd::LongType one = 1;
|
||||
sd::LongType len = indices->lengthOf();
|
||||
classesRangesBegs.assign(len);
|
||||
classesRangesLens.assign(zero);
|
||||
dim3 dims = getSegmentSumDims(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()) {
|
||||
unsortedSegmentSumLinearKernel<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, "unsortedSegmentSumLinearKernel failed");
|
||||
|
||||
} else {
|
||||
output->assign(zero);
|
||||
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 dims = segmentTad(input->sizeAt(0));
|
||||
segmentSumTadKernel<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, indices->lengthOf());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
dimensions = nullptr;
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
void unsortedSegmentSumFunctor(LaunchContext* context, NDArray* input, NDArray* indices, LongType numOfClasses,
|
||||
NDArray* output) {
|
||||
NDArray::prepareSpecialUse({output}, {input, indices});
|
||||
output->nullify();
|
||||
auto indicesDType = indices->dataType();
|
||||
auto outputDType = input ->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), unsortedSegmentSumFunctor_,
|
||||
(context, input, indices, numOfClasses, output), SD_NUMERIC_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices});
|
||||
}
|
||||
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Backpropagate ops
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
// Sorted sum backpropagate
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSumBPLinearKernel(const void* inputBuf, const LongType* inputShape, const void* eps,
|
||||
const LongType* epsShape, const void* indicesBuf,
|
||||
const LongType* indicesShape, void* outputBuf,
|
||||
const LongType* outputShape) {
|
||||
__shared__ LongType xLen, gradLen;
|
||||
__shared__ sd::LongType inputRank, outputRank, indicesRank, epsRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* epsShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
__shared__ const sd::LongType* epsStridePtr;
|
||||
|
||||
auto x = reinterpret_cast<const T*>(inputBuf);
|
||||
auto y = reinterpret_cast<const I*>(indicesBuf);
|
||||
auto z = reinterpret_cast<T*>(outputBuf);
|
||||
auto gradOut = reinterpret_cast<const T*>(eps);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
gradLen = shape::length(epsShape);
|
||||
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
epsRank = shape::rank(epsShape);
|
||||
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
epsShapePtr = shape::shapeOf(epsShape);
|
||||
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
epsStridePtr = shape::stride(epsShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto step = gridDim.x * blockDim.x;
|
||||
|
||||
for (auto e = start; e < xLen; e += step) {
|
||||
LongType zCoords[SD_MAX_RANK];
|
||||
LongType xCoords[SD_MAX_RANK];
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType zOffset;
|
||||
LongType xOffset;
|
||||
LongType yOffset;
|
||||
LongType gradOffsetO;
|
||||
|
||||
INDEX2COORDS(e, outputRank, outputShapePtr, zCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, zCoords, zOffset);
|
||||
INDEX2COORDS(e, inputRank, inputShapePtr, xCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(e, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yOffset);
|
||||
auto classIndex = y[yOffset];
|
||||
INDEX2COORDS(classIndex, epsRank, epsShapePtr, zCoords);
|
||||
COORDS2INDEX(epsRank, epsStridePtr, zCoords, gradOffsetO);
|
||||
|
||||
z[zOffset] = gradOut[gradOffsetO];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static SD_KERNEL void segmentSumBPTadKernel(const void* inputBuf, const LongType* inputShape, const void* eps,
|
||||
const LongType* epsShape, const void* indicesBuf,
|
||||
const LongType* indicesShape, void* outputBuf,
|
||||
const LongType* outputShape, const LongType* inputTad,
|
||||
const LongType* inputOffsets, const LongType* gradOutTad,
|
||||
const LongType* gradOutOffsets, const LongType* outTad,
|
||||
const LongType* outOffsets) {
|
||||
__shared__ const T* x;
|
||||
__shared__ const T* gradOut;
|
||||
__shared__ const I* y;
|
||||
__shared__ T* z;
|
||||
__shared__ LongType xLen, yLen, gradLen, currentLen;
|
||||
__shared__ sd::LongType indicesRank;
|
||||
__shared__ const sd::LongType* indicesShapePtr;
|
||||
__shared__ const sd::LongType* indicesStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLen = shape::length(inputShape);
|
||||
x = reinterpret_cast<const T*>(inputBuf);
|
||||
y = reinterpret_cast<const I*>(indicesBuf);
|
||||
z = reinterpret_cast<T*>(outputBuf);
|
||||
yLen = shape::length(indicesShape);
|
||||
gradOut = reinterpret_cast<const T*>(eps);
|
||||
gradLen = shape::length(epsShape);
|
||||
currentLen = shape::length(outTad);
|
||||
|
||||
indicesRank = shape::rank(indicesShape);
|
||||
indicesShapePtr = shape::shapeOf(indicesShape);
|
||||
indicesStridePtr = shape::stride(indicesShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (auto i = blockIdx.x; i < yLen; i += gridDim.x) {
|
||||
LongType yCoords[SD_MAX_RANK];
|
||||
LongType yIndex;
|
||||
INDEX2COORDS(i, indicesRank, indicesShapePtr, yCoords);
|
||||
COORDS2INDEX(indicesRank, indicesStridePtr, yCoords, yIndex);
|
||||
auto segment = y[yIndex];
|
||||
auto currentOut = z + outOffsets[i];
|
||||
auto outGrad = gradOut + gradOutOffsets[segment];
|
||||
|
||||
for (auto e = threadIdx.x; e < currentLen; e += blockDim.x) {
|
||||
currentOut[e] = outGrad[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
template <typename T, typename I>
|
||||
Status segmentSumFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* gradOut,
|
||||
NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
LongType loop_size = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf();
|
||||
segmentSumBPLinearKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumBPLinearKernel 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 packGradOut = ConstantTadHelper::getInstance().tadForDimensions(gradOut->shapeInfo(), dimensions);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
|
||||
segmentSumBPTadKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
inputTads, inputTadOffsets, gradOutTads, gradOutTadOffsets, outputTads, outputTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
|
||||
Status segmentSumFunctorBP(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 segmentSumFunctorBP_,
|
||||
(context, input, indices, gradOut, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
}
|
||||
|
||||
template <typename T, typename I>
|
||||
static Status unsortedSegmentSumFunctorBP_(LaunchContext* context, NDArray* input, NDArray* indices,
|
||||
NDArray* gradOut,
|
||||
LongType numOfClasses, NDArray* output) {
|
||||
auto stream = context->getCudaStream();
|
||||
NDArray::prepareSpecialUse({output}, {input, indices, gradOut});
|
||||
if (input->isVector() || input->isScalar()) {
|
||||
LongType loop_size = input->lengthOf();
|
||||
auto numOfClasses = gradOut->lengthOf();
|
||||
segmentSumBPLinearKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo());
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumBPLinearKernel 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 packGradOut = ConstantTadHelper::getInstance().tadForDimensions(gradOut->shapeInfo(), dimensions);
|
||||
auto inputTads = packX->specialShapeInfo();
|
||||
auto inputTadOffsets = packX->specialOffsets();
|
||||
auto outputTads = packZ->specialShapeInfo();
|
||||
auto outputTadOffsets = packZ->specialOffsets();
|
||||
auto gradOutTads = packGradOut->specialShapeInfo();
|
||||
auto gradOutTadOffsets = packGradOut->specialOffsets();
|
||||
|
||||
segmentSumBPTadKernel<T, I><<<gradOut->lengthOf(), input->lengthOf(), 256, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), gradOut->specialBuffer(), gradOut->specialShapeInfo(),
|
||||
indices->specialBuffer(), indices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(),
|
||||
inputTads, inputTadOffsets, gradOutTads, gradOutTadOffsets, outputTads, outputTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "segmentSumBPTadKernel failed");
|
||||
|
||||
delete dimensions;
|
||||
}
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
return Status::OK;
|
||||
}
|
||||
// -------------------------------------------------------------------------------------------------------------- //
|
||||
Status unsortedSegmentSumFunctorBP(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 unsortedSegmentSumFunctorBP_,
|
||||
(context, input, indices, gradOut, numOfClasses, output), SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
||||
NDArray::registerSpecialUse({output}, {input, indices, gradOut});
|
||||
|
||||
}
|
||||
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,101 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <execution/cuda/LaunchDims.h>
|
||||
#include <ops/declarable/helpers/sequence_mask.h>
|
||||
|
||||
|
||||
#include "helpers/DebugHelper.h"
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace helpers {
|
||||
|
||||
template <typename I, typename B>
|
||||
static SD_KERNEL void sequenceMaskKernel(const void* inputBuf, const LongType* inputShape, void* outputBuf,
|
||||
const LongType* outputShape, int maxIndex) {
|
||||
__shared__ const I* input;
|
||||
__shared__ B* output;
|
||||
__shared__ LongType inputLen, outputLen;
|
||||
|
||||
// Cache shape information
|
||||
__shared__ sd::LongType inputRank, outputRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* outputShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ const sd::LongType* outputStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
input = reinterpret_cast<const I*>(inputBuf);
|
||||
output = reinterpret_cast<B*>(outputBuf);
|
||||
inputLen = shape::length(inputShape);
|
||||
outputLen = shape::length(outputShape);
|
||||
|
||||
// Cache shape information
|
||||
inputRank = shape::rank(inputShape);
|
||||
outputRank = shape::rank(outputShape);
|
||||
inputShapePtr = shape::shapeOf(inputShape);
|
||||
outputShapePtr = shape::shapeOf(outputShape);
|
||||
inputStridePtr = shape::stride(inputShape);
|
||||
outputStridePtr = shape::stride(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType inputCoords[SD_MAX_RANK];
|
||||
LongType outputCoords[SD_MAX_RANK];
|
||||
LongType inputOffset;
|
||||
LongType outputOffset;
|
||||
|
||||
for (auto i = blockIdx.x; i < maxIndex; i += gridDim.x)
|
||||
for (auto k = threadIdx.x; k < inputLen; k += blockDim.x) {
|
||||
INDEX2COORDS(k, inputRank, inputShapePtr, inputCoords);
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inputCoords, inputOffset);
|
||||
if (i < input[inputOffset]) {
|
||||
INDEX2COORDS(k * maxIndex + i, outputRank, outputShapePtr, outputCoords);
|
||||
COORDS2INDEX(outputRank, outputStridePtr, outputCoords, outputOffset);
|
||||
output[outputOffset] = B(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
template <typename I, typename B>
|
||||
static void sequenceMask_(LaunchContext* context, NDArray* input, NDArray* output, int maxIndex) {
|
||||
dim3 launchDims = getSequenceMaskLaunchDims(maxIndex,*input);
|
||||
NDArray::prepareSpecialUse({output}, {input});
|
||||
auto stream = context->getCudaStream();
|
||||
sequenceMaskKernel<I, B><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
input->specialBuffer(), input->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo(), maxIndex);
|
||||
sd::DebugHelper::checkErrorCode(stream, "sequenceMaskKernel failed");
|
||||
|
||||
NDArray::registerSpecialUse({output}, {input});
|
||||
}
|
||||
|
||||
void sequenceMask(LaunchContext* context, NDArray* input, NDArray* output, int maxIndex) {
|
||||
auto inputDType = input->dataType();
|
||||
auto outputDType = output->dataType();
|
||||
BUILD_DOUBLE_SELECTOR(input->dataType(), output->dataType(), sequenceMask_, (context, input, output, maxIndex),
|
||||
SD_INTEGER_TYPES, SD_COMMON_TYPES_EXTENDED);
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( void sequenceMask_,
|
||||
(sd::LaunchContext * context, NDArray* input, NDArray* output, int maxIndex), SD_INTEGER_TYPES,
|
||||
SD_COMMON_TYPES_EXTENDED);
|
||||
} // namespace helpers
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user