trt_session.cpp

/************************************************
 * Copyright MaybeShewill-CV. All Rights Reserved.
 * Author: MaybeShewill-CV
 * File: trt_session.cpp
 * Date: 26-8-20
 ************************************************/

#include "models/backend/trt_session.h"

#include <algorithm>
#include <cstring>
#include <fstream>
#include <sstream>

#include "NvInferPlugin.h"
#include "NvInferVersion.h"
#include "glog/logging.h"

#include "common/file_path_util.h"
#include "common/stage_timing.h"
#include "models/backend/session_io.h"

static_assert(NV_TENSORRT_MAJOR >= 10, "Mortred GPU line requires TensorRT 10 headers");

namespace mortred {
namespace models {
namespace backend {

using mortred::common::StatusCode;

namespace {

nvinfer1::Dims to_trt_dims(const std::vector<int64_t>& shape) {
    nvinfer1::Dims dims{};
    dims.nbDims = static_cast<int32_t>(shape.size());
    for (size_t idx = 0; idx < shape.size(); ++idx) {
        dims.d[idx] = static_cast<int64_t>(shape[idx]);
    }
    return dims;
}

std::vector<int64_t> from_trt_dims(const nvinfer1::Dims& dims) {
    std::vector<int64_t> shape;
    shape.reserve(static_cast<size_t>(dims.nbDims));
    for (int32_t idx = 0; idx < dims.nbDims; ++idx) {
        shape.push_back(dims.d[idx]);
    }
    return shape;
}

bool dtype_from_trt(const nvinfer1::DataType& dtype, DType* out, std::string* err) {
    switch (dtype) {
        case nvinfer1::DataType::kFLOAT:
            *out = DType::F32;
            return true;
        case nvinfer1::DataType::kHALF:
            *out = DType::F16;
            return true;
        case nvinfer1::DataType::kINT32:
            *out = DType::I32;
            return true;
        case nvinfer1::DataType::kUINT8:
            *out = DType::U8;
            return true;
#if NV_TENSORRT_MAJOR >= 10
        case nvinfer1::DataType::kINT64:
            *out = DType::I64;
            return true;
#endif
        default:
            if (err != nullptr) {
                *err = "unsupported tensorrt io dtype: " + std::to_string(static_cast<int>(dtype));
            }
            return false;
    }
}

}  // namespace

namespace trt_detail {

void SessionLogger::log(Severity severity, const char* msg) noexcept {
    if (msg == nullptr) {
        return;
    }
    switch (severity) {
        case Severity::kINTERNAL_ERROR:
        case Severity::kERROR:
            LOG(ERROR) << "[tensorrt] " << msg;
            break;
        case Severity::kWARNING:
            LOG(WARNING) << "[tensorrt] " << msg;
            break;
        case Severity::kINFO:
            LOG(INFO) << "[tensorrt] " << msg;
            break;
        default:
            break;
    }
}

}  // namespace trt_detail

/***
 * TensorRT calls this allocator when an output shape depends on runtime data
 * (for example the lightglue extractor's NMS output).  It owns the device
 * allocation and records the shape communicated through notifyShape.
 */
class TrtSession::DynamicOutputAllocator final : public nvinfer1::IOutputAllocator {
  public:
    DynamicOutputAllocator() = default;
    ~DynamicOutputAllocator() override {
        release();
    }

    void* reallocateOutputAsync(char const* tensor_name, void* current_memory, uint64_t size,
                                uint64_t alignment, cudaStream_t stream) noexcept override {
        (void)tensor_name;
        (void)current_memory;
        (void)alignment;
        (void)stream;
        allocation_failed = false;
        if (size == 0) {
            allocation_failed = true;
            return nullptr;
        }
        if (memory != nullptr && allocated_bytes >= size) {
            return memory;
        }
        release();
        const auto status = cudaMalloc(&memory, size);
        if (status != cudaSuccess || memory == nullptr) {
            LOG(ERROR) << "allocate dynamic tensorrt output failed: "
                       << cudaGetErrorString(status);
            memory = nullptr;
            allocation_failed = true;
            return nullptr;
        }
        allocated_bytes = size;
        return memory;
    }

