wehub-resource-sync
305 lines
9.4 KiB
Plaintext
305 lines
9.4 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
|
|
//
|
|
|
|
#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
|