gemini-code-assist[bot] commented on code in PR #19789:
URL: https://github.com/apache/tvm/pull/19789#discussion_r3418231087
##########
src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc:
##########
@@ -372,11 +384,20 @@ class TensorRTRuntime : public JSONRuntimeBase {
infile.close();
std::string serialized_engine;
LoadBinaryFromFile(path, &serialized_engine);
- // Deserialize engine
+ // Deserialize engine. TensorRT 10 dropped the trailing IPluginFactory*
argument and the runtime
+ // must outlive the engine, so it is owned by the cached
TensorRTEngineAndContext.
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger_);
TensorRTEngineAndContext engine_and_context;
+ engine_and_context.runtime = runtime;
engine_and_context.engine =
- runtime->deserializeCudaEngine(&serialized_engine[0],
serialized_engine.size(), nullptr);
+ runtime->deserializeCudaEngine(&serialized_engine[0],
serialized_engine.size());
Review Comment:
The `nvinfer1::createInferRuntime` call can return `nullptr` if
initialization fails. We should check `runtime` for `nullptr` before calling
`deserializeCudaEngine` to prevent a potential null pointer dereference.
```c
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger_);
if (runtime == nullptr) {
LOG(WARNING) << "Failed to create TensorRT runtime for
deserialization.";
return false;
}
TensorRTEngineAndContext engine_and_context;
engine_and_context.runtime = runtime;
engine_and_context.engine =
runtime->deserializeCudaEngine(&serialized_engine[0],
serialized_engine.size());
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc:
##########
@@ -410,31 +406,27 @@ class DenseOpConverter : public TensorRTOpConverter {
void Convert(TensorRTOpConverterParams* params) const {
auto input_tensor = params->inputs.at(0).tensor;
- auto input_dims = TrtDimsToVector(input_tensor->getDimensions());
- TVM_FFI_ICHECK(input_dims.size() > 0 && input_dims.size() <= 3);
- const size_t required_rank = TRT_HAS_IMPLICIT_BATCH(params) ? 3 : 4;
- const bool need_reshape_on_input = input_dims.size() != required_rank;
- if (need_reshape_on_input) {
- // Add dims of size 1 until rank is required_rank.
- std::vector<int> new_shape(input_dims);
- while (new_shape.size() < required_rank)
new_shape.insert(new_shape.end(), 1);
- input_tensor = Reshape(params, input_tensor, new_shape);
- }
- // Weights are in KC format.
+ // Weights are in KC (out_units x in_features) format.
TVM_FFI_ICHECK_EQ(params->inputs.at(1).weight_shape.size(), 2);
- const int num_units = params->inputs.at(1).weight_shape[0];
- const nvinfer1::DataType weight_type = params->inputs.at(1).weight.type;
- nvinfer1::Weights bias{weight_type, nullptr, 0};
- nvinfer1::IFullyConnectedLayer* fc_layer =
params->network->addFullyConnected(
- *input_tensor, num_units, params->inputs.at(1).weight, bias);
- TVM_FFI_ICHECK(fc_layer != nullptr);
- auto output_tensor = fc_layer->getOutput(0);
- if (need_reshape_on_input) {
- // Remove added dims.
- input_dims[input_dims.size() - 1] = num_units;
- output_tensor = Reshape(params, output_tensor, input_dims);
- }
- params->outputs.push_back(output_tensor);
+ // addMatrixMultiply requires the input to have at least 2 dimensions
(rows x K); the old
+ // FullyConnected path padded the rank, so guard explicitly now that it is
gone.
+ TVM_FFI_ICHECK_GE(input_tensor->getDimensions().nbDims, 2)
+ << "TensorRT dense expects an input of rank >= 2 (got "
+ << input_tensor->getDimensions().nbDims << ")";
+ // TensorRT 10 removed IFullyConnectedLayer/addFullyConnected. Implement
dense as a matrix
+ // multiply: out[.., O] = in[.., K] * weightįµ, with weight a constant of
shape [O, K].
+ // IMatrixMultiplyLayer contracts the last dim of `input` (K) with the
last dim of the
+ // transposed weight (also K) and broadcasts the remaining leading
dimensions, which matches
+ // nn.dense semantics for any input rank >= 2 without the rank-padding
reshape FC required.
+ auto* weight_tensor = params->network
+
->addConstant(VectorToTrtDims(params->inputs.at(1).weight_shape),
+ params->inputs.at(1).weight)
+ ->getOutput(0);
Review Comment:
The `addConstant` call can return `nullptr` if layer creation fails.
Chaining `->getOutput(0)` directly on it can cause a null pointer dereference.
We should check the returned layer for `nullptr` first.
```suggestion
auto* weight_layer = params->network->addConstant(
VectorToTrtDims(params->inputs.at(1).weight_shape),
params->inputs.at(1).weight);
TVM_FFI_ICHECK(weight_layer != nullptr);
auto* weight_tensor = weight_layer->getOutput(0);
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc:
##########
@@ -219,34 +231,31 @@ class TensorRTRuntime : public JSONRuntimeBase {
for (size_t i = 0; i < outputs_.size(); ++i) {
uint32_t eid = EntryID(outputs_[i]);
const std::string& name = engine_and_context.outputs[i];
- int binding_index = engine->getBindingIndex(name.c_str());
- TVM_FFI_ICHECK_NE(binding_index, -1);
+ void* device_ptr = nullptr;
if (data_entry_[eid]->device.device_type == kDLCUDA) {
- bindings[binding_index] = data_entry_[eid]->data;
+ device_ptr = data_entry_[eid]->data;
} else {
- auto device_buffer = GetOrAllocateDeviceBuffer(eid, binding_index);
- bindings[binding_index] = device_buffer->data;
+ auto device_buffer = GetOrAllocateDeviceBuffer(name, eid);
+ device_ptr = device_buffer->data;
}
+ TVM_FFI_ICHECK(context->setTensorAddress(name.c_str(), device_ptr));
}
-#if TRT_VERSION_GE(6, 0, 1)
- if (use_implicit_batch_) {
- TVM_FFI_ICHECK(context->execute(batch_size, bindings.data())) <<
"Running TensorRT failed.";
- } else {
- TVM_FFI_ICHECK(context->executeV2(bindings.data())) << "Running TensorRT
failed.";
- }
-#else
- TVM_FFI_ICHECK(context->execute(batch_size, bindings.data())) << "Running
TensorRT failed.";
-#endif
+ // Run on TVM's current CUDA stream so the engine is ordered after the
inputs produced upstream
+ // (and to avoid TensorRT's default-stream synchronization warning).
enqueueV3 is async-only in
+ // TRT10, so synchronize afterwards to preserve Run()'s blocking semantics.
+ const DLDevice& dev = data_entry_[input_var_eid_[0]]->device;
+ const int device_id = dev.device_type == kDLCUDA ? dev.device_id : 0;
Review Comment:
If `input_var_eid_` is empty, accessing `input_var_eid_[0]` will result in
out-of-bounds access. We should safely resolve the device ID by checking if
`input_var_eid_` is empty first, and falling back to the output device if
available.
```c
int device_id = 0;
if (!input_var_eid_.empty()) {
const DLDevice& dev = data_entry_[input_var_eid_[0]]->device;
device_id = dev.device_type == kDLCUDA ? dev.device_id : 0;
} else if (!outputs_.empty()) {
const DLDevice& dev = data_entry_[EntryID(outputs_[0])]->device;
device_id = dev.device_type == kDLCUDA ? dev.device_id : 0;
}
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_ops.cc:
##########
@@ -265,21 +270,22 @@ class Conv1DOpConverter : public TensorRTOpConverter {
nvinfer1::Weights bias{weight_type, nullptr, 0};
- auto conv_layer = params->network->addConvolution(*input_tensor, channels,
kernel_size,
-
params->inputs.at(1).weight, bias);
+ auto conv_layer = params->network->addConvolutionNd(*input_tensor,
channels, kernel_size,
+
params->inputs.at(1).weight, bias);
TVM_FFI_ICHECK(conv_layer != nullptr);
- conv_layer->setPadding(nvinfer1::DimsHW(static_cast<int>(padding[0]), 0));
+ conv_layer->setPaddingNd(nvinfer1::DimsHW(static_cast<int>(padding[0]),
0));
TVM_FFI_ICHECK_EQ(strides.size(), 1);
const auto trt_strides = nvinfer1::DimsHW(static_cast<int>(strides[0]), 1);
- conv_layer->setStride(trt_strides);
+ conv_layer->setStrideNd(trt_strides);
TVM_FFI_ICHECK_EQ(dilation.size(), 1);
const auto trt_dilation = nvinfer1::DimsHW(static_cast<int>(dilation[0]),
1);
- conv_layer->setDilation(trt_dilation);
+ conv_layer->setDilationNd(trt_dilation);
conv_layer->setNbGroups(groups);
input_tensor = conv_layer->getOutput(0);
- auto conv_output_dims = TrtDimsToVector(input_tensor->getDimensions());
- std::vector<int> back_shape = {0, 0};
+ // Drop the trailing unit dimension (NOW1 -> NOW); 0 copies the
corresponding input dimension,
+ // so the number of leading dims to keep matches the original input rank.
+ std::vector<int> back_shape(input_dims.size(), 0);
auto shuffle_back_layer = params->network->addShuffle(*input_tensor);
shuffle_back_layer->setReshapeDimensions(VectorToTrtDims(back_shape));
params->outputs.push_back(shuffle_back_layer->getOutput(0));
Review Comment:
The `addShuffle` call can return `nullptr` if layer creation fails. We
should check `shuffle_back_layer` for `nullptr` before calling
`setReshapeDimensions` to prevent a potential null pointer dereference.
```c
auto shuffle_back_layer = params->network->addShuffle(*input_tensor);
TVM_FFI_ICHECK(shuffle_back_layer != nullptr);
shuffle_back_layer->setReshapeDimensions(VectorToTrtDims(back_shape));
params->outputs.push_back(shuffle_back_layer->getOutput(0));
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_builder.cc:
##########
@@ -184,40 +168,48 @@ TensorRTEngineAndContext TensorRTBuilder::BuildEngine() {
LOG(INFO) << "config finishes setting up calibrator as INT8 mode ... ";
}
- // Add profiles.
- if (!use_implicit_batch_) {
- auto profile = builder_->createOptimizationProfile();
- for (int i = 0; i < network_->getNbInputs(); ++i) {
- auto name = network_->getInput(i)->getName();
- const uint32_t entry_id = entry_id_map_[name];
- std::vector<int64_t> shape(data_entry_[entry_id]->shape,
- data_entry_[entry_id]->shape +
data_entry_[entry_id]->ndim);
- auto dims = VectorToTrtDims(shape);
+ // 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).
+ auto profile = builder_->createOptimizationProfile();
+ for (int i = 0; i < network_->getNbInputs(); ++i) {
+ auto name = network_->getInput(i)->getName();
+ const uint32_t entry_id = entry_id_map_[name];
+ std::vector<int64_t> shape(data_entry_[entry_id]->shape,
+ data_entry_[entry_id]->shape +
data_entry_[entry_id]->ndim);
+ auto dims = VectorToTrtDims(shape);
- profile->setDimensions(name, nvinfer1::OptProfileSelector::kOPT, dims);
- profile->setDimensions(name, nvinfer1::OptProfileSelector::kMAX, dims);
- // Set minimum batch size to 1 when dynamic batching is used.
- if (network_->getInput(i)->getDimensions().nbDims >= 1 &&
- network_->getInput(i)->getDimensions().d[0] == -1) {
- dims.d[0] = 1;
- }
- profile->setDimensions(name, nvinfer1::OptProfileSelector::kMIN, dims);
+ profile->setDimensions(name, nvinfer1::OptProfileSelector::kOPT, dims);
+ profile->setDimensions(name, nvinfer1::OptProfileSelector::kMAX, dims);
+ // The network inputs are built with static shapes, so the profile must
match them exactly; only
+ // lower kMIN for a genuinely dynamic (-1) leading dimension.
+ if (network_->getInput(i)->getDimensions().nbDims >= 1 &&
+ network_->getInput(i)->getDimensions().d[0] == -1) {
+ dims.d[0] = 1;
}
- config_->addOptimizationProfile(profile);
+ profile->setDimensions(name, nvinfer1::OptProfileSelector::kMIN, dims);
}
- nvinfer1::ICudaEngine* engine = builder_->buildEngineWithConfig(*network_,
*config_);
-#else
- nvinfer1::ICudaEngine* engine = builder_->buildCudaEngine(*network_);
-#endif
- TVM_FFI_ICHECK_EQ(engine->getNbBindings(),
- network_input_names_.size() +
network_output_names_.size());
+ config_->addOptimizationProfile(profile);
+
+ // 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.";
+ }
Review Comment:
The `nvinfer1::createInferRuntime` call can return `nullptr` if
initialization fails. We should check `runtime` for `nullptr` before calling
`deserializeCudaEngine` to prevent a potential null pointer dereference.
```c
nvinfer1::IHostMemory* plan = builder_->buildSerializedNetwork(*network_,
*config_);
TVM_FFI_ICHECK(plan) << "Failed to build TensorRT serialized network.";
nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(*trt_logger_);
TVM_FFI_ICHECK(runtime != nullptr) << "Failed to create TensorRT runtime.";
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.";
}
```
##########
src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc:
##########
@@ -166,40 +176,42 @@ class TensorRTRuntime : public JSONRuntimeBase {
auto& engine_and_context = GetOrBuildEngine();
int batch_size = GetBatchSize();
if (batch_size == 0) return;
- auto engine = engine_and_context.engine;
auto context = engine_and_context.context;
- const int num_bindings = engine->getNbBindings();
- std::vector<void*> bindings(num_bindings, nullptr);
- std::vector<size_t> binding_sizes(num_bindings, 0);
+
+ // TensorRT 10 uses named-tensor I/O
(setInputShape/setTensorAddress/enqueueV3, no binding
+ // indices). Track input device pointers and per-sample element counts for
the INT8 calibrator.
+ std::vector<void*> input_bindings;
+ std::vector<size_t> input_binding_sizes;
+
// Setup input bindings.
for (size_t i = 0; i < input_nodes_.size(); ++i) {
auto nid = input_nodes_[i];
if (nodes_[nid].GetOpType() == "input") {
for (size_t j = 0; j < nodes_[nid].GetOpShape().size(); ++j) {
uint32_t eid = EntryID(nid, j);
const std::string name = nodes_[nid].GetOpName() + "_" +
std::to_string(j);
- int binding_index = engine->getBindingIndex(name.c_str());
- TVM_FFI_ICHECK_NE(binding_index, -1);
-#if TRT_VERSION_GE(6, 0, 1)
- if (!use_implicit_batch_) {
- std::vector<int64_t> shape(data_entry_[eid]->shape,
- data_entry_[eid]->shape +
data_entry_[eid]->ndim);
- auto dims = VectorToTrtDims(shape);
- TVM_FFI_ICHECK(context->setBindingDimensions(binding_index, dims));
- }
-#endif
+ std::vector<int64_t> shape(data_entry_[eid]->shape,
+ data_entry_[eid]->shape +
data_entry_[eid]->ndim);
+ auto dims = VectorToTrtDims(shape);
+ TVM_FFI_ICHECK(context->setInputShape(name.c_str(), dims));
+
+ void* device_ptr = nullptr;
if (data_entry_[eid]->device.device_type == kDLCUDA) {
- bindings[binding_index] = data_entry_[eid]->data;
+ device_ptr = data_entry_[eid]->data;
} else {
- auto device_buffer = GetOrAllocateDeviceBuffer(eid, binding_index);
+ auto device_buffer = GetOrAllocateDeviceBuffer(name, eid);
device_buffer.CopyFrom(data_entry_[eid]);
- bindings[binding_index] = device_buffer->data;
+ device_ptr = device_buffer->data;
}
+ TVM_FFI_ICHECK(context->setTensorAddress(name.c_str(), device_ptr));
- auto dims = engine->getBindingDimensions(binding_index);
+ // Per-sample element count (exclude the batch dimension d[0]); the
INT8 calibrator
+ // multiplies by the batch size itself when copying calibration
data, so including the
+ // batch dim here would over-read the device buffer by a factor of
batch_size.
int num_elements = 1;
- for (int i = 0; i < dims.nbDims; ++i) num_elements *= dims.d[i];
- binding_sizes[binding_index] = num_elements;
+ for (int k = 1; k < dims.nbDims; ++k) num_elements *= dims.d[k];
Review Comment:
If any dimension other than the batch dimension is dynamic (represented as
`-1` in TensorRT), `dims.d[k]` will be negative, causing `num_elements` to
overflow when cast to `size_t`. We should check that dimensions are
non-negative to prevent potential integer overflow or negative allocation sizes.
```c
int num_elements = 1;
for (int k = 1; k < dims.nbDims; ++k) {
TVM_FFI_ICHECK_GE(dims.d[k], 0) << "Dynamic dimensions are not
supported for INT8 calibration.";
num_elements *= dims.d[k];
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
