TensorRT/samples/common/buffers.h - GenericBuffer源碼研讀
buffers.h - GenericBuffer
buffers.h
定義了用於緩存管理的BufferManager
。這個類別的核心是ManagedBuffer
數據結構。ManagedBuffer
由DeviceBuffer
及HostBuffer
組成。這兩者都是GenericBuffer
的特例。GenericBuffer
則是一個RAII class,管理著緩存的申請,釋放及查詢。
本文將先介紹buffers.h
中的GenericBuffer
,下一篇將繼續介紹BufferManager
。
/*
* Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
*
* 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 TENSORRT_BUFFERS_H
#define TENSORRT_BUFFERS_H
#include "NvInfer.h"
#include "common.h"
#include "half.h"
#include <cassert>
#include <cuda_runtime_api.h>
#include <iostream>
#include <iterator>
#include <memory>
#include <new>
#include <numeric>
#include <string>
#include <vector>
using namespace std;
namespace samplesCommon
{
//!
//! \brief The GenericBuffer class is a templated class for buffers.
//!
//! \details This templated RAII (Resource Acquisition Is Initialization) class handles the allocation,
//! deallocation, querying of buffers on both the device and the host.
//! It can handle data of arbitrary types because it stores byte buffers.
//! The template parameters AllocFunc and FreeFunc are used for the
//! allocation and deallocation of the buffer.
//! AllocFunc must be a functor that takes in (void** ptr, size_t size)
//! and returns bool. ptr is a pointer to where the allocated buffer address should be stored.
//! size is the amount of memory in bytes to allocate.
//! The boolean indicates whether or not the memory allocation was successful.
//! FreeFunc must be a functor that takes in (void* ptr) and returns void.
//! ptr is the allocated buffer address. It must work with nullptr input.
//!
//GenericBuffer的核心是mBuffer這個指標,
//其所包含的成員函數都是圍繞著mBuffer來展開
template <typename AllocFunc, typename FreeFunc>
class GenericBuffer
{
public:
//!
//! \brief Construct an empty buffer.
//!
//nvinfer1::DataType定義於NvInferRuntimeCommon.h
/*
enum class DataType : int
{
kFLOAT = 0, //!< FP32 format.
kHALF = 1, //!< FP16 format.
kINT8 = 2, //!< quantized INT8 format.
kINT32 = 3, //!< INT32 format.
kBOOL = 4 //!< BOOL format.
};
*/
GenericBuffer(nvinfer1::DataType type = nvinfer1::DataType::kFLOAT) //type的默認值
: mSize(0)
, mCapacity(0)
, mType(type)
, mBuffer(nullptr)
{
}
//!
//! \brief Construct a buffer with the specified allocation size in bytes.
//!
GenericBuffer(size_t size, nvinfer1::DataType type)
: mSize(size)
, mCapacity(size)
, mType(type)
{
//allocFn第一個參數是指標的指標
if (!allocFn(&mBuffer, this->nbBytes()))
{
//failure to allocate storage
throw std::bad_alloc();
}
}
//move constructor
//此處的GenericBuffer&&是rvalue reference
GenericBuffer(GenericBuffer&& buf)
//獲取temporary object的成員屬性
: mSize(buf.mSize)
, mCapacity(buf.mCapacity)
, mType(buf.mType)
, mBuffer(buf.mBuffer)
{
//將temp object還原到一個valid state
buf.mSize = 0;
buf.mCapacity = 0;
buf.mType = nvinfer1::DataType::kFLOAT;
buf.mBuffer = nullptr;
}
//function overriding
//"="運算子也有回傳值,回傳的是該物件本身
//move assignment operator
GenericBuffer& operator=(GenericBuffer&& buf)
{
//this及&buf都是指向GenericBuffer物件的指標
if (this != &buf)
{
//清理現有的資料
freeFn(mBuffer);
//從temporary object中獲取資料
mSize = buf.mSize;
mCapacity = buf.mCapacity;
mType = buf.mType;
mBuffer = buf.mBuffer;
// Reset buf.
//將temporary state還原到一個valid state
//注意!在使用move assignment operator時等號右邊的GenericBuffer物件會被reset!
buf.mSize = 0;
buf.mCapacity = 0;
buf.mBuffer = nullptr;
}
return *this;
}
//!
//! \brief Returns pointer to underlying array.
//!
void* data()
{
return mBuffer;
}
//!
//! \brief Returns pointer to underlying array.
//!
const void* data() const
{
return mBuffer;
}
//!
//! \brief Returns the size (in number of elements) of the buffer.
//!
//回傳buffer的元素個數
size_t size() const
{
return mSize;
}
//!
//! \brief Returns the size (in bytes) of the buffer.
//!
//回傳buffer所佔用的byte數
size_t nbBytes() const
{
//samplesCommon::getElementSize:回傳各種數據類型所佔用的byte數
return this->size() * samplesCommon::getElementSize(mType);
}
//!
//! \brief Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity.
//!
void resize(size_t newSize)
{
mSize = newSize;
//只有當newSize大於現有容量mCapacity時才會進行如下操作
if (mCapacity < newSize)
{
//做resize會清空現有mBuffer?
freeFn(mBuffer);
//this->nbBytes是一個函數,回傳的值會隨著新設定的mSize改變
if (!allocFn(&mBuffer, this->nbBytes()))
{
//在分配記憶體的函數出錯時會拋出此異常
throw std::bad_alloc{};
}
mCapacity = newSize;
}
}
//!
//! \brief Overload of resize that accepts Dims
//!
//將參數由nvinfer1::Dims轉為size_t後,呼叫void resize(size_t newSize)
//nvinfer1::Dims定義於NvInferRuntimeCommon.h
/*
class Dims
{
public:
static const int MAX_DIMS = 8; //!< The maximum number of dimensions supported for a tensor.
int nbDims; //!< The number of dimensions.
int d[MAX_DIMS]; //!< The extent of each dimension.
TRT_DEPRECATED DimensionType type[MAX_DIMS]; //!< The type of each dimension.
};
*/
void resize(const nvinfer1::Dims& dims)
{
//samplesCommon::volume計算"dims"各維度的大小的連乘積,即體積
return this->resize(samplesCommon::volume(dims));
}
~GenericBuffer()
{
freeFn(mBuffer);
}
private:
//mSize:mBuffer擁有的元素個數
size_t mSize{0}, mCapacity{0};
nvinfer1::DataType mType;
void* mBuffer;
AllocFunc allocFn;
FreeFunc freeFn;
};
//用於顯卡上記憶體分配的functor
class DeviceAllocator
{
public:
//overload ()這個operator
//用於分配size大小的記憶體
//注意其參數是指標的指標void**
bool operator()(void** ptr, size_t size) const
{
/*
__host____device__cudaError_t cudaMalloc ( void** devPtr, size_t size )
Allocate memory on the device.
Parameters
devPtr
- Pointer to allocated device memory
size
- Requested allocation size in bytes
Returns
cudaSuccess, cudaErrorInvalidValue, cudaErrorMemoryAllocation
*/
return cudaMalloc(ptr, size) == cudaSuccess;
}
};
//用於顯卡上記憶體釋放的functor
class DeviceFree
{
public:
void operator()(void* ptr) const
{
/*
__host____device__cudaError_t cudaFree ( void* devPtr )
Frees memory on the device.
Parameters
devPtr
- Device pointer to memory to free
Returns
cudaSuccess, cudaErrorInvalidValue
*/
cudaFree(ptr);
}
};
//用於記憶體分配的functor
class HostAllocator
{
public:
//ptr是指向指標的指標
bool operator()(void** ptr, size_t size) const
{
//分配一塊記憶體,並將其位置指定給ptr指向的指標
/*
void* malloc (size_t size);
Allocate memory block
Allocates a block of size bytes of memory, returning a pointer to the beginning of the block.
*/
*ptr = malloc(size);
//確認拿到的記憶體位置非空
return *ptr != nullptr;
}
};
//用於記憶體釋放的functor
class HostFree
{
public:
void operator()(void* ptr) const
{
/*
void free (void* ptr);
Deallocate memory block
A block of memory previously allocated by a call to malloc, calloc or realloc is deallocated, making it available again for further allocations.
*/
free(ptr);
}
};
//GenericBuffer的兩個特例:在host上的buffer及在device上的buffer
using DeviceBuffer = GenericBuffer<DeviceAllocator, DeviceFree>;
using HostBuffer = GenericBuffer<HostAllocator, HostFree>;
//!
//! \brief The ManagedBuffer class groups together a pair of corresponding device and host buffers.
//!
//這個類別的作用只是把DeviceBuffer和HostBuffer綁在一起而已
class ManagedBuffer
{
public:
DeviceBuffer deviceBuffer;
HostBuffer hostBuffer;
};
//此處省略BufferManager...
} // namespace samplesCommon
#endif // TENSORRT_BUFFERS_H
RAII
在官方註釋中,GenericBuffer
及BufferManager
兩個類別都被標記為RAII (Resource Acquisition Is Initialization) class,這是什麼意思呢?詳見C++ RAII(Resource Acquisition Is Initialization)。
functor
定義DeviceBuffer
及HostBuffer
的過程中使用到了DeviceAllocator
, DeviceFree
, HostAllocator
及 HostFree
等四個functor。關於functor,詳見C++ functor。
const member function
在DeviceAllocator::operator()
函數的定義中,被加了一個const
,這使它成為了const member function。關於const member function,詳見C++ const member function。
回傳const指標/參考的函數
GenericBuffer::data()
及BufferManager::getDeviceBindings()
兩個函數的回傳值前都被加了一個const
,這是什麼意思呢?詳見C 回傳const指標/參考的函數。
using
關於using
的用法,詳見C++ using。
std::bad_alloc
std::bad_alloc is the type of the object thrown as exceptions
by the allocation functions to report failure to allocate storage.
如果一個用於分配記憶體的函數出錯了,我們會讓它拋出std::bad_alloc
的異常來通知使用者。
function overloading v.s. function overriding
function overloading:可以定義多個同名不同signature(參數個數,型別,回傳值型別)的函數。achieved at compile time。
function overriding:衍生類別可以重新定義基礎類別裡定義過的函數。achieved at run time。
operator overloading
自定義類別的物件可以overload()
,=
等運算子,如下:
class DeviceAllocator
{
public:
bool operator()(void** ptr, size_t size) const
{
return cudaMalloc(ptr, size) == cudaSuccess;
}
};
在定義DeviceBuffer
時,用到了上面定義的DeviceAllocator
:
using DeviceBuffer = GenericBuffer<DeviceAllocator, DeviceFree>;
DeviceAllocator
在下面的代碼中是以AllocFunc
的身份出現的。GenericBuffer
裡定義了一個DeviceAllocator
型別的變數allocFn
:
AllocFunc allocFn;
在GenericBuffer
的建構子中用到了allocFn()
函數:
//!
//! \brief Construct a buffer with the specified allocation size in bytes.
//!
GenericBuffer(size_t size, nvinfer1::DataType type)
: mSize(size)
, mCapacity(size)
, mType(type)
{
if (!allocFn(&mBuffer, this->nbBytes()))
{
//failure to allocate storage
throw std::bad_alloc();
}
}
參考連結
C++ RAII(Resource Acquisition Is Initialization)
Reference and Constant Parameters
Function Overloading vs Function Overriding in C++
Function Call Operator () Overloading in C++