#include "ForwardDeclarations.h"

#include "utils.h"

#include <chrono>

#include <iomanip>

#include <pybind11/stl.h>

#include "infer/pyCoreDoc.h"

#include <cuda.h>

#include <cuda_runtime_api.h>

namespace tensorrt

{

using namespace nvinfer1;

namespace lambdas

{

static auto const opt_profile_set_shape

= [](IOptimizationProfile& self, std::string const& inputName, Dims const& min, Dims const& opt, Dims const& max) {

PY_ASSERT_RUNTIME_ERROR(self.setDimensions(inputName.c_str(), OptProfileSelector::kMIN, min),

"Shape provided for min is inconsistent with other shapes.");

PY_ASSERT_RUNTIME_ERROR(self.setDimensions(inputName.c_str(), OptProfileSelector::kOPT, opt),

"Shape provided for opt is inconsistent with other shapes.");

PY_ASSERT_RUNTIME_ERROR(self.setDimensions(inputName.c_str(), OptProfileSelector::kMAX, max),

"Shape provided for max is inconsistent with other shapes.");

};

static auto const opt_profile_get_shape

= [](IOptimizationProfile& self, std::string const& inputName) -> std::vector<Dims> {

std::vector<Dims> shapes{};

Dims minShape = self.getDimensions(inputName.c_str(), OptProfileSelector::kMIN);

if (minShape.nbDims != -1)

{

shapes.emplace_back(minShape);

shapes.emplace_back(self.getDimensions(inputName.c_str(), OptProfileSelector::kOPT));

shapes.emplace_back(self.getDimensions(inputName.c_str(), OptProfileSelector::kMAX));

}

return shapes;

};

static auto const opt_profile_set_shape_input = [](IOptimizationProfile& self, std::string const& inputName,

std::vector<int64_t> const& min, std::vector<int64_t> const& opt,

std::vector<int64_t> const& max) {

PY_ASSERT_RUNTIME_ERROR(self.setShapeValuesV2(inputName.c_str(), OptProfileSelector::kMIN, min.data(), min.size()),

"min input provided for shape tensor is inconsistent with other inputs.");

PY_ASSERT_RUNTIME_ERROR(self.setShapeValuesV2(inputName.c_str(), OptProfileSelector::kOPT, opt.data(), opt.size()),

"opt input provided for shape tensor is inconsistent with other inputs.");

PY_ASSERT_RUNTIME_ERROR(self.setShapeValuesV2(inputName.c_str(), OptProfileSelector::kMAX, max.data(), max.size()),

"max input provided for shape tensor is inconsistent with other inputs.");

};

static auto const opt_profile_get_shape_input

= [](IOptimizationProfile& self, std::string const& inputName) -> std::vector<std::vector<int64_t>> {

std::vector<std::vector<int64_t>> shapes{};

int32_t const shapeSize = self.getNbShapeValues(inputName.c_str());

int64_t const* shapePtr = self.getShapeValuesV2(inputName.c_str(), OptProfileSelector::kMIN);

// In the Python bindings, it is impossible to set only one shape in an optimization profile.

if (shapePtr && shapeSize >= 0)

{

shapes.emplace_back(shapePtr, shapePtr + shapeSize);

shapePtr = self.getShapeValuesV2(inputName.c_str(), OptProfileSelector::kOPT);

PY_ASSERT_RUNTIME_ERROR(shapePtr != nullptr, "Invalid shape for OPT.");

shapes.emplace_back(shapePtr, shapePtr + shapeSize);

shapePtr = self.getShapeValuesV2(inputName.c_str(), OptProfileSelector::kMAX);

PY_ASSERT_RUNTIME_ERROR(shapePtr != nullptr, "Invalid shape for MAX.");

shapes.emplace_back(shapePtr, shapePtr + shapeSize);

}

return shapes;

};

static auto const execute_v2 = [](IExecutionContext& self, std::vector<size_t>& bindings) {

return self.executeV2(reinterpret_cast<void**>(bindings.data()));

};

std::vector<char const*> infer_shapes(IExecutionContext& self)

}

bool execute_async_v3(IExecutionContext& self, size_t streamHandle)

{

return self.enqueueV3(reinterpret_cast<cudaStream_t>(streamHandle));

}

bool set_tensor_address(IExecutionContext& self, char const* tensor_name, size_t memory)

{

return self.setTensorAddress(tensor_name, reinterpret_cast<void*>(memory));

}

size_t get_tensor_address(IExecutionContext& self, char const* tensor_name)

{

return reinterpret_cast<size_t>(self.getTensorAddress(tensor_name));

}

bool set_input_consumed_event(IExecutionContext& self, size_t inputConsumed)

{

return self.setInputConsumedEvent(reinterpret_cast<cudaEvent_t>(inputConsumed));

}

size_t get_input_consumed_event(IExecutionContext& self)

{

return reinterpret_cast<size_t>(self.getInputConsumedEvent());

}

void set_aux_streams(IExecutionContext& self, std::vector<size_t> streamHandle)

{

self.setAuxStreams(reinterpret_cast<cudaStream_t*>(streamHandle.data()), static_cast<int32_t>(streamHandle.size()));

}

template <typename PyIterable>

Dims castDimsFromPyIterable(PyIterable& in)

}

