NvInferSafeRuntime.h

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-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.
 */
 
// @brief Main header file for the NVIDIA Safe Runtime API.
// This file provides the primary interface for users to interact with the NVIDIA Safe Runtime API.
// It includes the necessary definitions, classes, and functions for creating and managing safe graphs,
// executing inference, and handling errors.
// Users should include this header file in their application to access the Safe Runtime API functionality.
 
#ifndef NV_INFER_SAFE_RUNTIME_H
#define NV_INFER_SAFE_RUNTIME_H
#include "NvInferSafeMemAllocator.h"
#include "NvInferSafePlugin.h"
#include "NvInferSafeRecorder.h"
#include <algorithm>
#include <cuda_fp16.h>
 
namespace nvinfer2
{
namespace safe
{
using half_t = __half;
 
class TypedArray
{
public:
    TypedArray() noexcept
        : mType(DataType::kFLOAT)
        , mData(nullptr)
        , mBufferSize(0U)
    {
    }
    // Shallow copy of TypedArray is allowed
    TypedArray(TypedArray const&) = default;
    TypedArray(TypedArray&&) = default;
    TypedArray& operator=(TypedArray const&) & = default;
    TypedArray& operator=(TypedArray&&) & = default;
    ~TypedArray() noexcept = default;
 
    TypedArray(float* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kFLOAT)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    TypedArray(half_t* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kHALF)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    TypedArray(int64_t* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kINT64)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    TypedArray(int32_t* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kINT32)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    TypedArray(int8_t* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kINT8)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    TypedArray(uint8_t* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kUINT8)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
    TypedArray(bool* ptr, uint64_t const bufferSize) noexcept
        : mType(DataType::kBOOL)
        , mData(ptr)
        , mBufferSize(bufferSize)
    {
    }
 
    // NOLINTBEGIN to avoid clang-tidy requiring [[nodiscard]]
    DataType getType() const noexcept
    {
        return mType;
    }
 
    float* getFloat() const noexcept
    {
        if (mType == DataType::kFLOAT)
        {
            return static_cast<float*>(mData);
        }
        return nullptr;
    }
 
    half_t* getHalf() const noexcept
    {
        if (mType == DataType::kHALF)
        {
            return static_cast<half_t*>(mData);
        }
        return nullptr;
    }
 
    int64_t* getInt64() const noexcept
    {
        if (mType == DataType::kINT64)
        {
            return static_cast<int64_t*>(mData);
        }
        return nullptr;
    }
 
    int32_t* getInt32() const noexcept
    {
        if (mType == DataType::kINT32)
        {
            return static_cast<int32_t*>(mData);
        }
        return nullptr;
    }
 
    int8_t* getInt8() const noexcept
    {
        if (mType == DataType::kINT8)
        {
            return static_cast<int8_t*>(mData);
        }
        return nullptr;
    }
 
    uint8_t* getUint8() const noexcept
    {
        if (mType == DataType::kUINT8)
        {
            return static_cast<uint8_t*>(mData);
        }
        return nullptr;
    }
 
    bool* getBool() const noexcept
    {
        if (mType == DataType::kBOOL)
        {
            return static_cast<bool*>(mData);
        }
        return nullptr;
    }
 
    uint64_t getSize() const noexcept
    {
        return mBufferSize;
    }
 
    void* getData() const noexcept
    {
        return mData;
    }
    // NOLINTEND
 
private:
    DataType mType;       // This is the current type of the data.
    void* mData;          // This is the pointer that holds the data.
    uint64_t mBufferSize; // This is the size of the buffer in bytes that holds the data.
};
 
class PhysicalDims
{
public:
    static constexpr int32_t MAX_DIMS{9};
 
    int32_t nbDims;
 
