wehub-resource-sync
267 lines
11 KiB
Plaintext
267 lines
11 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 07.03.2019
|
|
//
|
|
|
|
#include <helpers/PointersManager.h>
|
|
#include <helpers/ShapeUtils.h>
|
|
#include <ops/declarable/helpers/gather.h>
|
|
|
|
#include <numeric>
|
|
|
|
#include "execution/cuda/LaunchDims.h"
|
|
#include "helpers/DebugHelper.h"
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
template <typename X, typename Y>
|
|
SD_KERNEL static void gatherCudaLinearKernel(const void* vx, const LongType* xShapeInfo, const void* vy,
|
|
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo) {
|
|
__shared__ const X* x;
|
|
__shared__ const Y* y;
|
|
__shared__ X* z;
|
|
__shared__ LongType xLen, yLen, zLen;
|
|
__shared__ LongType xRank, yRank, zRank;
|
|
__shared__ const LongType *xShapePtr, *xStridePtr;
|
|
__shared__ const LongType *yShapePtr, *yStridePtr;
|
|
__shared__ const LongType *zShapePtr, *zStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
x = reinterpret_cast<const X*>(vx);
|
|
z = reinterpret_cast<X*>(vz);
|
|
y = reinterpret_cast<const Y*>(vy);
|
|
|
|
xLen = shape::length(xShapeInfo);
|
|
yLen = shape::length(yShapeInfo);
|
|
zLen = shape::length(zShapeInfo);
|
|
|
|
xRank = shape::rank(xShapeInfo);
|
|
yRank = shape::rank(yShapeInfo);
|
|
zRank = shape::rank(zShapeInfo);
|
|
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
|
|
yShapePtr = shape::shapeOf(yShapeInfo);
|
|
yStridePtr = shape::stride(yShapeInfo);
|
|
|
|
zShapePtr = shape::shapeOf(zShapeInfo);
|
|
zStridePtr = shape::stride(zShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
auto start = blockIdx.x * blockDim.x + threadIdx.x;
|
|
auto step = blockDim.x * gridDim.x;
|
|
|
|
for (LongType j = start; j < zLen; j += step) {
|
|
LongType zIndex, yIndex, xIndex;
|
|
LongType zCoords[SD_MAX_RANK], yCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
|
|
|
// Compute z coordinates and offset
|
|
INDEX2COORDS(j, zRank, zShapePtr, zCoords);
|
|
COORDS2INDEX(zRank, zStridePtr, zCoords, zIndex);
|
|
|
|
// Compute y coordinates and offset
|
|
INDEX2COORDS(j, yRank, yShapePtr, yCoords);
|
|
COORDS2INDEX(yRank, yStridePtr, yCoords, yIndex);
|
|
|
|
// Compute x coordinates and offset
|
|
INDEX2COORDS(y[yIndex], xRank, xShapePtr, xCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, xCoords, xIndex);
|
|
|
|
// Assign value to z
|
|
z[zIndex] = x[xIndex];
|
|
}
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Y>
|
|
SD_KERNEL static void gatherCuda(const int numOfSubArrs, const void* vx, const LongType* xShapeInfo,
|
|
const LongType* xOffsets, const void* vy, const LongType* yShapeInfo, void* vz,
|
|
const LongType* zShapeInfo, const LongType* zOffsets) {
|
|
const Y* y = reinterpret_cast<const Y*>(vy);
|
|
|
|
__shared__ const X* x;
|
|
__shared__ X* z;
|
|
__shared__ LongType xLen, yRank, xRank, zRank;
|
|
__shared__ const LongType *xShapePtr, *xStridePtr, *yShapePtr, *yStridePtr, *zShapePtr, *zStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
xLen = shape::length(xShapeInfo);
|
|
yRank = shape::rank(yShapeInfo);
|
|
xRank = shape::rank(xShapeInfo);
|
|
zRank = shape::rank(zShapeInfo);
|
|
|
|
xShapePtr = shape::shapeOf(xShapeInfo);
|
|
xStridePtr = shape::stride(xShapeInfo);
|
|
yShapePtr = shape::shapeOf(yShapeInfo);
|
|
yStridePtr = shape::stride(yShapeInfo);
|
|
zShapePtr = shape::shapeOf(zShapeInfo);
|
|
zStridePtr = shape::stride(zShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
for (LongType i = blockIdx.x; i < numOfSubArrs; i += gridDim.x) {
|
|
if (threadIdx.x == 0) {
|
|
LongType yIndex, xOffset, zOffset;
|
|
LongType yCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK], zCoords[SD_MAX_RANK];
|
|
|
|
// Calculate y index
|
|
INDEX2COORDS(i, yRank, yShapePtr, yCoords);
|
|
COORDS2INDEX(yRank, yStridePtr, yCoords, yIndex);
|
|
|
|
// Calculate x offset
|
|
INDEX2COORDS(y[yIndex], xRank, xShapePtr, xCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
|
|
|
|
// Calculate z offset
|
|
INDEX2COORDS(i, zRank, zShapePtr, zCoords);
|
|
COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
|
|
|
|
x = reinterpret_cast<const X*>(vx) + xOffsets[xOffset];
|
|
z = reinterpret_cast<X*>(vz) + zOffsets[zOffset];
|
|
}
|
|
__syncthreads();
|
|
|
|
// Copy data from x to z
|
|
for (LongType j = threadIdx.x; j < xLen; j += blockDim.x) {
|
|
LongType zIndex, xIndex;
|
|
LongType zCoords[SD_MAX_RANK], xCoords[SD_MAX_RANK];
|
|
|
|
// Calculate z index
|
|
INDEX2COORDS(j, zRank, zShapePtr, zCoords);
|
|
COORDS2INDEX(zRank, zStridePtr, zCoords, zIndex);
|
|
|
|
// Calculate x index
|
|
INDEX2COORDS(j, xRank, xShapePtr, xCoords);
|
|
COORDS2INDEX(xRank, xStridePtr, xCoords, xIndex);
|
|
|
|
// Copy value
|
|
z[zIndex] = x[xIndex];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
}
|
|
|
|
|
|
template <typename X, typename Y>
|
|
SD_HOST static void gatherCudaLinear(const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
|
const void* vy, const LongType* yShapeInfo, void* vz,
|
|
const LongType* zShapeInfo) {
|
|
//note gather linear and gather are different kernels
|
|
dim3 gatherLinear = getLaunchDims("gather_linear");
|
|
gatherCudaLinearKernel<X, Y><<<gatherLinear.x, gatherLinear.y, gatherLinear.z, *stream>>>(vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo);
|
|
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"gatherCudaLinearKernel failed");
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
template <typename X, typename Y>
|
|
SD_HOST static void gatherCudaLauncher(const cudaStream_t* stream, const int numOfSubArrs, const void* vx,
|
|
const LongType* xShapeInfo, const LongType* xOffsets, const void* vy,
|
|
const LongType* yShapeInfo, void* vz, const LongType* zShapeInfo,
|
|
const LongType* zOffsets) {
|
|
dim3 gatherLinear = getGatherLinear(numOfSubArrs);
|
|
gatherCuda<X, Y><<<gatherLinear.y, gatherLinear.x, gatherLinear.z, *stream>>>(numOfSubArrs, vx, xShapeInfo, xOffsets, vy,
|
|
yShapeInfo, vz, zShapeInfo, zOffsets);
|
|
DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream),"gatherCudaLauncher failed");
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////
|
|
void gather(LaunchContext* context, NDArray* input, NDArray* indices, NDArray* output,
|
|
const std::vector<LongType>& intArgs) {
|
|
const LongType inputRank = input->rankOf();
|
|
const LongType numOfIntArgs = intArgs.size();
|
|
|
|
LongType axis = numOfIntArgs > 0 ? intArgs[0] : 0;
|
|
if (axis < 0) axis += inputRank;
|
|
|
|
if (indices == nullptr && numOfIntArgs == 2) { // scalar case
|
|
NDArray scalar = (*input)(intArgs[1], {axis});
|
|
output->assign(&scalar);
|
|
} else if (indices != nullptr && indices->isScalar()) {
|
|
if (input->rankOf() <= 1) { // For scalar indices, rank 0 or 1 input: can't do tensor along dimension 0 as this is
|
|
// whole array... instead, we want to get a scalar
|
|
auto idx = indices->e<LongType>(0);
|
|
auto scalarNDArray = input->e(idx);
|
|
output->assign(&scalarNDArray);
|
|
} else {
|
|
NDArray inSubArr = (*input)(indices->e<LongType>(0), {axis});
|
|
output->assign(&inSubArr);
|
|
}
|
|
} else {
|
|
NDArray* pIndices = const_cast<NDArray*>(indices);
|
|
if (indices == nullptr) {
|
|
std::vector<LongType> firstShape = {numOfIntArgs - 1};
|
|
std::vector<double> data = std::vector<double>(intArgs.begin() + 1, intArgs.end());
|
|
pIndices = new NDArray(input->ordering(),firstShape,
|
|
data, INT64, input->getContext());
|
|
}
|
|
std::vector<LongType> dimsOut(pIndices->rankOf());
|
|
std::iota(dimsOut.begin(), dimsOut.end(), axis); // fill with axis, axis+1, ... axis+pIndices->rankOf()-1
|
|
|
|
const LongType numOfSubArrs = pIndices->lengthOf();
|
|
|
|
LongType *outSubArrShapeInfo(nullptr), *inSubArrShapeInfo(nullptr), *outSubArrOffsets(nullptr),
|
|
*inSubArrOffsets(nullptr);
|
|
input->getSubArrShapeAndOffsets({axis}, inSubArrShapeInfo, inSubArrOffsets);
|
|
output->getSubArrShapeAndOffsets(dimsOut, outSubArrShapeInfo, outSubArrOffsets);
|
|
if (output->rankOf() > 1) {
|
|
PointersManager manager(context, "gather");
|
|
auto xShapeInfo = reinterpret_cast<LongType*>(
|
|
manager.replicatePointer(inSubArrShapeInfo, shape::shapeInfoByteLength(inSubArrShapeInfo)));
|
|
auto zShapeInfo = reinterpret_cast<LongType*>(
|
|
manager.replicatePointer(outSubArrShapeInfo, shape::shapeInfoByteLength(outSubArrShapeInfo)));
|
|
auto xOffsets = reinterpret_cast<LongType*>(manager.replicatePointer(
|
|
inSubArrOffsets, (input->lengthOf() / shape::length(inSubArrShapeInfo)) * sizeof(LongType)));
|
|
auto zOffsets = reinterpret_cast<LongType*>(manager.replicatePointer(
|
|
outSubArrOffsets, (output->lengthOf() / shape::length(outSubArrShapeInfo)) * sizeof(LongType)));
|
|
|
|
NDArray::prepareSpecialUse({output}, {input, pIndices});
|
|
BUILD_DOUBLE_SELECTOR(
|
|
input->dataType(), pIndices->dataType(), gatherCudaLauncher,
|
|
(context->getCudaStream(), numOfSubArrs, input->specialBuffer(), xShapeInfo, xOffsets,
|
|
pIndices->specialBuffer(), pIndices->specialShapeInfo(), output->specialBuffer(), zShapeInfo, zOffsets),
|
|
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, pIndices});
|
|
manager.synchronize();
|
|
} else {
|
|
NDArray::prepareSpecialUse({output}, {input, pIndices});
|
|
BUILD_DOUBLE_SELECTOR(
|
|
input->dataType(), pIndices->dataType(), gatherCudaLinear,
|
|
(context->getCudaStream(), input->specialBuffer(), input->specialShapeInfo(), pIndices->specialBuffer(),
|
|
pIndices->specialShapeInfo(), output->specialBuffer(), output->specialShapeInfo()),
|
|
SD_COMMON_TYPES, SD_INDEXING_TYPES);
|
|
NDArray::registerSpecialUse({output}, {input, pIndices});
|
|
}
|
|
|
|
if (indices == nullptr) delete pIndices;
|
|
}
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|