This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 84f3dad2a9 [Relax][TensorRT] Build and embed engines during code
generation (#20301)
84f3dad2a9 is described below
commit 84f3dad2a94f4dcf793a43e158a68c3713741ce6
Author: Zupeng Wang <[email protected]>
AuthorDate: Fri Sep 11 11:42:44 2026 +0800
[Relax][TensorRT] Build and embed engines during code generation (#20301)
TensorRT BYOC exports currently store graph metadata, so the first
inference builds an engine unless a separate disk cache is managed. Add
opt-in `relax.ext.tensorrt.options.build_at_compile_time=True` to build
engines during `RunCodegen` and embed their plans in the exported
library, addressing #20040. No warm-up inference is needed before
export.
The default lazy path retains its existing byte format. Prebuilt modules
store a versioned TensorRT-specific trailer and deserialize on the
execution GPU without a builder fallback. The mode supports FP32/FP16
with positive static input dimensions; dynamic input shapes and INT8
calibration are rejected. Plans require a compatible GPU, platform, and
TensorRT version. Resource cleanup covers failed construction, CPU
tensor staging, and CUDA shutdown.
Validation on RTX 3090, TensorRT 10.13.3, CUDA 13.0, and LLVM 18.1.8:
- New feature tests: **14 passed, 1 skipped**. Fresh subprocesses
interpose `createInferBuilder_INTERNAL`: lazy exports hit the guard,
while prebuilt exports produce correct outputs and exit normally. Cases
cover weighted multi-partition graphs, FP32/FP16, GPU 1, CPU tensor
staging, and a public-API Conv2d+ReLU partition, along with
serialization compatibility, malformed plans, and rejected calibration
settings.
- Existing TensorRT codegen/inference tests: **31 passed**; backend
partition tests: **7 passed**.
- Separate baseline and patched builds completed; scoped pre-commit and
patch-application checks passed. Runtime-disabled branches passed syntax
compilation; their full integration build was not run, and the
corresponding test was skipped.
These results were obtained on base
`40c2f54908e96875fa19e92db83f123854d97991` with tvm-ffi
`9d784c4da74ff7360d76c79a89fe60a63516f880`. Main subsequently advanced
through #20294 and #20299, updating tvm-ffi; the local results do not
cover that newer integration. This Draft preserves the reviewed and
tested patch. No startup-latency or throughput claim is made.
Signed-off-by: Zupeng Wang <[email protected]>
---
python/tvm/relax/backend/contrib/tensorrt.py | 50 +++
src/relax/backend/contrib/tensorrt/codegen.cc | 54 ++-
.../extra/contrib/tensorrt/tensorrt_builder.cc | 56 +--
.../extra/contrib/tensorrt/tensorrt_builder.h | 4 +
.../extra/contrib/tensorrt/tensorrt_runtime.cc | 269 ++++++++++++-
.../relax/test_tensorrt_engine_serialization.py | 422 +++++++++++++++++++++
6 files changed, 814 insertions(+), 41 deletions(-)
diff --git a/python/tvm/relax/backend/contrib/tensorrt.py
b/python/tvm/relax/backend/contrib/tensorrt.py
index 21ab01b24d..3711f29de1 100644
--- a/python/tvm/relax/backend/contrib/tensorrt.py
+++ b/python/tvm/relax/backend/contrib/tensorrt.py
@@ -203,6 +203,56 @@ def partition_for_tensorrt(mod: IRModule) -> IRModule:
mod : tvm.ir.IRModule
The module with TensorRT-supported subgraphs grouped into composite
functions annotated for the ``tensorrt`` codegen.
+
+ Notes
+ -----
+ By default, TensorRT builds engines on the first inference call.
+ To build and embed engines before exporting the library, enable
+ ``build_at_compile_time`` in the ``relax.ext.tensorrt.options`` PassContext
+ configuration while running ``RunCodegen``.
+
+ Compile-time building requires the TensorRT runtime and a CUDA GPU on the
+ compilation host. Select the deployment GPU before starting code
generation,
+ for example with ``CUDA_VISIBLE_DEVICES``. Tensor inputs must have static,
+ positive dimensions and FP16 or FP32 dtype. Dynamic inputs and INT8
+ calibration are not supported by this mode.
+
+ The exported library contains the serialized engines. No inference call or
+ disk engine cache is needed before export. On the first inference call,
+ the runtime deserializes these plans instead of rebuilding engines.
+ Deployment still requires the TensorRT-enabled TVM runtime, CUDA, and a
+ compatible GPU and TensorRT version; the plans are not portable across
+ arbitrary GPU architectures or TensorRT versions.
+
+ Examples
+ --------
+ After binding model weights as constants, partition and generate code on
the
+ compilation GPU before exporting::
+
+ import tvm
+ from tvm import relax
+ from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
+
+ partitioned = partition_for_tensorrt(mod)
+ with tvm.transform.PassContext(
+ config={
+ "relax.ext.tensorrt.options": {
+ "build_at_compile_time": True,
+ }
+ }
+ ):
+ offloaded = relax.transform.RunCodegen()(partitioned)
+ executable = tvm.compile(offloaded, target="cuda")
+ executable.export_library("model.so")
+
+ Load the library and run inference in a new process on the deployment GPU::
+
+ import tvm
+ from tvm import relax
+
+ executable = tvm.runtime.load_module("model.so")
+ vm = relax.VirtualMachine(executable, tvm.cuda(0))
+ result = vm["main"](*inputs)
"""
patterns = get_patterns_with_prefix("tensorrt")
mod = FuseOpsByPattern(patterns, bind_constants=True,
annotate_codegen=False)(mod)
diff --git a/src/relax/backend/contrib/tensorrt/codegen.cc
b/src/relax/backend/contrib/tensorrt/codegen.cc
index 2b5d4d3db7..4d40ec9185 100644
--- a/src/relax/backend/contrib/tensorrt/codegen.cc
+++ b/src/relax/backend/contrib/tensorrt/codegen.cc
@@ -60,6 +60,7 @@ struct TensorRTCompilerConfigNode : public ffi::Object {
bool remove_no_mac_subgraphs;
bool use_fp16;
bool use_uint8;
+ bool build_at_compile_time;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -77,6 +78,9 @@ struct TensorRTCompilerConfigNode : public ffi::Object {
.def_ro("use_fp16", &TensorRTCompilerConfigNode::use_fp16, "Use FP16",
refl::DefaultValue(false))
.def_ro("use_uint8", &TensorRTCompilerConfigNode::use_uint8, "Use
uint8",
+ refl::DefaultValue(false))
+ .def_ro("build_at_compile_time",
&TensorRTCompilerConfigNode::build_at_compile_time,
+ "Build and embed TensorRT engines during code generation",
refl::DefaultValue(false));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ext.attrs.TensorRTCompilerConfig",
@@ -220,8 +224,8 @@ class CollectFromCompositeFunctionBody : public ExprVisitor
{
/*!
* \brief Generates an TensorRTModule from a relax expression by serializing
the expression to a
- * json representation. TensorRT is not required here because use of TensorRT
APIs is deferred until
- * runtime.
+ * json representation. TensorRT APIs are deferred until runtime unless
compile-time engine
+ * building is explicitly enabled.
*/
class TensorRTJSONSerializer : public JSONSerializer {
public:
@@ -335,18 +339,60 @@ void CollectFromCompositeFunctionBody::VisitExpr_(const
CallNode* call_node) {
ffi::Array<ffi::Module> TensorRTCompiler(ffi::Array<Function> functions,
ffi::Map<ffi::String, ffi::Any>
/*unused*/,
ffi::Map<Constant, ffi::String>
constant_names) {
+ auto cfg =
transform::PassContext::Current()->GetConfig<TensorRTCompilerConfig>(
+ "relax.ext.tensorrt.options");
+ bool build_at_compile_time = cfg.has_value() &&
cfg.value()->build_at_compile_time;
+ ffi::Map<ffi::String, runtime::Tensor> constant_tensors;
+ if (build_at_compile_time) {
+ for (const auto& entry : constant_names) {
+ constant_tensors.Set(entry.second, entry.first->data);
+ }
+ }
+
ffi::Array<ffi::Module> compiled_functions;
for (const auto& func : functions) {
VLOG(1) << "TensorRT partition:" << std::endl << func;
+ if (build_at_compile_time) {
+ // Reject dynamic inputs before the JSON serializer requires integer
shapes.
+ auto check_input_type = [&](const auto& self, const Type& type) -> void {
+ if (const auto* tuple = type.as<TupleTypeNode>()) {
+ for (const auto& field : tuple->fields) self(self, field);
+ return;
+ }
+ const auto* tensor = type.as<TensorTypeNode>();
+ TVM_FFI_CHECK(tensor != nullptr && tensor->shape.has_value(),
ValueError)
+ << "TensorRT compile-time engine building requires static positive
input dimensions";
+ const auto* shape = tensor->shape.value().as<ShapeExprNode>();
+ TVM_FFI_CHECK(shape != nullptr, ValueError)
+ << "TensorRT compile-time engine building requires static positive
input dimensions";
+ for (const auto& dim : shape->values) {
+ const auto* value = dim.as<IntImmNode>();
+ TVM_FFI_CHECK(value != nullptr && value->value > 0, ValueError)
+ << "TensorRT compile-time engine building requires static
positive input dimensions";
+ }
+ };
+ for (const auto& param : func->params)
check_input_type(check_input_type, GetType(param));
+ }
TensorRTJSONSerializer serializer(constant_names, AnalyzeVar2Value(func));
serializer.serialize(func);
std::string graph_json = serializer.GetJSON();
VLOG(1) << "TensorRT JSON:" << std::endl << graph_json;
- auto constant_names = serializer.GetConstantNames();
+ auto ordered_constant_names = serializer.GetConstantNames();
const auto pf =
tvm::ffi::Function::GetGlobalRequired("runtime.tensorrt_runtime_create");
std::string func_name = GetExtSymbol(func);
VLOG(1) << "Creating tensorrt ffi::Module for '" << func_name << "'";
- compiled_functions.push_back(pf(func_name, graph_json,
constant_names).cast<ffi::Module>());
+ auto module = pf(func_name, graph_json,
ordered_constant_names).cast<ffi::Module>();
+ if (build_at_compile_time) {
+ auto build_engine = module->GetFunction("build_engine");
+ TVM_FFI_CHECK(build_engine.has_value(), RuntimeError)
+ << "TensorRT compile-time engine building requires the TensorRT
runtime";
+ ffi::Array<runtime::Tensor> ordered_constants;
+ for (const auto& name : ordered_constant_names) {
+ ordered_constants.push_back(constant_tensors[name]);
+ }
+ build_engine.value()(ordered_constants);
+ }
+ compiled_functions.push_back(module);
}
return compiled_functions;
}
diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc
b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc
index 10b3bdf447..94a03dabcc 100644
--- a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc
+++ b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc
@@ -49,7 +49,8 @@ TensorRTBuilder::TensorRTBuilder(TensorRTLogger* logger,
use_int8_(false),
calibrator_(calibrator) {
// Create TRT builder and network.
- builder_ = nvinfer1::createInferBuilder(*trt_logger_);
+ std::unique_ptr<nvinfer1::IBuilder>
builder(nvinfer1::createInferBuilder(*trt_logger_));
+ TVM_FFI_ICHECK(builder != nullptr) << "Creating the TensorRT builder failed";
// TensorRT 10 removed implicit-batch mode and the kEXPLICIT_BATCH creation
flag; every network is
// explicit-batch, so the batch dimension is simply dimension 0 of each
binding and is varied
@@ -57,9 +58,14 @@ TensorRTBuilder::TensorRTBuilder(TensorRTLogger* logger,
if (calibrator_ != nullptr) {
use_int8_ = true;
}
- network_ = builder_->createNetworkV2(0U);
+ std::unique_ptr<nvinfer1::INetworkDefinition>
network(builder->createNetworkV2(0U));
+ TVM_FFI_ICHECK(network != nullptr) << "Creating the TensorRT network failed";
+ builder_ = builder.release();
+ network_ = network.release();
}
+TensorRTBuilder::~TensorRTBuilder() { CleanUp(); }
+
nvinfer1::DataType DLDataType2NVDataType(DLDataType data_type) {
TVM_FFI_ICHECK(data_type.code == kDLFloat && (data_type.bits == 16 ||
data_type.bits == 32))
<< "Invalid input Tensor type. Only float16 and float32 are supported";
@@ -155,6 +161,7 @@ void TensorRTBuilder::AddLayer(int nid, const
JSONGraphNode& node) {
TensorRTEngineAndContext TensorRTBuilder::BuildEngine() {
// Build engine.
config_ = builder_->createBuilderConfig();
+ TVM_FFI_ICHECK(config_ != nullptr) << "Creating the TensorRT builder config
failed";
// TensorRT 10 replaced IBuilderConfig::setMaxWorkspaceSize with a tunable
memory pool.
config_->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE,
max_workspace_size_);
// Disable TF32 (on by default on Ampere+) so FP32 layers match TVM's
full-precision reference.
@@ -172,7 +179,9 @@ TensorRTEngineAndContext TensorRTBuilder::BuildEngine() {
// Every network is explicit-batch in TRT10, so always add an optimization
profile that pins each
// input to its concrete shape (with a minimum batch of 1 for dynamic batch
dimensions).
+ // The builder owns this profile and releases it during CleanUp.
auto profile = builder_->createOptimizationProfile();
+ TVM_FFI_ICHECK(profile != nullptr) << "Creating the TensorRT optimization
profile failed";
for (int i = 0; i < network_->getNbInputs(); ++i) {
auto name = network_->getInput(i)->getName();
const uint32_t entry_id = entry_id_map_[name];
@@ -194,23 +203,27 @@ TensorRTEngineAndContext TensorRTBuilder::BuildEngine() {
// TensorRT 10 removed buildEngineWithConfig; build a serialized engine and
deserialize it through
// an IRuntime that is kept alive alongside the engine
(TensorRTEngineAndContext::runtime).
- nvinfer1::IHostMemory* plan = builder_->buildSerializedNetwork(*network_,
*config_);
- TVM_FFI_ICHECK(plan) << "Failed to build TensorRT serialized network.";
- nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(*trt_logger_);
- nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(plan->data(),
plan->size());
- delete plan;
- if (engine == nullptr) {
- delete runtime;
- TVM_FFI_THROW(InternalError) << "Failed to deserialize the TensorRT
engine.";
- }
+ std::unique_ptr<nvinfer1::IHostMemory> plan(
+ builder_->buildSerializedNetwork(*network_, *config_));
+ TVM_FFI_ICHECK(plan != nullptr) << "Failed to build TensorRT serialized
network.";
+ std::unique_ptr<nvinfer1::IRuntime>
runtime(nvinfer1::createInferRuntime(*trt_logger_));
+ TVM_FFI_ICHECK(runtime != nullptr) << "Creating the TensorRT deserialization
runtime failed";
+ std::unique_ptr<nvinfer1::ICudaEngine> engine(
+ runtime->deserializeCudaEngine(plan->data(), plan->size()));
+ TVM_FFI_ICHECK(engine != nullptr) << "Failed to deserialize the TensorRT
engine.";
TVM_FFI_ICHECK_EQ(engine->getNbIOTensors(),
static_cast<int32_t>(network_input_names_.size() +
network_output_names_.size()));
- nvinfer1::IExecutionContext* context = engine->createExecutionContext();
+ std::unique_ptr<nvinfer1::IExecutionContext>
context(engine->createExecutionContext());
+ TVM_FFI_ICHECK(context != nullptr) << "Creating the TensorRT execution
context failed";
CleanUp();
- TVM_FFI_ICHECK(context);
-
- return {runtime, engine, context, network_input_names_,
network_output_names_};
+ TensorRTEngineAndContext result;
+ result.inputs = network_input_names_;
+ result.outputs = network_output_names_;
+ result.runtime = runtime.release();
+ result.engine = engine.release();
+ result.context = context.release();
+ return result;
}
nvinfer1::Weights TensorRTBuilder::GetDLTensorAsWeights(const DLTensor* dptr,
@@ -228,11 +241,12 @@ nvinfer1::Weights
TensorRTBuilder::GetDLTensorAsWeights(const DLTensor* dptr,
count *= dptr->shape[i];
}
weight.count = count;
- weight.values = new float[count];
- // Tensor::CopyToBytes throws on failure (the old C API
TVMTensorCopyToBytes/TVMGetLastError
- // were removed during the tvm-ffi refactor).
- Tensor::CopyToBytes(dptr, const_cast<void*>(weight.values), weight_bytes);
+ std::unique_ptr<float[]> values(new float[count]);
+ weight.values = values.get();
+ // Keep temporary storage owned until both the copy and registration succeed.
+ Tensor::CopyToBytes(dptr, values.get(), weight_bytes);
trt_weights_.push_back(weight);
+ values.release();
return weight;
}
@@ -246,23 +260,19 @@ nvinfer1::ITensor*
TensorRTBuilder::GetInputAsTensor(const TensorRTOpInput& inpu
void TensorRTBuilder::CleanUp() {
// TensorRT 10 removed obj->destroy(); objects are released with the delete
operator.
VLOG(1) << "Destroying TensorRT network";
- TVM_FFI_ICHECK(network_);
delete network_;
network_ = nullptr;
VLOG(1) << "Destroying TensorRT config";
- TVM_FFI_ICHECK(config_);
delete config_;
config_ = nullptr;
VLOG(1) << "Destroying TensorRT builder";
- TVM_FFI_ICHECK(builder_);
delete builder_;
builder_ = nullptr;
VLOG(1) << "Destroying TensorRT weights";
for (auto weight : trt_weights_) {
- TVM_FFI_ICHECK(weight.values);
if (weight.type == nvinfer1::DataType::kFLOAT || weight.type ==
nvinfer1::DataType::kHALF) {
delete[] static_cast<const float*>(weight.values);
} else {
diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.h
b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.h
index 108f56b9f3..0fd06acc24 100644
--- a/src/runtime/extra/contrib/tensorrt/tensorrt_builder.h
+++ b/src/runtime/extra/contrib/tensorrt/tensorrt_builder.h
@@ -76,6 +76,10 @@ class TensorRTBuilder {
size_t max_workspace_size, bool use_fp16,
nvinfer1::IInt8Calibrator* calibrator = nullptr);
+ ~TensorRTBuilder();
+ TensorRTBuilder(const TensorRTBuilder&) = delete;
+ TensorRTBuilder& operator=(const TensorRTBuilder&) = delete;
+
/*!
* \brief Add TensorRT input(s) for input node in network definition.
* \param nid The input node id.
diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
b/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
index b421e096d7..1f177f7605 100644
--- a/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
+++ b/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
@@ -28,10 +28,13 @@
#include <tvm/runtime/logging.h>
#include <tvm/runtime/tensor.h>
+#include <cstdint>
#include <fstream>
+#include <limits>
#include <memory>
#include <string>
#include <unordered_map>
+#include <utility>
#include <vector>
#include "../../../../support/env.h"
@@ -42,6 +45,7 @@
#ifdef TVM_GRAPH_EXECUTOR_TENSORRT
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/cuda/base.h>
+#include <tvm/ffi/extra/cuda/device_guard.h>
#include "NvInfer.h"
#include "tensorrt_builder.h"
@@ -108,6 +112,70 @@ class TensorRTRuntime : public JSONRuntimeBase {
return ffi::Module::kBinarySerializable | ffi::Module::kRunnable;
}
+ ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) override {
+ if (name == "build_engine") {
+ ffi::ObjectPtr<ffi::Object> self = ffi::GetObjectPtr<ffi::Object>(this);
+ return ffi::Function([self, this](ffi::PackedArgs args, ffi::Any* rv) {
+ TVM_FFI_ICHECK_EQ(args.size(), 1U);
+ std::lock_guard<std::mutex> guard(initialize_mutex_);
+ TVM_FFI_ICHECK(!initialized_) << "Build the TensorRT engine before
initializing the module";
+ BuildEngineAtCompileTime(args[0].cast<ffi::Array<Tensor>>());
+ *rv = 0;
+ });
+ }
+ return JSONRuntimeBase::GetFunction(name);
+ }
+
+ ffi::Bytes SaveToBytes() const override {
+ auto graph_bytes = JSONRuntimeBase::SaveToBytes();
+ if (serialized_engine_.empty()) return graph_bytes;
+
+ // Keep the JSON runtime prefix unchanged. Only opt-in, prebuilt modules
have this trailer.
+ std::string result(graph_bytes.data(), graph_bytes.size());
+ support::BytesOutStream stream(&result);
+ stream.Write(kEngineMagic);
+ stream.Write(kEngineVersion);
+ stream.Write(serialized_engine_);
+ stream.Write(serialized_inputs_);
+ stream.Write(serialized_outputs_);
+ stream.Write(serialized_batch_size_);
+ return ffi::Bytes(std::move(result));
+ }
+
+ static ffi::Module LoadFromBytes(const ffi::Bytes& bytes) {
+ support::BytesInStream stream(bytes);
+ std::string symbol, graph_json;
+ std::vector<std::string> consts;
+ TVM_FFI_ICHECK(stream.Read(&symbol)) << "Loading symbol name failed";
+ TVM_FFI_ICHECK(stream.Read(&graph_json)) << "Loading graph json failed";
+ TVM_FFI_ICHECK(stream.Read(&consts)) << "Loading the const name list
failed";
+ ffi::Array<ffi::String> const_names;
+ for (const auto& name : consts) const_names.push_back(name);
+ auto n = ffi::make_object<TensorRTRuntime>(symbol, graph_json,
const_names);
+
+ uint64_t magic;
+ size_t nread = stream.Read(&magic, sizeof(magic));
+ if (nread == 0) return ffi::Module(n); // Legacy JSON-only module.
+ TVM_FFI_ICHECK_EQ(nread, sizeof(magic)) << "Truncated TensorRT engine
trailer";
+ TVM_FFI_ICHECK_EQ(magic, kEngineMagic) << "Invalid TensorRT engine
trailer";
+ uint32_t version;
+ TVM_FFI_ICHECK(stream.Read(&version)) << "Loading TensorRT engine version
failed";
+ TVM_FFI_ICHECK_EQ(version, kEngineVersion) << "Unsupported TensorRT engine
version";
+ TVM_FFI_ICHECK(stream.Read(&n->serialized_engine_)) << "Loading TensorRT
engine failed";
+ TVM_FFI_ICHECK(!n->serialized_engine_.empty()) << "The embedded TensorRT
engine is empty";
+ TVM_FFI_ICHECK(stream.Read(&n->serialized_inputs_)) << "Loading TensorRT
input names failed";
+ TVM_FFI_ICHECK(stream.Read(&n->serialized_outputs_)) << "Loading TensorRT
output names failed";
+ TVM_FFI_ICHECK(stream.Read(&n->serialized_batch_size_)) << "Loading
TensorRT batch size failed";
+ TVM_FFI_ICHECK_GT(n->serialized_batch_size_, 0) << "Invalid TensorRT batch
size";
+ TVM_FFI_ICHECK_EQ(n->serialized_inputs_.size(), n->input_var_eid_.size());
+ TVM_FFI_ICHECK_EQ(n->serialized_outputs_.size(), n->outputs_.size());
+ char trailing;
+ TVM_FFI_ICHECK_EQ(stream.Read(&trailing, sizeof(trailing)), 0U)
+ << "Unexpected data after the TensorRT engine";
+ // Defer GPU-dependent deserialization until the first call supplies the
execution device.
+ return ffi::Module(n);
+ }
+
/*!
* \brief Initialize runtime. Create TensorRT layer from JSON
* representation.
@@ -117,9 +185,15 @@ class TensorRTRuntime : public JSONRuntimeBase {
void Init(const ffi::Array<Tensor>& consts) override {
TVM_FFI_ICHECK_EQ(consts.size(), const_idx_.size())
<< "The number of input constants must match the number of required.";
+ if (!serialized_engine_.empty()) {
+ TVM_FFI_ICHECK(!support::GetEnv("TVM_TENSORRT_USE_INT8", false) &&
+ support::GetEnv("TENSORRT_NUM_CALI_INT8", 0) == 0 &&
+ num_calibration_batches_remaining_ == 0)
+ << "A prebuilt TensorRT engine cannot perform INT8 calibration";
+ }
LoadGlobalAttributes();
SetupConstants(consts);
- GetCachedEnginesFromDisk();
+ if (serialized_engine_.empty()) GetCachedEnginesFromDisk();
}
void LoadGlobalAttributes() {
@@ -149,8 +223,19 @@ class TensorRTRuntime : public JSONRuntimeBase {
}
#ifdef TVM_GRAPH_EXECUTOR_TENSORRT
- /*! \brief Destroy engines and contexts. */
+ /*! \brief Destroy engines and contexts on the device that owns them. */
void DestroyEngines() {
+ if (embedded_device_id_ >= 0) {
+ ffi::CUDADeviceGuard guard(embedded_device_id_);
+ DestroyEnginesOnCurrentDevice();
+ device_buffers_.reset();
+ embedded_device_id_ = -1;
+ } else {
+ DestroyEnginesOnCurrentDevice();
+ }
+ }
+
+ void DestroyEnginesOnCurrentDevice() {
for (auto& it : trt_engine_cache_) {
// TensorRT 10 removed obj->destroy(); release with delete. The
deserialization runtime must
// outlive the engine it produced, so delete the context, then the
engine, then the runtime.
@@ -166,13 +251,143 @@ class TensorRTRuntime : public JSONRuntimeBase {
}
~TensorRTRuntime() {
+ // FFI can retain imported modules until process shutdown, after CUDA has
already unloaded.
+ // TensorRT destructors and CUDADeviceGuard cannot run then; the process
reclaims the resources.
+ int current_device;
+ if ((!trt_engine_cache_.empty() || device_buffers_ || calibrator_) &&
+ cudaGetDevice(¤t_device) == cudaErrorCudartUnloading) {
+ // CPU-input staging tensors and the calibrator also own CUDA
allocations. Release their
+ // containers so member destruction cannot enter CUDA after this
confirmed shutdown state.
+ (void)device_buffers_.release();
+ (void)calibrator_.release();
+ return;
+ }
VLOG(1) << "Destroying TensorRT runtime";
DestroyEngines();
VLOG(1) << "Destroyed TensorRT runtime";
}
- /*! \brief Run inference using built engine. */
+ /*! \brief Build from static input metadata and weights, without running
inference. */
+ void BuildEngineAtCompileTime(const ffi::Array<Tensor>& consts) {
+ TVM_FFI_ICHECK_EQ(consts.size(), const_idx_.size())
+ << "The number of input constants must match the number of required.";
+ TVM_FFI_ICHECK(!support::GetEnv("TVM_TENSORRT_USE_INT8", false) &&
+ support::GetEnv("TENSORRT_NUM_CALI_INT8", 0) == 0 &&
+ num_calibration_batches_remaining_ == 0)
+ << "TensorRT build_at_compile_time does not support INT8 calibration";
+ TVM_FFI_ICHECK(!input_var_eid_.empty())
+ << "TensorRT build_at_compile_time requires at least one tensor input";
+ TVM_FFI_ICHECK(serialized_engine_.empty()) << "The TensorRT engine has
already been built";
+
+ std::vector<std::vector<int64_t>> shapes(NumEntries());
+ std::vector<DLTensor> inputs(NumEntries());
+ for (auto nid : input_nodes_) {
+ const auto& node = nodes_[nid];
+ if (node.GetOpType() != "input") continue;
+ auto node_shapes = node.GetOpShape();
+ auto dtypes = node.GetOpDataType();
+ for (size_t j = 0; j < node_shapes.size(); ++j) {
+ TVM_FFI_ICHECK(dtypes[j].code == kDLFloat && dtypes[j].lanes == 1 &&
+ (dtypes[j].bits == 16 || dtypes[j].bits == 32))
+ << "TensorRT build_at_compile_time requires float16 or float32
inputs";
+ uint32_t eid = EntryID(nid, j);
+ for (int64_t dim : node_shapes[j]) {
+ TVM_FFI_ICHECK_GT(dim, 0)
+ << "TensorRT build_at_compile_time requires static positive
input dimensions";
+ TVM_FFI_ICHECK_LE(dim, std::numeric_limits<int32_t>::max())
+ << "TensorRT input dimensions must fit in int32";
+ shapes[eid].push_back(dim);
+ }
+ inputs[eid].device = DLDevice{kDLCPU, 0};
+ inputs[eid].ndim = static_cast<int32_t>(shapes[eid].size());
+ inputs[eid].dtype = dtypes[j];
+ inputs[eid].shape = shapes[eid].data();
+ }
+ }
+
+ for (const auto& constant : consts) {
+ TVM_FFI_ICHECK_EQ(constant->device.device_type, kDLCPU)
+ << "TensorRT build_at_compile_time requires constants on the CPU";
+ const auto dtype = constant->dtype;
+ TVM_FFI_ICHECK(dtype.code == kDLFloat && dtype.lanes == 1 &&
+ (dtype.bits == 16 || dtype.bits == 32))
+ << "TensorRT build_at_compile_time requires float16 or float32
constants";
+ }
+ LoadGlobalAttributes();
+ auto original_entries = data_entry_;
+ try {
+ SetupConstants(consts);
+ for (auto eid : input_var_eid_) data_entry_[eid] = &inputs[eid];
+ int batch_size = GetBatchSize();
+ BuildEngineFromJson(batch_size);
+ const auto& built = trt_engine_cache_.at(std::make_pair(symbol_name_,
batch_size));
+ std::unique_ptr<nvinfer1::IHostMemory> plan(built.engine->serialize());
+ TVM_FFI_ICHECK(plan != nullptr) << "Serializing the TensorRT engine
failed";
+ TVM_FFI_ICHECK_GT(plan->size(), 0U) << "The serialized TensorRT engine
is empty";
+ serialized_engine_.assign(static_cast<const char*>(plan->data()),
plan->size());
+ serialized_inputs_ = built.inputs;
+ serialized_outputs_ = built.outputs;
+ serialized_batch_size_ = batch_size;
+ } catch (...) {
+ data_entry_ = std::move(original_entries);
+ DestroyEngines();
+ serialized_engine_.clear();
+ serialized_inputs_.clear();
+ serialized_outputs_.clear();
+ serialized_batch_size_ = 0;
+ throw;
+ }
+ // Neither input metadata nor the supplied constants are owned by the
runtime at this point.
+ data_entry_ = std::move(original_entries);
+ DestroyEngines();
+ }
+
+ void LoadEmbeddedEngine(int device_id) {
+ TVM_FFI_ICHECK(!support::GetEnv("TVM_TENSORRT_USE_INT8", false) &&
+ support::GetEnv("TENSORRT_NUM_CALI_INT8", 0) == 0 &&
+ num_calibration_batches_remaining_ == 0)
+ << "A prebuilt TensorRT engine cannot perform INT8 calibration";
+ std::unique_ptr<nvinfer1::IRuntime>
runtime(nvinfer1::createInferRuntime(GetTensorRTLogger()));
+ TVM_FFI_ICHECK(runtime != nullptr) << "Creating the TensorRT
deserialization runtime failed";
+ std::unique_ptr<nvinfer1::ICudaEngine> engine(
+ runtime->deserializeCudaEngine(serialized_engine_.data(),
serialized_engine_.size()));
+ TVM_FFI_ICHECK(engine != nullptr)
+ << "Failed to deserialize the embedded TensorRT engine. Recompile the
model for this "
+ "GPU, platform, and TensorRT version; compile-time engines are not
rebuilt at runtime.";
+ std::unique_ptr<nvinfer1::IExecutionContext>
context(engine->createExecutionContext());
+ TVM_FFI_ICHECK(context != nullptr) << "Creating the embedded TensorRT
execution context failed";
+ TensorRTEngineAndContext loaded;
+ loaded.inputs = serialized_inputs_;
+ loaded.outputs = serialized_outputs_;
+ auto key = std::make_pair(symbol_name_, serialized_batch_size_);
+ auto result = trt_engine_cache_.emplace(key, std::move(loaded));
+ TVM_FFI_ICHECK(result.second) << "The embedded TensorRT engine has already
been initialized";
+ result.first->second.runtime = runtime.release();
+ result.first->second.engine = engine.release();
+ result.first->second.context = context.release();
+ max_batch_size_ = serialized_batch_size_;
+ embedded_device_id_ = device_id;
+ }
+
+ /*! \brief Run inference using a prebuilt or lazily built engine. */
void Run() override {
+ if (serialized_engine_.empty()) {
+ RunWithEngine();
+ return;
+ }
+ // Const-loader initialization has no execution device. The first
invocation does, so load the
+ // plan here without invoking the TensorRT builder and restore the
caller's current device.
+ const DLDevice& dev = data_entry_[input_var_eid_[0]]->device;
+ const int device_id = dev.device_type == kDLCUDA ? dev.device_id : 0;
+ ffi::CUDADeviceGuard guard(device_id);
+ if (embedded_device_id_ != device_id) {
+ DestroyEngines();
+ LoadEmbeddedEngine(device_id);
+ }
+ RunWithEngine();
+ }
+
+ void RunWithEngine() {
auto& engine_and_context = GetOrBuildEngine();
int batch_size = GetBatchSize();
if (batch_size == 0) return;
@@ -297,6 +512,12 @@ class TensorRTRuntime : public JSONRuntimeBase {
int batch_size = GetBatchSize();
int compatible_engine_batch_size = -1;
bool find_engine_flag = FindCompatibleEngine(batch_size,
&compatible_engine_batch_size);
+ if (!serialized_engine_.empty()) {
+ TVM_FFI_ICHECK_EQ(batch_size, serialized_batch_size_)
+ << "The input batch size does not match the compile-time TensorRT
engine";
+ TVM_FFI_ICHECK(find_engine_flag) << "The embedded TensorRT engine has
not been initialized";
+ return trt_engine_cache_.at(std::make_pair(symbol_name_,
serialized_batch_size_));
+ }
const bool use_int8 = (support::GetEnv("TVM_TENSORRT_USE_INT8", 0) != 0);
const bool int8_calibration_not_used_or_not_complete =
(calibrator_ != nullptr && num_calibration_batches_remaining_ != 0);
@@ -478,24 +699,27 @@ class TensorRTRuntime : public JSONRuntimeBase {
/*! \brief Retreive a GPU buffer for input or output or allocate if needed.
Keyed by TensorRT IO
* tensor name (TRT10 has no binding indices). */
Tensor GetOrAllocateDeviceBuffer(const std::string& name, int entry_id) {
+ if (!device_buffers_) {
+ device_buffers_ = std::make_unique<std::unordered_map<std::string,
Tensor>>();
+ }
+ auto& device_buffers = *device_buffers_;
+ DLDevice device{kDLCUDA, embedded_device_id_ >= 0 ? embedded_device_id_ :
0};
std::vector<int64_t> shape(data_entry_[entry_id]->shape,
data_entry_[entry_id]->shape +
data_entry_[entry_id]->ndim);
- if (device_buffers_.count(name)) {
+ if (device_buffers.count(name)) {
// Buffer is already initialized.
- if (shape[0] > device_buffers_[name]->shape[0]) {
+ if (shape[0] > device_buffers[name]->shape[0]) {
// Buffer is too small. Need to allocate bigger buffer.
- device_buffers_[name] =
- runtime::Tensor::Empty(shape, data_entry_[entry_id]->dtype,
{kDLCUDA, 0});
- } else if (shape[0] < device_buffers_[name]->shape[0]) {
+ device_buffers[name] = runtime::Tensor::Empty(shape,
data_entry_[entry_id]->dtype, device);
+ } else if (shape[0] < device_buffers[name]->shape[0]) {
// Buffer is too large. Create view.
- return device_buffers_[name].CreateView(shape,
data_entry_[entry_id]->dtype);
+ return device_buffers[name].CreateView(shape,
data_entry_[entry_id]->dtype);
}
} else {
// Buffer not initialized yet.
- device_buffers_[name] =
- runtime::Tensor::Empty(shape, data_entry_[entry_id]->dtype,
{kDLCUDA, 0});
+ device_buffers[name] = runtime::Tensor::Empty(shape,
data_entry_[entry_id]->dtype, device);
}
- return device_buffers_.at(name);
+ return device_buffers.at(name);
}
void CreateInt8Calibrator(const TensorRTEngineAndContext&
engine_and_context) {
@@ -513,6 +737,9 @@ class TensorRTRuntime : public JSONRuntimeBase {
std::unordered_map<std::pair<std::string, int>, TensorRTEngineAndContext,
PairHash>
trt_engine_cache_;
+ /*! \brief Device owning the deserialized compile-time engine, or -1 before
the first call. */
+ int embedded_device_id_ = -1;
+
/*! \brief Calibrator for INT8 mode. */
std::unique_ptr<TensorRTCalibrator> calibrator_;
@@ -520,7 +747,7 @@ class TensorRTRuntime : public JSONRuntimeBase {
* is not "cuda". Since TensorRT execution can only read data from GPU, we
need to copy data from
* the runtime device to these buffers first. These will be allocated for
the highest batch size
* used by all engines. */
- std::unordered_map<std::string, Tensor> device_buffers_;
+ std::unique_ptr<std::unordered_map<std::string, Tensor>> device_buffers_;
#else // TVM_GRAPH_EXECUTOR_TENSORRT
void Run() override {
@@ -533,11 +760,25 @@ class TensorRTRuntime : public JSONRuntimeBase {
<< "Please build with USE_TENSORRT_RUNTIME.";
}
+ void BuildEngineAtCompileTime(const ffi::Array<Tensor>& consts) {
+ TVM_FFI_THROW(InternalError)
+ << "TensorRT build_at_compile_time requires TVM to be built with
USE_TENSORRT_RUNTIME";
+ }
+
bool GetCachedEnginesFromDisk() { return false; }
void CacheEngineToDisk() {}
#endif // TVM_GRAPH_EXECUTOR_TENSORRT
+ static constexpr uint64_t kEngineMagic = 0x31545254564d5445ULL;
+ static constexpr uint32_t kEngineVersion = 1;
+
+ // Retain serialized bytes independently of initialized CUDA objects,
including after reloading.
+ std::string serialized_engine_;
+ std::vector<std::string> serialized_inputs_;
+ std::vector<std::string> serialized_outputs_;
+ int serialized_batch_size_ = 0;
+
bool use_implicit_batch_;
size_t max_workspace_size_;
@@ -570,7 +811,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("runtime.tensorrt_runtime_create", TensorRTRuntimeCreate)
- .def("ffi.Module.load_from_bytes.tensorrt",
JSONRuntimeBase::LoadFromBytes<TensorRTRuntime>);
+ .def("ffi.Module.load_from_bytes.tensorrt",
TensorRTRuntime::LoadFromBytes);
}
} // namespace contrib
diff --git a/tests/python/relax/test_tensorrt_engine_serialization.py
b/tests/python/relax/test_tensorrt_engine_serialization.py
new file mode 100644
index 0000000000..a527c63880
--- /dev/null
+++ b/tests/python/relax/test_tensorrt_engine_serialization.py
@@ -0,0 +1,422 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you 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.
+"""Compile-time TensorRT engines and their exported runtime artifacts."""
+
+import json
+import os
+import struct
+import subprocess
+import sys
+import textwrap
+
+import numpy as np
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import relax, tirx
+from tvm.relax.backend.contrib.tensorrt import partition_for_tensorrt
+from tvm.relax.dpl import is_op, wildcard
+from tvm.support import cc
+from tvm.testing import env
+
+_has_codegen = tvm.get_global_func("relax.ext.tensorrt", True) is not None
+_runtime_enabled = tvm.get_global_func("relax.is_tensorrt_runtime_enabled",
True)
+_has_runtime = _runtime_enabled is not None and _runtime_enabled()
+
+pytestmark = pytest.mark.skipif(not _has_codegen, reason="TensorRT codegen is
not enabled")
+requires_runtime = pytest.mark.skipif(not _has_runtime, reason="TensorRT
runtime is not enabled")
+requires_gpu = pytest.mark.skipif(not env.has_cuda(), reason="CUDA is not
available")
+
+# Version 1 extends the legacy JSON runtime prefix with a private engine
trailer.
+_ENGINE_MAGIC = 0x31545254564D5445
+
+
[email protected](autouse=True)
+def isolated_tensorrt_options(monkeypatch):
+ for name in (
+ "TVM_TENSORRT_CACHE_DIR",
+ "TVM_TENSORRT_USE_INT8",
+ "TENSORRT_NUM_CALI_INT8",
+ "TVM_TENSORRT_MULTI_ENGINE",
+ "TVM_TENSORRT_USE_FP16",
+ "TVM_TENSORRT_MAX_WORKSPACE_SIZE",
+ ):
+ monkeypatch.delenv(name, raising=False)
+
+
+def _string(value):
+ value = value.encode() if isinstance(value, str) else value
+ return struct.pack("<Q", len(value)) + value
+
+
+def _strings(values):
+ return struct.pack("<Q", len(values)) + b"".join(_string(value) for value
in values)
+
+
+def _legacy_bytes(module):
+ return (
+ _string(module["get_symbol"]())
+ + _string(module.inspect_source())
+ + _strings(module["get_const_vars"]())
+ )
+
+
+def _module_bytes(module):
+ """Extract a leaf's payload from the same import packer used by
export_library."""
+ assert not module.imports
+ packed =
tvm.get_global_func("runtime.ModulePackImportsToTensor")(module).numpy().tobytes()
+ assert struct.unpack_from("<Q", packed)[0] == len(packed) - 8
+ offset = 8
+ # The import tree consists of row pointers followed by child indices.
+ for _ in range(2):
+ count = struct.unpack_from("<Q", packed, offset)[0]
+ offset += 8 + 8 * count
+ kind_size = struct.unpack_from("<Q", packed, offset)[0]
+ offset += 8
+ assert packed[offset : offset + kind_size] == b"tensorrt"
+ offset += kind_size
+ payload_size = struct.unpack_from("<Q", packed, offset)[0]
+ offset += 8
+ assert offset + payload_size == len(packed)
+ return packed[offset:]
+
+
+def _load_bytes(payload):
+ return tvm.get_global_func("ffi.Module.load_from_bytes.tensorrt")(payload)
+
+
+def _relu_partition(shape=(2, 4)):
+ data = relax.Var("data", relax.TensorType(shape, "float32"))
+ builder = relax.BlockBuilder()
+ with builder.function("main", [data]):
+ with builder.dataflow():
+ output = builder.emit_output(relax.op.nn.relu(data))
+ builder.emit_func_output(output)
+ return tvm.transform.Sequential(
+ [
+ relax.transform.FuseOpsByPattern(
+ [("tensorrt.nn.relu", is_op("relax.nn.relu")(wildcard()))]
+ ),
+ relax.transform.MergeCompositeFunctions(),
+ ]
+ )(builder.get())
+
+
+def _codegen(partitioned, eager=None):
+ config = {}
+ if eager is not None:
+ config["relax.ext.tensorrt.options"] = {"build_at_compile_time": eager}
+ with tvm.transform.PassContext(config=config):
+ return relax.transform.RunCodegen()(partitioned)
+
+
+def _external_modules(mod):
+ modules = list(mod.attrs["external_mods"])
+ assert modules and all(module.kind == "tensorrt" for module in modules)
+ return modules
+
+
+def test_default_preserves_legacy_bytes():
+ partitioned = _relu_partition()
+ default = _external_modules(_codegen(partitioned))[0]
+ disabled = _external_modules(_codegen(partitioned, eager=False))[0]
+ payload = _legacy_bytes(default)
+ assert _module_bytes(default) == payload
+ assert _module_bytes(disabled) == payload
+ # A real old-format record must remain loadable and exportable without a
GPU.
+ restored = _load_bytes(payload)
+ assert restored["get_symbol"]() == default["get_symbol"]()
+ assert restored.inspect_source() == default.inspect_source()
+ assert _module_bytes(restored) == payload
+
+
[email protected](_has_runtime, reason="This test requires a codegen-only
TensorRT build")
+def test_eager_requires_runtime():
+ with pytest.raises(RuntimeError, match="USE_TENSORRT_RUNTIME"):
+ _codegen(_relu_partition(), eager=True)
+
+
[email protected](
+ "trailer, message",
+ [
+ (b"partial", "Truncated TensorRT engine trailer"),
+ (struct.pack("<Q", 0), "Invalid TensorRT engine trailer"),
+ (
+ struct.pack("<QI", _ENGINE_MAGIC, 999),
+ "Unsupported TensorRT engine version",
+ ),
+ (
+ struct.pack("<QI", _ENGINE_MAGIC, 1) + _string(b""),
+ "embedded TensorRT engine is empty",
+ ),
+ ],
+)
+def test_reject_malformed_engine_trailer(trailer, message):
+ module = _external_modules(_codegen(_relu_partition()))[0]
+ with pytest.raises(tvm.error.InternalError, match=message):
+ _load_bytes(_legacy_bytes(module) + trailer)
+
+
+@requires_runtime
[email protected]("batch", ["dynamic", 0])
+def test_eager_requires_static_positive_dimensions(batch):
+ dimension = tirx.Var("batch", "int64") if batch == "dynamic" else batch
+ with pytest.raises(ValueError, match="static positive"):
+ _codegen(_relu_partition((dimension, 4)), eager=True)
+
+
+@requires_runtime
+def test_eager_rejects_int8_calibration(monkeypatch):
+ monkeypatch.setenv("TVM_TENSORRT_USE_INT8", "1")
+ monkeypatch.setenv("TENSORRT_NUM_CALI_INT8", "1")
+ with pytest.raises(tvm.error.InternalError, match="INT8 calibration"):
+ _codegen(_relu_partition(), eager=True)
+
+
+def _two_partitions(dtype):
+ data, first, second, third = [
+ relax.Var(name, relax.TensorType((2, 4), dtype))
+ for name in ("data", "first", "second", "third")
+ ]
+ builder = relax.BlockBuilder()
+ with builder.function("main", [data, first, second, third]):
+ with builder.dataflow():
+ difference = builder.emit(relax.op.subtract(second, data))
+ shifted = builder.emit(relax.op.add(difference, first))
+ separated = builder.emit(relax.op.sin(shifted))
+ result = builder.emit_output(relax.op.multiply(separated, third))
+ builder.emit_func_output(result)
+
+ parameters = {
+ "first": np.linspace(-0.3, 0.4, 8, dtype=dtype).reshape(2, 4),
+ "second": np.linspace(0.8, 2.2, 8, dtype=dtype).reshape(2, 4),
+ "third": np.linspace(1.2, 2.9, 8, dtype=dtype).reshape(2, 4),
+ }
+ patterns = [
+ ("tensorrt.subtract", is_op("relax.subtract")(wildcard(), wildcard())),
+ ("tensorrt.add", is_op("relax.add")(wildcard(), wildcard())),
+ ("tensorrt.multiply", is_op("relax.multiply")(wildcard(), wildcard())),
+ ]
+ partitioned = tvm.transform.Sequential(
+ [
+ relax.transform.BindParams("main", parameters),
+ # Leave sin on the TVM side to force two TensorRT partitions.
+ relax.transform.FuseOpsByPattern(patterns, bind_constants=True),
+ relax.transform.MergeCompositeFunctions(),
+ ]
+ )(builder.get())
+ return partitioned, parameters
+
+
+def _conv2d_relu_partition():
+ data = relax.Var("data", relax.TensorType((1, 2, 5, 5), "float32"))
+ weight = relax.Var("weight", relax.TensorType((3, 2, 3, 3), "float32"))
+ builder = relax.BlockBuilder()
+ with builder.function("main", [data, weight]):
+ with builder.dataflow():
+ convolution = builder.emit(relax.op.nn.conv2d(data, weight,
padding=(1, 1)))
+ result = builder.emit_output(relax.op.nn.relu(convolution))
+ builder.emit_func_output(result)
+
+ weight_np = np.linspace(-0.3, 0.4, 54, dtype="float32").reshape(3, 2, 3, 3)
+ bound = relax.transform.BindParams("main", {"weight":
weight_np})(builder.get())
+ partitioned = partition_for_tensorrt(bound)
+ data_np = np.linspace(-1.7, 1.3, 100, dtype="float32").reshape(2, 1, 2, 5,
5)
+ padded = np.pad(data_np, ((0, 0), (0, 0), (0, 0), (1, 1), (1, 1)))
+ windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3),
axis=(-2, -1))
+ expected = np.maximum(np.einsum("snihwkl,oikl->snohw", windows,
weight_np), 0)
+ return partitioned, data_np, expected
+
+
+@requires_runtime
+@requires_gpu
[email protected]
+def test_failed_engine_build_can_be_retried():
+ """A failure after copying weights must not leave a partial engine or
borrowed inputs."""
+ partitioned, _ = _two_partitions("float32")
+ offloaded = _codegen(partitioned, eager=False)
+ module = max(_external_modules(offloaded), key=lambda mod:
len(mod["get_const_vars"]()))
+ graph = json.loads(module.inspect_source())
+ kernels = [node for node in graph["nodes"] if node["op"] == "kernel"]
+ kernels[-1]["name"] = "tensorrt.unsupported_test_op"
+ names = module["get_const_vars"]()
+ assert len(names) == 2
+ failing = tvm.get_global_func("runtime.tensorrt_runtime_create")(
+ module["get_symbol"](), json.dumps(graph), names
+ )
+ constants_by_name = dict(offloaded.attrs["const_name_to_constant"])
+ constants = [constants_by_name[name] for name in names]
+ legacy = _legacy_bytes(failing)
+
+ def check():
+ for _ in range(2):
+ with pytest.raises(tvm.error.InternalError, match="Unsupported
operator"):
+ failing["build_engine"](constants)
+ assert _module_bytes(failing) == legacy
+ assert _module_bytes(_load_bytes(legacy)) == legacy
+
+ tvm.testing.run_with_gpu_lock(check)
+
+
+@requires_runtime
+@requires_gpu
[email protected](not env.has_llvm(), reason="LLVM is not available")
[email protected]
[email protected](sys.platform != "linux", reason="The builder guard uses
LD_PRELOAD")
[email protected](
+ "dtype, device_id, model, execution_device",
+ [
+ ("float32", 0, "elementwise", "cuda"),
+ ("float16", 0, "elementwise", "cuda"),
+ ("float16", 1, "elementwise", "cuda"),
+ ("float32", 0, "conv2d_relu", "cuda"),
+ ("float32", 0, "elementwise", "cpu"),
+ ],
+)
+def test_exported_engines_run_without_builder(
+ tmp_path, monkeypatch, dtype, device_id, model, execution_device
+):
+ """Build weighted partitions before VM creation, then load and run without
a builder."""
+ if not tvm.cuda(device_id).exist:
+ pytest.skip(f"CUDA device {device_id} is not available")
+
+ target = "llvm" if execution_device == "cpu" else "cuda"
+ guard_source = tmp_path / "builder_guard.cc"
+ guard_library = tmp_path / "builder_guard.so"
+ guard_source.write_text(
+ "#include <cstdlib>\n"
+ 'extern "C" void* createInferBuilder_INTERNAL(void*, int) {
std::_Exit(86); }\n'
+ )
+ cc.create_shared(str(guard_library), [str(guard_source)])
+
+ if model == "conv2d_relu":
+ partitioned, data, expected = _conv2d_relu_partition()
+ expected_partition_count = 1
+ # Match the numerical tolerance of the existing TensorRT convolution
tests.
+ tolerance = 1e-3
+ else:
+ partitioned, parameters = _two_partitions(dtype)
+ data = np.linspace(-1.7, 1.3, 16, dtype=dtype).reshape(2, 2, 4)
+ expected = np.sin(parameters["second"] - data + parameters["first"]) *
parameters["third"]
+ expected_partition_count = 2
+ tolerance = 1e-2 if dtype == "float16" else 1e-5
+ np.savez(tmp_path / "reference.npz", data=data, expected=expected)
+
+ runner = textwrap.dedent(
+ """
+ import sys
+ import numpy as np
+ import tvm
+ from tvm import relax
+
+ artifact = tvm.runtime.load_module(sys.argv[1])
+ device = tvm.cpu() if sys.argv[5] == "cpu" else
tvm.cuda(int(sys.argv[3]))
+ tolerance = float(sys.argv[4])
+ vm = relax.VirtualMachine(artifact, device)
+ reference = np.load(sys.argv[2])
+ for data, expected in zip(reference["data"], reference["expected"]):
+ result = vm["main"](tvm.runtime.tensor(data, device))
+ assert result.device == device
+ np.testing.assert_allclose(result.numpy(), expected,
rtol=tolerance, atol=tolerance)
+ """
+ )
+ child_env = os.environ.copy()
+ child_env.pop("TVM_TENSORRT_CACHE_DIR", None)
+ previous_preload = child_env.get("LD_PRELOAD", "")
+ child_env["LD_PRELOAD"] = str(guard_library) + (
+ ":" + previous_preload if previous_preload else ""
+ )
+
+ def run_child(path):
+ return subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ runner,
+ str(path),
+ str(tmp_path / "reference.npz"),
+ str(device_id),
+ str(tolerance),
+ execution_device,
+ ],
+ env=child_env,
+ capture_output=True,
+ text=True,
+ timeout=120,
+ check=False,
+ )
+
+ def check():
+ # Negative control: prove the preload intercepts the existing lazy
builder.
+ lazy = _codegen(partitioned, eager=False)
+ assert len(_external_modules(lazy)) == expected_partition_count
+ lazy_path = tmp_path / "lazy.so"
+ tvm.compile(lazy, target).export_library(str(lazy_path))
+ result = run_child(lazy_path)
+ assert result.returncode == 86, result.stdout + result.stderr
+
+ eager = _codegen(partitioned, eager=True)
+ modules = _external_modules(eager)
+ assert len(modules) == expected_partition_count
+ for module in modules:
+ payload = _module_bytes(module)
+ # No VirtualMachine has been created in this process.
+ assert len(payload) > len(_legacy_bytes(module))
+ restored = _load_bytes(payload)
+ assert _module_bytes(restored) == payload
+
+ eager_path = tmp_path / "eager.so"
+ tvm.compile(eager, target).export_library(str(eager_path))
+ result = run_child(eager_path)
+ assert result.returncode == 0, result.stdout + result.stderr
+
+ module = modules[0]
+ constants_by_name = dict(eager.attrs["const_name_to_constant"])
+ constants = [constants_by_name[name] for name in
module["get_const_vars"]()]
+ # Embedded plans cannot accept either an INT8 flag or a captured
calibration count.
+ for use_int8 in ("1", "0"):
+ with monkeypatch.context() as calibration:
+ calibration.setenv("TVM_TENSORRT_USE_INT8", use_int8)
+ calibration.setenv("TENSORRT_NUM_CALI_INT8", "1")
+ loaded = _load_bytes(_module_bytes(module))
+ with pytest.raises(
+ tvm.error.InternalError, match="prebuilt TensorRT.*INT8
calibration"
+ ):
+ loaded["__init_" + loaded["get_symbol"]()](constants)
+
+ # A corrupt plan must fail deserialization instead of rebuilding from
the JSON.
+ payload = bytearray(_module_bytes(module))
+ plan_start = len(_legacy_bytes(module)) + 8 + 4 + 8
+ payload[plan_start : plan_start + 32] = b"\0" * 32
+ corrupt = _load_bytes(bytes(payload))
+ symbol = corrupt["get_symbol"]()
+ corrupt["__init_" + symbol](constants)
+ # Engine deserialization is deferred until the input device is known.
+ device = tvm.cpu() if execution_device == "cpu" else
tvm.cuda(device_id)
+ argument = tvm.runtime.tensor(np.zeros(data.shape[1:], dtype=dtype),
device)
+ output = tvm.runtime.tensor(np.zeros(expected.shape[1:], dtype=dtype),
device)
+ with pytest.raises(
+ tvm.error.InternalError, match="deserialize the embedded TensorRT
engine"
+ ):
+ corrupt[symbol](argument, output)
+
+ tvm.testing.run_with_gpu_lock(check)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
