chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,441 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <system/op_boilerplate.h>
|
||||
#include <loops/broadcasting.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <types/types.h>
|
||||
#include <system/Environment.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <string>
|
||||
#include <stdexcept>
|
||||
#include <helpers/StringUtils.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
template<typename X, typename Y, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastSimple(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
functions::broadcast::Broadcast<X,Y,Z>::template transformCuda<OpClass>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ);
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastSimple(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* y,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
functions::broadcast::Broadcast<X,Y,Z>::template transformCuda<OpClass>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo);
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastInverseSimple(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, long long int dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
functions::broadcast::Broadcast<X,Y,Z>::template transformInverseCuda<OpClass>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace broadcast {
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void Broadcast<X,Y,Z>::intermediateBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
void const* x, sd::LongType const* xShapeInfo,
|
||||
void const* y, sd::LongType const* yShapeInfo,
|
||||
void* z, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
broadcastSimple<X, Y, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ);
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void Broadcast<X,Y,Z>::intermediateBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
const void* x, const sd::LongType* xShapeInfo,
|
||||
const void* y, const sd::LongType* yShapeInfo,
|
||||
void* z, const sd::LongType* zShapeInfo) {
|
||||
|
||||
broadcastSimple<X, Y, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo);
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
SD_HOST void Broadcast<X,Y,Z>::execBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x, sd::LongType const* xShapeInfo,
|
||||
void const* y, sd::LongType const* yShapeInfo,
|
||||
void* z, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateBroadcast,
|
||||
PARAMS(launchDims, stream,
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ),
|
||||
OPS_A(BROADCAST_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "execBroadcast(...) failed");
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
SD_HOST void Broadcast<X,Y,Z>::execBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x, const sd::LongType* xShapeInfo,
|
||||
const void* y, const sd::LongType* yShapeInfo,
|
||||
void* z, const sd::LongType* zShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateBroadcast,
|
||||
PARAMS(launchDims, stream,
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo),
|
||||
OPS_A(BROADCAST_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcast(...) failed");
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void Broadcast<X,Y,Z>::intermediateInverseBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
void const* x, sd::LongType const* xShapeInfo,
|
||||
void const* y, sd::LongType const* yShapeInfo,
|
||||
void* z, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
broadcastInverseSimple<X, Y, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ);
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
SD_HOST void Broadcast<X,Y,Z>::execInverseBroadcast(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x, sd::LongType const* xShapeInfo,
|
||||
void const* y, sd::LongType const* yShapeInfo,
|
||||
void* z, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateInverseBroadcast,
|
||||
PARAMS(launchDims, stream,
|
||||
x, xShapeInfo,
|
||||
y, yShapeInfo,
|
||||
z, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ),
|
||||
OPS_A(BROADCAST_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "execInverseBroadcast(...) failed");
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Broadcast<X,Y,Z>::transformInverseCuda(
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
auto x = reinterpret_cast<X const*>(vx);
|
||||
auto y = reinterpret_cast<Y const*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ int numTads;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShape;
|
||||
__shared__ const sd::LongType* xStride;
|
||||
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ const sd::LongType* yShape;
|
||||
__shared__ const sd::LongType* yStride;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShape;
|
||||
__shared__ const sd::LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(yShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
|
||||
yRank = shape::rank(tadOnlyShapeInfo);
|
||||
yShape = shape::shapeOf(tadOnlyShapeInfo);
|
||||
yStride = shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
zRank = shape::rank(tadOnlyShapeInfoZ);
|
||||
zShape = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
zStride = shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto rY = y + tadOffsets[r];
|
||||
auto rZ = z + tadOffsetsZ[r];
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShape, yCoords);
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
|
||||
rZ[zOffset] = OpType::op(x[xOffset], rY[yOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Broadcast<X,Y,Z>::transformCuda(
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ, sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
auto x = reinterpret_cast<X const*>(vx);
|
||||
auto y = reinterpret_cast<Y const*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShape;
|
||||
__shared__ const sd::LongType* xStride;
|
||||
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ const sd::LongType* yShape;
|
||||
__shared__ const sd::LongType* yStride;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShape;
|
||||
__shared__ const sd::LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(tadOnlyShapeInfo);
|
||||
xShape = shape::shapeOf(tadOnlyShapeInfo);
|
||||
xStride = shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(tadOnlyShapeInfoZ);
|
||||
zShape = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
zStride = shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto rX = x + tadOffsets[r];
|
||||
auto rZ = z + tadOffsetsZ[r];
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShape, coordsX);
|
||||
COORDS2INDEX(xRank, xStride, coordsX, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShape, coordsY);
|
||||
COORDS2INDEX(yRank, yStride, coordsY, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShape, coordsZ);
|
||||
COORDS2INDEX(zRank, zStride, coordsZ, zOffset);
|
||||
|
||||
rZ[zOffset] = OpType::op(rX[xOffset], y[yOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Broadcast<X,Y,Z>::transformCuda(
|
||||
const void* vx, const sd::LongType* xShapeInfo,
|
||||
const void* vy, const sd::LongType* yShapeInfo,
|
||||
void* vz, const sd::LongType* zShapeInfo) {
|
||||
|
||||
const X* x = reinterpret_cast<const X*>(vx);
|
||||
const Y* y = reinterpret_cast<const Y*>(vy);
|
||||
Z* z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ sd::LongType zLen;
|
||||
__shared__ int rank;
|
||||
__shared__ bool xzSameOffsets, yzSameOffsets;
|
||||
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = shape::rank(zShapeInfo);
|
||||
xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
|
||||
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (sd::LongType i = tid; i < zLen; i += blockDim.x * gridDim.x) {
|
||||
INDEX2COORDS(i, rank, zShapePtr, coords);
|
||||
|
||||
sd::LongType zOffset;
|
||||
COORDS2INDEX(rank, zStridePtr, coords, zOffset);
|
||||
|
||||
sd::LongType xOffset, yOffset;
|
||||
if (xzSameOffsets) {
|
||||
xOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(rank, xStridePtr, coords, xOffset);
|
||||
}
|
||||
|
||||
if (yzSameOffsets) {
|
||||
yOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(rank, yStridePtr, coords, yOffset);
|
||||
}
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], y[yOffset]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
} // namespace broadcast
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,570 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <helpers/StringUtils.h>
|
||||
#include <loops/broadcasting_bool.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastBoolSimple(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
functions::broadcast::BroadcastBool<X, Z>::template transformCuda<OpClass>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastBoolSimple(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* y,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* extraParams) {
|
||||
|
||||
functions::broadcast::BroadcastBool<X, Z>::template transformCuda<OpClass>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpClass>
|
||||
static SD_KERNEL void broadcastBoolInverseSimple(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
functions::broadcast::BroadcastBool<X, Z>::template transformInverseCuda<OpClass>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace broadcast {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastBool<X, Z>::intermediateBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
broadcastBoolSimple<X, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastBool<X, Z>::intermediateBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* y,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* extraParams) {
|
||||
|
||||
broadcastBoolSimple<X, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void BroadcastBool<X, Y>::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateBroadcast,
|
||||
PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams,
|
||||
dimension, dimensionLength, tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ),
|
||||
OPS_A(BROADCAST_BOOL_OPS));
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "execBroadcast(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void BroadcastBool<X, Y>::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* y,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* extraParams) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateBroadcast,
|
||||
PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams),
|
||||
OPS_A(BROADCAST_BOOL_OPS));
|
||||
|
||||
DEBUG_KERNEL(stream, opNum);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastBool<X, Z>::intermediateInverseBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
broadcastBoolInverseSimple<X, Z, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
extraParams,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcastBool(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void BroadcastBool<X, Y>::execInverseBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateInverseBroadcast,
|
||||
PARAMS(launchDims, stream, x, xShapeInfo, y, yShapeInfo, z, zShapeInfo, extraParams,
|
||||
dimension, dimensionLength, tadOnlyShapeInfo, tadOffsets,
|
||||
tadOnlyShapeInfoZ, tadOffsetsZ),
|
||||
OPS_A(BROADCAST_BOOL_OPS));
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "execInverseBroadcast(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void BroadcastBool<X, Z>::transformInverseCuda(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams= reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ int numTads;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
|
||||
__shared__ sd::LongType tadRankZ;
|
||||
__shared__ const sd::LongType* tadShapePtrZ;
|
||||
__shared__ const sd::LongType* tadStridePtrZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(yShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr= shape::stride(xShapeInfo);
|
||||
|
||||
tadRank = shape::rank(tadOnlyShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadOnlyShapeInfo);
|
||||
tadStridePtr= shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
tadRankZ = shape::rank(tadOnlyShapeInfoZ);
|
||||
tadShapePtrZ = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
tadStridePtrZ= shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto rZ = z + tadOffsetsZ[r];
|
||||
auto rY = y + tadOffsets[r];
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
// for x
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
|
||||
|
||||
// for y (tad)
|
||||
INDEX2COORDS(i, tadRank, tadShapePtr, coordsY);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, coordsY, yOffset);
|
||||
|
||||
// for z (tadZ)
|
||||
INDEX2COORDS(i, tadRankZ, tadShapePtrZ, coordsZ);
|
||||
COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, zOffset);
|
||||
|
||||
rZ[zOffset] = OpType::op(x[xOffset], rY[yOffset], extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void BroadcastBool<X, Z>::transformCuda(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams= reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(tadOnlyShapeInfo);
|
||||
xShapePtr = shape::shapeOf(tadOnlyShapeInfo);
|
||||
xStridePtr= shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr= shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(tadOnlyShapeInfoZ);
|
||||
zShapePtr = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
zStridePtr= shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto rX = x + tadOffsets[r];
|
||||
auto rZ = z + tadOffsetsZ[r];
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, zOffset);
|
||||
|
||||
rZ[zOffset] = OpType::op(rX[xOffset], y[yOffset], extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void BroadcastBool<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* vy,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
const auto x = reinterpret_cast<const X*>(vx);
|
||||
const auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ sd::LongType zLen;
|
||||
__shared__ int xRank, yRank, zRank;
|
||||
__shared__ bool xzSameOffsets, yzSameOffsets;
|
||||
|
||||
// We'll store stride data for each shape in shared mem
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zLen = shape::length(zShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = blockDim.x * gridDim.x;
|
||||
|
||||
for (sd::LongType i2 = tid; i2 < zLen; i2 += totalThreads) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
INDEX2COORDS(i2, zRank, zShapePtr, coords);
|
||||
|
||||
sd::LongType zOffset;
|
||||
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
||||
|
||||
sd::LongType xOffset;
|
||||
if (xzSameOffsets) {
|
||||
xOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(xRank, xStridePtr, coords, xOffset);
|
||||
}
|
||||
|
||||
sd::LongType yOffset;
|
||||
if (yzSameOffsets) {
|
||||
yOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(yRank, yStridePtr, coords, yOffset);
|
||||
}
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
|
||||
}
|
||||
}
|
||||
|
||||
// build the class
|
||||
BUILD_DOUBLE_TEMPLATE( class BroadcastBool, , SD_COMMON_TYPES, SD_BOOL_TYPES);
|
||||
|
||||
} // namespace broadcast
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,476 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <helpers/StringUtils.h>
|
||||
#include <loops/broadcasting_int.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <loops/pairwise_instantiations.h>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Cached kernel that caches shape info in shared memory and performs the broadcast
|
||||
template <typename X, typename OpClass>
|
||||
__global__ void broadcastIntSimpleCached(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// Delegate the broadcast operation to the transformCuda method with cached shape info
|
||||
functions::broadcast::BroadcastInt<X>::template transformCuda<OpClass>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Cached kernel that caches shape info in shared memory and performs the inverse broadcast
|
||||
template <typename X, typename OpClass>
|
||||
__global__ void broadcastIntInverseSimpleCached(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// Delegate the inverse broadcast operation to the transformInverseCuda method with cached shape info
|
||||
functions::broadcast::BroadcastInt<X>::template transformInverseCuda<OpClass>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace broadcast {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the intermediateBroadcast function that launches the cached kernel with dimensions
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastInt<X>::intermediateBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// Launch the cached broadcastIntSimpleCached kernel with all parameters
|
||||
broadcastIntSimpleCached<X, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkErrorCode(stream, "intermediateBroadcast(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the intermediateBroadcast function that launches the cached kernel without dimensions
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastInt<X>::intermediateBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* y,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo)
|
||||
{
|
||||
// Launch the cached broadcastIntSimpleCached kernel without dimensions
|
||||
broadcastIntSimpleCached<X, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
nullptr, // dimension
|
||||
0, // dimensionLength
|
||||
nullptr, // tadOnlyShapeInfo
|
||||
nullptr, // tadOffsets
|
||||
nullptr, // tadOnlyShapeInfoZ
|
||||
nullptr // tadOffsetsZ
|
||||
);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkGlobalErrorCode("broadcastIntSimpleCached(...) failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the intermediateInverseBroadcast function that launches the cached inverse kernel with dimensions
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_HOST void BroadcastInt<X>::intermediateInverseBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// Launch the cached broadcastIntInverseSimpleCached kernel with all parameters
|
||||
broadcastIntInverseSimpleCached<X, OpClass>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
y,
|
||||
yShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets,
|
||||
tadOnlyShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkGlobalErrorCode("broadcastIntInverseSimpleCached(...) failed");
|
||||
}
|
||||
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the transformCuda device function for BroadcastInt with cached shape info
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_DEVICE void BroadcastInt<X>::transformCuda(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// If tadOnlyShapeInfoZ is null, set it to tadOnlyShapeInfo and tadOffsetsZ to tadOffsets
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
// Cast pointers to appropriate types
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ int numTads;
|
||||
__shared__ int xRank;
|
||||
__shared__ int yRank;
|
||||
__shared__ int zRank;
|
||||
|
||||
__shared__ const sd::LongType* tadShape;
|
||||
__shared__ const sd::LongType* tadStride;
|
||||
__shared__ const sd::LongType* tadShapeZ;
|
||||
__shared__ const sd::LongType* tadStrideZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache essential shape information
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
tadShape = shape::shapeOf(tadOnlyShapeInfo);
|
||||
tadStride = shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
tadShapeZ = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
tadStrideZ = shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each block handles a subset of TADs
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto xTad = x + tadOffsets[r];
|
||||
auto zTad = z + tadOffsetsZ[r];
|
||||
|
||||
// Loop over TAD elements
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType xOffset, yOffset, zOffset;
|
||||
|
||||
// Convert index to coordinates using cached shape info
|
||||
INDEX2COORDS(i, xRank, tadShape, coords);
|
||||
COORDS2INDEX(xRank, tadStride, coords, xOffset);
|
||||
|
||||
COORDS2INDEX(yRank, shape::stride(yShapeInfo), coords, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, tadShapeZ, coords);
|
||||
COORDS2INDEX(zRank, tadStrideZ, coords, zOffset);
|
||||
|
||||
// Apply the operation
|
||||
zTad[zOffset] = OpClass::op(x[xOffset], y[yOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the transformInverseCuda device function for BroadcastInt with cached shape info
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_DEVICE void BroadcastInt<X>::transformInverseCuda(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// If tadOnlyShapeInfoZ is null, set it to tadOnlyShapeInfo and tadOffsetsZ to tadOffsets
|
||||
if (tadOnlyShapeInfoZ == nullptr) {
|
||||
tadOnlyShapeInfoZ = tadOnlyShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
// Cast pointers to appropriate types
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ int numTads;
|
||||
__shared__ int xRank;
|
||||
__shared__ int yRank;
|
||||
__shared__ int zRank;
|
||||
|
||||
__shared__ const sd::LongType* tadShape;
|
||||
__shared__ const sd::LongType* tadStride;
|
||||
__shared__ const sd::LongType* tadShapeZ;
|
||||
__shared__ const sd::LongType* tadStrideZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache essential shape information
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(yShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
|
||||
tadShape = shape::shapeOf(tadOnlyShapeInfo);
|
||||
tadStride = shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
tadShapeZ = shape::shapeOf(tadOnlyShapeInfoZ);
|
||||
tadStrideZ = shape::stride(tadOnlyShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each block handles a subset of TADs
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto zTad = z + tadOffsetsZ[r];
|
||||
auto yTad = y + tadOffsets[r];
|
||||
|
||||
// Loop over TAD elements
|
||||
for (sd::LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
// Derive coordinates and offsets
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType xOffset, yOffset, zOffset;
|
||||
|
||||
// Convert index to coordinates using cached shape info
|
||||
INDEX2COORDS(i, xRank, tadShape, coords);
|
||||
COORDS2INDEX(xRank, tadStride, coords, xOffset);
|
||||
|
||||
COORDS2INDEX(yRank, shape::stride(yShapeInfo), coords, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, tadShapeZ, coords);
|
||||
COORDS2INDEX(zRank, tadStrideZ, coords, zOffset);
|
||||
|
||||
// Apply the inverse operation
|
||||
zTad[zOffset] = OpClass::op(x[xOffset], yTad[yOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the transformCuda device function for BroadcastInt without dimensions
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_DEVICE void BroadcastInt<X>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const void* vy,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo)
|
||||
{
|
||||
const X* x = reinterpret_cast<const X*>(vx);
|
||||
const X* y = reinterpret_cast<const X*>(vy);
|
||||
X* z = reinterpret_cast<X*>(vz);
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType zLen;
|
||||
__shared__ int rank;
|
||||
__shared__ bool xzSameOffsets, yzSameOffsets;
|
||||
|
||||
__shared__ const sd::LongType* xShapeCached;
|
||||
__shared__ const sd::LongType* yShapeCached;
|
||||
__shared__ const sd::LongType* zShapeCached;
|
||||
|
||||
__shared__ const sd::LongType* xStrideCached;
|
||||
__shared__ const sd::LongType* yStrideCached;
|
||||
__shared__ const sd::LongType* zStrideCached;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache essential shape information
|
||||
zLen = shape::length(zShapeInfo);
|
||||
rank = shape::rank(zShapeInfo);
|
||||
|
||||
xzSameOffsets = shape::haveSameShapeAndStrides(xShapeInfo, zShapeInfo);
|
||||
yzSameOffsets = shape::haveSameShapeAndStrides(yShapeInfo, zShapeInfo);
|
||||
|
||||
xShapeCached = shape::shapeOf(xShapeInfo);
|
||||
yShapeCached = shape::shapeOf(yShapeInfo);
|
||||
zShapeCached = shape::shapeOf(zShapeInfo);
|
||||
|
||||
xStrideCached = shape::stride(xShapeInfo);
|
||||
yStrideCached = shape::stride(yShapeInfo);
|
||||
zStrideCached = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = blockDim.x * gridDim.x;
|
||||
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (sd::LongType i = tid; i < zLen; i += totalThreads) {
|
||||
// Quick coordinate transform
|
||||
INDEX2COORDS(i, rank, zShapeCached, coords);
|
||||
|
||||
sd::LongType zOffset, xOffset, yOffset;
|
||||
COORDS2INDEX(rank, zStrideCached, coords, zOffset);
|
||||
|
||||
if (xzSameOffsets) {
|
||||
xOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(rank, xStrideCached, coords, xOffset);
|
||||
}
|
||||
|
||||
if (yzSameOffsets) {
|
||||
yOffset = zOffset;
|
||||
} else {
|
||||
COORDS2INDEX(rank, yStrideCached, coords, yOffset);
|
||||
}
|
||||
|
||||
z[zOffset] = OpClass::op(x[xOffset], y[yOffset]);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
// Instantiate templates for common integer types
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template class BroadcastInt, ,
|
||||
SD_INTEGER_TYPES);
|
||||
|
||||
} // namespace broadcast
|
||||
} // namespace functions
|
||||
|
||||
|
||||
|
||||
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/broadcasting.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function instantiation for Broadcast::execBroadcast with dimension
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo);
|
||||
);
|
||||
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/broadcasting.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function instantiation for Broadcast::execInverseBroadcast
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execInverseBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ);
|
||||
);
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/broadcasting.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* Broadcast::execBroadcast instantiated for types in @COMB1@, @COMB2@, @COMB3@
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::broadcast::Broadcast,
|
||||
::execInverseBroadcast(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* y,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadOnlyShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ);
|
||||
);
|
||||
+81
@@ -0,0 +1,81 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/indexreduce.cu>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* IndexReduce::executeIndexReduce instantiated for types in @COMB1@, @COMB2@
|
||||
* Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_INDEXING_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::indexreduce::IndexReduce,
|
||||
::executeIndexReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
void const* dx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* result,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationBuffer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_INDEXING_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::indexreduce::IndexReduce,
|
||||
::executeIndexReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
void const* dx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* result,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationBuffer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets);
|
||||
);
|
||||
@@ -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
|
||||
* *****************************************************************************
|
||||
*/
|
||||
|
||||
//
|
||||
// Created by agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/pairwise.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function instantiation for PairWiseTransform::executeCudaShaped
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::pairwise_transforms::PairWiseTransform,
|
||||
::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
const void *vy,
|
||||
const sd::LongType *yShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
void *vextraParams);
|
||||
);
|
||||
+49
@@ -0,0 +1,49 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/pairwise.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* PairWiseTransform::executeCudaShaped instantiated for types in @COMB1@, @COMB2@, @COMB3@
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::pairwise_transforms::PairWiseTransform,
|
||||
::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
const void *vy,
|
||||
const sd::LongType *yShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
void *vextraParams);
|
||||
);
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/reduce3.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* Reduce3::exec instantiated for types in @COMB1@, @COMB2@
|
||||
* Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_FLOAT_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce3::Reduce3,
|
||||
::exec(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_FLOAT_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce3::Reduce3,
|
||||
::execAll(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_FLOAT_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce3::Reduce3,
|
||||
::execScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo);
|
||||
);
|
||||
+73
@@ -0,0 +1,73 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/reduce/reduce_float.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* ReduceFloatFunction::execReduce instantiated for types in @COMB1@, @COMB2@
|
||||
* Note: Using @COMB1@ for X type and @COMB2@ for Z type (2-type combinations)
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_NUMERIC_TYPES_PART_@COMB1@,
|
||||
SD_FLOAT_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceFloatFunction,
|
||||
::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* dXShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
const sd::LongType* dims);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
SD_NUMERIC_TYPES_PART_@COMB1@,
|
||||
SD_FLOAT_TYPES_PART_@COMB2@,
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceFloatFunction,
|
||||
::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo);
|
||||
);
|
||||
+55
@@ -0,0 +1,55 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/scalar.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function instantiation for ScalarTransform::executeCudaAlongDimension
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::scalar::ScalarTransform,
|
||||
::executeCudaAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
const void *vscalars,
|
||||
void *vextraParams,
|
||||
sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets,
|
||||
const sd::LongType *tadShapeInfoZ,
|
||||
const sd::LongType *tadOffsetsZ);
|
||||
);
|
||||
@@ -0,0 +1,51 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/scalar.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function instantiation for ScalarTransform::executeCudaShaped
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_SINGLE_TYPE_@COMB1@,
|
||||
SD_SINGLE_TYPE_@COMB2@,
|
||||
SD_SINGLE_TYPE_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::scalar::ScalarTransform,
|
||||
::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
const sd::LongType *hxShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
const sd::LongType *hzShapeInfo,
|
||||
const void *vscalar,
|
||||
void *vextraParams);
|
||||
);
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 agibsonccc on 3/10/25.
|
||||
//
|
||||
|
||||
#include <loops/cuda/scalar.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#include <loops/pairwise_instantiations_single.h>
|
||||
|
||||
/*
|
||||
* Function Instantiation:
|
||||
* ScalarTransform::executeCudaShaped instantiated for types in @COMB1@, @COMB2@, @COMB3@
|
||||
*/
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::scalar::ScalarTransform,
|
||||
::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
const sd::LongType *hxShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
const sd::LongType *hzShapeInfo,
|
||||
const void *vscalar,
|
||||
void *vextraParams);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS_3(
|
||||
SD_COMMON_TYPES_PART_@COMB1@,
|
||||
SD_COMMON_TYPES_PART_@COMB2@,
|
||||
SD_COMMON_TYPES_PART_@COMB3@,
|
||||
INSTANT_PROCESS_COMBINATION_3,
|
||||
functions::scalar::ScalarTransform,
|
||||
::executeCudaAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void *vx,
|
||||
const sd::LongType *xShapeInfo,
|
||||
void *vz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
const void *vscalars,
|
||||
void *vextraParams,
|
||||
sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets,
|
||||
const sd::LongType *tadShapeInfoZ,
|
||||
const sd::LongType *tadOffsetsZ);
|
||||
);
|
||||
@@ -0,0 +1,30 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <loops/cuda/broadcasting.chpp>
|
||||
#cmakedefine SD_PAIRWISE_TYPES_GEN
|
||||
#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
|
||||
namespace functions {
|
||||
namespace broadcast {
|
||||
BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN Broadcast, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,30 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <loops/cuda/pairwise.chpp>
|
||||
#cmakedefine SD_PAIRWISE_TYPES_GEN
|
||||
#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
|
||||
namespace functions {
|
||||
namespace pairwise_transforms {
|
||||
BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN PairWiseTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,30 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <loops/cuda/reduce3.chpp>
|
||||
#cmakedefine SD_FLOAT_TYPES_GEN
|
||||
#if defined(SD_FLOAT_TYPES_GEN) && defined(SD_FLOAT_TYPES_@FL_TYPE_INDEX@)
|
||||
namespace functions {
|
||||
namespace reduce3 {
|
||||
BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN Reduce3, , SD_COMMON_TYPES, SD_FLOAT_TYPES_@FL_TYPE_INDEX@);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,30 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <loops/cuda/reduce/reduce_float.chpp>
|
||||
#cmakedefine SD_FLOAT_TYPES_GEN
|
||||
#if defined(SD_FLOAT_TYPES_GEN) && defined(SD_FLOAT_TYPES_@FL_TYPE_INDEX@)
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
BUILD_DOUBLE_TEMPLATE(template class SD_LIB_HIDDEN ReduceFloatFunction, , SD_COMMON_TYPES, SD_FLOAT_TYPES_@FL_TYPE_INDEX@);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,31 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <loops/cuda/scalar.chpp>
|
||||
#include <system/type_boilerplate.h>
|
||||
#cmakedefine SD_PAIRWISE_TYPES_GEN
|
||||
#if defined(SD_PAIRWISE_TYPES_GEN) && defined(SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@)
|
||||
namespace functions {
|
||||
namespace scalar {
|
||||
BUILD_PAIRWISE_TEMPLATE(extern template class SD_LIB_HIDDEN ScalarTransform, , SD_PAIRWISE_TYPES_@FL_TYPE_INDEX@);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,331 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 raver on 4/9/2018.
|
||||
//
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
//note: keep this. It's required for proper linker work
|
||||
|
||||
#include "../indexreduce.h"
|
||||
#include "../legacy_ops.h"
|
||||
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
template <typename X, typename Z>
|
||||
static SD_KERNEL void simpleIndexReduceGeneric(const int op, void const *dx, sd::LongType const *xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType zRank,
|
||||
sd::LongType *dimension, sd::LongType dimensionLength, int postProcessOrNot, sd::LongType *allocationBuffer, void *reductionBuffer,
|
||||
sd::LongType const *tadOnlyShapeInfo, sd::LongType const *tadOffsets) {
|
||||
functions::indexreduce::IndexReduce<X, Z>::transform(op, dx, xShapeInfo, extraParams, result, zShapeInfo, dimension,
|
||||
dimensionLength, postProcessOrNot, allocationBuffer,
|
||||
reductionBuffer, tadOnlyShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace indexreduce {
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void IndexReduce<X, Z>::executeIndexReduceScalar(
|
||||
dim3 launchDims, cudaStream_t *stream, const int opNum, void const *dx, sd::LongType const *xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType zRank,
|
||||
sd::LongType *dimension, sd::LongType dimensionLength,
|
||||
int postProcessOrNot,sd::LongType *allocationBuffer, void *reductionBuffer, sd::LongType const *tadOnlyShapeInfo,
|
||||
sd::LongType const *tadOffsets) {
|
||||
simpleIndexReduceGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum, dx, xShapeInfo, xRank, extraParams, result, zShapeInfo, 0, nullptr, 0, 1, allocationBuffer,
|
||||
reductionBuffer, tadOnlyShapeInfo, tadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "executeIndexReduceScalar(...) failed");
|
||||
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void IndexReduce<X, Z>::executeIndexReduce(dim3 launchDims,
|
||||
cudaStream_t *stream,
|
||||
const int opNum,
|
||||
void const *dx,
|
||||
sd::LongType const *xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void *extraParams,
|
||||
void *result,
|
||||
sd::LongType const *zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType *allocationBuffer,
|
||||
void *reductionBuffer,
|
||||
sd::LongType const *tadOnlyShapeInfo,
|
||||
sd::LongType const *tadOffsets) {
|
||||
simpleIndexReduceGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum, dx, xShapeInfo, xRank, extraParams, result, zShapeInfo, zRank, dimension, dimensionLength, postProcessOrNot,
|
||||
allocationBuffer, reductionBuffer, tadOnlyShapeInfo, tadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "executeIndexReduce(...) failed");
|
||||
|
||||
}
|
||||
|
||||
// This is the un-specialized struct. Note that we prevent instantiation of this
|
||||
// struct by putting an undefined symbol in the function body so it won't compile.
|
||||
template <typename T>
|
||||
struct SharedIndexValue {
|
||||
// Ensure that we won't compile any un-specialized types
|
||||
SD_DEVICE T *getPointer() {
|
||||
extern SD_DEVICE void error(void);
|
||||
error();
|
||||
return 0;
|
||||
}
|
||||
};
|
||||
|
||||
// Following are the specializations for the following types.
|
||||
// int, sd::Unsigned, char, uchar, short, ushort, long long, ulong long, bool, float, and double
|
||||
// One could also specialize it for user-defined types.
|
||||
|
||||
template <>
|
||||
struct SharedIndexValue<float> {
|
||||
SD_DEVICE IndexValue<float> *getPointer() {
|
||||
extern __shared__ IndexValue<float> s_int2[];
|
||||
return s_int2;
|
||||
}
|
||||
};
|
||||
// Following are the specializations for the following types.
|
||||
// int, sd::Unsigned, char, uchar, short, ushort, long long, ulong long, bool, float, and double
|
||||
// One could also specialize it for user-defined types.
|
||||
|
||||
template <>
|
||||
struct SharedIndexValue<double> {
|
||||
SD_DEVICE IndexValue<double> *getPointer() {
|
||||
extern __shared__ IndexValue<double> s_int6[];
|
||||
return s_int6;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void IndexReduce<X, Z>::aggregatePartials(IndexValue<X> *sPartials, sd::LongType tid,
|
||||
sd::LongType numElements, void *vextraParams) {
|
||||
// start the shared memory loop on the next power of 2 less
|
||||
// than the block size. If block size is not a power of 2,
|
||||
// accumulate the intermediate sums in the remainder range.
|
||||
auto extraParams = static_cast<X *>(vextraParams);
|
||||
sd::LongType floorPow2 = static_cast<sd::LongType>(blockDim.x);
|
||||
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= floorPow2 - 1;
|
||||
}
|
||||
|
||||
if (tid >= floorPow2) {
|
||||
IndexValue<X> prev = sPartials[tid - floorPow2];
|
||||
IndexValue<X> curr = sPartials[tid];
|
||||
sPartials[tid - floorPow2] = OpType::update(prev, curr, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && tid + activeThreads < numElements) {
|
||||
IndexValue<X> curr = sPartials[tid];
|
||||
IndexValue<X> next = sPartials[tid + activeThreads];
|
||||
sPartials[tid] = OpType::update(curr, next, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_DEVICE void IndexReduce<X, Y>::transform(int opNum, void const *x, sd::LongType const *xShapeInfo,
|
||||
void *extraParams, void *result, sd::LongType const *zShapeInfo, sd::LongType *dimension,
|
||||
sd::LongType dimensionLength, int postProcessOrNot,
|
||||
sd::LongType *allocationBuffer, void *reductionBuffer,
|
||||
sd::LongType const *tadShapeInfo, sd::LongType const *tadOffset) {
|
||||
DISPATCH_BY_OPNUM_TT(transform,
|
||||
PARAMS(x, xShapeInfo, extraParams, result, zShapeInfo, dimension, dimensionLength,
|
||||
postProcessOrNot, allocationBuffer, reductionBuffer, tadShapeInfo, tadOffset),
|
||||
INDEX_REDUCE_OPS);
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void IndexReduce<X, Z>::transform(void const *vdx, sd::LongType const *xShapeInfo, void *vextraParams,
|
||||
void *vz, sd::LongType const *zShapeInfo, sd::LongType *dimension,
|
||||
sd::LongType dimensionLength, int postProcessOrNot,
|
||||
sd::LongType *allocationBuffer,
|
||||
void *vreductionBuffer, sd::LongType const *tadOnlyShapeInfo,
|
||||
sd::LongType const *tadOffsets) {
|
||||
auto dx = reinterpret_cast<X const *>(vdx);
|
||||
auto z = reinterpret_cast<Z *>(vz);
|
||||
auto extraParams = static_cast<X *>(vextraParams);
|
||||
auto reductionBuffer = static_cast<unsigned int *>(vreductionBuffer);
|
||||
auto order = shape::order(xShapeInfo);
|
||||
sd::LongType tid = static_cast<sd::LongType>(blockIdx.x * blockDim.x + threadIdx.x);
|
||||
__shared__ volatile bool resultScalar;
|
||||
|
||||
__shared__ IndexValue<X> sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingIndexValue(dx);
|
||||
|
||||
__shared__ volatile sd::LongType xLength;
|
||||
|
||||
__shared__ volatile sd::LongType zLen;
|
||||
|
||||
IndexValue<X> reduction = OpType::startingIndexValue(dx);
|
||||
sd::LongType threadIdxX = static_cast<sd::LongType>(threadIdx.x);
|
||||
sd::LongType blockDimX = static_cast<sd::LongType>(blockDim.x);
|
||||
sd::LongType blockIdxX = static_cast<sd::LongType>(blockIdx.x);
|
||||
sd::LongType gridDimX = static_cast<sd::LongType>(gridDim.x);
|
||||
|
||||
if (threadIdxX == 0) {
|
||||
if (zShapeInfo != nullptr)
|
||||
zLen = shape::length(zShapeInfo);
|
||||
else
|
||||
zLen = 1;
|
||||
|
||||
if (zLen == 1)
|
||||
resultScalar = true;
|
||||
else
|
||||
resultScalar = false;
|
||||
|
||||
xLength = shape::length(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (!resultScalar) {
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ sd::LongType tadEWS;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
tadEWS = shape::elementWiseStride(tadOnlyShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (dimensionLength > 1 || tadEWS < 1) {
|
||||
for (sd::LongType r = blockIdxX; r < numTads; r += gridDimX) {
|
||||
auto tadOffsetForBlock = tadOffsets[r];
|
||||
sPartials[threadIdxX] = OpType::startingIndexValue(dx);
|
||||
|
||||
for (sd::LongType i = threadIdxX; i < tadLength; i += blockDimX) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(i, shape::rank(tadOnlyShapeInfo), shape::shapeOf(tadOnlyShapeInfo), coords);
|
||||
COORDS2INDEX(shape::rank(tadOnlyShapeInfo), shape::stride(tadOnlyShapeInfo), coords, xOffset);
|
||||
xOffset += tadOffsetForBlock;
|
||||
IndexValue<X> comp{dx[xOffset], i};
|
||||
sPartials[threadIdxX] = OpType::update(sPartials[threadIdxX], comp, extraParams);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials,threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDimX, tadLength), extraParams);
|
||||
|
||||
__syncthreads();
|
||||
if (threadIdxX == 0) {
|
||||
z[r] = static_cast<Z>(sPartials[threadIdxX].index);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
} else {
|
||||
for (sd::LongType i = blockIdxX; i < numTads; i += gridDimX) {
|
||||
sd::LongType tadOffsetForBlock = tadOffsets[i];
|
||||
|
||||
sPartials[threadIdxX] = OpType::startingIndexValue(dx);
|
||||
|
||||
for (sd::LongType x = threadIdxX; x < tadLength; x += blockDimX) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(x, shape::rank(tadOnlyShapeInfo), shape::shapeOf(tadOnlyShapeInfo), coords);
|
||||
COORDS2INDEX(shape::rank(tadOnlyShapeInfo), shape::stride(tadOnlyShapeInfo), coords, xOffset);
|
||||
IndexValue<X> comp{dx[tadOffsetForBlock + xOffset], x};
|
||||
sPartials[threadIdxX] = OpType::update(sPartials[threadIdxX], comp, extraParams);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDim.x, tadLength), extraParams);
|
||||
|
||||
__syncthreads();
|
||||
if (threadIdxX == 0) {
|
||||
z[i] = static_cast<Z>(sPartials[threadIdxX].index);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
} else {
|
||||
auto n = shape::length(xShapeInfo);
|
||||
|
||||
for (sd::LongType i = tid; i < n; i += (gridDimX * blockDimX)) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(i, shape::rank(xShapeInfo), shape::shapeOf(xShapeInfo), coords);
|
||||
COORDS2INDEX(shape::rank(xShapeInfo), shape::stride(xShapeInfo), coords, xOffset);
|
||||
IndexValue<X> comp{dx[xOffset], i};
|
||||
reduction = OpType::update(reduction, comp, extraParams);
|
||||
}
|
||||
|
||||
sPartials[threadIdxX] = reduction;
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdxX, sd::math::sd_min<sd::LongType,sd::LongType>(blockDim.x, n), extraParams);
|
||||
if (gridDimX > 1) {
|
||||
__shared__ bool amLast;
|
||||
unsigned int *unsignedSharedMemory = (unsigned int *)reductionBuffer;
|
||||
tid = threadIdx.x;
|
||||
if (threadIdx.x == 0)
|
||||
reductionBuffer[blockIdx.x] = sPartials[threadIdx.x].index;
|
||||
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
unsigned int ticket = atomicInc(&unsignedSharedMemory[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
sPartials[threadIdx.x] = OpType::startingIndexValue(dx);
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
IndexValue<X> comp{static_cast<X>(0), reductionBuffer[i]};
|
||||
sPartials[threadIdx.x] = OpType::update(sPartials[threadIdx.x], comp, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdxX, gridDim.x, extraParams);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = static_cast<Z>(sPartials[threadIdx.x].index);
|
||||
unsignedSharedMemory[16384] = 0;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = static_cast<Z>(sPartials[threadIdx.x].index);
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace indexreduce
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,242 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
|
||||
#ifndef DEV_TESTS_REDUCE_SAME_LOOPS_H
|
||||
#define DEV_TESTS_REDUCE_SAME_LOOPS_H
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <helpers/shape.h>
|
||||
#include <ops/ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
#include <execution/LaunchContext.h>
|
||||
#include "execution/cuda/LaunchDims.h"
|
||||
using namespace simdOps;
|
||||
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
|
||||
template <typename X>
|
||||
class ReduceSameInplace {
|
||||
public:
|
||||
// Dispatch method for old-style calls
|
||||
static SD_INLINE void SD_DEVICE execScalarCudaLegacy(
|
||||
int opNum,
|
||||
void* vx,
|
||||
sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo);
|
||||
|
||||
// Template for real call
|
||||
template <typename OpClass>
|
||||
static SD_INLINE void SD_DEVICE execScalarCuda(
|
||||
void* vx,
|
||||
sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo);
|
||||
|
||||
template <typename OpClass>
|
||||
static SD_INLINE void SD_DEVICE aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams);
|
||||
};
|
||||
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_INLINE void SD_DEVICE ReduceSameInplace<X>::aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
auto sPartials = static_cast<X*>(vsPartials);
|
||||
auto extraParams = static_cast<X*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= (floorPow2 - 1);
|
||||
}
|
||||
if (tid >= floorPow2) {
|
||||
sPartials[tid - floorPow2] =
|
||||
OpClass::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && (tid + activeThreads) < numItems) {
|
||||
sPartials[tid] =
|
||||
OpClass::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
SD_INLINE void SD_DEVICE ReduceSameInplace<X>::execScalarCudaLegacy(
|
||||
int opNum,
|
||||
void* vx,
|
||||
sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
execScalarCuda,
|
||||
PARAMS(vx, xShapeInfo, vextraParams, vz, zShapeInfo, vreductionBuffer, tadOnlyShapeInfo),
|
||||
REDUCE_SAME_OPS);
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_INLINE void SD_DEVICE ReduceSameInplace<X>::execScalarCuda(
|
||||
void* vx,
|
||||
sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<X*>(vx);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
auto reductionBuff = reinterpret_cast<X*>(vreductionBuffer);
|
||||
|
||||
// We'll cache relevant shape info in shared memory so we don't call them repeatedly
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int rank;
|
||||
__shared__ const sd::LongType* shapePtr;
|
||||
__shared__ const sd::LongType* stridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
rank = shape::rank(xShapeInfo);
|
||||
shapePtr = shape::shapeOf(xShapeInfo);
|
||||
stridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto gridSize = gridDim.x * blockDim.x;
|
||||
|
||||
// We'll use some shared memory for partial sums
|
||||
__shared__ X* sPartials;
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
sPartials = reinterpret_cast<X*>(shmem);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each thread gets a starting value
|
||||
sPartials[threadIdx.x] = OpClass::startingValue(x);
|
||||
|
||||
// We'll stride over the entire array
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType offset;
|
||||
|
||||
INDEX2COORDS(i, rank, shapePtr, coords);
|
||||
COORDS2INDEX(rank, stridePtr, coords, offset);
|
||||
|
||||
sPartials[threadIdx.x] =
|
||||
OpClass::update(sPartials[threadIdx.x], OpClass::op(x[offset], extraParams), extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Next: reduce partial sums in the block
|
||||
aggregatePartials<OpClass>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
// If gridDim.x > 1, we do a multi-block reduce using the global buffer
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuff);
|
||||
__shared__ bool amLast;
|
||||
|
||||
// each block's sum is stored in the 'reductionBuff'
|
||||
if (threadIdx.x == 0) {
|
||||
reductionBuff[blockIdx.x] = sPartials[0];
|
||||
}
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
// The 16384 is a special "counter" location in the reductionBuff
|
||||
if (threadIdx.x == 0) {
|
||||
unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
// We do a final block-level reduce
|
||||
tc[16384] = 0; // reset
|
||||
sPartials[threadIdx.x] = OpClass::startingValue(x);
|
||||
|
||||
// accumulate partial sums from each block
|
||||
for (int i = threadIdx.x; i < static_cast<int>(gridDim.x); i += blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpClass::update(sPartials[threadIdx.x], reductionBuff[i], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpClass>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpClass::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
// single-block case
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuff);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpClass::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace reduce
|
||||
} // namespace functions
|
||||
|
||||
#endif // DEV_TESTS_REDUCE_SAME_LOOPS_H
|
||||
@@ -0,0 +1,138 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
#ifndef DEV_TESTS_SCALAR_INPLACE_H
|
||||
#define DEV_TESTS_SCALAR_INPLACE_H
|
||||
|
||||
#include <helpers/shape.h>
|
||||
#include <ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
namespace functions {
|
||||
namespace scalar {
|
||||
|
||||
template <typename X, typename Y, typename Z>
|
||||
class ScalarInplace {
|
||||
public:
|
||||
static SD_INLINE SD_DEVICE void transformCudaLegacy(
|
||||
int opNum,
|
||||
void* vscalar,
|
||||
void* vy,
|
||||
sd::LongType* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationBuffer);
|
||||
|
||||
template <typename OpClass>
|
||||
static SD_INLINE SD_DEVICE void transformCuda(
|
||||
void* vscalar,
|
||||
void* vy,
|
||||
sd::LongType* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationBuffer);
|
||||
};
|
||||
|
||||
template <typename X, typename Y, typename Z>
|
||||
SD_INLINE SD_DEVICE void ScalarInplace<X, Y, Z>::transformCudaLegacy(
|
||||
int opNum,
|
||||
void* vscalar,
|
||||
void* vy,
|
||||
sd::LongType* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationBuffer) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
transformCuda,
|
||||
PARAMS(vscalar, vy, yShapeInfo, vparams, vz, zShapeInfo, allocationBuffer),
|
||||
SCALAR_OPS);
|
||||
}
|
||||
|
||||
template <typename X, typename Y, typename Z>
|
||||
template <typename OpClass>
|
||||
SD_INLINE SD_DEVICE void ScalarInplace<X, Y, Z>::transformCuda(
|
||||
void* vscalar,
|
||||
void* vy,
|
||||
sd::LongType* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationBuffer) {
|
||||
|
||||
auto scalar = reinterpret_cast<X*>(vscalar)[0];
|
||||
auto y = reinterpret_cast<Y*>(vy);
|
||||
auto params = reinterpret_cast<Z*>(vparams);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
// We'll store relevant shape info for y & z in shared memory
|
||||
__shared__ sd::LongType yLen;
|
||||
__shared__ int yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
yLen = shape::length(yShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
// Each thread does part of the length
|
||||
for (sd::LongType i = tid; i < yLen; i += totalThreads) {
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType offsetY;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
// apply scalar op in place to produce z
|
||||
z[offsetZ] = OpClass::op(y[offsetY], scalar, params);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace scalar
|
||||
} // namespace functions
|
||||
|
||||
#endif // DEV_TESTS_SCALAR_INPLACE_H
|
||||
@@ -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 raver119@gmail.com
|
||||
//
|
||||
#ifndef DEV_TESTS_TRANSFORM_FLOAT_INPLACE_H
|
||||
#define DEV_TESTS_TRANSFORM_FLOAT_INPLACE_H
|
||||
|
||||
#include <helpers/shape.h>
|
||||
#include <ops.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
#define LOCAL_TRANSFORM_STRICT_OPS (23, Exp), (24, Log)
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
template <typename X>
|
||||
class TransformStrictInplace {
|
||||
public:
|
||||
static SD_INLINE SD_DEVICE void transformCudaLegacy(
|
||||
int opNum,
|
||||
void* dy,
|
||||
sd::LongType* shapeInfo,
|
||||
void* params,
|
||||
void* result,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationPointer,
|
||||
void* reductionPointer,
|
||||
sd::LongType* tadShapeInfo,
|
||||
sd::LongType* tadOffsets);
|
||||
|
||||
template <typename OpClass>
|
||||
static SD_INLINE SD_DEVICE void transformCuda(
|
||||
void* vdy,
|
||||
sd::LongType* shapeInfo,
|
||||
void* vparams,
|
||||
void* vresult,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
sd::LongType* tadShapeInfo,
|
||||
sd::LongType* tadOffsets);
|
||||
};
|
||||
|
||||
template <typename X>
|
||||
template <typename OpClass>
|
||||
SD_INLINE SD_DEVICE void TransformStrictInplace<X>::transformCuda(
|
||||
void* vdy,
|
||||
sd::LongType* shapeInfo,
|
||||
void* vparams,
|
||||
void* vresult,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
sd::LongType* tadShapeInfo,
|
||||
sd::LongType* tadOffsets) {
|
||||
|
||||
auto dy = static_cast<X*>(vdy);
|
||||
auto result= static_cast<X*>(vresult);
|
||||
auto params= static_cast<X*>(vparams);
|
||||
// reductionPointer is unused in this basic transform?
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Cache shape info in shared memory to avoid repeated calls
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int rank;
|
||||
__shared__ const sd::LongType* shapePtr;
|
||||
__shared__ const sd::LongType* stridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(shapeInfo);
|
||||
rank = shape::rank(shapeInfo);
|
||||
shapePtr = shape::shapeOf(shapeInfo);
|
||||
stridePtr = shape::stride(shapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType inCoords[SD_MAX_RANK];
|
||||
sd::LongType outCoords[SD_MAX_RANK];
|
||||
sd::LongType inOffset;
|
||||
sd::LongType outOffset;
|
||||
|
||||
INDEX2COORDS(i, rank, shapePtr, inCoords);
|
||||
COORDS2INDEX(rank, stridePtr, inCoords, inOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, outCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, outCoords, outOffset);
|
||||
|
||||
// Apply transform op in place
|
||||
result[outOffset] = OpClass::op(dy[inOffset], params);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
SD_INLINE SD_DEVICE void TransformStrictInplace<X>::transformCudaLegacy(
|
||||
int opNum,
|
||||
void* dy,
|
||||
sd::LongType* shapeInfo,
|
||||
void* params,
|
||||
void* result,
|
||||
sd::LongType* zShapeInfo,
|
||||
int* allocationPointer,
|
||||
void* reductionPointer,
|
||||
sd::LongType* tadShapeInfo,
|
||||
sd::LongType* tadOffsets) {
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
transformCuda,
|
||||
PARAMS(dy, shapeInfo, params, result, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets),
|
||||
LOCAL_TRANSFORM_STRICT_OPS);
|
||||
}
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
|
||||
#undef LOCAL_TRANSFORM_STRICT_OPS
|
||||
#endif // DEV_TESTS_TRANSFORM_FLOAT_INPLACE_H
|
||||
@@ -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
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
|
||||
#ifndef PAIRWISE_CU
|
||||
#define PAIRWISE_CU
|
||||
#include <loops/pairwise_instantiations.h>
|
||||
|
||||
#include "../pairwise_transform.h"
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z, typename OpType>
|
||||
SD_KERNEL static void pairwiseSimpleShaped(void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
auto x = static_cast<X const*>(vx);
|
||||
auto y = static_cast<Y const*>(vy);
|
||||
auto z = static_cast<Z*>(vz);
|
||||
auto extraParams= static_cast<Z*>(vextraParams);
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
__shared__ sd::LongType len;
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ sd::LongType zRank;
|
||||
|
||||
__shared__ const sd::LongType* xShape;
|
||||
__shared__ const sd::LongType* xStride;
|
||||
__shared__ const sd::LongType* yShape;
|
||||
__shared__ const sd::LongType* yStride;
|
||||
__shared__ const sd::LongType* zShape;
|
||||
__shared__ const sd::LongType* zStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
len = shape::length(xShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShape = shape::shapeOf(yShapeInfo);
|
||||
yStride = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType i = tid; i < len; i += (gridDim.x * blockDim.x)) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShape, yCoords);
|
||||
COORDS2INDEX(yRank, yStride, yCoords, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShape, zCoords);
|
||||
COORDS2INDEX(zRank, zStride, zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
|
||||
}
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace pairwise_transforms {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
void SD_HOST PairWiseTransform<X, Y, Z>::intermediateShaped(dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
pairwiseSimpleShaped<X, Y, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
vz, zShapeInfo,
|
||||
vextraParams);
|
||||
sd::DebugHelper::checkErrorCode(stream, "PairWiseTransform intermediateShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
void SD_HOST PairWiseTransform<X, Y, Z>::executeCudaShaped(dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
|
||||
PAIRWISE_TRANSFORM_OPS);
|
||||
}
|
||||
|
||||
} // namespace pairwise_transforms
|
||||
} // namespace functions
|
||||
|
||||
#endif // PAIRWISE_CU
|
||||
@@ -0,0 +1,158 @@
|
||||
/********************************************************************************
|
||||
*
|
||||
* 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), created on 08.11.2018
|
||||
//
|
||||
#ifndef PAIRWISE_BOOL_CU
|
||||
#define PAIRWISE_BOOL_CU
|
||||
|
||||
#include "../pairwise_bool.h"
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL static void pairwiseSimpleShaped(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// Cache shape info
|
||||
__shared__ int xRank;
|
||||
__shared__ sd::LongType* xShapePtr;
|
||||
__shared__ sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ int yRank;
|
||||
__shared__ sd::LongType* yShapePtr;
|
||||
__shared__ sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ sd::LongType* zShapePtr;
|
||||
__shared__ sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
sd::LongType offsetX;
|
||||
sd::LongType offsetY;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, offsetX);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
z[offsetZ] = OpType::op(x[offsetX], y[offsetY], extraParams);
|
||||
}
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace pairwise_transforms {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
void SD_HOST PairWiseBoolTransform<X,Z>::intermediateShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
pairwiseSimpleShaped<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
xShapeInfo,
|
||||
vy,
|
||||
yShapeInfo,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
vextraParams);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(
|
||||
stream,
|
||||
"PairWiseBoolTransform intermediateShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
void PairWiseBoolTransform<X,Y>::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
|
||||
PAIRWISE_BOOL_OPS);
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class PairWiseBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
|
||||
|
||||
} // namespace pairwise_transforms
|
||||
} // namespace functions
|
||||
|
||||
#endif // PAIRWISE_BOOL_CU
|
||||
@@ -0,0 +1,163 @@
|
||||
/********************************************************************************
|
||||
*
|
||||
*
|
||||
* 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), created on 08.11.2018
|
||||
//
|
||||
#ifndef PAIRWISE_INT_CU
|
||||
#define PAIRWISE_INT_CU
|
||||
|
||||
#include "../pairwise_int.h"
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename OpType>
|
||||
SD_KERNEL static void pairwiseSimpleShaped(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRank;
|
||||
__shared__ int yRank;
|
||||
__shared__ int zRank;
|
||||
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, yOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, zOffset);
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
|
||||
}
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace pairwise_transforms {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
void SD_HOST PairWiseIntTransform<X>::intermediateShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
pairwiseSimpleShaped<X, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
xShapeInfo,
|
||||
vy,
|
||||
yShapeInfo,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
vextraParams);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(
|
||||
stream,
|
||||
"PairWiseIntTransform intermediateShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
void PairWiseIntTransform<X>::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void* vextraParams) {
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
|
||||
PAIRWISE_INT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(
|
||||
stream,
|
||||
"PairWiseIntTransform executeCudaShaped(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( class PairWiseIntTransform, , SD_INTEGER_TYPES);
|
||||
|
||||
} // namespace pairwise_transforms
|
||||
} // namespace functions
|
||||
|
||||
#endif // PAIRWISE_INT_CU
|
||||
@@ -0,0 +1,491 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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/DebugHelper.h>
|
||||
#include <loops/random.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
#include <system/common.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
|
||||
|
||||
using namespace randomOps;
|
||||
|
||||
template <typename T, typename OpClass>
|
||||
static SD_INLINE SD_DEVICE void randomSingleGeneric(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer,
|
||||
void* extraArguments) {
|
||||
functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, z, zShapeBuffer, extraArguments);
|
||||
}
|
||||
|
||||
template <typename T, typename OpClass>
|
||||
static SD_INLINE SD_DEVICE void randomDoubleGeneric(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer,
|
||||
void* z, sd::LongType const* zShapeBuffer, void* extraArguments) {
|
||||
functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, x, xShapeBuffer, z, zShapeBuffer,
|
||||
extraArguments);
|
||||
}
|
||||
|
||||
template <typename T, typename OpClass>
|
||||
static SD_INLINE SD_DEVICE void randomTripleGeneric(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer,
|
||||
void const* y, sd::LongType const* yShapeBuffer, void* z,
|
||||
sd::LongType const* zShapeBuffer, void* extraArguments) {
|
||||
functions::random::RandomFunction<T>::template execTransformCuda<OpClass>(state, x, xShapeBuffer, y, yShapeBuffer, z,
|
||||
zShapeBuffer, extraArguments);
|
||||
}
|
||||
|
||||
// here we generate kernels for target operations
|
||||
DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, float,
|
||||
INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, double,
|
||||
INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, float16,
|
||||
INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomSingle_, randomSingleGeneric, bfloat16,
|
||||
INPUT(sd::Pointer state, void* z, sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, float,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
|
||||
sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, double,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
|
||||
sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, float16,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
|
||||
sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomDouble_, randomDoubleGeneric, bfloat16,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void* z,
|
||||
sd::LongType const* zShapeBuffer, void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, float,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
|
||||
sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
|
||||
void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, double,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
|
||||
sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
|
||||
void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, float16,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
|
||||
sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
|
||||
void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
DISPATCH_KERNEL_SIMPLE(randomTriple_, randomTripleGeneric, bfloat16,
|
||||
INPUT(sd::Pointer state, void const* x, sd::LongType const* xShapeBuffer, void const* y,
|
||||
sd::LongType const* yShapeBuffer, void* z, sd::LongType const* zShapeBuffer,
|
||||
void* extraArguments),
|
||||
PARAMS(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
|
||||
namespace functions {
|
||||
namespace random {
|
||||
template <typename T>
|
||||
template <typename OpClass>
|
||||
void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void const* vx, sd::LongType const* xShapeBuffer,
|
||||
void const* vy, sd::LongType const* yShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<T const*>(vx);
|
||||
auto y = reinterpret_cast<T const*>(vy);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto extraArguments = reinterpret_cast<T*>(vextraArguments);
|
||||
|
||||
if (OpClass::requiresSpecial) {
|
||||
OpClass::specialOpCuda(state, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments);
|
||||
return;
|
||||
} else {
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ char xOrder;
|
||||
__shared__ char yOrder;
|
||||
__shared__ char zOrder;
|
||||
|
||||
// Cache shape information for x buffer
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
// Cache shape information for y buffer
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
// Cache shape information for z buffer
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
__shared__ sd::graph::RandomGenerator* buffer;
|
||||
__shared__ unsigned char* cB;
|
||||
__shared__ unsigned char* dB;
|
||||
sd::graph::RandomGenerator* devBuffer;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(zShapeBuffer);
|
||||
xOrder = shape::order(xShapeBuffer);
|
||||
yOrder = shape::order(yShapeBuffer);
|
||||
zOrder = shape::order(zShapeBuffer);
|
||||
|
||||
// Cache all shape information in thread 0
|
||||
xRank = shape::rank(xShapeBuffer);
|
||||
xShapePtr = shape::shapeOf(xShapeBuffer);
|
||||
xStridePtr = shape::stride(xShapeBuffer);
|
||||
|
||||
yRank = shape::rank(yShapeBuffer);
|
||||
yShapePtr = shape::shapeOf(yShapeBuffer);
|
||||
yStridePtr = shape::stride(yShapeBuffer);
|
||||
|
||||
zRank = shape::rank(zShapeBuffer);
|
||||
zShapePtr = shape::shapeOf(zShapeBuffer);
|
||||
zStridePtr = shape::stride(zShapeBuffer);
|
||||
|
||||
extern __shared__ unsigned char shmem[];
|
||||
buffer = (sd::graph::RandomGenerator*)shmem;
|
||||
cB = shmem;
|
||||
devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
|
||||
dB = reinterpret_cast<unsigned char*>(state);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// using this loop instead of memcpy
|
||||
for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x)
|
||||
cB[e] = dB[e];
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += blockDim.x * gridDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
|
||||
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpClass::op(x[xOffset], y[yOffset], i, length, buffer, extraArguments);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
template <typename OpClass>
|
||||
void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void const* vx, sd::LongType const* xShapeBuffer,
|
||||
void* vz, sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<T const*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto extraArguments = reinterpret_cast<T*>(vextraArguments);
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ char xOrder;
|
||||
__shared__ char zOrder;
|
||||
|
||||
__shared__ sd::graph::RandomGenerator* buffer;
|
||||
__shared__ unsigned char* cB;
|
||||
__shared__ unsigned char* dB;
|
||||
__shared__ sd::graph::RandomGenerator* devBuffer;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
buffer = (sd::graph::RandomGenerator*)shmem;
|
||||
cB = shmem;
|
||||
devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
|
||||
dB = reinterpret_cast<unsigned char*>(state);
|
||||
|
||||
length = shape::length(zShapeBuffer);
|
||||
xOrder = shape::order(xShapeBuffer);
|
||||
zOrder = shape::order(zShapeBuffer);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// using this loop instead of memcpy
|
||||
for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x) cB[e] = dB[e];
|
||||
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType i = blockIdx.x * blockDim.x + threadIdx.x; i < length; i += blockDim.x * gridDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, shape::rank(xShapeBuffer), shape::shapeOf(xShapeBuffer), xCoords);
|
||||
COORDS2INDEX(shape::rank(xShapeBuffer), shape::stride(xShapeBuffer), xCoords, xOffset);
|
||||
INDEX2COORDS(i, shape::rank(zShapeBuffer), shape::shapeOf(zShapeBuffer), zCoords);
|
||||
COORDS2INDEX(shape::rank(zShapeBuffer), shape::stride(zShapeBuffer), zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpClass::op(x[xOffset], i, length, buffer, extraArguments);
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
template <typename OpClass>
|
||||
void SD_DEVICE RandomFunction<T>::execTransformCuda(sd::Pointer state, void* vz, sd::LongType const* zShapeBuffer,
|
||||
void* vextraArguments) {
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto extraArguments = reinterpret_cast<T*>(vextraArguments);
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ sd::graph::RandomGenerator* buffer;
|
||||
__shared__ unsigned char* cB;
|
||||
__shared__ unsigned char* dB;
|
||||
__shared__ sd::graph::RandomGenerator* devBuffer;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char shmem[];
|
||||
buffer = (sd::graph::RandomGenerator*)shmem;
|
||||
cB = shmem;
|
||||
devBuffer = reinterpret_cast<sd::graph::RandomGenerator*>(state);
|
||||
dB = reinterpret_cast<unsigned char*>(state);
|
||||
length = shape::length(zShapeBuffer);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// using this loop instead of memcpy
|
||||
for (int e = threadIdx.x; e < sizeof(sd::graph::RandomGenerator); e += blockDim.x) cB[e] = dB[e];
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += blockDim.x * gridDim.x) {
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, shape::rank(zShapeBuffer), shape::shapeOf(zShapeBuffer), zCoords);
|
||||
COORDS2INDEX(shape::rank(zShapeBuffer), shape::stride(zShapeBuffer), zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpClass::op(i, length, buffer, extraArguments);
|
||||
}
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void* vz, sd::LongType const* zShapeBuffer,
|
||||
void* vextraArguments) {
|
||||
auto z = reinterpret_cast<float*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomSingle, float, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float16>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto z = reinterpret_cast<float16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomSingle, float16, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<bfloat16>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto z = reinterpret_cast<bfloat16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomSingle, bfloat16, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<double>::executeCudaSingle(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto z = reinterpret_cast<double*>(vz);
|
||||
auto extraArguments = reinterpret_cast<double*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomSingle, double, PARAMS(stateHost, z, zShapeBuffer, extraArguments), OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<float const*>(vx);
|
||||
auto z = reinterpret_cast<float*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomDouble, float, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float16>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<float16 const*>(vx);
|
||||
auto z = reinterpret_cast<float16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomDouble, float16, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<bfloat16>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<bfloat16 const*>(vx);
|
||||
auto z = reinterpret_cast<bfloat16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomDouble, bfloat16, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<double>::executeCudaDouble(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<double const*>(vx);
|
||||
auto z = reinterpret_cast<double*>(vz);
|
||||
auto extraArguments = reinterpret_cast<double*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomDouble, double, PARAMS(stateHost, x, xShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void const* vy,
|
||||
sd::LongType const* yShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<float const*>(vx);
|
||||
auto y = reinterpret_cast<float const*>(vy);
|
||||
auto z = reinterpret_cast<float*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomTriple, float,
|
||||
PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<float16>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void const* vy,
|
||||
sd::LongType const* yShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<float16 const*>(vx);
|
||||
auto y = reinterpret_cast<float16 const*>(vy);
|
||||
auto z = reinterpret_cast<float16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<float16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomTriple, float16,
|
||||
PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<bfloat16>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void const* vy,
|
||||
sd::LongType const* yShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<bfloat16 const*>(vx);
|
||||
auto y = reinterpret_cast<bfloat16 const*>(vy);
|
||||
auto z = reinterpret_cast<bfloat16*>(vz);
|
||||
auto extraArguments = reinterpret_cast<bfloat16*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomTriple, bfloat16,
|
||||
PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
template <>
|
||||
SD_HOST void RandomFunction<double>::executeCudaTriple(dim3& launchDims, cudaStream_t* stream, int opNum,
|
||||
sd::Pointer stateHost, void const* vx,
|
||||
sd::LongType const* xShapeBuffer, void const* vy,
|
||||
sd::LongType const* yShapeBuffer, void* vz,
|
||||
sd::LongType const* zShapeBuffer, void* vextraArguments) {
|
||||
auto x = reinterpret_cast<double const*>(vx);
|
||||
auto y = reinterpret_cast<double const*>(vy);
|
||||
auto z = reinterpret_cast<double*>(vz);
|
||||
auto extraArguments = reinterpret_cast<double*>(vextraArguments);
|
||||
|
||||
// this macro builds bunch of IF/ELSE selectors for kernel launch
|
||||
DISPATCH_SIMPLE(randomTriple, double,
|
||||
PARAMS(stateHost, x, xShapeBuffer, y, yShapeBuffer, z, zShapeBuffer, extraArguments),
|
||||
OPS_A(RANDOM_OPS))
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "RandomFunction executeCudaSingle(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( class RandomFunction, , SD_COMMON_TYPES);
|
||||
|
||||
|
||||
} // namespace random
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,558 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <exceptions/cuda_exception.h>
|
||||
#include <execution/LaunchContext.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/reduce_bool.h>
|
||||
#include <loops/scalar.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <system/common.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleReduce(
|
||||
const void* x,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceBoolFunction<X, Z>::template transformCuda<OpType>(
|
||||
x,
|
||||
outerXTadShapeInfo,
|
||||
innerXTadShapeInfo,
|
||||
vreductionBuffer,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleScalar(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceBoolFunction<X, Z>::template execScalarCuda<OpType>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceBoolFunction<X, Z>::aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
auto sPartials = reinterpret_cast<Z*>(vsPartials);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= (floorPow2 - 1);
|
||||
}
|
||||
if (tid >= floorPow2) {
|
||||
sPartials[tid - floorPow2] =
|
||||
OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && (tid + activeThreads) < numItems) {
|
||||
sPartials[tid] = OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceBoolFunction<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ int tadLen;
|
||||
__shared__ int numTads;
|
||||
__shared__ bool sameOffsets;
|
||||
|
||||
// Cache shape info
|
||||
__shared__ sd::LongType outerRank;
|
||||
__shared__ sd::LongType innerRank;
|
||||
__shared__ sd::LongType zRank;
|
||||
|
||||
__shared__ const sd::LongType* outerShapePtr;
|
||||
__shared__ const sd::LongType* outerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* innerShapePtr;
|
||||
__shared__ const sd::LongType* innerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outerRank = shape::rank(outerXTadShapeInfo);
|
||||
outerShapePtr = shape::shapeOf(outerXTadShapeInfo);
|
||||
outerStridePtr = shape::stride(outerXTadShapeInfo);
|
||||
|
||||
innerRank = shape::rank(innerXTadShapeInfo);
|
||||
innerShapePtr = shape::shapeOf(innerXTadShapeInfo);
|
||||
innerStridePtr = shape::stride(innerXTadShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
|
||||
tadLen = shape::length(innerXTadShapeInfo);
|
||||
numTads = shape::length(outerXTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
INDEX2COORDS(r, outerRank, outerShapePtr, coords);
|
||||
|
||||
sd::LongType outerOffset;
|
||||
COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
|
||||
|
||||
sd::LongType zOffset;
|
||||
if (sameOffsets) {
|
||||
zOffset = outerOffset;
|
||||
} else {
|
||||
INDEX2COORDS(r, zRank, zShapePtr, coords);
|
||||
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
||||
}
|
||||
|
||||
const X* xTad = x + outerOffset;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(xTad);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType innerOffset;
|
||||
|
||||
INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
|
||||
COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(xTad[innerOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, tadLen),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceBoolFunction<X, Z>::execScalarCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// Cache shape info
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
auto reductionBuffer = reinterpret_cast<Z *>(vreductionBuffer);
|
||||
|
||||
int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
sd::LongType gridSize = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(x[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
reductionBuffer[blockIdx.x] = sPartials[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;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
reinterpret_cast<Z*>(vreductionBuffer)[i],
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceBoolFunction<X, Z>::intermediate(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
sd::LongType* dXShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
sd::LongType* dZShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dims) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal =
|
||||
static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
|
||||
auto res = cudaMemcpyAsync(
|
||||
sd::LaunchContext::defaultContext()->getScalarPointer(),
|
||||
&startingVal,
|
||||
sizeof(Z),
|
||||
cudaMemcpyHostToDevice,
|
||||
*stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceBoolFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
|
||||
}
|
||||
|
||||
auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
|
||||
// scalar assign
|
||||
scalar::ScalarTransform<Z,Z,Z>::executeCudaShaped(
|
||||
launchDims,
|
||||
stream,
|
||||
14,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hZShapeInfo,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hZShapeInfo,
|
||||
ptr,
|
||||
nullptr);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reduceBoolDim empty(...) failed");
|
||||
} else {
|
||||
const sd::LongType zRank = shape::rank(hZShapeInfo);
|
||||
const sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
|
||||
|
||||
auto outerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(hXShapeInfo, dims, zRank);
|
||||
auto innerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(hXShapeInfo, dims + zRank, tadRank);
|
||||
|
||||
simpleReduce<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
outerPack->special(),
|
||||
innerPack->special(),
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
dZShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reduceBoolDim(...) failed");
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceBoolFunction<X, Z>::intermediateScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal =
|
||||
static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
auto res = cudaMemcpyAsync(
|
||||
z,
|
||||
&startingVal,
|
||||
sizeof(Z),
|
||||
cudaMemcpyHostToDevice,
|
||||
*stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceBoolFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "reduceBoolScalar empty(...) failed");
|
||||
} else {
|
||||
simpleScalar<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reduceBoolScalar(...) failed");
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST SD_INLINE void ReduceBoolFunction<X, Y>::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateScalar,
|
||||
PARAMS(
|
||||
launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z,
|
||||
zShapeInfo, hZShapeInfo, dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
|
||||
OPS_A(REDUCE_BOOL_OPS));
|
||||
sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST SD_INLINE void ReduceBoolFunction<X, Y>::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
sd::LongType* dXShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
sd::LongType* dZShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dims) {
|
||||
|
||||
if (shape::length(hZShapeInfo) == 1) {
|
||||
execReduceScalar(
|
||||
launchDims, stream, opNum,
|
||||
x, dXShapeInfo, hXShapeInfo,
|
||||
extraParams, z,
|
||||
dZShapeInfo, hZShapeInfo,
|
||||
nullptr, 0,
|
||||
vreductionBuffer,
|
||||
nullptr);
|
||||
} else {
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediate,
|
||||
PARAMS(
|
||||
launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams,
|
||||
vreductionBuffer, z, dZShapeInfo, hZShapeInfo, dims),
|
||||
OPS_A(REDUCE_BOOL_OPS));
|
||||
}
|
||||
DEBUG_KERNEL(stream, opNum);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
|
||||
int sPartialsLength = sMemSize / sizeof(X);
|
||||
X* sPartialsDeref = reinterpret_cast<X*>(*sPartials);
|
||||
for (int i = 0; i < sPartialsLength; i++) {
|
||||
sPartialsDeref[i] = extraParams[0];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
(SD_COMMON_TYPES),
|
||||
(SD_BOOL_TYPES),
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceBoolFunction,
|
||||
::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
sd::LongType* dXShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
sd::LongType* dZShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dims);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
(SD_COMMON_TYPES),
|
||||
(SD_BOOL_TYPES),
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceBoolFunction,
|
||||
::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo);
|
||||
);
|
||||
|
||||
|
||||
} // namespace reduce
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,502 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <execution/LaunchContext.h>
|
||||
#include <exceptions/cuda_exception.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <loops/reduce_float.h>
|
||||
#include <loops/scalar.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <types/types.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL void simpleReduce(
|
||||
const void* x,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceFloatFunction<X,Z>::template transformCuda<OpType>(
|
||||
x,
|
||||
outerXTadShapeInfo,
|
||||
innerXTadShapeInfo,
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
zShapeInfo);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL void simpleScalar(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceFloatFunction<X, Z>::template execScalarCuda<OpType>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void ReduceFloatFunction<X,Z>::aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
using Y = typename OpType::InterType;
|
||||
auto sPartials = reinterpret_cast<Y*>(vsPartials);
|
||||
auto extraParams = reinterpret_cast<Z*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= (floorPow2 - 1);
|
||||
}
|
||||
if (tid >= floorPow2) {
|
||||
sPartials[tid - floorPow2] =
|
||||
OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && (tid + activeThreads) < numItems) {
|
||||
sPartials[tid] =
|
||||
OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void ReduceFloatFunction<X,Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vreductionBuffer,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams= reinterpret_cast<Z*>(vextraParams);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ sd::LongType tadLen;
|
||||
__shared__ sd::LongType numTads;
|
||||
__shared__ bool sameOffsets;
|
||||
|
||||
// Cache ranks/shape/stride for outer and z
|
||||
__shared__ sd::LongType outerRank;
|
||||
__shared__ const sd::LongType* outerStridePtr;
|
||||
__shared__ const sd::LongType* outerShapePtr;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
|
||||
// Cache ranks/shape/stride for inner as well if needed
|
||||
__shared__ sd::LongType innerRank;
|
||||
__shared__ const sd::LongType* innerStridePtr;
|
||||
__shared__ const sd::LongType* innerShapePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outerRank = shape::rank(outerXTadShapeInfo);
|
||||
outerShapePtr = shape::shapeOf(outerXTadShapeInfo);
|
||||
outerStridePtr = shape::stride(outerXTadShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
innerRank = shape::rank(innerXTadShapeInfo);
|
||||
innerShapePtr = shape::shapeOf(innerXTadShapeInfo);
|
||||
innerStridePtr = shape::stride(innerXTadShapeInfo);
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
|
||||
tadLen = shape::length(innerXTadShapeInfo);
|
||||
numTads = shape::length(outerXTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
INDEX2COORDS(r, outerRank, outerShapePtr, coords);
|
||||
|
||||
sd::LongType outerOffset;
|
||||
COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
|
||||
|
||||
sd::LongType zOffset;
|
||||
if (sameOffsets) {
|
||||
zOffset = outerOffset;
|
||||
} else {
|
||||
INDEX2COORDS(r, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
}
|
||||
|
||||
const X* xTad = x + outerOffset;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(xTad);
|
||||
|
||||
// For the inner dimension
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType innerOffset;
|
||||
|
||||
INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
|
||||
COORDS2INDEX(innerRank, innerStridePtr, iCoords, innerOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(xTad[innerOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, tadLen),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[zOffset] =
|
||||
OpType::postProcess(sPartials[0], tadLen, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void ReduceFloatFunction<X, Z>::execScalarCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<Z*>(vextraParams);
|
||||
auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);
|
||||
|
||||
using Y = typename OpType::InterType;
|
||||
|
||||
__shared__ Y sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// Cache rank/shape/stride
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
sd::LongType gridSize = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
if(xOffset < length) {
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(x[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
reductionBuffer[blockIdx.x] = sPartials[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;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::update(sPartials[threadIdx.x],
|
||||
reductionBuffer[i],
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template<typename OpType>
|
||||
SD_HOST void ReduceFloatFunction<X,Z>::intermediate(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* dXShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
const sd::LongType* dims) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
|
||||
? sd::DataTypeUtils::nanOrZero<Z>()
|
||||
: static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
|
||||
auto res = cudaMemcpyAsync(
|
||||
sd::LaunchContext::defaultContext()->getScalarPointer(),
|
||||
&startingVal,
|
||||
sizeof(Z),
|
||||
cudaMemcpyHostToDevice,
|
||||
*stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceFloatFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
|
||||
}
|
||||
|
||||
auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
|
||||
// scalar assign
|
||||
functions::scalar::ScalarTransform<Z, Z, Z>::executeCudaShaped(
|
||||
launchDims,
|
||||
stream,
|
||||
14,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hZShapeInfo,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hZShapeInfo,
|
||||
ptr,
|
||||
nullptr);
|
||||
} else {
|
||||
const int zRank = shape::rank(hZShapeInfo);
|
||||
const int tadRank = shape::rank(hXShapeInfo) - zRank;
|
||||
|
||||
auto outerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
|
||||
const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
|
||||
auto innerPack = sd::ConstantShapeHelper::getInstance().createSubArrShapeInfo(
|
||||
const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);
|
||||
|
||||
simpleReduce<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
reinterpret_cast<const sd::LongType*>(outerPack->special()),
|
||||
reinterpret_cast<const sd::LongType*>(innerPack->special()),
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
dZShapeInfo);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template<typename OpType>
|
||||
SD_HOST void ReduceFloatFunction<X,Z>::intermediateScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
const auto startingVal = std::is_same<OpType, simdOps::Mean<X,Z>>::value
|
||||
? sd::DataTypeUtils::nanOrZero<Z>()
|
||||
: static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
|
||||
auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z), cudaMemcpyHostToDevice, *stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceFloatFunction<X,Z>::intermediateScalar: failed to copy resulting scalar", res);
|
||||
}
|
||||
} else {
|
||||
simpleScalar<X, Z, OpType><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction intermediateScalar(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void ReduceFloatFunction<X,Y>::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateScalar,
|
||||
PARAMS(
|
||||
launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z, dZShapeInfo, hZShapeInfo, dimension,
|
||||
dimensionLength, reductionBuffer, tadOnlyShapeInfo),
|
||||
OPS_A(REDUCE_FLOAT_OPS));
|
||||
sd::DebugHelper::checkErrorCode(stream, "execReduceScalarFloat(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void ReduceFloatFunction<X,Y>::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* dXShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
const sd::LongType* dims) {
|
||||
|
||||
if (shape::length(hZShapeInfo) == 1) {
|
||||
ReduceFloatFunction<X,Y>::execReduceScalar(
|
||||
launchDims, stream, opNum, x, dXShapeInfo, hXShapeInfo,
|
||||
extraParams, z, dZShapeInfo, hZShapeInfo,
|
||||
nullptr, 0, vreductionBuffer, nullptr);
|
||||
} else {
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediate,
|
||||
PARAMS(
|
||||
launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
|
||||
hZShapeInfo, dims),
|
||||
OPS_A(REDUCE_FLOAT_OPS));
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceFloatFunction execReduce(...) failed");
|
||||
}
|
||||
|
||||
} // namespace reduce
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,552 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <exceptions/cuda_exception.h>
|
||||
#include <execution/LaunchContext.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/reduce_long.h>
|
||||
#include <loops/scalar.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleReduce(
|
||||
const void* x,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceLongFunction<X, Z>::template transformCuda<OpType>(
|
||||
x,
|
||||
outerXTadShapeInfo,
|
||||
innerXTadShapeInfo,
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
zShapeInfo);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_DEVICE SD_INLINE void reduceScalarGeneric(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceLongFunction<X, Z>::template execScalarCuda<OpType>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleScalar(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
reduceScalarGeneric<X, Z, OpType>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
auto sPartials = reinterpret_cast<Z*>(vsPartials);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= floorPow2 - 1;
|
||||
}
|
||||
if (tid >= floorPow2) {
|
||||
sPartials[tid - floorPow2] =
|
||||
OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && (tid + activeThreads) < numItems) {
|
||||
sPartials[tid] = OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vreductionBuffer,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ int tadLen;
|
||||
__shared__ int numTads;
|
||||
__shared__ bool sameOffsets;
|
||||
|
||||
// Cache shape info for outer/inner and z if needed
|
||||
__shared__ sd::LongType outerRank;
|
||||
__shared__ sd::LongType innerRank;
|
||||
__shared__ sd::LongType zRank;
|
||||
|
||||
__shared__ const sd::LongType* outerShapePtr;
|
||||
__shared__ const sd::LongType* outerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* innerShapePtr;
|
||||
__shared__ const sd::LongType* innerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outerRank = shape::rank(outerXTadShapeInfo);
|
||||
outerShapePtr = shape::shapeOf(outerXTadShapeInfo);
|
||||
outerStridePtr = shape::stride(outerXTadShapeInfo);
|
||||
|
||||
innerRank = shape::rank(innerXTadShapeInfo);
|
||||
innerShapePtr = shape::shapeOf(innerXTadShapeInfo);
|
||||
innerStridePtr = shape::stride(innerXTadShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
|
||||
tadLen = shape::length(innerXTadShapeInfo);
|
||||
numTads = shape::length(outerXTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
INDEX2COORDS(r, outerRank, outerShapePtr, coords);
|
||||
|
||||
sd::LongType outerOffset;
|
||||
COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
|
||||
|
||||
sd::LongType zOffset;
|
||||
if (sameOffsets) {
|
||||
zOffset = outerOffset;
|
||||
} else {
|
||||
INDEX2COORDS(r, zRank, zShapePtr, coords);
|
||||
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
||||
}
|
||||
|
||||
const X* xTad = x + outerOffset;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(xTad);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
|
||||
COORDS2INDEX(innerRank, innerStridePtr, iCoords, xOffset);
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(xTad[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, tadLen),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceLongFunction<X, Z>::execScalarCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
auto reductionBuffer = reinterpret_cast<Z*>(vreductionBuffer);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// Cache shape info
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
sd::LongType gridSize = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(x[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
reductionBuffer[blockIdx.x] = sPartials[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;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::update(sPartials[threadIdx.x], reductionBuffer[i], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceLongFunction<X, Z>::intermediate(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
sd::LongType* dXShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
sd::LongType* dZShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dims) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
// If input is empty, we skip unless z is also empty
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal = static_cast<Z>(
|
||||
OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
auto res = cudaMemcpyAsync(
|
||||
sd::LaunchContext::defaultContext()->getScalarPointer(),
|
||||
&startingVal,
|
||||
sizeof(Z),
|
||||
cudaMemcpyHostToDevice,
|
||||
*stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceLongFunction<X,Z>::intermediate: failed to copy temporary scalar", res);
|
||||
}
|
||||
|
||||
auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
|
||||
// scalar assign
|
||||
scalar::ScalarTransform<Z,Z,Z>::executeCudaShaped(
|
||||
launchDims, stream, 14, z, dZShapeInfo, hXShapeInfo,
|
||||
z, dZShapeInfo, hZShapeInfo,
|
||||
ptr, nullptr);
|
||||
} else {
|
||||
sd::LongType zRank = shape::rank(hZShapeInfo);
|
||||
sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
|
||||
|
||||
auto outerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(hXShapeInfo, dims, zRank);
|
||||
auto innerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(hXShapeInfo, dims + zRank, tadRank);
|
||||
|
||||
simpleReduce<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
outerPack->special(),
|
||||
innerPack->special(),
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
dZShapeInfo);
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction intermediate(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceLongFunction<X, Z>::intermediateScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal =
|
||||
static_cast<Z>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
auto res = cudaMemcpyAsync(z, &startingVal, sizeof(Z),
|
||||
cudaMemcpyHostToDevice, *stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceLongFunction<X,Z>::intermediateScalar: failed to copy resulting scalar",
|
||||
res);
|
||||
}
|
||||
} else {
|
||||
simpleScalar<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction intermediateScalar(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST SD_INLINE void ReduceLongFunction<X, Y>::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
sd::LongType* xShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
sd::LongType* zShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateScalar,
|
||||
PARAMS(launchDims, stream, x, xShapeInfo, hXShapeInfo,
|
||||
extraParams, z, zShapeInfo, hZShapeInfo,
|
||||
dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
|
||||
OPS_A(REDUCE_LONG_OPS));
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceLongFunction execReduceScalar(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST SD_INLINE void ReduceLongFunction<X, Y>::execReduce(dim3 launchDims, cudaStream_t *stream, int opNum, const void *vx,
|
||||
sd::LongType *dXShapeInfo, sd::LongType *hXShapeInfo, void *extraParams,
|
||||
void *vreductionBuffer, void *vz, sd::LongType *dZShapeInfo,
|
||||
sd::LongType *hZShapeInfo, sd::LongType *dims) {
|
||||
|
||||
if (shape::length(hZShapeInfo) == 1) {
|
||||
execReduceScalar(
|
||||
launchDims, stream, opNum,
|
||||
vx, dXShapeInfo, hXShapeInfo,
|
||||
extraParams, vz,
|
||||
dZShapeInfo, hZShapeInfo,
|
||||
nullptr, 0, vreductionBuffer, nullptr);
|
||||
} else {
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediate,
|
||||
PARAMS(launchDims, stream, vx, dXShapeInfo, hXShapeInfo,
|
||||
extraParams, vreductionBuffer, vz, dZShapeInfo, hZShapeInfo, dims),
|
||||
OPS_A(REDUCE_LONG_OPS));
|
||||
}
|
||||
DEBUG_KERNEL(stream, opNum);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
|
||||
int sPartialsLength = sMemSize / sizeof(X);
|
||||
X* sPartialsDeref = reinterpret_cast<X*>(*sPartials);
|
||||
for (int i = 0; i < sPartialsLength; i++) {
|
||||
sPartialsDeref[i] = extraParams[0];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
(SD_COMMON_TYPES),
|
||||
(SD_LONG_TYPES),
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceLongFunction,
|
||||
::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void* vx,
|
||||
sd::LongType* dXShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* vz,
|
||||
sd::LongType* dZShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dims);
|
||||
);
|
||||
|
||||
ITERATE_COMBINATIONS(
|
||||
(SD_COMMON_TYPES),
|
||||
(SD_LONG_TYPES),
|
||||
INSTANT_PROCESS_COMBINATION,
|
||||
functions::reduce::ReduceLongFunction,
|
||||
::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
sd::LongType* xShapeInfo,
|
||||
sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
sd::LongType* zShapeInfo,
|
||||
sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
sd::LongType* tadOnlyShapeInfo);
|
||||
);
|
||||
|
||||
} // namespace reduce
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,553 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <exceptions/cuda_exception.h>
|
||||
#include <execution/LaunchContext.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/reduce_same.h>
|
||||
#include <loops/scalar.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleReduce(
|
||||
const void* x,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceSameFunction<X>::template transformCuda<OpType>(
|
||||
x,
|
||||
outerXTadShapeInfo,
|
||||
innerXTadShapeInfo,
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
zShapeInfo);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename OpType>
|
||||
SD_KERNEL SD_INLINE void simpleScalar(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
functions::reduce::ReduceSameFunction<X>::template execScalarCuda<OpType>(
|
||||
x, xShapeInfo, extraParams, z, zShapeInfo, reductionBuffer, tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace reduce {
|
||||
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceSameFunction<X>::aggregatePartials(
|
||||
void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
auto sPartials = reinterpret_cast<X*>(vsPartials);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= floorPow2 - 1;
|
||||
}
|
||||
if (tid >= floorPow2) {
|
||||
sPartials[tid - floorPow2] =
|
||||
OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = (floorPow2 >> 1); activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && (tid + activeThreads) < numItems) {
|
||||
sPartials[tid] =
|
||||
OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceSameFunction<X>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* outerXTadShapeInfo,
|
||||
const sd::LongType* innerXTadShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vreductionBuffer,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ X sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ int tadLen;
|
||||
__shared__ int numTads;
|
||||
__shared__ bool sameOffsets;
|
||||
|
||||
// Cache shape info for outer/inner and z if needed
|
||||
__shared__ sd::LongType outerRank;
|
||||
__shared__ sd::LongType innerRank;
|
||||
__shared__ sd::LongType zRank;
|
||||
|
||||
__shared__ const sd::LongType* outerShapePtr;
|
||||
__shared__ const sd::LongType* outerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* innerShapePtr;
|
||||
__shared__ const sd::LongType* innerStridePtr;
|
||||
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
outerRank = shape::rank(outerXTadShapeInfo);
|
||||
outerShapePtr = shape::shapeOf(outerXTadShapeInfo);
|
||||
outerStridePtr = shape::stride(outerXTadShapeInfo);
|
||||
|
||||
innerRank = shape::rank(innerXTadShapeInfo);
|
||||
innerShapePtr = shape::shapeOf(innerXTadShapeInfo);
|
||||
innerStridePtr = shape::stride(innerXTadShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(zShapeInfo, outerXTadShapeInfo);
|
||||
tadLen = shape::length(innerXTadShapeInfo);
|
||||
numTads = shape::length(outerXTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
INDEX2COORDS(r, outerRank, outerShapePtr, coords);
|
||||
|
||||
sd::LongType outerOffset;
|
||||
COORDS2INDEX(outerRank, outerStridePtr, coords, outerOffset);
|
||||
|
||||
sd::LongType zOffset;
|
||||
if (sameOffsets) {
|
||||
zOffset = outerOffset;
|
||||
} else {
|
||||
INDEX2COORDS(r, zRank, zShapePtr, coords);
|
||||
COORDS2INDEX(zRank, zStridePtr, coords, zOffset);
|
||||
}
|
||||
|
||||
const X* xTad = x + outerOffset;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(xTad);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
|
||||
INDEX2COORDS(i, innerRank, innerShapePtr, iCoords);
|
||||
COORDS2INDEX(innerRank, innerStridePtr, iCoords, xOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(xTad[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, tadLen),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[zOffset] = OpType::postProcess(sPartials[0], tadLen, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_DEVICE SD_INLINE void ReduceSameFunction<X>::execScalarCudaLegacy(
|
||||
int opNum,
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
execScalarCuda,
|
||||
PARAMS(vx, xShapeInfo, vextraParams, vz, zShapeInfo, vreductionBuffer, tadOnlyShapeInfo),
|
||||
REDUCE_SAME_OPS);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_DEVICE SD_INLINE void ReduceSameFunction<X>::execScalarCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void* vreductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<X*>(vz);
|
||||
auto extraParams = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
__shared__ X sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ sd::LongType length;
|
||||
auto reductionBuffer = reinterpret_cast<X *>(vreductionBuffer);
|
||||
|
||||
// Cache shape info
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
sd::LongType gridSize = gridDim.x * blockDim.x;
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::op(x[xOffset], extraParams),
|
||||
extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
reductionBuffer[blockIdx.x] = sPartials[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;
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::update(sPartials[threadIdx.x], reinterpret_cast<X*>(vreductionBuffer)[i], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
sPartials,
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (threadIdx.x == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(vreductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraParams);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceSameFunction<X>::intermediate(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* dXShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
const sd::LongType* dims) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
// If input is empty, skip unless z is also empty
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal =
|
||||
static_cast<X>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
auto res = cudaMemcpyAsync(
|
||||
sd::LaunchContext::defaultContext()->getScalarPointer(),
|
||||
&startingVal,
|
||||
sizeof(X),
|
||||
cudaMemcpyHostToDevice,
|
||||
*stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceSameFunction<X>::intermediate: failed to copy temporary scalar", res);
|
||||
}
|
||||
|
||||
auto ptr = sd::LaunchContext::defaultContext()->getScalarPointer();
|
||||
// scalar assign
|
||||
scalar::ScalarTransform<X,X,X>::executeCudaShaped(
|
||||
launchDims,
|
||||
stream,
|
||||
14,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hXShapeInfo,
|
||||
z,
|
||||
dZShapeInfo,
|
||||
hZShapeInfo,
|
||||
ptr,
|
||||
nullptr);
|
||||
} else {
|
||||
const sd::LongType zRank = shape::rank(hZShapeInfo);
|
||||
const sd::LongType tadRank = shape::rank(hXShapeInfo) - zRank;
|
||||
|
||||
auto outerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(
|
||||
const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims), zRank);
|
||||
auto innerPack = sd::ConstantShapeHelper::getInstance()
|
||||
.createSubArrShapeInfo(
|
||||
const_cast<sd::LongType*>(hXShapeInfo), const_cast<sd::LongType*>(dims + zRank), tadRank);
|
||||
|
||||
simpleReduce<X, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
outerPack->special(),
|
||||
innerPack->special(),
|
||||
extraParams,
|
||||
vreductionBuffer,
|
||||
z,
|
||||
dZShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceSameFunction intermediate(...) failed");
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_HOST SD_INLINE void ReduceSameFunction<X>::intermediateScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
if (shape::isEmptyConst(hXShapeInfo)) {
|
||||
if (shape::isEmptyConst(hZShapeInfo)) return;
|
||||
|
||||
const auto startingVal =
|
||||
static_cast<X>(OpType::startingValue(reinterpret_cast<const X*>(x)));
|
||||
auto res = cudaMemcpyAsync(z, &startingVal, sizeof(X),
|
||||
cudaMemcpyHostToDevice, *stream);
|
||||
if (res != 0) {
|
||||
throw sd::cuda_exception::build(
|
||||
"ReduceSameFunction<X>::intermediateScalar: failed to copy resulting scalar",
|
||||
res);
|
||||
}
|
||||
} else {
|
||||
simpleScalar<X, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
sd::DebugHelper::checkErrorCode(stream, "ReduceSameFunction intermediateScalar(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_HOST SD_INLINE void ReduceSameFunction<X>::execReduceScalar(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
void* reductionBuffer,
|
||||
const sd::LongType* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
intermediateScalar,
|
||||
PARAMS(
|
||||
launchDims, stream, x, xShapeInfo, hXShapeInfo, extraParams, z,
|
||||
zShapeInfo, hZShapeInfo, dimension, dimensionLength, reductionBuffer, tadOnlyShapeInfo),
|
||||
REDUCE_SAME_OPS);
|
||||
sd::DebugHelper::checkErrorCode(stream, "execReduceScalarSame(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_HOST SD_INLINE void ReduceSameFunction<X>::execReduce(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* dXShapeInfo,
|
||||
const sd::LongType* hXShapeInfo,
|
||||
void* extraParams,
|
||||
void* vreductionBuffer,
|
||||
void* z,
|
||||
const sd::LongType* dZShapeInfo,
|
||||
const sd::LongType* hZShapeInfo,
|
||||
const sd::LongType* dims) {
|
||||
|
||||
if (shape::length(hZShapeInfo) == 1) {
|
||||
execReduceScalar(
|
||||
launchDims, stream, opNum,
|
||||
x, dXShapeInfo, hXShapeInfo,
|
||||
extraParams, z,
|
||||
dZShapeInfo, hZShapeInfo,
|
||||
nullptr, 0, vreductionBuffer, nullptr);
|
||||
} else {
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
intermediate,
|
||||
PARAMS(
|
||||
launchDims, stream, x, dXShapeInfo, hXShapeInfo, extraParams, vreductionBuffer, z, dZShapeInfo,
|
||||
hZShapeInfo, dims),
|
||||
REDUCE_SAME_OPS);
|
||||
}
|
||||
DEBUG_KERNEL(stream, opNum);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
SD_DEVICE void initializeShared(X* extraParams, X** sPartials, int sMemSize) {
|
||||
int sPartialsLength = sMemSize / sizeof(X);
|
||||
X* sPartialsDeref = reinterpret_cast<X*>(*sPartials);
|
||||
for (int i = 0; i < sPartialsLength; i++) {
|
||||
sPartialsDeref[i] = extraParams[0];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#define INSTANT_PROCESS_SINGLE(a1) \
|
||||
template void functions::reduce::ReduceSameFunction<GET_SECOND(a1)>::execReduce( \
|
||||
dim3 launchDims, \
|
||||
cudaStream_t* stream, \
|
||||
const int opNum, \
|
||||
const void* x, \
|
||||
const sd::LongType* dXShapeInfo, \
|
||||
const sd::LongType* hXShapeInfo, \
|
||||
void* extraParams, \
|
||||
void* vreductionBuffer, \
|
||||
void* z, \
|
||||
const sd::LongType* dZShapeInfo, \
|
||||
const sd::LongType* hZShapeInfo, \
|
||||
const sd::LongType* dims); \
|
||||
\
|
||||
template void functions::reduce::ReduceSameFunction<GET_SECOND(a1)>::execReduceScalar( \
|
||||
dim3 launchDims, \
|
||||
cudaStream_t* stream, \
|
||||
int opNum, \
|
||||
const void* x, \
|
||||
const sd::LongType* xShapeInfo, \
|
||||
const sd::LongType* hXShapeInfo, \
|
||||
void* extraParams, \
|
||||
void* z, \
|
||||
const sd::LongType* zShapeInfo, \
|
||||
const sd::LongType* hZShapeInfo, \
|
||||
sd::LongType* dimension, \
|
||||
sd::LongType dimensionLength, \
|
||||
void* reductionBuffer, \
|
||||
const sd::LongType* tadOnlyShapeInfo);
|
||||
|
||||
ITERATE_LIST((SD_COMMON_TYPES), INSTANT_PROCESS_SINGLE)
|
||||
|
||||
} // namespace reduce
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,703 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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), created on 19.11.2018
|
||||
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <loops/reduce3.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <types/types.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
namespace functions {
|
||||
namespace reduce3 {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL void execScalarGeneric(const int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void *extraParams,
|
||||
void *vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void *reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo) {
|
||||
|
||||
Reduce3<X, Z>::execScalarCuda(opNum,
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
allocationPointer,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL void execAllGeneric(const int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void *extraParams,
|
||||
void *vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets) {
|
||||
|
||||
Reduce3<X, Z>::execAllCuda(opNum,
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
postProcessOrNot,
|
||||
allocationPointer,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL void execGeneric(const int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void *extraParams,
|
||||
void *vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets) {
|
||||
|
||||
Reduce3<X, Z>::execCuda(opNum,
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
postProcessOrNot,
|
||||
allocationPointer,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Reduce3<X, Z>::aggregatePartials(void* vsPartials,
|
||||
sd::LongType tid,
|
||||
sd::LongType numItems,
|
||||
void* vextraParams) {
|
||||
|
||||
auto sPartials = reinterpret_cast<Z*>(vsPartials);
|
||||
auto extraParams = reinterpret_cast<Z*>(vextraParams);
|
||||
|
||||
sd::LongType floorPow2 = numItems;
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1))
|
||||
floorPow2 &= floorPow2 - 1;
|
||||
|
||||
if (tid >= floorPow2)
|
||||
sPartials[tid - floorPow2] =
|
||||
OpType::update(sPartials[tid - floorPow2], sPartials[tid], extraParams);
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads) {
|
||||
sPartials[tid] =
|
||||
OpType::update(sPartials[tid], sPartials[tid + activeThreads], extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Reduce3<X, Z>::execScalarCuda(
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo) {
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ Z extraZ[3]; // just 3 values used in logic below
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
// Cache rank/shape/stride for x and y in __shared__
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
extraZ[0] = (Z) 0.0f;
|
||||
extraZ[1] = (Z) 0.0f;
|
||||
extraZ[2] = extraParams != nullptr
|
||||
? (reinterpret_cast<Z*>(extraParams))[2]
|
||||
: (Z) 0.0f;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
char xOrder = shape::order(xShapeInfo);
|
||||
char yOrder = shape::order(yShapeInfo);
|
||||
sd::LongType gridSize = gridDim.x * blockDim.x;
|
||||
|
||||
// fill partial sums
|
||||
for (sd::LongType i = tid; i < length; i += gridSize) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
|
||||
COORDS2INDEX(yRank, yStridePtr, yCoords, yOffset);
|
||||
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::update(sPartials[threadIdx.x],
|
||||
OpType::opAtomic(x[xOffset], y[yOffset], extraZ),
|
||||
extraZ);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
// reduce partial sums
|
||||
aggregatePartials<OpType>(
|
||||
reinterpret_cast<void*>(sPartials),
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, length),
|
||||
extraZ);
|
||||
__syncthreads();
|
||||
|
||||
// multi-block reduce
|
||||
if (gridDim.x > 1) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
__shared__ bool amLast;
|
||||
tid = threadIdx.x;
|
||||
Z* extraBuffer = reinterpret_cast<Z*>(allocationPointer);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
reinterpret_cast<Z*>(reductionBuffer)[blockIdx.x] = sPartials[0];
|
||||
extraBuffer[blockIdx.x] = extraZ[0];
|
||||
extraBuffer[gridDim.x + blockIdx.x] = extraZ[1];
|
||||
}
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
sPartials[tid] = OpType::startingValue(x);
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
tc[16384] = 0; // reset
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
|
||||
if (tid == 0 && extraZ[0] != static_cast<Z>(0) && extraZ[1] != static_cast<Z>(0)) {
|
||||
extraZ[0] = 0.0;
|
||||
extraZ[1] = 0.0;
|
||||
for (int i = 0; i < gridDim.x; i++) {
|
||||
extraZ[0] += extraBuffer[i];
|
||||
extraZ[1] += extraBuffer[gridDim.x + i];
|
||||
}
|
||||
}
|
||||
for (sd::LongType i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::update(sPartials[threadIdx.x],
|
||||
(reinterpret_cast<Z*>(reductionBuffer))[i],
|
||||
extraZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
aggregatePartials<OpType>(
|
||||
reinterpret_cast<void*>(sPartials),
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(gridDim.x, blockDim.x),
|
||||
extraZ);
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraZ);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (tid == 0) {
|
||||
auto tc = reinterpret_cast<unsigned int*>(reductionBuffer);
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::postProcess(sPartials[0], length, extraZ);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Reduce3<X, Z>::transformAll(
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
int postProcessOrNot, sd::LongType* allocationPointer,
|
||||
sd::LongType const* xTadShapeInfo, sd::LongType const* xOffsets,
|
||||
sd::LongType const* yTadShapeInfo, sd::LongType const* yOffsets) {
|
||||
|
||||
auto dx = reinterpret_cast<const X*>(vx);
|
||||
auto dy = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
__shared__ Z extraZ[OpType::extraParamsLen > 0 ? OpType::extraParamsLen : 1];
|
||||
|
||||
__shared__ int xTadLength;
|
||||
__shared__ int yTadLength;
|
||||
__shared__ int xTads;
|
||||
__shared__ int yTads;
|
||||
|
||||
// Cache shape info for xTadShapeInfo and yTadShapeInfo
|
||||
__shared__ sd::LongType xTadRank;
|
||||
__shared__ const sd::LongType* xTadShapePtr;
|
||||
__shared__ const sd::LongType* xTadStridePtr;
|
||||
|
||||
__shared__ sd::LongType yTadRank;
|
||||
__shared__ const sd::LongType* yTadShapePtr;
|
||||
__shared__ const sd::LongType* yTadStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
sPartials[threadIdx.x] = OpType::startingValue(dx);
|
||||
|
||||
xTadLength = shape::length(xTadShapeInfo);
|
||||
yTadLength = shape::length(yTadShapeInfo);
|
||||
|
||||
xTads = shape::length(xShapeInfo) / xTadLength;
|
||||
yTads = shape::length(yShapeInfo) / yTadLength;
|
||||
|
||||
xTadRank = shape::rank(xTadShapeInfo);
|
||||
xTadShapePtr = shape::shapeOf(xTadShapeInfo);
|
||||
xTadStridePtr = shape::stride(xTadShapeInfo);
|
||||
|
||||
yTadRank = shape::rank(yTadShapeInfo);
|
||||
yTadShapePtr = shape::shapeOf(yTadShapeInfo);
|
||||
yTadStridePtr = shape::stride(yTadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
Z startingVal = OpType::startingValue(dx);
|
||||
|
||||
const int maxBlock = blockDim.x;
|
||||
const int limit = (xTadLength + maxBlock - 1) / maxBlock; // ceiling division
|
||||
|
||||
for (int r = blockIdx.x; r < xTads; r += (blockDim.x * gridDim.x)) {
|
||||
// load partial x
|
||||
auto xLocal = dx + xOffsets[r];
|
||||
|
||||
// Fill tile for x once if thread < xTadLength
|
||||
__shared__ X sXCache[SD_CUDA_BLOCK_SIZE];
|
||||
if (threadIdx.x < xTadLength && threadIdx.x < maxBlock) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOff;
|
||||
INDEX2COORDS(threadIdx.x, xTadRank, xTadShapePtr, xCoords);
|
||||
COORDS2INDEX(xTadRank, xTadStridePtr, xCoords, xOff);
|
||||
sXCache[threadIdx.x] = xLocal[xOff];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int g = 0; g < yTads; g++) {
|
||||
auto yLocal = dy + yOffsets[g];
|
||||
int ri = (r * yTads) + g;
|
||||
|
||||
sPartials[threadIdx.x] = startingVal;
|
||||
if (OpType::extraParamsLen > 0 && threadIdx.x < OpType::extraParamsLen) {
|
||||
extraZ[threadIdx.x] = startingVal;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Possibly multiple tiles per x
|
||||
for (int t = 0; t < limit; t++) {
|
||||
// re-fetch x tile if needed
|
||||
if (t >= 1 && (threadIdx.x + t * maxBlock < xTadLength)) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOff;
|
||||
INDEX2COORDS(threadIdx.x + t * maxBlock,
|
||||
xTadRank, xTadShapePtr, xCoords);
|
||||
COORDS2INDEX(xTadRank, xTadStridePtr, xCoords, xOff);
|
||||
sXCache[threadIdx.x] = xLocal[xOff];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// compute partials
|
||||
for (int f = threadIdx.x + t * maxBlock; (f < xTadLength) && (f < (t + 1) * maxBlock); f += (blockDim.x * gridDim.x)) {
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType yOff;
|
||||
INDEX2COORDS(f, yTadRank, yTadShapePtr, yCoords);
|
||||
COORDS2INDEX(yTadRank, yTadStridePtr, yCoords, yOff);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::opAtomic(sXCache[threadIdx.x], yLocal[yOff], extraZ),
|
||||
extraZ);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// reduce partials
|
||||
aggregatePartials<OpType>(
|
||||
reinterpret_cast<void*>(sPartials),
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, xTadLength),
|
||||
extraZ);
|
||||
__syncthreads();
|
||||
|
||||
// store final
|
||||
if (threadIdx.x == 0) {
|
||||
z[ri] = OpType::postProcess(sPartials[0], xTadLength, extraZ);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void Reduce3<X, Z>::transform(
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets) {
|
||||
|
||||
if (shape::isScalar(zShapeInfo))
|
||||
return; // no-op if scalar
|
||||
|
||||
if (yTadOnlyShapeInfo == nullptr)
|
||||
yTadOnlyShapeInfo = yShapeInfo; // execReduce3TAD case
|
||||
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto y = reinterpret_cast<const X*>(vy);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
|
||||
// For partial sums
|
||||
__shared__ Z extraZ[OpType::extraParamsLen > 0 ? OpType::extraParamsLen : 1];
|
||||
__shared__ Z sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
__shared__ sd::LongType tadLen;
|
||||
__shared__ sd::LongType zLen;
|
||||
__shared__ sd::LongType yTadNum;
|
||||
__shared__ sd::LongType xTadEws;
|
||||
__shared__ sd::LongType yTadEws;
|
||||
__shared__ char xTadOrder;
|
||||
__shared__ char yTadOrder;
|
||||
|
||||
// Cache shape info
|
||||
__shared__ sd::LongType xTadRank;
|
||||
__shared__ const sd::LongType* xTadShape;
|
||||
__shared__ const sd::LongType* xTadStride;
|
||||
|
||||
__shared__ sd::LongType yTadRank;
|
||||
__shared__ const sd::LongType* yTadShape;
|
||||
__shared__ const sd::LongType* yTadStride;
|
||||
|
||||
__shared__ sd::LongType zRank; // Might not be used, but let's keep consistent
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLen = shape::length(tadOnlyShapeInfo);
|
||||
zLen = shape::length(zShapeInfo);
|
||||
yTadNum = shape::length(yShapeInfo) / tadLen;
|
||||
xTadEws = shape::elementWiseStride(tadOnlyShapeInfo);
|
||||
yTadEws = shape::elementWiseStride(yTadOnlyShapeInfo);
|
||||
xTadOrder = shape::order(tadOnlyShapeInfo);
|
||||
yTadOrder = shape::order(yTadOnlyShapeInfo);
|
||||
|
||||
xTadRank = shape::rank(tadOnlyShapeInfo);
|
||||
xTadShape = shape::shapeOf(tadOnlyShapeInfo);
|
||||
xTadStride = shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
yTadRank = shape::rank(yTadOnlyShapeInfo);
|
||||
yTadShape = shape::shapeOf(yTadOnlyShapeInfo);
|
||||
yTadStride = shape::stride(yTadOnlyShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
sPartials[threadIdx.x] = OpType::startingValue(x);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
Z startingVal = OpType::startingValue(x);
|
||||
|
||||
// either a direct ews approach or fallback
|
||||
for (int i = blockIdx.x; i < zLen; i += gridDim.x) {
|
||||
sd::LongType xBaseOffset = tadOffsets[i];
|
||||
sd::LongType yBaseOffset = (yTadNum == 1 ? 0 : yTadOffsets[i]);
|
||||
|
||||
if (OpType::extraParamsLen > 0 && threadIdx.x < OpType::extraParamsLen) {
|
||||
extraZ[threadIdx.x] = startingVal;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// partial sums
|
||||
for (int j = threadIdx.x; j < tadLen; j += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType xOff;
|
||||
sd::LongType yOff;
|
||||
|
||||
INDEX2COORDS(j, xTadRank, xTadShape, xCoords);
|
||||
COORDS2INDEX(xTadRank, xTadStride, xCoords, xOff);
|
||||
|
||||
INDEX2COORDS(j, yTadRank, yTadShape, yCoords);
|
||||
COORDS2INDEX(yTadRank, yTadStride, yCoords, yOff);
|
||||
|
||||
// update partial
|
||||
if (j < blockDim.x) {
|
||||
sPartials[threadIdx.x] =
|
||||
OpType::opAtomic(x[xBaseOffset + xOff],
|
||||
y[yBaseOffset + yOff],
|
||||
extraZ);
|
||||
} else {
|
||||
sPartials[threadIdx.x] = OpType::update(
|
||||
sPartials[threadIdx.x],
|
||||
OpType::opAtomic(x[xBaseOffset + xOff],
|
||||
y[yBaseOffset + yOff],
|
||||
extraZ),
|
||||
extraZ);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// reduce partials
|
||||
aggregatePartials<OpType>(
|
||||
reinterpret_cast<void*>(sPartials),
|
||||
threadIdx.x,
|
||||
sd::math::sd_min<int>(blockDim.x, tadLen),
|
||||
extraZ);
|
||||
__syncthreads();
|
||||
|
||||
// write final
|
||||
if (threadIdx.x == 0) {
|
||||
z[i] = OpType::postProcess(sPartials[0], tadLen, extraZ);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_DEVICE void Reduce3<X, Y>::execCuda(
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
int postProcessOrNot, sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo, sd::LongType const* yTadOffsets) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
transform,
|
||||
PARAMS(vx, xShapeInfo, vy, yShapeInfo,
|
||||
extraParams, vz, zShapeInfo,
|
||||
dimension, dimensionLength, postProcessOrNot,
|
||||
allocationPointer, tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets),
|
||||
REDUCE3_OPS);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_DEVICE void Reduce3<X, Y>::execAllCuda(
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength,
|
||||
int postProcessOrNot, sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo, sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo, sd::LongType const* yTadOffsets) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
transformAll,
|
||||
PARAMS(vx, xShapeInfo, vy, yShapeInfo,
|
||||
extraParams, vz, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
postProcessOrNot, allocationPointer,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets),
|
||||
REDUCE3_OPS);
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_DEVICE void Reduce3<X, Y>::execScalarCuda(
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
execScalarCuda,
|
||||
PARAMS(vx, xShapeInfo, vy, yShapeInfo,
|
||||
extraParams, vz, zShapeInfo,
|
||||
allocationPointer, reductionBuffer, tadOnlyShapeInfo),
|
||||
REDUCE3_OPS);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void Reduce3<X, Z>::exec(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets) {
|
||||
|
||||
execGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum,
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
postProcessOrNot,
|
||||
allocationPointer,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "reduce3exec(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void Reduce3<X, Z>::execAll(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* dimension, long long int dimensionLength,
|
||||
int postProcessOrNot,
|
||||
sd::LongType* allocationPointer,
|
||||
sd::LongType const* tadOnlyShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* yTadOnlyShapeInfo,
|
||||
sd::LongType const* yTadOffsets) {
|
||||
|
||||
execAllGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum, vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
dimension, dimensionLength,
|
||||
postProcessOrNot,
|
||||
allocationPointer,
|
||||
tadOnlyShapeInfo, tadOffsets,
|
||||
yTadOnlyShapeInfo, yTadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "execAllGeneric(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void Reduce3<X, Z>::execScalar(
|
||||
dim3 launchDims, cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx, sd::LongType const* xShapeInfo,
|
||||
void const* vy, sd::LongType const* yShapeInfo,
|
||||
void* extraParams,
|
||||
void* vz, sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionBuffer,
|
||||
sd::LongType const* tadOnlyShapeInfo) {
|
||||
|
||||
execScalarGeneric<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum,
|
||||
vx, xShapeInfo,
|
||||
vy, yShapeInfo,
|
||||
extraParams,
|
||||
vz, zShapeInfo,
|
||||
allocationPointer,
|
||||
reductionBuffer,
|
||||
tadOnlyShapeInfo);
|
||||
sd::DebugHelper::checkErrorCode(stream, "execScalarGeneric(...) failed");
|
||||
}
|
||||
|
||||
} // namespace reduce3
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,304 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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
|
||||
//
|
||||
|
||||
#ifndef SCALAR_CU
|
||||
#define SCALAR_CU
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <loops/pairwise_instantiations.h>
|
||||
|
||||
#include "../../types/types.h"
|
||||
#include "loops/scalar.h"
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z, typename OpType>
|
||||
SD_KERNEL static void scalarSimpleShaped(
|
||||
void const* x,
|
||||
void const* y,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* params,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationBuffer) {
|
||||
|
||||
auto xVals = reinterpret_cast<X const*>(x);
|
||||
auto scalar = reinterpret_cast<Y const*>(y)[0];
|
||||
auto pVals = reinterpret_cast<Z*>(params);
|
||||
auto zVals = reinterpret_cast<Z*>(z);
|
||||
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
|
||||
//need special handling for scalars as strides may not be set for scalars
|
||||
if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
|
||||
z[0] = OpType::op(x[0], scalar[0], extraParams);
|
||||
return;
|
||||
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
//need special handling for scalars as strides may not be set for scalars
|
||||
if(shape::length(xShapeInfo) <= 1 && shape::length(zShapeInfo) <= 1) {
|
||||
return;
|
||||
|
||||
}
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
zVals[zOffset] = OpType::op(xVals[xOffset], scalar, pVals);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z, typename OpType>
|
||||
SD_KERNEL static void scalarAlongDimension(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* scalars,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
auto xVals = reinterpret_cast<X const*>(x);
|
||||
auto zVals = reinterpret_cast<Z*>(z);
|
||||
auto pVals = reinterpret_cast<Z*>(extraParams);
|
||||
auto sVals = reinterpret_cast<Y const*>(scalars);
|
||||
|
||||
if (tadShapeInfoZ == nullptr) {
|
||||
tadShapeInfoZ = tadShapeInfo;
|
||||
tadOffsetsZ = tadOffsets;
|
||||
}
|
||||
|
||||
__shared__ sd::LongType tadLength;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
|
||||
__shared__ sd::LongType tadRankZ;
|
||||
__shared__ const sd::LongType* tadShapePtrZ;
|
||||
__shared__ const sd::LongType* tadStridePtrZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadShapeInfo);
|
||||
tadStridePtr = shape::stride(tadShapeInfo);
|
||||
|
||||
tadRankZ = shape::rank(tadShapeInfoZ);
|
||||
tadShapePtrZ = shape::shapeOf(tadShapeInfoZ);
|
||||
tadStridePtrZ = shape::stride(tadShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto zPart = zVals + tadOffsetsZ[r];
|
||||
auto xPart = xVals + tadOffsets[r];
|
||||
auto scalarV = sVals[r];
|
||||
|
||||
for (sd::LongType f = threadIdx.x; f < tadLength; f += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(f, tadRank, tadShapePtr, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(f, tadRankZ, tadShapePtrZ, zCoords);
|
||||
COORDS2INDEX(tadRankZ, tadStridePtrZ, zCoords, zOffset);
|
||||
|
||||
zPart[zOffset] = OpType::op(xPart[xOffset], scalarV, pVals);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace scalar {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
void SD_HOST ScalarTransform<X, Y, Z>::intermediateShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
sd::LongType const* hxShapeInfo, // Unused in the kernel
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType const* hzShapeInfo, // Unused in the kernel
|
||||
void const* vscalar,
|
||||
void* vextraParams,
|
||||
sd::LongType* allocPointer) {
|
||||
|
||||
scalarSimpleShaped<X, Y, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
vscalar,
|
||||
xShapeInfo,
|
||||
vextraParams,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
allocPointer);
|
||||
sd::DebugHelper::checkErrorCode(stream, "scalarSimpleShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
template <typename OpType>
|
||||
void SD_HOST ScalarTransform<X, Y, Z>::intermediateAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* scalars,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
scalarAlongDimension<X, Y, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
scalars,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadShapeInfo,
|
||||
tadOffsets,
|
||||
tadShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
sd::DebugHelper::checkErrorCode(stream, "scalarAlongDimension(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
void ScalarTransform<X, Y, Z>::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
sd::LongType const* hxShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType const* hzShapeInfo,
|
||||
void const* vscalar,
|
||||
void* vextraParams) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, hxShapeInfo, vz, zShapeInfo, hzShapeInfo, vscalar, vextraParams, nullptr),
|
||||
SCALAR_OPS);
|
||||
}
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y, typename Z>
|
||||
void ScalarTransform<X, Y, Z>::executeCudaAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* vscalars,
|
||||
void* vextraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ) {
|
||||
|
||||
DISPATCH_BY_OPNUM_TTT(
|
||||
intermediateAlongDimension,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo, vscalars, vextraParams, dimension, dimensionLength,
|
||||
tadShapeInfo, tadOffsets, tadShapeInfoZ, tadOffsetsZ),
|
||||
SCALAR_OPS);
|
||||
}
|
||||
|
||||
} // namespace scalar
|
||||
} // namespace functions
|
||||
|
||||
#endif // SCALAR_CU
|
||||
@@ -0,0 +1,364 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
* 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 08.11.2018
|
||||
// @author raver119@gmail.com
|
||||
|
||||
#ifndef SCALAR_BOOL_CU
|
||||
#define SCALAR_BOOL_CU
|
||||
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
#include "../legacy_ops.h"
|
||||
#include "../scalar_bool.h"
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <system/Environment.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// A kernel that applies a scalar bool transform along a specific dimension (TAD).
|
||||
// It uses shared memory caching for relevant shape information to reduce overhead.
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void scalarAlongDimensionCachedKernel(
|
||||
void const* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void const* scalars, // per-TAD scalars
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets,
|
||||
const sd::LongType* tadShapeInfoZ, // if null, same as x TAD
|
||||
const sd::LongType* tadOffsetsZ) {
|
||||
|
||||
auto xTyped = reinterpret_cast<const X*>(x);
|
||||
auto zTyped = reinterpret_cast<Z*>(z);
|
||||
auto extra = reinterpret_cast<X*>(extraParams);
|
||||
auto scalarsTyped = reinterpret_cast<const X*>(scalars);
|
||||
|
||||
// If not provided, fallback
|
||||
const auto* actualTadShapeInfoZ = (tadShapeInfoZ == nullptr ? tadShapeInfo : tadShapeInfoZ);
|
||||
const auto* actualTadOffsetsZ = (tadShapeInfoZ == nullptr ? tadOffsets : tadOffsetsZ);
|
||||
|
||||
// Cache shape info in shared memory
|
||||
__shared__ sd::LongType tadLen;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
__shared__ int tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
|
||||
__shared__ int tadRankZ;
|
||||
__shared__ const sd::LongType* tadShapePtrZ;
|
||||
__shared__ const sd::LongType* tadStridePtrZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLen = shape::length(tadShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLen;
|
||||
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadShapeInfo);
|
||||
tadStridePtr= shape::stride(tadShapeInfo);
|
||||
|
||||
tadRankZ = shape::rank(actualTadShapeInfoZ);
|
||||
tadShapePtrZ = shape::shapeOf(actualTadShapeInfoZ);
|
||||
tadStridePtrZ = shape::stride(actualTadShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each block handles multiple TADs
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
Z* zTad = zTyped + actualTadOffsetsZ[r];
|
||||
const X* xTad = xTyped + tadOffsets[r];
|
||||
X scalar = scalarsTyped[r];
|
||||
|
||||
// Each thread processes part of a single TAD
|
||||
for (sd::LongType f = threadIdx.x; f < tadLen; f += blockDim.x) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
sd::LongType offsetX;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
// Compute offset for X TAD
|
||||
INDEX2COORDS(f, tadRank, tadShapePtr, coordsX);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, coordsX, offsetX);
|
||||
|
||||
// Compute offset for Z TAD
|
||||
INDEX2COORDS(f, tadRankZ, tadShapePtrZ, coordsZ);
|
||||
COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, offsetZ);
|
||||
|
||||
zTad[offsetZ] = OpType::op(xTad[offsetX], scalar, extra);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// A kernel to apply a scalar transform to a shaped buffer, with caching logic
|
||||
// for shape info in shared memory to reduce overhead.
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void scalarSimpleShapedCachedKernel(
|
||||
void const* x, // the "scalar" input
|
||||
void const* y, // the "array" input
|
||||
const sd::LongType* xShapeInfo,// we just read rank from here if needed
|
||||
void* params,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationBuffer) {
|
||||
|
||||
auto scalar = reinterpret_cast<const X*>(x)[0];
|
||||
auto yTyped = reinterpret_cast<const X*>(y);
|
||||
auto zTyped = reinterpret_cast<Z*>(z);
|
||||
auto extra = reinterpret_cast<X*>(params);
|
||||
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo); // or maybe shape::length(zShapeInfo)
|
||||
// Actually we only need length from either input array's shape
|
||||
// but let's assume x is scalar, so let's do it from z shape if that's a shaped array
|
||||
// For now we keep as is.
|
||||
|
||||
yRank = shape::rank(xShapeInfo); // or we do: shape::rank(some-other-shape)
|
||||
yShapePtr = shape::shapeOf(xShapeInfo);
|
||||
yStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType offsetY;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
// get offset for Y
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
|
||||
|
||||
// get offset for Z
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, extra);
|
||||
}
|
||||
}
|
||||
|
||||
// *********************************************************************//
|
||||
// *********************************************************************//
|
||||
namespace functions {
|
||||
namespace scalar {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
__device__ void ScalarBoolTransform<X,Z>::transformCuda(
|
||||
void const* vscalar,
|
||||
void const* vy,
|
||||
const sd::LongType* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationBuffer) {
|
||||
|
||||
auto scalar = reinterpret_cast<const X*>(vscalar)[0];
|
||||
auto yTyped = reinterpret_cast<const X*>(vy);
|
||||
auto zTyped = reinterpret_cast<Z*>(vz);
|
||||
auto extra = reinterpret_cast<const X*>(vparams);
|
||||
|
||||
// store shape info in shared memory
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(yShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = blockDim.x * gridDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType offsetY;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, const_cast<X*>(extra));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
__host__ void ScalarBoolTransform<X,Z>::intermediateAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void const* scalars,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets,
|
||||
const sd::LongType* tadShapeInfoZ,
|
||||
const sd::LongType* tadOffsetsZ)
|
||||
{
|
||||
scalarAlongDimensionCachedKernel<X,Z,OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
scalars,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadShapeInfo,
|
||||
tadOffsets,
|
||||
tadShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "scalarAlongDimensionCachedKernel(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
__host__ void ScalarBoolTransform<X,Z>::intermediateShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void const* vscalar,
|
||||
void* vextraParams,
|
||||
sd::LongType* allocPointer)
|
||||
{
|
||||
scalarSimpleShapedCachedKernel<X,Z,OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
vscalar,
|
||||
xShapeInfo,
|
||||
vextraParams,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
allocPointer);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "scalarSimpleShapedCachedKernel(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
__host__ void ScalarBoolTransform<X, Y>::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void const* vscalar,
|
||||
void const* vextraParams)
|
||||
{
|
||||
if (sd::Environment::getInstance().isDebugAndVerbose()) {
|
||||
printf("H14 opNum:[%i]\n", opNum);
|
||||
}
|
||||
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
|
||||
vscalar, const_cast<void*>(vextraParams), nullptr),
|
||||
SCALAR_BOOL_OPS);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
__host__ void ScalarBoolTransform<X, Y>::executeCudaAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
void const* vscalars,
|
||||
void* vextraParams,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets,
|
||||
const sd::LongType* tadShapeInfoZ,
|
||||
const sd::LongType* tadOffsetsZ)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateAlongDimension,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
|
||||
vscalars, vextraParams, dimension, dimensionLength,
|
||||
tadShapeInfo, tadOffsets,
|
||||
tadShapeInfoZ, tadOffsetsZ),
|
||||
SCALAR_BOOL_OPS);
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class ScalarBoolTransform, , SD_COMMON_TYPES, SD_BOOL_TYPES);
|
||||
|
||||
} // namespace scalar
|
||||
} // namespace functions
|
||||
|
||||
#endif // SCALAR_BOOL_CU
|
||||
@@ -0,0 +1,388 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
* 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 08.11.2018
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
#ifndef SCALAR_INT_CU
|
||||
#define SCALAR_INT_CU
|
||||
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
#include "../legacy_ops.h"
|
||||
#include "../scalar_int.h"
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <system/Environment.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// A kernel that applies an integer-based scalar transform along specified dimension (via TAD).
|
||||
// We'll cache shape info in shared memory to reduce repeated calls to shapeOf, strideOf, etc.
|
||||
template <typename X, typename OpType>
|
||||
__global__ void scalarAlongDimensionCachedKernel(
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* scalars,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// delegate the actual transform to the transformCuda method,
|
||||
// caching shape info inside that method.
|
||||
functions::scalar::ScalarIntTransform<X>::template transformCuda<OpType>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
scalars,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadShapeInfo,
|
||||
tadOffsets,
|
||||
tadShapeInfoZ,
|
||||
tadOffsetsZ
|
||||
);
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
// A kernel to handle shaped transforms: x is "scalar," y is "array," but
|
||||
// we store shape data in shared memory.
|
||||
template <typename X, typename OpType>
|
||||
__global__ void scalarSimpleShapedCachedKernel(
|
||||
void const* x,
|
||||
void const* y,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* params,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationBuffer)
|
||||
{
|
||||
// We'll call the transformCuda method (which will do caching).
|
||||
functions::scalar::ScalarIntTransform<X>::template transformCuda<OpType>(
|
||||
y,
|
||||
x,
|
||||
xShapeInfo,
|
||||
params,
|
||||
z,
|
||||
zShapeInfo,
|
||||
allocationBuffer
|
||||
);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace scalar {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
__device__ void ScalarIntTransform<X>::transformCuda(
|
||||
void const* vscalar,
|
||||
void const* vy,
|
||||
sd::LongType const* yShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
sd::LongType* allocationBuffer)
|
||||
{
|
||||
auto scalar = reinterpret_cast<const X*>(vscalar)[0];
|
||||
auto yTyped = reinterpret_cast<const X*>(vy);
|
||||
auto zTyped = reinterpret_cast<X*>(vz);
|
||||
auto extra = reinterpret_cast<X*>(vparams);
|
||||
|
||||
// cache shape info in shared memory
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
__shared__ int yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(yShapeInfo);
|
||||
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr= shape::stride(yShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr= shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// now we do the transform
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsY[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
sd::LongType offsetY;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, coordsY, offsetY);
|
||||
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
zTyped[offsetZ] = OpType::op(yTyped[offsetY], scalar, extra);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
__device__ void ScalarIntTransform<X>::transformCuda(
|
||||
sd::LongType len,
|
||||
void const* vx,
|
||||
void const* vy,
|
||||
sd::LongType yEWS,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
sd::LongType zEWS,
|
||||
sd::LongType* allocationBuffer)
|
||||
{
|
||||
auto x = reinterpret_cast<const X*>(vx)[0]; // scalar
|
||||
auto yTyped= reinterpret_cast<const X*>(vy);
|
||||
auto zTyped= reinterpret_cast<X*>(vz);
|
||||
auto extra = reinterpret_cast<X*>(vparams);
|
||||
|
||||
const int tid = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
const int totalThreads = blockDim.x * gridDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < len; i += totalThreads) {
|
||||
zTyped[i * zEWS] = OpType::op(yTyped[i * yEWS], x, extra);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
__device__ void ScalarIntTransform<X>::transformCuda(
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* vextraParams,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* vscalars,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto scalars = reinterpret_cast<const X*>(vscalars);
|
||||
auto zTyped = reinterpret_cast<X*>(vz);
|
||||
auto extra = reinterpret_cast<X*>(vextraParams);
|
||||
|
||||
// if z TAD not provided, fallback
|
||||
const auto* actualTadShapeInfoZ = (tadShapeInfoZ == nullptr ? tadShapeInfo : tadShapeInfoZ);
|
||||
const auto* actualTadOffsetsZ = (tadShapeInfoZ == nullptr ? tadOffsets : tadOffsetsZ);
|
||||
|
||||
// cache shape info in shared memory
|
||||
__shared__ sd::LongType tadLen;
|
||||
__shared__ sd::LongType numTads;
|
||||
|
||||
__shared__ int tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
|
||||
__shared__ int tadRankZ;
|
||||
__shared__ const sd::LongType* tadShapePtrZ;
|
||||
__shared__ const sd::LongType* tadStridePtrZ;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLen = shape::length(tadShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLen;
|
||||
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadShapeInfo);
|
||||
tadStridePtr = shape::stride(tadShapeInfo);
|
||||
|
||||
tadRankZ = shape::rank(actualTadShapeInfoZ);
|
||||
tadShapePtrZ = shape::shapeOf(actualTadShapeInfoZ);
|
||||
tadStridePtrZ = shape::stride(actualTadShapeInfoZ);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
X* zTad = zTyped + actualTadOffsetsZ[r];
|
||||
const X* xTad = x + tadOffsets[r];
|
||||
X scalar = scalars[r];
|
||||
|
||||
// each thread processes part of TAD
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
|
||||
sd::LongType offsetX;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
INDEX2COORDS(i, tadRank, tadShapePtr, coordsX);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, coordsX, offsetX);
|
||||
|
||||
INDEX2COORDS(i, tadRankZ, tadShapePtrZ, coordsZ);
|
||||
COORDS2INDEX(tadRankZ, tadStridePtrZ, coordsZ, offsetZ);
|
||||
|
||||
zTad[offsetZ] = OpType::op(xTad[offsetX], scalar, extra);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
__host__ void ScalarIntTransform<X>::intermediateAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* x,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* z,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* scalars,
|
||||
void* extraParams,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
// we use the new, cached version
|
||||
scalarAlongDimensionCachedKernel<X,OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShapeInfo,
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
scalars,
|
||||
dimension,
|
||||
dimensionLength,
|
||||
tadShapeInfo,
|
||||
tadOffsets,
|
||||
tadShapeInfoZ,
|
||||
tadOffsetsZ);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform intermediateAlongDimension(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
__host__ void ScalarIntTransform<X>::intermediateShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* vscalar,
|
||||
void* vextraParams,
|
||||
sd::LongType* allocPointer)
|
||||
{
|
||||
// call the new cached kernel
|
||||
scalarSimpleShapedCachedKernel<X,OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
vscalar,
|
||||
xShapeInfo,
|
||||
vextraParams,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
allocPointer);
|
||||
|
||||
sd::DebugHelper::checkGlobalErrorCode("scalarSimpleShapedCachedKernel(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
__host__ void ScalarIntTransform<X>::executeCudaShaped(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* vscalar,
|
||||
void* vextraParams)
|
||||
{
|
||||
if (sd::Environment::getInstance().isDebugAndVerbose()) {
|
||||
printf("H14 scalar int transform opNum:[%i]\n", opNum);
|
||||
}
|
||||
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo, vscalar, vextraParams, nullptr),
|
||||
SCALAR_INT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform executeCudaShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename X>
|
||||
__host__ void ScalarIntTransform<X>::executeCudaAlongDimension(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
void const* vx,
|
||||
sd::LongType const* xShapeInfo,
|
||||
void* vz,
|
||||
sd::LongType const* zShapeInfo,
|
||||
void const* vscalars,
|
||||
void* vextraParams,
|
||||
sd::LongType* dimension,
|
||||
long long int dimensionLength,
|
||||
sd::LongType const* tadShapeInfo,
|
||||
sd::LongType const* tadOffsets,
|
||||
sd::LongType const* tadShapeInfoZ,
|
||||
sd::LongType const* tadOffsetsZ)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_T(
|
||||
intermediateAlongDimension,
|
||||
PARAMS(launchDims, stream, vx, xShapeInfo, vz, zShapeInfo,
|
||||
vscalars, vextraParams, dimension, dimensionLength,
|
||||
tadShapeInfo, tadOffsets, tadShapeInfoZ, tadOffsetsZ),
|
||||
SCALAR_INT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "ScalarIntTransform executeCudaAlongDimension(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( class ScalarIntTransform, , SD_INTEGER_TYPES);
|
||||
|
||||
} // namespace scalar
|
||||
} // namespace functions
|
||||
|
||||
#endif // SCALAR_INT_CU
|
||||
@@ -0,0 +1,88 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
/**
|
||||
* This kernel accumulates X arrays, and stores z into Z
|
||||
*
|
||||
* @tparam T
|
||||
* @param x
|
||||
* @param z
|
||||
* @param n
|
||||
* @param length
|
||||
*/
|
||||
template <typename T>
|
||||
SD_DEVICE void accumulateKernel(void **vx, void *vz, int n, const LongType length) {
|
||||
auto x = reinterpret_cast<T **>(vx);
|
||||
auto z = reinterpret_cast<T *>(vz);
|
||||
|
||||
__shared__ T *shmem;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char sharedmem[];
|
||||
shmem = (T *)sharedmem;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockDim.x * blockIdx.x; r < length; r += blockDim.x * gridDim.x) {
|
||||
shmem[threadIdx.x] = 0.0f;
|
||||
|
||||
LongType baseIdx = r;
|
||||
|
||||
// aggregation step, we roll over all arrays
|
||||
for (int ar = 0; ar < n; ar++) {
|
||||
T *cdata = (T *)x[ar];
|
||||
cdata += baseIdx;
|
||||
|
||||
if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] += cdata[threadIdx.x];
|
||||
}
|
||||
|
||||
T *wdata = z + baseIdx;
|
||||
|
||||
// saving accumulated values
|
||||
if (baseIdx + threadIdx.x < length) wdata[threadIdx.x] = shmem[threadIdx.x];
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execAccumulateKernel(void **vx, void *vz, int n, const LongType length) {
|
||||
accumulateKernel<T>(vx, vz, n, length);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void accumulateKernelGeneric(dim3 &launchDims, cudaStream_t *stream, void **vx, void *vz, int n,
|
||||
const LongType length) {
|
||||
execAccumulateKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vx, vz, n, length);
|
||||
DebugHelper::checkErrorCode(stream, "accumulate(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void accumulateKernelGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void **vx, void *vz, int n, const sd::LongType length),
|
||||
SD_COMMON_TYPES);
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,97 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_DEVICE void averagingKernel(void **vdx, void *vdz, int n, LongType length, bool propagate) {
|
||||
auto dx = reinterpret_cast<T **>(vdx);
|
||||
auto dz = reinterpret_cast<T *>(vdz);
|
||||
|
||||
__shared__ T *shmem;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ unsigned char sharedmem[];
|
||||
shmem = (T *)sharedmem;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// each block cycles over it's own part of arrays
|
||||
for (int r = blockDim.x * blockIdx.x; r < length; r += blockDim.x * gridDim.x) {
|
||||
shmem[threadIdx.x] = (T)0.0f;
|
||||
|
||||
LongType baseIdx = r;
|
||||
|
||||
// aggregation step, we roll over all arrays
|
||||
for (int ar = 0; ar < n; ar++) {
|
||||
T *cdata = (T *)dx[ar];
|
||||
cdata += baseIdx;
|
||||
|
||||
if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] += cdata[threadIdx.x];
|
||||
}
|
||||
|
||||
// average data in shared memory
|
||||
if (baseIdx + threadIdx.x < length) shmem[threadIdx.x] /= n;
|
||||
|
||||
// div step & write out step
|
||||
if (dz != nullptr) {
|
||||
T *wdata = dz + baseIdx;
|
||||
|
||||
if (baseIdx + threadIdx.x < length) {
|
||||
wdata[threadIdx.x] = shmem[threadIdx.x];
|
||||
}
|
||||
}
|
||||
|
||||
// propagate averaged data to all arrays
|
||||
if (propagate)
|
||||
for (int ar = 0; ar < n; ar++) {
|
||||
T *cdata = (T *)dx[ar];
|
||||
cdata += baseIdx;
|
||||
|
||||
if (baseIdx + threadIdx.x < length) cdata[threadIdx.x] = shmem[threadIdx.x];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execAveragingKernel(void **vdx, void *vdz, int n, LongType length, bool propagate) {
|
||||
averagingKernel<T>(vdx, vdz, n, length, propagate);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void averagingKernelGeneric(dim3 &launchDims, cudaStream_t *stream, void **vdx, void *vdz, int n,
|
||||
LongType length, bool propagate) {
|
||||
execAveragingKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vdx, vdz, n, length, propagate);
|
||||
DebugHelper::checkErrorCode(stream, "averaging(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void averagingKernelGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void **vdx, void *vdz, int n, sd::LongType length,
|
||||
bool propagate),
|
||||
SD_COMMON_TYPES);
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,294 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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);
|
||||
@@ -0,0 +1,257 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 bitonicSortStepKernelKey(
|
||||
void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vy,
|
||||
const sd::LongType* yShapeInfo,
|
||||
int j,
|
||||
int k,
|
||||
int length,
|
||||
bool descending) {
|
||||
|
||||
auto x = static_cast<X*>(vx);
|
||||
auto y = static_cast<Y*>(vy);
|
||||
const unsigned int i = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ sd::LongType yRank;
|
||||
__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();
|
||||
|
||||
if (i >= static_cast<unsigned int>(length)) return;
|
||||
|
||||
const unsigned int ixj = i ^ j;
|
||||
if (ixj <= i) return;
|
||||
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType ixjCoords[SD_MAX_RANK];
|
||||
sd::LongType iOffset;
|
||||
sd::LongType ixjOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, iCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, iCoords, iOffset);
|
||||
|
||||
INDEX2COORDS(ixj, xRank, xShapePtr, ixjCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, ixjCoords, ixjOffset);
|
||||
|
||||
const bool ascending = ((i & k) == 0);
|
||||
X xi = x[iOffset];
|
||||
X xixj = x[ixjOffset];
|
||||
|
||||
if (ascending) {
|
||||
// Sort ascending
|
||||
if (!descending == (xi > xixj)) {
|
||||
x[iOffset] = xixj;
|
||||
x[ixjOffset] = xi;
|
||||
|
||||
sd::LongType iCoordsY[SD_MAX_RANK];
|
||||
sd::LongType ixjCoordsY[SD_MAX_RANK];
|
||||
sd::LongType iOffsetY;
|
||||
sd::LongType ixjOffsetY;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, iCoordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, iCoordsY, iOffsetY);
|
||||
|
||||
INDEX2COORDS(ixj, yRank, yShapePtr, ixjCoordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, ixjCoordsY, ixjOffsetY);
|
||||
|
||||
Y yi = y[iOffsetY];
|
||||
Y yixj = y[ixjOffsetY];
|
||||
y[iOffsetY] = yixj;
|
||||
y[ixjOffsetY] = yi;
|
||||
}
|
||||
}
|
||||
else {
|
||||
// Sort descending
|
||||
if (!descending == (xi < xixj)) {
|
||||
x[iOffset] = xixj;
|
||||
x[ixjOffset] = xi;
|
||||
|
||||
sd::LongType iCoordsY[SD_MAX_RANK];
|
||||
sd::LongType ixjCoordsY[SD_MAX_RANK];
|
||||
sd::LongType iOffsetY;
|
||||
sd::LongType ixjOffsetY;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, iCoordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, iCoordsY, iOffsetY);
|
||||
|
||||
INDEX2COORDS(ixj, yRank, yShapePtr, ixjCoordsY);
|
||||
COORDS2INDEX(yRank, yStridePtr, ixjCoordsY, ixjOffsetY);
|
||||
|
||||
Y yi = y[iOffsetY];
|
||||
Y yixj = y[ixjOffsetY];
|
||||
y[iOffsetY] = yixj;
|
||||
y[ixjOffsetY] = yi;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void bitonicSortStepKernel(
|
||||
void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
int j,
|
||||
int k,
|
||||
int length,
|
||||
bool descending) {
|
||||
|
||||
auto x = static_cast<T*>(vx);
|
||||
const unsigned int i = threadIdx.x + blockDim.x * blockIdx.x;
|
||||
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
__shared__ sd::LongType xLength;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
xLength = shape::length(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (i >= static_cast<unsigned int>(length)) return;
|
||||
|
||||
const unsigned int ixj = i ^ j;
|
||||
if (ixj <= i) return;
|
||||
|
||||
sd::LongType iCoords[SD_MAX_RANK];
|
||||
sd::LongType ixjCoords[SD_MAX_RANK];
|
||||
sd::LongType iOffset;
|
||||
sd::LongType ixjOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, iCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, iCoords, iOffset);
|
||||
|
||||
INDEX2COORDS(ixj, xRank, xShapePtr, ixjCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, ixjCoords, ixjOffset);
|
||||
|
||||
const bool ascending = ((i & k) == 0);
|
||||
T xi = x[iOffset];
|
||||
T xixj = x[ixjOffset];
|
||||
|
||||
if (ascending) {
|
||||
// Sort ascending
|
||||
if (!descending == (xi > xixj)) {
|
||||
x[iOffset] = xixj;
|
||||
x[ixjOffset] = xi;
|
||||
}
|
||||
}
|
||||
else {
|
||||
// Sort descending
|
||||
if (!descending == (xi < xixj)) {
|
||||
x[iOffset] = xixj;
|
||||
x[ixjOffset] = xi;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void bitonicSortStepGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
int j,
|
||||
int k,
|
||||
int length,
|
||||
bool descending) {
|
||||
|
||||
bitonicSortStepKernel<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
xShapeInfo,
|
||||
j,
|
||||
k,
|
||||
length,
|
||||
descending);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "bitonicSortStepGeneric failed");
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void bitonicSortStepGenericKey(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vy,
|
||||
const sd::LongType* yShapeInfo,
|
||||
int j,
|
||||
int k,
|
||||
int length,
|
||||
bool descending) {
|
||||
|
||||
bitonicSortStepKernelKey<X, Y>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx,
|
||||
xShapeInfo,
|
||||
vy,
|
||||
yShapeInfo,
|
||||
j,
|
||||
k,
|
||||
length,
|
||||
descending);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "bitonicSortStepGenericKey failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template void bitonicSortStepGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
|
||||
int j, int k, int length, bool descending),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE(
|
||||
template void bitonicSortStepGenericKey,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
|
||||
void *vy, sd::LongType const *yShapeInfo, int j, int k, int length, bool descending),
|
||||
SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
@@ -0,0 +1,300 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_DEVICE void concatKernel(
|
||||
int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* resultShapeInfo,
|
||||
Pointer* tadPointers,
|
||||
Pointer* offsetPointers,
|
||||
LongType* zTadShape,
|
||||
LongType* zOffsets) {
|
||||
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// Shared variables for shape data
|
||||
__shared__ int zRank;
|
||||
__shared__ const LongType* zShape;
|
||||
__shared__ const LongType* zStride;
|
||||
__shared__ int zTadRank;
|
||||
__shared__ const LongType* zTadShapeOf;
|
||||
__shared__ const LongType* zTadStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zRank = shape::rank(resultShapeInfo);
|
||||
zShape = shape::shapeOf(resultShapeInfo);
|
||||
zStride = shape::stride(resultShapeInfo);
|
||||
zTadRank = shape::rank(zTadShape);
|
||||
zTadShapeOf = shape::shapeOf(zTadShape);
|
||||
zTadStride = shape::stride(zTadShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
auto result = reinterpret_cast<T*>(vz);
|
||||
auto dataT = reinterpret_cast<T**>(data);
|
||||
auto shapeInfoPtrs = reinterpret_cast<LongType**>(inputShapeInfos);
|
||||
auto tadShapes = reinterpret_cast<LongType**>(tadPointers);
|
||||
auto tadOffsets = reinterpret_cast<LongType**>(offsetPointers);
|
||||
|
||||
__shared__ bool _vec;
|
||||
__shared__ int baseIdx;
|
||||
__shared__ int arrOffset;
|
||||
__shared__ int yLength;
|
||||
__shared__ char yOrder;
|
||||
__shared__ int numTads;
|
||||
__shared__ char zOrder;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zOrder = shape::order(resultShapeInfo);
|
||||
_vec = shape::isVector(resultShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int zLength = shape::length(resultShapeInfo);
|
||||
|
||||
// Special case for when result is a vector
|
||||
if (_vec) {
|
||||
for (int r = blockIdx.x; r < numArrays; r += gridDim.x) {
|
||||
auto currShapeInfo = shapeInfoPtrs[r];
|
||||
if (shape::isVector(currShapeInfo) || shape::order(currShapeInfo) == zOrder) {
|
||||
if (threadIdx.x == 0) {
|
||||
yLength = shape::length(currShapeInfo);
|
||||
baseIdx = 0;
|
||||
for (int f = 0; f < r; f++) {
|
||||
baseIdx += shape::length(shapeInfoPtrs[f]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int i = threadIdx.x; i < yLength && baseIdx + i < zLength; i += blockDim.x) {
|
||||
result[baseIdx + i] = dataT[r][i];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// For all non-vector results
|
||||
for (int r = 0; r < numArrays; r++) {
|
||||
auto currShape = shapeInfoPtrs[r];
|
||||
auto currData = dataT[r];
|
||||
auto currTad = tadShapes[r];
|
||||
auto currOffsets = tadOffsets[r];
|
||||
|
||||
// Shared variables for current TAD shape data
|
||||
__shared__ int currTadRank;
|
||||
__shared__ const LongType* currTadShape;
|
||||
__shared__ const LongType* currTadStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
currTadRank = shape::rank(currTad);
|
||||
currTadShape = shape::shapeOf(currTad);
|
||||
currTadStride = shape::stride(currTad);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
yLength = shape::length(currTad);
|
||||
yOrder = shape::order(currTad);
|
||||
numTads = shape::length(currShape) / yLength;
|
||||
|
||||
arrOffset = 0;
|
||||
for (int f = 0; f < r; f++) {
|
||||
arrOffset += shape::length(tadShapes[f]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Case for single element
|
||||
if (yLength == 1 && _vec) {
|
||||
for (LongType j = tid; j < numTads; j += blockDim.x * gridDim.x) {
|
||||
LongType inputOffset = currOffsets[j];
|
||||
LongType resultOffset = zOffsets[j];
|
||||
|
||||
T* dataTAD = currData + inputOffset;
|
||||
T* resultTAD = result + resultOffset;
|
||||
|
||||
LongType sub[SD_MAX_RANK];
|
||||
INDEX2COORDS(arrOffset, zTadRank, zTadShapeOf, sub);
|
||||
|
||||
LongType baseOffset;
|
||||
COORDS2INDEX(zTadRank, zTadStride, sub, baseOffset);
|
||||
resultTAD += baseOffset;
|
||||
|
||||
INDEX2COORDS(0, currTadRank, currTadShape, sub);
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(currTadRank, currTadStride, sub, yOffset);
|
||||
|
||||
COORDS2INDEX(zTadRank, zTadStride, sub, resultOffset);
|
||||
resultTAD[resultOffset] = dataTAD[yOffset];
|
||||
}
|
||||
}
|
||||
else {
|
||||
for (LongType j = blockIdx.x; j < numTads; j += gridDim.x) {
|
||||
const LongType inputOffset = currOffsets[j];
|
||||
const LongType resultOffset = zOffsets[j];
|
||||
|
||||
T* dataTAD = currData + inputOffset;
|
||||
T* resultTAD = result + resultOffset;
|
||||
|
||||
LongType sub[SD_MAX_RANK];
|
||||
INDEX2COORDS(arrOffset, zTadRank, zTadShapeOf, sub);
|
||||
|
||||
LongType baseOffset;
|
||||
COORDS2INDEX(zTadRank, zTadStride, sub, baseOffset);
|
||||
resultTAD += baseOffset;
|
||||
|
||||
if (zOrder == yOrder) {
|
||||
for (int i = threadIdx.x; i < yLength; i += blockDim.x) {
|
||||
resultTAD[i] = dataTAD[i];
|
||||
}
|
||||
}
|
||||
else {
|
||||
if (shape::order(resultShapeInfo) == shape::order(currTad)) {
|
||||
if (threadIdx.x == 0) {
|
||||
baseIdx = 0;
|
||||
for (int f = 0; f < r; f++) {
|
||||
baseIdx += shape::length(shapeInfoPtrs[f]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (numTads == 1) {
|
||||
for (int k = threadIdx.x; k < yLength; k += blockDim.x) {
|
||||
resultTAD[baseIdx + k] = dataTAD[k];
|
||||
}
|
||||
}
|
||||
else {
|
||||
LongType yIdx[SD_MAX_RANK];
|
||||
for (LongType i = threadIdx.x; i < yLength; i += blockDim.x) {
|
||||
INDEX2COORDS(i, currTadRank, currTadShape, yIdx);
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(currTadRank, currTadStride, yIdx, yOffset);
|
||||
|
||||
resultTAD[baseIdx + i] = dataTAD[yOffset];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
else {
|
||||
LongType zIdx[SD_MAX_RANK];
|
||||
LongType yIdx[SD_MAX_RANK];
|
||||
|
||||
for (LongType i = threadIdx.x; i < yLength; i += blockDim.x) {
|
||||
INDEX2COORDS(i, currTadRank, currTadShape, yIdx);
|
||||
INDEX2COORDS(i, zTadRank, zTadShapeOf, zIdx);
|
||||
|
||||
LongType yOffset;
|
||||
COORDS2INDEX(currTadRank, currTadStride, yIdx, yOffset);
|
||||
|
||||
LongType rOffset;
|
||||
COORDS2INDEX(zTadRank, zTadStride, zIdx, rOffset);
|
||||
resultTAD[rOffset] = dataTAD[yOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execConcatKernel(
|
||||
int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* zShapeInfo,
|
||||
Pointer* tadPointers,
|
||||
Pointer* offsetPointers,
|
||||
LongType* zTadShape,
|
||||
LongType* zOffsets) {
|
||||
|
||||
concatKernel<T>(
|
||||
numArrays,
|
||||
data,
|
||||
inputShapeInfos,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
tadPointers,
|
||||
offsetPointers,
|
||||
zTadShape,
|
||||
zOffsets);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void concatKernelGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* zShapeInfo,
|
||||
Pointer* tadPointers,
|
||||
Pointer* offsetPointers,
|
||||
LongType* zTadShape,
|
||||
LongType* zOffsets) {
|
||||
|
||||
execConcatKernel<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
numArrays,
|
||||
data,
|
||||
inputShapeInfos,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
tadPointers,
|
||||
offsetPointers,
|
||||
zTadShape,
|
||||
zOffsets);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "concatGenericLegacy(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void concatKernelGeneric,
|
||||
(dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
int numArrays,
|
||||
sd::Pointer* data,
|
||||
sd::Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
sd::Pointer* tadPointers,
|
||||
sd::Pointer* offsetPointers,
|
||||
sd::LongType* zTadShape,
|
||||
sd::LongType* zOffsets),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,130 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
template <typename T>
|
||||
SD_DEVICE void concatKernelHStack(int numArrays, Pointer *data, Pointer *inputShapeInfos, void *vz,
|
||||
LongType *zShapeInfo) {
|
||||
// We expect each input array to be a vector, and the result (z) to be a 2D matrix for horizontal stacking.
|
||||
// The row dimension is presumably 1 for each input, with the column dimension the length of each vector,
|
||||
// stacked horizontally.
|
||||
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto inputData = reinterpret_cast<T**>(data);
|
||||
auto shapes = reinterpret_cast<LongType**>(inputShapeInfos);
|
||||
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
// Output shape data in shared memory
|
||||
__shared__ int zRank;
|
||||
__shared__ const LongType* zShape;
|
||||
__shared__ const LongType* zStride;
|
||||
|
||||
// Current input array shape data in shared memory
|
||||
__shared__ int inRank;
|
||||
__shared__ const LongType* inShape;
|
||||
__shared__ const LongType* inStride;
|
||||
|
||||
// Working variables in shared memory
|
||||
__shared__ int inputLength;
|
||||
__shared__ int baseIdx;
|
||||
|
||||
// Initialize output shape data once
|
||||
if (threadIdx.x == 0) {
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShape = shape::shapeOf(zShapeInfo);
|
||||
zStride = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Loop over all input arrays
|
||||
for (int r = blockIdx.x; r < numArrays; r += gridDim.x) {
|
||||
// Cache the current input array's shape data and compute offsets
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache current input shape data
|
||||
inRank = shape::rank(shapes[r]);
|
||||
inShape = shape::shapeOf(shapes[r]);
|
||||
inStride = shape::stride(shapes[r]);
|
||||
|
||||
// Compute base offset
|
||||
baseIdx = 0;
|
||||
for (int f = 0; f < r; f++) {
|
||||
baseIdx += shape::length(shapes[f]);
|
||||
}
|
||||
inputLength = shape::length(shapes[r]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each thread will copy a subset of data
|
||||
for (int i = tid; i < inputLength; i += blockDim.x * gridDim.x) {
|
||||
// Coordinates in the input vector
|
||||
LongType inCoords[SD_MAX_RANK];
|
||||
// Coordinates in the output 2D shape
|
||||
LongType outCoords[SD_MAX_RANK];
|
||||
|
||||
// 1) Get input coordinates using cached shape data
|
||||
INDEX2COORDS(i, inRank, inShape, inCoords);
|
||||
LongType inOffset;
|
||||
COORDS2INDEX(inRank, inStride, inCoords, inOffset);
|
||||
|
||||
// 2) The output coordinate index is baseIdx + i in the horizontal dimension
|
||||
const LongType outIndex = baseIdx + i;
|
||||
|
||||
// Get output coordinates using cached shape data
|
||||
INDEX2COORDS(outIndex, zRank, zShape, outCoords);
|
||||
LongType outOffset;
|
||||
COORDS2INDEX(zRank, zStride, outCoords, outOffset);
|
||||
|
||||
z[outOffset] = inputData[r][inOffset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execConcatKernelHStack(int numArrays, Pointer *data, Pointer *inputShapeInfos, void *vz,
|
||||
LongType *zShapeInfo) {
|
||||
concatKernelHStack<T>(numArrays, data, inputShapeInfos, vz, zShapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void concatKernelHStackGeneric(dim3 &launchDims, cudaStream_t *stream,
|
||||
int numArrays,
|
||||
Pointer *data,
|
||||
Pointer *inputShapeInfos,
|
||||
void *vz,
|
||||
LongType *zShapeInfo) {
|
||||
execConcatKernelHStack<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
numArrays, data, inputShapeInfos, vz, zShapeInfo);
|
||||
DebugHelper::checkErrorCode(stream, "concatHStack(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void concatKernelHStackGeneric,
|
||||
(dim3 &launchDims, cudaStream_t *stream, int numArrays, sd::Pointer *data,
|
||||
sd::Pointer *inputShapeInfos, void *vz, sd::LongType *zShapeInfo),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,55 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_DEVICE void concatKernelScalar(int numArrays, Pointer *data, void *vz) {
|
||||
auto z = static_cast<T *>(vz);
|
||||
LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
auto input = reinterpret_cast<T **>(data);
|
||||
|
||||
for (int i = tid; i < numArrays; i += blockDim.x * gridDim.x) z[i] = input[i][0];
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execConcatKernelScalar(int numArrays, Pointer *data, void *vz) {
|
||||
concatKernelScalar<T>(numArrays, data, vz);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void concatKernelScalarGeneric(dim3 &launchDims, cudaStream_t *stream, int numArrays, Pointer *data,
|
||||
void *vz) {
|
||||
execConcatKernelScalar<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(numArrays, data, vz);
|
||||
DebugHelper::checkErrorCode(stream, "concatScalar(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void concatKernelScalarGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, int numArrays, sd::Pointer *data, void *vz),
|
||||
SD_COMMON_TYPES);
|
||||
} // 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
|
||||
// @author Yurii Shyrma, created on 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_DEVICE void concatKernelVStack(int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* zShapeInfo) {
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto inputShapes = reinterpret_cast<LongType**>(inputShapeInfos);
|
||||
auto inputData = reinterpret_cast<T**>(data);
|
||||
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
__shared__ sd::LongType zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
// We'll store the rank/shape/stride of each input vector once per array
|
||||
// to avoid repeated calls inside the loop
|
||||
__shared__ sd::LongType inputRank;
|
||||
__shared__ const sd::LongType* inputShapePtr;
|
||||
__shared__ const sd::LongType* inputStridePtr;
|
||||
__shared__ sd::LongType rowLength; // length of each input vector
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// For each array, we assume it is a vector that will form one row of z
|
||||
for (int r = blockIdx.x; r < numArrays; r += gridDim.x) {
|
||||
// Single thread loads shape info for the current input array
|
||||
if (threadIdx.x == 0) {
|
||||
inputRank = shape::rank(inputShapes[r]);
|
||||
inputShapePtr = shape::shapeOf(inputShapes[r]);
|
||||
inputStridePtr = shape::stride(inputShapes[r]);
|
||||
rowLength = shape::length(inputShapes[r]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each thread copies part of the row for this array
|
||||
for (sd::LongType i = tid; i < rowLength; i += blockDim.x * gridDim.x) {
|
||||
// We'll do coordinate transforms to find the correct offsets:
|
||||
|
||||
// 1) Input offset
|
||||
sd::LongType inCoords[SD_MAX_RANK];
|
||||
INDEX2COORDS(i, inputRank, inputShapePtr, inCoords);
|
||||
|
||||
sd::LongType inOffset;
|
||||
COORDS2INDEX(inputRank, inputStridePtr, inCoords, inOffset);
|
||||
|
||||
// 2) Output offset
|
||||
// The "row" dimension is r, the "column" dimension is i.
|
||||
sd::LongType outCoords[SD_MAX_RANK];
|
||||
outCoords[0] = r; // row
|
||||
outCoords[1] = i; // column
|
||||
|
||||
sd::LongType outOffset;
|
||||
COORDS2INDEX(zRank, zStridePtr, outCoords, outOffset);
|
||||
|
||||
z[outOffset] = inputData[r][inOffset];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execConcatKernelVStack(
|
||||
int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* zShapeInfo) {
|
||||
|
||||
concatKernelVStack<T>(
|
||||
numArrays,
|
||||
data,
|
||||
inputShapeInfos,
|
||||
vz,
|
||||
zShapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void concatKernelVStackGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
int numArrays,
|
||||
Pointer* data,
|
||||
Pointer* inputShapeInfos,
|
||||
void* vz,
|
||||
LongType* zShapeInfo) {
|
||||
|
||||
execConcatKernelVStack<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
numArrays,
|
||||
data,
|
||||
inputShapeInfos,
|
||||
vz,
|
||||
zShapeInfo);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "concatVStack(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void concatKernelVStackGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, int numArrays, sd::Pointer *data,
|
||||
sd::Pointer *inputShapeInfos, void *vz, sd::LongType *zShapeInfo),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,46 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execConvertHalfs(half *dx, LongType n, void *dz) {
|
||||
auto z = reinterpret_cast<T *>(dz);
|
||||
int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
for (LongType i = tid; i < n; i += blockDim.x * gridDim.x) z[i] = static_cast<T>(__half2float(dx[i]));
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void convertHalfsToGeneric(dim3 &launchDims, cudaStream_t *stream, half *dx, LongType n, void *dz) {
|
||||
execConvertHalfs<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, n, dz);
|
||||
DebugHelper::checkErrorCode(stream, "convertHalfsToGeneric(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void convertHalfsToGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, half *dx, sd::LongType n, void *dz), SD_COMMON_TYPES);
|
||||
} // 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 raver119@gmail.com
|
||||
// @author Yurii Shyrma, created on 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
|
||||
namespace sd {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execConvertToHalf(void *dx, LongType n, half *dz) {
|
||||
auto x = reinterpret_cast<T *>(dx);
|
||||
int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
|
||||
for (LongType i = tid; i < n; i += blockDim.x * gridDim.x) dz[i] = __float2half(static_cast<T>(x[i]));
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void convertToHalfGeneric(dim3 &launchDims, cudaStream_t *stream, void *dx, LongType n, half *dz) {
|
||||
execConvertToHalf<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, n, dz);
|
||||
DebugHelper::checkErrorCode(stream, "convertToHalfs(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void convertToHalfGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *dx, sd::LongType n, half *dz), SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,131 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_DEVICE void fillDimensionalIsMax(
|
||||
const void* vdX,
|
||||
void* vdZ,
|
||||
const LongType* zShapeInfo,
|
||||
const LongType* tadOnlyShapeInfo,
|
||||
LongType* dimension,
|
||||
LongType dimensionLength,
|
||||
const LongType* tadOffsets) {
|
||||
|
||||
const auto dX = reinterpret_cast<const LongType*>(vdX);
|
||||
auto dZ = reinterpret_cast<T*>(vdZ);
|
||||
|
||||
__shared__ int tadLen;
|
||||
__shared__ int numTads;
|
||||
__shared__ int tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLen = static_cast<int>(shape::length(tadOnlyShapeInfo));
|
||||
numTads = static_cast<int>(shape::length(zShapeInfo) / tadLen);
|
||||
|
||||
tadRank = shape::rank(tadOnlyShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadOnlyShapeInfo);
|
||||
tadStridePtr= shape::stride(tadOnlyShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// each block handles some portion of the TADs
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
const auto tadOffsetForBlock = tadOffsets[r];
|
||||
const auto highestElement = dX[r]; // this is presumably the index in [0..tadLen)
|
||||
|
||||
// Each thread does part of the tad's length
|
||||
// dimensionLength determines if we have multiple dims, but code is the same
|
||||
for (LongType e = threadIdx.x; e < tadLen; e += blockDim.x) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType offset;
|
||||
|
||||
INDEX2COORDS(e, tadRank, tadShapePtr, coords);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, coords, offset);
|
||||
|
||||
const auto finalOffset = tadOffsetForBlock + offset;
|
||||
dZ[finalOffset] = (e == highestElement ? static_cast<T>(1) : static_cast<T>(0));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execfillDimensionalIsMax(
|
||||
const void* dX,
|
||||
void* dZ,
|
||||
const LongType* zShapeInfo,
|
||||
const LongType* tadOnlyShapeInfo,
|
||||
LongType* dimension,
|
||||
LongType dimensionLength,
|
||||
const LongType* tadOffsets) {
|
||||
|
||||
fillDimensionalIsMax<T>(
|
||||
dX, dZ, zShapeInfo, tadOnlyShapeInfo,
|
||||
dimension, dimensionLength, tadOffsets);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void fillDimensionalIsMaxGeneric(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* dX,
|
||||
void* dZ,
|
||||
const LongType* zShapeInfo,
|
||||
const LongType* tadOnlyShapeInfo,
|
||||
LongType* dimension,
|
||||
LongType dimensionLength,
|
||||
const LongType* tadOffsets) {
|
||||
|
||||
execfillDimensionalIsMax<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
dX, dZ, zShapeInfo, tadOnlyShapeInfo,
|
||||
dimension, dimensionLength, tadOffsets);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "fillDimensionalIsMax(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template void fillDimensionalIsMaxGeneric,
|
||||
(dim3 & launchDims,
|
||||
cudaStream_t *stream,
|
||||
const void* dX,
|
||||
void* dZ,
|
||||
const sd::LongType* zShapeInfo,
|
||||
const sd::LongType* tadOnlyShapeInfo,
|
||||
sd::LongType* dimension,
|
||||
sd::LongType dimensionLength,
|
||||
const sd::LongType* tadOffsets),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -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 raver119@gmail.com
|
||||
// @author Yurii Shyrma, created on 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execFillIsMax(
|
||||
void* vdZ,
|
||||
const LongType* xShapeInfo,
|
||||
LongType length,
|
||||
long idx) {
|
||||
|
||||
auto dz = reinterpret_cast<T*>(vdZ);
|
||||
|
||||
__shared__ int rank;
|
||||
__shared__ const sd::LongType* shapePtr;
|
||||
__shared__ const sd::LongType* stridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(xShapeInfo);
|
||||
shapePtr = shape::shapeOf(xShapeInfo);
|
||||
stridePtr = shape::stride(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int totalThreads = blockDim.x * gridDim.x;
|
||||
|
||||
for (LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType offset;
|
||||
|
||||
INDEX2COORDS(i, rank, shapePtr, coords);
|
||||
COORDS2INDEX(rank, stridePtr, coords, offset);
|
||||
|
||||
dz[offset] = (i == idx ? static_cast<T>(1) : static_cast<T>(0));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void fillIsMaxGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
void* dz,
|
||||
const LongType* zShapeInfo,
|
||||
LongType length,
|
||||
long idx) {
|
||||
|
||||
execFillIsMax<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
dz,
|
||||
zShapeInfo,
|
||||
length,
|
||||
idx);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "fillIsMax(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template void fillIsMaxGeneric,
|
||||
(dim3 & launchDims,
|
||||
cudaStream_t *stream,
|
||||
void *dz,
|
||||
const sd::LongType *zShapeInfo,
|
||||
sd::LongType length,
|
||||
long idx),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // 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 raver119@gmail.com
|
||||
// @author Yurii Shyrma, created on 27.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
#include <ops/declarable/helpers/flatten.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void flattenKernel(
|
||||
Pointer* extraPointers,
|
||||
int dOffset,
|
||||
char order,
|
||||
void* vz,
|
||||
LongType* zShapeInfo,
|
||||
void* vy,
|
||||
LongType* yShapeInfo) {
|
||||
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
auto y = reinterpret_cast<T*>(vy);
|
||||
|
||||
__shared__ sd::LongType yLen;
|
||||
__shared__ int yRank;
|
||||
__shared__ const sd::LongType* yShapePtr;
|
||||
__shared__ const sd::LongType* yStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
yLen = shape::length(yShapeInfo);
|
||||
yRank = shape::rank(yShapeInfo);
|
||||
yShapePtr = shape::shapeOf(yShapeInfo);
|
||||
yStridePtr = shape::stride(yShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (auto i = tid; i < yLen; i += totalThreads) {
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType yOffset;
|
||||
|
||||
INDEX2COORDS(i, yRank, yShapePtr, coords);
|
||||
COORDS2INDEX(yRank, yStridePtr, coords, yOffset);
|
||||
|
||||
z[i + dOffset] = y[yOffset];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void flattenKernelGeneric(
|
||||
dim3& launchDims,
|
||||
cudaStream_t* stream,
|
||||
Pointer* extraPointers,
|
||||
int dOffset,
|
||||
char order,
|
||||
void* vz,
|
||||
LongType* zShapeInfo,
|
||||
void* vy,
|
||||
LongType* yShapeInfo) {
|
||||
|
||||
flattenKernel<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
extraPointers,
|
||||
dOffset,
|
||||
order,
|
||||
vz,
|
||||
zShapeInfo,
|
||||
vy,
|
||||
yShapeInfo);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "flattenGeneric(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void flattenKernelGeneric,
|
||||
(dim3 & launchDims,
|
||||
cudaStream_t* stream,
|
||||
sd::Pointer* extraPointers,
|
||||
int dOffset,
|
||||
char order,
|
||||
void* vz,
|
||||
sd::LongType* zShapeInfo,
|
||||
void* vy,
|
||||
sd::LongType* yShapeInfo),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,250 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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/specials_cuda.h>
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL void execOesTadKernelKey(void *vx, sd::LongType const *xShapeInfo, void *vy, sd::LongType const *yShapeInfo,
|
||||
sd::LongType *dimension, long long int dimensionLength, sd::LongType const *tadShapeInfo,
|
||||
sd::LongType const *tadOffsets, bool descending) {
|
||||
auto x = static_cast<X *>(vx);
|
||||
auto y = static_cast<Y *>(vy);
|
||||
|
||||
__shared__ int xLength;
|
||||
__shared__ int xTadLength;
|
||||
__shared__ int numTads;
|
||||
__shared__ int tadRank;
|
||||
__shared__ sd::LongType *tadShape;
|
||||
__shared__ sd::LongType *tadStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLength = shape::length(xShapeInfo);
|
||||
xTadLength = shape::length(tadShapeInfo);
|
||||
numTads = xLength / xTadLength;
|
||||
|
||||
// Cache shape information
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShape = shape::shapeOf(tadShapeInfo);
|
||||
tadStride = shape::stride(tadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto dx = x + tadOffsets[r];
|
||||
auto dy = y + tadOffsets[r];
|
||||
|
||||
// this is general loop, we go uncached
|
||||
int iterations = xTadLength;
|
||||
|
||||
for (int i = 0; i < iterations; i++) {
|
||||
if (i % 2 == 0) {
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
auto top = 2 * tid + 1;
|
||||
if (top < xTadLength) {
|
||||
sd::LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
||||
sd::LongType t0Offset, t1Offset;
|
||||
|
||||
INDEX2COORDS(top - 1, tadRank, tadShape, t0Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t0Coords, t0Offset);
|
||||
INDEX2COORDS(top, tadRank, tadShape, t1Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t1Coords, t1Offset);
|
||||
|
||||
if (!descending == (dx[t0Offset] > dx[t1Offset])) {
|
||||
X dt0 = dx[t0Offset];
|
||||
dx[t0Offset] = dx[t1Offset];
|
||||
dx[t1Offset] = dt0;
|
||||
|
||||
Y dy0 = dy[t0Offset];
|
||||
dy[t0Offset] = dy[t1Offset];
|
||||
dy[t1Offset] = dy0;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
auto top = 2 * tid + 2;
|
||||
if (top < xTadLength) {
|
||||
sd::LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
||||
sd::LongType t0Offset, t1Offset;
|
||||
|
||||
INDEX2COORDS(top - 1, tadRank, tadShape, t0Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t0Coords, t0Offset);
|
||||
INDEX2COORDS(top, tadRank, tadShape, t1Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t1Coords, t1Offset);
|
||||
|
||||
if (!descending == (dx[t0Offset] > dx[t1Offset])) {
|
||||
X dt0 = dx[t0Offset];
|
||||
dx[t0Offset] = dx[t1Offset];
|
||||
dx[t1Offset] = dt0;
|
||||
|
||||
Y dy0 = dy[t0Offset];
|
||||
dy[t0Offset] = dy[t1Offset];
|
||||
dy[t1Offset] = dy0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL void execOesTadKernel(void *vx, sd::LongType const *xShapeInfo, sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const *tadShapeInfo, sd::LongType const *tadOffsets, bool descending) {
|
||||
auto x = static_cast<T *>(vx);
|
||||
const int sharedSize = 32768;
|
||||
|
||||
__shared__ int xLength;
|
||||
__shared__ int xTadLength;
|
||||
__shared__ int numTads;
|
||||
__shared__ T *shmem;
|
||||
__shared__ bool cached;
|
||||
__shared__ int tadRank;
|
||||
__shared__ sd::LongType *tadShape;
|
||||
__shared__ sd::LongType *tadStride;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
xLength = shape::length(xShapeInfo);
|
||||
xTadLength = shape::length(tadShapeInfo);
|
||||
numTads = xLength / xTadLength;
|
||||
|
||||
extern __shared__ unsigned char shrd[];
|
||||
shmem = (T *)shrd;
|
||||
|
||||
cached = xTadLength <= (sharedSize / sizeof(T));
|
||||
|
||||
// Cache shape information
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShape = shape::shapeOf(tadShapeInfo);
|
||||
tadStride = shape::stride(tadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto dx = x + tadOffsets[r];
|
||||
|
||||
// this is general loop, we go uncached
|
||||
int iterations = xTadLength;
|
||||
if (cached) {
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(tid, tadRank, tadShape, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStride, xCoords, xOffset);
|
||||
shmem[tid] = dx[xOffset];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
dx = shmem;
|
||||
}
|
||||
|
||||
for (int i = 0; i < iterations; i++) {
|
||||
if (i % 2 == 0) {
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
auto top = 2 * tid + 1;
|
||||
if (top < xTadLength) {
|
||||
sd::LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
||||
sd::LongType t0Offset, t1Offset;
|
||||
|
||||
INDEX2COORDS(top - 1, tadRank, tadShape, t0Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t0Coords, t0Offset);
|
||||
INDEX2COORDS(top, tadRank, tadShape, t1Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t1Coords, t1Offset);
|
||||
|
||||
if (!descending == (dx[t0Offset] > dx[t1Offset])) {
|
||||
T dt0 = dx[t0Offset];
|
||||
dx[t0Offset] = dx[t1Offset];
|
||||
dx[t1Offset] = dt0;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
auto top = 2 * tid + 2;
|
||||
if (top < xTadLength) {
|
||||
sd::LongType t0Coords[SD_MAX_RANK], t1Coords[SD_MAX_RANK];
|
||||
sd::LongType t0Offset, t1Offset;
|
||||
|
||||
INDEX2COORDS(top - 1, tadRank, tadShape, t0Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t0Coords, t0Offset);
|
||||
INDEX2COORDS(top, tadRank, tadShape, t1Coords);
|
||||
COORDS2INDEX(tadRank, tadStride, t1Coords, t1Offset);
|
||||
|
||||
if (!descending == (dx[t0Offset] > dx[t1Offset])) {
|
||||
T dt0 = dx[t0Offset];
|
||||
dx[t0Offset] = dx[t1Offset];
|
||||
dx[t1Offset] = dt0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (cached) {
|
||||
dx = x + tadOffsets[r];
|
||||
for (int tid = threadIdx.x; tid < xTadLength; tid += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(tid, tadRank, tadShape, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStride, xCoords, xOffset);
|
||||
dx[xOffset] = shmem[tid];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void oesTadGeneric(dim3 &launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
|
||||
sd::LongType *dimension, sd::LongType dimensionLength, sd::LongType const *tadShapeInfo,
|
||||
sd::LongType const *tadOffsets, bool descending) {
|
||||
execOesTadKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(vx, xShapeInfo, dimension, dimensionLength,
|
||||
tadShapeInfo, tadOffsets, descending);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "execOesTadKernel failed");
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void oesTadGenericKey(dim3 &launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
|
||||
void *vy, sd::LongType const *yShapeInfo, sd::LongType *dimension,
|
||||
sd::LongType dimensionLength,
|
||||
sd::LongType const *tadShapeInfo, sd::LongType const *tadOffsets, bool descending) {
|
||||
execOesTadKernelKey<X, Y><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(
|
||||
vx, xShapeInfo, vy, yShapeInfo, dimension, dimensionLength, tadShapeInfo, tadOffsets, descending);
|
||||
sd::DebugHelper::checkErrorCode(stream, "execOesTadKernelKey failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( void oesTadGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo,
|
||||
sd::LongType *dimension, sd::LongType dimensionLength, sd::LongType const *tadShapeInfo,
|
||||
sd::LongType const *tadOffsets, bool descending),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( void oesTadGenericKey,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *vx, sd::LongType const *xShapeInfo, void *vy,
|
||||
sd::LongType const *yShapeInfo, sd::LongType *dimension, sd::LongType dimensionLength,
|
||||
sd::LongType const *tadShapeInfo, sd::LongType const *tadOffsets, bool descending),
|
||||
SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
@@ -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
|
||||
// @author Yurii Shyrma, created on 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_DEVICE void pullRowsKernel(
|
||||
void* vx,
|
||||
void* vz,
|
||||
LongType len,
|
||||
LongType* indexes,
|
||||
LongType* tadShapeInfo,
|
||||
LongType* tadOffsets,
|
||||
LongType* zTadShapeInfo,
|
||||
LongType* zTadOffsets) {
|
||||
|
||||
auto x = reinterpret_cast<T*>(vx);
|
||||
auto z = reinterpret_cast<T*>(vz);
|
||||
|
||||
__shared__ int tadRank;
|
||||
__shared__ const sd::LongType* tadShapePtr;
|
||||
__shared__ const sd::LongType* tadStridePtr;
|
||||
__shared__ int zTadRank;
|
||||
__shared__ const sd::LongType* zTadShapePtr;
|
||||
__shared__ const sd::LongType* zTadStridePtr;
|
||||
__shared__ sd::LongType tadLen;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadRank = shape::rank(tadShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadShapeInfo);
|
||||
tadStridePtr = shape::stride(tadShapeInfo);
|
||||
|
||||
zTadRank = shape::rank(zTadShapeInfo);
|
||||
zTadShapePtr = shape::shapeOf(zTadShapeInfo);
|
||||
zTadStridePtr = shape::stride(zTadShapeInfo);
|
||||
|
||||
tadLen = shape::length(tadShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each block handles some subset of the total 'len' (the number of TAD pulls)
|
||||
for (sd::LongType idx = blockIdx.x; idx < len; idx += gridDim.x) {
|
||||
const auto xTadOffset = tadOffsets[indexes[idx]];
|
||||
const auto zTadOffset = zTadOffsets[idx];
|
||||
|
||||
auto rX = x + xTadOffset;
|
||||
auto rZ = z + zTadOffset;
|
||||
|
||||
for (sd::LongType i = threadIdx.x; i < tadLen; i += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, tadRank, tadShapePtr, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, zTadRank, zTadShapePtr, zCoords);
|
||||
COORDS2INDEX(zTadRank, zTadStridePtr, zCoords, zOffset);
|
||||
|
||||
rZ[zOffset] = rX[xOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void execPullRowsKernel(
|
||||
void* vx,
|
||||
void* vz,
|
||||
LongType len,
|
||||
LongType* indexes,
|
||||
LongType* tadShapeInfo,
|
||||
LongType* tadOffsets,
|
||||
LongType* zTadShapeInfo,
|
||||
LongType* zTadOffsets) {
|
||||
|
||||
pullRowsKernel<T>(
|
||||
vx, vz, len, indexes, tadShapeInfo, tadOffsets,
|
||||
zTadShapeInfo, zTadOffsets);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
SD_HOST void pullRowsKernelGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t* stream,
|
||||
void* vx,
|
||||
void* vz,
|
||||
LongType len,
|
||||
LongType* indexes,
|
||||
LongType* tadShapeInfo,
|
||||
LongType* tadOffsets,
|
||||
LongType* zTadShapeInfo,
|
||||
LongType* zTadOffsets) {
|
||||
|
||||
execPullRowsKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vx, vz, len, indexes, tadShapeInfo, tadOffsets,
|
||||
zTadShapeInfo, zTadOffsets);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "pullRows(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template void pullRowsKernelGeneric,
|
||||
(dim3 &launchDims,
|
||||
cudaStream_t* stream,
|
||||
void* vx,
|
||||
void* vz,
|
||||
sd::LongType len,
|
||||
sd::LongType* indexes,
|
||||
sd::LongType* tadShapeInfo,
|
||||
sd::LongType* tadOffsets,
|
||||
sd::LongType* zTadShapeInfo,
|
||||
sd::LongType* zTadOffsets),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,187 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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>, created on 21.01.2019
|
||||
//
|
||||
#include <array/NDArray.h>
|
||||
#include <loops/special_kernels.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// set up a given value above (or on) a specified diagonal in a 2D matrix
|
||||
template <typename T>
|
||||
static SD_KERNEL void setDiagValueUpperKernel(
|
||||
void* buffer,
|
||||
LongType* shape,
|
||||
T value,
|
||||
int diagonal,
|
||||
LongType rows,
|
||||
LongType cols) {
|
||||
|
||||
__shared__ LongType rank;
|
||||
__shared__ const sd::LongType* stridePtr;
|
||||
__shared__ T* arr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(shape);
|
||||
stridePtr = shape::stride(shape);
|
||||
arr = reinterpret_cast<T*>(buffer);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType r = blockIdx.x; r < rows; r += gridDim.x) {
|
||||
for (LongType c = threadIdx.x; c < cols; c += blockDim.x) {
|
||||
sd::LongType coords[2] = {r, c};
|
||||
sd::LongType offset;
|
||||
|
||||
COORDS2INDEX(rank, stridePtr, coords, offset);
|
||||
|
||||
// If c >= r + diagonal
|
||||
// means c - r >= diagonal
|
||||
// i.e. we are on/above the diagonal
|
||||
if (r + diagonal <= c) {
|
||||
arr[offset] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
// set up a given value below (or on) a specified diagonal in a 2D matrix
|
||||
template <typename T>
|
||||
static SD_KERNEL void setDiagValueLowerKernel(
|
||||
void* buffer,
|
||||
LongType* shape,
|
||||
T value,
|
||||
int diagonal,
|
||||
LongType rows,
|
||||
LongType cols) {
|
||||
|
||||
__shared__ LongType rank;
|
||||
__shared__ const sd::LongType* stridePtr;
|
||||
__shared__ T* arr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
rank = shape::rank(shape);
|
||||
stridePtr = shape::stride(shape);
|
||||
arr = reinterpret_cast<T*>(buffer);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (LongType r = blockIdx.x; r < rows; r += gridDim.x) {
|
||||
for (LongType c = threadIdx.x; c < cols; c += blockDim.x) {
|
||||
sd::LongType coords[2] = {r, c};
|
||||
sd::LongType offset;
|
||||
|
||||
COORDS2INDEX(rank, stridePtr, coords, offset);
|
||||
|
||||
// If c <= r + diagonal
|
||||
// means c - r <= diagonal
|
||||
// i.e. we are on/below the diagonal
|
||||
if (r + diagonal >= c) {
|
||||
arr[offset] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Below are template instantiations for various data types
|
||||
template SD_KERNEL void setDiagValueUpperKernel<double>(
|
||||
void* buffer, LongType* shape, double value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<double>(
|
||||
void* buffer, LongType* shape, double value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<float>(
|
||||
void* buffer, LongType* shape, float value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<float>(
|
||||
void* buffer, LongType* shape, float value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<int>(
|
||||
void* buffer, LongType* shape, int value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<int>(
|
||||
void* buffer, LongType* shape, int value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<float16>(
|
||||
void* buffer, LongType* shape, float16 value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<float16>(
|
||||
void* buffer, LongType* shape, float16 value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<bfloat16>(
|
||||
void* buffer, LongType* shape, bfloat16 value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<bfloat16>(
|
||||
void* buffer, LongType* shape, bfloat16 value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<LongType>(
|
||||
void* buffer, LongType* shape, LongType value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<LongType>(
|
||||
void* buffer, LongType* shape, LongType value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<int16_t>(
|
||||
void* buffer, LongType* shape, int16_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<int16_t>(
|
||||
void* buffer, LongType* shape, int16_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<uint8_t>(
|
||||
void* buffer, LongType* shape, uint8_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<uint8_t>(
|
||||
void* buffer, LongType* shape, uint8_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<int8_t>(
|
||||
void* buffer, LongType* shape, int8_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<int8_t>(
|
||||
void* buffer, LongType* shape, int8_t value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueUpperKernel<bool>(
|
||||
void* buffer, LongType* shape, bool value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
template SD_KERNEL void setDiagValueLowerKernel<bool>(
|
||||
void* buffer, LongType* shape, bool value,
|
||||
int diagonal, LongType rows, LongType cols);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,182 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 15.11.2018
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
|
||||
namespace sd {
|
||||
template <typename T>
|
||||
SD_KERNEL void execShuffleKernel(
|
||||
void** vdX,
|
||||
LongType** dxShapeInfo,
|
||||
void** vdZ,
|
||||
int N,
|
||||
int* shuffleMap,
|
||||
LongType** tadOnlyShapeInfo,
|
||||
LongType** tadOffsets) {
|
||||
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Main array shape data
|
||||
__shared__ int xRank;
|
||||
__shared__ LongType xLength;
|
||||
__shared__ const LongType* xShape;
|
||||
__shared__ const LongType* xStride;
|
||||
|
||||
// TAD shape data
|
||||
__shared__ int tadRank;
|
||||
__shared__ int tadLength;
|
||||
__shared__ int numTads;
|
||||
__shared__ const LongType* tadShape;
|
||||
__shared__ const LongType* tadStride;
|
||||
|
||||
// Current shape info pointer
|
||||
__shared__ LongType* xShapeInfo;
|
||||
|
||||
// Process each array
|
||||
for (int arrIndex = 0; arrIndex < N; arrIndex++) {
|
||||
auto x = reinterpret_cast<T*>(vdX[arrIndex]);
|
||||
auto z = reinterpret_cast<T*>(vdZ[arrIndex]);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// Cache main array shape data
|
||||
xShapeInfo = dxShapeInfo[arrIndex];
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xLength = shape::length(xShapeInfo);
|
||||
xShape = shape::shapeOf(xShapeInfo);
|
||||
xStride = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache TAD shape data
|
||||
tadLength = static_cast<int>(shape::length(tadOnlyShapeInfo[arrIndex]));
|
||||
tadRank = shape::rank(tadOnlyShapeInfo[arrIndex]);
|
||||
tadShape = shape::shapeOf(tadOnlyShapeInfo[arrIndex]);
|
||||
tadStride = shape::stride(tadOnlyShapeInfo[arrIndex]);
|
||||
numTads = static_cast<int>(xLength / tadLength);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Rank-1 case: treat as vector
|
||||
if (xRank == 1) {
|
||||
for (LongType elem = tid; elem < xLength; elem += totalThreads) {
|
||||
const int swapIndex = shuffleMap[elem];
|
||||
if (swapIndex >= 0 && swapIndex < xLength) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType swapCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType swapOffset;
|
||||
|
||||
// Use cached shape data for coordinate transforms
|
||||
INDEX2COORDS(elem, xRank, xShape, xCoords);
|
||||
COORDS2INDEX(xRank, xStride, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(swapIndex, xRank, xShape, swapCoords);
|
||||
COORDS2INDEX(xRank, xStride, swapCoords, swapOffset);
|
||||
|
||||
// Swap values
|
||||
T oldVal = x[xOffset];
|
||||
x[xOffset] = x[swapOffset];
|
||||
x[swapOffset] = oldVal;
|
||||
}
|
||||
}
|
||||
}
|
||||
else {
|
||||
// TAD-based processing
|
||||
for (LongType r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
const int swapTarget = shuffleMap[r];
|
||||
if (swapTarget >= 0) {
|
||||
const auto oldOffset = tadOffsets[arrIndex][r];
|
||||
const auto newOffset = tadOffsets[arrIndex][swapTarget];
|
||||
|
||||
// Pointers to TADs
|
||||
auto rX = x + oldOffset;
|
||||
auto rY = x + newOffset;
|
||||
auto zX = z + oldOffset;
|
||||
auto zY = z + newOffset;
|
||||
|
||||
// Process elements within TAD
|
||||
for (LongType i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType yCoords[SD_MAX_RANK];
|
||||
sd::LongType yOffset;
|
||||
|
||||
// Use cached TAD shape data for coordinate transforms
|
||||
INDEX2COORDS(i, tadRank, tadShape, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStride, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, tadRank, tadShape, yCoords);
|
||||
COORDS2INDEX(tadRank, tadStride, yCoords, yOffset);
|
||||
|
||||
// Add TAD base offsets
|
||||
xOffset += oldOffset;
|
||||
yOffset += newOffset;
|
||||
|
||||
// Perform swap via z array
|
||||
T oldVal = x[xOffset];
|
||||
z[xOffset] = x[yOffset];
|
||||
z[yOffset] = oldVal;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
template <typename T>
|
||||
SD_HOST void shuffleKernelGeneric(
|
||||
dim3 &launchDims,
|
||||
cudaStream_t *stream,
|
||||
void** vdX,
|
||||
LongType** xShapeInfo,
|
||||
void** vdZ,
|
||||
int N,
|
||||
int* shuffleMap,
|
||||
LongType** tadOnlyShapeInfo,
|
||||
LongType** tadOffsets) {
|
||||
|
||||
execShuffleKernel<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
vdX,
|
||||
xShapeInfo,
|
||||
vdZ,
|
||||
N,
|
||||
shuffleMap,
|
||||
tadOnlyShapeInfo,
|
||||
tadOffsets);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "shuffleGeneric(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
template void shuffleKernelGeneric,
|
||||
(dim3 & launchDims,
|
||||
cudaStream_t *stream,
|
||||
void** vdX,
|
||||
sd::LongType** xShapeInfo,
|
||||
void** vdZ,
|
||||
int N,
|
||||
int* shuffleMap,
|
||||
sd::LongType** tadOnlyShapeInfo,
|
||||
sd::LongType** tadOffsets),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // 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 GS <sgazeos@gmail.com>, created on 25.01.2019
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void swapUnsafeKernel(
|
||||
void* theFirstBuffer,
|
||||
LongType const* theFirstShape,
|
||||
void* theSecondBuffer,
|
||||
LongType const* theSecondShape) {
|
||||
|
||||
// thread & grid info
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// cache relevant info in shared memory
|
||||
__shared__ bool sameOffsets, sameOrders;
|
||||
__shared__ sd::LongType resultLength;
|
||||
|
||||
__shared__ sd::LongType firstRank;
|
||||
__shared__ const sd::LongType* firstShapePtr;
|
||||
__shared__ const sd::LongType* firstStridePtr;
|
||||
|
||||
__shared__ sd::LongType secondRank;
|
||||
__shared__ const sd::LongType* secondShapePtr;
|
||||
__shared__ const sd::LongType* secondStridePtr;
|
||||
|
||||
__shared__ T* outPtr;
|
||||
__shared__ T* inPtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
resultLength = shape::length(theFirstShape);
|
||||
|
||||
// first shape/stride info
|
||||
firstRank = shape::rank(theFirstShape);
|
||||
firstShapePtr = shape::shapeOf(theFirstShape);
|
||||
firstStridePtr = shape::stride(theFirstShape);
|
||||
|
||||
// second shape/stride info
|
||||
secondRank = shape::rank(theSecondShape);
|
||||
secondShapePtr = shape::shapeOf(theSecondShape);
|
||||
secondStridePtr = shape::stride(theSecondShape);
|
||||
|
||||
outPtr = reinterpret_cast<T*>(theFirstBuffer);
|
||||
inPtr = reinterpret_cast<T*>(theSecondBuffer);
|
||||
|
||||
sameOffsets = shape::haveSameShapeAndStrides(theFirstShape, theSecondShape);
|
||||
sameOrders = (shape::order(theFirstShape) == shape::order(theSecondShape));
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType i = tid; i < resultLength; i += totalThreads) {
|
||||
sd::LongType firstCoords[SD_MAX_RANK];
|
||||
sd::LongType secondCoords[SD_MAX_RANK];
|
||||
|
||||
// offsets in each array
|
||||
sd::LongType firstOffset;
|
||||
sd::LongType secondOffset;
|
||||
|
||||
// compute coordinates in the first array
|
||||
INDEX2COORDS(i, firstRank, firstShapePtr, firstCoords);
|
||||
COORDS2INDEX(firstRank, firstStridePtr, firstCoords, firstOffset);
|
||||
|
||||
// compute coordinates in the second array
|
||||
INDEX2COORDS(i, secondRank, secondShapePtr, secondCoords);
|
||||
COORDS2INDEX(secondRank, secondStridePtr, secondCoords, secondOffset);
|
||||
|
||||
if (sameOrders && firstOffset >= 0 && secondOffset >= 0) {
|
||||
// direct swap with the known offsets
|
||||
math::sd_swap(outPtr[firstOffset], inPtr[secondOffset]);
|
||||
}
|
||||
else if (sameOffsets) {
|
||||
// same shape/strides => same offset for both
|
||||
math::sd_swap(outPtr[firstOffset], inPtr[firstOffset]);
|
||||
}
|
||||
else {
|
||||
math::sd_swap(outPtr[firstOffset], inPtr[secondOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
SD_KERNEL void swapUnsafeKernel,
|
||||
(void* theFirstBuffer,
|
||||
sd::LongType const* theFirstShape,
|
||||
void* theSecondBuffer,
|
||||
sd::LongType const* theSecondShape),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
template <typename T>
|
||||
void templatedSwapUnsafe(
|
||||
void* theFirstBuffer,
|
||||
LongType const* theFirstShape,
|
||||
void* theSecondBuffer,
|
||||
LongType const* theSecondShape,
|
||||
cudaStream_t* theStream) {
|
||||
|
||||
dim3 launchDims = getLaunchDims("swap_unsafe");
|
||||
|
||||
swapUnsafeKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *theStream>>>(
|
||||
theFirstBuffer,
|
||||
theFirstShape,
|
||||
theSecondBuffer,
|
||||
theSecondShape);
|
||||
|
||||
DebugHelper::checkGlobalErrorCode("templatedSwapUnsafe(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void templatedSwapUnsafe,
|
||||
(void* theFirstBuffer,
|
||||
sd::LongType const* theFirstShape,
|
||||
void* theSecondBuffer,
|
||||
sd::LongType const* theSecondShape,
|
||||
cudaStream_t* theStream),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,251 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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>, created on 16.01.2019
|
||||
//
|
||||
#include <loops/special_kernels.h>
|
||||
#include <execution/cuda/LaunchDims.h>
|
||||
|
||||
namespace sd {
|
||||
|
||||
template <typename T>
|
||||
SD_KERNEL void tileKernel(void const* inputBuffer,
|
||||
LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
LongType const* outputShape,
|
||||
LongType resultLength) {
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Cache shape info to avoid repeated calls
|
||||
__shared__ sd::LongType inRank;
|
||||
__shared__ const sd::LongType* inShapePtr;
|
||||
__shared__ const sd::LongType* inStridePtr;
|
||||
|
||||
__shared__ sd::LongType outRank;
|
||||
__shared__ const sd::LongType* outShapePtr;
|
||||
__shared__ const sd::LongType* outStridePtr;
|
||||
__shared__ char outOrder;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inRank = shape::rank(inputShape);
|
||||
inShapePtr = shape::shapeOf(inputShape);
|
||||
inStridePtr = shape::stride(inputShape);
|
||||
|
||||
outRank = shape::rank(outputShape);
|
||||
outShapePtr = shape::shapeOf(outputShape);
|
||||
outStridePtr= shape::stride(outputShape);
|
||||
|
||||
outOrder = shape::order(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto inData = reinterpret_cast<const T*>(inputBuffer);
|
||||
auto outData = reinterpret_cast<T*>(outputBuffer);
|
||||
|
||||
if (outOrder == 'c') {
|
||||
// If the output is in 'c' order, we do direct linear indexing in output
|
||||
for (LongType i = tid; i < resultLength; i += totalThreads) {
|
||||
// We compute the input offset by using the output coordinate
|
||||
// to index into the input shape/stride
|
||||
sd::LongType coords[SD_MAX_RANK];
|
||||
sd::LongType inOffset;
|
||||
|
||||
INDEX2COORDS(i, outRank, outShapePtr, coords);
|
||||
COORDS2INDEX(outRank, inStridePtr, coords, inOffset);
|
||||
|
||||
// outData[i] = inData[inOffset]
|
||||
// The linear output index is i, so the input is
|
||||
// determined by the coords from the output shape
|
||||
outData[i] = inData[inOffset];
|
||||
}
|
||||
}
|
||||
else {
|
||||
// If the output has some other order, we do a more general coordinate transform
|
||||
for (LongType i = tid; i < resultLength; i += totalThreads) {
|
||||
// We map the linear index i into coordinates for the output shape
|
||||
sd::LongType outCoords[SD_MAX_RANK];
|
||||
sd::LongType outOffset;
|
||||
|
||||
INDEX2COORDS(i, outRank, outShapePtr, outCoords);
|
||||
COORDS2INDEX(outRank, outStridePtr, outCoords, outOffset);
|
||||
|
||||
// Then we interpret i as an index for the input as well, or use outCoords
|
||||
// Actually, the kernel code as written uses the same index i for input coords,
|
||||
// but let's remain consistent with the original logic:
|
||||
sd::LongType inCoords[SD_MAX_RANK];
|
||||
sd::LongType inOffset;
|
||||
|
||||
INDEX2COORDS(i, inRank, inShapePtr, inCoords);
|
||||
COORDS2INDEX(inRank, inStridePtr, inCoords, inOffset);
|
||||
|
||||
outData[outOffset] = inData[inOffset];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// We build specialized versions of tileKernel for all SD_COMMON_TYPES
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
SD_KERNEL void tileKernel,
|
||||
(void const* inputBuffer,
|
||||
sd::LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
sd::LongType const* outputShape,
|
||||
sd::LongType resultLength),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
template <typename T>
|
||||
void tileKernelH(void const* inputBuffer,
|
||||
LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
LongType const* outputShape,
|
||||
LongType resultLength,
|
||||
cudaStream_t* stream) {
|
||||
|
||||
dim3 launchDims = getLaunchDims("tile");
|
||||
tileKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
inputBuffer, inputShape, outputBuffer, outputShape, resultLength);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "tileKernel failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE(
|
||||
void tileKernelH,
|
||||
(void const* inputBuffer,
|
||||
sd::LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
sd::LongType const* outputShape,
|
||||
sd::LongType resultLength,
|
||||
cudaStream_t* stream),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
// Enhancement for different input (Y) and output (X) data types
|
||||
template <typename X, typename Y>
|
||||
SD_KERNEL void tileKernelDouble(
|
||||
void const* inputBuffer,
|
||||
LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
LongType const* outputShape,
|
||||
LongType resultLength) {
|
||||
|
||||
const int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
__shared__ sd::LongType inRank;
|
||||
__shared__ const sd::LongType* inShapePtr;
|
||||
__shared__ const sd::LongType* inStridePtr;
|
||||
|
||||
__shared__ sd::LongType outRank;
|
||||
__shared__ const sd::LongType* outShapePtr;
|
||||
__shared__ const sd::LongType* outStridePtr;
|
||||
__shared__ char outOrder;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
inRank = shape::rank(inputShape);
|
||||
inShapePtr = shape::shapeOf(inputShape);
|
||||
inStridePtr = shape::stride(inputShape);
|
||||
|
||||
outRank = shape::rank(outputShape);
|
||||
outShapePtr = shape::shapeOf(outputShape);
|
||||
outStridePtr= shape::stride(outputShape);
|
||||
|
||||
outOrder = shape::order(outputShape);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
const auto inData = reinterpret_cast<const Y*>(inputBuffer);
|
||||
auto outData = reinterpret_cast<X*>(outputBuffer);
|
||||
|
||||
if (outOrder == 'c') {
|
||||
for (LongType i = tid; i < resultLength; i += totalThreads) {
|
||||
sd::LongType outCoords[SD_MAX_RANK];
|
||||
sd::LongType inCoords[SD_MAX_RANK];
|
||||
sd::LongType inOffset;
|
||||
|
||||
// Get output coordinates
|
||||
INDEX2COORDS(i, outRank, outShapePtr, outCoords);
|
||||
|
||||
// Map to input coordinates (using modulo for tiling)
|
||||
for (int d = 0; d < inRank; d++) {
|
||||
inCoords[d] = outCoords[d] % inShapePtr[d];
|
||||
}
|
||||
|
||||
// Get input offset from input coordinates
|
||||
COORDS2INDEX(inRank, inStridePtr, inCoords, inOffset);
|
||||
|
||||
outData[i] = inData[inOffset];
|
||||
}
|
||||
}
|
||||
else {
|
||||
for (LongType i = tid; i < resultLength; i += totalThreads) {
|
||||
sd::LongType outCoords[SD_MAX_RANK];
|
||||
sd::LongType outOffset;
|
||||
sd::LongType inCoords[SD_MAX_RANK];
|
||||
sd::LongType inOffset;
|
||||
|
||||
INDEX2COORDS(i, outRank, outShapePtr, outCoords);
|
||||
COORDS2INDEX(outRank, outStridePtr, outCoords, outOffset);
|
||||
|
||||
// The original logic does a symmetrical approach for input.
|
||||
// We'll maintain that for consistency:
|
||||
INDEX2COORDS(i, inRank, inShapePtr, inCoords);
|
||||
COORDS2INDEX(inRank, inStridePtr, inCoords, inOffset);
|
||||
|
||||
outData[outOffset] = static_cast<X>(inData[inOffset]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE_TWICE(
|
||||
SD_KERNEL void tileKernelDouble,
|
||||
(void const* inputBuffer,
|
||||
sd::LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
sd::LongType const* outputShape,
|
||||
sd::LongType resultLength),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
// The host wrapper for tileKernelDouble
|
||||
template <typename X, typename Y>
|
||||
void tileKernelHH(void const* inputBuffer,
|
||||
LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
LongType const* outputShape,
|
||||
LongType resultLength,
|
||||
cudaStream_t* stream) {
|
||||
|
||||
dim3 launchDims = getLaunchDims("tile");
|
||||
tileKernelDouble<X, Y><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
inputBuffer, inputShape, outputBuffer, outputShape, resultLength);
|
||||
|
||||
DebugHelper::checkErrorCode(stream, "tileKernelDouble(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE_TWICE(
|
||||
void tileKernelHH,
|
||||
(void const* inputBuffer,
|
||||
sd::LongType const* inputShape,
|
||||
void* outputBuffer,
|
||||
sd::LongType const* outputShape,
|
||||
sd::LongType resultLength,
|
||||
cudaStream_t* stream),
|
||||
SD_COMMON_TYPES);
|
||||
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,363 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <helpers/DebugHelper.h>
|
||||
|
||||
#include <helpers/shape.h>
|
||||
#include <loops/summarystatsreduce.h>
|
||||
#include <ops/specials_cuda.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/float16.h>
|
||||
#include <types/types.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
namespace functions {
|
||||
namespace summarystats {
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_KERNEL void summaryStatsReduceKernel(
|
||||
int op, void * dx, sd::LongType * xShapeInfo, sd::LongType xRank,
|
||||
void* extraParams, void* z, sd::LongType * zShapeInfo, sd::LongType zRank,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength, int postProcessOrNot,
|
||||
bool biasCorrected, sd::LongType* allocationBuffer, void* reductionBuffer,
|
||||
sd::LongType * tadOnlyShapeInfo, sd::LongType * tadOffsets) {
|
||||
|
||||
SummaryStatsReduce<X, Z>::transform(
|
||||
op, dx, xShapeInfo, extraParams, z, zShapeInfo, dimension, dimensionLength,
|
||||
postProcessOrNot, allocationBuffer, reductionBuffer, tadOnlyShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param sPartialsRef
|
||||
* @param tid
|
||||
* @param extraParams
|
||||
*/
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void SummaryStatsReduce<X, Z>::aggregatePartials(SummaryStatsData<X>* sPartials, sd::LongType tid,
|
||||
sd::LongType numElements, void* vextraParams) {
|
||||
// start the shared memory loop on the next power of 2 less
|
||||
// than the block size. If block size is not a power of 2,
|
||||
// accumulate the intermediate sums in the remainder range.
|
||||
|
||||
auto extraParams = static_cast<Z*>(vextraParams);
|
||||
sd::LongType floorPow2 = numElements;
|
||||
|
||||
if (floorPow2 & (floorPow2 - 1)) {
|
||||
while (floorPow2 & (floorPow2 - 1)) {
|
||||
floorPow2 &= floorPow2 - 1;
|
||||
}
|
||||
|
||||
if (tid >= floorPow2) {
|
||||
SummaryStatsData<X> prev = sPartials[tid - floorPow2];
|
||||
SummaryStatsData<X> curr = sPartials[tid];
|
||||
sPartials[tid - floorPow2] = update(prev, curr, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
for (sd::LongType activeThreads = floorPow2 >> 1; activeThreads; activeThreads >>= 1) {
|
||||
if (tid < activeThreads && tid + activeThreads < numElements) {
|
||||
SummaryStatsData<X> curr = sPartials[tid];
|
||||
SummaryStatsData<X> next = sPartials[tid + activeThreads];
|
||||
sPartials[tid] = update(curr, next, extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @param n n is the number of
|
||||
* elements to loop through
|
||||
* @param dx the data to operate on
|
||||
* @param xVectorInfo the meta data for the vector:
|
||||
* 0 is the offset
|
||||
* 1 is the increment/stride
|
||||
* 2 is the real length of the buffer (n and dx.length won't always be the same)
|
||||
* 3 is the element wise stride for the buffer
|
||||
* 4 is the number of elements it takes to get to the next row/column/tensor
|
||||
* @param gpuInformation
|
||||
* 0 is the block size
|
||||
* 1 is the grid size
|
||||
* 2 is the shared memory size
|
||||
* @param problemDefinition
|
||||
* 0 is the number of elements per vector
|
||||
* 1 is the number of vectors
|
||||
*/
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void SummaryStatsReduce<X, Z>::transform(void * vx, sd::LongType * xShapeInfo, void* vextraParams,
|
||||
void* vz, sd::LongType * zShapeInfo, sd::LongType* dimension,
|
||||
sd::LongType dimensionLength, int postProcessOrNot,
|
||||
sd::LongType* allocationBuffer,
|
||||
void* vreductionBuffer, sd::LongType * tadOnlyShapeInfo,
|
||||
sd::LongType * tadOffsets) {
|
||||
auto dx = static_cast<X *>(vx);
|
||||
auto z = static_cast<Z*>(vz);
|
||||
auto extraParams = static_cast<Z*>(vextraParams);
|
||||
auto reductionBuffer = static_cast<Z*>(vreductionBuffer);
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
__shared__ volatile bool resultScalar;
|
||||
|
||||
int numElements = blockDim.x;
|
||||
// shared memory space for storing intermediate results
|
||||
__shared__ SummaryStatsData<X> sPartials[SD_CUDA_BLOCK_SIZE];
|
||||
|
||||
// Cache shape information for x buffer
|
||||
__shared__ sd::LongType xRank;
|
||||
__shared__ sd::LongType* xShapePtr;
|
||||
__shared__ sd::LongType* xStridePtr;
|
||||
|
||||
// Cache shape information for TAD
|
||||
__shared__ sd::LongType tadRank;
|
||||
__shared__ sd::LongType* tadShapePtr;
|
||||
__shared__ sd::LongType* tadStridePtr;
|
||||
|
||||
Z startingVal = startingValue(dx);
|
||||
|
||||
SummaryStatsData<X> val;
|
||||
val.initWithValue(static_cast<X>(startingVal));
|
||||
val.n = 0;
|
||||
sPartials[threadIdx.x] = val;
|
||||
|
||||
// length for the tad
|
||||
__shared__ volatile int xLength;
|
||||
__shared__ volatile int resultLength;
|
||||
|
||||
SummaryStatsData<X> reduction;
|
||||
reduction.initWithValue(static_cast<X>(0.0));
|
||||
reduction.n = 0;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
if (zShapeInfo != nullptr)
|
||||
resultLength = shape::length(zShapeInfo);
|
||||
else
|
||||
resultLength = 1;
|
||||
|
||||
|
||||
|
||||
if (resultLength <= 1)
|
||||
resultScalar = 1;
|
||||
|
||||
xLength = shape::length(xShapeInfo);
|
||||
|
||||
// Cache x shape information
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
// Cache TAD shape information
|
||||
if (tadOnlyShapeInfo != nullptr && !resultScalar) {
|
||||
tadRank = shape::rank(tadOnlyShapeInfo);
|
||||
tadShapePtr = shape::shapeOf(tadOnlyShapeInfo);
|
||||
tadStridePtr = shape::stride(tadOnlyShapeInfo);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (!resultScalar) {
|
||||
__shared__ int tadLength;
|
||||
__shared__ int numTads;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
tadLength = shape::length(tadOnlyShapeInfo);
|
||||
numTads = shape::length(xShapeInfo) / tadLength;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int r = blockIdx.x; r < numTads; r += gridDim.x) {
|
||||
auto tadOffsetForBlock = tadOffsets[r];
|
||||
|
||||
val.initWithValue(static_cast<X>(startingVal));
|
||||
val.n = 0;
|
||||
sPartials[threadIdx.x] = val;
|
||||
|
||||
for (int i = threadIdx.x; i < tadLength; i += blockDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(i, tadRank, tadShapePtr, xCoords);
|
||||
COORDS2INDEX(tadRank, tadStridePtr, xCoords, xOffset);
|
||||
auto xOffsetFinal = tadOffsetForBlock + xOffset;
|
||||
SummaryStatsData<X> indexVal2;
|
||||
indexVal2.initWithValue(dx[xOffsetFinal]);
|
||||
|
||||
sPartials[threadIdx.x] = update(sPartials[threadIdx.x], OpType::op(indexVal2, extraParams), extraParams);
|
||||
}
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdx.x, sd::math::sd_min<int>(blockDim.x, tadLength), extraParams);
|
||||
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
z[r] = OpType::getValue(postProcessOrNot, sPartials[threadIdx.x]);
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
} else if (resultScalar) {
|
||||
__shared__ int n;
|
||||
if (threadIdx.x == 0) {
|
||||
n = shape::length(xShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (sd::LongType i = tid; i < n; i += blockDim.x * gridDim.x) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
SummaryStatsData<X> indexVal2;
|
||||
indexVal2.initWithValue(dx[xOffset]);
|
||||
reduction = update(reduction, indexVal2, extraParams);
|
||||
}
|
||||
sPartials[threadIdx.x] = reduction;
|
||||
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdx.x, blockDim.x, extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (gridDim.x > 1) {
|
||||
__shared__ bool amLast;
|
||||
unsigned int* tc = (unsigned int*)reductionBuffer;
|
||||
tid = threadIdx.x;
|
||||
if (threadIdx.x == 0) {
|
||||
SummaryStatsData<X>* pBuffer = (SummaryStatsData<X>*)reductionBuffer;
|
||||
pBuffer[blockIdx.x] = sPartials[0];
|
||||
}
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
unsigned int ticket = atomicInc(&tc[16384], gridDim.x);
|
||||
amLast = (ticket == gridDim.x - 1);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (amLast) {
|
||||
tc[16384] = 0;
|
||||
SummaryStatsData<X>* pBuffer = (SummaryStatsData<X>*)reductionBuffer;
|
||||
|
||||
Z startingVal = startingValue(dx);
|
||||
|
||||
SummaryStatsData<X> val;
|
||||
val.initWithValue(static_cast<X>(startingVal));
|
||||
val.n = 0;
|
||||
sPartials[threadIdx.x] = val;
|
||||
|
||||
for (int i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
|
||||
sPartials[threadIdx.x] = update(sPartials[threadIdx.x], pBuffer[i], extraParams);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
aggregatePartials<OpType>(sPartials, threadIdx.x, gridDim.x, extraParams);
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
z[0] = OpType::getValue(postProcessOrNot, sPartials[0]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if (tid == 0) {
|
||||
unsigned int* tc = (unsigned*)reductionBuffer;
|
||||
tc[16384] = 0;
|
||||
z[0] = OpType::getValue(postProcessOrNot, sPartials[0]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename X, typename Y>
|
||||
SD_DEVICE void SummaryStatsReduce<X, Y>::transform( int opNum, void * dx, sd::LongType * xShapeInfo,
|
||||
void* extraParams, void* z, sd::LongType * zShapeInfo,
|
||||
sd::LongType* dimension, sd::LongType dimensionLength, int postProcessOrNot, sd::LongType* allocationBuffer, void* reductionBuffer,
|
||||
sd::LongType * tadOnlyShapeInfo,
|
||||
sd::LongType * tadOffsets) {
|
||||
DISPATCH_BY_OPNUM_TT(transform,
|
||||
PARAMS(dx, xShapeInfo, extraParams, z, zShapeInfo, dimension, dimensionLength, postProcessOrNot,
|
||||
allocationBuffer, reductionBuffer, tadOnlyShapeInfo, tadOffsets),
|
||||
SUMMARY_STATS_OPS);
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void SummaryStatsReduce<X, Z>::execSummaryStatsReduceScalar(
|
||||
dim3& launchDims, cudaStream_t* stream, int opNum, void * vx, sd::LongType * xShapeInfo,
|
||||
sd::LongType * hxShapeInfo, void* vextraParams, void* vz, sd::LongType * zShapeInfo,
|
||||
sd::LongType * hzShapeInfo, sd::LongType * tadShapeInfo, sd::LongType * tadOffsets,
|
||||
bool biasCorrected, void* reductionBuffer) {
|
||||
auto x = static_cast<X *>(vx);
|
||||
auto extraParams = static_cast<Z*>(vextraParams);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto reductionPointerA = reinterpret_cast<Z*>(reductionBuffer);
|
||||
|
||||
if (sd::Environment::getInstance().isDebugAndVerbose()) printf("D16 opNum:[%i]\n", opNum);
|
||||
|
||||
summaryStatsReduceKernel<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum,
|
||||
x,
|
||||
xShapeInfo,
|
||||
shape::rank(hxShapeInfo),
|
||||
extraParams,
|
||||
z,
|
||||
zShapeInfo,
|
||||
shape::rank(hzShapeInfo),
|
||||
nullptr,
|
||||
0,
|
||||
1,
|
||||
biasCorrected,
|
||||
nullptr,
|
||||
reductionPointerA,
|
||||
tadShapeInfo,
|
||||
tadOffsets);
|
||||
|
||||
// this is blocking method since method should return scalar
|
||||
sd::DebugHelper::checkErrorCode(stream, "execSSReduceScalar(...) failed");
|
||||
}
|
||||
|
||||
template <typename X, typename Z>
|
||||
SD_HOST void SummaryStatsReduce<X, Z>::execSummaryStatsReduce(
|
||||
dim3& launchDims, cudaStream_t* stream, int opNum, void * vx, sd::LongType * xShapeInfo,
|
||||
sd::LongType * hxShapeInfo, void* vextraParams, void* vz, sd::LongType * zShapeInfo,
|
||||
sd::LongType * hzShapeInfo, sd::LongType * tadShapeInfo, sd::LongType * tadOffsets,
|
||||
bool biasCorrected, void* reductionBuffer) {
|
||||
auto x = static_cast<X *>(vx);
|
||||
auto z = static_cast<Z*>(vz);
|
||||
auto extraParams = static_cast<Z*>(vextraParams);
|
||||
|
||||
if (sd::Environment::getInstance().isDebugAndVerbose()) printf("F17 opNum:[%i]\n", opNum);
|
||||
|
||||
auto reductionPointerA = reinterpret_cast<Z*>(reductionBuffer);
|
||||
|
||||
summaryStatsReduceKernel<X, Z><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
opNum, x, xShapeInfo, shape::rank(hxShapeInfo), extraParams, z, zShapeInfo, shape::rank(hzShapeInfo), nullptr, 1,
|
||||
1, biasCorrected, nullptr, reductionPointerA, tadShapeInfo, tadOffsets);
|
||||
|
||||
DEBUG_KERNEL(stream, opNum);
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class SummaryStatsReduce, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
|
||||
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
* 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/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_any.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <execution/cuda/DeviceValidator.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The kernel that calls the transform CUDA method,
|
||||
// caching shape info in shared memory for offset computations.
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void transformAnySimpleCached(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* params,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Just delegate to transformCuda,
|
||||
// which will do the shape caching logic for coords->offset conversions.
|
||||
functions::transform::TransformAny<X, Z>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the "executeTransformShaped" that calls the new cached kernel
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void TransformAny<X, Y>::executeTransformShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_ANY_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformAny executeTransformShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The transformCuda method that uses shared memory for shape/stride caching,
|
||||
// then does coords->offset conversions.
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformAny<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// cast pointers
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto params = reinterpret_cast<X*>(vparams);
|
||||
|
||||
if (x == nullptr || z == nullptr) return;
|
||||
|
||||
// cache shape info in shared memory
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ int zRank;
|
||||
__shared__ const sd::LongType* zShapePtr;
|
||||
__shared__ const sd::LongType* zStridePtr;
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRank = shape::rank(xShapeInfo);
|
||||
xShapePtr = shape::shapeOf(xShapeInfo);
|
||||
xStridePtr = shape::stride(xShapeInfo);
|
||||
|
||||
zRank = shape::rank(zShapeInfo);
|
||||
zShapePtr = shape::shapeOf(zShapeInfo);
|
||||
zStridePtr = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// do the transform
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType offsetX;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
// convert i -> coords -> offset for x
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, offsetX);
|
||||
|
||||
// convert i -> coords -> offset for z
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
z[offsetZ] = OpType::op(x[offsetX], params);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformAny<X, Z>::intermediateShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// We call the new transformAnySimpleCached kernel
|
||||
transformAnySimpleCached<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShape,
|
||||
xRank,
|
||||
extraParams,
|
||||
z,
|
||||
zShape,
|
||||
zRank,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformAny(...) cached kernel failed");
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class TransformAny, , SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,210 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
* 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/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_bool.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <execution/cuda/DeviceValidator.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Cached kernel that caches shape info in shared memory and uses cached variables
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void transformBoolSimpleCached(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* params,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Delegate the transform to the device function with cached shape info
|
||||
functions::transform::TransformBool<X, Z>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the "executeTransformShaped" that launches the cached kernel
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void TransformBool<X, Y>::executeTransformShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
long long int xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
long long int zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_BOOL_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformBool executeTransformShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Device function that caches shape info and uses cached variables for computations
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformBool<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Cast pointers to appropriate types
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto params = reinterpret_cast<X*>(vparams);
|
||||
auto reductionPointer = reinterpret_cast<Z*>(vreductionPointer);
|
||||
|
||||
// Check for special operations
|
||||
if (OpType::requiresSpecial) {
|
||||
OpType::execSpecialCuda(x,
|
||||
xShapeInfo,
|
||||
z,
|
||||
zShapeInfo,
|
||||
params,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets);
|
||||
return;
|
||||
}
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRankCached;
|
||||
__shared__ const sd::LongType* xShapePtrCached;
|
||||
__shared__ const sd::LongType* xStridePtrCached;
|
||||
|
||||
__shared__ int zRankCached;
|
||||
__shared__ const sd::LongType* zShapePtrCached;
|
||||
__shared__ const sd::LongType* zStridePtrCached;
|
||||
|
||||
// Thread 0 caches the shape information
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRankCached = shape::rank(xShapeInfo);
|
||||
xShapePtrCached = shape::shapeOf(xShapeInfo);
|
||||
xStridePtrCached = shape::stride(xShapeInfo);
|
||||
|
||||
zRankCached = shape::rank(zShapeInfo);
|
||||
zShapePtrCached = shape::shapeOf(zShapeInfo);
|
||||
zStridePtrCached = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate thread ID and total threads
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Loop over all elements using cached shape info
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
// Convert index to coordinates using cached shape info
|
||||
INDEX2COORDS(i, xRankCached, xShapePtrCached, xCoords);
|
||||
COORDS2INDEX(xRankCached, xStridePtrCached, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, zRankCached, zShapePtrCached, zCoords);
|
||||
COORDS2INDEX(zRankCached, zStridePtrCached, zCoords, zOffset);
|
||||
|
||||
// Apply the operation using cached offsets
|
||||
z[zOffset] = OpType::op(x[xOffset], params);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Host function that launches the cached kernel
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformBool<X, Z>::intermediateShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Launch the cached kernel
|
||||
transformBoolSimpleCached<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShape,
|
||||
xRank,
|
||||
extraParams,
|
||||
z,
|
||||
zShape,
|
||||
zRank,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets
|
||||
);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformBool(...) cached kernel failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Macro to instantiate templates for TransformBool with common and bool types
|
||||
BUILD_DOUBLE_TEMPLATE( class TransformBool, , SD_COMMON_TYPES, SD_BOOL_TYPES);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,202 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_float.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <execution/cuda/DeviceValidator.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The cached kernel that caches shape info in shared memory and uses cached variables
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void transformFloatSimpleCached(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* params,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Delegate the transform to the transformCuda method with cached shape info
|
||||
functions::transform::TransformFloat<X, Z>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the "executeTransformShaped" that launches the cached kernel
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void TransformFloat<X, Y>::executeTransformShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_FLOAT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformFloat executeTransformShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Device function that caches shape info and uses cached variables for computations
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformFloat<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Cast pointers to appropriate types
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto params = reinterpret_cast<Z*>(vparams);
|
||||
auto reductionPointer = reinterpret_cast<Z*>(vreductionPointer);
|
||||
|
||||
// Check for special operations
|
||||
if (OpType::requiresSpecial) {
|
||||
OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
|
||||
tadOffsets);
|
||||
return;
|
||||
}
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRankCached;
|
||||
__shared__ const sd::LongType* xShapePtrCached;
|
||||
__shared__ const sd::LongType* xStridePtrCached;
|
||||
|
||||
__shared__ int zRankCached;
|
||||
__shared__ const sd::LongType* zShapePtrCached;
|
||||
__shared__ const sd::LongType* zStridePtrCached;
|
||||
|
||||
// Thread 0 caches the shape information
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRankCached = shape::rank(xShapeInfo);
|
||||
xShapePtrCached = shape::shapeOf(xShapeInfo);
|
||||
xStridePtrCached = shape::stride(xShapeInfo);
|
||||
|
||||
zRankCached = shape::rank(zShapeInfo);
|
||||
zShapePtrCached = shape::shapeOf(zShapeInfo);
|
||||
zStridePtrCached = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate thread ID and total threads
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Loop over all elements using cached shape info
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
// Convert index to coordinates using cached shape info
|
||||
INDEX2COORDS(i, xRankCached, xShapePtrCached, xCoords);
|
||||
COORDS2INDEX(xRankCached, xStridePtrCached, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, zRankCached, zShapePtrCached, zCoords);
|
||||
COORDS2INDEX(zRankCached, zStridePtrCached, zCoords, zOffset);
|
||||
|
||||
// Apply the operation using cached offsets
|
||||
z[zOffset] = OpType::op(x[xOffset], params);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Host function that launches the cached kernel
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformFloat<X, Z>::intermediateShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Launch the cached kernel
|
||||
transformFloatSimpleCached<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShape,
|
||||
xRank,
|
||||
extraParams,
|
||||
z,
|
||||
zShape,
|
||||
zRank,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets
|
||||
);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformFloat(...) cached kernel failed");
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class TransformFloat, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
@@ -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 raver119@gmail.com
|
||||
//
|
||||
#include <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_same.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
template <typename X, typename OpType>
|
||||
SD_KERNEL void transformSameSimple(const void *x, const sd::LongType *xShapeInfo, long long int xRank, void *params, void *z,
|
||||
const sd::LongType *zShapeInfo, long long int zRank,
|
||||
sd::LongType *allocationPointer,
|
||||
void *reductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
functions::transform::TransformSame<X>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
template <typename X>
|
||||
SD_HOST void TransformSame<X>::executeTransformShaped(dim3 launchDims, cudaStream_t *stream, const int opNum,
|
||||
const void *x, const sd::LongType *xShape, sd::LongType xRank,
|
||||
void *extraParams, void *z, const sd::LongType *zShape,
|
||||
sd::LongType zRank, sd::LongType *allocationPointer, void *reductionPointer,
|
||||
const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
DISPATCH_BY_OPNUM_T(intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_SAME_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformAny(...) failed");
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformSame<X>::transformCuda(const void *vx, const sd::LongType *xShapeInfo, void *vparams, void *vz,
|
||||
const sd::LongType *zShapeInfo, sd::LongType *allocationPointer,
|
||||
void *vreductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
auto x = static_cast<const X *>(vx);
|
||||
auto z = static_cast<X *>(vz);
|
||||
auto params = static_cast<X *>(vparams);
|
||||
auto reductionPointer = static_cast<X *>(vreductionPointer);
|
||||
|
||||
if (OpType::requiresSpecial) {
|
||||
OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
|
||||
tadOffsets);
|
||||
return;
|
||||
} else {
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// 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) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
// 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 = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], params);
|
||||
}
|
||||
}
|
||||
};
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformSame<X>::intermediateShaped(dim3 launchDims, cudaStream_t *stream, const void *x,
|
||||
const sd::LongType *xShape, sd::LongType xRank, void *extraParams, void *z,
|
||||
const sd::LongType *zShape, sd::LongType zRank,
|
||||
sd::LongType *allocationPointer,
|
||||
void *reductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
transformSameSimple<X, OpType><<<launchDims.x, launchDims.x, launchDims.z, *stream>>>(
|
||||
x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformSame(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( class TransformSame, , SD_COMMON_TYPES);
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
@@ -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 <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_strict.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
template <typename X, typename OpType>
|
||||
SD_KERNEL void transformStrictSimple(const void *x, const sd::LongType *xShapeInfo, long long int xRank, void *params, void *z,
|
||||
const sd::LongType *zShapeInfo, long long int zRank,
|
||||
sd::LongType *allocationPointer,
|
||||
void *reductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
functions::transform::TransformStrict<X>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
template <typename X>
|
||||
SD_HOST void TransformStrict<X>::executeTransformShaped(dim3 launchDims, cudaStream_t *stream, const int opNum,
|
||||
const void *x, const sd::LongType *xShape, sd::LongType xRank,
|
||||
void *extraParams, void *z, const sd::LongType *zShape,
|
||||
sd::LongType zRank, sd::LongType *allocationPointer, void *reductionPointer,
|
||||
const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
DISPATCH_BY_OPNUM_T(intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_STRICT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformStrict(...) failed");
|
||||
}
|
||||
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformStrict<X>::transformCuda(const void *vx, const sd::LongType *xShapeInfo, void *vparams,
|
||||
void *vz, const sd::LongType *zShapeInfo,
|
||||
sd::LongType *allocationPointer,
|
||||
void *vreductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
auto x = static_cast<const X *>(vx);
|
||||
auto z = static_cast<X *>(vz);
|
||||
auto params = static_cast<X *>(vparams);
|
||||
auto reductionPointer = static_cast<X *>(vreductionPointer);
|
||||
|
||||
if (OpType::requiresSpecial) {
|
||||
OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
|
||||
tadOffsets);
|
||||
return;
|
||||
} else {
|
||||
__shared__ sd::LongType length;
|
||||
|
||||
// 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) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
// 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 = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType xCoords[SD_MAX_RANK];
|
||||
sd::LongType zCoords[SD_MAX_RANK];
|
||||
sd::LongType xOffset;
|
||||
sd::LongType zOffset;
|
||||
|
||||
INDEX2COORDS(i, xRank, xShapePtr, xCoords);
|
||||
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
||||
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
||||
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
||||
|
||||
z[zOffset] = OpType::op(x[xOffset], params);
|
||||
}
|
||||
}
|
||||
};
|
||||
template <typename X>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformStrict<X>::intermediateShaped(dim3 launchDims, cudaStream_t *stream, const void *x,
|
||||
const sd::LongType *xShape, sd::LongType xRank, void *extraParams, void *z,
|
||||
const sd::LongType *zShape, sd::LongType zRank,
|
||||
sd::LongType *allocationPointer,
|
||||
void *reductionPointer, const sd::LongType *tadShapeInfo,
|
||||
const sd::LongType *tadOffsets) {
|
||||
transformStrictSimple<X, OpType><<<launchDims.x, launchDims.x, launchDims.z, *stream>>>(
|
||||
x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformStrict(...) failed");
|
||||
}
|
||||
|
||||
BUILD_SINGLE_TEMPLATE( class TransformStrict, , SD_FLOAT_TYPES);
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
@@ -0,0 +1,541 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/DebugHelper.h>
|
||||
#include <loops/type_conversions.h>
|
||||
#include <types/types.h>
|
||||
|
||||
namespace sd {
|
||||
template <typename S, typename T>
|
||||
void TypeCast::convertGenericCuda(Pointer *extras, void *dx, LongType N, void *dz) {
|
||||
auto stream = reinterpret_cast<cudaStream_t *>(&extras[1]);
|
||||
|
||||
sd::convertKernel<S, T><<<256, 1024, 1024, *stream>>>(dx, N, dz);
|
||||
DebugHelper::checkErrorCode(stream, "convertGeneric(...) failed");
|
||||
};
|
||||
|
||||
template <typename S, typename T>
|
||||
SD_DEVICE void convertKernelGeneric(S *x, LongType N, T *z) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
for (LongType i = tid; i < N; i += blockDim.x * gridDim.x) {
|
||||
// despite it's stupid, it simplifies conversion to bottom dtypes
|
||||
// FIXME: get rid of through-float though
|
||||
z[i] = static_cast<T>(static_cast<float>(x[i]));
|
||||
}
|
||||
};
|
||||
|
||||
// Define this to more rigorously avoid bank conflicts, even at the lower (root) levels of the tree
|
||||
//#define ZERO_BANK_CONFLICTS
|
||||
|
||||
#ifdef ZERO_BANK_CONFLICTS
|
||||
#define CONFLICT_FREE_OFFSET(index) ((index) >> LOG_NUM_BANKS + (index) >> (2 * LOG_NUM_BANKS))
|
||||
#else
|
||||
#define CONFLICT_FREE_OFFSET(index) ((index) >> LOG_NUM_BANKS)
|
||||
#endif
|
||||
|
||||
#ifdef CHECK_BANK_CONFLICTS
|
||||
#define TEMP(index) CUT_BANK_CHECKER(temp, index)
|
||||
#else
|
||||
#define TEMP(index) temp[index]
|
||||
#endif
|
||||
|
||||
template <bool isNP2>
|
||||
SD_DEVICE void loadSharedChunkFromMem(int *s_data, const int *g_idata, int n, int baseIndex, int &ai, int &bi,
|
||||
int &mem_ai, int &mem_bi, int &bankOffsetA, int &bankOffsetB) {
|
||||
int thid = threadIdx.x;
|
||||
mem_ai = baseIndex + threadIdx.x;
|
||||
mem_bi = mem_ai + blockDim.x;
|
||||
|
||||
ai = thid;
|
||||
bi = thid + blockDim.x;
|
||||
|
||||
// compute spacing to avoid bank conflicts
|
||||
bankOffsetA = CONFLICT_FREE_OFFSET(ai);
|
||||
bankOffsetB = CONFLICT_FREE_OFFSET(bi);
|
||||
|
||||
// Cache the computational window in shared memory
|
||||
// pad values beyond n with zeros
|
||||
s_data[ai + bankOffsetA] = g_idata[mem_ai];
|
||||
|
||||
if (isNP2) { // compile-time decision
|
||||
s_data[bi + bankOffsetB] = (bi < n) ? g_idata[mem_bi] : static_cast<LongType>(0) ;
|
||||
} else {
|
||||
s_data[bi + bankOffsetB] = g_idata[mem_bi];
|
||||
}
|
||||
}
|
||||
|
||||
template <bool isNP2>
|
||||
SD_DEVICE void storeSharedChunkToMem(int *g_odata, int *s_data, int n, int ai, int bi, int mem_ai, int mem_bi,
|
||||
int bankOffsetA, int bankOffsetB) {
|
||||
__syncthreads();
|
||||
|
||||
// write results to global memory
|
||||
g_odata[mem_ai] = s_data[ai + bankOffsetA];
|
||||
if (isNP2) { // compile-time decision
|
||||
if (bi < n) g_odata[mem_bi] = s_data[bi + bankOffsetB];
|
||||
} else {
|
||||
g_odata[mem_bi] = s_data[bi + bankOffsetB];
|
||||
}
|
||||
}
|
||||
|
||||
template <bool storeSum>
|
||||
SD_DEVICE void clearLastElement(int *s_data, int *g_blockSums, int blockIndex) {
|
||||
if (threadIdx.x == 0) {
|
||||
int index = (blockDim.x << 1) - 1;
|
||||
index += CONFLICT_FREE_OFFSET(index);
|
||||
|
||||
if (storeSum) { // compile-time decision
|
||||
// write this block's total sum to the corresponding index in the blockSums array
|
||||
g_blockSums[blockIndex] = s_data[index];
|
||||
}
|
||||
|
||||
// zero the last element in the scan so it will propagate back to the front
|
||||
s_data[index] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
SD_DEVICE unsigned int buildSum(int *s_data) {
|
||||
unsigned int thid = threadIdx.x;
|
||||
unsigned int stride = 1;
|
||||
|
||||
// build the sum in place up the tree
|
||||
for (int d = blockDim.x; d > 0; d >>= 1) {
|
||||
__syncthreads();
|
||||
|
||||
if (thid < d) {
|
||||
int i = __mul24(__mul24(2, stride), thid);
|
||||
int ai = i + stride - 1;
|
||||
int bi = ai + stride;
|
||||
|
||||
ai += CONFLICT_FREE_OFFSET(ai);
|
||||
bi += CONFLICT_FREE_OFFSET(bi);
|
||||
|
||||
s_data[bi] += s_data[ai];
|
||||
}
|
||||
|
||||
stride *= 2;
|
||||
}
|
||||
|
||||
return stride;
|
||||
}
|
||||
|
||||
SD_DEVICE void scanRootToLeaves(int *s_data, unsigned int stride) {
|
||||
unsigned int thid = threadIdx.x;
|
||||
|
||||
// traverse down the tree building the scan in place
|
||||
for (int d = 1; d <= blockDim.x; d *= 2) {
|
||||
stride >>= 1;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (thid < d) {
|
||||
int i = __mul24(__mul24(2, stride), thid);
|
||||
int ai = i + stride - 1;
|
||||
int bi = ai + stride;
|
||||
|
||||
ai += CONFLICT_FREE_OFFSET(ai);
|
||||
bi += CONFLICT_FREE_OFFSET(bi);
|
||||
|
||||
float t = s_data[ai];
|
||||
s_data[ai] = s_data[bi];
|
||||
s_data[bi] += t;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <bool storeSum>
|
||||
SD_DEVICE void prescanBlock(int *data, int blockIndex, int *blockSums) {
|
||||
int stride = buildSum(data); // build the sum in place up the tree
|
||||
clearLastElement<storeSum>(data, blockSums, (blockIndex == 0) ? blockIdx.x : blockIndex);
|
||||
scanRootToLeaves(data, stride); // traverse down tree to build the scan
|
||||
}
|
||||
|
||||
template <bool storeSum, bool isNP2>
|
||||
SD_KERNEL void prescan(int *g_odata, const int *g_idata, int *g_blockSums, int n, int blockIndex, int baseIndex) {
|
||||
int ai, bi, mem_ai, mem_bi, bankOffsetA, bankOffsetB;
|
||||
extern __shared__ int s_data[];
|
||||
|
||||
// load data into shared memory
|
||||
loadSharedChunkFromMem<isNP2>(reinterpret_cast<int *>(s_data), g_idata, n,
|
||||
(baseIndex == 0) ? __mul24(blockIdx.x, (blockDim.x << 1)) : baseIndex, ai, bi, mem_ai,
|
||||
mem_bi, bankOffsetA, bankOffsetB);
|
||||
|
||||
// scan the data in each block
|
||||
prescanBlock<storeSum>(s_data, blockIndex, g_blockSums);
|
||||
|
||||
// write results to device memory
|
||||
storeSharedChunkToMem<isNP2>(g_odata, s_data, n, ai, bi, mem_ai, mem_bi, bankOffsetA, bankOffsetB);
|
||||
}
|
||||
|
||||
SD_KERNEL void uniformAdd(int *g_data, int *uniforms, int n, int blockOffset, int baseIndex) {
|
||||
__shared__ float uni;
|
||||
if (threadIdx.x == 0) uni = uniforms[blockIdx.x + blockOffset];
|
||||
|
||||
unsigned int address = __mul24(blockIdx.x, (blockDim.x << 1)) + baseIndex + threadIdx.x;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// note two adds per thread
|
||||
g_data[address] += uni;
|
||||
g_data[address + blockDim.x] += (threadIdx.x + blockDim.x < n) * uni;
|
||||
}
|
||||
|
||||
/*
|
||||
* This kernel does prefix sum in parallel, to calculate offsets for each block
|
||||
*/
|
||||
template <typename T>
|
||||
SD_DEVICE inline void encoderKernelP2Generic(void *dx, LongType n, void *dz) {
|
||||
// TODO: to be remove
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
/*
|
||||
* PLEASE NOTE: This kernel doesn't allow loop for data. Basically: grid will be huge.
|
||||
*/
|
||||
template <typename T>
|
||||
SD_KERNEL static void execEncoderKernelP1(const void *dx, LongType N, void *dz, float threshold) {
|
||||
auto x = reinterpret_cast<const T *>(dx);
|
||||
auto z = reinterpret_cast<int *>(dz);
|
||||
|
||||
// basically, for phase One we want do calculation: how many eligible values we have, and which blocks will be holding
|
||||
// data
|
||||
LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
int pass = tid < N && math::sd_abs<T,T>(x[tid]) >= static_cast<T>(threshold) ? static_cast<LongType>(1) : static_cast<LongType>(0) ;
|
||||
int bp = __syncthreads_count(pass);
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
// saving out per-block passes
|
||||
z[blockIdx.x + 1] = bp;
|
||||
|
||||
// saving out sum
|
||||
atomicAdd(&z[0], bp);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void encoderKernelP1Generic(dim3 &launchDims, cudaStream_t *stream, const void *dx, LongType N, void *dz,
|
||||
float threshold) {
|
||||
execEncoderKernelP1<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, N, dz, threshold);
|
||||
DebugHelper::checkErrorCode(stream, "encoderP1(...) failed");
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void encoderKernelP1Generic,
|
||||
(dim3 & launchDims, cudaStream_t *stream, const void *dx, sd::LongType N, void *dz,
|
||||
float threshold),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
/*
|
||||
* PLEASE NOTE: This kernel doesn't allow loop for data. Basically: grid will be huge.
|
||||
*
|
||||
* Based on: https://github.com/knotman90/cuStreamComp <-- efficient CUDA stream compaction algorithm
|
||||
*/
|
||||
template <typename T>
|
||||
SD_KERNEL static void execEncoderKernelP3(void *dx, int *offsets, LongType N, void *dz) {
|
||||
auto x = reinterpret_cast<T *>(dx);
|
||||
auto z = reinterpret_cast<int *>(dz);
|
||||
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
extern __shared__ int warpTotals[];
|
||||
|
||||
// fetch block offset only once
|
||||
__shared__ float threshold;
|
||||
__shared__ FloatBits fb;
|
||||
__shared__ int bo;
|
||||
__shared__ int limit;
|
||||
if (threadIdx.x == 0) {
|
||||
limit = z[0];
|
||||
fb.i_ = z[2];
|
||||
threshold = fb.f_;
|
||||
bo = offsets[blockIdx.x];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// out-of-limit threads do not play here
|
||||
auto value = tid < N ? x[tid] : (T)0.f;
|
||||
|
||||
// out-of-limit threads just declare they have no changes
|
||||
auto pred = tid >= N ? static_cast<LongType>(0) : math::sd_abs<T,T>(value) >= static_cast<T>(threshold) ? static_cast<LongType>(1) : static_cast<LongType>(0) ;
|
||||
auto w_i = threadIdx.x / warpSize; // warp index (or, warp number) - index of the Warp within TOTAL_WARPS
|
||||
auto t_i = threadIdx.x % warpSize; // thread index within a warp
|
||||
unsigned int t_m = INT_MAX >> (warpSize - t_i - 1); // thread mask (ERROR IN THE PAPER minus one is required)
|
||||
|
||||
int b = __ballot_sync(t_m, pred); // balres = number whose ith bit isone if the ith's thread pred is true masked up
|
||||
// to the current index in warp
|
||||
auto t_u = __popc(b); // popc count the number of bit one. simply count the number predicated true BEFORE MY INDEX
|
||||
|
||||
if (t_i == warpSize - 1) warpTotals[w_i] = t_u + pred;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
int w_i_u = 0;
|
||||
for (int j = 0; j <= 5; j++) {
|
||||
unsigned int b_j =
|
||||
__ballot_sync(t_m, warpTotals[t_i] & pow2i(j)); //# of the ones in the j'th digit of the warp offsets
|
||||
w_i_u += (__popc(b_j) << j);
|
||||
}
|
||||
|
||||
// we just ignore all results coming from non-0 threads
|
||||
if (w_i == 0 && t_i < blockDim.x / warpSize) warpTotals[t_i] = w_i_u;
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// pred is always false if we're out-of-limits
|
||||
if (pred) {
|
||||
int idx = t_u + warpTotals[w_i] + bo + 4;
|
||||
if (idx < limit + 4) {
|
||||
z[idx] = value > static_cast<T>(0.0f) ? tid + 1 : -(tid + 1);
|
||||
x[tid] = value > static_cast<T>(0.0f) ? x[tid] - threshold : x[tid] + threshold;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void encoderKernelP3Generic(dim3 &launchDims, cudaStream_t *stream, void *dx, int *offsets, LongType N,
|
||||
void *dz) {
|
||||
execEncoderKernelP3<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, offsets, N, dz);
|
||||
DebugHelper::checkErrorCode(stream, "encoderP3(...) failed");
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void encoderKernelP3Generic,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *dx, int *offsets, sd::LongType N, void *dz),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
/*
|
||||
* This kernel handles decode from sparse threshold array, to dense array
|
||||
*
|
||||
* PLEASE NOTE: Z is expected to be memset to 0
|
||||
*/
|
||||
template <typename T>
|
||||
SD_KERNEL static void execDecoderKernel(const void *dx, LongType N, void *dz) {
|
||||
auto x = reinterpret_cast<const int *>(dx);
|
||||
auto z = reinterpret_cast<T *>(dz);
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
__shared__ float threshold;
|
||||
__shared__ int limit;
|
||||
|
||||
__shared__ FloatBits fb;
|
||||
if (threadIdx.x == 0) {
|
||||
limit = x[0];
|
||||
fb.i_ = x[2];
|
||||
threshold = fb.f_;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int e = tid; e < limit; e += blockDim.x * gridDim.x) {
|
||||
int el = x[e + 4];
|
||||
int ael = sd::math::sd_abs<int,int>(el) - 1;
|
||||
|
||||
// TODO: investigate, if += would work better here, as in "decoded accumulation"
|
||||
z[ael] += el > 0 ? threshold : -threshold;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void decoderKernelGeneric(dim3 &launchDims, cudaStream_t *stream, const void *dx, LongType N, void *dz) {
|
||||
execDecoderKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, N, dz);
|
||||
DebugHelper::checkErrorCode(stream, "execDecoder(...) failed");
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void decoderKernelGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, const void *dx, sd::LongType N, void *dz),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void execCudaEncodeBitmapKernel(void *vdx, LongType N, int *dz, int *scalar, int *reductionBuffer,
|
||||
float threshold) {
|
||||
auto dx = reinterpret_cast<T *>(vdx);
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
T off(0.0f);
|
||||
__shared__ int counter;
|
||||
__shared__ int *shmem;
|
||||
__shared__ T *vals;
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ char mem[];
|
||||
shmem = reinterpret_cast<int *>(mem);
|
||||
vals = reinterpret_cast<T *>(shmem + blockDim.x);
|
||||
counter = 0;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
LongType loopRemainder = N % (blockDim.x * gridDim.x);
|
||||
LongType loopLimit = N + (blockDim.x * gridDim.x - loopRemainder);
|
||||
|
||||
for (LongType i = tid; i < loopLimit; i += blockDim.x * gridDim.x) {
|
||||
// all threads in block reading stuff
|
||||
T val = i < N ? dx[i] : off;
|
||||
T abs = math::sd_abs<T,T>(val);
|
||||
|
||||
int byteId = i / 16 + 4;
|
||||
int bitId = i % 16;
|
||||
|
||||
shmem[threadIdx.x] = 0;
|
||||
vals[threadIdx.x] = val;
|
||||
|
||||
if (abs >= static_cast<T>(threshold) && i < N) {
|
||||
shmem[threadIdx.x] = 1 << (bitId);
|
||||
atomicAdd(&counter, 1);
|
||||
if (val < static_cast<T>(0.0f)) {
|
||||
shmem[threadIdx.x] |= 1 << (bitId + 16);
|
||||
vals[threadIdx.x] += static_cast<T>(threshold);
|
||||
} else {
|
||||
vals[threadIdx.x] -= static_cast<T>(threshold);
|
||||
}
|
||||
} else if (abs >= static_cast<T>(threshold) / static_cast<T>(2.0f) && val < static_cast<T>(0.0f) && i < N) {
|
||||
atomicAdd(&counter, 1);
|
||||
shmem[threadIdx.x] = 1 << (bitId + 16);
|
||||
|
||||
vals[threadIdx.x] += static_cast<T>(threshold) / static_cast<T>(2.0f);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x % 16 == 0 && i < N) {
|
||||
int byte = 0;
|
||||
for (int e = 0; e < 16; e++) {
|
||||
if (i + e >= N) continue;
|
||||
|
||||
byte |= shmem[threadIdx.x + e];
|
||||
}
|
||||
dz[byteId] = byte;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (i < N) dx[i] = vals[threadIdx.x];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
atomicAdd(scalar, counter);
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void cudaEncodeBitmapGeneric(dim3 &launchDims, cudaStream_t *stream, void *vdx, LongType N, int *dz,
|
||||
int *scalar, int *reductionBuffer, float threshold) {
|
||||
execCudaEncodeBitmapKernel<T>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(vdx, N, dz, scalar, reductionBuffer, threshold);
|
||||
DebugHelper::checkErrorCode(stream, "encodeBitmap(...) failed");
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void cudaEncodeBitmapGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, void *vdx, sd::LongType N, int *dz, int *scalar,
|
||||
int *reductionBuffer, float threshold),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_KERNEL static void execCudaDecodeBitmapKernel(const void *dx, LongType N, void *vdz) {
|
||||
auto dz = static_cast<T *>(vdz);
|
||||
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
__shared__ T *shmem;
|
||||
__shared__ FloatBits fb;
|
||||
__shared__ float threshold;
|
||||
__shared__ const int *x;
|
||||
if (threadIdx.x == 0) {
|
||||
extern __shared__ char mem[];
|
||||
shmem = reinterpret_cast<T *>(mem);
|
||||
x = reinterpret_cast<const int *>(dx);
|
||||
fb.i_ = x[2];
|
||||
threshold = fb.f_;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int lim = N / 16 + 5;
|
||||
for (int i = tid; i < N; i += blockDim.x * gridDim.x) {
|
||||
int byteId = i / 16 + 4;
|
||||
shmem[threadIdx.x] = dz[i];
|
||||
__syncthreads();
|
||||
|
||||
if (threadIdx.x % 16 == 0) {
|
||||
int byte = x[byteId];
|
||||
|
||||
for (int e = 0; e < 16; e++) {
|
||||
if (i + e >= N) continue;
|
||||
|
||||
int bitId = (i + e) % 16;
|
||||
|
||||
bool hasBit = (byte & 1 << (bitId)) != 0;
|
||||
bool hasSign = (byte & 1 << (bitId + 16)) != 0;
|
||||
|
||||
if (hasBit) {
|
||||
if (hasSign)
|
||||
shmem[threadIdx.x + bitId] -= threshold;
|
||||
else
|
||||
shmem[threadIdx.x + bitId] += threshold;
|
||||
} else if (hasSign) {
|
||||
shmem[threadIdx.x + bitId] -= threshold / 2;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
dz[i] = shmem[threadIdx.x];
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
template <typename T>
|
||||
SD_HOST void cudaDecodeBitmapGeneric(dim3 &launchDims, cudaStream_t *stream, const void *dx, LongType N,
|
||||
void *vdz) {
|
||||
execCudaDecodeBitmapKernel<T><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(dx, N, vdz);
|
||||
DebugHelper::checkErrorCode(stream, "cudeDecodeBitmap(...) failed");
|
||||
}
|
||||
BUILD_SINGLE_TEMPLATE( void cudaDecodeBitmapGeneric,
|
||||
(dim3 & launchDims, cudaStream_t *stream, const void *dx, sd::LongType N, void *vdz),
|
||||
SD_FLOAT_TYPES);
|
||||
|
||||
template <bool storeSum, bool isNP2>
|
||||
SD_HOST void prescanLauncher(dim3 &blocks, dim3 &threads, int shmem, cudaStream_t *stream, int *g_odata,
|
||||
const int *g_idata, int *g_blockSums, int n, int blockIndex, int baseIndex) {
|
||||
shmem = sd::math::sd_max<int>(shmem, 16384);
|
||||
prescan<storeSum, isNP2>
|
||||
<<<blocks, threads, shmem, *stream>>>(g_odata, g_idata, g_blockSums, n, blockIndex, baseIndex);
|
||||
sd::DebugHelper::checkErrorCode(stream, "prescanLauncher failed");
|
||||
|
||||
};
|
||||
|
||||
template <typename S, typename T>
|
||||
SD_KERNEL void convertKernel(void *dx, LongType N, void *dz) {
|
||||
auto x = reinterpret_cast<S *>(dx);
|
||||
auto z = reinterpret_cast<T *>(dz);
|
||||
|
||||
sd::convertKernelGeneric(x, N, z);
|
||||
}
|
||||
|
||||
#define LIBND4J_BOOLS_LOCAL (randomName0, 0), (randomName1, 1)
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( void TypeCast::convertGenericCuda,
|
||||
(sd::Pointer * extras, void *dx, sd::LongType N, void *dz), SD_COMMON_TYPES,
|
||||
SD_COMMON_TYPES);
|
||||
BUILD_DOUBLE_TEMPLATE( void prescanLauncher,
|
||||
(dim3 & blocks, dim3 &threads, int shmem, cudaStream_t *stream, int *g_odata, const int *g_idata,
|
||||
int *g_blockSums, int n, int blockIndex, int baseIndex),
|
||||
LIBND4J_BOOLS_LOCAL, LIBND4J_BOOLS_LOCAL);
|
||||
|
||||
#undef LIBND4J_BOOLS_LOCAL
|
||||
} // namespace sd
|
||||
Reference in New Issue
Block a user