154 lines
4.9 KiB
Plaintext
154 lines
4.9 KiB
Plaintext
/*
|
|
* 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 "maxPoolKernel.h"
|
|
|
|
#include <cstdint>
|
|
|
|
template <typename T>
|
|
__device__ __forceinline__ const T& max(const T& a, const T& b);
|
|
|
|
template <>
|
|
__device__ __forceinline__ const half& max(const half& a, const half& b)
|
|
{
|
|
#if __CUDA_ARCH__ >= 700 || !defined(__CUDA_ARCH__)
|
|
return __hgt(a, b) ? a : b;
|
|
#else
|
|
return (static_cast<float>(a) > static_cast<float>(b)) ? a : b;
|
|
#endif
|
|
}
|
|
|
|
template <>
|
|
__device__ __forceinline__ const float& max(const float& a, const float& b)
|
|
{
|
|
return (a > b) ? a : b;
|
|
}
|
|
|
|
template <>
|
|
__device__ __forceinline__ int8_t const& max(int8_t const& a, int8_t const& b)
|
|
{
|
|
return (a > b) ? a : b;
|
|
}
|
|
|
|
// Cuda kernel to find maximum in the kernelsize matrix
|
|
template <typename T>
|
|
__global__ void maxKernel(
|
|
int32_t B, int32_t C, int32_t H, int32_t W,
|
|
const T* input,
|
|
T* output, int32_t kernsize, int32_t stride, int32_t pad)
|
|
{
|
|
// Total input volume
|
|
int32_t const N = B * C * H * W;
|
|
int32_t out_id;
|
|
int32_t b, c, h, w;
|
|
|
|
int32_t const H_out = (H + 2 * pad - kernsize) / stride + 1;
|
|
int32_t const W_out = (W + 2 * pad - kernsize) / stride + 1;
|
|
|
|
// Index in the output tensor
|
|
out_id = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (out_id > B * C * H_out * W_out - 1)
|
|
{
|
|
return;
|
|
}
|
|
T maxim = static_cast<T>(0);
|
|
|
|
// Output index of batch
|
|
b = out_id / (C * H_out * W_out);
|
|
int32_t const temp = out_id % (C * H_out * W_out);
|
|
|
|
// Output index of channels
|
|
c = temp / (H_out * W_out);
|
|
int32_t const x = temp % (H_out * W_out);
|
|
|
|
// Output index of height
|
|
h = x / W_out; // row major format
|
|
|
|
// Output index of width
|
|
w = x % W_out;
|
|
|
|
// Index in input tensor considering stride
|
|
int32_t k = (b * C * H * W) + (c * (H * W)) + (h * stride * W) + (w * stride);
|
|
|
|
maxim = input[k];
|
|
// Find maximum value in the kernelsize matrix
|
|
for (int32_t i = k; i < k + kernsize; i++)
|
|
{
|
|
for (int32_t j = 0; j < kernsize; j++)
|
|
{
|
|
if ((i + (j * W)) < N)
|
|
maxim = max(maxim, input[i + (j * W)]);
|
|
}
|
|
}
|
|
|
|
output[out_id] = maxim;
|
|
}
|
|
|
|
int32_t maxPoolFloat(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
|
int32_t kernsize, int32_t stride, int32_t pad)
|
|
{
|
|
int32_t const blocksize = 512;
|
|
// Compute number of entries in output
|
|
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
|
int32_t const g = ((K + blocksize - 1) / blocksize);
|
|
|
|
maxKernel<float><<<g, blocksize, 0, stream>>>(
|
|
batch_size, C, H, W, static_cast<const float*>(input), static_cast<float*>(output), kernsize, stride, pad);
|
|
auto retVal = cudaStreamSynchronize(stream);
|
|
if (retVal != cudaSuccess)
|
|
{
|
|
return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int32_t maxPoolHalf(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
|
int32_t kernsize, int32_t stride, int32_t pad)
|
|
{
|
|
int32_t const blocksize = 512;
|
|
// Compute number of entries in output
|
|
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
|
int32_t const g = ((K + blocksize - 1) / blocksize);
|
|
|
|
maxKernel<half><<<g, blocksize, 0, stream>>>(
|
|
batch_size, C, H, W, static_cast<const half*>(input), static_cast<half*>(output), kernsize, stride, pad);
|
|
auto retVal = cudaStreamSynchronize(stream);
|
|
if (retVal != cudaSuccess)
|
|
{
|
|
return 1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
int32_t maxPoolInt8(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
|
int32_t kernsize, int32_t stride, int32_t pad)
|
|
{
|
|
int32_t const blocksize = 512;
|
|
// Compute number of entries in output
|
|
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
|
int32_t const g = ((K + blocksize - 1) / blocksize);
|
|
|
|
maxKernel<int8_t><<<g, blocksize, 0, stream>>>(
|
|
batch_size, C, H, W, static_cast<int8_t const*>(input), static_cast<int8_t*>(output), kernsize, stride, pad);
|
|
auto retVal = cudaStreamSynchronize(stream);
|
|
if (retVal != cudaSuccess)
|
|
{
|
|
return 1;
|
|
}
|
|
return 0;
|
|
}
|