template <typename PyIterable>

bool setInputShape(IExecutionContext& self, char const* tensorName, PyIterable& in)

{

return self.setInputShape(tensorName, castDimsFromPyIterable<PyIterable>(in));

}

static auto const runtime_deserialize_cuda_engine = [](IRuntime& self, py::buffer& serializedEngine) {

py::buffer_info info = serializedEngine.request();

py::gil_scoped_release releaseGil{};

return self.deserializeCudaEngine(info.ptr, info.size * info.itemsize);

};

static auto const reader_v2_read = [](IStreamReaderV2& self, void* destination, int64_t nbBytes, size_t stream) {

return self.read(destination, nbBytes, reinterpret_cast<cudaStream_t>(stream));

};

static auto const engine_getitem = [](ICudaEngine& self, int32_t pyIndex) {

// Support python's negative indexing

int32_t const index = (pyIndex < 0) ? static_cast<int32_t>(self.getNbIOTensors()) + pyIndex : pyIndex;

PY_ASSERT_INDEX_ERROR(index < self.getNbIOTensors());

return self.getIOTensorName(index);

};

std::vector<Dims> get_tensor_profile_shape(ICudaEngine& self, std::string const& tensorName, int32_t profileIndex)

}

std::vector<std::vector<int64_t>> get_tensor_profile_values(

}

void* allocate_async(

IGpuAllocator& self, uint64_t const size, uint64_t const alignment, AllocatorFlags const flags, size_t streamHandle)

{

return self.allocateAsync(size, alignment, flags, reinterpret_cast<cudaStream_t>(streamHandle));

}

bool deallocate_async(IGpuAllocator& self, void* const memory, size_t streamHandle)

{

return self.deallocateAsync(memory, reinterpret_cast<cudaStream_t>(streamHandle));

}

void* reallocate_output_async(IOutputAllocator& self, char const* tensorName, void* currentMemory, uint64_t size,

uint64_t alignment, size_t streamHandle)

{

return self.reallocateOutputAsync(

tensorName, currentMemory, size, alignment, reinterpret_cast<cudaStream_t>(streamHandle));

}

#if EXPORT_ALL_BINDINGS

static auto const netconfig_get_profile_stream

= [](IBuilderConfig& self) -> size_t { return reinterpret_cast<size_t>(self.getProfileStream()); };

static auto const netconfig_set_profile_stream = [](IBuilderConfig& self, size_t streamHandle) {

self.setProfileStream(reinterpret_cast<cudaStream_t>(streamHandle));

};

static auto const netconfig_create_timing_cache = [](IBuilderConfig& self, py::buffer& serializedTimingCache) {

py::buffer_info info = serializedTimingCache.request();

py::gil_scoped_release releaseGil{};

return self.createTimingCache(info.ptr, info.size * info.itemsize);

};

