wehub-resource-sync
365 lines
14 KiB
Plaintext
365 lines
14 KiB
Plaintext
/*
|
|
* ******************************************************************************
|
|
* *
|
|
* *
|
|
* * This program and the accompanying materials are made available under the
|
|
* * terms of the Apache License, Version 2.0 which is available at
|
|
* * https://www.apache.org/licenses/LICENSE-2.0.
|
|
* *
|
|
* * See the NOTICE file distributed with this work for additional
|
|
* * information regarding copyright ownership.
|
|
* * Unless required by applicable law or agreed to in writing, software
|
|
* * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
* * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
* * License for the specific language governing permissions and limitations
|
|
* * under the License.
|
|
* *
|
|
* * SPDX-License-Identifier: Apache-2.0
|
|
* *****************************************************************************
|
|
*/
|
|
|
|
//
|
|
// @author 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 invertPermutationCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo) {
|
|
const T* x = reinterpret_cast<const T*>(vx);
|
|
T* z = reinterpret_cast<T*>(vz);
|
|
|
|
__shared__ LongType len, totalThreads;
|
|
|
|
if (threadIdx.x == 0) {
|
|
len = shape::length(xShapeInfo);
|
|
totalThreads = gridDim.x * blockDim.x;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xOffset;
|
|
LongType zOffset;
|
|
|
|
for (LongType i = tid; i < len; i += totalThreads) {
|
|
INDEX2COORDS(i, shape::rank(xShapeInfo), shape::shapeOf(xShapeInfo), xCoords);
|
|
COORDS2INDEX(shape::rank(xShapeInfo), shape::stride(xShapeInfo), xCoords, xOffset);
|
|
const LongType index = x[xOffset];
|
|
INDEX2COORDS(index, shape::rank(zShapeInfo), shape::shapeOf(zShapeInfo), zCoords);
|
|
COORDS2INDEX(shape::rank(zShapeInfo), shape::stride(zShapeInfo), zCoords, zOffset);
|
|
z[zOffset] = i;
|
|
}
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_HOST static void invertPermutationCudaLauncher(const int blocksPerGrid, const int threadsPerBlock,
|
|
const int sharedMemory, const cudaStream_t* stream, const void* vx,
|
|
const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo) {
|
|
invertPermutationCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory, *stream>>>(vx, xShapeInfo, vz, zShapeInfo);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "invertPermutationCuda failed");
|
|
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////
|
|
void invertPermutation(LaunchContext* context, NDArray& input, NDArray& output) {
|
|
dim3 invertPermuteDims = invertPermutationDims(input.lengthOf());
|
|
PointersManager manager(context, "invertPermutation");
|
|
|
|
NDArray::prepareSpecialUse({&output}, {&input});
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), invertPermutationCudaLauncher,
|
|
(invertPermuteDims.x, invertPermuteDims.y, invertPermuteDims.z,context->getCudaStream(), input.specialBuffer(),
|
|
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo()),
|
|
SD_COMMON_TYPES);
|
|
NDArray::registerSpecialUse({&output}, {&input});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_KERNEL static void traceCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo, const LongType diagLen) {
|
|
const auto x = reinterpret_cast<const T*>(vx);
|
|
auto z = reinterpret_cast<T*>(vz);
|
|
|
|
__shared__ T sharedMem[SD_CUDA_BLOCK_SIZE];
|
|
|
|
// Shared variables for ranks, shapes, and strides
|
|
__shared__ sd::LongType xRank, zRank;
|
|
__shared__ const sd::LongType* xShapePtr;
|
|
__shared__ const sd::LongType* xStridePtr;
|
|
__shared__ const sd::LongType* zShapePtr;
|
|
__shared__ const sd::LongType* zStridePtr;
|
|
__shared__ LongType xLen, zLen;
|
|
|
|
// Cache all shape-related values in thread 0
|
|
if (threadIdx.x == 0) {
|
|
xRank = shape::rank(xShapeInfo);
|
|
zRank = shape::rank(zShapeInfo);
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
zShapePtr = shape::shapeOf(zShapeInfo);
|
|
zStridePtr = shape::stride(zShapeInfo);
|
|
xLen = shape::length(xShapeInfo);
|
|
zLen = shape::length(zShapeInfo); // corresponds to number of matrices
|
|
}
|
|
__syncthreads();
|
|
|
|
LongType coords[SD_MAX_RANK];
|
|
|
|
// One block per each element of z, that is per each matrix
|
|
for (LongType m = blockIdx.x; m < zLen; m += gridDim.x) {
|
|
INDEX2COORDS(m, zRank, zShapePtr, coords);
|
|
LongType zOffset;
|
|
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
|
|
|
sharedMem[threadIdx.x] = 0;
|
|
|
|
for (LongType i = threadIdx.x; i < diagLen; i += blockDim.x) {
|
|
coords[zRank] = coords[zRank + 1] = i;
|
|
LongType xOffset;
|
|
COORDS2INDEX(xRank, xStridePtr, coords, xOffset);
|
|
sharedMem[threadIdx.x] += x[xOffset];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// Aggregate sum
|
|
for (LongType activeThreads = blockDim.x / 2; activeThreads > 0; activeThreads /= 2) {
|
|
if (threadIdx.x < activeThreads) {
|
|
sharedMem[threadIdx.x] += sharedMem[threadIdx.x + activeThreads];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
z[zOffset] = *sharedMem;
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void traceCudaLauncher(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 diagLen) {
|
|
traceCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, diagLen);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "traceCuda failed");
|
|
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
void trace(LaunchContext* context, NDArray& input, NDArray& output) {
|
|
PointersManager manager(context, "trace");
|
|
|
|
const LongType diagLen = input.sizeAt(-1) < input.sizeAt(-2) ? input.sizeAt(-1) : input.sizeAt(-2);
|
|
const int threadsPerBlock = SD_CUDA_BLOCK_SIZE;
|
|
const int blocksPerGrid = (output.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
|
const int sharedMem = 1024;
|
|
|
|
dim3 traceDims2 = traceDims(output.lengthOf());
|
|
NDArray::prepareSpecialUse({&output}, {&input});
|
|
BUILD_SINGLE_SELECTOR(input.dataType(), traceCudaLauncher,
|
|
(traceDims2.y, traceDims2.x, traceDims2.z, context->getCudaStream(), input.specialBuffer(),
|
|
input.specialShapeInfo(), output.specialBuffer(), output.specialShapeInfo(), diagLen),
|
|
SD_COMMON_TYPES);
|
|
NDArray::registerSpecialUse({&output}, {&input});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_KERNEL static void triuBPCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo, const int diag) {
|
|
const auto x = reinterpret_cast<const T*>(vx); // gradO
|
|
auto z = reinterpret_cast<T*>(vz); // gradI
|
|
|
|
__shared__ int rank, areSameOffsets;
|
|
__shared__ LongType len, totalThreads; // xLen = zLen
|
|
|
|
// Cache shape information
|
|
__shared__ const sd::LongType* xShapePtr;
|
|
__shared__ const sd::LongType* zShapePtr;
|
|
__shared__ const sd::LongType* xStridePtr;
|
|
__shared__ const sd::LongType* zStridePtr;
|
|
__shared__ int xRank, zRank;
|
|
|
|
if (threadIdx.x == 0) {
|
|
areSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
|
rank = shape::rank(xShapeInfo);
|
|
len = shape::length(zShapeInfo);
|
|
totalThreads = gridDim.x * blockDim.x;
|
|
|
|
// Cache shape information
|
|
xRank = shape::rank(xShapeInfo);
|
|
zRank = shape::rank(zShapeInfo);
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
zShapePtr = shape::shapeOf(zShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
zStridePtr = shape::stride(zShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
LongType coords[SD_MAX_RANK];
|
|
const LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
for (LongType i = tid; i < len; i += totalThreads) {
|
|
INDEX2COORDS(i, zRank, zShapePtr, coords);
|
|
|
|
sd::LongType zOffset;
|
|
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
|
|
|
if ((coords[rank - 2] + diag > coords[rank - 1])) // row + diag > col
|
|
z[zOffset] = 0;
|
|
else {
|
|
if (areSameOffsets) {
|
|
z[zOffset] = x[zOffset];
|
|
} else {
|
|
sd::LongType xOffset;
|
|
COORDS2INDEX(xRank, xStridePtr, coords, xOffset);
|
|
z[zOffset] = x[xOffset];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void triuBPCudaLauncher(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 diag) {
|
|
triuBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, diag);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "triuBP failed");
|
|
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
void triuBP(LaunchContext* context, NDArray& input, NDArray& gradO, NDArray& gradI,
|
|
const int diagonal) {
|
|
const int threadsPerBlock = SD_MAX_NUM_THREADS / 4;
|
|
const int blocksPerGrid = (gradO.lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
|
|
const int sharedMem = threadsPerBlock * sizeof(LongType) * gradO.rankOf() + 128;
|
|
dim3 triuDims2 = triuDims(gradO.lengthOf(),gradO.rankOf());
|
|
PointersManager manager(context, "triuBP");
|
|
|
|
NDArray::prepareSpecialUse({&gradI}, {&gradO});
|
|
BUILD_SINGLE_SELECTOR(gradI.dataType(), triuBPCudaLauncher,
|
|
(triuDims2.y, triuDims2.x, triuDims2.z, context->getCudaStream(), gradO.specialBuffer(),
|
|
gradO.specialShapeInfo(), gradI.specialBuffer(), gradI.specialShapeInfo(), diagonal),
|
|
SD_COMMON_TYPES);
|
|
NDArray::registerSpecialUse({&gradI}, {&gradO});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_KERNEL static void tileBPCuda(const void* vx, const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo,
|
|
LongType* globMem) {
|
|
const auto x = reinterpret_cast<const T*>(vx); // gradO
|
|
auto z = reinterpret_cast<T*>(vz); // gradI
|
|
|
|
__shared__ int xRank, zRank;
|
|
__shared__ LongType numOfXOffsets, zLen, totalThreads;
|
|
|
|
// Cache shape information
|
|
__shared__ const sd::LongType* zShapePtr;
|
|
__shared__ const sd::LongType* zStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
xRank = shape::rank(zShapeInfo);
|
|
zRank = shape::rank(zShapeInfo);
|
|
zLen = shape::length(zShapeInfo);
|
|
numOfXOffsets = shape::length(xShapeInfo) / zLen;
|
|
totalThreads = gridDim.x * blockDim.x;
|
|
|
|
// Cache shape information
|
|
zShapePtr = shape::shapeOf(zShapeInfo);
|
|
zStridePtr = shape::stride(zShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
LongType memBuff[SD_MAX_RANK * 2];
|
|
auto xOffsets = globMem + tid * numOfXOffsets;
|
|
|
|
for (LongType i = tid; i < zLen; i += totalThreads) {
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType zOffset;
|
|
|
|
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
|
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
|
|
|
shape::outerArrayOffsets(xOffsets, i, xShapeInfo, zShapeInfo, memBuff, nullptr);
|
|
|
|
z[zOffset] = x[xOffsets[0]]; // first offset
|
|
for (LongType j = 1; j < numOfXOffsets; ++j) // rest offsets
|
|
z[zOffset] += x[xOffsets[j]];
|
|
}
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void tileBPCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
|
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, void* vz,
|
|
const LongType* zShapeInfo, LongType* globMem) {
|
|
tileBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vz, zShapeInfo, globMem);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "tileBPCudaLauncher failed");
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void tileBP(LaunchContext* context, NDArray gradO /*input*/, NDArray& gradI /*output*/,
|
|
const std::vector<LongType> reps) {
|
|
auto grad0Shape = gradO.getShapeAsVector();
|
|
NDArray memBuff(
|
|
'c', grad0Shape, INT64,
|
|
context); // empty auxiliary array for storing device memory which will be used in kernel calculations
|
|
|
|
dim3 tileDims2 = tileDims(gradI.lengthOf(),gradI.rankOf());
|
|
PointersManager manager(context, "tileBP");
|
|
|
|
NDArray::prepareSpecialUse({&gradI}, {&gradO, &memBuff});
|
|
BUILD_SINGLE_SELECTOR(gradI.dataType(), tileBPCudaLauncher,
|
|
(tileDims2.y, tileDims2.x, tileDims2.z, context->getCudaStream(), gradO.specialBuffer(),
|
|
gradO.specialShapeInfo(), gradI.specialBuffer(), gradI.specialShapeInfo(),
|
|
reinterpret_cast<sd::LongType*>(memBuff.specialBuffer())),
|
|
SD_FLOAT_TYPES);
|
|
NDArray::registerSpecialUse({&gradI}, {&gradO, &memBuff});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void eye(LaunchContext* context, NDArray& output) { output.setIdentity(); }
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|