chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,192 @@
|
||||
/******************************************************************************
|
||||
*
|
||||
* 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_any.h>
|
||||
#include <system/Environment.h>
|
||||
#include <system/op_boilerplate.h>
|
||||
#include <types/types.h>
|
||||
#include <execution/cuda/DeviceValidator.h>
|
||||
|
||||
using namespace simdOps;
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The kernel that calls the transform CUDA method,
|
||||
// caching shape info in shared memory for offset computations.
|
||||
template <typename X, typename Z, typename OpType>
|
||||
__global__ void transformAnySimpleCached(
|
||||
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)
|
||||
{
|
||||
// Just delegate to transformCuda,
|
||||
// which will do the shape caching logic for coords->offset conversions.
|
||||
functions::transform::TransformAny<X, Z>::template transformCuda<OpType>(
|
||||
x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
|
||||
}
|
||||
|
||||
namespace functions {
|
||||
namespace transform {
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// Implementation of the "executeTransformShaped" that calls the new cached kernel
|
||||
template <typename X, typename Y>
|
||||
SD_HOST void TransformAny<X, Y>::executeTransformShaped(
|
||||
dim3 launchDims,
|
||||
cudaStream_t* stream,
|
||||
const 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_ANY_OPS);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformAny executeTransformShaped(...) failed");
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
// The transformCuda method that uses shared memory for shape/stride caching,
|
||||
// then does coords->offset conversions.
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_DEVICE void TransformAny<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
|
||||
auto x = reinterpret_cast<const X*>(vx);
|
||||
auto z = reinterpret_cast<Z*>(vz);
|
||||
auto params = reinterpret_cast<X*>(vparams);
|
||||
|
||||
if (x == nullptr || z == nullptr) return;
|
||||
|
||||
// cache shape info in shared memory
|
||||
__shared__ sd::LongType length;
|
||||
__shared__ int xRank;
|
||||
__shared__ const sd::LongType* xShapePtr;
|
||||
__shared__ const sd::LongType* xStridePtr;
|
||||
|
||||
__shared__ int 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);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// do the transform
|
||||
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const auto totalThreads = gridDim.x * blockDim.x;
|
||||
|
||||
for (sd::LongType i = tid; i < length; i += totalThreads) {
|
||||
sd::LongType coordsX[SD_MAX_RANK];
|
||||
sd::LongType coordsZ[SD_MAX_RANK];
|
||||
sd::LongType offsetX;
|
||||
sd::LongType offsetZ;
|
||||
|
||||
// convert i -> coords -> offset for x
|
||||
INDEX2COORDS(i, xRank, xShapePtr, coordsX);
|
||||
COORDS2INDEX(xRank, xStridePtr, coordsX, offsetX);
|
||||
|
||||
// convert i -> coords -> offset for z
|
||||
INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
|
||||
COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
|
||||
|
||||
z[offsetZ] = OpType::op(x[offsetX], params);
|
||||
}
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
template <typename X, typename Z>
|
||||
template <typename OpType>
|
||||
SD_HOST void TransformAny<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)
|
||||
{
|
||||
// We call the new transformAnySimpleCached kernel
|
||||
transformAnySimpleCached<X, Z, OpType>
|
||||
<<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
|
||||
x,
|
||||
xShape,
|
||||
xRank,
|
||||
extraParams,
|
||||
z,
|
||||
zShape,
|
||||
zRank,
|
||||
allocationPointer,
|
||||
reductionPointer,
|
||||
tadShapeInfo,
|
||||
tadOffsets);
|
||||
|
||||
sd::DebugHelper::checkErrorCode(stream, "transformAny(...) cached kernel failed");
|
||||
}
|
||||
|
||||
BUILD_DOUBLE_TEMPLATE( class TransformAny, , SD_COMMON_TYPES, SD_COMMON_TYPES);
|
||||
|
||||
} // namespace transform
|
||||
} // namespace functions
|
||||
Reference in New Issue
Block a user