chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,202 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/DebugHelper.h>
|
||||
#include <loops/legacy_ops.h>
|
||||
#include <loops/transform_float.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <execution/cuda/DeviceValidator.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The cached kernel that caches shape info in shared memory and uses cached variables
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void transformFloatSimpleCached(
|
||||
const void* x,
|
||||
const sd::LongType* xShapeInfo,
|
||||
sd::LongType xRank,
|
||||
void* params,
|
||||
void* z,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Delegate the transform to the transformCuda method with cached shape info
|
||||
functions::transform::TransformFloat<X, Z>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the "executeTransformShaped" that launches the cached kernel
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void TransformFloat<X, Y>::executeTransformShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
int opNum,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
DISPATCH_BY_OPNUM_TT(
|
||||
intermediateShaped,
|
||||
PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
|
||||
reductionPointer, tadShapeInfo, tadOffsets),
|
||||
TRANSFORM_FLOAT_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformFloat executeTransformShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Device function that caches shape info and uses cached variables for computations
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformFloat<X, Z>::transformCuda(
|
||||
const void* vx,
|
||||
const sd::LongType* xShapeInfo,
|
||||
void* vparams,
|
||||
void* vz,
|
||||
const sd::LongType* zShapeInfo,
|
||||
sd::LongType* allocationPointer,
|
||||
void* vreductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Cast pointers to appropriate types
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto params = reinterpret_cast<Z*>(vparams);
|
||||
auto reductionPointer = reinterpret_cast<Z*>(vreductionPointer);
|
||||
|
||||
// Check for special operations
|
||||
if (OpType::requiresSpecial) {
|
||||
OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
|
||||
tadOffsets);
|
||||
return;
|
||||
}
|
||||
|
||||
// Shared memory for caching shape information
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRankCached;
|
||||
__shared__ const sd::LongType* xShapePtrCached;
|
||||
__shared__ const sd::LongType* xStridePtrCached;
|
||||
|
||||
__shared__ int zRankCached;
|
||||
__shared__ const sd::LongType* zShapePtrCached;
|
||||
__shared__ const sd::LongType* zStridePtrCached;
|
||||
|
||||
// Thread 0 caches the shape information
|
||||
if (threadIdx.x == 0) {
|
||||
length = shape::length(xShapeInfo);
|
||||
|
||||
xRankCached = shape::rank(xShapeInfo);
|
||||
xShapePtrCached = shape::shapeOf(xShapeInfo);
|
||||
xStridePtrCached = shape::stride(xShapeInfo);
|
||||
|
||||
zRankCached = shape::rank(zShapeInfo);
|
||||
zShapePtrCached = shape::shapeOf(zShapeInfo);
|
||||
zStridePtrCached = shape::stride(zShapeInfo);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Calculate thread ID and total threads
|
||||
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
// Loop over all elements using cached shape info
|
||||
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;
|
||||
|
||||
// Convert index to coordinates using cached shape info
|
||||
INDEX2COORDS(i, xRankCached, xShapePtrCached, xCoords);
|
||||
COORDS2INDEX(xRankCached, xStridePtrCached, xCoords, xOffset);
|
||||
|
||||
INDEX2COORDS(i, zRankCached, zShapePtrCached, zCoords);
|
||||
COORDS2INDEX(zRankCached, zStridePtrCached, zCoords, zOffset);
|
||||
|
||||
// Apply the operation using cached offsets
|
||||
z[zOffset] = OpType::op(x[xOffset], params);
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Host function that launches the cached kernel
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformFloat<X, Z>::intermediateShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const void* x,
|
||||
const sd::LongType* xShape,
|
||||
sd::LongType xRank,
|
||||
void* extraParams,
|
||||
void* z,
|
||||
const sd::LongType* zShape,
|
||||
sd::LongType zRank,
|
||||
sd::LongType* allocationPointer,
|
||||
void* reductionPointer,
|
||||
const sd::LongType* tadShapeInfo,
|
||||
const sd::LongType* tadOffsets)
|
||||
{
|
||||
// Launch the cached kernel
|
||||
transformFloatSimpleCached<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShape,
|
||||
xRank,
|
||||
extraParams,
|
||||
z,
|
||||
zShape,
|
||||
zRank,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets
|
||||
);
|
||||
|
||||
// Check for any errors during kernel execution
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformFloat(...) cached kernel failed");
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class TransformFloat, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
Reference in New Issue
Block a user