wehub-resource-sync
477 lines
14 KiB
Plaintext
477 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 <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
|
|
|
|
|
|
|
|
|