CE_MemoryPool.h

Go to the documentation of this file.

/*
 * PROPRIETARY INFORMATION.  This software is proprietary to
 * Side Effects Software Inc., and is not to be reproduced,
 * transmitted, or disclosed in any way without written permission.
 *
 * NAME: CE_MemoryPool.h ( CE Library, C++)
 *
 * COMMENTS: Compute Engine Memory Pool.
 */

#ifndef __CE_MemoryPool__
#define __CE_MemoryPool__

#include "CE_API.h"

#include <UT/UT_Array.h>
#include <UT/UT_Lock.h>
#include <UT/UT_LRUCache.h>
#include <UT/UT_Map.h>
#include <UT/UT_MemoryResource.h>
#include <UT/UT_UniquePtr.h>
#include <SYS/SYS_Types.h>

#include <iosfwd>

#if UT_ASSERT_LEVEL >= UT_ASSERT_LEVEL_NORMAL
#include <set>
#endif

class CE_Context;

/// This class implements a memory pool for OpenCL buffer objects, using an
/// LRUCache to limit the size of the pool.
class CE_API CE_MemoryPool
{
public:
    /// Initialize the pool with specified maximum size.
    CE_MemoryPool(int64 size);
    virtual ~CE_MemoryPool();

    /// Empty the memory pool of all held buffers.
    void        clear();

    /// Set the maximums size of the pool.  This will prune the LRUCache as
    /// needed to fit the size.
    void        setMaxSize(int64 size) { myLRUCache.setMaxSize(size); }

    /// Returns the maximum size of the pool.
    int64       getMaxSize() const { return myLRUCache.maxSize(); }

    /// Dump memory usage statistics to stderr.
    void        reportUsage() const;

    /// Dump memory to a stream
    void        reportUsage(std::ostream &os) const;

    /// Query the current usage; total is the total amount of memory held by the
    /// pool, not_in_use is the amount that can be freed immediately on demand.
    void        getUsage(exint& total, exint& not_in_use) const
    {
        not_in_use = myLRUCache.currentSize();
        total = not_in_use + myInUseSize;
    }

    /// Allocate a buffer for the provided context.  If a buffer of the
    /// specified size already exists in the pool, it will be returned.
    /// Otherwise a new one will be created on the device.
    cl::Buffer  allocBuffer(CE_Context *, int64 size);

    /// Check the provided buffer's reference count.  If it is about to be
    /// released, return it to the pool if it will fit, else release it.  If
    /// the reference count is higher than 1, meaning it is still active in
    /// an asynchronous operation like a kernel, put it on the inUse list to
    /// check later in sweepInUse().
    void        returnBuffer(cl::Buffer &&buf, bool use_pool);

    /// Sweep the inUse list, looking for buffers that are no longer being used
    /// in an asynchronous operation and can be returned to the pool. Returns
    /// the amount of released memory (i.e. memory that's returned back to the
    /// system).
    exint       sweepInUse();

    /// Attempts to release the requested amount of memory back to the system.
    /// Returns the amount actually freed.
    exint       freeMemory(exint amount);

    /// Registers this memory pool as a client for the given resource. This must
    /// be called with the appropriate memory resource to ensure that the memory
    /// pool can be asked to unload when device memory is tight.
    void        registerClient(UT_MemoryResource* resource)
    {
        resource->registerClient(&myMemoryClient, UT_MemoryClient::NORMAL_FREE);
    }
    /// Unregisters this memory pool as a client for the given resource. This
    /// must be called on destruction with the same resource used with
    /// registerClient().
    void        unregisterClient(UT_MemoryResource* resource)
    {
        resource->unregisterClient(&myMemoryClient);
    }

private:
    // This class just holds a cl::Buffer while it's in the LRUCache.  It's also
    // responsible for alerting the CE_MemoryPool object when it's deleted
    // by pruning or other means.  Each cl::Buffer is owned by only one of
    // these objects, as maintained by being non-copyable and having move
    // semantics.
    class ut_clBuffer : public UT_NonCopyable
    {
    public:
        ut_clBuffer(CE_MemoryPool *pool, const cl::Buffer &buf, int64 size);

        ut_clBuffer(ut_clBuffer &&b);

        ~ut_clBuffer();

        ut_clBuffer &operator=(ut_clBuffer &&b);

        static exint size(const ut_clBuffer &b) { return b.mySize; }

    private:
        int64 mySize;
        cl::Buffer myBuffer;
        CE_MemoryPool *myPool;
    };

    /// This memory client is used to ensure that this memory pool plays nicely
    /// with other users of memory on the OpenCL device.
    class MemoryClient : public UT_MemoryClient
    {
    public:
        MemoryClient(CE_MemoryPool& pool) : myMemoryPool(pool) {}

        const char* name() const override
        {
            return "OpenCL Buffer Pool";
        }

        bool freeMemoryRequest(const UT_MemoryResource* resource,
                               RequestSeverity severity, Niceness niceness,
                               exint amount, exint& freed_amount) override
        {
            freed_amount = myMemoryPool.sweepInUse();
            if (freed_amount < amount)
                freed_amount += myMemoryPool.freeMemory(amount - freed_amount);
            return freed_amount > 0;
        }

        bool memoryUse(const UT_MemoryResource* resource, exint& in_use,
                       exint& cache) override
        {
            myMemoryPool.getUsage(in_use, cache);
            return true;
        }

    private:
        CE_MemoryPool&          myMemoryPool;
    };
    
    friend class ut_clBuffer;
    friend class CE_Context;

    bool        getBufferFromPool(int64 size, cl::Buffer &buf);
    
    /// Add the buffer to the pool. Returns false if it operation fails.
    bool        appendToPool(const cl::Buffer &buf, int64 size);

    /// Remove the the buffer from the pool.
    void        releaseBuffer(const cl::Buffer &buf, int64 size);

    /// Used for tracing and debugging buffer allocations.
    void        recordAlloc(cl_mem buffer, int64 size);
    void        recordDealloc(cl_mem buffer, int64 size, bool returning=true);

    using BufferCache = UT_LRUCache<cl_mem, ut_clBuffer, ut_clBuffer::size>;
    using BufferTable = UT_Map<exint, UT_Array<cl::Buffer *>>;
    using BufferPtr = UT_UniquePtr<cl::Buffer>;

    BufferCache                 myLRUCache;
    UT_Array<BufferPtr>         myInUseList;
    BufferTable                 myBufferTable;
    int64                       myInUseSize;
    int64                       myTotalAllocSize;
    mutable UT_Lock             myMemoryPoolLock;

    /// The memory client that's registered with the OpenCL device.
    MemoryClient                myMemoryClient;

#if UT_ASSERT_LEVEL >= UT_ASSERT_LEVEL_NORMAL
    // Used for catching dangling buffers in debug/opt builds.
    void validateAndFreeBuffer(cl_mem buf);
    static void memDestructorCallback(cl_mem, void *);
    UT_Map<cl_mem, int64>       myActiveBuffers;
    std::multiset<cl_mem>       myReturnedBuffers;
#endif
};


#endif