    int64_t d[MAX_DIMS];
};
 
inline bool operator==(PhysicalDims const& d0, PhysicalDims const& d1) noexcept
{
    return d0.nbDims == d1.nbDims && (d0.nbDims <= 0 || std::equal(d0.d, d0.d + d0.nbDims, d1.d));
}
 
inline bool operator!=(PhysicalDims const& d0, PhysicalDims const& d1) noexcept
{
    return !(d0 == d1);
}
 
struct TensorDescriptor
{
    AsciiChar const* tensorName{nullptr};
    PhysicalDims shape{-1, {}};
    PhysicalDims stride{-1, {}};
    DataType dataType{DataType::kFLOAT};
    uint64_t bytesPerComponent{0U};
    uint64_t componentsPerVector{1U};
    int64_t vectorizedDim{-1};
    uint64_t sizeInBytes{0U};
    TensorIOMode ioMode{TensorIOMode::kNONE};
    MemoryPlacement memPlacement{MemoryPlacement::kNONE};
    PhysicalDims strideOrder{-1, {}};
    Dims userShape{-1, {}};
};
 
class ITRTGraph
{
public:
    ITRTGraph(ITRTGraph const&) = delete;
    ITRTGraph(ITRTGraph&&) = delete;
    ITRTGraph& operator=(ITRTGraph const&) & = delete;
    ITRTGraph& operator=(ITRTGraph&&) & = delete;
 
    virtual ~ITRTGraph() noexcept = default;
 
    virtual ErrorCode clone(ITRTGraph*& graph, ISafeRecorder& recorder) noexcept = 0;
 
    virtual ErrorCode getScratchMemorySize(size_t& size) const noexcept = 0;
 
    virtual ErrorCode getTRTManagedScratch(bool& flag) const noexcept = 0;
 
    virtual ErrorCode setScratchMemory(void* memory) noexcept = 0;
 
    virtual ErrorCode getScratchMemory(void*& memory) noexcept = 0;
 
    virtual ErrorCode getNbIOTensors(int64_t& nb) const noexcept = 0;
 
    virtual ErrorCode getIOTensorName(AsciiChar const*& name, size_t const index) const noexcept = 0;
 
    virtual ErrorCode getIOTensorDescriptor(
        TensorDescriptor& desc, AsciiChar const* const tensorName) const noexcept = 0;
 
    virtual ErrorCode getIOTensorDescriptor(TensorDescriptor& desc, int32_t const index) const noexcept = 0;
 
    virtual ErrorCode setIOTensorAddress(AsciiChar const* const tensorName, TypedArray const& tensor) noexcept = 0;
 
    virtual ErrorCode setIOTensorAddress(int32_t const index, TypedArray const& tensor) noexcept = 0;
 
    virtual ErrorCode getIOTensorAddress(AsciiChar const* const tensorName, TypedArray& tensor) noexcept = 0;
 
    virtual ErrorCode getIOTensorAddress(int32_t const index, TypedArray& tensor) noexcept = 0;
 
    virtual ErrorCode setInputConsumedEvent(cudaEvent_t event) noexcept = 0;
 
    virtual ErrorCode getInputConsumedEvent(cudaEvent_t& event) const noexcept = 0;
 
    virtual ErrorCode getErrorBuffer(RuntimeErrorInformation*& buffer) const noexcept = 0;
 
    virtual ErrorCode getSafeRecorder(ISafeRecorder*& recorder) const noexcept = 0;
 
    virtual ErrorCode getNbIOProfiles(int64_t& nb) const noexcept = 0;
 
    virtual ErrorCode setIOProfile(int64_t profileIndex) noexcept = 0;
 
    virtual ErrorCode getIOProfile(int64_t& profileIndex) const noexcept = 0;
 
    virtual ErrorCode getNbAuxStreams(int32_t& nbStreams) const noexcept = 0;
 
    virtual ErrorCode setAuxStreams(cudaStream_t* auxStreams, int32_t nbStreams) noexcept = 0;
 
    virtual ErrorCode executeAsync(cudaStream_t stream) noexcept = 0;
 
    virtual ErrorCode sync() noexcept = 0;
 
protected:
    ITRTGraph() = default;
};
 
extern "C" ErrorCode createTRTGraph(ITRTGraph*& graph, void const* buffer, int64_t bufferSize, ISafeRecorder& recorder,
    bool trtManagedScratch = true, ISafeMemAllocator* allocator = nullptr) noexcept;
 
extern "C" ErrorCode destroyTRTGraph(ITRTGraph*& graph) noexcept;
 
} // namespace safe
} // namespace nvinfer2
#endif /* NV_INFER_SAFE_RUNTIME_H */