sources/libs/nvdsinferserver/infer_utils.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-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.
 */
 
#ifndef __NVDSINFER_SERVER_INFER_UTILS_H__
#define __NVDSINFER_SERVER_INFER_UTILS_H__
 
#include <infer_batch_buffer.h>
#include <infer_common.h>
#include <infer_datatypes.h>
 
#include <algorithm>
#include <iomanip>
#include <iostream>
#include <limits>
#include <numeric>
 
namespace nvdsinferserver INFER_EXPORT_API {
 
void dsInferLogPrint__(NvDsInferLogLevel level, const char* fmt, ...);
 
void dsInferLogVPrint__(NvDsInferLogLevel level, const char* fmt, va_list args);
 
 
inline const char* safeStr(const char* str) {
    return !str ? "" : str;
}
 
inline const char* safeStr(const std::string& str) {
    return str.c_str();
}
inline bool string_empty(const char* str) {
    return !str || strlen(str) == 0;
}
 
inline bool file_accessible(const char* path) {
    assert(path);
    return (access(path, F_OK) != -1);
}
 
inline bool file_accessible(const std::string& path) {
    return (!path.empty()) && file_accessible(path.c_str());
}
template <typename T>
inline bool
isNonBatch(T b)
{
    return b == 0;
}
 
std::string dims2Str(const InferDims& d);
std::string batchDims2Str(const InferBatchDims& d);
std::string dataType2Str(const InferDataType type);
std::string dataType2GrpcStr(const InferDataType type);
InferDataType grpcStr2DataType(const std::string &type);
NvDsInferNetworkInfo dims2ImageInfo(
    const InferDims& d, InferTensorOrder order);
std::string tensorOrder2Str(InferTensorOrder order);
bool fEqual(float a, float b);
 
class DlLibHandle {
public:
    DlLibHandle(const std::string& path, int mode = RTLD_LAZY);
 
    /*
     * @brief Destructor. Close the dynamically loaded library.
     */
    ~DlLibHandle();
 
    /*
     * @brief Check that the library handle is valid.
     */
    bool isValid() const { return m_LibHandle; }
 
    /*
     * @brief Get the filename of the library.
     */
    const std::string& getPath() const { return m_LibPath; }
 
    /*
     * @brief Get the function pointer from the library for given function
     * name.
     */
    template <typename FuncPtr>
    FuncPtr symbol(const char* func) {
        assert(!string_empty(func));
        if (!m_LibHandle)
            return nullptr;
        InferDebug("lib: %s dlsym :%s", safeStr(m_LibPath), safeStr(func));
        return (FuncPtr)dlsym(m_LibHandle, func);
    }
 
    template <typename FuncPtr>
    FuncPtr symbol(const std::string& func) {
        return symbol<FuncPtr>(func.c_str());
    }
private:
    /*
     * @brief Handle for the library returned by dlopen().
     */
    void* m_LibHandle{nullptr};
    /*
     * @brief Filename of the dynamically loaded library.
     */
    const std::string m_LibPath;
};
 
class WakeupException : public std::exception {
    std::string m_Msg;
 
public:
    WakeupException(const std::string& s) : m_Msg(s) {}
    const char* what() const noexcept override { return m_Msg.c_str(); }
};
 
template <typename Container>
class GuardQueue {
public:
    typedef typename Container::value_type T;
    void push(T data) {
        std::unique_lock<std::mutex> lock(m_Mutex);
        m_Queue.emplace_back(std::move(data));
        m_Cond.notify_one();
    }
    T pop() {
        std::unique_lock<std::mutex> lock(m_Mutex);
        m_Cond.wait(
            lock, [this]() { return m_WakeupOnce || !m_Queue.empty(); });
        if (m_WakeupOnce) {
            m_WakeupOnce = false;
            InferDebug("GuardQueue pop end on wakeup signal");
            throw WakeupException("GuardQueue stopped");
        }
        assert(!m_Queue.empty());
        T ret = std::move(*m_Queue.begin());
        m_Queue.erase(m_Queue.begin());
        return ret;
    }
    void wakeupOnce() {
        InferDebug("GuardQueue trigger wakeup once");
        std::unique_lock<std::mutex> lock(m_Mutex);
        m_WakeupOnce = true;
        m_Cond.notify_all();
    }
    void clear() {
        InferDebug("GuardQueue clear");
        std::unique_lock<std::mutex> lock(m_Mutex);
        m_Queue.clear();
        m_WakeupOnce = false;
    }
    int size() {
        std::unique_lock<std::mutex> lock(m_Mutex);
        return m_Queue.size();
    }
 
private:
    std::mutex m_Mutex;
    std::condition_variable m_Cond;
    Container m_Queue;
    bool m_WakeupOnce = false;
};
 
template <typename Container>
class QueueThread {
public:
    using Item = typename Container::value_type;
    using RunFunc = std::function<bool(Item)>;
 
