C++ functor

functor

假想有一個函數的功能是將傳進來的參數加一之後回傳：

int add_one(int x){
    return x + 1;
}

如果我們想把該函數的功能改為加二，一種方法是再定義一個同樣的函數，但是把加一改成加二：

int add_two(int x){
    return x + 2;
}

雖然上面這種方法也能達到目的，但總顯得不夠簡潔。

另一種比較簡潔的方法則是定義一個functor（即function object）“類別”。包括三部份：

1. 定義一個私有成員變數（用來當作"加數"）
2. 建構子接受一個參數，把剛定義的私有成員變數（即"加數"）設為該值
3. overload ()這個運算子：其功能是把參數加上"加數"後回傳

class Adder{
private:
    int addend; //加數
public:
    Adder(int x) : addend(x) {}
    
    int operator() (int augend) const {
        return augend + addend;
    }
};

如此一來，我們每定義一個Adder類別的物件，就能得到為參數加一個特定"加數"的functor，而不必像上一個方法中為每一個"加數"都定義一個函數了：

#include <iostream>

// Driver code 
int main() 
{ 
    int arr[] = {1, 2, 3};
    
    //this is a functor that adds a number by 3
    Adder myadder(3); 
    
    for(int e : arr){
        std::cout << myadder(e) << " ";
    }
    std::cout << std::endl;
    // 4 5 6 

    myadder = Adder(5);
    
    for(int e : arr){
        std::cout << myadder(e) << " ";
    }
    std::cout << std::endl;
    // 6 7 8
} 

在TensorRT/samples/common/buffers.h中：

//!          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.
template <typename AllocFunc, typename FreeFunc>
class GenericBuffer
{
public:
    GenericBuffer(size_t size, nvinfer1::DataType type)
        : mSize(size)
        , mCapacity(size)
        , mType(type)
    {
        //調用allocFn的operator()函數來分配記憶體
        if (!allocFn(&mBuffer, this->nbBytes()))
        {
            //failure to allocate storage
            throw std::bad_alloc();
        }
    }

    ~GenericBuffer()
    {
        //調用freeFn的operator()函數來釋放記憶體
        freeFn(mBuffer);
    }

private:
    //...
    void* mBuffer;
    //兩個function object
    AllocFunc allocFn;
    FreeFunc freeFn;
};

class DeviceAllocator
{
public:
    //注意其參數是指標的指標void**
    //overload ()這個operator
    bool operator()(void** ptr, size_t size) const
    {
        return cudaMalloc(ptr, size) == cudaSuccess;
    }
};

class DeviceFree
{
public:
    void operator()(void* ptr) const
    {
        cudaFree(ptr);
    }
};

class HostAllocator
{
public:
    //ptr是指向指標的指標
    bool operator()(void** ptr, size_t size) const
    {
        //分配一塊記憶體，並將其位置指定給ptr指向的指標
        *ptr = malloc(size);
        //確認拿到的記憶體位置非空
        return *ptr != nullptr;
    }
};

class HostFree
{
public:
    void operator()(void* ptr) const
    {
        free(ptr);
    }
};

using DeviceBuffer = GenericBuffer<DeviceAllocator, DeviceFree>;
using HostBuffer = GenericBuffer<HostAllocator, HostFree>;

首先定義 DeviceAllocator， DeviceFree， HostAllocator及 HostFree等四個functor。

其次在GenericBuffer中定義了兩個相對應的物件allocFn及freeFn。

最後因為()已經被overload過了，所以可以透過 allocFn(&mBuffer, this->nbBytes()) 及 freeFn(mBuffer) 來調用。

參考連結

Functors in C++