post-25.06: Update default branch (#114)

* ci: Fix shape tensor segfault

* ci: Update SetDataFromShapeValues implementation after changing shape_values type (#113)

---------

Co-authored-by: Yingge He <[email protected]>

Author: mc-nv

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_);

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) {

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;
@@ -99,14 +99,14 @@ ShapeTensor::SetDataFromShapeValues(
     datatype_ = ShapeTensorDataType::INT32;
     data_.reset(new char[size_]);
     int32_t* data_ptr = reinterpret_cast<int32_t*>(data_.get());
-    std::memcpy(data_ptr, shape_values, size_);
+    for (size_t i = 0; i < nb_shape_values_; ++i) {
+      data_ptr[i] = static_cast<int32_t>(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<int64_t*>(data_.get());
-    for (size_t i = 0; i < nb_shape_values_; ++i) {
-      data_ptr[i] = static_cast<int64_t>(shape_values[i]);
-    }
+    std::memcpy(data_ptr, shape_values, size_);
   } else {
     return TRITONSERVER_ErrorNew(
         TRITONSERVER_ERROR_INVALID_ARG,

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
@@ -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;