    QueueThread(RunFunc runFunc, const std::string& name) : m_Run(runFunc) {
        std::promise<void> p;
        std::future<void> f = p.get_future();
        InferDebug("QueueThread starting new thread");
        m_Thread = std::thread([&p, this]() {
            p.set_value();
            this->threadLoop();
        });
        setThreadName(name);
        f.wait();
    }
    void setThreadName(const std::string &name) {
        assert(!name.empty());
        m_Name = name;
        if (m_Thread.joinable()) {
            const int kMakLen = 16;
            char cName[kMakLen];
            strncpy(cName, name.c_str(), kMakLen);
            cName[kMakLen - 1] = 0;
            if (pthread_setname_np(m_Thread.native_handle(), cName) != 0) {
                InferError("set thread name: %s failed", safeStr(name));
                return;
            }
            InferDebug("QueueThread set new thread name:%s", cName);
        }
    }
    ~QueueThread() { join(); }
    void join() {
        InferDebug("QueueThread: %s join", safeStr(m_Name));
        if (m_Thread.joinable()) {
            m_Queue.wakeupOnce();
            m_Thread.join();
        }
        m_Queue.clear();
    }
    bool queueItem(Item item) {
        m_Queue.push(std::move(item));
        return true;
    }
 
private:
    void threadLoop() {
        while (true) {
            try {
                Item item = m_Queue.pop();
                if (!m_Run(std::move(item))) {
                    InferDebug("QueueThread:%s return and stop", safeStr(m_Name));
                    return;
                }
            }
            catch (const WakeupException& e) {
                InferDebug("QueueThread:%s stopped", safeStr(m_Name));
                return;
            }
            catch (...) {  // unexpected
                InferError(
                    "QueueThread:%s internal unexpected error, may cause stop",
                    safeStr(m_Name));
                // Usually can move on to next, but need developer to check
                continue;
            }
        }
    }
 
private:
    std::thread m_Thread;
    std::string m_Name;
    RunFunc m_Run;
    GuardQueue<Container> m_Queue;
};
 
template <class UniPtr>
class BufferPool : public std::enable_shared_from_this<BufferPool<UniPtr> > {
public:
    using ItemType = typename UniPtr::element_type;
    using RecylePtr = std::unique_ptr<ItemType, std::function<void(ItemType*)>>;
    BufferPool(const std::string &name)
        : m_Name(name) {}
    virtual ~BufferPool()
    {
        InferDebug(
            "BufferPool: %s deleted with free buffer size:%d", safeStr(m_Name),
            m_FreeBuffers.size());
    }
    bool setBuffer(UniPtr buf) {
        assert(buf);
        buf->reuse();
        m_FreeBuffers.push(std::move(buf));
        InferDebug("BufferPool: %s set buf to free, available size:%d",
            safeStr(m_Name), m_FreeBuffers.size());
        return true;
    }
    int size() { return m_FreeBuffers.size(); }
 
    RecylePtr acquireBuffer() {
        try {
            UniPtr p = m_FreeBuffers.pop();
            auto deleter = p.get_deleter();
            std::weak_ptr<BufferPool<UniPtr>> poolPtr =
                this->shared_from_this();
            RecylePtr recBuf(
                p.release(), [poolPtr, d = deleter](ItemType* buf) {
                assert(buf);
                UniPtr data(buf, d);
                auto pool = poolPtr.lock();
                if (pool) {
                    InferDebug("BufferPool: %s release a buffer", safeStr(pool->m_Name));
                    pool->setBuffer(std::move(data));
                } else {
                    InferError("BufferPool is deleted, check internal error.");
                    assert(false);
                }
            });
            InferDebug("BufferPool: %s acquired buffer, available free buffer left:%d",
                safeStr(m_Name), m_FreeBuffers.size());
            return recBuf;
        } catch (...) {
            InferDebug(
                "BufferPool: %s acquired buffer failed, queue maybe waked up.",
                safeStr(m_Name));
            assert(false);
            return nullptr;
        }
    }
 
private:
    GuardQueue<std::deque<UniPtr>> m_FreeBuffers;
    const std::string m_Name;
};
 
template <class UniPtr>
using SharedBufPool = std::shared_ptr<BufferPool<UniPtr>>;
 
template<typename Key, typename UniqBuffer>
class MapBufferPool {
public:
    using SharedPool = SharedBufPool<UniqBuffer>;
    using RecylePtr = typename BufferPool<UniqBuffer>::RecylePtr;
public:
    MapBufferPool(const std::string &name): m_Name(name) {}
    virtual ~MapBufferPool()
    {
        InferDebug(
            "MapBufferPool: %s deleted with buffer pool size:%d",
            safeStr(m_Name), (int)m_MapPool.size());
    }
 
    MapBufferPool(const MapBufferPool& other) = delete;
    MapBufferPool& operator=(const MapBufferPool& other) = delete;
    bool setBuffer(const Key &key, UniqBuffer buf) {
        std::unique_lock<std::shared_timed_mutex> uniqLock(m_MapPoolMutex);
        assert(buf);
        SharedPool &pool = m_MapPool[key];
        if(!pool) {
            uint32_t id = m_MapPool.size() - 1;
            std::string poolName = m_Name + std::to_string(id);
            pool = std::make_shared<BufferPool<UniqBuffer>>(poolName);
            assert(pool);
            InferDebug("MapBufferPool: %s create new pool id:%d",
                safeStr(m_Name), id);
        }
        if(!pool) {
            return false;
        }
        return pool->setBuffer(std::move(buf));
    }
    uint32_t getPoolSize(const Key &key) {
        SharedPool pool = findPool(key);
        if (!pool)
            return 0;
        return pool->size();
    }
 
