wehub-resource-sync
126 lines
4.9 KiB
Plaintext
126 lines
4.9 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
|
|
******************************************************************************/
|
|
|
|
//
|
|
// Created by raver119 on 30.11.17.
|
|
//
|
|
#include <helpers/PointersManager.h>
|
|
#include <ops/declarable/helpers/im2col.h>
|
|
|
|
#include <execution/cuda/LaunchDims.h>
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
// input [bS, iC, iH, iW] is convoluted to output [bS, iC, kH, kW, oH, oW]
|
|
template <typename T>
|
|
SD_KERNEL static void im2colCuda(const void *image, void *columns, const LongType *imShapeInfo,
|
|
const LongType *colShapeInfo, const LongType sH, const LongType sW, const LongType pH,
|
|
const LongType pW, const LongType dH, const LongType dW, const double zeroPadValD) {
|
|
T zeroPadVal = static_cast<T>(zeroPadValD); // Value to use when value is padding
|
|
const auto im = reinterpret_cast<const T *>(image);
|
|
auto col = reinterpret_cast<T *>(columns);
|
|
|
|
// Shared memory caching
|
|
__shared__ LongType colLen, imLen, iH, iW;
|
|
__shared__ LongType imRank, colRank;
|
|
__shared__ const LongType *imShapePtr, *imStridePtr;
|
|
__shared__ const LongType *colShapePtr, *colStridePtr;
|
|
|
|
if (threadIdx.x == 0) {
|
|
colRank = 6;
|
|
imRank = 4;
|
|
|
|
colLen = shape::length(colShapeInfo);
|
|
imLen = shape::length(imShapeInfo);
|
|
|
|
iH = shape::shapeOf(imShapeInfo)[2];
|
|
iW = shape::shapeOf(imShapeInfo)[3];
|
|
|
|
imShapePtr = shape::shapeOf(imShapeInfo);
|
|
imStridePtr = shape::stride(imShapeInfo);
|
|
|
|
colShapePtr = shape::shapeOf(colShapeInfo);
|
|
colStridePtr = shape::stride(colShapeInfo);
|
|
}
|
|
__syncthreads();
|
|
|
|
const auto colInd = threadIdx.x + blockIdx.x * blockDim.x;
|
|
|
|
if (colInd >= colLen) return; // Boundary check for threads
|
|
|
|
LongType coords[SD_MAX_RANK];
|
|
|
|
// Calculate coordinates and offsets
|
|
INDEX2COORDS(colInd, colRank, colShapePtr, coords);
|
|
|
|
LongType colOffset;
|
|
COORDS2INDEX(colRank, colStridePtr, coords, colOffset);
|
|
|
|
coords[2] = (-pH + coords[2] * dH) + coords[4] * sH; // imH
|
|
coords[3] = (-pW + coords[3] * dW) + coords[5] * sW; // imW
|
|
|
|
// Check bounds and assign appropriate values
|
|
if (coords[2] >= iH || coords[3] >= iW || coords[2] < 0 || coords[3] < 0) {
|
|
if (colOffset < colLen)
|
|
col[colOffset] = zeroPadVal;
|
|
} else {
|
|
LongType imOffset;
|
|
COORDS2INDEX(imRank, imStridePtr, coords, imOffset);
|
|
if (imOffset < imLen && colOffset < colLen)
|
|
col[colOffset] = im[imOffset];
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void im2colCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMemory,
|
|
LaunchContext &context, const void *image, void *columns,
|
|
const LongType *imShapeInfo, const LongType *colShapeInfo, LongType sH,
|
|
LongType sW, LongType pH, LongType pW, LongType dH, LongType dW, double zeroPadVal) {
|
|
im2colCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMemory /* rank of columns = 6 */, *context.getCudaStream()>>>(
|
|
image, columns, imShapeInfo, colShapeInfo, sH, sW, pH, pW, dH, dW, zeroPadVal);
|
|
DebugHelper::checkErrorCode(context.getCudaStream(), "im2colCuda(...) failed");
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void im2col(LaunchContext &context, NDArray&image, NDArray &columns, const LongType kH, const LongType kW,
|
|
const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW,
|
|
NDArray&arrZeroPadVal) {
|
|
PointersManager manager(&context, "im2col");
|
|
|
|
dim3 im2colDevs = getim2ColLaunchParams(columns);
|
|
NDArray::prepareSpecialUse({&columns}, {&image});
|
|
BUILD_SINGLE_SELECTOR(
|
|
columns.dataType(), im2colCudaLauncher,
|
|
(im2colDevs.x, im2colDevs.y,im2colDevs.z, context, image.specialBuffer(), columns.specialBuffer(),
|
|
image.specialShapeInfo(), columns.specialShapeInfo(), sH, sW, pH, pW, dH, dW, arrZeroPadVal.e<double>(0)),
|
|
SD_FLOAT_TYPES);
|
|
NDArray::registerSpecialUse({&columns}, {&image});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|