wehub-resource-sync
130 lines
5.0 KiB
Plaintext
130 lines
5.0 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 sgazeos@gmail.com
|
|
//
|
|
#include <helpers/ConstantTadHelper.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
|
|
#include <legacy/NativeOps.h>
|
|
#include <ops/declarable/helpers/nth_element.h>
|
|
|
|
#include "array/NDArrayFactory.h"
|
|
#include "execution/cuda/LaunchDims.h"
|
|
#include "helpers/DebugHelper.h"
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
template <typename T>
|
|
static SD_KERNEL void fillUpElementKernel(void* outputBuffer, const LongType* outputShapeInfo, void* inputBuffer,
|
|
const LongType* inputShapeInfo, const LongType* pTadShape,
|
|
const LongType* pTadOffsets, LongType n) {
|
|
__shared__ LongType bufferLength;
|
|
__shared__ int rankOutput, rankTad;
|
|
__shared__ const LongType *shapeOutput, *strideOutput, *shapeTad, *strideTad;
|
|
|
|
auto z = reinterpret_cast<T*>(outputBuffer);
|
|
auto x = reinterpret_cast<T*>(inputBuffer);
|
|
|
|
if (threadIdx.x == 0) {
|
|
bufferLength = shape::length(outputShapeInfo);
|
|
rankOutput = shape::rank(outputShapeInfo);
|
|
rankTad = shape::rank(pTadShape);
|
|
shapeOutput = shape::shapeOf(outputShapeInfo);
|
|
strideOutput = shape::stride(outputShapeInfo);
|
|
shapeTad = shape::shapeOf(pTadShape);
|
|
strideTad = shape::stride(pTadShape);
|
|
}
|
|
__syncthreads();
|
|
|
|
const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
const auto step = gridDim.x * blockDim.x;
|
|
|
|
LongType zCoords[SD_MAX_RANK];
|
|
LongType xCoords[SD_MAX_RANK];
|
|
|
|
for (LongType t = tid; t < bufferLength; t += step) {
|
|
// Compute output coordinates and offset
|
|
INDEX2COORDS(t, rankOutput, shapeOutput, zCoords);
|
|
LongType zOffset;
|
|
COORDS2INDEX(rankOutput, strideOutput, zCoords, zOffset);
|
|
|
|
// Compute input coordinates and offset
|
|
INDEX2COORDS(n, rankTad, shapeTad, xCoords);
|
|
LongType xOffset;
|
|
COORDS2INDEX(rankTad, strideTad, xCoords, xOffset);
|
|
|
|
// Access and assign the value
|
|
z[zOffset] = x[pTadOffsets[t] + xOffset];
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
void nthElementFunctor_(LaunchContext* context, NDArray* input, LongType n, NDArray* output, bool reverse) {
|
|
NDArray::prepareSpecialUse({output}, {input});
|
|
NDArray sortedVals(*input);
|
|
Pointer params[2];
|
|
params[0] = context;
|
|
params[1] = context->getCudaStream();
|
|
// Nth element in sorted sequence : basic algorithm sort and retrieve nth element in sorted
|
|
if (input->isVector()) {
|
|
sort(params, &sortedVals, reverse);
|
|
|
|
cudaMemcpy(reinterpret_cast<T*>(output->specialBuffer()), reinterpret_cast<T*>(sortedVals.specialBuffer()) + n,
|
|
sizeof(T), cudaMemcpyDeviceToDevice);
|
|
} else { // rank greater than 1
|
|
std::vector<LongType> lastDims(
|
|
{input->rankOf() - 1});
|
|
NDArray *dimData = NDArrayFactory::create_<LongType>('c',{2},lastDims, context);
|
|
auto packX = ConstantTadHelper::getInstance().tadForDimensions(sortedVals.shapeInfo(), &lastDims);
|
|
|
|
auto pTadShape = packX->specialShapeInfo();
|
|
auto pTadShapeH = packX->primaryShapeInfo();
|
|
auto pTadOffsets = packX->specialOffsets();
|
|
sortTad(params, &sortedVals,
|
|
reinterpret_cast<sd::LongType *>(lastDims.data()),
|
|
lastDims.size(),
|
|
const_cast<sd::LongType *>(pTadShape),
|
|
const_cast<sd::LongType *>(pTadOffsets),
|
|
reverse);
|
|
sortedVals.tickWriteDevice();
|
|
sortedVals.syncToHost();
|
|
auto stream = context->getCudaStream();
|
|
dim3 launchDims = getLaunchDims("nth_element_fill");
|
|
fillUpElementKernel<T><<<launchDims.y, launchDims.x, launchDims.z, *stream>>>(output->specialBuffer(), output->specialShapeInfo(),
|
|
sortedVals.specialBuffer(), sortedVals.specialShapeInfo(),
|
|
pTadShape, pTadOffsets, n);
|
|
sd::DebugHelper::checkErrorCode(stream, "fillUpElementKernel failed");
|
|
|
|
}
|
|
NDArray::registerSpecialUse({output}, {input});
|
|
}
|
|
void nthElementFunctor(LaunchContext* context, NDArray* input, LongType n, NDArray* output, bool reverse) {
|
|
auto inputDType = input->dataType();
|
|
|
|
BUILD_SINGLE_SELECTOR(input->dataType(), nthElementFunctor_, (context, input, n, output, reverse), SD_COMMON_TYPES);
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|