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