ANIKET SHIVAM
644
tools/util/include/cutlass/util/device_layernorm.h
Normal file
644
tools/util/include/cutlass/util/device_layernorm.h
Normal file
@@ -0,0 +1,644 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
/**
|
||||
* \file
|
||||
* \brief cuda kernels to do layernorm on a device memory tensor with RowMajor layout.
|
||||
*/
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "device_utils.h"
|
||||
#include <float.h>
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
/** \brief interface to do layernorm on a device memory tensor with RowMajor layout.
|
||||
* \tparam T: data type
|
||||
*/
|
||||
template <typename T>
|
||||
void layernorm(cutlass::MatrixCoord tensor_size,
|
||||
TensorRef<T, layout::RowMajor> ref_output,
|
||||
TensorRef<T, layout::RowMajor> ref_input,
|
||||
TensorRef<T, layout::RowMajor> ref_gamma,
|
||||
TensorRef<T, layout::RowMajor> ref_beta,
|
||||
cudaStream_t stream);
|
||||
|
||||
/**
|
||||
* output [m, n] row-major
|
||||
* input [m, n] row-major
|
||||
* gamma [n]
|
||||
* beta [n]
|
||||
* grid(m)
|
||||
* block(block_size) -- each block deals with n elements ; each thread deals with ITEM_PER_THREAD elements
|
||||
*/
|
||||
template<typename T, int ITEM_PER_THREAD>
|
||||
__global__ void layernorm_twoPassAlgo_stored_locally_e1(T* output,
|
||||
const T* input,
|
||||
const T* gamma,
|
||||
const T* beta,
|
||||
const int m,
|
||||
const int n)
|
||||
{
|
||||
const int m_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int bdimx = blockDim.x;
|
||||
__shared__ float s_mean, s_variance;
|
||||
T local_val[ITEM_PER_THREAD];
|
||||
float local_sums[1] = {0.0f};
|
||||
int offset = m_idx * n;
|
||||
input += offset;
|
||||
output += offset;
|
||||
|
||||
const T zero = T(0.0f);
|
||||
#pragma unroll
|
||||
for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
|
||||
int index = tid + i*bdimx;
|
||||
local_val[i] = index < n ? input[index] : zero;
|
||||
local_sums[0] += static_cast<float>(local_val[i]);
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_mean = local_sums[0] / n;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
local_sums[0] = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
|
||||
int index = tid + i*bdimx;
|
||||
if (index < n){
|
||||
const float tmp = static_cast<float>(local_val[i]) - s_mean;
|
||||
local_sums[0] += tmp * tmp;
|
||||
}
|
||||
}
|
||||
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_variance = rsqrtf(local_sums[0] / n + 1e-5);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
|
||||
int index = tid + i*bdimx;
|
||||
if (index < n) {
|
||||
const T gamma_val = gamma[index];
|
||||
const T beta_val = beta[index];
|
||||
output[index] = T((static_cast<float>(local_val[i]) - s_mean) * s_variance * static_cast<float>(gamma_val) + static_cast<float>(beta_val));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* output [m, n] row-major
|
||||
* input [m, n] row-major
|
||||
* gamma [n]
|
||||
* beta [n]
|
||||
* grid(m)
|
||||
* block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*2 elements;
|
||||
*/
|
||||
template<typename T2, typename T, int ITEM_PER_THREAD>
|
||||
__global__ void layernorm_twoPassAlgo_stored_locally_e2(T2* output,
|
||||
const T2* input,
|
||||
const T2* gamma,
|
||||
const T2* beta,
|
||||
const int m,
|
||||
const int n)
|
||||
{
|
||||
const int m_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int bdimx = blockDim.x;
|
||||
__shared__ float s_mean, s_variance;
|
||||
float local_sums[1] = {0.0f};
|
||||
T2 local_val[ITEM_PER_THREAD];
|
||||
const int n_2 = n / 2;
|
||||
int offset = m_idx * n_2;
|
||||
input += offset;
|
||||
output += offset;
|
||||
|
||||
const T2 zero = {T(0.0f), T(0.0f)};
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
local_val[i] = index < n_2 ? input[index] : zero;
|
||||
local_sums[0] += static_cast<float>(local_val[i].x) + static_cast<float>(local_val[i].y);
|
||||
}
|
||||
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_mean = local_sums[0] / n;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
local_sums[0] = 0.0f;
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
if (index < n_2){
|
||||
const float2 tmp = {static_cast<float>(local_val[i].x) - s_mean,
|
||||
static_cast<float>(local_val[i].y) - s_mean};
|
||||
local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y;
|
||||
}
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_variance = rsqrtf(local_sums[0] / n + 1e-5);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
if (index < n_2){
|
||||
const T2 gamma_val = gamma[index];
|
||||
const T2 beta_val = beta[index];
|
||||
T2 tmp;
|
||||
tmp.x = T((static_cast<float>(local_val[i].x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
|
||||
tmp.y = T((static_cast<float>(local_val[i].y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
|
||||
output[index] = tmp;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* output [m, n] row-major
|
||||
* input [m, n] row-major
|
||||
* gamma [n]
|
||||
* beta [n]
|
||||
* grid(m)
|
||||
* block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*4 elements;
|
||||
*/
|
||||
template<typename T4, typename T, int ITEM_PER_THREAD>
|
||||
__global__ void layernorm_twoPassAlgo_stored_locally_e4(T4* output,
|
||||
const T4* input,
|
||||
const T4* gamma,
|
||||
const T4* beta,
|
||||
const int m,
|
||||
const int n)
|
||||
{
|
||||
const int m_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int bdimx = blockDim.x;
|
||||
__shared__ float s_mean, s_variance;
|
||||
float local_sums[1] = {0.0f};
|
||||
T4 local_val[ITEM_PER_THREAD];
|
||||
const int n_4 = n / 4;
|
||||
int offset = m_idx * n_4;
|
||||
input += offset;
|
||||
output += offset;
|
||||
|
||||
const T4 zero = {T(0.0f), T(0.0f), T(0.0f), T(0.0f)};
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
local_val[i] = index < n_4 ? input[index] : zero;
|
||||
local_sums[0] += static_cast<float>(local_val[i].x) + static_cast<float>(local_val[i].y) +
|
||||
static_cast<float>(local_val[i].z) + static_cast<float>(local_val[i].w);
|
||||
}
|
||||
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_mean = local_sums[0] / n;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
local_sums[0] = 0.0f;
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
if (index < n_4){
|
||||
const float4 tmp = {static_cast<float>(local_val[i].x) - s_mean,
|
||||
static_cast<float>(local_val[i].y) - s_mean,
|
||||
static_cast<float>(local_val[i].z) - s_mean,
|
||||
static_cast<float>(local_val[i].w) - s_mean};
|
||||
local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y + tmp.z * tmp.z + tmp.w * tmp.w;
|
||||
}
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_variance = rsqrtf(local_sums[0] / n + 1e-5);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#pragma UNROLL
|
||||
for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
|
||||
const int index = i*bdimx + tid;
|
||||
if (index < n_4){
|
||||
const T4 gamma_val = gamma[index];
|
||||
const T4 beta_val = beta[index];
|
||||
T4 tmp;
|
||||
tmp.x = T((static_cast<float>(local_val[i].x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
|
||||
tmp.y = T((static_cast<float>(local_val[i].y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
|
||||
tmp.z = T((static_cast<float>(local_val[i].z) - s_mean)*s_variance*static_cast<float>(gamma_val.z) + static_cast<float>(beta_val.z));
|
||||
tmp.w = T((static_cast<float>(local_val[i].w) - s_mean)*s_variance*static_cast<float>(gamma_val.w) + static_cast<float>(beta_val.w));
|
||||
output[index] = tmp;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* output [m, n] row-major
|
||||
* input [m, n] row-major
|
||||
* gamma [n]
|
||||
* beta [n]
|
||||
* grid(m)
|
||||
* block(block_size) -- each block deals with n elements ; each thread deals with ITEM_PER_THREAD elements
|
||||
*/
|
||||
template<typename T>
|
||||
__global__ void layernorm_twoPassAlgo_e1(T* output,
|
||||
const T* input,
|
||||
const T* gamma,
|
||||
const T* beta,
|
||||
const int m,
|
||||
const int n)
|
||||
{
|
||||
const int m_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int bdimx = blockDim.x;
|
||||
__shared__ float s_mean, s_variance;
|
||||
float local_sums[1] = {0.0f};
|
||||
int offset = m_idx * n;
|
||||
input += offset;
|
||||
output += offset;
|
||||
|
||||
for (int index = tid ; index < n ; index += bdimx){
|
||||
float local_val = static_cast<float>(input[index]);
|
||||
local_sums[0] += local_val;
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_mean = local_sums[0] / n;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
local_sums[0] = 0.0f;
|
||||
for (int index = tid ; index < n ; index += bdimx){
|
||||
float local_val = static_cast<float>(input[index]);
|
||||
local_val = local_val - s_mean;
|
||||
local_sums[0] += local_val * local_val;
|
||||
}
|
||||
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_variance = rsqrtf(local_sums[0] / n + 1e-5);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int index = tid ; index < n ; index += bdimx){
|
||||
const T gamma_val = gamma[index];
|
||||
const T beta_val = beta[index];
|
||||
const T local_val = input[index];
|
||||
output[index] = T((static_cast<float>(local_val) - s_mean) * s_variance * static_cast<float>(gamma_val) + static_cast<float>(beta_val));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* output [m, n] row-major
|
||||
* input [m, n] row-major
|
||||
* gamma [n]
|
||||
* beta [n]
|
||||
* grid(m)
|
||||
* block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*2 elements;
|
||||
*/
|
||||
template<typename T2, typename T>
|
||||
__global__ void layernorm_twoPassAlgo_e2(T2* output,
|
||||
const T2* input,
|
||||
const T2* gamma,
|
||||
const T2* beta,
|
||||
const int m,
|
||||
const int n)
|
||||
{
|
||||
const int m_idx = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
const int bdimx = blockDim.x;
|
||||
__shared__ float s_mean, s_variance;
|
||||
float local_sums[1] = {0.0f};
|
||||
const int n_2 = n / 2;
|
||||
int offset = m_idx * n_2;
|
||||
input += offset;
|
||||
output += offset;
|
||||
|
||||
for (int index = tid; index < n_2; index += bdimx) {
|
||||
const T2 local_val = input[index];
|
||||
local_sums[0] += static_cast<float>(local_val.x) + static_cast<float>(local_val.y);
|
||||
}
|
||||
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_mean = local_sums[0] / n;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
local_sums[0] = 0.0f;
|
||||
for (int index = tid; index < n_2; index += bdimx) {
|
||||
const T2 local_val = input[index];
|
||||
const float2 tmp = {static_cast<float>(local_val.x) - s_mean,
|
||||
static_cast<float>(local_val.y) - s_mean};
|
||||
local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y;
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
warpReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
else {
|
||||
blockReduceSum<float, 1>(local_sums);
|
||||
}
|
||||
if (threadIdx.x == 0) {
|
||||
s_variance = rsqrtf(local_sums[0] / n + 1e-5);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int index = tid; index < n_2; index += bdimx) {
|
||||
const T2 local_val = input[index];
|
||||
const T2 gamma_val = gamma[index];
|
||||
const T2 beta_val = beta[index];
|
||||
T2 tmp;
|
||||
tmp.x = T((static_cast<float>(local_val.x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
|
||||
tmp.y = T((static_cast<float>(local_val.y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
|
||||
output[index] = tmp;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void layernorm(cutlass::MatrixCoord tensor_size,
|
||||
TensorRef<T, layout::RowMajor> ref_output,
|
||||
TensorRef<T, layout::RowMajor> ref_input,
|
||||
TensorRef<T, layout::RowMajor> ref_gamma,
|
||||
TensorRef<T, layout::RowMajor> ref_beta,
|
||||
cudaStream_t stream){
|
||||
const int m = tensor_size.row();
|
||||
const int n = tensor_size.column();
|
||||
T* output = ref_output.data();
|
||||
const T* input = ref_input.data();
|
||||
const T* gamma = ref_gamma.data();
|
||||
const T* beta = ref_beta.data();
|
||||
dim3 grid(m);
|
||||
dim3 block((n + 31)/32*32);
|
||||
if (block.x > 1024){
|
||||
block.x = 1024;
|
||||
}
|
||||
// TODO : There should be better configs for different cases, we only use several samples to show how to use here
|
||||
// TODO : using registers to store values locally can reduce the ldgs from global memory and speedup the kernels.
|
||||
if ((n % 4 == 0) && (n >= 128) && (n <= 4096)) {
|
||||
block.x = (n/4 + 31)/32*32;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_stored_locally_e4<float4, float, 1><<<grid, block, 0, stream>>>(
|
||||
(float4*)output,
|
||||
(const float4*)input,
|
||||
(const float4*)gamma,
|
||||
(const float4*)beta,
|
||||
m,
|
||||
n);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_stored_locally_e4<half4, half, 1><<<grid, block, 0, stream>>>(
|
||||
(half4*)output,
|
||||
(const half4*)input,
|
||||
(const half4*)gamma,
|
||||
(const half4*)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
} //if ((n % 4 == 0) && (n >= 128) && (n <= 4096))
|
||||
else if (n % 2 == 0) {
|
||||
if (n / 2 <= 1024) {
|
||||
block.x = (n/2 + 31)/32*32;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<float2, float, 1><<<grid, block, 0, stream>>>(
|
||||
(float2*)output,
|
||||
(const float2*)input,
|
||||
(const float2*)gamma,
|
||||
(const float2*)beta,
|
||||
m,
|
||||
n);
|
||||
} //if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<half2, half, 1><<<grid, block, 0, stream>>>(
|
||||
(half2*)output,
|
||||
(const half2*)input,
|
||||
(const half2*)gamma,
|
||||
(const half2*)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
} // if (n / 2 <= 1024)
|
||||
else if (n <= 8192) {
|
||||
block.x = ((n + 7)/8 + 31)/32*32;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<float2, float, 4><<<grid, block, 0, stream>>>(
|
||||
(float2*)output,
|
||||
(const float2*)input,
|
||||
(const float2*)gamma,
|
||||
(const float2*)beta,
|
||||
m,
|
||||
n);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<half2, half, 4><<<grid, block, 0, stream>>>(
|
||||
(half2*)output,
|
||||
(const half2*)input,
|
||||
(const half2*)gamma,
|
||||
(const half2*)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
} // if (n <= 8192)
|
||||
else if (n <= 16384) {
|
||||
block.x = ((n + 15)/ 16 + 31)/32*32;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<float2, float, 8><<<grid, block, 0, stream>>>(
|
||||
(float2*)output,
|
||||
(const float2*)input,
|
||||
(const float2*)gamma,
|
||||
(const float2*)beta,
|
||||
m,
|
||||
n);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<half2, half, 8><<<grid, block, 0, stream>>>(
|
||||
(half2*)output,
|
||||
(const half2*)input,
|
||||
(const half2*)gamma,
|
||||
(const half2*)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
} // if (n <= 16384)
|
||||
else if (n <= 32768) {
|
||||
block.x = ((n + 31)/32 + 31)/32*32;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<float2, float, 16><<<grid, block, 0, stream>>>(
|
||||
(float2*)output,
|
||||
(const float2*)input,
|
||||
(const float2*)gamma,
|
||||
(const float2*)beta,
|
||||
m,
|
||||
n);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_stored_locally_e2<half2, half, 16><<<grid, block, 0, stream>>>(
|
||||
(half2*)output,
|
||||
(const half2*)input,
|
||||
(const half2*)gamma,
|
||||
(const half2*)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
} // if (n <= 32768)
|
||||
else {
|
||||
if (block.x > 512)
|
||||
block.x = 512;
|
||||
if (std::is_same<T, float>::value) {
|
||||
layernorm_twoPassAlgo_e2<float2, float><<<grid, block, 0, stream>>>(
|
||||
(float2 *)output,
|
||||
(const float2 *)input,
|
||||
(const float2 *)gamma,
|
||||
(const float2 *)beta,
|
||||
m,
|
||||
n);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
layernorm_twoPassAlgo_e2<half2, half><<<grid, block, 0, stream>>>(
|
||||
(half2 *)output,
|
||||
(const half2 *)input,
|
||||
(const half2 *)gamma,
|
||||
(const half2 *)beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
}
|
||||
} // if (n % 2 == 0)
|
||||
else {
|
||||
if (n <= 1024) {
|
||||
layernorm_twoPassAlgo_stored_locally_e1<T, 1><<<grid, block, 0, stream>>>(
|
||||
output,
|
||||
input,
|
||||
gamma,
|
||||
beta,
|
||||
m,
|
||||
n);
|
||||
} // if (n <= 1024)
|
||||
else if (n <= 8192) {
|
||||
block.x = ((n + 7)/8 + 31)/32*32;
|
||||
layernorm_twoPassAlgo_stored_locally_e1<T, 8><<<grid, block, 0, stream>>>(
|
||||
output,
|
||||
input,
|
||||
gamma,
|
||||
beta,
|
||||
m,
|
||||
n);
|
||||
} // if (n <= 8192)
|
||||
else if (n <= 16384) {
|
||||
block.x = ((n + 15)/16 + 32)/32*32;
|
||||
layernorm_twoPassAlgo_stored_locally_e1<T, 16><<<grid, block, 0, stream>>>(
|
||||
output,
|
||||
input,
|
||||
gamma,
|
||||
beta,
|
||||
m,
|
||||
n);
|
||||
} // if (n <= 16384)
|
||||
else if (n <= 32768) {
|
||||
block.x = ((n + 31)/32 + 31)/32*32;
|
||||
layernorm_twoPassAlgo_stored_locally_e1<T, 32><<<grid, block, 0, stream>>>(
|
||||
output,
|
||||
input,
|
||||
gamma,
|
||||
beta,
|
||||
m,
|
||||
n);
|
||||
} // if (n <= 32768)
|
||||
else{
|
||||
if (block.x > 512) {
|
||||
block.x = 512;
|
||||
}
|
||||
layernorm_twoPassAlgo_e1<<<grid, block, 0, stream>>>(
|
||||
output,
|
||||
input,
|
||||
gamma,
|
||||
beta,
|
||||
m,
|
||||
n);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} //namespace cutlass
|
||||
576
tools/util/include/cutlass/util/device_nhwc_pooling.h
Normal file
576
tools/util/include/cutlass/util/device_nhwc_pooling.h
Normal file
@@ -0,0 +1,576 @@
|
||||
/******************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
******************************************************************************/
|
||||
|
||||
#pragma once
|
||||
|
||||
/**
|
||||
* \file
|
||||
* \brief cuda kernels to do avg/max pooling on a device memory tensor with NHWC layout.
|
||||
*/
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/tensor.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/tensor_coord.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "device_utils.h"
|
||||
#include <float.h>
|
||||
|
||||
namespace cutlass {
|
||||
|
||||
/** \brief interface to do avg/max pooling on a device memory tensor with NHWC layout.
|
||||
* \tparam T: data type
|
||||
*/
|
||||
template <typename T>
|
||||
void pooling_nhwc(cutlass::Tensor4DCoord input_tensor_size,
|
||||
cutlass::Tensor4DCoord filter_tensor_size,
|
||||
cutlass::Tensor4DCoord output_tensor_size,
|
||||
cutlass::MatrixCoord padding,
|
||||
cutlass::MatrixCoord stride,
|
||||
TensorRef<T, layout::TensorNHWC> ref_input,
|
||||
TensorRef<T, layout::TensorNHWC> ref_output,
|
||||
int poolingType, //0 for avg pooling ; 1 for max pooling
|
||||
cudaStream_t stream);
|
||||
|
||||
/** get the output size of pooling
|
||||
*/
|
||||
inline int getOutputSize(int H_W, int padding, int kernel_size, int stride)
|
||||
{
|
||||
return (H_W + 2 * padding - kernel_size) / stride + 1;
|
||||
}
|
||||
|
||||
/**
|
||||
* input is [N, H, W, C]
|
||||
* assume stride == kernel_size
|
||||
* output_h = (H + 2*padding_H - kernel_H)/stride_H
|
||||
* output_w = (W + 2*padding_W - kernel_W)/stride_W
|
||||
* output is [N, output_h, output_w, C]
|
||||
* grid(N, output_h, output_w)
|
||||
* block(min(C, 256)) :
|
||||
* each block deals with C elements of output when each thread deals with ((C + 255)/256 element of output)
|
||||
*/
|
||||
template<typename T, bool IS_AVG_POOLING>
|
||||
__global__ void pooling_nhwc_element1_kernel(T* output,
|
||||
const T* input,
|
||||
const int N,
|
||||
const int H,
|
||||
const int W,
|
||||
const int C,
|
||||
const int output_H,
|
||||
const int output_W,
|
||||
const int kernel_H,
|
||||
const int kernel_W,
|
||||
const int stride_H,
|
||||
const int stride_W,
|
||||
const int padding_H,
|
||||
const int padding_W)
|
||||
{
|
||||
const int tid = threadIdx.x;
|
||||
const int n_idx = blockIdx.x;
|
||||
const int output_h_idx = blockIdx.y;
|
||||
const int output_w_idx = blockIdx.z;
|
||||
|
||||
int h_start_idx = output_h_idx * stride_H - padding_H;
|
||||
int h_end_idx = h_start_idx + kernel_H;
|
||||
h_start_idx = (h_start_idx < 0) ? 0 : h_start_idx;
|
||||
h_end_idx = h_end_idx > H ? H : h_end_idx;
|
||||
|
||||
int w_start_idx = output_w_idx * stride_W - padding_W;
|
||||
int w_end_idx = w_start_idx + kernel_W;
|
||||
w_start_idx = (w_start_idx < 0) ? 0 : w_start_idx;
|
||||
w_end_idx = w_end_idx > W ? W : w_end_idx;
|
||||
|
||||
input += n_idx * H * W * C;
|
||||
output += ((n_idx * output_H + output_h_idx) * output_W + output_w_idx) * C;
|
||||
const int kernel_size2 = kernel_H * kernel_W;
|
||||
for (int c_idx = tid; c_idx < C; c_idx += blockDim.x) {
|
||||
float pooling;
|
||||
if (IS_AVG_POOLING){
|
||||
pooling = 0.0f;
|
||||
}
|
||||
else{
|
||||
pooling = -FLT_MAX;
|
||||
}
|
||||
for (int h = h_start_idx; h < h_end_idx; h++) {
|
||||
for (int w = w_start_idx; w < w_end_idx; w++) {
|
||||
const int idx = (h * W + w) * C;
|
||||
const float tmp = static_cast<float>(input[idx + c_idx]);
|
||||
if (IS_AVG_POOLING){
|
||||
pooling = pooling + tmp;
|
||||
}
|
||||
else{
|
||||
pooling = pooling > tmp ? pooling : tmp;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
T output_val;
|
||||
if (IS_AVG_POOLING){
|
||||
output_val = T(pooling/kernel_size2);
|
||||
}
|
||||
else{
|
||||
output_val = T(pooling);
|
||||
}
|
||||
output[c_idx] = output_val;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T2, typename T, bool IS_AVG_POOLING>
|
||||
__global__ void pooling_nhwc_element2_kernel(T2* output,
|
||||
const T2* input,
|
||||
const int N,
|
||||
const int H,
|
||||
const int W,
|
||||
const int C,
|
||||
const int output_H,
|
||||
const int output_W,
|
||||
const int kernel_H,
|
||||
const int kernel_W,
|
||||
const int stride_H,
|
||||
const int stride_W,
|
||||
const int padding_H,
|
||||
const int padding_W)
|
||||
{
|
||||
const int tid = threadIdx.x;
|
||||
const int n_idx = blockIdx.x;
|
||||
const int output_h_idx = blockIdx.y;
|
||||
const int output_w_idx = blockIdx.z;
|
||||
|
||||
int h_start_idx = output_h_idx * stride_H - padding_H;
|
||||
int h_end_idx = h_start_idx + kernel_H;
|
||||
h_start_idx = (h_start_idx < 0) ? 0 : h_start_idx;
|
||||
h_end_idx = h_end_idx > H ? H : h_end_idx;
|
||||
|
||||
int w_start_idx = output_w_idx * stride_W - padding_W;
|
||||
int w_end_idx = w_start_idx + kernel_W;
|
||||
w_start_idx = (w_start_idx < 0) ? 0 : w_start_idx;
|
||||
w_end_idx = w_end_idx > W ? W : w_end_idx;
|
||||
|
||||
input += n_idx * H * W * C;
|
||||
output += ((n_idx * output_H + output_h_idx) * output_W + output_w_idx) * C;
|
||||
const int kernel_size2 = kernel_H * kernel_W;
|
||||
for (int c_idx = tid; c_idx < C; c_idx += blockDim.x) {
|
||||
float2 pooling;
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling = {0.0f, 0.0f};
|
||||
}
|
||||
else {
|
||||
pooling = {-FLT_MAX, -FLT_MAX};
|
||||
}
|
||||
for (int h = h_start_idx; h < h_end_idx; h++) {
|
||||
for (int w = w_start_idx; w < w_end_idx; w++) {
|
||||
const int idx = (h * W + w) * C;
|
||||
const T2 tmp = input[idx + c_idx];
|
||||
const float2 tmp_flt2 = {static_cast<float>(tmp.x), static_cast<float>(tmp.y)};
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling.x += tmp_flt2.x;
|
||||
pooling.y += tmp_flt2.y;
|
||||
}
|
||||
else {
|
||||
pooling.x = pooling.x > tmp_flt2.x ? pooling.x : tmp_flt2.x;
|
||||
pooling.y = pooling.y > tmp_flt2.y ? pooling.y : tmp_flt2.y;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
T2 output_val;
|
||||
if (IS_AVG_POOLING) {
|
||||
output_val.x = T(pooling.x/kernel_size2);
|
||||
output_val.y = T(pooling.y/kernel_size2);
|
||||
}
|
||||
else {
|
||||
output_val.x = T(pooling.x);
|
||||
output_val.y = T(pooling.y);
|
||||
}
|
||||
output[c_idx] = output_val;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* output [N, 1, 1, C]
|
||||
* input [N, H, W, C]
|
||||
* grid(C, N)
|
||||
* block(block_size) -- each block deals with H*W/block_size elements;
|
||||
*/
|
||||
template<typename T, bool IS_AVG_POOLING>
|
||||
__global__ void pooling_nxhTo1x1_element1_kernel(
|
||||
T* output, const T* input, const int N, const int HW, const int C)
|
||||
{
|
||||
const int c_idx = blockIdx.x;
|
||||
const int n_idx = blockIdx.y;
|
||||
float pooling[1];
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling[0] = 0.0f;
|
||||
}
|
||||
else {
|
||||
pooling[0] = -FLT_MAX;
|
||||
}
|
||||
const size_t input_offset = n_idx * HW * C + c_idx;
|
||||
input += input_offset;
|
||||
const size_t output_offset = n_idx * C + c_idx;
|
||||
output += output_offset;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
for (int index = tid; index < HW; index += blockDim.x) {
|
||||
float val = static_cast<float>(input[index * C]);
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling[0] += val;
|
||||
}
|
||||
else {
|
||||
pooling[0] = pooling[0] > val ? pooling[0] : val;
|
||||
}
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
if (IS_AVG_POOLING) {
|
||||
warpReduceSum<float, 1>(pooling);
|
||||
}
|
||||
else {
|
||||
warpReduceMax<float, 1>(pooling);
|
||||
}
|
||||
}
|
||||
else {
|
||||
if (IS_AVG_POOLING) {
|
||||
blockReduceSum<float, 1>(pooling);
|
||||
}
|
||||
else {
|
||||
blockReduceMax<float, 1>(pooling);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
T output_val;
|
||||
if (IS_AVG_POOLING) {
|
||||
output_val = T(pooling[0] / HW);
|
||||
}
|
||||
else {
|
||||
output_val = T(pooling[0]);
|
||||
}
|
||||
output[0] = output_val;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* output [N, 1, 1, C]
|
||||
* input [N, H, W, C]
|
||||
* grid(C/2, N)
|
||||
* block(block_size) -- each thread deals with H*W/block_size * 2 elements;
|
||||
*/
|
||||
template<typename T2, typename T, bool IS_AVG_POOLING>
|
||||
__global__ void pooling_nxhTo1x1_element2_kernel(
|
||||
T2* output, const T2* input, const int N, const int HW, const int C)
|
||||
{
|
||||
const int c_idx = blockIdx.x;
|
||||
const int n_idx = blockIdx.y;
|
||||
float pooling[2];
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling[0] = pooling[1] = 0.0f;
|
||||
}
|
||||
else {
|
||||
pooling[0] = pooling[1] = -FLT_MAX;
|
||||
}
|
||||
const int C_2 = C / 2;
|
||||
const size_t input_offset = n_idx * HW * C_2 + c_idx;
|
||||
input += input_offset;
|
||||
const size_t output_offset = n_idx * C_2 + c_idx;
|
||||
output += output_offset;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
for (int index = tid; index < HW; index += blockDim.x) {
|
||||
T2 val = input[index * C_2];
|
||||
float2 val_flt2 = {static_cast<float>(val.x), static_cast<float>(val.y)};
|
||||
if (IS_AVG_POOLING) {
|
||||
pooling[0] += val_flt2.x;
|
||||
pooling[1] += val_flt2.y;
|
||||
}
|
||||
else {
|
||||
pooling[0] = pooling[0] > val_flt2.x ? pooling[0] : val_flt2.x;
|
||||
pooling[1] = pooling[1] > val_flt2.y ? pooling[1] : val_flt2.y;
|
||||
}
|
||||
}
|
||||
if (blockDim.x <= 32) {
|
||||
if (IS_AVG_POOLING) {
|
||||
warpReduceSum<float, 2>(pooling);
|
||||
}
|
||||
else {
|
||||
warpReduceMax<float, 2>(pooling);
|
||||
}
|
||||
}
|
||||
else {
|
||||
if (IS_AVG_POOLING) {
|
||||
blockReduceSum<float, 2>(pooling);
|
||||
}
|
||||
else {
|
||||
blockReduceMax<float, 2>(pooling);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (threadIdx.x == 0) {
|
||||
T2 output_val;
|
||||
if (IS_AVG_POOLING) {
|
||||
output_val.x = T(pooling[0] / HW);
|
||||
output_val.y = T(pooling[1] / HW);
|
||||
}
|
||||
else {
|
||||
output_val.x = T(pooling[0]);
|
||||
output_val.y = T(pooling[1]);
|
||||
}
|
||||
output[0] = output_val;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void pooling_nhwc(cutlass::Tensor4DCoord input_tensor_size,
|
||||
cutlass::Tensor4DCoord filter_tensor_size,
|
||||
cutlass::Tensor4DCoord output_tensor_size,
|
||||
cutlass::Tensor4DCoord padding,
|
||||
cutlass::MatrixCoord stride,
|
||||
TensorRef<T, layout::TensorNHWC> ref_input,
|
||||
TensorRef<T, layout::TensorNHWC> ref_output,
|
||||
int poolingType, //0 for avg pooling ; 1 for max pooling
|
||||
cudaStream_t stream) {
|
||||
|
||||
assert(input_tensor_size.n() == output_tensor_size.n() &&
|
||||
input_tensor_size.c() == output_tensor_size.c());
|
||||
|
||||
assert(filter_tensor_size.h() == stride.row() &&
|
||||
filter_tensor_size.w() == stride.column());
|
||||
|
||||
const int N = input_tensor_size.n();
|
||||
const int H = input_tensor_size.h();
|
||||
const int W = input_tensor_size.w();
|
||||
const int C = input_tensor_size.c();
|
||||
const int padding_H = padding.h();
|
||||
const int padding_W = padding.w();
|
||||
const int kernel_H = filter_tensor_size.h();
|
||||
const int kernel_W = filter_tensor_size.w();
|
||||
const int stride_H = stride.row();
|
||||
const int stride_W = stride.column();
|
||||
|
||||
const int output_H = getOutputSize(H, padding_H, kernel_H, stride_H);
|
||||
const int output_W = getOutputSize(W, padding_W, kernel_W, stride_W);
|
||||
|
||||
assert(output_tensor_size.h() == output_H &&
|
||||
output_tensor_size.w() == output_W);
|
||||
|
||||
if (C % 2 != 0) {
|
||||
if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0)) {
|
||||
dim3 grid(C, N);
|
||||
dim3 block(256);
|
||||
if (H*W < block.x){
|
||||
block.x = (H*W + 31)/32*32;
|
||||
}
|
||||
if (poolingType == 0) {
|
||||
pooling_nxhTo1x1_element1_kernel<T, true><<<grid, block, 0, stream>>>(
|
||||
ref_output.data(),
|
||||
ref_input.data(),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
} // if (poolingType == 0)
|
||||
else {
|
||||
pooling_nxhTo1x1_element1_kernel<T, false><<<grid, block, 0, stream>>>(
|
||||
ref_output.data(),
|
||||
ref_input.data(),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
}
|
||||
} // if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0))
|
||||
else {
|
||||
dim3 grid(N, output_H, output_W);
|
||||
dim3 block(256);
|
||||
if (C < block.x) {
|
||||
block.x = C;
|
||||
}
|
||||
if (poolingType == 0) {
|
||||
pooling_nhwc_element1_kernel<T, true><<<grid, block, 0, stream>>>(
|
||||
ref_output.data(),
|
||||
ref_input.data(),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
} // if (poolingType == 0)
|
||||
else {
|
||||
pooling_nhwc_element1_kernel<T, false><<<grid, block, 0, stream>>>(
|
||||
ref_output.data(),
|
||||
ref_input.data(),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
}
|
||||
}
|
||||
} // if (C % 2 != 0))
|
||||
else {
|
||||
if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0)) {
|
||||
dim3 grid(C/2, N);
|
||||
dim3 block(256);
|
||||
if (H*W < block.x){
|
||||
block.x = (H*W + 31)/32*32;
|
||||
}
|
||||
if (poolingType == 0) {
|
||||
if (std::is_same<T, float>::value) {
|
||||
pooling_nxhTo1x1_element2_kernel<float2, float, true><<<grid, block, 0, stream>>>(
|
||||
(float2*)(ref_output.data()),
|
||||
(const float2*)(ref_input.data()),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
pooling_nxhTo1x1_element2_kernel<half2, half, true><<<grid, block, 0, stream>>>(
|
||||
(half2*)(ref_output.data()),
|
||||
(const half2*)(ref_input.data()),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
}
|
||||
} // if (poolingType == 0)
|
||||
else {
|
||||
if (std::is_same<T, float>::value) {
|
||||
pooling_nxhTo1x1_element2_kernel<float2, float, false><<<grid, block, 0, stream>>>(
|
||||
(float2*)(ref_output.data()),
|
||||
(const float2*)(ref_input.data()),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
pooling_nxhTo1x1_element2_kernel<half2, half, false><<<grid, block, 0, stream>>>(
|
||||
(half2*)(ref_output.data()),
|
||||
(const half2*)(ref_input.data()),
|
||||
N,
|
||||
H*W,
|
||||
C);
|
||||
}
|
||||
}
|
||||
} // if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0))
|
||||
else {
|
||||
dim3 grid(N, output_H, output_W);
|
||||
dim3 block(256);
|
||||
if (C/2 < block.x) {
|
||||
block.x = C/2;
|
||||
}
|
||||
if (poolingType == 0) {
|
||||
if (std::is_same<T, float>::value) {
|
||||
pooling_nhwc_element2_kernel<float2, float, true><<<grid, block, 0, stream>>>(
|
||||
(float2*)(ref_output.data()),
|
||||
(const float2*)(ref_input.data()),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C/2,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
pooling_nhwc_element2_kernel<half2, half, true><<<grid, block, 0, stream>>>(
|
||||
(half2*)(ref_output.data()),
|
||||
(const half2*)(ref_input.data()),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C/2,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
}
|
||||
} // if (poolingType == 0)
|
||||
else {
|
||||
if (std::is_same<T, float>::value) {
|
||||
pooling_nhwc_element2_kernel<float2, float, false><<<grid, block, 0, stream>>>(
|
||||
(float2*)(ref_output.data()),
|
||||
(const float2*)(ref_input.data()),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C/2,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
} // if (std::is_same<T, float>::value)
|
||||
else {
|
||||
pooling_nhwc_element2_kernel<half2, half, false><<<grid, block, 0, stream>>>(
|
||||
(half2*)(ref_output.data()),
|
||||
(const half2*)(ref_input.data()),
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
C/2,
|
||||
output_H,
|
||||
output_W,
|
||||
kernel_H,
|
||||
kernel_W,
|
||||
stride_H,
|
||||
stride_W,
|
||||
padding_H,
|
||||
padding_W);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} //namespace cutlass
|
||||
127
tools/util/include/cutlass/util/device_utils.h
Normal file
127
tools/util/include/cutlass/util/device_utils.h
Normal file
@@ -0,0 +1,127 @@
|
||||
/***************************************************************************************************
|
||||
* Copyright (c) 2017 - 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: BSD-3-Clause
|
||||
*
|
||||
* Redistribution and use in source and binary forms, with or without
|
||||
* modification, are permitted provided that the following conditions are met:
|
||||
*
|
||||
* 1. Redistributions of source code must retain the above copyright notice, this
|
||||
* list of conditions and the following disclaimer.
|
||||
*
|
||||
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
||||
* this list of conditions and the following disclaimer in the documentation
|
||||
* and/or other materials provided with the distribution.
|
||||
*
|
||||
* 3. Neither the name of the copyright holder nor the names of its
|
||||
* contributors may be used to endorse or promote products derived from
|
||||
* this software without specific prior written permission.
|
||||
*
|
||||
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
|
||||
/*! \file
|
||||
\brief utils code for device cutlass code
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <float.h>
|
||||
#define FINAL_MASK 0xffffffff
|
||||
|
||||
struct half4 {
|
||||
half x, y, z, w;
|
||||
};
|
||||
|
||||
template<typename T, int NUM>
|
||||
__inline__ __device__ T warpReduceSum(T* val)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
|
||||
}
|
||||
return (T)(0.0f);
|
||||
}
|
||||
|
||||
template<typename T, int NUM>
|
||||
__inline__ __device__ T blockReduceSum(T* val)
|
||||
{
|
||||
__shared__ T shared[NUM][33];
|
||||
int lane = threadIdx.x & 0x1f;
|
||||
int wid = threadIdx.x >> 5;
|
||||
|
||||
warpReduceSum<T, NUM>(val);
|
||||
|
||||
if (lane == 0) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
shared[i][wid] = val[i];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
bool is_mask = threadIdx.x < (blockDim.x / 32.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
|
||||
}
|
||||
warpReduceSum<T, NUM>(val);
|
||||
return (T)0.0f;
|
||||
}
|
||||
|
||||
template<typename T, int NUM>
|
||||
__inline__ __device__ T warpReduceMax(T* val)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val[i] = max(val[i], __shfl_xor_sync(FINAL_MASK, val[i], mask, 32));
|
||||
}
|
||||
return (T)(0.0f);
|
||||
}
|
||||
|
||||
template<typename T, int NUM>
|
||||
__inline__ __device__ T blockReduceMax(T* val)
|
||||
{
|
||||
static __shared__ T shared[32][NUM];
|
||||
int lane = threadIdx.x & 0x1f; // in-warp idx
|
||||
int wid = threadIdx.x >> 5; // warp idx
|
||||
|
||||
warpReduceMax<T, NUM>(val); // get maxx in each warp
|
||||
|
||||
if (lane == 0) // record in-warp maxx by warp Idx
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
shared[wid][i] = val[i];
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
|
||||
// blockDim.x is not divided by 32
|
||||
bool is_mask = threadIdx.x < (blockDim.x / 32.f);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < NUM; i++) {
|
||||
val[i] = is_mask ? shared[lane][i] : (T)(-FLT_MAX);
|
||||
}
|
||||
warpReduceMax<T, NUM>(val);
|
||||
|
||||
return (T)0.0f;
|
||||
}
|
||||
|
||||
@@ -123,11 +123,17 @@ __global__ void Conv2dFprop(
|
||||
}
|
||||
}
|
||||
|
||||
int c_per_group = problem_size.C / problem_size.groups;
|
||||
int k_per_group = problem_size.K / problem_size.groups;
|
||||
|
||||
// Compute convolution
|
||||
for (int R = 0; R < problem_size.R; ++R) {
|
||||
for (int S = 0; S < problem_size.S; ++S) {
|
||||
for (int C = 0; C < problem_size.C; ++C) {
|
||||
|
||||
// Get group id of currnet channel
|
||||
int c_group_idx = C / c_per_group;
|
||||
|
||||
// Load from activations tensor
|
||||
int filter_r = R;
|
||||
int filter_s = S;
|
||||
@@ -154,9 +160,10 @@ __global__ void Conv2dFprop(
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < kThreadN; ++n) {
|
||||
int thread_k = k_start + n;
|
||||
int k_group_idx = thread_k / k_per_group;
|
||||
|
||||
if (thread_k < problem_size.K) {
|
||||
element_B[n] = ElementAccumulator(tensor_w.at({thread_k, R, S, C}));
|
||||
if (thread_k < problem_size.K && k_group_idx == c_group_idx) {
|
||||
element_B[n] = ElementAccumulator(tensor_w.at({thread_k, R, S, C % c_per_group}));
|
||||
}
|
||||
else {
|
||||
element_B[n] = ElementAccumulator();
|
||||
|
||||
@@ -86,11 +86,14 @@ void Conv2dFprop(
|
||||
for (int q = 0; q < problem_size.Q; ++q) {
|
||||
for (int k = 0; k < problem_size.K; ++k) {
|
||||
|
||||
int group_idx = k / (problem_size.K / problem_size.groups);
|
||||
int channels_per_group = problem_size.C / problem_size.groups;
|
||||
|
||||
ElementAccumulator acc = ElementAccumulator();
|
||||
|
||||
for (int r = 0; r < problem_size.R; ++r) {
|
||||
for (int s = 0; s < problem_size.S; ++s) {
|
||||
for (int c = 0; c < problem_size.C; ++c) {
|
||||
for (int c = 0; c < channels_per_group; ++c) {
|
||||
|
||||
int filter_r = r;
|
||||
int filter_s = s;
|
||||
@@ -105,7 +108,7 @@ void Conv2dFprop(
|
||||
|
||||
if (h >= 0 && h < problem_size.H && w >= 0 && w < problem_size.W) {
|
||||
|
||||
ElementA a = tensor_x.at({n, h, w, c});
|
||||
ElementA a = tensor_x.at({n, h, w, c + group_idx * channels_per_group});
|
||||
ElementB b = tensor_w.at({k, r, s, c});
|
||||
|
||||
acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
|
||||
@@ -137,21 +140,21 @@ template <typename ElementA,
|
||||
typename LayoutB,
|
||||
typename ElementC,
|
||||
typename LayoutC,
|
||||
typename ElementAccumulator,
|
||||
typename ElementCompute,
|
||||
typename ElementAccumulator = ElementCompute,
|
||||
typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
|
||||
typename InnerProductOp = multiply_add<ElementAccumulator> >
|
||||
void Depsep_Fprop(
|
||||
cutlass::TensorView<ElementA, LayoutA> tensor_A,
|
||||
void Depsep_Fprop(cutlass::TensorView<ElementA, LayoutA> tensor_A,
|
||||
cutlass::TensorView<ElementB, LayoutB> tensor_B,
|
||||
cutlass::TensorView<ElementC, LayoutC> tensor_C,
|
||||
cutlass::TensorView<ElementC, LayoutC> tensor_D,
|
||||
ElementCompute alpha,
|
||||
ElementCompute beta,
|
||||
cutlass::Tensor4DCoord padding,
|
||||
cutlass::Coord<2> conv_stride,
|
||||
cutlass::Coord<2> dilation,
|
||||
cutlass::Tensor4DCoord padding = cutlass::Tensor4DCoord(),
|
||||
cutlass::Coord<2> conv_stride = cutlass::Coord<2>(),
|
||||
cutlass::Coord<2> dilation = cutlass::Coord<2>(),
|
||||
cutlass::conv::Mode mode = cutlass::conv::Mode::kCrossCorrelation) {
|
||||
|
||||
|
||||
ConvertOp convert_op;
|
||||
InnerProductOp inner_product_op;
|
||||
|
||||
@@ -163,15 +166,13 @@ void Depsep_Fprop(
|
||||
ElementAccumulator acc = ElementAccumulator();
|
||||
for (int r = 0; r < tensor_B.extent().h(); ++r) {
|
||||
for (int s = 0; s < tensor_B.extent().w(); ++s) {
|
||||
if ((p * conv_stride[0] - padding[0] + r * dilation[0]) < tensor_A.extent().h() &&
|
||||
(p * conv_stride[0] - padding[0] + r * dilation[0]) >= 0 &&
|
||||
(q * conv_stride[1] - padding[2] + s * dilation[1]) < tensor_A.extent().w() &&
|
||||
(q * conv_stride[1] - padding[2] + s * dilation[1]) >= 0) {
|
||||
ElementA a = tensor_A.at(
|
||||
cutlass::make_Coord(n,
|
||||
p * conv_stride[0] - padding[0] + r * dilation[0],
|
||||
q * conv_stride[1] - padding[2] + s * dilation[1],
|
||||
g));
|
||||
|
||||
// input activation H and W
|
||||
int h = p * conv_stride[0] - padding[0] + r * dilation[0];
|
||||
int w = q * conv_stride[1] - padding[2] + s * dilation[1];
|
||||
|
||||
if (h < tensor_A.extent().h() && h >= 0 && w < tensor_A.extent().w() && w >= 0) {
|
||||
ElementA a = tensor_A.at(cutlass::make_Coord(n, h, w, g));
|
||||
|
||||
ElementB b = (mode == cutlass::conv::Mode::kCrossCorrelation)
|
||||
? tensor_B.at(cutlass::make_Coord(g, r, s, 0))
|
||||
@@ -185,7 +186,7 @@ void Depsep_Fprop(
|
||||
|
||||
// Apply Epilogue, compute ElementCompute, convert and store ElementC
|
||||
ElementC c_ref = tensor_C.at(cutlass::make_Coord(n, p, q, g));
|
||||
tensor_C.at(cutlass::make_Coord(n, p, q, g)) =
|
||||
tensor_D.at(cutlass::make_Coord(n, p, q, g)) =
|
||||
convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -694,10 +694,6 @@ struct TensorFillSymmetricRandomUniformFunc {
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
//
|
||||
// We expect to release this with CUTLASS 2.4. -akerr
|
||||
|
||||
/// Computes a random Uniform distribution and pads diagonal with zeros
|
||||
template <
|
||||
typename Element, ///< Element type
|
||||
|
||||
Reference in New Issue
Block a user