wehub-resource-sync
295 lines
8.5 KiB
Plaintext
295 lines
8.5 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 raver119@gmail.com
|
|
// @author Yurii Shyrma, created on 28.11.2018
|
|
//
|
|
#include <ops/specials_cuda.h>
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Y>
|
|
SD_KERNEL void bitonicArbitraryStepKernelKey(
|
|
void* vx,
|
|
const sd::LongType* xShapeInfo,
|
|
void* vy,
|
|
const sd::LongType* yShapeInfo,
|
|
int window,
|
|
int length,
|
|
int reverse,
|
|
bool descending) {
|
|
|
|
auto x = static_cast<X*>(vx);
|
|
auto y = static_cast<Y*>(vy);
|
|
const int tid = threadIdx.x + blockDim.x * blockIdx.x;
|
|
const int half = window >> 1;
|
|
|
|
__shared__ sd::LongType xRank;
|
|
__shared__ const sd::LongType* xShapePtr;
|
|
__shared__ const sd::LongType* xStridePtr;
|
|
|
|
__shared__ sd::LongType yRank; // Potentially unused for direct indexing, but let's keep the pattern consistent
|
|
__shared__ const sd::LongType* yShapePtr;
|
|
__shared__ const sd::LongType* yStridePtr;
|
|
|
|
__shared__ sd::LongType xLength;
|
|
|
|
if (threadIdx.x == 0) {
|
|
xRank = shape::rank(xShapeInfo);
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
|
|
yRank = shape::rank(yShapeInfo);
|
|
yShapePtr = shape::shapeOf(yShapeInfo);
|
|
yStridePtr = shape::stride(yShapeInfo);
|
|
|
|
xLength = shape::length(xShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
const int WARP_SIZE = 32;
|
|
const int numWarps = (gridDim.x * blockDim.x) / WARP_SIZE;
|
|
const int warpId = tid / WARP_SIZE;
|
|
const int warpIdx = tid % WARP_SIZE;
|
|
|
|
int firstPosition;
|
|
int firstStep;
|
|
int secondPosition;
|
|
int secondStep;
|
|
|
|
if (half >= 128) {
|
|
firstPosition = blockIdx.x * window;
|
|
firstStep = gridDim.x * window;
|
|
|
|
secondPosition = threadIdx.x;
|
|
secondStep = blockDim.x;
|
|
}
|
|
else if (half >= 32) {
|
|
firstPosition = warpId * window;
|
|
firstStep = numWarps * window;
|
|
|
|
secondPosition = warpIdx;
|
|
secondStep = WARP_SIZE;
|
|
}
|
|
else {
|
|
firstPosition = tid * window;
|
|
firstStep = blockDim.x * gridDim.x * window;
|
|
|
|
secondPosition = 0;
|
|
secondStep = 1;
|
|
}
|
|
|
|
for (int i = firstPosition; i < length; i += firstStep) {
|
|
for (int j = secondPosition; j < half; j += secondStep) {
|
|
const int it = (reverse) ? i + j + half : i + window - j - 1;
|
|
const int ij = i + j;
|
|
if (it < length && ij < length) {
|
|
sd::LongType itCoords[SD_MAX_RANK];
|
|
sd::LongType ijCoords[SD_MAX_RANK];
|
|
sd::LongType itOffset;
|
|
sd::LongType ijOffset;
|
|
|
|
INDEX2COORDS(it, xRank, xShapePtr, itCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, itCoords, itOffset);
|
|
|
|
INDEX2COORDS(ij, xRank, xShapePtr, ijCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, ijCoords, ijOffset);
|
|
|
|
X v0 = x[ijOffset];
|
|
X v1 = x[itOffset];
|
|
|
|
const bool condition = (!descending == (v0 > v1));
|
|
if (condition) {
|
|
x[ijOffset] = v1;
|
|
x[itOffset] = v0;
|
|
|
|
sd::LongType itCoordsY[SD_MAX_RANK];
|
|
sd::LongType ijCoordsY[SD_MAX_RANK];
|
|
sd::LongType itOffsetY;
|
|
sd::LongType ijOffsetY;
|
|
|
|
INDEX2COORDS(it, yRank, yShapePtr, itCoordsY);
|
|
COORDS2INDEX(yRank, yStridePtr, itCoordsY, itOffsetY);
|
|
|
|
INDEX2COORDS(ij, yRank, yShapePtr, ijCoordsY);
|
|
COORDS2INDEX(yRank, yStridePtr, ijCoordsY, ijOffsetY);
|
|
|
|
Y ytemp = y[ijOffsetY];
|
|
y[ijOffsetY] = y[itOffsetY];
|
|
y[itOffsetY] = ytemp;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_KERNEL void execBitonicArbitraryStepKernel(
|
|
void* vx,
|
|
const sd::LongType* xShapeInfo,
|
|
int window,
|
|
int length,
|
|
int reverse,
|
|
bool descending) {
|
|
|
|
auto x = static_cast<T*>(vx);
|
|
const int tid = threadIdx.x + blockDim.x * blockIdx.x;
|
|
const int half = window >> 1;
|
|
|
|
__shared__ sd::LongType xRank;
|
|
__shared__ const sd::LongType* xShapePtr;
|
|
__shared__ const sd::LongType* xStridePtr;
|
|
|
|
__shared__ sd::LongType xLength;
|
|
|
|
// We'll omit using shared memory for x data except for small merges,
|
|
// but keep the pattern of caching shape info
|
|
if (threadIdx.x == 0) {
|
|
xRank = shape::rank(xShapeInfo);
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
|
|
xLength = shape::length(xShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
const int WARP_SIZE = 32;
|
|
const int numWarps = (gridDim.x * blockDim.x) / WARP_SIZE;
|
|
const int warpId = tid / WARP_SIZE;
|
|
const int warpIdx = tid % WARP_SIZE;
|
|
|
|
int firstPosition;
|
|
int firstStep;
|
|
int secondPosition;
|
|
int secondStep;
|
|
|
|
if (half >= 128) {
|
|
firstPosition = blockIdx.x * window;
|
|
firstStep = gridDim.x * window;
|
|
|
|
secondPosition = threadIdx.x;
|
|
secondStep = blockDim.x;
|
|
}
|
|
else if (half >= 32) {
|
|
firstPosition = warpId * window;
|
|
firstStep = numWarps * window;
|
|
|
|
secondPosition = warpIdx;
|
|
secondStep = WARP_SIZE;
|
|
}
|
|
else {
|
|
firstPosition = tid * window;
|
|
firstStep = blockDim.x * gridDim.x * window;
|
|
|
|
secondPosition = 0;
|
|
secondStep = 1;
|
|
}
|
|
|
|
for (int i = firstPosition; i < length; i += firstStep) {
|
|
for (int j = secondPosition; j < half; j += secondStep) {
|
|
const int it = (reverse) ? i + j + half : i + window - j - 1;
|
|
const int ij = i + j;
|
|
if (it < length && ij < length) {
|
|
sd::LongType itCoords[SD_MAX_RANK];
|
|
sd::LongType ijCoords[SD_MAX_RANK];
|
|
sd::LongType itOffset;
|
|
sd::LongType ijOffset;
|
|
|
|
INDEX2COORDS(it, xRank, xShapePtr, itCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, itCoords, itOffset);
|
|
|
|
INDEX2COORDS(ij, xRank, xShapePtr, ijCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, ijCoords, ijOffset);
|
|
|
|
T v0 = x[ijOffset];
|
|
T v1 = x[itOffset];
|
|
|
|
const bool condition = (!descending == (v0 > v1));
|
|
if (condition) {
|
|
x[ijOffset] = v1;
|
|
x[itOffset] = v0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
SD_HOST void bitonicArbitraryStepGeneric(
|
|
dim3 &launchDims,
|
|
cudaStream_t *stream,
|
|
void* vx,
|
|
const sd::LongType* xShapeInfo,
|
|
int window,
|
|
int length,
|
|
int reverse,
|
|
bool descending) {
|
|
|
|
execBitonicArbitraryStepKernel<T>
|
|
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
|
vx,
|
|
xShapeInfo,
|
|
window,
|
|
length,
|
|
reverse,
|
|
descending);
|
|
|
|
sd::DebugHelper::checkErrorCode(stream, "execBitonicArbitraryStepKernel failed");
|
|
}
|
|
|
|
template <typename X, typename Y>
|
|
SD_HOST void bitonicArbitraryStepGenericKey(
|
|
dim3 &launchDims,
|
|
cudaStream_t *stream,
|
|
void* vx,
|
|
const sd::LongType* xShapeInfo,
|
|
void* vy,
|
|
const sd::LongType* yShapeInfo,
|
|
int window,
|
|
int length,
|
|
int reverse,
|
|
bool descending) {
|
|
|
|
bitonicArbitraryStepKernelKey<X, Y>
|
|
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
|
vx,
|
|
xShapeInfo,
|
|
vy,
|
|
yShapeInfo,
|
|
window,
|
|
length,
|
|
reverse,
|
|
descending);
|
|
|
|
sd::DebugHelper::checkErrorCode(stream, "bitonicArbitraryStepKernelKey failed");
|
|
}
|
|
|
|
BUILD_SINGLE_TEMPLATE(
|
|
void bitonicArbitraryStepGeneric,
|
|
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo, int window,
|
|
int length, int reverse, bool descending),
|
|
SD_COMMON_TYPES);
|
|
|
|
BUILD_DOUBLE_TEMPLATE(
|
|
void bitonicArbitraryStepGenericKey,
|
|
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo, void *vy,
|
|
sd::LongType const *yShapeInfo, int window, int length, int reverse, bool descending),
|
|
SD_COMMON_TYPES, SD_COMMON_TYPES);
|