chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,26 @@
|
||||
#
|
||||
# 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.
|
||||
#
|
||||
|
||||
add_plugin_source(
|
||||
instanceNormalizationPlugin.cu
|
||||
instanceNormalizationPlugin.h
|
||||
instanceNormalizationPluginLegacy.cu
|
||||
instanceNormalizationPluginLegacy.h
|
||||
instanceNormCommon.h
|
||||
instanceNormFwd.h
|
||||
instanceNormFwdImpl.cu
|
||||
)
|
||||
@@ -0,0 +1,67 @@
|
||||
# InstanceNormalizationPlugin
|
||||
|
||||
**Table Of Contents**
|
||||
- [Description](#description)
|
||||
* [Structure](#structure)
|
||||
- [Parameters](#parameters)
|
||||
- [Additional resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
## Description
|
||||
|
||||
> NOTE: `InstanceNormalization_TRT` version 1 is deprecated since TensorRT 10.3. `InstanceNormalization_TRT` version 2 is deprecated since TensorRT 10.12. Please use `InstanceNormalization_TRT` version 3. Alternatively, the native `INormalizationLayer` can also be used as appropriate to replace their functionality.
|
||||
|
||||
The `InstanceNormalizePlugin` is used for the InstanceNormalization layer, which is generally used in deep learning models that perform image generation. This plugin is based off the [ONNX opset 6 definition](https://github.com/onnx/onnx/blob/master/docs/Operators.md#InstanceNormalization), and is used in any ONNX model that uses this operation.
|
||||
|
||||
Specifically, given an array of values `x = [x_0, x_1, ..., x_n]` , a scale factor, a bias factor, and an epsilon, the InstanceNormalization of x is `scale * (x-mean) / sqrt(variance + epsilon) + bias` where the mean and variance are computed per instance per channel.
|
||||
|
||||
### Structure
|
||||
|
||||
This plugin takes one input and generates one output. The first input is the data from the last layer that is going to be normalized. It has a shape of `[N, C, H, W]`, where `N` is the batch size, `C` is the number of channels, `H` is the height, `W` is the width.
|
||||
|
||||
The dimensions of the output are exactly the same as the input.
|
||||
|
||||
## Parameters
|
||||
|
||||
This plugin consists of the plugin creator class `InstanceNormalizationPluginCreator` and the plugin class `InstanceNormalizationPlugin`. To create the plugin instance, the following parameters are used:
|
||||
|
||||
| Type | Parameter | Description
|
||||
|------------|--------------------------|--------------------------------------------------------
|
||||
|`float` |`epsilon` |A small number to prevent being divided by zero during normalization.
|
||||
|`Weights *` |`scale` |A pointer to weights which contains information about scale factors for normalization. The definition of `Weights` can be found in the [NvInfer.h](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/_nv_infer_8h_source.html) header.
|
||||
|`Weights *` |`bias` |A pointer to weights which contains information about the bias values for normalization. The definition of `Weights` can be found in the [NvInfer.h](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/_nv_infer_8h_source.html) header.
|
||||
|`int` |`relu` |A value used to enable leaky relu activation
|
||||
|`float` |`alpha` |A small negative slope for the leaky relu activation
|
||||
|
||||
|
||||
## Additional resources
|
||||
|
||||
The following resources provide a deeper understanding of the `InstanceNormalizationPlugin` plugin:
|
||||
|
||||
**Networks**
|
||||
- [ONNX Operator Definition](https://github.com/onnx/onnx/blob/master/docs/Operators.md#InstanceNormalization)
|
||||
- [Instance Normalization Paper](https://arxiv.org/abs/1607.08022)
|
||||
|
||||
## License
|
||||
|
||||
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
|
||||
documentation.
|
||||
|
||||
|
||||
## Changelog
|
||||
|
||||
May 2025
|
||||
Deprecated version 2 of this plugin.
|
||||
|
||||
July 2024
|
||||
Deprecated version 1 of this plugin.
|
||||
|
||||
September 2019
|
||||
This is the first release of this `README.md` file.
|
||||
|
||||
|
||||
## Known issues
|
||||
|
||||
There are no known issues in this plugin.
|
||||
@@ -0,0 +1,803 @@
|
||||
/*
|
||||
* 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 INSTANCE_NORM_COMMON_H
|
||||
#define INSTANCE_NORM_COMMON_H
|
||||
|
||||
#include "common/plugin.h"
|
||||
#include <stdint.h>
|
||||
using namespace nvinfer1::pluginInternal;
|
||||
|
||||
#define DEVICE_FUNCTION static inline __device__
|
||||
|
||||
template <typename T, int32_t ELEMENTS_PER_LDG>
|
||||
struct PackedStorage
|
||||
{
|
||||
enum
|
||||
{
|
||||
PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG
|
||||
};
|
||||
typedef T Type;
|
||||
};
|
||||
|
||||
template <int32_t ELEMENTS_PER_LDG>
|
||||
struct PackedStorage<uint16_t, ELEMENTS_PER_LDG>
|
||||
{
|
||||
enum
|
||||
{
|
||||
PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG / 2
|
||||
};
|
||||
typedef int32_t Type;
|
||||
};
|
||||
|
||||
template <int32_t ELEMENTS_PER_LDG>
|
||||
struct PackedStorage<int8_t, ELEMENTS_PER_LDG>
|
||||
{
|
||||
enum
|
||||
{
|
||||
PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG / 4
|
||||
};
|
||||
typedef int32_t Type;
|
||||
};
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void fromFloat(int32_t (&dst)[N], float const (&src)[2 * N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
uint16_t lo, hi;
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(lo) : "f"(src[2 * i + 0]));
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(hi) : "f"(src[2 * i + 1]));
|
||||
asm volatile("mov.b32 %0, {%1, %2};" : "=r"(dst[i]) : "h"(lo), "h"(hi));
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void fromFloat(int32_t (&dst)[N], float const (&src)[4 * N], float scale)
|
||||
{
|
||||
union Pack_t
|
||||
{
|
||||
int8_t x[4];
|
||||
int32_t val;
|
||||
};
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
Pack_t packed;
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < 4; ii++)
|
||||
{
|
||||
packed.x[ii] = __float_as_int(min(max(src[4 * i + ii] * scale + 12582912.0F, 12582785.0F), 12583039.0F));
|
||||
}
|
||||
dst[i] = packed.val;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void fromFloat(float (&dst)[N], float const (&src)[N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
dst[i] = src[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N, bool DO_SCALE = false>
|
||||
DEVICE_FUNCTION void toFloat(float (&dst)[2 * N], int32_t (&src)[N], float scale = 1.f)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
uint16_t lo, hi;
|
||||
asm volatile("mov.b32 {%0, %1}, %2;" : "=h"(lo), "=h"(hi) : "r"(src[i]));
|
||||
asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 0]) : "h"(lo));
|
||||
asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 1]) : "h"(hi));
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N, bool DO_SCALE = false>
|
||||
DEVICE_FUNCTION void toFloat(float (&dst)[4 * N], int32_t (&src)[N], float scale = 1.f)
|
||||
{
|
||||
union Pack_t
|
||||
{
|
||||
int8_t x[4];
|
||||
int32_t val;
|
||||
};
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
Pack_t packed;
|
||||
packed.val = src[i];
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < 4; ++ii)
|
||||
{
|
||||
dst[4 * i + ii]
|
||||
= (DO_SCALE) ? __int2float_rn((int32_t) packed.x[ii]) * scale : __int2float_rn((int32_t) packed.x[ii]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N, bool DO_SCALE = false>
|
||||
DEVICE_FUNCTION void toFloat(float (&dst)[N], float (&src)[N], float scale = 1.f)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
dst[i] = (DO_SCALE) ? src[i] * scale : src[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void ldg(int32_t (&dst)[1], T const* gmem)
|
||||
{
|
||||
dst[0] = __ldg((int32_t const*) gmem);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void ldgStream(int32_t (&dst)[1], T const* gmem)
|
||||
{
|
||||
uint32_t tmp;
|
||||
asm volatile("ld.global.cs.nc.s32 %0, [%1];" : "=r"(tmp) : "l"((uint32_t const*) gmem));
|
||||
dst[0] = tmp;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void ldg(int32_t (&dst)[2], T const* gmem)
|
||||
{
|
||||
int2 tmp = __ldg((int2 const*) gmem);
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void ldgStream(int32_t (&dst)[2], T const* gmem)
|
||||
{
|
||||
int2 tmp;
|
||||
asm volatile("ld.global.cs.nc.v2.s32 {%0,%1}, [%2];" : "=r"(tmp.x), "=r"(tmp.y) : "l"((int2 const*) gmem));
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void ldg(int32_t (&dst)[2], uint16_t const* gmem)
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 320
|
||||
int2 tmp = __ldg((int2 const*) gmem);
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
#endif
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void ldgStream(int32_t (&dst)[2], uint16_t const* gmem)
|
||||
{
|
||||
int2 tmp;
|
||||
asm volatile("ld.global.cs.nc.v2.s32 {%0,%1}, [%2];" : "=r"(tmp.x), "=r"(tmp.y) : "l"((int2 const*) gmem));
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void ldg(float (&dst)[N], uint16_t const* gmem)
|
||||
{
|
||||
int32_t tmp[N / 2];
|
||||
ldg(tmp, gmem);
|
||||
toFloat(dst, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void ldgStream(float (&dst)[N], uint16_t const* gmem)
|
||||
{
|
||||
int32_t tmp[N / 2];
|
||||
ldgStream(tmp, gmem);
|
||||
toFloat(dst, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void ldg(float (&dst)[N], int8_t const* gmem)
|
||||
{
|
||||
int32_t tmp[N / 4];
|
||||
ldg(tmp, gmem);
|
||||
toFloat(dst, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void ldgStream(float (&dst)[N], int8_t const* gmem)
|
||||
{
|
||||
int32_t tmp[N / 4];
|
||||
ldgStream(tmp, gmem);
|
||||
toFloat(dst, tmp);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void stg(T* gmem, int32_t (&src)[1])
|
||||
{
|
||||
reinterpret_cast<int32_t*>(gmem)[0] = src[0];
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void stgStream(T* gmem, int32_t (&src)[1])
|
||||
{
|
||||
uint32_t tmp = src[0];
|
||||
asm volatile("st.global.cs.s32 [%0], %1;" ::"l"((uint32_t*) gmem), "r"(tmp));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void stg(T* gmem, int32_t (&src)[2])
|
||||
{
|
||||
reinterpret_cast<int2*>(gmem)[0] = make_int2(src[0], src[1]);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DEVICE_FUNCTION void stgStream(T* gmem, int32_t (&src)[2])
|
||||
{
|
||||
asm volatile("st.global.cs.v2.s32 [%0], {%1,%2};" ::"l"((uint32_t*) gmem), "r"(src[0]), "r"(src[1]));
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void stg(uint16_t* gmem, float (&src)[N], float scale)
|
||||
{
|
||||
int32_t tmp[N / 2];
|
||||
fromFloat(tmp, src);
|
||||
stg(gmem, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void stgStream(uint16_t* gmem, float (&src)[N], float scale)
|
||||
{
|
||||
int32_t tmp[N / 2];
|
||||
fromFloat(tmp, src);
|
||||
stgStream(gmem, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void stg(int8_t* gmem, float (&src)[N], float scale)
|
||||
{
|
||||
int32_t tmp[N / 4];
|
||||
fromFloat(tmp, src, scale);
|
||||
stg(gmem, tmp);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void stgStream(int8_t* gmem, float (&src)[N], float scale)
|
||||
{
|
||||
int32_t tmp[N / 4];
|
||||
fromFloat(tmp, src, scale);
|
||||
stg(gmem, tmp);
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromGmem(float (&dst)[2], float const* gmem, int32_t idx)
|
||||
{
|
||||
float2 tmp = __ldg((float2*) &gmem[2 * idx]);
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromGmem(float (&dst)[4], float const* gmem, int32_t idx)
|
||||
{
|
||||
float4 tmp = __ldg((float4*) &gmem[4 * idx]);
|
||||
dst[0] = tmp.x;
|
||||
dst[1] = tmp.y;
|
||||
dst[2] = tmp.z;
|
||||
dst[3] = tmp.w;
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void readFromGmem(float (&dst)[N], __half const* gmem, int32_t idx)
|
||||
{
|
||||
int32_t ival[N / 2];
|
||||
if (N == 4)
|
||||
reinterpret_cast<int2*>(ival)[0] = __ldg((int2*) &gmem[4 * idx]);
|
||||
else
|
||||
reinterpret_cast<int32_t*>(ival)[0] = __ldg((int32_t*) &gmem[2 * idx]);
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N / 2; ++i)
|
||||
{
|
||||
uint16_t lo, hi;
|
||||
asm volatile("mov.b32 {%0, %1}, %2;" : "=h"(lo), "=h"(hi) : "r"(ival[i]));
|
||||
asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 0]) : "h"(lo));
|
||||
asm volatile("cvt.f32.f16 %0, %1;" : "=f"(dst[2 * i + 1]) : "h"(hi));
|
||||
}
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromSmem(float (&x)[2], float const* smem, int32_t idx)
|
||||
{
|
||||
float2 tmp = *(float2 const*) &smem[2 * idx];
|
||||
x[0] = tmp.x;
|
||||
x[1] = tmp.y;
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromSmem(float (&x)[4], float const* smem, int32_t idx)
|
||||
{
|
||||
float4 tmp = *(float4 const*) &smem[4 * idx];
|
||||
x[0] = tmp.x;
|
||||
x[1] = tmp.y;
|
||||
x[2] = tmp.z;
|
||||
x[3] = tmp.w;
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromSmem(int32_t (&x)[1], int32_t const* smem, int32_t idx)
|
||||
{
|
||||
x[0] = smem[idx];
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void readFromSmem(int32_t (&x)[2], int32_t const* smem, int32_t idx)
|
||||
{
|
||||
int2 tmp = *(int2 const*) &smem[2 * idx];
|
||||
x[0] = tmp.x;
|
||||
x[1] = tmp.y;
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void writeToGmem(float* gmem, int32_t idx, float const (&src)[2])
|
||||
{
|
||||
reinterpret_cast<float2*>(&gmem[2 * idx])[0] = make_float2(src[0], src[1]);
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void writeToGmem(float* gmem, int32_t idx, float const (&src)[4])
|
||||
{
|
||||
reinterpret_cast<float4*>(&gmem[4 * idx])[0] = make_float4(src[0], src[1], src[2], src[3]);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void writeToGmem(__half* gmem, int32_t idx, float const (&src)[N])
|
||||
{
|
||||
int32_t ival[N / 2];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N / 2; ++i)
|
||||
{
|
||||
uint16_t lo;
|
||||
uint16_t hi;
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(lo) : "f"(src[2 * i + 0]));
|
||||
asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(hi) : "f"(src[2 * i + 1]));
|
||||
asm volatile("mov.b32 %0, {%1, %2};" : "=r"(ival[i]) : "h"(lo), "h"(hi));
|
||||
}
|
||||
if (N == 4)
|
||||
{
|
||||
reinterpret_cast<int2*>(&gmem[4 * idx])[0] = make_int2(ival[0], ival[1]);
|
||||
}
|
||||
else
|
||||
{
|
||||
reinterpret_cast<int32_t*>(&gmem[2 * idx])[0] = ival[0];
|
||||
}
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void writeToSmem(float* smem, int32_t idx, float const (&x)[2])
|
||||
{
|
||||
reinterpret_cast<float2*>(&smem[2 * idx])[0] = make_float2(x[0], x[1]);
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void writeToSmem(float* smem, int32_t idx, float const (&x)[4])
|
||||
{
|
||||
reinterpret_cast<float4*>(&smem[4 * idx])[0] = make_float4(x[0], x[1], x[2], x[3]);
|
||||
}
|
||||
|
||||
DEVICE_FUNCTION void writeToSmem(int32_t* smem, int32_t idx, int32_t const (&x)[1])
|
||||
{
|
||||
smem[idx] = x[0];
|
||||
}
|
||||
|
||||
static inline __device__ void writeToSmem(int32_t* smem, int32_t idx, int32_t const (&x)[2])
|
||||
{
|
||||
reinterpret_cast<int2*>(&smem[2 * idx])[0] = make_int2(x[0], x[1]);
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void zero(int32_t (&dst)[N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
dst[i] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void zero(float (&dst)[N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
dst[i] = 0.f;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void add(float (&x)[N], float const (&y)[N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
x[i] += y[i];
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void normalize(float (&x)[N], float const (&bias)[N], float const (&scale)[N], float const (&m1)[N])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
x[i] = bias[i] + scale[i] * (x[i] - m1[i]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Storage>
|
||||
DEVICE_FUNCTION Storage relu(Storage in, Storage alpha)
|
||||
{
|
||||
Storage zero = (Storage) 0.f;
|
||||
return (in < zero) ? in * alpha : in;
|
||||
}
|
||||
|
||||
template <int32_t N>
|
||||
DEVICE_FUNCTION void reluActivation(float (&x)[N], float alpha)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < N; ++i)
|
||||
{
|
||||
x[i] = relu(x[i], alpha);
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t THREADS_PER_CTA>
|
||||
DEVICE_FUNCTION void parallelSums_16x2(float* smem, float (&x)[4], int32_t nhw)
|
||||
{
|
||||
|
||||
// The size of a warp.
|
||||
int32_t const THREADS_PER_WARP = 32;
|
||||
// The number of warps in a CTA.
|
||||
int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
|
||||
// The number of threads per pixel.
|
||||
int32_t const THREADS_PER_PIXEL = 16;
|
||||
// The number of elements per ldg.
|
||||
int32_t const ELEMENTS_PER_LDG = 4;
|
||||
// The warp decomposition.
|
||||
int32_t const warp_id = threadIdx.x / THREADS_PER_WARP;
|
||||
int32_t const lane_id = threadIdx.x % THREADS_PER_WARP;
|
||||
|
||||
// Store the values to shared memory.
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
|
||||
// Compute the parallel sum inside the warp. Use SHFL and reduce the amount of SMEM by 2x?
|
||||
__syncwarp();
|
||||
|
||||
// Read the running sum from the other thread in the warp.
|
||||
float y[ELEMENTS_PER_LDG];
|
||||
if (lane_id < THREADS_PER_PIXEL)
|
||||
{
|
||||
readFromSmem(y, smem, threadIdx.x + THREADS_PER_PIXEL);
|
||||
}
|
||||
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
|
||||
// The data is in SMEM. Do the final reduction.
|
||||
__syncthreads();
|
||||
|
||||
// The warp leaders, write to SMEM.
|
||||
if (lane_id < THREADS_PER_PIXEL)
|
||||
{
|
||||
writeToSmem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);
|
||||
}
|
||||
|
||||
// The data is in SMEM. Do the final reduction.
|
||||
__syncthreads();
|
||||
|
||||
// The 1st warp does all the work.
|
||||
if (warp_id == 0)
|
||||
{
|
||||
readFromSmem(x, smem, threadIdx.x);
|
||||
}
|
||||
|
||||
// We do the final reduction each half-warp sequentially reduces the final values.
|
||||
#pragma unroll
|
||||
for (int32_t offset = 1; offset < WARPS_PER_CTA / 2; ++offset)
|
||||
{
|
||||
|
||||
// Read the mean and variance from the other pixel.
|
||||
if (warp_id == 0)
|
||||
{
|
||||
readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);
|
||||
}
|
||||
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
}
|
||||
|
||||
// Make sure the data is in SMEM.
|
||||
__syncwarp();
|
||||
|
||||
// Store the mean/var for the different pixels. TODO: Use SHFL?
|
||||
if (warp_id == 0)
|
||||
{
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
}
|
||||
|
||||
// Make sure the data is in SMEM.
|
||||
__syncwarp();
|
||||
|
||||
// The first half warp finishes the work.
|
||||
if (threadIdx.x < THREADS_PER_PIXEL)
|
||||
{
|
||||
readFromSmem(y, smem, threadIdx.x + THREADS_PER_PIXEL);
|
||||
}
|
||||
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
|
||||
// Make sure the data was read from SMEM.
|
||||
__syncwarp();
|
||||
|
||||
// Store the final values.
|
||||
if (threadIdx.x < THREADS_PER_PIXEL)
|
||||
{
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t THREADS_PER_CTA>
|
||||
static inline __device__ void parallelSums_8x4(float* smem, float (&x)[4], int32_t nhw)
|
||||
{
|
||||
// The size of a warp.
|
||||
int32_t const THREADS_PER_WARP = 32;
|
||||
// The number of warps in a CTA.
|
||||
int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
|
||||
// The number of threads per pixel.
|
||||
int32_t const THREADS_PER_PIXEL = 8;
|
||||
// The number of elements per ldg.
|
||||
int32_t const ELEMENTS_PER_LDG = 4;
|
||||
// The warp decomposition.
|
||||
int32_t const warp_id = threadIdx.x / THREADS_PER_WARP;
|
||||
int32_t const lane_id = threadIdx.x % THREADS_PER_WARP;
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);
|
||||
x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);
|
||||
}
|
||||
|
||||
// The warp leaders, write to SMEM.
|
||||
if (lane_id < THREADS_PER_PIXEL)
|
||||
{
|
||||
writeToSmem(smem, warp_id * THREADS_PER_PIXEL + lane_id, x);
|
||||
}
|
||||
|
||||
// The data is in SMEM. Do the final reduction.
|
||||
__syncthreads();
|
||||
|
||||
// The 1st warp does all the work.
|
||||
// We do the final reduction each half-warp sequentially reduces the final values.
|
||||
if (warp_id == 0)
|
||||
{
|
||||
readFromSmem(x, smem, threadIdx.x);
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t offset = 1; offset < WARPS_PER_CTA / (THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset)
|
||||
{
|
||||
float y[ELEMENTS_PER_LDG];
|
||||
// Read the mean and variance from the other pixel.
|
||||
readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_WARP);
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
}
|
||||
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL + lane_id);
|
||||
x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL * 2 + lane_id);
|
||||
}
|
||||
|
||||
// Make sure the data was read from SMEM.
|
||||
__syncwarp();
|
||||
|
||||
// Store the final values.
|
||||
if (threadIdx.x < THREADS_PER_PIXEL)
|
||||
{
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t THREADS_PER_CTA, int32_t THREADS_PER_PIXEL, int32_t ELEMENTS_PER_LDG>
|
||||
DEVICE_FUNCTION void parallelSums(float* smem, float (&x)[ELEMENTS_PER_LDG], int32_t nhw)
|
||||
{
|
||||
|
||||
// The size of a warp.
|
||||
int32_t const THREADS_PER_WARP = 32;
|
||||
// The number of warps in a CTA.
|
||||
int32_t const WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
|
||||
// The number of pixels computed by a single warp.
|
||||
int32_t const PIXELS_PER_WARP = THREADS_PER_WARP / THREADS_PER_PIXEL;
|
||||
|
||||
// The position in the warp.
|
||||
int32_t const nhw_in_warp = nhw % PIXELS_PER_WARP;
|
||||
// The C in the warp.
|
||||
int32_t const c_in_warp = threadIdx.x % THREADS_PER_PIXEL;
|
||||
|
||||
// Store the values to shared memory.
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
|
||||
// Compute the parallel sums.
|
||||
for (int32_t offset = PIXELS_PER_WARP / 2; offset > 0; offset /= 2)
|
||||
{
|
||||
|
||||
if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
|
||||
{
|
||||
__syncthreads();
|
||||
}
|
||||
else
|
||||
{
|
||||
// NOP.
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Read the running sum from the other thread.
|
||||
float y[ELEMENTS_PER_LDG];
|
||||
if (nhw_in_warp < offset)
|
||||
{
|
||||
readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);
|
||||
}
|
||||
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
|
||||
if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
|
||||
{
|
||||
__syncthreads();
|
||||
}
|
||||
else
|
||||
{
|
||||
// NOP.
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Update the sum in SMEM.
|
||||
if (offset > 1 && nhw_in_warp < offset)
|
||||
{
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
}
|
||||
}
|
||||
|
||||
// The warps are done. Do the final reduction at the CTA level.
|
||||
__syncthreads();
|
||||
|
||||
// The warp leaders, write to SMEM.
|
||||
int32_t const idx = (threadIdx.x / THREADS_PER_WARP) * THREADS_PER_PIXEL + c_in_warp;
|
||||
if (nhw_in_warp == 0)
|
||||
{
|
||||
writeToSmem(smem, idx, x);
|
||||
}
|
||||
|
||||
// The data is in SMEM. Do the final reduction.
|
||||
__syncthreads();
|
||||
|
||||
// Read the 1st element to prepare the work.
|
||||
if (nhw < WARPS_PER_CTA / 2)
|
||||
{
|
||||
readFromSmem(x, smem, threadIdx.x);
|
||||
}
|
||||
|
||||
// We have the running mean and running m2. Let's build the mean/var of the CTA.
|
||||
for (int32_t offset = WARPS_PER_CTA / 2; offset > 0; offset /= 2)
|
||||
{
|
||||
|
||||
if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
|
||||
{
|
||||
__syncthreads();
|
||||
}
|
||||
else
|
||||
{
|
||||
// NOP.
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Read the mean and variance from the other pixel.
|
||||
float y[ELEMENTS_PER_LDG];
|
||||
if (nhw < offset)
|
||||
{
|
||||
readFromSmem(y, smem, threadIdx.x + offset * THREADS_PER_PIXEL);
|
||||
}
|
||||
|
||||
// Compute the updated sum.
|
||||
add(x, y);
|
||||
|
||||
if ((WARPS_PER_CTA * THREADS_PER_WARP) / THREADS_PER_PIXEL > THREADS_PER_WARP)
|
||||
{
|
||||
__syncthreads();
|
||||
}
|
||||
else
|
||||
{
|
||||
// NOP.
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// Store the mean/var for the different pixels.
|
||||
if (nhw < offset)
|
||||
{
|
||||
writeToSmem(smem, threadIdx.x, x);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t THREADS_PER_PIXEL, int32_t ELEMENTS_PER_LDG>
|
||||
struct ParallelSums
|
||||
{
|
||||
template <int32_t THREADS_PER_CTA>
|
||||
DEVICE_FUNCTION void dispatch(float* smem, float (&x)[ELEMENTS_PER_LDG], int32_t nhw)
|
||||
{
|
||||
parallelSums<THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG>(smem, x, nhw);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct ParallelSums<16, 4>
|
||||
{
|
||||
template <int32_t THREADS_PER_CTA>
|
||||
DEVICE_FUNCTION void dispatch(float* smem, float (&x)[4], int32_t nhw)
|
||||
{
|
||||
parallelSums_16x2<THREADS_PER_CTA>(smem, x, nhw);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct ParallelSums<8, 4>
|
||||
{
|
||||
template <int32_t THREADS_PER_CTA>
|
||||
static inline __device__ void dispatch(float* smem, float (&x)[4], int32_t nhw)
|
||||
{
|
||||
parallelSums_8x4<THREADS_PER_CTA>(smem, x, nhw);
|
||||
}
|
||||
};
|
||||
|
||||
namespace
|
||||
{
|
||||
int32_t divUp(int32_t m, int32_t n)
|
||||
{
|
||||
PLUGIN_ASSERT(m >= 0);
|
||||
PLUGIN_ASSERT(n > 0);
|
||||
// Use unsigned arithmetic to preclude overflow.
|
||||
auto const mu = static_cast<uint32_t>(m);
|
||||
auto const nu = static_cast<uint32_t>(n);
|
||||
return (mu + nu - 1U) / nu;
|
||||
}
|
||||
|
||||
cudnnStatus_t convertTrt2cudnnDtype(nvinfer1::DataType trt_dtype, cudnnDataType_t* cudnn_dtype)
|
||||
{
|
||||
switch (trt_dtype)
|
||||
{
|
||||
case nvinfer1::DataType::kFLOAT: *cudnn_dtype = CUDNN_DATA_FLOAT; break;
|
||||
case nvinfer1::DataType::kHALF: *cudnn_dtype = CUDNN_DATA_HALF; break;
|
||||
default: return CUDNN_STATUS_BAD_PARAM;
|
||||
}
|
||||
return CUDNN_STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
template <typename T, int32_t THREADS_PER_CTA>
|
||||
__global__ __launch_bounds__(THREADS_PER_CTA) void in3dReluActivation(T* dst, T const* src, T alpha, int32_t count)
|
||||
{
|
||||
int32_t idx = blockIdx.x * THREADS_PER_CTA + threadIdx.x;
|
||||
if (idx >= count)
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
T val = src[idx];
|
||||
dst[idx] = (val < static_cast<T>(0.F)) ? val * alpha : val;
|
||||
}
|
||||
|
||||
#endif // INSTANCE_NORM_COMMON_H
|
||||
@@ -0,0 +1,165 @@
|
||||
/*
|
||||
* 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 INSTANCE_NORM_FWD_H
|
||||
#define INSTANCE_NORM_FWD_H
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime_api.h>
|
||||
|
||||
namespace instance_norm_impl
|
||||
{
|
||||
#define PLUGIN_CHECK_CUDA(call) \
|
||||
do \
|
||||
{ \
|
||||
cudaError_t status = call; \
|
||||
if (status != cudaSuccess) \
|
||||
{ \
|
||||
return status; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define PLUGIN_CHECK_CUDNN(call) \
|
||||
do \
|
||||
{ \
|
||||
cudnnStatus_t status = call; \
|
||||
if (status != CUDNN_STATUS_SUCCESS) \
|
||||
{ \
|
||||
return status; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
typedef float GMEM_SUMS_TYPE;
|
||||
|
||||
#define ACCUM_MEAN_VAR_IN_FLOAT 1
|
||||
|
||||
template <typename StorageType, int32_t SM>
|
||||
constexpr int32_t getPixelsPerThreadInRegisters()
|
||||
{
|
||||
return (sizeof(StorageType) == 4 || sizeof(StorageType) == 2)
|
||||
? 6 - sizeof(StorageType)
|
||||
: (SM < 800 ? (SM == 750 ? 16 : 8) : (SM == 860 ? 16 : 24));
|
||||
}
|
||||
|
||||
template <typename StorageType, int32_t SM>
|
||||
constexpr int32_t getPixelsPerThreadInSmem()
|
||||
{
|
||||
return (sizeof(StorageType) == 4 || sizeof(StorageType) == 2)
|
||||
? (sizeof(StorageType) == 4 ? 4 : 8)
|
||||
: (SM < 800 ? (SM == 750 ? 7 : 8) : (SM == 860 ? 16 : 24));
|
||||
}
|
||||
|
||||
template <typename Input_Data_Type_ = uint16_t, typename Output_Data_Type_ = uint16_t, typename StorageType_ = float,
|
||||
int32_t THREADS_PER_CTA_ = 512, int32_t THREADS_PER_PIXEL_ = 16, int32_t C_ELEMENTS_PER_CTA_ = 64,
|
||||
int32_t SM_ = 700>
|
||||
struct Instance_norm_kernel_params
|
||||
{
|
||||
static constexpr int32_t USE_ONLINE_APPROACH = 1;
|
||||
|
||||
static constexpr int32_t THREADS_PER_CTA = THREADS_PER_CTA_;
|
||||
|
||||
//! 8 or 16
|
||||
static constexpr int32_t THREADS_PER_PIXEL = THREADS_PER_PIXEL_;
|
||||
|
||||
static constexpr int32_t SM = SM_;
|
||||
|
||||
typedef Input_Data_Type_ Input_Data_Type;
|
||||
typedef Output_Data_Type_ Output_Data_Type;
|
||||
|
||||
typedef StorageType_ StorageType;
|
||||
|
||||
static constexpr int32_t PIXELS_PER_THREAD_IN_REGISTERS = getPixelsPerThreadInRegisters<StorageType, SM>();
|
||||
|
||||
static constexpr int32_t PIXELS_PER_THREAD_IN_SMEM = getPixelsPerThreadInSmem<StorageType, SM>();
|
||||
|
||||
//! 64
|
||||
static constexpr int32_t C_ELEMENTS_PER_CTA = C_ELEMENTS_PER_CTA_;
|
||||
|
||||
//! 4 default
|
||||
static constexpr int32_t ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL;
|
||||
|
||||
// Derived params.
|
||||
static constexpr int32_t PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
|
||||
|
||||
static constexpr int32_t MIN_PIXELS_PER_CTA = PIXELS_PER_LDG * PIXELS_PER_THREAD_IN_REGISTERS;
|
||||
};
|
||||
|
||||
struct InstanceNormFwdContext
|
||||
{
|
||||
InstanceNormFwdContext()
|
||||
: sm_count(0)
|
||||
, sm_shared_size(0)
|
||||
, sm_version(0)
|
||||
{
|
||||
}
|
||||
int32_t sm_count;
|
||||
int32_t sm_shared_size;
|
||||
int32_t sm_version;
|
||||
};
|
||||
|
||||
struct InstanceNormFwdParams
|
||||
{
|
||||
// The input/output tensors.
|
||||
void const* gmem_src;
|
||||
void* gmem_dst;
|
||||
// The bias/scale.
|
||||
float* gmem_bias;
|
||||
float* gmem_scale;
|
||||
// running mean/var (refer BN API from cudnn doc)
|
||||
float* gmem_running_mean;
|
||||
float* gmem_running_var;
|
||||
// saved mean/var (refer BN API from cudnn doc)
|
||||
float* gmem_saved_mean;
|
||||
float* gmem_saved_var;
|
||||
// The dimensions.
|
||||
int32_t nhw;
|
||||
int32_t c;
|
||||
int32_t n;
|
||||
// The buffer to do the reduction for mean, stddev and count.
|
||||
GMEM_SUMS_TYPE* gmem_sums;
|
||||
// The buffer to count items in the different CTAs.
|
||||
int32_t* gmem_counts;
|
||||
// The counters of retired CTAs.
|
||||
int32_t* gmem_retired_ctas;
|
||||
// The epsilon to apply to the computation of the variance.
|
||||
float var_eps;
|
||||
// outer loop count
|
||||
int32_t outer_loops;
|
||||
// exponential average factor
|
||||
float exp_avg_factor;
|
||||
bool use_relu;
|
||||
float relu_alpha;
|
||||
|
||||
int32_t c_blks;
|
||||
|
||||
float in_scale;
|
||||
|
||||
float out_scale;
|
||||
};
|
||||
|
||||
void instanceNormBufferSizesDispatch(InstanceNormFwdContext const& context, InstanceNormFwdParams const& params,
|
||||
size_t& size_sums, size_t& size_counts, size_t& size_retired_ctas, int32_t input_data_type = 1,
|
||||
int32_t output_data_type = 1);
|
||||
|
||||
int32_t instanceNormFwdDispatch(InstanceNormFwdContext const& context, InstanceNormFwdParams& params,
|
||||
cudaStream_t stream, int32_t input_data_type = 1, int32_t output_data_type = 1);
|
||||
|
||||
} // namespace instance_norm_impl
|
||||
|
||||
#endif // INSTANCE_NORM_FWD_H
|
||||
@@ -0,0 +1,856 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#include <stdio.h>
|
||||
#include <assert.h>
|
||||
#include <type_traits>
|
||||
#include <cuda.h>
|
||||
|
||||
#include "instanceNormFwd.h"
|
||||
#include "instanceNormCommon.h"
|
||||
|
||||
namespace instance_norm_impl
|
||||
{
|
||||
|
||||
static inline int32_t divUp(int32_t m, int32_t n)
|
||||
{
|
||||
return (m + n - 1) / n;
|
||||
}
|
||||
|
||||
using kernel_params_32 = Instance_norm_kernel_params<uint16_t, uint16_t, uint16_t, 512, 8, 32>;
|
||||
using kernel_params_64 = Instance_norm_kernel_params<uint16_t, uint16_t, uint16_t, 512, 16, 64>;
|
||||
|
||||
using kernel_params_32_int8 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32>;
|
||||
using kernel_params_32_int8_sm_700 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 700>;
|
||||
using kernel_params_32_int8_sm_720 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 720>;
|
||||
using kernel_params_32_int8_sm_750 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 750>;
|
||||
using kernel_params_32_int8_sm_800 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 800>;
|
||||
using kernel_params_32_int8_sm_860 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 860>;
|
||||
using kernel_params_32_int8_sm_870 = Instance_norm_kernel_params<int8_t, int8_t, int8_t, 512, 8, 32, 870>;
|
||||
using kernel_params_32_fp16_int8 = Instance_norm_kernel_params<uint16_t, int8_t, float, 512, 8, 32>;
|
||||
|
||||
template <typename Storage, typename Input_Data_Type, typename Output_Data_Type, int32_t THREADS_PER_CTA,
|
||||
int32_t THREADS_PER_PIXEL, int32_t PIXELS_PER_THREAD_IN_REGISTERS, int32_t PIXELS_PER_THREAD_IN_SMEM,
|
||||
int32_t ELEMENTS_PER_LDG, int32_t USE_ONLINE_APPROACH, int32_t OUTER_LOOPS_, int32_t DESIRED_OCCUPANCY>
|
||||
__global__ __launch_bounds__(THREADS_PER_CTA, DESIRED_OCCUPANCY) void instanceNormFwd(InstanceNormFwdParams params)
|
||||
{
|
||||
|
||||
// Single pass numerically stable algorithm, see:
|
||||
// https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm
|
||||
//
|
||||
// n = 0, mean = 0.0, M2 = 0.0
|
||||
//
|
||||
// for x in data:
|
||||
// n += 1
|
||||
// delta = x - mean
|
||||
// mean += delta/n
|
||||
// delta2 = x - mean
|
||||
// M2 += delta*delta2
|
||||
//
|
||||
// if n < 2:
|
||||
// return float('nan')
|
||||
// else:
|
||||
// return M2 / (n - 1)
|
||||
|
||||
bool const IS_INPUT_INT8 = std::is_same_v<Input_Data_Type, int8_t>;
|
||||
bool const IS_OUTPUT_INT8 = std::is_same_v<Output_Data_Type, int8_t>;
|
||||
|
||||
// The number of pixels loaded in a single LDG.
|
||||
int32_t const PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
|
||||
// The number of pixels computed per CTA stored in registers.
|
||||
int32_t const PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;
|
||||
// The number of pixels computed per CTA stored in SMEM.
|
||||
int32_t const PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM * PIXELS_PER_LDG;
|
||||
// The number of C elements per CTA.
|
||||
int32_t const C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL * ELEMENTS_PER_LDG;
|
||||
|
||||
// Shared memory to do CTA-wide parallel sums.
|
||||
__shared__ float smem[ELEMENTS_PER_LDG * THREADS_PER_CTA];
|
||||
|
||||
// The position in the NHW dimension where the CTA starts.
|
||||
int32_t cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;
|
||||
// The position in the NHW dimension where the CTA starts for the portion in SMEM.
|
||||
int32_t cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;
|
||||
// Compute the NHW coordinate of the thread in the CTA.
|
||||
int32_t const thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
|
||||
|
||||
for (int32_t nc_blk_index = blockIdx.y; nc_blk_index < params.c_blks * params.n; nc_blk_index += gridDim.y)
|
||||
{
|
||||
|
||||
int32_t n_blk_index = nc_blk_index / params.c_blks;
|
||||
int32_t c_blk_index = nc_blk_index % params.c_blks;
|
||||
|
||||
// The position in the C dimension where the CTA starts.
|
||||
int32_t const cta_c = c_blk_index * C_ELEMENTS_PER_CTA;
|
||||
// Compute the C coordinate of the thread in the CTA.
|
||||
int32_t const thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
|
||||
// Compute the C coordinate of the thread.
|
||||
int32_t const thread_c = cta_c + thread_in_cta_c * ELEMENTS_PER_LDG;
|
||||
|
||||
// Is the thread working on a valid C dimension?
|
||||
int32_t const is_valid_c = thread_c < params.c;
|
||||
|
||||
// The adapter for the storage.
|
||||
typedef PackedStorage<Storage, ELEMENTS_PER_LDG> PackedStorage_;
|
||||
// The data type for packed storage in SMEM.
|
||||
typedef typename PackedStorage_::Type PackedStorageType;
|
||||
// The number of elements in the packed storage.
|
||||
int32_t const PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;
|
||||
// Registers to keep the data live for the persistent approach.
|
||||
PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
|
||||
|
||||
// Shared memory buffer to store the extra pixels.
|
||||
extern __shared__ char smem_storage_[];
|
||||
PackedStorageType* smem_storage = reinterpret_cast<PackedStorageType*>(smem_storage_);
|
||||
|
||||
float int8_in_scale = params.in_scale;
|
||||
float int8_out_scale = params.out_scale;
|
||||
|
||||
// Register to store the number of elements read so far.
|
||||
float count = 0.f, mean[ELEMENTS_PER_LDG], m2[ELEMENTS_PER_LDG];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
mean[i] = 0.f;
|
||||
m2[i] = 0.f;
|
||||
}
|
||||
|
||||
// The number of elements loaded by this CTA.
|
||||
int32_t cta_count = 0;
|
||||
int32_t global_batch_offset = n_blk_index * params.nhw * params.c;
|
||||
// int8 relevant
|
||||
// int8 output implies we have NC/32DHW32 input for bath fp16 and int8
|
||||
int32_t global_thread_c_input = (IS_INPUT_INT8 || IS_OUTPUT_INT8)
|
||||
? thread_in_cta_c * ELEMENTS_PER_LDG + (cta_c % 32) // handle C_ELEMENTS_PER_CTA == 16 case
|
||||
+ (cta_c / 32) * 32 * params.nhw
|
||||
: thread_c;
|
||||
int32_t stride_c_input = (IS_INPUT_INT8 || IS_OUTPUT_INT8) ? 32 : params.c;
|
||||
int32_t global_thread_c_output = (IS_OUTPUT_INT8)
|
||||
? thread_in_cta_c * ELEMENTS_PER_LDG + (cta_c % 32) // handle C_ELEMENTS_PER_CTA == 16 case
|
||||
+ (cta_c / 32) * 32 * params.nhw
|
||||
: thread_c;
|
||||
int32_t stride_c_output = (IS_OUTPUT_INT8) ? 32 : params.c;
|
||||
// The base pointer to load from.
|
||||
Input_Data_Type const* gmem_src
|
||||
= &reinterpret_cast<Input_Data_Type const*>(params.gmem_src)[global_thread_c_input + global_batch_offset];
|
||||
|
||||
// Load the batch of elements. Compute the mean/var across those elements.
|
||||
int32_t const pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS * gridDim.x;
|
||||
|
||||
// outer loops
|
||||
int32_t OUTER_LOOPS = OUTER_LOOPS_ == 1 ? 1 : params.outer_loops;
|
||||
|
||||
#pragma unroll 1
|
||||
for (int32_t loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i)
|
||||
{
|
||||
// The nhw position.
|
||||
int32_t nhw_regs = cta_nhw_regs + loop_i * pixels_per_iteration;
|
||||
|
||||
cta_count += max(min(nhw_regs + PIXELS_PER_CTA_IN_REGISTERS, params.nhw) - max(nhw_regs, 0), 0);
|
||||
|
||||
// Load the data and compute the local mean/sum and the variance.
|
||||
if (USE_ONLINE_APPROACH)
|
||||
{
|
||||
// Read the elements from memory.
|
||||
float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
int32_t const idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
zero(x_storage[i]);
|
||||
is_valid[i] = 0.f;
|
||||
if (idx < params.nhw && is_valid_c)
|
||||
{
|
||||
ldgStream(x_storage[i], &gmem_src[idx * stride_c_input]);
|
||||
is_valid[i] = 1.f;
|
||||
}
|
||||
}
|
||||
|
||||
// Do the math.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
// Convert to float.
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage[i], int8_in_scale);
|
||||
|
||||
// Update the count.
|
||||
count += is_valid[i];
|
||||
// Invert the count.
|
||||
float inv_count = is_valid[i] ? 1.f / count : 0.f;
|
||||
|
||||
// Update the mean and m2 using deltas.
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < ELEMENTS_PER_LDG; ++j)
|
||||
{
|
||||
float delta0 = x_math[j] - mean[j];
|
||||
mean[j] += delta0 * inv_count;
|
||||
float delta1 = x_math[j] - mean[j];
|
||||
m2[j] += delta0 * delta1 * is_valid[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Read the elements from memory.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
int32_t const idx = nhw_regs + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
zero(x_storage[i]);
|
||||
if (idx < params.nhw && is_valid_c)
|
||||
{
|
||||
ldgStream(x_storage[i], &gmem_src[idx * stride_c_input]);
|
||||
count += 1.f;
|
||||
}
|
||||
}
|
||||
|
||||
// Sum the elements in registers.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
// Convert to float.
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage[i], int8_in_scale);
|
||||
|
||||
// Update the mean and m2 using deltas.
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < ELEMENTS_PER_LDG; ++j)
|
||||
{
|
||||
mean[j] += x_math[j];
|
||||
}
|
||||
}
|
||||
|
||||
// Compute the mean.
|
||||
float inv_count = 1.f / count;
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < ELEMENTS_PER_LDG; ++j)
|
||||
{
|
||||
mean[j] *= inv_count;
|
||||
}
|
||||
|
||||
// Compute the variance.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
// Convert to float.
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage[i], int8_in_scale);
|
||||
|
||||
// Is it a valid pixel?
|
||||
float is_valid = i < (int32_t) count ? 1.f : 0.f;
|
||||
// Update the mean and m2 using deltas.
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < ELEMENTS_PER_LDG; ++j)
|
||||
{
|
||||
m2[j] += (x_math[j] - mean[j]) * (x_math[j] - mean[j]) * is_valid;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The elements to load and store in SMEM.
|
||||
int32_t smem_nhw = OUTER_LOOPS * pixels_per_iteration + cta_nhw_smem;
|
||||
// Load elements from SMEM, update the CTA count.
|
||||
int32_t pixels_in_smem = min(smem_nhw + PIXELS_PER_CTA_IN_SMEM, params.nhw) - max(smem_nhw, 0);
|
||||
if (pixels_in_smem > 0)
|
||||
{
|
||||
cta_count += pixels_in_smem;
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i)
|
||||
{
|
||||
int32_t const idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
float is_pixel_valid = (idx < params.nhw && is_valid_c) ? 1.f : 0.f;
|
||||
|
||||
PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];
|
||||
ldgStream(x_storage_local, &gmem_src[(is_pixel_valid ? idx : 0) * stride_c_input]);
|
||||
|
||||
// The offset to store in SMEM.
|
||||
int32_t const offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;
|
||||
// Store in SMEM.
|
||||
writeToSmem(&smem_storage[offset], threadIdx.x, x_storage_local);
|
||||
// Update the count.
|
||||
count += is_pixel_valid;
|
||||
// Invert the count.
|
||||
float inv_count = is_pixel_valid ? 1.f / count : 0.f;
|
||||
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage_local, int8_in_scale);
|
||||
// Update the mean and m2 using deltas.
|
||||
#pragma unroll
|
||||
for (int32_t j = 0; j < ELEMENTS_PER_LDG; ++j)
|
||||
{
|
||||
float delta0 = x_math[j] - mean[j];
|
||||
mean[j] += delta0 * inv_count;
|
||||
float delta1 = x_math[j] - mean[j];
|
||||
m2[j] += delta0 * delta1 * is_pixel_valid;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// We scale the mean by the number of elements. It brings more stability.
|
||||
float m1[ELEMENTS_PER_LDG];
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m1[i] = mean[i] * count;
|
||||
}
|
||||
|
||||
// Run the parallel sum accross the CTA to get the local sum.
|
||||
ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(smem, m1, thread_in_cta_nhw);
|
||||
__syncthreads();
|
||||
|
||||
// The values in shared memory correspond to the CTA-wide sums.
|
||||
readFromSmem(m1, smem, thread_in_cta_c);
|
||||
__syncthreads();
|
||||
|
||||
// Adjust the variance.
|
||||
float inv_cta_count = 1.f / (float) cta_count;
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
float mean_diff = m1[i] * inv_cta_count - mean[i];
|
||||
m2[i] = m2[i] + mean_diff * mean_diff * count;
|
||||
}
|
||||
|
||||
// Run the parallel sum accross the CTA to get the local adjusted variance.
|
||||
ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(smem, m2, thread_in_cta_nhw);
|
||||
|
||||
// The workspace in global memory is distributed across the different CTA.
|
||||
int32_t gmem_sums_offset = nc_blk_index * gridDim.x * C_ELEMENTS_PER_CTA * 2;
|
||||
|
||||
// Write the data for the CTA to global memory.
|
||||
GMEM_SUMS_TYPE* gmem_sums = ¶ms.gmem_sums[gmem_sums_offset];
|
||||
if (threadIdx.x < THREADS_PER_PIXEL)
|
||||
{
|
||||
int32_t const idx = blockIdx.x * THREADS_PER_PIXEL + threadIdx.x;
|
||||
writeToGmem(&gmem_sums[0], idx, m1);
|
||||
writeToGmem(&gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx, m2);
|
||||
}
|
||||
|
||||
// The memory location to store the number of pixels per CTA.
|
||||
int32_t* gmem_counts = ¶ms.gmem_counts[nc_blk_index * gridDim.x];
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
// gmem_counts[0] = cta_count;
|
||||
gmem_counts[blockIdx.x] = cta_count;
|
||||
}
|
||||
|
||||
// Read the bias and scale.
|
||||
float bias[ELEMENTS_PER_LDG];
|
||||
float scale[ELEMENTS_PER_LDG];
|
||||
if (is_valid_c)
|
||||
{
|
||||
readFromGmem(bias, ¶ms.gmem_bias[cta_c], thread_in_cta_c);
|
||||
readFromGmem(scale, ¶ms.gmem_scale[cta_c], thread_in_cta_c);
|
||||
}
|
||||
|
||||
// The counters to count how many CTAs have retired at this point. One per chunk of C.
|
||||
int32_t* gmem_retired_ctas = ¶ms.gmem_retired_ctas[nc_blk_index];
|
||||
|
||||
// Make sure the threads are done and reconverged.
|
||||
__syncthreads();
|
||||
|
||||
// Register the CTA.
|
||||
int32_t expected_count = gridDim.x;
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
// Issue the membar.
|
||||
__threadfence();
|
||||
// Notify that the CTA is done.
|
||||
int32_t val_to_add = 1;
|
||||
if (blockIdx.x == 0)
|
||||
{
|
||||
val_to_add = -(expected_count - 1);
|
||||
}
|
||||
atomicAdd(gmem_retired_ctas, val_to_add);
|
||||
}
|
||||
|
||||
// Are all CTAs done?
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
int32_t retired_ctas = -1;
|
||||
do
|
||||
{
|
||||
__threadfence();
|
||||
asm volatile("ld.global.cg.b32 %0, [%1];" : "=r"(retired_ctas) : "l"(gmem_retired_ctas));
|
||||
} while (retired_ctas != 0);
|
||||
}
|
||||
__threadfence();
|
||||
__syncthreads();
|
||||
|
||||
// Reset the mean to compute the global mean.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m1[i] = 0.f;
|
||||
}
|
||||
|
||||
// Build the global mean.
|
||||
#pragma unroll 1
|
||||
for (int32_t idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA)
|
||||
{
|
||||
float tmp[ELEMENTS_PER_LDG];
|
||||
readFromGmem(tmp, gmem_sums, idx);
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m1[i] += tmp[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Run the parallel sum accross the CTA to get the local sum.
|
||||
ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(smem, m1, thread_in_cta_nhw);
|
||||
__syncthreads();
|
||||
|
||||
// The values in shared memory correspond to the CTA-wide sums.
|
||||
readFromSmem(m1, smem, thread_in_cta_c);
|
||||
__syncthreads();
|
||||
|
||||
// Normalize the mean.
|
||||
float inv_count = 1.f / (float) params.nhw;
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m1[i] = m1[i] * inv_count;
|
||||
}
|
||||
|
||||
// Reset the variance.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m2[i] = 0.f;
|
||||
}
|
||||
|
||||
// Build the global variance.
|
||||
#pragma unroll 1
|
||||
for (int32_t idx = threadIdx.x; idx < THREADS_PER_PIXEL * gridDim.x; idx += THREADS_PER_CTA)
|
||||
{
|
||||
|
||||
// Read the means computed by different CTAs (again). Reuse tmp if we have 1 iteration.
|
||||
float tmp_mean[ELEMENTS_PER_LDG], tmp_var[ELEMENTS_PER_LDG];
|
||||
readFromGmem(tmp_mean, &gmem_sums[0], idx);
|
||||
readFromGmem(tmp_var, &gmem_sums[C_ELEMENTS_PER_CTA * gridDim.x], idx);
|
||||
|
||||
// Read the number of pixels visited by a given CTA.
|
||||
cta_count = __ldg(&gmem_counts[idx / THREADS_PER_PIXEL]);
|
||||
|
||||
// Compute the diff to update the variance.
|
||||
float mean_diff[ELEMENTS_PER_LDG], inv_cta_count = 1.f / (float) cta_count;
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
mean_diff[i] = m1[i] - tmp_mean[i] * inv_cta_count;
|
||||
}
|
||||
|
||||
// Update the variance.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m2[i] += tmp_var[i] + mean_diff[i] * mean_diff[i] * (float) cta_count;
|
||||
}
|
||||
}
|
||||
|
||||
// Run the parallel sum accross the CTA to get the local sum.
|
||||
ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(smem, m2, thread_in_cta_nhw);
|
||||
__syncthreads();
|
||||
|
||||
readFromSmem(m2, smem, thread_in_cta_c);
|
||||
__syncthreads();
|
||||
|
||||
// Finalize the stddev.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
m2[i] *= inv_count;
|
||||
}
|
||||
|
||||
// store the saved mean/var
|
||||
float svarinv[ELEMENTS_PER_LDG];
|
||||
bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
svarinv[i] = rsqrtf(m2[i] + params.var_eps);
|
||||
}
|
||||
|
||||
#if ACCUM_MEAN_VAR_IN_FLOAT
|
||||
int32_t global_stats_offset = n_blk_index * params.c;
|
||||
if (is_valid_for_saving)
|
||||
{
|
||||
writeToGmem(params.gmem_saved_mean + global_stats_offset, thread_c / ELEMENTS_PER_LDG, m1);
|
||||
writeToGmem(params.gmem_saved_var + global_stats_offset, thread_c / ELEMENTS_PER_LDG, svarinv);
|
||||
}
|
||||
|
||||
// store the running mean/var
|
||||
float rmean[ELEMENTS_PER_LDG];
|
||||
float rvar[ELEMENTS_PER_LDG];
|
||||
zero(rmean);
|
||||
zero(rvar);
|
||||
if (params.exp_avg_factor != 1.f && is_valid_for_saving)
|
||||
{
|
||||
readFromGmem(rmean, params.gmem_running_mean + global_stats_offset, thread_c / ELEMENTS_PER_LDG);
|
||||
readFromGmem(rvar, params.gmem_running_var + global_stats_offset, thread_c / ELEMENTS_PER_LDG);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
rmean[i] = (1.f - params.exp_avg_factor) * rmean[i] + params.exp_avg_factor * m1[i];
|
||||
rvar[i] = (1.f - params.exp_avg_factor) * rvar[i] + params.exp_avg_factor * m2[i];
|
||||
}
|
||||
if (is_valid_for_saving)
|
||||
{
|
||||
writeToGmem(params.gmem_running_mean + global_stats_offset, thread_c / ELEMENTS_PER_LDG, rmean);
|
||||
writeToGmem(params.gmem_running_var + global_stats_offset, thread_c / ELEMENTS_PER_LDG, rvar);
|
||||
}
|
||||
#endif
|
||||
|
||||
// Update the scale with the stddev and eps.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < ELEMENTS_PER_LDG; ++i)
|
||||
{
|
||||
scale[i] *= svarinv[i];
|
||||
}
|
||||
|
||||
// The base pointer to write to.
|
||||
Output_Data_Type* const gmem_dst
|
||||
= &reinterpret_cast<Output_Data_Type*>(params.gmem_dst)[global_thread_c_output + global_batch_offset];
|
||||
|
||||
// Store the elements in registers.
|
||||
#pragma unroll 1
|
||||
for (int32_t loop_i = OUTER_LOOPS - 1; loop_i >= 0; --loop_i)
|
||||
{
|
||||
|
||||
// The value for nhw.
|
||||
int32_t out_nhw = cta_nhw_regs + loop_i * pixels_per_iteration;
|
||||
|
||||
// On CUDA-11.5 or above, full unrolling caused compiler to panic about register pressure and the perf
|
||||
// dropped significantly. Therefore, limit the extent of unrolling on CUDA-11.5 or above. The number "8" is
|
||||
// chosen based on experiments.
|
||||
#if CUDA_VERSION >= 11050
|
||||
#pragma unroll 8
|
||||
#else
|
||||
#pragma unroll
|
||||
#endif
|
||||
// Normalize the elements and write to memory.
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
// Convert to float.
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage[i], int8_in_scale);
|
||||
|
||||
// Normalize and apply activation function
|
||||
normalize(x_math, bias, scale, m1);
|
||||
if (params.use_relu)
|
||||
{
|
||||
reluActivation(x_math, params.relu_alpha);
|
||||
}
|
||||
|
||||
// Write back.
|
||||
int32_t const idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
if ((unsigned) idx < params.nhw && is_valid_c)
|
||||
{
|
||||
stgStream(&gmem_dst[idx * stride_c_output], x_math, int8_out_scale);
|
||||
}
|
||||
}
|
||||
|
||||
// The next value of nhw.
|
||||
out_nhw -= pixels_per_iteration;
|
||||
|
||||
// Read the next elements from memory.
|
||||
#pragma unroll
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i)
|
||||
{
|
||||
int32_t const idx = out_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
if ((unsigned) idx < params.nhw && is_valid_c)
|
||||
{
|
||||
ldgStream(x_storage[i], &gmem_src[idx * stride_c_output]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize the elements from SMEM and write them out.
|
||||
if (pixels_in_smem > 0)
|
||||
{
|
||||
for (int32_t i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i)
|
||||
{
|
||||
// Read from SMEM.
|
||||
int32_t const offset = i * THREADS_PER_CTA * PACKED_ELEMENTS_PER_LDG;
|
||||
float x_math[ELEMENTS_PER_LDG];
|
||||
PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];
|
||||
readFromSmem(x_storage_local, &smem_storage[offset], threadIdx.x);
|
||||
toFloat<PACKED_ELEMENTS_PER_LDG, IS_INPUT_INT8>(x_math, x_storage_local, int8_in_scale);
|
||||
|
||||
// Normalize and apply activation function
|
||||
normalize(x_math, bias, scale, m1);
|
||||
if (params.use_relu)
|
||||
{
|
||||
reluActivation(x_math, params.relu_alpha);
|
||||
}
|
||||
|
||||
// Write back.
|
||||
int32_t const idx = smem_nhw + thread_in_cta_nhw + i * PIXELS_PER_LDG;
|
||||
if ((unsigned) idx < params.nhw && is_valid_c)
|
||||
{
|
||||
stgStream(&gmem_dst[idx * stride_c_output], x_math, int8_out_scale);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
} // blockIdx.y loop
|
||||
}
|
||||
|
||||
template <typename Kernel_params>
|
||||
dim3 estimateInGridDim(InstanceNormFwdParams const& params)
|
||||
{
|
||||
dim3 grid_dim;
|
||||
grid_dim.x = divUp(params.nhw, Kernel_params::MIN_PIXELS_PER_CTA); // PIXELS_PER_CTA
|
||||
grid_dim.y = divUp(params.c, Kernel_params::C_ELEMENTS_PER_CTA) * params.n;
|
||||
grid_dim.z = 1; // params.n;
|
||||
|
||||
return grid_dim;
|
||||
}
|
||||
|
||||
template <typename Kernel_params>
|
||||
void instanceNormBufferSizes(
|
||||
InstanceNormFwdParams const& params, size_t& size_sums, size_t& size_counts, size_t& size_retired_ctas)
|
||||
{
|
||||
dim3 grid_dim = estimateInGridDim<Kernel_params>(params);
|
||||
|
||||
size_sums = grid_dim.z * grid_dim.y * grid_dim.x * Kernel_params::THREADS_PER_PIXEL
|
||||
* Kernel_params::ELEMENTS_PER_LDG * 2 * sizeof(GMEM_SUMS_TYPE);
|
||||
size_counts = grid_dim.z * grid_dim.y * grid_dim.x * sizeof(int32_t);
|
||||
size_retired_ctas = grid_dim.z * grid_dim.y * sizeof(int32_t);
|
||||
|
||||
size_sums = divUp(size_sums, 256) * 256;
|
||||
size_counts = divUp(size_counts, 256) * 256;
|
||||
size_retired_ctas = divUp(size_retired_ctas, 256) * 256;
|
||||
}
|
||||
|
||||
template <typename Kernel_params>
|
||||
int32_t instance_norm_fwd_launch(
|
||||
InstanceNormFwdContext const& context, InstanceNormFwdParams& params, cudaStream_t stream)
|
||||
{
|
||||
|
||||
size_t smem_size = Kernel_params::PIXELS_PER_THREAD_IN_SMEM * Kernel_params::THREADS_PER_CTA
|
||||
* Kernel_params::ELEMENTS_PER_LDG * sizeof(typename Kernel_params::StorageType);
|
||||
|
||||
dim3 grid_dim = estimateInGridDim<Kernel_params>(params);
|
||||
|
||||
params.c_blks = divUp(params.c, Kernel_params::C_ELEMENTS_PER_CTA);
|
||||
|
||||
size_t size_retired_ctas = grid_dim.z * grid_dim.y * sizeof(int32_t);
|
||||
|
||||
#define KERNEL_RUN(OUTER_LOOPS, DESIRED_OCCUPANCY) \
|
||||
{ \
|
||||
PLUGIN_CHECK_CUDA(cudaMemsetAsync(params.gmem_retired_ctas, 0, size_retired_ctas, stream)); \
|
||||
if (smem_size > 0) \
|
||||
PLUGIN_CHECK_CUDA(cudaFuncSetAttribute( \
|
||||
instanceNormFwd<typename Kernel_params::StorageType, typename Kernel_params::Input_Data_Type, \
|
||||
typename Kernel_params::Output_Data_Type, Kernel_params::THREADS_PER_CTA, \
|
||||
Kernel_params::THREADS_PER_PIXEL, Kernel_params::PIXELS_PER_THREAD_IN_REGISTERS, \
|
||||
Kernel_params::PIXELS_PER_THREAD_IN_SMEM, Kernel_params::ELEMENTS_PER_LDG, \
|
||||
Kernel_params::USE_ONLINE_APPROACH, OUTER_LOOPS, DESIRED_OCCUPANCY>, \
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
|
||||
instanceNormFwd<typename Kernel_params::StorageType, typename Kernel_params::Input_Data_Type, \
|
||||
typename Kernel_params::Output_Data_Type, Kernel_params::THREADS_PER_CTA, \
|
||||
Kernel_params::THREADS_PER_PIXEL, Kernel_params::PIXELS_PER_THREAD_IN_REGISTERS, \
|
||||
Kernel_params::PIXELS_PER_THREAD_IN_SMEM, Kernel_params::ELEMENTS_PER_LDG, \
|
||||
Kernel_params::USE_ONLINE_APPROACH, OUTER_LOOPS, DESIRED_OCCUPANCY> \
|
||||
<<<grid_dim, Kernel_params::THREADS_PER_CTA, smem_size, stream>>>(params); \
|
||||
}
|
||||
|
||||
size_t total_smem_bytes
|
||||
= smem_size + Kernel_params::ELEMENTS_PER_LDG * Kernel_params::THREADS_PER_CTA * sizeof(float);
|
||||
int32_t smem_driven_fwd_occupancy = min(int32_t(context.sm_shared_size) / (int32_t) total_smem_bytes, (int32_t) 2);
|
||||
int32_t max_grid = context.sm_count * smem_driven_fwd_occupancy;
|
||||
if ((context.sm_version >= 700) && (context.sm_version < 800))
|
||||
{
|
||||
max_grid = max_grid - 4;
|
||||
}
|
||||
|
||||
if (max_grid / int32_t(grid_dim.x) > 1)
|
||||
{
|
||||
grid_dim.y = max_grid / int32_t(grid_dim.x);
|
||||
grid_dim.y = int32_t(grid_dim.y) > params.c_blks * params.n ? params.c_blks * params.n : int32_t(grid_dim.y);
|
||||
}
|
||||
else
|
||||
{
|
||||
grid_dim.y = 1;
|
||||
}
|
||||
|
||||
int32_t loop = 1;
|
||||
if (int32_t(grid_dim.x) <= max_grid)
|
||||
{
|
||||
if (smem_driven_fwd_occupancy >= 2)
|
||||
{
|
||||
KERNEL_RUN(1, 2);
|
||||
}
|
||||
else
|
||||
{
|
||||
KERNEL_RUN(1, 1);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
grid_dim.x = max_grid;
|
||||
int32_t nhw_in_regs
|
||||
= params.nhw - Kernel_params::PIXELS_PER_THREAD_IN_SMEM * Kernel_params::PIXELS_PER_LDG * grid_dim.x;
|
||||
int32_t pixels_per_iteration
|
||||
= Kernel_params::PIXELS_PER_THREAD_IN_REGISTERS * Kernel_params::PIXELS_PER_LDG * grid_dim.x;
|
||||
nhw_in_regs = (nhw_in_regs <= 0) ? pixels_per_iteration : nhw_in_regs;
|
||||
if (nhw_in_regs < 0)
|
||||
{
|
||||
nhw_in_regs = pixels_per_iteration;
|
||||
// make PIXELS_PER_THREAD_IN_SMEM <= PIXELS_PER_THREAD_IN_REGISTERS if the assert fails
|
||||
assert(pixels_per_iteration >= params.nhw);
|
||||
}
|
||||
|
||||
loop = divUp(nhw_in_regs, pixels_per_iteration);
|
||||
params.outer_loops = loop;
|
||||
assert(loop >= 1);
|
||||
|
||||
if (loop == 1)
|
||||
{
|
||||
if (smem_driven_fwd_occupancy >= 2)
|
||||
{
|
||||
KERNEL_RUN(1, 2);
|
||||
}
|
||||
else
|
||||
{
|
||||
KERNEL_RUN(1, 1);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
if (smem_driven_fwd_occupancy >= 2)
|
||||
{
|
||||
KERNEL_RUN(0, 2);
|
||||
}
|
||||
else
|
||||
{
|
||||
KERNEL_RUN(0, 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
return loop;
|
||||
}
|
||||
|
||||
static int32_t c_cond_g = 32;
|
||||
|
||||
void instanceNormBufferSizesDispatch(InstanceNormFwdContext const& context, InstanceNormFwdParams const& params,
|
||||
size_t& size_sums, size_t& size_counts, size_t& size_retired_ctas, int32_t input_data_type,
|
||||
int32_t output_data_type)
|
||||
{
|
||||
if (input_data_type == 2 && output_data_type == 2)
|
||||
{
|
||||
switch (context.sm_version)
|
||||
{
|
||||
case 700:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_700>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
case 720:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_720>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
case 750:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_750>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
case 800:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_800>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
case 860:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_860>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
case 870:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8_sm_870>(
|
||||
params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
default:
|
||||
return instanceNormBufferSizes<kernel_params_32_int8>(params, size_sums, size_counts, size_retired_ctas);
|
||||
break;
|
||||
}
|
||||
return instanceNormBufferSizes<kernel_params_32_int8>(params, size_sums, size_counts, size_retired_ctas);
|
||||
}
|
||||
else if (input_data_type == 1 && output_data_type == 2)
|
||||
{
|
||||
return instanceNormBufferSizes<kernel_params_32_fp16_int8>(params, size_sums, size_counts, size_retired_ctas);
|
||||
}
|
||||
else if (input_data_type == 1 && output_data_type == 1)
|
||||
{
|
||||
if (params.c <= c_cond_g)
|
||||
{
|
||||
return instanceNormBufferSizes<kernel_params_32>(params, size_sums, size_counts, size_retired_ctas);
|
||||
}
|
||||
else
|
||||
{
|
||||
return instanceNormBufferSizes<kernel_params_64>(params, size_sums, size_counts, size_retired_ctas);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
fprintf(stderr, "Unsupported format combination by the instance norm kernel\n");
|
||||
assert(0);
|
||||
}
|
||||
}
|
||||
|
||||
int32_t instanceNormFwdDispatch(InstanceNormFwdContext const& context, InstanceNormFwdParams& params,
|
||||
cudaStream_t stream, int32_t input_data_type, int32_t output_data_type)
|
||||
{
|
||||
assert(context.sm_version >= 600);
|
||||
if (input_data_type == 2 && output_data_type == 2)
|
||||
{
|
||||
switch (context.sm_version)
|
||||
{
|
||||
case 700: return instance_norm_fwd_launch<kernel_params_32_int8_sm_700>(context, params, stream); break;
|
||||
case 720: return instance_norm_fwd_launch<kernel_params_32_int8_sm_720>(context, params, stream); break;
|
||||
case 750: return instance_norm_fwd_launch<kernel_params_32_int8_sm_750>(context, params, stream); break;
|
||||
case 800: return instance_norm_fwd_launch<kernel_params_32_int8_sm_800>(context, params, stream); break;
|
||||
case 860: return instance_norm_fwd_launch<kernel_params_32_int8_sm_860>(context, params, stream); break;
|
||||
case 870: return instance_norm_fwd_launch<kernel_params_32_int8_sm_870>(context, params, stream); break;
|
||||
default: return instance_norm_fwd_launch<kernel_params_32_int8>(context, params, stream); break;
|
||||
}
|
||||
}
|
||||
else if (input_data_type == 1 && output_data_type == 2)
|
||||
{
|
||||
return instance_norm_fwd_launch<kernel_params_32_fp16_int8>(context, params, stream);
|
||||
}
|
||||
else if (input_data_type == 1 && output_data_type == 1)
|
||||
{
|
||||
if (params.c <= c_cond_g)
|
||||
{
|
||||
return instance_norm_fwd_launch<kernel_params_32>(context, params, stream);
|
||||
}
|
||||
else
|
||||
{
|
||||
return instance_norm_fwd_launch<kernel_params_64>(context, params, stream);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
fprintf(stderr, "Unsupported format combination by the instance norm kernel\n");
|
||||
assert(0);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // namespace instance_norm_impl
|
||||
@@ -0,0 +1,713 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
|
||||
*/
|
||||
#include "common/checkMacrosPlugin.h"
|
||||
#include "instanceNormalizationPlugin.h"
|
||||
#include "instanceNormCommon.h"
|
||||
#include <algorithm>
|
||||
#include <cuda_fp16.h>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string_view>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::pluginInternal;
|
||||
using namespace instance_norm_impl;
|
||||
using nvinfer1::plugin::InstanceNormalizationV3Plugin;
|
||||
using nvinfer1::plugin::InstanceNormalizationV3PluginCreator;
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr char const* gInstancePluginVersion{"3"};
|
||||
constexpr char const* gInstancePluginName{"InstanceNormalization_TRT"};
|
||||
} // namespace
|
||||
|
||||
InstanceNormalizationV3Plugin::InstanceNormalizationV3Plugin(
|
||||
float epsilon, std::vector<float> const& scale, std::vector<float> const& bias, int32_t relu, float alpha)
|
||||
: mEpsilon(epsilon)
|
||||
, mAlpha(alpha)
|
||||
, mRelu(relu)
|
||||
, mNchan(scale.size())
|
||||
, mHostScale(scale)
|
||||
, mHostBias(bias)
|
||||
{
|
||||
PLUGIN_VALIDATE(scale.size() == bias.size());
|
||||
}
|
||||
|
||||
InstanceNormalizationV3Plugin::InstanceNormalizationV3Plugin(
|
||||
float epsilon, nvinfer1::Weights const& scale, nvinfer1::Weights const& bias, int32_t relu, float alpha)
|
||||
: mEpsilon(epsilon)
|
||||
, mAlpha(alpha)
|
||||
, mRelu(relu)
|
||||
, mNchan(scale.count)
|
||||
{
|
||||
PLUGIN_VALIDATE(scale.count == bias.count);
|
||||
auto const copyWeights = [](nvinfer1::Weights const& input, std::vector<float>& output)
|
||||
{
|
||||
output.reserve(input.count);
|
||||
if (input.type == nvinfer1::DataType::kFLOAT)
|
||||
{
|
||||
output.assign(
|
||||
static_cast<float const*>(input.values), static_cast<float const*>(input.values) + input.count);
|
||||
}
|
||||
else if (input.type == nvinfer1::DataType::kHALF)
|
||||
{
|
||||
for (int32_t c = 0; c < input.count; ++c)
|
||||
{
|
||||
auto const value = static_cast<unsigned short const*>(input.values);
|
||||
output.push_back(__internal_half2float(value[c]));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_ERROR("Unsupported scale/bias dtype");
|
||||
}
|
||||
};
|
||||
|
||||
copyWeights(scale, mHostScale);
|
||||
copyWeights(bias, mHostBias);
|
||||
}
|
||||
|
||||
InstanceNormalizationV3Plugin::~InstanceNormalizationV3Plugin()
|
||||
{
|
||||
exitContext();
|
||||
}
|
||||
|
||||
// InstanceNormalizationV3Plugin returns one output.
|
||||
int32_t InstanceNormalizationV3Plugin::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
IPluginCapability* InstanceNormalizationV3Plugin::getCapabilityInterface(PluginCapabilityType type) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
if (type == PluginCapabilityType::kBUILD)
|
||||
{
|
||||
return static_cast<IPluginV3OneBuild*>(this);
|
||||
}
|
||||
if (type == PluginCapabilityType::kRUNTIME)
|
||||
{
|
||||
return static_cast<IPluginV3OneRuntime*>(this);
|
||||
}
|
||||
PLUGIN_ASSERT(type == PluginCapabilityType::kCORE);
|
||||
return static_cast<IPluginV3OneCore*>(this);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::initializeContext()
|
||||
{
|
||||
if (!mInitialized)
|
||||
{
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreate(&mCudnnHandle));
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mBDescriptor));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mXDescriptor));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mYDescriptor));
|
||||
|
||||
// NDHWC path
|
||||
// Device info.
|
||||
int32_t device;
|
||||
PLUGIN_CUASSERT(cudaGetDevice(&device));
|
||||
cudaDeviceProp props;
|
||||
PLUGIN_CUASSERT(cudaGetDeviceProperties(&props, device));
|
||||
|
||||
mContext.sm_count = props.multiProcessorCount;
|
||||
mContext.sm_shared_size = props.sharedMemPerMultiprocessor;
|
||||
mContext.sm_version = props.major * 100 + props.minor * 10;
|
||||
|
||||
PLUGIN_CUASSERT(cudaMalloc(&mDeviceScale, mNchan * sizeof(float)));
|
||||
PLUGIN_ASSERT(mDeviceScale != nullptr);
|
||||
PLUGIN_CUASSERT(cudaMalloc(&mDeviceBias, mNchan * sizeof(float)));
|
||||
PLUGIN_ASSERT(mDeviceBias != nullptr);
|
||||
PLUGIN_CUASSERT(cudaMemcpy(mDeviceScale, mHostScale.data(), mNchan * sizeof(float), cudaMemcpyHostToDevice));
|
||||
PLUGIN_CUASSERT(cudaMemcpy(mDeviceBias, mHostBias.data(), mNchan * sizeof(float), cudaMemcpyHostToDevice));
|
||||
|
||||
PLUGIN_CUASSERT(cudaDriverGetVersion(&mCudaDriverVersion));
|
||||
}
|
||||
mInitialized = true;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
void InstanceNormalizationV3Plugin::exitContext()
|
||||
{
|
||||
if (mInitialized)
|
||||
{
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mYDescriptor));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mXDescriptor));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mBDescriptor));
|
||||
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroy(mCudnnHandle));
|
||||
|
||||
PLUGIN_CUASSERT(cudaFree(mDeviceBias));
|
||||
PLUGIN_CUASSERT(cudaFree(mDeviceScale));
|
||||
}
|
||||
mInitialized = false;
|
||||
}
|
||||
|
||||
size_t InstanceNormalizationV3Plugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputs[0].desc.dims;
|
||||
PLUGIN_ASSERT(input_dims.nbDims == 4 || input_dims.nbDims == 5);
|
||||
|
||||
if (inputs[0].desc.format == nvinfer1::PluginFormat::kLINEAR)
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputs[0].desc.dims;
|
||||
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
size_t scale_size = n * nchan_bytes;
|
||||
size_t bias_size = n * nchan_bytes;
|
||||
|
||||
size_t total_wss = scale_size + bias_size;
|
||||
|
||||
return total_wss;
|
||||
}
|
||||
else if (inputs[0].desc.format == nvinfer1::PluginFormat::kDHWC8 || inputs[0].desc.format == nvinfer1::PluginFormat::kCDHW32)
|
||||
{
|
||||
PLUGIN_ASSERT(input_dims.nbDims == 5);
|
||||
int32_t input_data_type = (inputs[0].desc.type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
int32_t output_data_type = (outputs[0].desc.type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
nvinfer1::Dims input_dims = inputs[0].desc.dims;
|
||||
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
|
||||
InstanceNormFwdParams params{};
|
||||
// only these parameters are required for workspace computation
|
||||
params.nhw = d * h * w;
|
||||
params.c = c;
|
||||
params.n = n;
|
||||
// Reserve memory for the workspaces.
|
||||
size_t size_sums, size_counts, size_retired_ctas;
|
||||
instanceNormBufferSizesDispatch(
|
||||
mContext, params, size_sums, size_counts, size_retired_ctas, input_data_type, output_data_type);
|
||||
size_t size_nc = n * c * sizeof(float);
|
||||
size_nc = ((size_nc + 256 - 1) / 256) * 256;
|
||||
return size_sums + size_counts + size_retired_ctas + 4 * size_nc;
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_ASSERT(0);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr && workspace != nullptr);
|
||||
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
// early return for empty tensor
|
||||
if (std::any_of(input_dims.d, input_dims.d + input_dims.nbDims, [](int32_t d) { return d == 0; }))
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
auto const callRelu = [this, &stream](void* inOut, int32_t count, nvinfer1::DataType type) {
|
||||
if (mRelu > 0)
|
||||
{
|
||||
int32_t constexpr kBLOCK_SZ = 256;
|
||||
switch (type)
|
||||
{
|
||||
case nvinfer1::DataType::kFLOAT:
|
||||
in3dReluActivation<float, kBLOCK_SZ><<<divUp(count, kBLOCK_SZ), kBLOCK_SZ, 0, stream>>>(
|
||||
static_cast<float*>(inOut), static_cast<float*>(inOut), mAlpha, count);
|
||||
break;
|
||||
case nvinfer1::DataType::kHALF:
|
||||
in3dReluActivation<__half, kBLOCK_SZ><<<divUp(count, kBLOCK_SZ), kBLOCK_SZ, 0, stream>>>(
|
||||
static_cast<__half*>(inOut), static_cast<__half*>(inOut), mAlpha, count);
|
||||
break;
|
||||
default: PLUGIN_ASSERT(0);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if (input_dims.nbDims <= 4)
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t h = input_dims.d[2];
|
||||
int32_t w = input_dims.nbDims > 3 ? input_dims.d[3] : 1;
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
|
||||
float* _d_array = static_cast<float*>(workspace);
|
||||
float* d_scale = &_d_array[0];
|
||||
float* d_bias = &_d_array[n * c];
|
||||
for (int32_t i = 0; i < n; ++i)
|
||||
{
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_scale + i * c, mDeviceScale, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_bias + i * c, mDeviceBias, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
}
|
||||
|
||||
PLUGIN_CUDNNASSERT(
|
||||
mCudnnWrapper.cudnnSetTensor4dDescriptor(mBDescriptor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, n * c, 1, 1));
|
||||
cudnnDataType_t cudnn_dtype{};
|
||||
PLUGIN_CUDNNASSERT(convertTrt2cudnnDtype(inputDesc[0].type, &cudnn_dtype));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetTensor4dDescriptor(mXDescriptor, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetTensor4dDescriptor(mYDescriptor, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w));
|
||||
float alpha = 1;
|
||||
float beta = 0;
|
||||
void const* x_ptr = inputs[0];
|
||||
void* y_ptr = outputs[0];
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetStream(mCudnnHandle, stream));
|
||||
// Note: Use of CUDNN_BATCHNORM_SPATIAL_PERSISTENT can cause numerical
|
||||
// overflows (NaNs) for fp32 data in some circumstances. The lower-
|
||||
// performance CUDNN_BATCHNORM_SPATIAL should be used if this is not
|
||||
// acceptable.
|
||||
|
||||
cudnnBatchNormMode_t cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
|
||||
|
||||
cudaStreamCaptureStatus streamStatus;
|
||||
PLUGIN_CUASSERT(cudaStreamIsCapturing(stream, &streamStatus));
|
||||
|
||||
if (streamStatus != cudaStreamCaptureStatusNone && mCudaDriverVersion < 11000)
|
||||
{
|
||||
gLogVerbose << "Using CUDNN_BATCHNORM_SPATIAL as a CUDA graph capture is in progress but the CUDA version "
|
||||
"may have issues with using CUDNN_BATCHNORM_SPATIAL_PERSISTENT"
|
||||
<< std::endl;
|
||||
cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL;
|
||||
}
|
||||
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnBatchNormalizationForwardTraining(mCudnnHandle, cudnnBatchNormMode,
|
||||
&alpha, &beta, mXDescriptor, x_ptr, mYDescriptor, y_ptr, mBDescriptor, d_scale, d_bias, 1., nullptr,
|
||||
nullptr, mEpsilon, nullptr, nullptr));
|
||||
|
||||
callRelu(y_ptr, n * c * h * w, inputDesc[0].type);
|
||||
}
|
||||
else
|
||||
{
|
||||
if (inputDesc[0].format == nvinfer1::PluginFormat::kLINEAR)
|
||||
{
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetStream(mCudnnHandle, stream));
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
|
||||
// Note: We repeat the data for each batch entry so that we can do the full
|
||||
// computation in a single CUDNN call in enqueue().
|
||||
float* _d_array = (float*) workspace;
|
||||
float* d_scale = &_d_array[0];
|
||||
float* d_bias = &_d_array[n * c];
|
||||
for (int32_t i = 0; i < n; ++i)
|
||||
{
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_scale + i * c, mDeviceScale, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_bias + i * c, mDeviceBias, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
}
|
||||
|
||||
int32_t nc_dimA[] = {1, n * c, 1, 1, 1};
|
||||
int32_t nc_strideA[] = {nc_dimA[1] * nc_dimA[2] * nc_dimA[3] * nc_dimA[4],
|
||||
nc_dimA[2] * nc_dimA[3] * nc_dimA[4], nc_dimA[3] * nc_dimA[4], nc_dimA[4], 1};
|
||||
int32_t img_dimA[] = {1, n * c, d, h, w};
|
||||
int32_t img_strideA[] = {img_dimA[1] * img_dimA[2] * img_dimA[3] * img_dimA[4],
|
||||
img_dimA[2] * img_dimA[3] * img_dimA[4], img_dimA[3] * img_dimA[4], img_dimA[4], 1};
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mBDescriptor, CUDNN_DATA_FLOAT, 5, nc_dimA, nc_strideA));
|
||||
cudnnDataType_t cudnn_dtype;
|
||||
PLUGIN_CHECK_CUDNN(convertTrt2cudnnDtype(inputDesc[0].type, &cudnn_dtype));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mXDescriptor, cudnn_dtype, 5, img_dimA, img_strideA));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mYDescriptor, cudnn_dtype, 5, img_dimA, img_strideA));
|
||||
float alpha = 1;
|
||||
float beta = 0;
|
||||
|
||||
void const* x_ptr = inputs[0];
|
||||
void* y_ptr = outputs[0];
|
||||
// Note: Use of CUDNN_BATCHNORM_SPATIAL_PERSISTENT can cause numerical
|
||||
// overflows (NaNs) for fp32 data in some circumstances. The lower-
|
||||
// performance CUDNN_BATCHNORM_SPATIAL should be used if this is not
|
||||
// acceptable.
|
||||
|
||||
cudnnBatchNormMode_t cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
|
||||
|
||||
cudaStreamCaptureStatus streamStatus;
|
||||
PLUGIN_CUASSERT(cudaStreamIsCapturing(stream, &streamStatus));
|
||||
|
||||
if (streamStatus != cudaStreamCaptureStatusNone && mCudaDriverVersion < 11000)
|
||||
{
|
||||
gLogVerbose
|
||||
<< "Using CUDNN_BATCHNORM_SPATIAL as a CUDA graph capture is in progress but the CUDA version "
|
||||
"may have issues with using CUDNN_BATCHNORM_SPATIAL_PERSISTENT"
|
||||
<< std::endl;
|
||||
cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL;
|
||||
}
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnBatchNormalizationForwardTraining(mCudnnHandle, cudnnBatchNormMode,
|
||||
&alpha, &beta, mXDescriptor, x_ptr, mYDescriptor, y_ptr, mBDescriptor, d_scale, d_bias, 1., nullptr,
|
||||
nullptr, mEpsilon, nullptr, nullptr));
|
||||
|
||||
callRelu(y_ptr, n * c * d * h * w, inputDesc[0].type);
|
||||
}
|
||||
else if (inputDesc[0].format == nvinfer1::PluginFormat::kDHWC8
|
||||
|| inputDesc[0].format == nvinfer1::PluginFormat::kCDHW32)
|
||||
{
|
||||
int32_t input_data_type = (inputDesc[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
int32_t output_data_type = (outputDesc[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
|
||||
InstanceNormFwdParams params{};
|
||||
params.nhw = d * h * w;
|
||||
params.c = c;
|
||||
params.n = n;
|
||||
|
||||
size_t size_sums, size_counts, size_retired_ctas;
|
||||
instanceNormBufferSizesDispatch(
|
||||
mContext, params, size_sums, size_counts, size_retired_ctas, input_data_type, output_data_type);
|
||||
|
||||
size_t size_nc = n * c * sizeof(float);
|
||||
size_nc = ((size_nc + 256 - 1) / 256) * 256;
|
||||
|
||||
char* d_buf = static_cast<char*>(workspace);
|
||||
|
||||
params.gmem_sums = reinterpret_cast<GMEM_SUMS_TYPE*>(d_buf);
|
||||
d_buf += size_sums;
|
||||
params.gmem_counts = reinterpret_cast<int32_t*>(d_buf);
|
||||
d_buf += size_counts;
|
||||
params.gmem_retired_ctas = reinterpret_cast<int32_t*>(d_buf);
|
||||
d_buf += size_retired_ctas;
|
||||
params.gmem_running_mean = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_running_var = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_saved_mean = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_saved_var = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
|
||||
params.gmem_src = inputs[0];
|
||||
params.gmem_dst = outputs[0];
|
||||
params.gmem_bias = mDeviceBias;
|
||||
params.gmem_scale = mDeviceScale;
|
||||
|
||||
params.var_eps = mEpsilon;
|
||||
params.exp_avg_factor = 1.F; //(float)exp_avg_factor;
|
||||
params.use_relu = mRelu; // use_relu;
|
||||
params.relu_alpha = mAlpha; // relu_alpha;
|
||||
|
||||
params.in_scale = inputDesc[0].scale;
|
||||
PLUGIN_ASSERT(outputDesc[0].scale != 0.F);
|
||||
params.out_scale = 1.F / outputDesc[0].scale;
|
||||
|
||||
instanceNormFwdDispatch(mContext, params, stream, input_data_type, output_data_type);
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_FAIL("Unexpected input format");
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
bool InstanceNormalizationV3Plugin::supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(inOut && pos < (nbInputs + nbOutputs));
|
||||
PLUGIN_ASSERT(pos == 0 || pos == 1);
|
||||
|
||||
// For 4-D or 3-D tensor (nbSpatialDims == 1 or 2), only FP32_Linear and FP16_Linear are supported.
|
||||
// For 5-D tensor (nbSpatialDims == 3), FP32_Linear, FP16_Linear, FP16_DHWC8, and INT8_CDHW32 are supported.
|
||||
// This is because we have special InstanceNorm3D kernels for vectorized formats from MLPerf-Inference.
|
||||
|
||||
int32_t const nbDims = inOut[pos].desc.dims.nbDims;
|
||||
PLUGIN_ASSERT(nbDims >= 3);
|
||||
PLUGIN_ASSERT(nbDims <= 5);
|
||||
bool const is3DInstanceNorm = (nbDims == 5);
|
||||
|
||||
bool const isFP32Linear
|
||||
= (inOut[pos].desc.type == nvinfer1::DataType::kFLOAT && inOut[pos].desc.format == nvinfer1::PluginFormat::kLINEAR
|
||||
&& inOut[pos].desc.type == inOut[0].desc.type && inOut[pos].desc.format == inOut[0].desc.format);
|
||||
|
||||
bool const isFP16Linear
|
||||
= (inOut[pos].desc.type == nvinfer1::DataType::kHALF && inOut[pos].desc.format == nvinfer1::PluginFormat::kLINEAR
|
||||
&& inOut[pos].desc.type == inOut[0].desc.type && inOut[pos].desc.format == inOut[0].desc.format);
|
||||
|
||||
bool const isFP16DHWC8
|
||||
= (inOut[pos].desc.type == nvinfer1::DataType::kHALF && inOut[pos].desc.format == nvinfer1::PluginFormat::kDHWC8
|
||||
&& inOut[pos].desc.type == inOut[0].desc.type && inOut[pos].desc.format == inOut[0].desc.format);
|
||||
|
||||
bool const isINT8CDHW32
|
||||
= (inOut[pos].desc.type == nvinfer1::DataType::kINT8 && inOut[pos].desc.format == nvinfer1::PluginFormat::kCDHW32
|
||||
&& inOut[pos].desc.type == inOut[0].desc.type && inOut[pos].desc.format == inOut[0].desc.format);
|
||||
|
||||
bool const isFormatOK = isFP32Linear || isFP16Linear || (is3DInstanceNorm && (isFP16DHWC8 || isINT8CDHW32));
|
||||
|
||||
// Kernels for vectorized formats only support the case of C % spv == 0.
|
||||
int32_t spv{1};
|
||||
switch (inOut[pos].desc.format)
|
||||
{
|
||||
case nvinfer1::PluginFormat::kDHWC8: spv = 8; break;
|
||||
case nvinfer1::PluginFormat::kCDHW32: spv = 32; break;
|
||||
default: break;
|
||||
}
|
||||
int32_t const isAlignmentOK = (inOut[pos].desc.dims.d[1] % spv == 0);
|
||||
|
||||
return isFormatOK && isAlignmentOK;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3Plugin::getPluginName() const noexcept
|
||||
{
|
||||
return gInstancePluginName;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3Plugin::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersion;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3Plugin::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
InstanceNormalizationV3Plugin* InstanceNormalizationV3Plugin::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
auto plugin = std::make_unique<InstanceNormalizationV3Plugin>(mEpsilon, mHostScale, mHostBias, mRelu, mAlpha);
|
||||
plugin->setPluginNamespace(mPluginNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// Set plugin namespace
|
||||
void InstanceNormalizationV3Plugin::setPluginNamespace(char const* pluginNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(pluginNamespace != nullptr);
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(inputTypes != nullptr);
|
||||
PLUGIN_ASSERT(nbInputs == 1);
|
||||
PLUGIN_ASSERT(nbOutputs == 1);
|
||||
outputTypes[0] = inputTypes[0];
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
|
||||
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(inputs != nullptr);
|
||||
PLUGIN_ASSERT(nbInputs == 1);
|
||||
PLUGIN_ASSERT(nbOutputs == 1);
|
||||
outputs[0] = inputs[0];
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
||||
IPluginV3* InstanceNormalizationV3Plugin::attachToContext(IPluginResourceContext* context) noexcept
|
||||
{
|
||||
InstanceNormalizationV3Plugin* obj = clone();
|
||||
obj->initializeContext();
|
||||
return obj;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
||||
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
||||
{
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationV3Plugin::onShapeChange(PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(in != nullptr);
|
||||
PLUGIN_ASSERT(out != nullptr);
|
||||
PLUGIN_ASSERT(nbOutputs == 1);
|
||||
PLUGIN_ASSERT(nbInputs == 1);
|
||||
// Not support dynamic shape in C dimension
|
||||
PLUGIN_ASSERT(in[0].dims.d[1] != -1);
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* InstanceNormalizationV3Plugin::getFieldsToSerialize() noexcept
|
||||
{
|
||||
mDataToSerialize.clear();
|
||||
mDataToSerialize.emplace_back("epsilon", &mEpsilon, PluginFieldType::kFLOAT32, 1);
|
||||
mDataToSerialize.emplace_back("scales", mHostScale.data(), PluginFieldType::kFLOAT32, mHostScale.size());
|
||||
mDataToSerialize.emplace_back("bias", mHostBias.data(), PluginFieldType::kFLOAT32, mHostBias.size());
|
||||
mDataToSerialize.emplace_back("relu", &mRelu, PluginFieldType::kINT32, 1);
|
||||
mDataToSerialize.emplace_back("alpha", &mAlpha, PluginFieldType::kFLOAT32, 1);
|
||||
|
||||
mFCToSerialize.nbFields = mDataToSerialize.size();
|
||||
mFCToSerialize.fields = mDataToSerialize.data();
|
||||
return &mFCToSerialize;
|
||||
}
|
||||
|
||||
|
||||
// InstanceNormalizationV3PluginCreator methods
|
||||
InstanceNormalizationV3PluginCreator::InstanceNormalizationV3PluginCreator()
|
||||
{
|
||||
static std::mutex sMutex;
|
||||
std::lock_guard<std::mutex> guard(sMutex);
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("epsilon", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("scales", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("relu", nullptr, PluginFieldType::kINT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3PluginCreator::getPluginName() const noexcept
|
||||
{
|
||||
return gInstancePluginName;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3PluginCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersion;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* InstanceNormalizationV3PluginCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV3* InstanceNormalizationV3PluginCreator::createPlugin(
|
||||
char const* name, nvinfer1::PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
std::vector<float> scaleValues;
|
||||
std::vector<float> biasValues;
|
||||
float epsilon{};
|
||||
int32_t relu{};
|
||||
float alpha{};
|
||||
PluginField const* fields = fc->fields;
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i)
|
||||
{
|
||||
std::string_view const attrName = fields[i].name;
|
||||
if (attrName == "epsilon"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
epsilon = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
else if (attrName == "scales"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
int32_t size = fields[i].length;
|
||||
scaleValues.reserve(size);
|
||||
auto const* w = static_cast<float const*>(fields[i].data);
|
||||
for (int32_t j = 0; j < size; j++)
|
||||
{
|
||||
scaleValues.push_back(*w);
|
||||
w++;
|
||||
}
|
||||
}
|
||||
else if (attrName == "bias"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
int32_t size = fields[i].length;
|
||||
biasValues.reserve(size);
|
||||
auto const* w = static_cast<float const*>(fields[i].data);
|
||||
for (int32_t j = 0; j < size; j++)
|
||||
{
|
||||
biasValues.push_back(*w);
|
||||
w++;
|
||||
}
|
||||
}
|
||||
else if (attrName == "relu"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
relu = *(static_cast<int32_t const*>(fields[i].data));
|
||||
}
|
||||
else if (attrName == "alpha"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
alpha = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
}
|
||||
|
||||
Weights scaleWeights{DataType::kFLOAT, scaleValues.data(), (int64_t) scaleValues.size()};
|
||||
Weights biasWeights{DataType::kFLOAT, biasValues.data(), (int64_t) biasValues.size()};
|
||||
|
||||
auto obj = std::make_unique<InstanceNormalizationV3Plugin>(epsilon, scaleWeights, biasWeights, relu, alpha);
|
||||
obj->setPluginNamespace(mNamespace.c_str());
|
||||
return obj.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void InstanceNormalizationV3PluginCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(libNamespace != nullptr);
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationV3PluginCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
#ifndef TRT_INSTANCE_NORMALIZATION_PLUGIN_H
|
||||
#define TRT_INSTANCE_NORMALIZATION_PLUGIN_H
|
||||
#include "NvInfer.h"
|
||||
#include "NvInferPlugin.h"
|
||||
#include "common/plugin.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include "instanceNormalizationPlugin/instanceNormFwd.h"
|
||||
#include <cuda_fp16.h>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
typedef uint16_t half_type;
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr char const* gInstancePluginFullNameV3{"InstanceNormalization_TRT, version:3"};
|
||||
} // namespace
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
class InstanceNormalizationV3Plugin : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
|
||||
public:
|
||||
InstanceNormalizationV3Plugin(float epsilon, nvinfer1::Weights const& scale, nvinfer1::Weights const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F);
|
||||
InstanceNormalizationV3Plugin(float epsilon, std::vector<float> const& scale, std::vector<float> const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F);
|
||||
InstanceNormalizationV3Plugin(void const* serialData, size_t serialLength);
|
||||
|
||||
InstanceNormalizationV3Plugin() = delete;
|
||||
|
||||
InstanceNormalizationV3Plugin(InstanceNormalizationV3Plugin const&) = default;
|
||||
|
||||
~InstanceNormalizationV3Plugin() override;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
|
||||
InstanceNormalizationV3Plugin* clone() noexcept override;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
|
||||
void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
|
||||
|
||||
// DynamicExt plugin supportsFormat update.
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept;
|
||||
|
||||
int32_t getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
|
||||
int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
|
||||
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out,
|
||||
int32_t nbOutputs) noexcept override;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
|
||||
int32_t initializeContext();
|
||||
|
||||
protected:
|
||||
void exitContext();
|
||||
|
||||
private:
|
||||
float mEpsilon{};
|
||||
float mAlpha{};
|
||||
int32_t mRelu{};
|
||||
int32_t mNchan{};
|
||||
std::vector<float> mHostScale;
|
||||
std::vector<float> mHostBias;
|
||||
float* mDeviceScale{nullptr};
|
||||
float* mDeviceBias{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnHandle_t mCudnnHandle{nullptr};
|
||||
nvinfer1::pluginInternal::CudnnWrapper& mCudnnWrapper
|
||||
= nvinfer1::pluginInternal::getCudnnWrapper(gInstancePluginFullNameV3);
|
||||
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mXDescriptor{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mYDescriptor{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mBDescriptor{nullptr};
|
||||
std::string mPluginNamespace;
|
||||
bool mInitialized{false};
|
||||
int32_t mCudaDriverVersion{-1};
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
|
||||
// NDHWC implementation
|
||||
instance_norm_impl::InstanceNormFwdContext mContext;
|
||||
};
|
||||
|
||||
class InstanceNormalizationV3PluginCreator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
InstanceNormalizationV3PluginCreator();
|
||||
|
||||
~InstanceNormalizationV3PluginCreator() override = default;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#endif // TRT_INSTANCE_NORMALIZATION_PLUGIN_H
|
||||
@@ -0,0 +1,729 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 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.
|
||||
*/
|
||||
#include "common/checkMacrosPlugin.h"
|
||||
#include "instanceNormalizationPluginLegacy.h"
|
||||
#include "instanceNormCommon.h"
|
||||
#include <algorithm>
|
||||
#include <cuda_fp16.h>
|
||||
#include <memory>
|
||||
#include <stdexcept>
|
||||
#include <string_view>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::pluginInternal;
|
||||
using namespace instance_norm_impl;
|
||||
using nvinfer1::plugin::InstanceNormalizationPlugin;
|
||||
using nvinfer1::plugin::InstanceNormalizationPluginV2;
|
||||
using nvinfer1::plugin::InstanceNormalizationPluginCreator;
|
||||
using nvinfer1::plugin::InstanceNormalizationPluginCreatorV2;
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr char const* gInstancePluginVersion{"1"};
|
||||
constexpr char const* gInstancePluginVersionV2{"2"};
|
||||
constexpr char const* gInstancePluginName{"InstanceNormalization_TRT"};
|
||||
} // namespace
|
||||
|
||||
InstanceNormalizationPlugin::InstanceNormalizationPlugin(
|
||||
float epsilon, std::vector<float> const& scale, std::vector<float> const& bias, int32_t relu, float alpha)
|
||||
: mEpsilon(epsilon)
|
||||
, mAlpha(alpha)
|
||||
, mRelu(relu)
|
||||
, mNchan(scale.size())
|
||||
, mHostScale(scale)
|
||||
, mHostBias(bias)
|
||||
{
|
||||
PLUGIN_VALIDATE(scale.size() == bias.size());
|
||||
}
|
||||
|
||||
InstanceNormalizationPlugin::InstanceNormalizationPlugin(
|
||||
float epsilon, nvinfer1::Weights const& scale, nvinfer1::Weights const& bias, int32_t relu, float alpha)
|
||||
: mEpsilon(epsilon)
|
||||
, mAlpha(alpha)
|
||||
, mRelu(relu)
|
||||
, mNchan(scale.count)
|
||||
{
|
||||
PLUGIN_VALIDATE(scale.count == bias.count);
|
||||
auto const copyWeights = [](nvinfer1::Weights const& input, std::vector<float>& output)
|
||||
{
|
||||
output.reserve(input.count);
|
||||
if (input.type == nvinfer1::DataType::kFLOAT)
|
||||
{
|
||||
output.assign(
|
||||
static_cast<float const*>(input.values), static_cast<float const*>(input.values) + input.count);
|
||||
}
|
||||
else if (input.type == nvinfer1::DataType::kHALF)
|
||||
{
|
||||
for (int32_t c = 0; c < input.count; ++c)
|
||||
{
|
||||
auto const value = static_cast<unsigned short const*>(input.values);
|
||||
output.push_back(__internal_half2float(value[c]));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_ERROR("Unsupported scale/bias dtype");
|
||||
}
|
||||
};
|
||||
|
||||
copyWeights(scale, mHostScale);
|
||||
copyWeights(bias, mHostBias);
|
||||
}
|
||||
|
||||
InstanceNormalizationPlugin::InstanceNormalizationPlugin(void const* serialData, size_t serialLength)
|
||||
{
|
||||
deserialize_value(&serialData, &serialLength, &mEpsilon);
|
||||
deserialize_value(&serialData, &serialLength, &mNchan);
|
||||
deserialize_value(&serialData, &serialLength, &mHostScale);
|
||||
deserialize_value(&serialData, &serialLength, &mHostBias);
|
||||
deserialize_value(&serialData, &serialLength, &mRelu);
|
||||
deserialize_value(&serialData, &serialLength, &mAlpha);
|
||||
}
|
||||
|
||||
InstanceNormalizationPlugin::~InstanceNormalizationPlugin()
|
||||
{
|
||||
terminate();
|
||||
}
|
||||
|
||||
// InstanceNormalizationPlugin returns one output.
|
||||
int32_t InstanceNormalizationPlugin::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
DimsExprs InstanceNormalizationPlugin::getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs,
|
||||
int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
nvinfer1::DimsExprs output(inputs[0]);
|
||||
return output;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationPlugin::initialize() noexcept
|
||||
{
|
||||
if (!mInitialized)
|
||||
{
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreate(&mCudnnHandle));
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mBDescriptor));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mXDescriptor));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnCreateTensorDescriptor(&mYDescriptor));
|
||||
|
||||
// NDHWC path
|
||||
// Device info.
|
||||
int32_t device;
|
||||
PLUGIN_CHECK_CUDA(cudaGetDevice(&device));
|
||||
cudaDeviceProp props;
|
||||
PLUGIN_CHECK_CUDA(cudaGetDeviceProperties(&props, device));
|
||||
|
||||
mContext.sm_count = props.multiProcessorCount;
|
||||
mContext.sm_shared_size = props.sharedMemPerMultiprocessor;
|
||||
mContext.sm_version = props.major * 100 + props.minor * 10;
|
||||
|
||||
PLUGIN_CHECK_CUDA(cudaMalloc(&mDeviceScale, mNchan * sizeof(float)));
|
||||
PLUGIN_CHECK_CUDA(cudaMalloc(&mDeviceBias, mNchan * sizeof(float)));
|
||||
PLUGIN_CHECK_CUDA(cudaMemcpy(mDeviceScale, mHostScale.data(), mNchan * sizeof(float), cudaMemcpyHostToDevice));
|
||||
PLUGIN_CHECK_CUDA(cudaMemcpy(mDeviceBias, mHostBias.data(), mNchan * sizeof(float), cudaMemcpyHostToDevice));
|
||||
|
||||
PLUGIN_CHECK_CUDA(cudaDriverGetVersion(&mCudaDriverVersion));
|
||||
}
|
||||
mInitialized = true;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
void InstanceNormalizationPlugin::terminate() noexcept
|
||||
{
|
||||
if (mInitialized)
|
||||
{
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mYDescriptor));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mXDescriptor));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroyTensorDescriptor(mBDescriptor));
|
||||
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnDestroy(mCudnnHandle));
|
||||
|
||||
PLUGIN_CUASSERT(cudaFree(mDeviceBias));
|
||||
PLUGIN_CUASSERT(cudaFree(mDeviceScale));
|
||||
}
|
||||
mInitialized = false;
|
||||
}
|
||||
|
||||
size_t InstanceNormalizationPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputs[0].dims;
|
||||
PLUGIN_ASSERT(input_dims.nbDims == 4 || input_dims.nbDims == 5);
|
||||
|
||||
if (inputs[0].format == nvinfer1::PluginFormat::kLINEAR)
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputs[0].dims;
|
||||
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
size_t scale_size = n * nchan_bytes;
|
||||
size_t bias_size = n * nchan_bytes;
|
||||
|
||||
size_t total_wss = scale_size + bias_size;
|
||||
|
||||
return total_wss;
|
||||
}
|
||||
else if (inputs[0].format == nvinfer1::PluginFormat::kDHWC8 || inputs[0].format == nvinfer1::PluginFormat::kCDHW32)
|
||||
{
|
||||
PLUGIN_ASSERT(input_dims.nbDims == 5);
|
||||
int32_t input_data_type = (inputs[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
int32_t output_data_type = (outputs[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
nvinfer1::Dims input_dims = inputs[0].dims;
|
||||
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
|
||||
InstanceNormFwdParams params{};
|
||||
// only these parameters are required for workspace computation
|
||||
params.nhw = d * h * w;
|
||||
params.c = c;
|
||||
params.n = n;
|
||||
// Reserve memory for the workspaces.
|
||||
size_t size_sums, size_counts, size_retired_ctas;
|
||||
instanceNormBufferSizesDispatch(
|
||||
mContext, params, size_sums, size_counts, size_retired_ctas, input_data_type, output_data_type);
|
||||
size_t size_nc = n * c * sizeof(float);
|
||||
size_nc = ((size_nc + 256 - 1) / 256) * 256;
|
||||
return size_sums + size_counts + size_retired_ctas + 4 * size_nc;
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_ASSERT(0);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t InstanceNormalizationPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
|
||||
nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr && workspace != nullptr);
|
||||
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
// early return for empty tensor
|
||||
if (std::any_of(input_dims.d, input_dims.d + input_dims.nbDims, [](int32_t d) { return d == 0; }))
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
auto const callRelu = [this, &stream](void* inOut, int32_t count, nvinfer1::DataType type) {
|
||||
if (mRelu > 0)
|
||||
{
|
||||
int32_t constexpr kBLOCK_SZ = 256;
|
||||
switch (type)
|
||||
{
|
||||
case nvinfer1::DataType::kFLOAT:
|
||||
in3dReluActivation<float, kBLOCK_SZ><<<divUp(count, kBLOCK_SZ), kBLOCK_SZ, 0, stream>>>(
|
||||
static_cast<float*>(inOut), static_cast<float*>(inOut), mAlpha, count);
|
||||
break;
|
||||
case nvinfer1::DataType::kHALF:
|
||||
in3dReluActivation<__half, kBLOCK_SZ><<<divUp(count, kBLOCK_SZ), kBLOCK_SZ, 0, stream>>>(
|
||||
static_cast<__half*>(inOut), static_cast<__half*>(inOut), mAlpha, count);
|
||||
break;
|
||||
default: PLUGIN_ASSERT(0);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if (input_dims.nbDims <= 4)
|
||||
{
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t h = input_dims.d[2];
|
||||
int32_t w = input_dims.nbDims > 3 ? input_dims.d[3] : 1;
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
|
||||
float* _d_array = static_cast<float*>(workspace);
|
||||
float* d_scale = &_d_array[0];
|
||||
float* d_bias = &_d_array[n * c];
|
||||
for (int32_t i = 0; i < n; ++i)
|
||||
{
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_scale + i * c, mDeviceScale, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
PLUGIN_CUASSERT(
|
||||
cudaMemcpyAsync(d_bias + i * c, mDeviceBias, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
}
|
||||
|
||||
PLUGIN_CUDNNASSERT(
|
||||
mCudnnWrapper.cudnnSetTensor4dDescriptor(mBDescriptor, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, n * c, 1, 1));
|
||||
cudnnDataType_t cudnn_dtype{};
|
||||
PLUGIN_CUDNNASSERT(convertTrt2cudnnDtype(inputDesc[0].type, &cudnn_dtype));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetTensor4dDescriptor(mXDescriptor, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w));
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetTensor4dDescriptor(mYDescriptor, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w));
|
||||
float alpha = 1;
|
||||
float beta = 0;
|
||||
void const* x_ptr = inputs[0];
|
||||
void* y_ptr = outputs[0];
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnSetStream(mCudnnHandle, stream));
|
||||
// Note: Use of CUDNN_BATCHNORM_SPATIAL_PERSISTENT can cause numerical
|
||||
// overflows (NaNs) for fp32 data in some circumstances. The lower-
|
||||
// performance CUDNN_BATCHNORM_SPATIAL should be used if this is not
|
||||
// acceptable.
|
||||
|
||||
cudnnBatchNormMode_t cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
|
||||
|
||||
cudaStreamCaptureStatus streamStatus;
|
||||
PLUGIN_CHECK_CUDA(cudaStreamIsCapturing(stream, &streamStatus));
|
||||
|
||||
if (streamStatus != cudaStreamCaptureStatusNone && mCudaDriverVersion < 11000)
|
||||
{
|
||||
gLogVerbose << "Using CUDNN_BATCHNORM_SPATIAL as a CUDA graph capture is in progress but the CUDA version "
|
||||
"may have issues with using CUDNN_BATCHNORM_SPATIAL_PERSISTENT"
|
||||
<< std::endl;
|
||||
cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL;
|
||||
}
|
||||
|
||||
PLUGIN_CUDNNASSERT(mCudnnWrapper.cudnnBatchNormalizationForwardTraining(mCudnnHandle, cudnnBatchNormMode,
|
||||
&alpha, &beta, mXDescriptor, x_ptr, mYDescriptor, y_ptr, mBDescriptor, d_scale, d_bias, 1., nullptr,
|
||||
nullptr, mEpsilon, nullptr, nullptr));
|
||||
|
||||
callRelu(y_ptr, n * c * h * w, inputDesc[0].type);
|
||||
}
|
||||
else
|
||||
{
|
||||
if (inputDesc[0].format == nvinfer1::PluginFormat::kLINEAR)
|
||||
{
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetStream(mCudnnHandle, stream));
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
size_t nchan_bytes = c * sizeof(float);
|
||||
|
||||
// Note: We repeat the data for each batch entry so that we can do the full
|
||||
// computation in a single CUDNN call in enqueue().
|
||||
float* _d_array = (float*) workspace;
|
||||
float* d_scale = &_d_array[0];
|
||||
float* d_bias = &_d_array[n * c];
|
||||
for (int32_t i = 0; i < n; ++i)
|
||||
{
|
||||
PLUGIN_CHECK_CUDA(
|
||||
cudaMemcpyAsync(d_scale + i * c, mDeviceScale, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
PLUGIN_CHECK_CUDA(
|
||||
cudaMemcpyAsync(d_bias + i * c, mDeviceBias, nchan_bytes, cudaMemcpyDeviceToDevice, stream));
|
||||
}
|
||||
|
||||
int32_t nc_dimA[] = {1, n * c, 1, 1, 1};
|
||||
int32_t nc_strideA[] = {nc_dimA[1] * nc_dimA[2] * nc_dimA[3] * nc_dimA[4],
|
||||
nc_dimA[2] * nc_dimA[3] * nc_dimA[4], nc_dimA[3] * nc_dimA[4], nc_dimA[4], 1};
|
||||
int32_t img_dimA[] = {1, n * c, d, h, w};
|
||||
int32_t img_strideA[] = {img_dimA[1] * img_dimA[2] * img_dimA[3] * img_dimA[4],
|
||||
img_dimA[2] * img_dimA[3] * img_dimA[4], img_dimA[3] * img_dimA[4], img_dimA[4], 1};
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mBDescriptor, CUDNN_DATA_FLOAT, 5, nc_dimA, nc_strideA));
|
||||
cudnnDataType_t cudnn_dtype;
|
||||
PLUGIN_CHECK_CUDNN(convertTrt2cudnnDtype(inputDesc[0].type, &cudnn_dtype));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mXDescriptor, cudnn_dtype, 5, img_dimA, img_strideA));
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnSetTensorNdDescriptor(mYDescriptor, cudnn_dtype, 5, img_dimA, img_strideA));
|
||||
float alpha = 1;
|
||||
float beta = 0;
|
||||
|
||||
void const* x_ptr = inputs[0];
|
||||
void* y_ptr = outputs[0];
|
||||
// Note: Use of CUDNN_BATCHNORM_SPATIAL_PERSISTENT can cause numerical
|
||||
// overflows (NaNs) for fp32 data in some circumstances. The lower-
|
||||
// performance CUDNN_BATCHNORM_SPATIAL should be used if this is not
|
||||
// acceptable.
|
||||
|
||||
cudnnBatchNormMode_t cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL_PERSISTENT;
|
||||
|
||||
cudaStreamCaptureStatus streamStatus;
|
||||
PLUGIN_CHECK_CUDA(cudaStreamIsCapturing(stream, &streamStatus));
|
||||
|
||||
if (streamStatus != cudaStreamCaptureStatusNone && mCudaDriverVersion < 11000)
|
||||
{
|
||||
gLogVerbose
|
||||
<< "Using CUDNN_BATCHNORM_SPATIAL as a CUDA graph capture is in progress but the CUDA version "
|
||||
"may have issues with using CUDNN_BATCHNORM_SPATIAL_PERSISTENT"
|
||||
<< std::endl;
|
||||
cudnnBatchNormMode = CUDNN_BATCHNORM_SPATIAL;
|
||||
}
|
||||
|
||||
PLUGIN_CHECK_CUDNN(mCudnnWrapper.cudnnBatchNormalizationForwardTraining(mCudnnHandle, cudnnBatchNormMode,
|
||||
&alpha, &beta, mXDescriptor, x_ptr, mYDescriptor, y_ptr, mBDescriptor, d_scale, d_bias, 1., nullptr,
|
||||
nullptr, mEpsilon, nullptr, nullptr));
|
||||
|
||||
callRelu(y_ptr, n * c * d * h * w, inputDesc[0].type);
|
||||
}
|
||||
else if (inputDesc[0].format == nvinfer1::PluginFormat::kDHWC8
|
||||
|| inputDesc[0].format == nvinfer1::PluginFormat::kCDHW32)
|
||||
{
|
||||
int32_t input_data_type = (inputDesc[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
int32_t output_data_type = (outputDesc[0].type == nvinfer1::DataType::kHALF) ? 1 : 2;
|
||||
|
||||
nvinfer1::Dims input_dims = inputDesc[0].dims;
|
||||
int32_t n = input_dims.d[0];
|
||||
int32_t c = input_dims.d[1];
|
||||
int32_t d = input_dims.d[2];
|
||||
int32_t h = input_dims.d[3];
|
||||
int32_t w = input_dims.d[4];
|
||||
|
||||
InstanceNormFwdParams params{};
|
||||
params.nhw = d * h * w;
|
||||
params.c = c;
|
||||
params.n = n;
|
||||
|
||||
size_t size_sums, size_counts, size_retired_ctas;
|
||||
instanceNormBufferSizesDispatch(
|
||||
mContext, params, size_sums, size_counts, size_retired_ctas, input_data_type, output_data_type);
|
||||
|
||||
size_t size_nc = n * c * sizeof(float);
|
||||
size_nc = ((size_nc + 256 - 1) / 256) * 256;
|
||||
|
||||
char* d_buf = static_cast<char*>(workspace);
|
||||
|
||||
params.gmem_sums = reinterpret_cast<GMEM_SUMS_TYPE*>(d_buf);
|
||||
d_buf += size_sums;
|
||||
params.gmem_counts = reinterpret_cast<int32_t*>(d_buf);
|
||||
d_buf += size_counts;
|
||||
params.gmem_retired_ctas = reinterpret_cast<int32_t*>(d_buf);
|
||||
d_buf += size_retired_ctas;
|
||||
params.gmem_running_mean = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_running_var = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_saved_mean = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
params.gmem_saved_var = reinterpret_cast<float*>(d_buf);
|
||||
d_buf += size_nc;
|
||||
|
||||
params.gmem_src = inputs[0];
|
||||
params.gmem_dst = outputs[0];
|
||||
params.gmem_bias = mDeviceBias;
|
||||
params.gmem_scale = mDeviceScale;
|
||||
|
||||
params.var_eps = mEpsilon;
|
||||
params.exp_avg_factor = 1.F; //(float)exp_avg_factor;
|
||||
params.use_relu = mRelu; // use_relu;
|
||||
params.relu_alpha = mAlpha; // relu_alpha;
|
||||
|
||||
params.in_scale = inputDesc[0].scale;
|
||||
PLUGIN_ASSERT(outputDesc[0].scale != 0.F);
|
||||
params.out_scale = 1.F / outputDesc[0].scale;
|
||||
|
||||
instanceNormFwdDispatch(mContext, params, stream, input_data_type, output_data_type);
|
||||
}
|
||||
else
|
||||
{
|
||||
PLUGIN_FAIL("Unexpected input format");
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
size_t InstanceNormalizationPlugin::getSerializationSize() const noexcept
|
||||
{
|
||||
return (serialized_size(mEpsilon) + serialized_size(mNchan) + serialized_size(mHostScale)
|
||||
+ serialized_size(mHostBias) + serialized_size(mRelu) + serialized_size(mAlpha));
|
||||
}
|
||||
|
||||
void InstanceNormalizationPlugin::serialize(void* buffer) const noexcept
|
||||
{
|
||||
serialize_value(&buffer, mEpsilon);
|
||||
serialize_value(&buffer, mNchan);
|
||||
serialize_value(&buffer, mHostScale);
|
||||
serialize_value(&buffer, mHostBias);
|
||||
serialize_value(&buffer, mRelu);
|
||||
serialize_value(&buffer, mAlpha);
|
||||
}
|
||||
|
||||
bool InstanceNormalizationPlugin::supportsFormatCombination(
|
||||
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(inOut && pos < (nbInputs + nbOutputs));
|
||||
PLUGIN_ASSERT(pos == 0 || pos == 1);
|
||||
|
||||
// For 4-D or 3-D tensor (nbSpatialDims == 1 or 2), only FP32_Linear and FP16_Linear are supported.
|
||||
// For 5-D tensor (nbSpatialDims == 3), FP32_Linear, FP16_Linear, FP16_DHWC8, and INT8_CDHW32 are supported.
|
||||
// This is because we have special InstanceNorm3D kernels for vectorized formats from MLPerf-Inference.
|
||||
|
||||
int32_t const nbDims = inOut[pos].dims.nbDims;
|
||||
PLUGIN_ASSERT(nbDims >= 3);
|
||||
PLUGIN_ASSERT(nbDims <= 5);
|
||||
bool const is3DInstanceNorm = (nbDims == 5);
|
||||
|
||||
bool const isFP32Linear
|
||||
= (inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR
|
||||
&& inOut[pos].type == inOut[0].type && inOut[pos].format == inOut[0].format);
|
||||
|
||||
bool const isFP16Linear
|
||||
= (inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR
|
||||
&& inOut[pos].type == inOut[0].type && inOut[pos].format == inOut[0].format);
|
||||
|
||||
bool const isFP16DHWC8
|
||||
= (inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == nvinfer1::PluginFormat::kDHWC8
|
||||
&& inOut[pos].type == inOut[0].type && inOut[pos].format == inOut[0].format);
|
||||
|
||||
bool const isINT8CDHW32
|
||||
= (inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == nvinfer1::PluginFormat::kCDHW32
|
||||
&& inOut[pos].type == inOut[0].type && inOut[pos].format == inOut[0].format);
|
||||
|
||||
bool const isFormatOK = isFP32Linear || isFP16Linear || (is3DInstanceNorm && (isFP16DHWC8 || isINT8CDHW32));
|
||||
|
||||
// Kernels for vectorized formats only support the case of C % spv == 0.
|
||||
int32_t spv{1};
|
||||
switch (inOut[pos].format)
|
||||
{
|
||||
case nvinfer1::PluginFormat::kDHWC8: spv = 8; break;
|
||||
case nvinfer1::PluginFormat::kCDHW32: spv = 32; break;
|
||||
default: break;
|
||||
}
|
||||
int32_t const isAlignmentOK = (inOut[pos].dims.d[1] % spv == 0);
|
||||
|
||||
return isFormatOK && isAlignmentOK;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPlugin::getPluginType() const noexcept
|
||||
{
|
||||
return gInstancePluginName;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPlugin::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersion;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPluginV2::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersionV2;
|
||||
}
|
||||
|
||||
void InstanceNormalizationPlugin::destroy() noexcept
|
||||
{
|
||||
delete this;
|
||||
}
|
||||
|
||||
template <class PluginType>
|
||||
IPluginV2DynamicExt* InstanceNormalizationPlugin::cloneBase() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
auto plugin = std::make_unique<PluginType>(mEpsilon, mHostScale, mHostBias, mRelu, mAlpha);
|
||||
plugin->setPluginNamespace(mPluginNamespace.c_str());
|
||||
plugin->initialize();
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPlugin::clone() const noexcept
|
||||
{
|
||||
return cloneBase<InstanceNormalizationPlugin>();
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginV2::clone() const noexcept
|
||||
{
|
||||
return cloneBase<InstanceNormalizationPluginV2>();
|
||||
}
|
||||
|
||||
// Set plugin namespace
|
||||
void InstanceNormalizationPlugin::setPluginNamespace(char const* pluginNamespace) noexcept
|
||||
{
|
||||
mPluginNamespace = pluginNamespace;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPlugin::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mPluginNamespace.c_str();
|
||||
}
|
||||
|
||||
nvinfer1::DataType InstanceNormalizationPlugin::getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(inputTypes && nbInputs > 0 && index == 0);
|
||||
return inputTypes[0];
|
||||
}
|
||||
|
||||
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
||||
void InstanceNormalizationPlugin::attachToContext(
|
||||
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
|
||||
{
|
||||
}
|
||||
|
||||
// Detach the plugin object from its execution context.
|
||||
void InstanceNormalizationPlugin::detachFromContext() noexcept {}
|
||||
|
||||
void InstanceNormalizationPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
||||
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// Not support dynamic shape in C dimension
|
||||
PLUGIN_ASSERT(nbInputs == 1 && in[0].desc.dims.d[1] != -1);
|
||||
}
|
||||
|
||||
// InstanceNormalizationPluginCreator methods
|
||||
InstanceNormalizationPluginCreator::InstanceNormalizationPluginCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("epsilon", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("scales", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("relu", nullptr, PluginFieldType::kINT32, 1));
|
||||
mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32, 1));
|
||||
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPluginCreator::getPluginName() const noexcept
|
||||
{
|
||||
return gInstancePluginName;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPluginCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersion;
|
||||
}
|
||||
|
||||
char const* InstanceNormalizationPluginCreatorV2::getPluginVersion() const noexcept
|
||||
{
|
||||
return gInstancePluginVersionV2;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* InstanceNormalizationPluginCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
template <class PluginType>
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreator::createPluginBase(
|
||||
char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
std::vector<float> scaleValues;
|
||||
std::vector<float> biasValues;
|
||||
float epsilon{};
|
||||
int32_t relu{};
|
||||
float alpha{};
|
||||
PluginField const* fields = fc->fields;
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i)
|
||||
{
|
||||
std::string_view const attrName = fields[i].name;
|
||||
if (attrName == "epsilon"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
epsilon = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
else if (attrName == "scales"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
int32_t size = fields[i].length;
|
||||
scaleValues.reserve(size);
|
||||
auto const* w = static_cast<float const*>(fields[i].data);
|
||||
for (int32_t j = 0; j < size; j++)
|
||||
{
|
||||
scaleValues.push_back(*w);
|
||||
w++;
|
||||
}
|
||||
}
|
||||
else if (attrName == "bias"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
int32_t size = fields[i].length;
|
||||
biasValues.reserve(size);
|
||||
auto const* w = static_cast<float const*>(fields[i].data);
|
||||
for (int32_t j = 0; j < size; j++)
|
||||
{
|
||||
biasValues.push_back(*w);
|
||||
w++;
|
||||
}
|
||||
}
|
||||
else if (attrName == "relu"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
|
||||
relu = *(static_cast<int32_t const*>(fields[i].data));
|
||||
}
|
||||
else if (attrName == "alpha"sv)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
|
||||
alpha = *(static_cast<float const*>(fields[i].data));
|
||||
}
|
||||
}
|
||||
|
||||
Weights scaleWeights{DataType::kFLOAT, scaleValues.data(), (int64_t) scaleValues.size()};
|
||||
Weights biasWeights{DataType::kFLOAT, biasValues.data(), (int64_t) biasValues.size()};
|
||||
|
||||
auto obj = std::make_unique<PluginType>(epsilon, scaleWeights, biasWeights, relu, alpha);
|
||||
obj->setPluginNamespace(mNamespace.c_str());
|
||||
obj->initialize();
|
||||
return obj.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreator::createPlugin(
|
||||
char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
return createPluginBase<InstanceNormalizationPlugin>(name, fc);
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreatorV2::createPlugin(
|
||||
char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
return createPluginBase<InstanceNormalizationPluginV2>(name, fc);
|
||||
}
|
||||
|
||||
template <class PluginType>
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreator::deserializePluginBase(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
auto obj = std::make_unique<PluginType>(serialData, serialLength);
|
||||
obj->setPluginNamespace(mNamespace.c_str());
|
||||
obj->initialize();
|
||||
return obj.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
return deserializePluginBase<InstanceNormalizationPlugin>(name, serialData, serialLength);
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* InstanceNormalizationPluginCreatorV2::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
return deserializePluginBase<InstanceNormalizationPluginV2>(name, serialData, serialLength);
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
#ifndef TRT_INSTANCE_NORMALIZATION_PLUGIN_LEGACY_H
|
||||
#define TRT_INSTANCE_NORMALIZATION_PLUGIN_LEGACY_H
|
||||
#include "common/plugin.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include "instanceNormalizationPlugin/instanceNormFwd.h"
|
||||
#include <cuda_fp16.h>
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
typedef uint16_t half_type;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr char const* gInstancePluginFullNameV1{"InstanceNormalization_TRT, version: 1"};
|
||||
constexpr char const* gInstancePluginFullNameV2{"InstanceNormalization_TRT, version: 2"};
|
||||
} // namespace
|
||||
|
||||
class InstanceNormalizationPlugin : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
|
||||
public:
|
||||
InstanceNormalizationPlugin(float epsilon, nvinfer1::Weights const& scale, nvinfer1::Weights const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F);
|
||||
InstanceNormalizationPlugin(float epsilon, std::vector<float> const& scale, std::vector<float> const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F);
|
||||
InstanceNormalizationPlugin(void const* serialData, size_t serialLength);
|
||||
|
||||
InstanceNormalizationPlugin() = delete;
|
||||
|
||||
~InstanceNormalizationPlugin() override;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
// DynamicExt plugins returns DimsExprs class instead of Dims
|
||||
using nvinfer1::IPluginV2::getOutputDimensions;
|
||||
DimsExprs getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs,
|
||||
nvinfer1::IExprBuilder& exprBuilder) noexcept override;
|
||||
|
||||
int32_t initialize() noexcept override;
|
||||
|
||||
void terminate() noexcept override;
|
||||
|
||||
using nvinfer1::IPluginV2::getWorkspaceSize;
|
||||
size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
using nvinfer1::IPluginV2::enqueue;
|
||||
int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
|
||||
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
|
||||
|
||||
size_t getSerializationSize() const noexcept override;
|
||||
|
||||
void serialize(void* buffer) const noexcept override;
|
||||
|
||||
// DynamicExt plugin supportsFormat update.
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
|
||||
|
||||
char const* getPluginType() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
void destroy() noexcept override;
|
||||
|
||||
nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept override;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
DataType getOutputDataType(
|
||||
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
|
||||
void attachToContext(
|
||||
cudnnContext* cudnn, cublasContext* cublas, nvinfer1::IGpuAllocator* allocator) noexcept override;
|
||||
|
||||
void detachFromContext() noexcept override;
|
||||
|
||||
using nvinfer1::IPluginV2Ext::configurePlugin;
|
||||
void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
||||
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
protected:
|
||||
template <class PluginType>
|
||||
nvinfer1::IPluginV2DynamicExt* cloneBase() const noexcept;
|
||||
|
||||
private:
|
||||
float mEpsilon{};
|
||||
float mAlpha{};
|
||||
int32_t mRelu{};
|
||||
int32_t mNchan{};
|
||||
std::vector<float> mHostScale;
|
||||
std::vector<float> mHostBias;
|
||||
float* mDeviceScale{nullptr};
|
||||
float* mDeviceBias{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnHandle_t mCudnnHandle{nullptr};
|
||||
nvinfer1::pluginInternal::CudnnWrapper& mCudnnWrapper
|
||||
= nvinfer1::pluginInternal::getCudnnWrapper(gInstancePluginFullNameV1);
|
||||
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mXDescriptor{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mYDescriptor{nullptr};
|
||||
nvinfer1::pluginInternal::cudnnTensorDescriptor_t mBDescriptor{nullptr};
|
||||
std::string mPluginNamespace;
|
||||
std::string mNamespace;
|
||||
bool mInitialized{false};
|
||||
int32_t mCudaDriverVersion{-1};
|
||||
|
||||
// NDHWC implementation
|
||||
instance_norm_impl::InstanceNormFwdContext mContext;
|
||||
};
|
||||
|
||||
class InstanceNormalizationPluginCreator : public nvinfer1::pluginInternal::BaseCreator
|
||||
{
|
||||
public:
|
||||
InstanceNormalizationPluginCreator();
|
||||
|
||||
~InstanceNormalizationPluginCreator() override = default;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV2DynamicExt* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
|
||||
|
||||
IPluginV2DynamicExt* deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
|
||||
protected:
|
||||
template <class PluginType>
|
||||
IPluginV2DynamicExt* createPluginBase(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept;
|
||||
|
||||
template <class PluginType>
|
||||
IPluginV2DynamicExt* deserializePluginBase(char const* name, void const* serialData, size_t serialLength) noexcept;
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
};
|
||||
|
||||
// For backward compatibility, create version "2" of the identical plugin.
|
||||
// Background: in TRT 8.0, we added 3D InstanceNorm plugin as the version 2 of the "InstanceNormalization_TRT" plugin.
|
||||
// However, in TRT 8.2, we have fused it into version 1, so a separate version 2 is no longer needed, but is only kept
|
||||
// for backward compatibility.
|
||||
class InstanceNormalizationPluginV2 final : public InstanceNormalizationPlugin
|
||||
{
|
||||
public:
|
||||
InstanceNormalizationPluginV2(float epsilon, nvinfer1::Weights const& scale, nvinfer1::Weights const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F)
|
||||
: InstanceNormalizationPlugin(epsilon, scale, bias, relu, alpha)
|
||||
{
|
||||
}
|
||||
InstanceNormalizationPluginV2(float epsilon, std::vector<float> const& scale, std::vector<float> const& bias,
|
||||
int32_t relu = 0, float alpha = 0.F)
|
||||
: InstanceNormalizationPlugin(epsilon, scale, bias, relu, alpha)
|
||||
{
|
||||
}
|
||||
InstanceNormalizationPluginV2(void const* serialData, size_t serialLength)
|
||||
: InstanceNormalizationPlugin(serialData, serialLength)
|
||||
{
|
||||
}
|
||||
InstanceNormalizationPluginV2() = delete;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
|
||||
|
||||
private:
|
||||
nvinfer1::pluginInternal::CudnnWrapper& mCudnnWrapper
|
||||
= nvinfer1::pluginInternal::getCudnnWrapper(gInstancePluginFullNameV2);
|
||||
};
|
||||
|
||||
class InstanceNormalizationPluginCreatorV2 final : public InstanceNormalizationPluginCreator
|
||||
{
|
||||
public:
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
IPluginV2DynamicExt* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
|
||||
IPluginV2DynamicExt* deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
};
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#endif // TRT_INSTANCE_NORMALIZATION_PLUGIN_LEGACY_H
|
||||
Reference in New Issue
Block a user