wehub-resource-sync
389 lines
12 KiB
Plaintext
389 lines
12 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 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
|