static auto const get_plugins_to_serialize = [](IBuilderConfig& self) {

std::vector<std::string> paths;

int64_t const nbPlugins = self.getNbPluginsToSerialize();

if (nbPlugins < 0)

{

utils::throwPyError(PyExc_RuntimeError, "Internal error");

}

paths.reserve(nbPlugins);

for (int64_t i = 0; i < nbPlugins; ++i)

{

paths.emplace_back(std::string{self.getPluginToSerialize(i)});

}

return paths;

};

static auto const set_plugins_to_serialize = [](IBuilderConfig& self, std::vector<std::string> const& paths) {

std::vector<char const*> cStrings;

cStrings.reserve(paths.size());

for (auto const& path : paths)

{

cStrings.push_back(path.c_str());

}

self.setPluginsToSerialize(reinterpret_cast<char const* const*>(cStrings.data()), cStrings.size());

};

static auto const get_remote_auto_tuning_config

= [](IBuilderConfig& self) { return std::string{self.getRemoteAutoTuningConfig()}; };

static auto const set_remote_auto_tuning_config

= [](IBuilderConfig& self, std::string const& config) { self.setRemoteAutoTuningConfig(config.c_str()); };

static auto const get_build_route = [](IBuilderConfig& self) { return std::string{self.getBuildRoute()}; };

static auto const set_build_route

= [](IBuilderConfig& self, std::string const& buildRoute) { self.setBuildRoute(buildRoute.c_str()); };

static auto const get_all_build_routes = [](IBuilderConfig& self) { return std::string{self.getAllBuildRoutes()}; };

#endif // EXPORT_ALL_BINDINGS

static auto const refitter_get_missing = [](IRefitter& self) {

// First get the number of missing weights.

int32_t const size{self.getMissing(0, nullptr, nullptr)};

// Now that we know how many weights are missing, we can create the buffers appropriately.

std::vector<const char*> layerNames(size);

std::vector<WeightsRole> roles(size);

self.getMissing(size, layerNames.data(), roles.data());

return std::pair<std::vector<const char*>, std::vector<WeightsRole>>{layerNames, roles};

};

static auto const refitter_get_missing_weights = [](IRefitter& self) {

// First get the number of missing weights.

int32_t const size{self.getMissingWeights(0, nullptr)};

// Now that we know how many weights are missing, we can create the buffers appropriately.

std::vector<char const*> names(size);

self.getMissingWeights(size, names.data());

return names;

};

static auto const refitter_get_all = [](IRefitter& self) {

int32_t const size{self.getAll(0, nullptr, nullptr)};

std::vector<char const*> layerNames(size);

std::vector<WeightsRole> roles(size);

self.getAll(size, layerNames.data(), roles.data());

return std::pair<std::vector<const char*>, std::vector<WeightsRole>>{layerNames, roles};

};

static auto const refitter_get_all_weights = [](IRefitter& self) {

int32_t const size{self.getAllWeights(0, nullptr)};

std::vector<char const*> names(size);

self.getAllWeights(size, names.data());

return names;

};

static auto const refitter_refit_cuda_engine_async = [](IRefitter& self, size_t streamHandle) {

return self.refitCudaEngineAsync(reinterpret_cast<cudaStream_t>(streamHandle));

};

static auto const context_set_optimization_profile_async

= [](IExecutionContext& self, int32_t const profileIndex, size_t streamHandle) {

PY_ASSERT_RUNTIME_ERROR(

self.setOptimizationProfileAsync(profileIndex, reinterpret_cast<cudaStream_t>(streamHandle)),

"Error in set optimization profile async.");

return true;

};

void context_set_device_memory(IExecutionContext& self, size_t memory)

{

self.setDeviceMemory(reinterpret_cast<void*>(memory));

}

void context_set_device_memory_v2(IExecutionContext& self, size_t memory, int64_t size)