    void notifyShape(char const* tensor_name, nvinfer1::Dims const& dims) noexcept override {
        (void)tensor_name;
        shape = dims;
        shape_notified = true;
    }

    void reset_for_run() {
        shape = nvinfer1::Dims{};
        shape_notified = false;
        allocation_failed = false;
    }

    bool failed() const {
        return allocation_failed;
    }

    bool has_shape() const {
        return shape_notified;
    }

    std::vector<int64_t> reported_shape() const {
        return from_trt_dims(shape);
    }

    size_t capacity() const {
        return allocated_bytes;
    }

    void* data() const {
        return memory;
    }

  private:
    void release() {
        if (memory != nullptr) {
            const auto status = cudaFree(memory);
            if (status != cudaSuccess) {
                LOG(ERROR) << "free dynamic tensorrt output failed: "
                           << cudaGetErrorString(status);
            }
            memory = nullptr;
        }
        allocated_bytes = 0;
    }

    void* memory = nullptr;
    size_t allocated_bytes = 0;
    nvinfer1::Dims shape{};
    bool shape_notified = false;
    bool allocation_failed = false;
};

TrtSession::DeviceBuffer& TrtSession::DeviceBuffer::operator=(DeviceBuffer&& other) noexcept {
    if (this != &other) {
        if (memory != nullptr) {
            cudaFree(memory);
        }
        memory = other.memory;
        bytes = other.bytes;
        other.memory = nullptr;
        other.bytes = 0;
    }
    return *this;
}

TrtSession::DeviceBuffer::~DeviceBuffer() {
    if (memory != nullptr) {
        const auto status = cudaFree(memory);
        if (status != cudaSuccess) {
            LOG(ERROR) << "free tensorrt device buffer failed: " << cudaGetErrorString(status);
        }
        memory = nullptr;
    }
}

StatusCode TrtSession::DeviceBuffer::ensure(size_t size_bytes) {
    if (size_bytes == 0) {
        return StatusCode::TRT_ALLOC_MEMO_FAILED;
    }
    if (memory != nullptr && bytes >= size_bytes) {
        return StatusCode::OK;
    }
    if (memory != nullptr) {
        const auto status = cudaFree(memory);
        if (status != cudaSuccess) {
            LOG(ERROR) << "free tensorrt device buffer failed: " << cudaGetErrorString(status);
            memory = nullptr;
            return StatusCode::TRT_ALLOC_MEMO_FAILED;
        }
        memory = nullptr;
        bytes = 0;
    }
    void* device_memory = nullptr;
    const auto status = cudaMalloc(&device_memory, size_bytes);
    if (status != cudaSuccess || device_memory == nullptr) {
        LOG(ERROR) << "allocate tensorrt device buffer failed: " << cudaGetErrorString(status);
        return StatusCode::TRT_ALLOC_MEMO_FAILED;
    }
    memory = device_memory;
    bytes = size_bytes;
    return StatusCode::OK;
}

TrtSession::~TrtSession() {
    // cudaFree/cudaStreamDestroy are context sensitive and destruction may run on
    // a thread whose current device differs from the one init selected.
    const auto set_dev = cudaSetDevice(_m_device_id);
    if (set_dev != cudaSuccess) {
        LOG(WARNING) << "tensorrt release cudaSetDevice(" << _m_device_id
                     << ") failed: " << cudaGetErrorString(set_dev);
    }
    release_pinned_staging();
    _m_device_buffers.clear();
    if (_m_stream != nullptr) {
        cudaStreamDestroy(_m_stream);
        _m_stream = nullptr;
    }
    // TRT 10 removed destroy(); public destructors are the supported path.
    delete _m_context;
    _m_context = nullptr;
    delete _m_engine;
    _m_engine = nullptr;
    delete _m_runtime;
    _m_runtime = nullptr;
    for (auto& item : _m_output_allocators) {
        delete item.second;
    }
    _m_output_allocators.clear();
}



void* TrtSession::pinned_h2d_stage(size_t bytes) {
    return _m_pinned_h2d_pool.acquire(bytes);
}

void TrtSession::pinned_h2d_release(void* p, size_t bytes) {
    _m_pinned_h2d_pool.release(p, bytes);
}

void* TrtSession::pinned_d2h_stage(size_t bytes) {
    return _m_pinned_d2h_pool.acquire(bytes);
}

void TrtSession::pinned_d2h_release(void* p, size_t bytes) {
    _m_pinned_d2h_pool.release(p, bytes);
}

void TrtSession::release_pinned_staging() {
    _m_pinned_h2d_pool.drain();
    _m_pinned_d2h_pool.drain();
}

StatusCode TrtSession::init(const BackendConfig& config, std::string* err) {
    if (err != nullptr) {
        err->clear();
    }
    if (!mortred::common::FilePathUtil::is_file_exist(config.model_file_path)) {
        if (err != nullptr) {
            *err = "tensorrt engine file not exist: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    std::ifstream engine_file(config.model_file_path, std::ios_base::in | std::ios_base::binary);
    if (!engine_file) {
        if (err != nullptr) {
            *err = "read tensorrt engine failed: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }
    std::stringstream buffer;
    buffer << engine_file.rdbuf();
    const std::string engine_stream = buffer.str();
    if (engine_stream.empty()) {
        if (err != nullptr) {
            *err = "tensorrt engine file is empty: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    const auto set_dev = cudaSetDevice(config.device_id);
    if (set_dev != cudaSuccess) {
        if (err != nullptr) {
            *err = std::string("cudaSetDevice failed: ") + cudaGetErrorString(set_dev);
        }
        return StatusCode::MODEL_INIT_FAILED;
    }
    _m_device_id = config.device_id;
    const auto cuda_init = cudaFree(nullptr);
    if (cuda_init != cudaSuccess) {
        if (err != nullptr) {
            *err = std::string("CUDA runtime init failed: ") + cudaGetErrorString(cuda_init);
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    if (!initLibNvInferPlugins(nullptr, "")) {
        if (err != nullptr) {
            *err = "init tensorrt plugin registry failed";
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    _m_runtime = nvinfer1::createInferRuntime(_m_logger);
    if (_m_runtime == nullptr) {
        if (err != nullptr) {
            *err = "create tensorrt runtime failed";
        }
        return StatusCode::MODEL_INIT_FAILED;
    }
    _m_engine = _m_runtime->deserializeCudaEngine(engine_stream.data(), engine_stream.size());
    if (_m_engine == nullptr) {
        if (err != nullptr) {
            *err = "deserialize tensorrt engine failed: " + config.model_file_path
                   + " (engine was built by an incompatible tensorrt version)";
        }
        delete _m_runtime;
        _m_runtime = nullptr;
        return StatusCode::MODEL_INIT_FAILED;
    }
    _m_context = _m_engine->createExecutionContext();
    if (_m_context == nullptr) {
        if (err != nullptr) {
            *err = "create tensorrt execution context failed";
        }
        return StatusCode::MODEL_INIT_FAILED;
    }

    const auto is_input = [this](const std::string& name) {
        return _m_engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT;
    };
    const auto wanted_input = [&is_input](const std::string& name) { return is_input(name); };
    const auto wanted_output = [this](const std::string& name) {
        return _m_engine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kOUTPUT;
    };

    const int32_t io_count = _m_engine->getNbIOTensors();
    for (int32_t idx = 0; idx < io_count; ++idx) {
        const char* raw_name = _m_engine->getIOTensorName(idx);
        if (raw_name == nullptr) {
            continue;
        }
        const std::string name(raw_name);
        const bool as_input = wanted_input(name);
        const bool as_output = wanted_output(name);
        if (!as_input && !as_output) {
            continue;
        }
        TensorInfo info;
        info.name = name;
        std::string dtype_err;
        if (!dtype_from_trt(_m_engine->getTensorDataType(name.c_str()), &info.dtype, &dtype_err)) {
            if (err != nullptr) {
                *err = "tensorrt io '" + name + "': " + dtype_err;
            }
            return StatusCode::MODEL_INIT_FAILED;
        }
        info.shape = from_trt_dims(_m_engine->getTensorShape(name.c_str()));
        info.dynamic = shape_is_dynamic(info.shape);
        info.layout = host_output_layout(info.shape);
        _m_device_buffers.emplace(name, DeviceBuffer{});
        if (as_input) {
            _m_input_infos.push_back(std::move(info));
        } else {
            _m_output_infos.push_back(std::move(info));
        }
    }
    if (_m_input_infos.empty() || _m_output_infos.empty()) {
        if (err != nullptr) {
            *err = "tensorrt engine exposes no io tensors: " + config.model_file_path;
        }
        return StatusCode::MODEL_INIT_FAILED;
    }
    if (apply_configured_io_names(config.input_names, &_m_input_infos, "input", err) !=
        StatusCode::OK) {
        return StatusCode::MODEL_INIT_FAILED;
    }
    if (apply_configured_io_names(config.output_names, &_m_output_infos, "output", err) !=
        StatusCode::OK) {
        return StatusCode::MODEL_INIT_FAILED;
    }

    // Outputs whose shape cannot be inferred from the input shape require an
    // IOutputAllocator. Without one, getTensorShape() remains dynamic and there
    // is no concrete output address to bind before enqueueV3().
    for (const auto& info : _m_output_infos) {
        if (!info.dynamic) {
            continue;
        }
        auto* allocator = new DynamicOutputAllocator();
        _m_output_allocators.emplace(info.name, allocator);
        const bool allocator_ready =
            _m_context->setOutputAllocator(info.name.c_str(), allocator) &&
            _m_context->setTensorAddress(info.name.c_str(), nullptr);
        if (!allocator_ready) {
            if (err != nullptr) {
                *err = "configure tensorrt output allocator failed: " + info.name;
            }
            return StatusCode::MODEL_INIT_FAILED;
        }
    }

    const auto cuda_status = cudaStreamCreate(&_m_stream);
    if (cuda_status != cudaSuccess) {
        if (err != nullptr) {
            *err = std::string("create cuda stream failed: ") + cudaGetErrorString(cuda_status);
        }
        return StatusCode::TRT_CUDA_ERROR;
    }

    _m_model_file_path = config.model_file_path;
    for (const auto& info : _m_input_infos) {
        // batch-capability self check: a dynamic input is only usable up to
        // the engine's optimization profile - surface min/opt/max so an
        // operator can immediately see whether a rebuilt engine really
        // carries the batch profile (and which N range it accepts)
        if (info.dynamic) {
            // dynamic dims require an optimization profile at build time, so
            // profile 0 always exists for a dynamic input
            const auto shape_min = from_trt_dims(_m_engine->getProfileShape(
                info.name.c_str(), 0, nvinfer1::OptProfileSelector::kMIN));
            const auto shape_opt = from_trt_dims(_m_engine->getProfileShape(
                info.name.c_str(), 0, nvinfer1::OptProfileSelector::kOPT));
            const auto shape_max = from_trt_dims(_m_engine->getProfileShape(
                info.name.c_str(), 0, nvinfer1::OptProfileSelector::kMAX));
            LOG(INFO) << "tensorrt session input: " << info.to_string()
                      << " profile min" << shape_to_string(shape_min)
                      << " opt" << shape_to_string(shape_opt)
                      << " max" << shape_to_string(shape_max);
        } else {
            LOG(INFO) << "tensorrt session input: " << info.to_string()
                      << " (static; batched shapes will be rejected)";
        }
    }
    for (const auto& info : _m_output_infos) {
        LOG(INFO) << "tensorrt session output: " << info.to_string();
    }
    return StatusCode::OK;
}

StatusCode TrtSession::run(const std::vector<NamedTensor>& inputs,
                           std::vector<NamedTensor>& outputs) {
    if (_m_context == nullptr || _m_engine == nullptr) {
        LOG(ERROR) << "tensorrt session is not initialized";
        return StatusCode::MODEL_INIT_FAILED;
    }
    // run() executes on whichever go/worker thread picked the session, and the
    // current CUDA device is per thread (default 0). Every allocation, copy and
    // stream call below is context sensitive, so re-select the configured device
    // instead of inheriting the caller's.
    const auto set_dev = cudaSetDevice(_m_device_id);
    if (set_dev != cudaSuccess) {
        LOG(ERROR) << "tensorrt run cudaSetDevice(" << _m_device_id
                   << ") failed: " << cudaGetErrorString(set_dev);
        return StatusCode::TRT_CUDA_ERROR;
    }
    std::vector<const NamedTensor*> ordered_inputs;
    std::string match_err;
    if (match_required_inputs(_m_input_infos, inputs, &ordered_inputs, &match_err) !=
        StatusCode::OK) {
        LOG(ERROR) << "tensorrt " << match_err;
        return StatusCode::MODEL_RUN_SESSION_FAILED;
    }

    // apply input shapes in session order (required set already matched)
    for (size_t idx = 0; idx < _m_input_infos.size(); ++idx) {
        const auto& info = _m_input_infos[idx];
        const auto& named = *ordered_inputs[idx];
        if (info.dtype != named.tensor.dtype) {
            LOG(ERROR) << "tensorrt input '" << named.name << "' dtype mismatch, expected "
                       << info.to_string() << ", got " << dtype_to_string(named.tensor.dtype);
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        if (!named.tensor.shape_is_concrete()) {
            LOG(ERROR) << "tensorrt input '" << named.name << "' has non concrete shape "
                       << shape_to_string(named.tensor.shape);
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        if (!_m_context->setInputShape(named.name.c_str(),
                                        to_trt_dims(named.tensor.shape))) {
            LOG(ERROR) << "tensorrt input '" << named.name << "' set shape "
                       << shape_to_string(named.tensor.shape)
                       << " rejected (outside the optimization profile?)";
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
    }

    // H2D + bind inputs (same session order). Pinned staging is sliced per
    // input (disjoint ranges, so all asyncs may be in flight together).
    size_t pinned_h2d_total = 0;
    for (size_t idx = 0; idx < ordered_inputs.size(); ++idx) {
        const auto& named = *ordered_inputs[idx];
        size_t input_bytes = 0;
        if (!checked_shape_nbytes(named.tensor.shape, named.tensor.dtype, &input_bytes)) {
            LOG(ERROR) << "tensorrt input '" << named.name
                       << "' shape byte size overflow or non-concrete: "
                       << shape_to_string(named.tensor.shape);
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        pinned_h2d_total += input_bytes;
    }
    // slot is owned for the whole request: the async H2D copies out of it
    // only provably executed once the request's stream sync completes
    char* pinned_h2d = pinned_h2d_total > 0 ? static_cast<char*>(pinned_h2d_stage(pinned_h2d_total)) : nullptr;
    // Error paths may return while cudaMemcpyAsync is still pending on
    // _m_stream. Releasing the slot into the pool (or cudaFreeHost) before
    // the stream finishes would UAF. On abandon we sync-then-release.
    struct StagingSlotGuard {
        TrtSession* session;
        void* slot;
        size_t bytes;
        cudaStream_t stream = nullptr;
        bool abandon = false;  // error path: queued copies may still be pending
        ~StagingSlotGuard() {
            if (session == nullptr || slot == nullptr) return;
            if (abandon && stream != nullptr) {
                cudaStreamSynchronize(stream);
            }
            session->pinned_h2d_release(slot, bytes);
        }
    } h2d_slot_guard{this, pinned_h2d, pinned_h2d_total, _m_stream, false};
    size_t pinned_h2d_offset = 0;
    for (size_t idx = 0; idx < ordered_inputs.size(); ++idx) {
        const auto& named = *ordered_inputs[idx];
        auto& buffer = _m_device_buffers.at(named.name);
        size_t input_bytes = 0;
        checked_shape_nbytes(named.tensor.shape, named.tensor.dtype, &input_bytes);
        // GPU zero-copy: input already in GPU memory, bind directly
        if (named.tensor.device_data != nullptr) {
            // the producer wrote this buffer on ANOTHER stream (the jpeggpu
            // decode stream): order this stream behind the recorded event
            // before inference reads it
            if (named.tensor.device_ready_event != nullptr) {
                cudaStreamWaitEvent(_m_stream, (cudaEvent_t)named.tensor.device_ready_event, 0);
            }
            if (!_m_context->setTensorAddress(named.name.c_str(), named.tensor.device_data)) {
                LOG(ERROR) << "tensorrt set device input address failed: " << named.name;
                h2d_slot_guard.abandon = true;
                return StatusCode::MODEL_RUN_SESSION_FAILED;
            }
            continue;
        }
        const auto status = buffer.ensure(input_bytes);
        if (status != StatusCode::OK) {
            // earlier inputs' H2D copies may already be queued on _m_stream
            h2d_slot_guard.abandon = true;
            return status;
        }
        const void* host_src = named.tensor.buffer.data();
        if (pinned_h2d != nullptr && named.tensor.byte_size() == input_bytes) {
            std::memcpy(pinned_h2d + pinned_h2d_offset, host_src, input_bytes);
            host_src = pinned_h2d + pinned_h2d_offset;
            pinned_h2d_offset += input_bytes;
        }
        const auto cuda_status = cudaMemcpyAsync(
            buffer.memory, host_src, named.tensor.byte_size(),
            cudaMemcpyHostToDevice, _m_stream);
        if (cuda_status != cudaSuccess) {
            LOG(ERROR) << "tensorrt H2D copy failed for '" << named.name
                       << "': " << cudaGetErrorString(cuda_status);
            h2d_slot_guard.abandon = true;
            return StatusCode::TRT_CUDA_ERROR;
        }
        if (!_m_context->setTensorAddress(named.name.c_str(), buffer.memory)) {
            LOG(ERROR) << "tensorrt set input tensor address failed: " << named.name;
            h2d_slot_guard.abandon = true;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
    }
    // note: cudaMemcpyAsync only enqueues; the actual transfer surfaces in
    // the first cudaStreamSynchronize below, so h2d measures launch cost and
    // exec absorbs the real H2D transfer
    mortred::common::stage_timing::mark("h2d");

    // Resolve and bind outputs whose shape is inferable from the inputs.
    // Runtime-data-dependent outputs are left to their IOutputAllocator.
    struct ResolvedOutput {
        const TensorInfo* info;
        std::vector<int64_t> shape;
        void* memory;
    };
    std::vector<ResolvedOutput> resolved_outputs;
    resolved_outputs.reserve(_m_output_infos.size());
    for (const auto& info : _m_output_infos) {
        resolved_outputs.push_back({&info, {}, nullptr});
    }
    for (auto& item : _m_output_allocators) {
        item.second->reset_for_run();
    }
    for (size_t idx = 0; idx < _m_output_infos.size(); ++idx) {
        const auto& info = _m_output_infos[idx];
        const auto allocator_iter = _m_output_allocators.find(info.name);
        if (allocator_iter != _m_output_allocators.end()) {
            continue;
        }
        auto shape = from_trt_dims(_m_context->getTensorShape(info.name.c_str()));
        if (shape_is_dynamic(shape)) {
            LOG(ERROR) << "tensorrt output '" << info.name << "' shape unresolved after input "
                       << "shapes: " << shape_to_string(shape);
            h2d_slot_guard.abandon = true;
            return StatusCode::TRT_ALLOC_DYNAMIC_SHAPE_MEMO;
        }
        auto& buffer = _m_device_buffers.at(info.name);
        size_t bytes = 0;
        if (!checked_shape_nbytes(shape, info.dtype, &bytes)) {
            LOG(ERROR) << "tensorrt output '" << info.name
                       << "' shape byte size overflow or non-concrete: "
                       << shape_to_string(shape);
            h2d_slot_guard.abandon = true;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        const auto status = buffer.ensure(bytes);
        if (status != StatusCode::OK) {
            // all input H2D copies are queued by now; sync before the slot
            // goes back to the pool
            h2d_slot_guard.abandon = true;
            return status;
        }
        if (!_m_context->setTensorAddress(info.name.c_str(), buffer.memory)) {
            LOG(ERROR) << "tensorrt set output tensor address failed: " << info.name;
            h2d_slot_guard.abandon = true;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        resolved_outputs[idx] = {&info, std::move(shape), buffer.memory};
    }

    if (!_m_context->enqueueV3(_m_stream)) {
        LOG(ERROR) << "tensorrt enqueueV3 inference failed";
        h2d_slot_guard.abandon = true;
        return StatusCode::MODEL_RUN_SESSION_FAILED;
    }
    const auto sync_status = cudaStreamSynchronize(_m_stream);
    if (sync_status != cudaSuccess) {
        LOG(ERROR) << "tensorrt stream sync failed: " << cudaGetErrorString(sync_status);
        h2d_slot_guard.abandon = true;
        return StatusCode::TRT_CUDA_ERROR;
    }
    mortred::common::stage_timing::mark("exec");

    for (size_t idx = 0; idx < _m_output_infos.size(); ++idx) {
        const auto& info = _m_output_infos[idx];
        const auto allocator_iter = _m_output_allocators.find(info.name);
        if (allocator_iter == _m_output_allocators.end()) {
            continue;
        }
        const auto* allocator = allocator_iter->second;
        if (allocator->failed() || allocator->data() == nullptr) {
            LOG(ERROR) << "tensorrt dynamic output '" << info.name << "' allocation failed";
            return StatusCode::TRT_ALLOC_DYNAMIC_SHAPE_MEMO;
        }
        auto shape = allocator->reported_shape();
        if (!allocator->has_shape() || shape_is_dynamic(shape)) {
            shape = from_trt_dims(_m_context->getTensorShape(info.name.c_str()));
        }
        if (shape_is_dynamic(shape)) {
            LOG(ERROR) << "tensorrt dynamic output '" << info.name
                       << "' shape unresolved after inference: " << shape_to_string(shape);
            return StatusCode::TRT_ALLOC_DYNAMIC_SHAPE_MEMO;
        }
        size_t bytes = 0;
        if (!checked_shape_nbytes(shape, info.dtype, &bytes) || bytes == 0) {
            LOG(ERROR) << "tensorrt dynamic output '" << info.name
                       << "' shape byte size overflow or non-concrete: "
                       << shape_to_string(shape);
            return StatusCode::TRT_ALLOC_DYNAMIC_SHAPE_MEMO;
        }
        if (allocator->capacity() < bytes) {
            LOG(ERROR) << "tensorrt dynamic output '" << info.name << "' allocation "
                       << allocator->capacity() << " bytes is smaller than output "
                       << bytes << " bytes";
            return StatusCode::TRT_ALLOC_DYNAMIC_SHAPE_MEMO;
        }
        resolved_outputs[idx] = {&info, std::move(shape), allocator->data()};
    }

    outputs.clear();
    outputs.reserve(resolved_outputs.size());
    // pinned D2H staging is only used when every output is static (shape
    // resolved before the copies are issued, so the total size is known and
    // one buffer can be sliced); dynamic outputs keep the pageable path
    size_t pinned_d2h_total = 0;
    if (_m_output_allocators.empty()) {
        for (const auto& item : resolved_outputs) {
            size_t bytes = 0;
            if (item.memory != nullptr && checked_shape_nbytes(item.shape, item.info->dtype, &bytes)) {
                pinned_d2h_total += bytes;
            } else {
                pinned_d2h_total = 0;
                break;
            }
        }
    }
    char* pinned_d2h = pinned_d2h_total > 0 ? static_cast<char*>(pinned_d2h_stage(pinned_d2h_total)) : nullptr;
    struct D2HSlotGuard {
        TrtSession* session;
        void* slot;
        size_t bytes;
        cudaStream_t stream = nullptr;
        bool abandon = false;
        ~D2HSlotGuard() {
            if (session == nullptr || slot == nullptr) return;
            if (abandon && stream != nullptr) {
                cudaStreamSynchronize(stream);
            }
            session->pinned_d2h_release(slot, bytes);
        }
    } d2h_slot_guard{this, pinned_d2h, pinned_d2h_total, _m_stream, false};
    size_t pinned_d2h_offset = 0;
    std::vector<std::pair<void*, size_t>> staged_d2h_copies;
    for (const auto& item : resolved_outputs) {
        if (item.memory == nullptr) {
            LOG(ERROR) << "tensorrt output '" << item.info->name
                       << "' was not bound to device memory";
            d2h_slot_guard.abandon = true;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        const auto& info = *item.info;
        const auto& shape = item.shape;
        NamedTensor named;
        named.name = info.name;
        named.tensor.dtype = info.dtype;
        named.tensor.shape = shape;
        named.tensor.layout = host_output_layout(shape);
        size_t bytes = 0;
        if (!checked_shape_nbytes(shape, info.dtype, &bytes)) {
            LOG(ERROR) << "tensorrt output '" << info.name
                       << "' shape byte size overflow or non-concrete: "
                       << shape_to_string(shape);
            d2h_slot_guard.abandon = true;
            return StatusCode::MODEL_RUN_SESSION_FAILED;
        }
        named.tensor.buffer.resize(bytes);
        void* host_dst = named.tensor.buffer.data();
        if (pinned_d2h != nullptr) {
            host_dst = pinned_d2h + pinned_d2h_offset;
            staged_d2h_copies.emplace_back(named.tensor.buffer.data(), bytes);
            pinned_d2h_offset += bytes;
        }
        const auto cuda_status = cudaMemcpyAsync(
            host_dst, item.memory, bytes, cudaMemcpyDeviceToHost, _m_stream);
        if (cuda_status != cudaSuccess) {
            LOG(ERROR) << "tensorrt D2H copy failed for '" << info.name
                       << "': " << cudaGetErrorString(cuda_status);
            d2h_slot_guard.abandon = true;
            return StatusCode::TRT_CUDA_ERROR;
        }
        outputs.push_back(std::move(named));
    }
    const auto final_sync = cudaStreamSynchronize(_m_stream);
    if (final_sync != cudaSuccess) {
        LOG(ERROR) << "tensorrt stream sync failed: " << cudaGetErrorString(final_sync);
        d2h_slot_guard.abandon = true;
        return StatusCode::TRT_CUDA_ERROR;
    }
    if (!staged_d2h_copies.empty()) {
        size_t staged_offset = 0;
        for (const auto& entry : staged_d2h_copies) {
            std::memcpy(entry.first, pinned_d2h + staged_offset, entry.second);
            staged_offset += entry.second;
        }
    }
    mortred::common::stage_timing::mark("d2h");
    return StatusCode::OK;
}

}  // namespace backend
}  // namespace models
}  // namespace mortred