From 1183626694529ff3938e77370a43e006c6c06164 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Tue, 17 Jun 2025 11:48:58 -0700 Subject: [PATCH 1/2] ci: Fix shape tensor segfault --- src/instance_state.cc | 6 +++--- src/model_state.cc | 2 +- src/shape_tensor.cc | 2 +- src/shape_tensor.h | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/instance_state.cc b/src/instance_state.cc index d3e00b0..8cd2a57 100644 --- a/src/instance_state.cc +++ b/src/instance_state.cc @@ -3095,21 +3095,21 @@ ModelInstanceState::InitializeShapeInputBinding( : context.max_dims_[io_index].d[0]; context.max_shapes_[io_index] = ShapeTensor(); context.max_shapes_[io_index].SetDataFromShapeValues( - engine_->getProfileTensorValues( + engine_->getProfileTensorValuesV2( input_name.c_str(), profile_index, nvinfer1::OptProfileSelector::kMAX), input_datatype, context.nb_shape_values_); context.min_shapes_[io_index] = ShapeTensor(); context.min_shapes_[io_index].SetDataFromShapeValues( - engine_->getProfileTensorValues( + engine_->getProfileTensorValuesV2( input_name.c_str(), profile_index, nvinfer1::OptProfileSelector::kMIN), input_datatype, context.nb_shape_values_); context.opt_shapes_[io_index] = ShapeTensor(); context.opt_shapes_[io_index].SetDataFromShapeValues( - engine_->getProfileTensorValues( + engine_->getProfileTensorValuesV2( input_name.c_str(), profile_index, nvinfer1::OptProfileSelector::kOPT), input_datatype, context.nb_shape_values_); diff --git a/src/model_state.cc b/src/model_state.cc index bc72d67..f7eebac 100644 --- a/src/model_state.cc +++ b/src/model_state.cc @@ -708,7 +708,7 @@ ModelState::GetProfileMaxBatchSize( } } else { - const int32_t* max_shapes = engine->getProfileTensorValues( + const int64_t* max_shapes = engine->getProfileTensorValuesV2( tensor_name.c_str(), profile_index, nvinfer1::OptProfileSelector::kMAX); if (*max_batch_size > *max_shapes) { diff --git a/src/shape_tensor.cc b/src/shape_tensor.cc index 8cca781..781dcb8 100644 --- a/src/shape_tensor.cc +++ b/src/shape_tensor.cc @@ -88,7 +88,7 @@ ShapeTensor::SetDataFromBuffer( TRITONSERVER_Error* ShapeTensor::SetDataFromShapeValues( - const int32_t* shape_values, TRITONSERVER_DataType datatype, + const int64_t* shape_values, TRITONSERVER_DataType datatype, size_t nb_shape_values) { nb_shape_values_ = nb_shape_values; diff --git a/src/shape_tensor.h b/src/shape_tensor.h index 1d4fc1a..2b5b24c 100644 --- a/src/shape_tensor.h +++ b/src/shape_tensor.h @@ -51,7 +51,7 @@ class ShapeTensor { const char* input_name, bool support_batching, size_t total_batch_size); TRITONSERVER_Error* SetDataFromShapeValues( - const int32_t* shape_values, TRITONSERVER_DataType datatype, + const int64_t* shape_values, TRITONSERVER_DataType datatype, size_t nb_shape_values); int64_t GetDistance(const ShapeTensor& other, int64_t total_batch_size) const; From 4dd5a32257222c0edc027a905c83930f901fe917 Mon Sep 17 00:00:00 2001 From: Yingge He <157551214+yinggeh@users.noreply.github.com> Date: Wed, 18 Jun 2025 10:59:18 -0700 Subject: [PATCH 2/2] ci: Update SetDataFromShapeValues implementation after changing shape_values type (#113) --- src/shape_tensor.cc | 8 ++++---- src/shape_tensor.h | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/shape_tensor.cc b/src/shape_tensor.cc index 781dcb8..9b9ba0f 100644 --- a/src/shape_tensor.cc +++ b/src/shape_tensor.cc @@ -99,14 +99,14 @@ ShapeTensor::SetDataFromShapeValues( datatype_ = ShapeTensorDataType::INT32; data_.reset(new char[size_]); int32_t* data_ptr = reinterpret_cast(data_.get()); - std::memcpy(data_ptr, shape_values, size_); + for (size_t i = 0; i < nb_shape_values_; ++i) { + data_ptr[i] = static_cast(shape_values[i]); + } } else if (datatype == TRITONSERVER_DataType::TRITONSERVER_TYPE_INT64) { datatype_ = ShapeTensorDataType::INT64; data_.reset(new char[size_]); int64_t* data_ptr = reinterpret_cast(data_.get()); - for (size_t i = 0; i < nb_shape_values_; ++i) { - data_ptr[i] = static_cast(shape_values[i]); - } + std::memcpy(data_ptr, shape_values, size_); } else { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, diff --git a/src/shape_tensor.h b/src/shape_tensor.h index 2b5b24c..74acd1f 100644 --- a/src/shape_tensor.h +++ b/src/shape_tensor.h @@ -1,4 +1,4 @@ -// Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Redistribution and use in source and binary forms, with or without // modification, are permitted provided that the following conditions