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：

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）

std::bad_alloc

Reference and Constant Parameters

Function Overloading vs Function Overriding in C++

Function Call Operator () Overloading in C++

Assignment Operators Overloading in C++

malloc

free

C++ functor

C++ const member function

C 回傳const指標/參考的函數

C++ using