wehub-resource-sync
485 lines
19 KiB
Plaintext
485 lines
19 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)
|
|
//
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <ops/declarable/helpers/top_k.h>
|
|
|
|
#include "execution/cuda/LaunchDims.h"
|
|
#include "helpers/DebugHelper.h"
|
|
#include <system/selective_rendering.h>
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Y>
|
|
__global__ static void inTopKCuda(const void* vx, const LongType* xShapeInfo, const void* vy,
|
|
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
|
const LongType* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType k) {
|
|
const auto y = reinterpret_cast<const Y*>(vy);
|
|
auto z = reinterpret_cast<bool*>(vz);
|
|
|
|
// Shared memory for caching shape information
|
|
__shared__ LongType shared_xRank;
|
|
__shared__ const LongType* shared_xShape;
|
|
__shared__ const LongType* shared_xStride;
|
|
|
|
__shared__ LongType shared_yRank;
|
|
__shared__ const LongType* shared_yShape;
|
|
__shared__ const LongType* shared_yStride;
|
|
|
|
__shared__ LongType shared_zRank;
|
|
__shared__ const LongType* shared_zShape;
|
|
__shared__ const LongType* shared_zStride;
|
|
|
|
__shared__ LongType shared_xTadRank;
|
|
__shared__ const LongType* shared_xTadShape;
|
|
__shared__ const LongType* shared_xTadStride;
|
|
|
|
__shared__ X elemToCompare;
|
|
__shared__ LongType xTadLen;
|
|
__shared__ LongType idx;
|
|
|
|
// Initialize shared memory
|
|
if (threadIdx.x == 0) {
|
|
// Cache ranks
|
|
shared_xRank = shape::rank(xShapeInfo);
|
|
shared_yRank = shape::rank(yShapeInfo);
|
|
shared_zRank = shape::rank(zShapeInfo);
|
|
shared_xTadRank = shape::rank(xTadShapeInfo);
|
|
|
|
// Cache shapes
|
|
shared_xShape = shape::shapeOf(xShapeInfo);
|
|
shared_yShape = shape::shapeOf(yShapeInfo);
|
|
shared_zShape = shape::shapeOf(zShapeInfo);
|
|
shared_xTadShape = shape::shapeOf(xTadShapeInfo);
|
|
|
|
// Cache strides
|
|
shared_xStride = shape::stride(xShapeInfo);
|
|
shared_yStride = shape::stride(yShapeInfo);
|
|
shared_zStride = shape::stride(zShapeInfo);
|
|
shared_xTadStride = shape::stride(xTadShapeInfo);
|
|
|
|
// Cache xTad length
|
|
xTadLen = shape::length(xTadShapeInfo);
|
|
|
|
// Initialize xTad pointer
|
|
// Assuming xTadOffsets is used to compute the starting point for each block
|
|
// Adjusted to point to the correct location in the xTad
|
|
// If xTadOffsets[blockIdx.x] is already in terms of elements, this is correct
|
|
// Otherwise, multiply by the size of X if xTadOffsets are byte offsets
|
|
// Here, we assume they are element offsets
|
|
// If not, use: xTad = reinterpret_cast<const X*>(vx) + xTadOffsets[blockIdx.x] / sizeof(X);
|
|
// Adjust accordingly based on how xTadOffsets are defined
|
|
const X* xTadPtr = reinterpret_cast<const X*>(vx) + xTadOffsets[blockIdx.x];
|
|
|
|
// Compute y coordinates from blockIdx.x
|
|
LongType yCoords[SD_MAX_RANK];
|
|
LongType yOffset;
|
|
INDEX2COORDS(blockIdx.x, shared_yRank, shared_yShape, yCoords);
|
|
COORDS2INDEX(shared_yRank, shared_yStride, yCoords, yOffset);
|
|
|
|
// Retrieve the index from y at the computed offset
|
|
idx = y[yOffset];
|
|
|
|
// Compute coordinates and offset for xTad using idx
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType xOffset;
|
|
INDEX2COORDS(idx, shared_xTadRank, shared_xTadShape, xCoords);
|
|
COORDS2INDEX(shared_xTadRank, shared_xTadStride, xCoords, xOffset);
|
|
|
|
// Store the element to compare
|
|
elemToCompare = xTadPtr[xOffset];
|
|
}
|
|
|
|
// Ensure all threads have access to the cached values
|
|
__syncthreads();
|
|
|
|
// Initialize shared memory for reduction
|
|
extern __shared__ LongType sharedMem[];
|
|
sharedMem[threadIdx.x] = 0;
|
|
__syncthreads();
|
|
|
|
// Pointer to xTad data
|
|
const X* xTad = reinterpret_cast<const X*>(vx) + xTadOffsets[blockIdx.x];
|
|
|
|
// Iterate over xTad elements using cached shape info
|
|
for (LongType i = threadIdx.x; i < xTadLen; i += blockDim.x) {
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType xOffset;
|
|
|
|
// Use cached rank, shape, and stride
|
|
INDEX2COORDS(i, shared_xTadRank, shared_xTadShape, xCoords);
|
|
COORDS2INDEX(shared_xTadRank, shared_xTadStride, xCoords, xOffset);
|
|
|
|
// Compare and update shared memory
|
|
if (elemToCompare < xTad[xOffset]) {
|
|
sharedMem[threadIdx.x]++;
|
|
}
|
|
}
|
|
|
|
// Ensure all threads have completed the counting
|
|
__syncthreads();
|
|
|
|
// Perform parallel reduction to sum counts
|
|
for (LongType activeThreads = blockDim.x / 2; activeThreads > 0; activeThreads /= 2) {
|
|
if (threadIdx.x < activeThreads) {
|
|
sharedMem[threadIdx.x] += sharedMem[threadIdx.x + activeThreads];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
// Write the result to z using cached shape info
|
|
if (threadIdx.x == 0) {
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType zOffset;
|
|
|
|
// Compute z coordinates from blockIdx.x
|
|
INDEX2COORDS(blockIdx.x, shared_zRank, shared_zShape, zCoords);
|
|
COORDS2INDEX(shared_zRank, shared_zStride, zCoords, zOffset);
|
|
|
|
// Compare the aggregated count with k and store the result
|
|
z[zOffset] = (sharedMem[0] < k);
|
|
}
|
|
}
|
|
//////////////////////////////////////////////////////////////
|
|
template <typename X, typename Y>
|
|
static void inTopKCudaLauncher(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* xTadShapeInfo, const LongType* xTadOffsets,
|
|
const LongType k) {
|
|
inTopKCuda<X, Y><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz,
|
|
zShapeInfo, xTadShapeInfo, xTadOffsets, k);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "inTopKCudaLauncher failed");
|
|
|
|
}
|
|
|
|
///////////////////////////////////////////////////////////////////
|
|
Status inTopKFunctor(LaunchContext* context, NDArray* predictions, NDArray* targets,
|
|
NDArray* output, const LongType k) {
|
|
PointersManager manager(context, "in_top_k");
|
|
|
|
const auto packX = ConstantTadHelper::getInstance().tadForDimensions(predictions->shapeInfo(), {1});
|
|
|
|
dim3 topkDims2 = topkDims(packX->numberOfTads());
|
|
const auto xType = predictions->dataType();
|
|
const auto yType = targets->dataType();
|
|
|
|
NDArray::prepareSpecialUse({output}, {predictions, targets});
|
|
BUILD_DOUBLE_SELECTOR(
|
|
xType, yType, inTopKCudaLauncher,
|
|
(topkDims2.y,topkDims2.x, topkDims2.z, context->getCudaStream(), predictions->specialBuffer(),
|
|
predictions->specialShapeInfo(), targets->specialBuffer(), targets->specialShapeInfo(), output->specialBuffer(),
|
|
output->specialShapeInfo(), packX->specialShapeInfo(), packX->specialOffsets(), k),
|
|
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {predictions, targets});
|
|
|
|
manager.synchronize();
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
template <typename X, typename Y>
|
|
static SD_KERNEL void topValuesMover(void const* vx, LongType const* xTadShapeInfo, LongType const* xTadOffsets,
|
|
void const* vi, LongType const* iTadShapeInfo, LongType const* iTadOffsets,
|
|
void* vz, LongType const* zTadShapeInfo, LongType const* zTadOffsets,
|
|
LongType tadLength, int numTads, int k) {
|
|
// Cache shape information in shared memory
|
|
__shared__ int xRank, iRank, zRank;
|
|
__shared__ LongType *xShape, *iShape, *zShape;
|
|
__shared__ LongType *xStride, *iStride, *zStride;
|
|
|
|
if (threadIdx.x == 0) {
|
|
// Cache ranks
|
|
xRank = shape::rank(xTadShapeInfo);
|
|
iRank = shape::rank(iTadShapeInfo);
|
|
zRank = shape::rank(zTadShapeInfo);
|
|
|
|
// Cache shapes
|
|
xShape = shape::shapeOf(xTadShapeInfo);
|
|
iShape = shape::shapeOf(iTadShapeInfo);
|
|
zShape = shape::shapeOf(zTadShapeInfo);
|
|
|
|
// Cache strides
|
|
xStride = shape::stride(xTadShapeInfo);
|
|
iStride = shape::stride(iTadShapeInfo);
|
|
zStride = shape::stride(zTadShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int t = blockIdx.x; t < numTads; t += gridDim.x) {
|
|
auto x = reinterpret_cast<X const*>(vx) + xTadOffsets[t];
|
|
auto i = reinterpret_cast<Y const*>(vi) + iTadOffsets[t];
|
|
auto z = reinterpret_cast<X*>(vz) + zTadOffsets[t];
|
|
|
|
LongType iCoords[SD_MAX_RANK];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType iOffset;
|
|
LongType zOffset;
|
|
LongType xOffset;
|
|
|
|
for (int e = threadIdx.x; e < k; e += blockDim.x) {
|
|
INDEX2COORDS(e, iRank, iShape, iCoords);
|
|
COORDS2INDEX(iRank, iStride, iCoords, iOffset);
|
|
auto idx = i[iOffset];
|
|
|
|
INDEX2COORDS(e, zRank, zShape, zCoords);
|
|
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
|
|
|
INDEX2COORDS(idx, xRank, xShape, xCoords);
|
|
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
|
|
|
z[zOffset] = x[xOffset];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename X, typename Y>
|
|
static SD_KERNEL void indicesAlongDimension(void const* vx, LongType const* xTadShapeInfo, LongType const* xTadOffsets, void* vi, LongType const* iTadShapeInfo, LongType const* iTadOffsets,
|
|
void* vz, LongType const* zTadShapeInfo, LongType const* zTadOffsets,
|
|
LongType tadLength, int numTads, int k,
|
|
int scanWidth, bool needSort) {
|
|
extern __shared__ char _shmem[];
|
|
|
|
X* tempValues = reinterpret_cast<X*>(_shmem) + threadIdx.x * scanWidth;
|
|
Y* tempIndices =
|
|
reinterpret_cast<Y*>(reinterpret_cast<X*>(_shmem) + blockDim.x * scanWidth) + threadIdx.x * scanWidth;
|
|
|
|
// Cache shape information in shared memory
|
|
__shared__ int xRank, iRank, zRank;
|
|
__shared__ LongType *xShape, *iShape, *zShape;
|
|
__shared__ LongType *xStride, *iStride, *zStride;
|
|
__shared__ X localMaximum;
|
|
|
|
if (threadIdx.x == 0) {
|
|
localMaximum = -DataTypeUtils::max<X>();
|
|
|
|
// Cache ranks
|
|
xRank = shape::rank(xTadShapeInfo);
|
|
iRank = shape::rank(iTadShapeInfo);
|
|
zRank = shape::rank(zTadShapeInfo);
|
|
|
|
// Cache shapes
|
|
xShape = shape::shapeOf(xTadShapeInfo);
|
|
iShape = shape::shapeOf(iTadShapeInfo);
|
|
zShape = shape::shapeOf(zTadShapeInfo);
|
|
|
|
// Cache strides
|
|
xStride = shape::stride(xTadShapeInfo);
|
|
iStride = shape::stride(iTadShapeInfo);
|
|
zStride = shape::stride(zTadShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int t = blockIdx.x; t < numTads; t += gridDim.x) {
|
|
auto x = reinterpret_cast<X const*>(vx) + xTadOffsets[t];
|
|
auto i = reinterpret_cast<Y*>(vi) + iTadOffsets[t];
|
|
auto z = reinterpret_cast<X*>(vz) + zTadOffsets[t];
|
|
|
|
// we'll do multiple reads here
|
|
for (int p = 0; p < k; p += scanWidth) {
|
|
// resetting temporary storage
|
|
for (int p = 0; p < scanWidth; p++) {
|
|
tempValues[p] = -DataTypeUtils::max<X>();
|
|
tempIndices[p] = DataTypeUtils::max<Y>();
|
|
}
|
|
|
|
// local max values/indices
|
|
for (int e = threadIdx.x; e < tadLength; e++) {
|
|
LongType xCoords[SD_MAX_RANK];
|
|
LongType xOffset;
|
|
INDEX2COORDS(e, xRank, xShape, xCoords);
|
|
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
|
auto value = x[xOffset];
|
|
|
|
// we'll compare this value to current stored ones
|
|
for (int f = 0; f < scanWidth; f++) {
|
|
if (value > tempValues[f] && (p == 0 || value < localMaximum)) {
|
|
tempValues[f] = value;
|
|
tempIndices[f] = e;
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// at this point we have local part ready for merge and define global maximum for this iteration
|
|
for (LongType activeThreads = blockDim.x / 2; activeThreads > 0; activeThreads /= 2) {
|
|
if (threadIdx.x < activeThreads) {
|
|
if (tempValues[0] < tempValues[0 + activeThreads * scanWidth]) {
|
|
tempValues[0] = tempValues[0 + activeThreads * scanWidth];
|
|
tempIndices[0] = tempIndices[0 + activeThreads * scanWidth];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
__syncthreads();
|
|
|
|
// at this point we know local minimum for next iteration
|
|
if (threadIdx.x == 0) {
|
|
localMaximum = tempValues[scanWidth - 1];
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType zOffset;
|
|
INDEX2COORDS(p, zRank, zShape, zCoords);
|
|
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
|
z[zOffset] = tempValues[scanWidth - 1];
|
|
LongType iCoords[SD_MAX_RANK];
|
|
LongType iOffset;
|
|
INDEX2COORDS(p, iRank, iShape, iCoords);
|
|
COORDS2INDEX(iRank, iStride, iCoords, iOffset);
|
|
i[iOffset] = tempIndices[scanWidth - 1];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
__syncthreads();
|
|
if (!needSort) {
|
|
// if we don't need sort, we need to return values based on their indices (ascending)
|
|
for (int m = 0; m < k; m++) {
|
|
if (m % 2 == 0) {
|
|
for (int tid = threadIdx.x; tid < k; tid += blockDim.x) {
|
|
auto top = 2 * tid + 1;
|
|
if (top < k) {
|
|
LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
|
LongType t0Offset, t1Offset;
|
|
|
|
INDEX2COORDS(top - 1, iRank, iShape, t0Coords);
|
|
COORDS2INDEX(iRank, iStride, t0Coords, t0Offset);
|
|
INDEX2COORDS(top, iRank, iShape, t1Coords);
|
|
COORDS2INDEX(iRank, iStride, t1Coords, t1Offset);
|
|
|
|
if (i[t0Offset] > i[t1Offset]) {
|
|
// swap indices first
|
|
Y di0 = i[t0Offset];
|
|
i[t0Offset] = i[t1Offset];
|
|
i[t1Offset] = di0;
|
|
|
|
// swap values next
|
|
LongType zT0Coords[SD_MAX_RANK], zT1Coords[SD_MAX_RANK];
|
|
LongType zT0Offset, zT1Offset;
|
|
|
|
INDEX2COORDS(top - 1, zRank, zShape, zT0Coords);
|
|
COORDS2INDEX(zRank, zStride, zT0Coords, zT0Offset);
|
|
INDEX2COORDS(top, zRank, zShape, zT1Coords);
|
|
COORDS2INDEX(zRank, zStride, zT1Coords, zT1Offset);
|
|
|
|
X dz0 = z[zT0Offset];
|
|
z[zT0Offset] = z[zT1Offset];
|
|
z[zT1Offset] = dz0;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
for (int tid = threadIdx.x; tid < k; tid += blockDim.x) {
|
|
auto top = 2 * tid + 2;
|
|
if (top < k) {
|
|
LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
|
LongType t0Offset, t1Offset;
|
|
|
|
INDEX2COORDS(top - 1, iRank, iShape, t0Coords);
|
|
COORDS2INDEX(iRank, iStride, t0Coords, t0Offset);
|
|
INDEX2COORDS(top, iRank, iShape, t1Coords);
|
|
COORDS2INDEX(iRank, iStride, t1Coords, t1Offset);
|
|
|
|
if (i[t0Offset] > i[t1Offset]) {
|
|
// swap indices first
|
|
Y di0 = i[t0Offset];
|
|
i[t0Offset] = i[t1Offset];
|
|
i[t1Offset] = di0;
|
|
|
|
// swap values next
|
|
LongType zT0Coords[SD_MAX_RANK], zT1Coords[SD_MAX_RANK];
|
|
LongType zT0Offset, zT1Offset;
|
|
|
|
INDEX2COORDS(top - 1, zRank, zShape, zT0Coords);
|
|
COORDS2INDEX(zRank, zStride, zT0Coords, zT0Offset);
|
|
INDEX2COORDS(top, zRank, zShape, zT1Coords);
|
|
COORDS2INDEX(zRank, zStride, zT1Coords, zT1Offset);
|
|
|
|
X dz0 = z[zT0Offset];
|
|
z[zT0Offset] = z[zT1Offset];
|
|
z[zT1Offset] = dz0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
template <typename X, typename Y>
|
|
static Status topKFunctor_(LaunchContext* context, NDArray* input, NDArray* values, NDArray* indices,
|
|
const LongType k, bool needSort) {
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(input->shapeInfo(), {input->rankOf() - 1});
|
|
auto packI = ConstantTadHelper::getInstance().tadForDimensions(indices->shapeInfo(), {input->rankOf() - 1});
|
|
auto packZ = ConstantTadHelper::getInstance().tadForDimensions(values->shapeInfo(), {input->rankOf() - 1});
|
|
|
|
auto tadLength = shape::length(packX->primaryShapeInfo());
|
|
|
|
// we get top K values first
|
|
if (k == 1) {
|
|
std::vector<LongType> dims = {input->rankOf() - 1};
|
|
input->applyIndexReduce(indexreduce::IndexMax, indices, &dims);
|
|
|
|
dim3 launchDims = getLaunchDims("top_k_mover");
|
|
// copy values on specified indices
|
|
topValuesMover<X, Y><<<launchDims.y, launchDims.x, launchDims.z, *context->getCudaStream()>>>(
|
|
input->specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), indices->specialBuffer(),
|
|
packI->platformShapeInfo(), packI->platformOffsets(), values->specialBuffer(), packZ->platformShapeInfo(),
|
|
packZ->platformOffsets(), tadLength, packX->numberOfTads(), k);
|
|
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "topValuesMover failed");
|
|
|
|
} else {
|
|
int scanWidth = 1;
|
|
dim3 topKIndices2 = topKIndices(scanWidth, sizeof(X), sizeof(Y));
|
|
indicesAlongDimension<X, Y><<<topKIndices2.y, topKIndices2.x, topKIndices2.z, *context->getCudaStream()>>>(
|
|
input->specialBuffer(), packX->platformShapeInfo(), packX->platformOffsets(), indices->specialBuffer(),
|
|
packI->platformShapeInfo(), packI->platformOffsets(), values->specialBuffer(), packZ->platformShapeInfo(),
|
|
packZ->platformOffsets(), tadLength, packX->numberOfTads(), k, scanWidth, needSort);
|
|
sd::DebugHelper::checkErrorCode(context->getCudaStream(), "indicesAlongDimension failed");
|
|
|
|
}
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
Status topKFunctor(LaunchContext* context, NDArray* input, NDArray* values, NDArray* indices,
|
|
const LongType k, bool needSort) {
|
|
input->syncToDevice();
|
|
|
|
BUILD_DOUBLE_SELECTOR(input->dataType(), indices->dataType(), topKFunctor_,
|
|
(context, input, values, indices, k, needSort), SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
|
|
|
values->tickWriteDevice();
|
|
indices->tickWriteDevice();
|
|
|
|
return Status::OK;
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|