{

self.setDeviceMemoryV2(reinterpret_cast<void*>(memory), size);

}

void serialization_config_set_flags(ISerializationConfig& self, uint32_t flags)

}

void docProcessDebugTensor(IDebugListener& self, void const* addr, TensorLocation location, DataType type,

}

uint64_t getTacticHash(TimingCacheValue const& value)

{

return value.tacticHash;

}

void setTacticHash(TimingCacheValue& value, uint64_t tacticHash)

{

value.tacticHash = tacticHash;

}

float getTimingMSec(TimingCacheValue const& value)

{

return value.timingMSec;

}

void setTimingMSec(TimingCacheValue& value, float timingMSec)

{

value.timingMSec = timingMSec;

}

namespace detail

{

constexpr int64_t kBYTES_PER_KEY = 16;

constexpr int64_t kCHARS_PER_BYTE = 2;

constexpr int64_t kPREFIX_CHARS = 2;

constexpr int64_t kTOTAL_CHARS = kPREFIX_CHARS + kBYTES_PER_KEY * kCHARS_PER_BYTE;

} // namespace detail

TimingCacheKey parseTimingCacheKey(std::string const& text)

}

std::string convertTimingCacheKeyToString(TimingCacheKey const& key)

}

std::vector<TimingCacheKey> queryTimingCacheKeys(ITimingCache const& cache)

}

} // namespace lambdas

namespace PyGpuAllocatorHelper

{

template <typename TAllocator, typename... Args>

void* allocHelper(TAllocator* allocator, char const* pyFuncName, bool showWarning, Args&&... args) noexcept

{

try

{

py::gil_scoped_acquire gil{};

py::function pyAllocFunc = utils::getOverride(static_cast<TAllocator*>(allocator), pyFuncName, showWarning);

if (!pyAllocFunc)

{

return nullptr;

}

py::object ptr = pyAllocFunc(std::forward<Args>(args)...);

try

{

return reinterpret_cast<void*>(ptr.cast<size_t>());

}

catch (py::cast_error const& e)

{

std::cerr << "[ERROR] Return value of allocate() could not be interpreted as an int" << std::endl;

}

}

catch (std::exception const& e)

{

std::cerr << "[ERROR] Exception caught in allocate(): " << e.what() << std::endl;

}

catch (...)

{

std::cerr << "[ERROR] Exception caught in allocate()" << std::endl;

return nullptr;

}

return nullptr;

}

} // namespace PyGpuAllocatorHelper

class PyGpuAllocator : public IGpuAllocator

}; // PyGpuAllocator

class PyGpuAsyncAllocator : public IGpuAsyncAllocator

}; // PyGpuAsyncAllocator

class PyOutputAllocator : public IOutputAllocator

};

class PyStreamReaderV2 : public IStreamReaderV2

};

class PyStreamWriter : public IStreamWriter

};

class PyDebugListener : public IDebugListener

};

void bindCore(py::module& m)

{

class PyLogger : public ILogger

};

py::class_<ILogger, PyLogger> baseLoggerBinding{m, "ILogger", ILoggerDoc::descr, py::module_local()};

py::enum_<ILogger::Severity>(

baseLoggerBinding, "Severity", py::arithmetic(), SeverityDoc::descr, py::module_local())

.value("INTERNAL_ERROR", ILogger::Severity::kINTERNAL_ERROR, SeverityDoc::internal_error)

.value("ERROR", ILogger::Severity::kERROR, SeverityDoc::error)

.value("WARNING", ILogger::Severity::kWARNING, SeverityDoc::warning)

.value("INFO", ILogger::Severity::kINFO, SeverityDoc::info)

.value("VERBOSE", ILogger::Severity::kVERBOSE, SeverityDoc::verbose)

// We export into the outer scope, so we can access with trt.ILogger.X.

.export_values();

baseLoggerBinding.def(py::init<>()).def("log", &ILogger::log, "severity"_a, "msg"_a, ILoggerDoc::log);

class DefaultLogger : public ILogger

};