wehub-resource-sync
280 lines
10 KiB
Plaintext
280 lines
10 KiB
Plaintext
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* 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.
|
|
*
|
|
* ************************************************************************
|
|
* Modified from Pytorch
|
|
* Copyright (c) 2016-present, Facebook, Inc.
|
|
*
|
|
* See https://github.com/pytorch/pytorch/blob/main/LICENSE for details
|
|
* ************************************************************************
|
|
* Modified from ONNX Runtime
|
|
* Copyright (c) Microsoft Corporation
|
|
*
|
|
* See https://github.com/microsoft/onnxruntime/blob/main/LICENSE for details
|
|
* ************************************************************************
|
|
*/
|
|
|
|
|
|
#include <cuda.h>
|
|
#include <cuda_fp16.h>
|
|
#include "common/common.cuh"
|
|
#include "roiAlignKernel.h"
|
|
|
|
using half = __half;
|
|
|
|
__device__ half floatMax(half a, half b)
|
|
{
|
|
#if __CUDA_ARCH__ >= 800
|
|
return __hmax(a, b);
|
|
#else
|
|
return __float2half(max(__half2float(a), __half2float(b)));
|
|
#endif
|
|
}
|
|
|
|
__device__ float floatMax(float a, float b)
|
|
{
|
|
return max(a, b);
|
|
}
|
|
|
|
template <typename T>
|
|
__device__ T bilinearInterpolate(T const* bottomData, int32_t const height, int32_t const width, T y, T x,
|
|
int32_t const isModeAvg, int32_t const index /* index for debug only*/)
|
|
{
|
|
// deal with cases that inverse elements are out of feature map boundary
|
|
if (y < static_cast<T>(-1.0) || y > static_cast<T>(height) || x < static_cast<T>(-1.0) || x > static_cast<T>(width))
|
|
{
|
|
// empty
|
|
return 0;
|
|
}
|
|
|
|
if (y <= static_cast<T>(0))
|
|
{
|
|
y = 0;
|
|
}
|
|
if (x <= static_cast<T>(0))
|
|
{
|
|
x = 0;
|
|
}
|
|
|
|
int32_t yLow = static_cast<int32_t>(y);
|
|
int32_t xLow = static_cast<int32_t>(x);
|
|
int32_t yHigh;
|
|
int32_t xHigh;
|
|
|
|
if (yLow >= height - 1)
|
|
{
|
|
yHigh = yLow = height - 1;
|
|
y = static_cast<T>(yLow);
|
|
}
|
|
else
|
|
{
|
|
yHigh = yLow + 1;
|
|
}
|
|
|
|
if (xLow >= width - 1)
|
|
{
|
|
xHigh = xLow = width - 1;
|
|
x = static_cast<T>(xLow);
|
|
}
|
|
else
|
|
{
|
|
xHigh = xLow + 1;
|
|
}
|
|
|
|
T ly = y - static_cast<T>(yLow);
|
|
T lx = x - static_cast<T>(xLow);
|
|
T hy = static_cast<T>(1.) - ly, hx = static_cast<T>(1.) - lx;
|
|
// do bilinear interpolation
|
|
T v1 = bottomData[yLow * width + xLow];
|
|
T v2 = bottomData[yLow * width + xHigh];
|
|
T v3 = bottomData[yHigh * width + xLow];
|
|
T v4 = bottomData[yHigh * width + xHigh];
|
|
T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
|
|
|
|
T val;
|
|
if (isModeAvg)
|
|
{
|
|
val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); // mode Avg
|
|
}
|
|
else
|
|
{
|
|
val = floatMax(floatMax(floatMax(w1 * v1, w2 * v2), w3 * v3), w4 * v4); // mode Max
|
|
}
|
|
|
|
return val;
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void RoIAlignForward(int32_t const nthreads, T const* bottomData, T const spatialScale, int32_t const channels,
|
|
int32_t const height, int32_t const width, int32_t const pooledHeight, int32_t const pooledWidth, int32_t const samplingRatio,
|
|
T const* bottomRois, T* topData, int32_t const isModeAvg, int32_t const* batchIndicesPtr,
|
|
int32_t const aligned)
|
|
{
|
|
for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x)
|
|
{
|
|
// (n, c, ph, pw) is an element in the pooled output
|
|
int32_t pw = index % pooledWidth;
|
|
int32_t ph = (index / pooledWidth) % pooledHeight;
|
|
int32_t c = (index / pooledWidth / pooledHeight) % channels;
|
|
int32_t n = index / pooledWidth / pooledHeight / channels;
|
|
|
|
T const* offsetBottomRois = bottomRois + n * 4;
|
|
auto const roiBatchInd = batchIndicesPtr[n];
|
|
|
|
bool continuousCoordinate = aligned;
|
|
// Do not using rounding; this implementation detail is critical
|
|
T roiOffset = static_cast<T>(continuousCoordinate ? 0.5 : 0);
|
|
T roiStartW = offsetBottomRois[0] * spatialScale - roiOffset;
|
|
T roiStartH = offsetBottomRois[1] * spatialScale - roiOffset;
|
|
T roiEndW = offsetBottomRois[2] * spatialScale - roiOffset;
|
|
T roiEndH = offsetBottomRois[3] * spatialScale - roiOffset;
|
|
|
|
T roiWidth = roiEndW - roiStartW;
|
|
T roiHeight = roiEndH - roiStartH;
|
|
if (!continuousCoordinate)
|
|
{ // backward compatibility
|
|
// Force malformed ROIs to be 1x1
|
|
roiWidth = floatMax(roiWidth, static_cast<T>(1.));
|
|
roiHeight = floatMax(roiHeight, static_cast<T>(1.));
|
|
}
|
|
T binSizeH = static_cast<T>(roiHeight) / static_cast<T>(pooledHeight);
|
|
T binSizeW = static_cast<T>(roiWidth) / static_cast<T>(pooledWidth);
|
|
|
|
T const* offsetBottomData = bottomData + static_cast<int32_t>((roiBatchInd * channels + c) * height * width);
|
|
|
|
// We use roiBinGrid to sample the grid and mimic integral
|
|
int32_t roiBinGridH;
|
|
if (samplingRatio > 0)
|
|
{
|
|
roiBinGridH = samplingRatio;
|
|
}
|
|
else
|
|
{
|
|
roiBinGridH = ceilf(roiHeight / static_cast<T>(pooledHeight));
|
|
}
|
|
|
|
int32_t roiBinGridW;
|
|
if (samplingRatio > 0)
|
|
{
|
|
roiBinGridW = samplingRatio;
|
|
}
|
|
else
|
|
{
|
|
roiBinGridW = ceilf(roiWidth / static_cast<T>(pooledWidth));
|
|
}
|
|
// We do average (integral) pooling inside a bin
|
|
T const count = roiBinGridH * roiBinGridW; // e.g. = 4
|
|
|
|
T const yOff = roiStartH + static_cast<T>(ph) * binSizeH;
|
|
T const yFac = binSizeH / static_cast<T>(roiBinGridH);
|
|
|
|
T const xOff = roiStartW + static_cast<T>(pw) * binSizeW;
|
|
T const xFac = binSizeW / static_cast<T>(roiBinGridW);
|
|
|
|
T outputVal = 0.;
|
|
bool maxFlag = false;
|
|
for (int32_t iy = 0; iy < roiBinGridH; iy++) // e.g., iy = 0, 1
|
|
{
|
|
T const y = yOff + static_cast<T>(iy + .5F) * yFac; // e.g., 0.5, 1.5
|
|
for (int32_t ix = 0; ix < roiBinGridW; ix++)
|
|
{
|
|
T const x = xOff + static_cast<T>(ix + .5F) * xFac;
|
|
|
|
T val = bilinearInterpolate(offsetBottomData, height, width, y, x, isModeAvg, index);
|
|
|
|
if (isModeAvg)
|
|
{
|
|
outputVal += val;
|
|
}
|
|
else
|
|
{
|
|
if (!maxFlag)
|
|
{
|
|
outputVal = val;
|
|
maxFlag = true;
|
|
}
|
|
else
|
|
{
|
|
outputVal = floatMax(outputVal, val);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (isModeAvg)
|
|
{
|
|
outputVal = outputVal / count;
|
|
}
|
|
|
|
topData[index] = outputVal;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
cudaError_t RoiAlignImpl(cudaStream_t stream, int32_t const maxThreadsPerBlock, T const* bottomData, T const spatialScale,
|
|
int32_t const numRois, int32_t const channels, int32_t const height, int32_t const width, int32_t const pooledHeight,
|
|
int32_t const pooledWidth, int32_t const samplingRatio, T const* bottomRois, T* topData, int32_t const isModeAvg,
|
|
int32_t const* batchIndicesPtr, int32_t const aligned)
|
|
{
|
|
PLUGIN_ASSERT(bottomData != nullptr);
|
|
PLUGIN_ASSERT(bottomRois != nullptr);
|
|
PLUGIN_ASSERT(batchIndicesPtr != nullptr);
|
|
PLUGIN_ASSERT(topData != nullptr);
|
|
|
|
PLUGIN_ASSERT(numRois >= 0);
|
|
PLUGIN_ASSERT(maxThreadsPerBlock > 0);
|
|
|
|
PLUGIN_ASSERT(height > 0);
|
|
PLUGIN_ASSERT(width > 0);
|
|
PLUGIN_ASSERT(pooledHeight > 0);
|
|
PLUGIN_ASSERT(pooledWidth > 0);
|
|
PLUGIN_ASSERT(samplingRatio >= 0);
|
|
PLUGIN_ASSERT(isModeAvg == 0 || isModeAvg == 1);
|
|
PLUGIN_ASSERT(static_cast<float>(spatialScale) > 0.0F);
|
|
PLUGIN_ASSERT(aligned == 0 || aligned == 1);
|
|
|
|
int32_t const outputSize = numRois * channels * pooledHeight * pooledWidth;
|
|
|
|
int32_t blocksPerGrid = static_cast<int32_t>(ceil(static_cast<float>(outputSize)
|
|
/ maxThreadsPerBlock));
|
|
|
|
RoIAlignForward<T><<<blocksPerGrid, maxThreadsPerBlock, 0, stream>>>(outputSize,// nthreads
|
|
bottomData, // bottomData
|
|
spatialScale, // spatialScale
|
|
channels, // channels
|
|
height, // height
|
|
width, // width
|
|
pooledHeight, // pooledHeight
|
|
pooledWidth, // pooledWidth
|
|
samplingRatio, // samplingRatio
|
|
bottomRois, // bottomRois
|
|
topData, // topData
|
|
isModeAvg, // isModeAvg
|
|
batchIndicesPtr, // batchIndicesPtr
|
|
aligned);
|
|
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
#define SPECIALIZED_IMPL(T) \
|
|
template cudaError_t RoiAlignImpl<T>(cudaStream_t stream, int32_t const maxThreadsPerBlock, T const* bottomData, \
|
|
T const spatialScale, int32_t const numRois, int32_t const channels, int32_t const height, \
|
|
int32_t const width, int32_t const pooledHeight, int32_t const pooledWidth, int32_t const samplingRatio, \
|
|
T const* bottomRois, T* topData, int32_t const isModeAvg, int32_t const* batchIndicesPtr, \
|
|
int32_t const aligned);
|
|
|
|
SPECIALIZED_IMPL(float)
|
|
SPECIALIZED_IMPL(half)
|