gemini-code-assist[bot] commented on code in PR #19792:
URL: https://github.com/apache/tvm/pull/19792#discussion_r3421040005
##########
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);
+ auto* matmul_layer = params->network->addMatrixMultiply(
+ *input_tensor, nvinfer1::MatrixOperation::kNONE, *weight_tensor,
+ nvinfer1::MatrixOperation::kTRANSPOSE);
+ TVM_FFI_ICHECK(matmul_layer != nullptr);
+ params->outputs.push_back(matmul_layer->getOutput(0));
Review Comment:
The `addMatrixMultiply` function requires the input tensor to have at least
2 dimensions. However, Relax `nn.dense` allows a 1D input tensor of shape
`[in_features]`. Guarding with `TVM_FFI_ICHECK_GE(..., 2)` will cause valid
Relax models with 1D dense inputs to fail compilation. We should instead handle
1D inputs by temporarily reshaping them to 2D (`[1, K]`) before the matrix
multiplication, and then reshaping the output back to 1D (`[O]`).
```c
auto input_tensor = params->inputs.at(0).tensor;
auto input_dims = input_tensor->getDimensions();
bool need_reshape = input_dims.nbDims == 1;
if (need_reshape) {
std::vector<int> new_shape = {1, input_dims.d[0]};
input_tensor = Reshape(params, input_tensor, new_shape);
}
TVM_FFI_ICHECK_EQ(params->inputs.at(1).weight_shape.size(), 2);
auto* weight_tensor = params->network
->addConstant(VectorToTrtDims(params->inputs.at(1).weight_shape),
params->inputs.at(1).weight)
->getOutput(0);
auto* matmul_layer = params->network->addMatrixMultiply(
*input_tensor, nvinfer1::MatrixOperation::kNONE, *weight_tensor,
nvinfer1::MatrixOperation::kTRANSPOSE);
TVM_FFI_ICHECK(matmul_layer != nullptr);
auto* output_tensor = matmul_layer->getOutput(0);
if (need_reshape) {
std::vector<int> back_shape =
{static_cast<int>(params->inputs.at(1).weight_shape[0])};
output_tensor = Reshape(params, output_tensor, back_shape);
}
params->outputs.push_back(output_tensor);
```
##########
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;
Review Comment:
Accessing `input_var_eid_[0]` directly without checking if `input_var_eid_`
is empty is unsafe. If a subgraph has zero inputs (e.g., a nullary function or
a subgraph containing only bound parameters/constants), `input_var_eid_` will
be empty, leading to an out-of-bounds access and crash. Since a JSON runtime
subgraph must have at least one output, we can safely fall back to
`outputs_[0]` to retrieve the device.
```suggestion
const DLDevice& dev = !input_var_eid_.empty() ?
data_entry_[input_var_eid_[0]]->device :
data_entry_[EntryID(outputs_[0])]->device;
```
--
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]
