wehub-resource-sync
107 lines
3.0 KiB
C++
107 lines
3.0 KiB
C++
/*
|
|
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 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.
|
|
*/
|
|
|
|
#ifndef _REDUCED_MATH_PLUGIN_H
|
|
#define _REDUCED_MATH_PLUGIN_H
|
|
#include <cstdint>
|
|
// Dynamically strength-reduced div and mod
|
|
//
|
|
// Ideas taken from Sean Baxter's MGPU library.
|
|
// These classes provide for reduced complexity division and modulus
|
|
// on integers, for the case where the same divisor or modulus will
|
|
// be used repeatedly.
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
namespace detail
|
|
{
|
|
|
|
void findDivisor(int32_t denom, uint32_t& mul_coeff, uint32_t& shift_coeff);
|
|
|
|
__host__ __device__ __forceinline__ uint32_t umulhi(uint32_t x, uint32_t y)
|
|
{
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 100
|
|
return __umulhi(x, y);
|
|
#else
|
|
uint64_t z = (uint64_t) x * (uint64_t) y;
|
|
return (uint32_t) (z >> 32);
|
|
#endif
|
|
}
|
|
|
|
// This is a weird implementation that returns div_up(0,1)=0 but
|
|
// div_up(0,2)=1 (wrong) -- just do not use it with a=0.
|
|
__host__ __device__ inline int32_t div_up(int32_t a, int32_t b)
|
|
{
|
|
return (a - 1) / b + 1;
|
|
}
|
|
|
|
} // end namespace detail
|
|
|
|
class ReducedDivisor
|
|
{
|
|
public:
|
|
ReducedDivisor() {}
|
|
__host__ __forceinline__ ReducedDivisor(int32_t _y)
|
|
: y(_y)
|
|
{
|
|
detail::findDivisor(y, mul_coeff, shift_coeff);
|
|
}
|
|
__host__ __device__ __forceinline__ ReducedDivisor(uint32_t _mul_coeff, uint32_t _shift_coeff, int32_t _y)
|
|
: mul_coeff(_mul_coeff)
|
|
, shift_coeff(_shift_coeff)
|
|
, y(_y)
|
|
{
|
|
}
|
|
__host__ __device__ __forceinline__ int32_t div(int32_t x) const
|
|
{
|
|
// if dividing by 1, then findDivisor wouldn't have worked because
|
|
// mul_coeff would have had to be 2^32, which can't be represented,
|
|
// so we have to special case that one.
|
|
return (y != 1) ? detail::umulhi((uint32_t) x, mul_coeff) >> shift_coeff : x;
|
|
}
|
|
__host__ __device__ __forceinline__ int32_t mod(int32_t x) const
|
|
{
|
|
return x - (div(x) * y);
|
|
}
|
|
__host__ __device__ __forceinline__ void divmod(int32_t x, int32_t& q, int32_t& mod) const
|
|
{
|
|
q = div(x);
|
|
mod = x - (q * y);
|
|
}
|
|
__host__ __device__ __forceinline__ int32_t get() const
|
|
{
|
|
return y;
|
|
}
|
|
inline __host__ void get_mul_shift(uint32_t& mul, uint32_t& shift)
|
|
{
|
|
mul = mul_coeff;
|
|
shift = shift_coeff;
|
|
}
|
|
|
|
protected:
|
|
uint32_t mul_coeff{};
|
|
uint32_t shift_coeff{};
|
|
int32_t y{};
|
|
};
|
|
|
|
} // namespace plugin
|
|
|
|
} // namespace nvinfer1
|
|
#endif /*_REDUCED_MATH_PLUGIN_H*/
|