wehub-resource-sync
442 lines
14 KiB
Plaintext
442 lines
14 KiB
Plaintext
/* ******************************************************************************
|
|
*
|
|
*
|
|
* This program and the accompanying materials are made available under the
|
|
* terms of the Apache License, Version 2.0 which is available at
|
|
* https://www.apache.org/licenses/LICENSE-2.0.
|
|
*
|
|
* See the NOTICE file distributed with this work for additional
|
|
* information regarding copyright ownership.
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See
|
|
* the License for the specific language governing permissions and limitations
|
|
* under the License.
|
|
*
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
******************************************************************************/
|
|
|
|
//
|
|
// @author 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
|