    RecylePtr acquireBuffer(const Key &key) {
        SharedPool pool = findPool(key);
        assert(pool);
        if (!pool) {
            InferWarning(
                "MapBufferPool: %s acquire buffer failed, no key found",
                safeStr(m_Name));
            return nullptr;
        }
        InferDebug("MapBufferPool: %s acquire buffer", safeStr(m_Name));
        return pool->acquireBuffer();
    }
    void clear() {
        InferDebug("MapBufferPool: %s clear all buffers", safeStr(m_Name));
        std::unique_lock<std::shared_timed_mutex> uniqLock(m_MapPoolMutex);
        m_MapPool.clear();
    }
 
private:
    SharedPool findPool(const Key& key) {
        std::shared_lock<std::shared_timed_mutex> sharedLock(m_MapPoolMutex);
        auto iter = m_MapPool.find(key);
        if (iter != m_MapPool.end()) {
            assert(iter->second);
            return iter->second;
        }
        return nullptr;
    }
 
private:
    std::map<Key, SharedPool> m_MapPool;
    std::shared_timed_mutex m_MapPoolMutex;
    const std::string m_Name;
};
 
inline uint32_t getElementSize(InferDataType t) {
    switch (t) {
    case InferDataType::kInt32:
    case InferDataType::kUint32:
    case InferDataType::kFp32:
        return 4;
    case InferDataType::kFp16:
    case InferDataType::kInt16:
    case InferDataType::kUint16:
        return 2;
    case InferDataType::kInt8:
    case InferDataType::kUint8:
    case InferDataType::kBool:
        return 1;
    case InferDataType::kString:
        return 0;
    case InferDataType::kFp64:
    case InferDataType::kInt64:
    case InferDataType::kUint64:
        return 8;
    default:
        InferError("Failed to get element size on Unknown datatype:%d",
            static_cast<int>(t));
        return 0;
    }
}
 
inline bool
hasWildcard(const InferDims& dims)
{
    return std::any_of(
        dims.d, dims.d + dims.numDims,
        [](int d) { return d <= INFER_WILDCARD_DIM_VALUE; });
}
 
inline size_t
dimsSize(const InferDims& dims)
{
    if (hasWildcard(dims) || !dims.numDims) {
        return 0;
    } else {
        return std::accumulate(
            dims.d, dims.d + dims.numDims, 1,
            [](int s, int i) { return s * i; });
    }
}
 
inline void
normalizeDims(InferDims& dims)
{
    dims.numElements = dimsSize(dims);
}
 
bool operator<=(const InferDims& a, const InferDims& b);
bool operator>(const InferDims& a, const InferDims& b);
bool operator==(const InferDims& a, const InferDims& b);
bool operator!=(const InferDims& a, const InferDims& b);
struct LayerDescription;
 
NvDsInferLayerInfo toCapi(const LayerDescription& desc, void* bufPtr);
 
NvDsInferDims toCapi(const InferDims &dims);
 
NvDsInferLayerInfo toCapiLayerInfo(
    const InferBufferDescription& desc, void* buf = nullptr);
 
NvDsInferDataType toCapiDataType(InferDataType dt);
 
bool intersectDims(
    const InferDims& a, const InferDims& b, InferDims& c);
 
bool isPrivateTensor(const std::string &tensorName);
 
std::string joinPath(const std::string& a, const std::string& b);
std::string dirName(const std::string& path);
bool isAbsolutePath(const std::string& path);
bool realPath(const std::string &inPath, std::string &absPath);
bool isCpuMem(InferMemType type);
 
std::string memType2Str(InferMemType type);
 
InferDims fullDims(int batchSize, const InferDims& in);
 
bool debatchFullDims(
    const InferDims& full, InferDims& debatched, uint32_t& batch);
 
bool squeezeMatch(const InferDims& a, const InferDims& b);
 
SharedBatchBuf ReshapeBuf(
    const SharedBatchBuf& in, uint32_t batch, const InferDims& dims,
    bool reCalcBytes = false);
 
SharedBatchBuf reshapeToFullDimsBuf(const SharedBatchBuf& buf,
        bool reCalcBytes = false);
 
NvDsInferStatus tensorBufferCopy(
    const SharedBatchBuf& in, const SharedBatchBuf& out,
    const SharedCuStream& stream);
 
} // namespace nvdsinferserver
 
extern "C" {
 
INFER_EXPORT_API const char* NvDsInferStatus2Str(NvDsInferStatus status);
 
INFER_EXPORT_API bool validateInferConfigStr(
    const std::string& configStr, const std::string& path, std::string& updated);
}
 
#endif