Strip invalid dim_param and dim_value values out. Allow re-use in event of shape mismatch if buffer is large enough (#1439)

* Remove invalid dim_param and dim_value values when creating a NodeArg.
* Allow re-use of a large enough buffer if there's a shape mismatch.
* Update handling in python to treat unset dimension the same as a dim_param (equivalent to None).
* Fix GetTensorShapeFromTensorShapeProto to handle neither dim_param and dim_value being set.

onnxruntime/core/framework/execution_frame.cc

@@ -316,14 +316,25 @@ Status ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer(OrtValue& ort_valu
   OrtValue& ort_value_reuse = GetMutableMLValue(ort_value_index_reuse);
 
   auto* reuse_tensor = ort_value_reuse.GetMutable<Tensor>();
+  auto buffer_num_elements = reuse_tensor->Shape().Size();
+  auto required_num_elements = shape.Size();
 
   // check number of elements matches. shape may not be an exact match (e.g. Reshape op)
-  if (reuse_tensor->Shape().Size() != shape.Size()) {
+  if (buffer_num_elements != required_num_elements) {
     // could be an allocation planner bug (less likely) or the model incorrectly uses something like 'None'
-    // as a dim_param in multiple places making the planner think those shapes are equal.
-    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
-                           "Shape mismatch attempting to re-use buffer. ", reuse_tensor->Shape(), " != ", shape,
-                           ". Validate usage of dim_param in shapes in the model.");
+    // as a dim_param, or -1 in dim_value in multiple places making the planner think those shapes are equal.
+    auto message = onnxruntime::MakeString(
+        "Shape mismatch attempting to re-use buffer. ",
+        reuse_tensor->Shape(), " != ", shape,
+        ". Validate usage of dim_value (values should be > 0) and "
+        "dim_param (all values with the same string should equate to the same size) in shapes in the model.");
+
+    // be generous and use the buffer if it's large enough. log a warning though as it indicates a bad model
+    if (buffer_num_elements >= required_num_elements) {
+      LOGS_DEFAULT(WARNING) << message;
+    } else {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, message);
+    }
   }
 
   void* reuse_buffer = reuse_tensor->MutableDataRaw();

onnxruntime/core/framework/tensorprotoutils.cc

@@ -278,8 +278,8 @@ TensorShape GetTensorShapeFromTensorShapeProto(const ONNX_NAMESPACE::TensorShape
   const auto& dims = tensor_shape_proto.dim();
   std::vector<int64_t> tensor_shape_vec(static_cast<size_t>(dims.size()));
   for (int i = 0; i < dims.size(); ++i) {
-    tensor_shape_vec[i] = dims[i].has_dim_param() ? -1 /* symbolic dimensions are represented as -1 in onnxruntime*/
-                                                  : dims[i].dim_value();
+    tensor_shape_vec[i] = dims[i].has_dim_value() ? dims[i].dim_value()
+                                                  : -1; /* symbolic dimensions are represented as -1 in onnxruntime*/
   }
   return TensorShape(std::move(tensor_shape_vec));
 }

onnxruntime/core/graph/graph.cc

@@ -51,13 +51,34 @@ static bool GraphLoadedFromModelFile(const GraphProto* graph_proto) {
                          graph_proto->value_info_size() != 0);
 }
 
-NodeArg::NodeArg(const std::string& name,
-                 const TypeProto* p_node_arg_type) {
+// there are some known invalid usages of dim_param and dim_value. remove them from the TypeProto so that
+// they don't affect shape inferencing or the allocation planner
+static void RemoveInvalidValues(ONNX_NAMESPACE::TypeProto& type) {
+  if (type.has_tensor_type() && type.tensor_type().has_shape()) {
+    auto* shape = type.mutable_tensor_type()->mutable_shape();
+    for (int i = 0, end = shape->dim_size(); i < end; ++i) {
+      auto& dim = *shape->mutable_dim(i);
+      if (dim.has_dim_param()) {
+        auto dim_param = dim.dim_param();
+        if (dim_param.empty() || dim_param == "None") {
+          dim.clear_dim_param();
+        }
+      } else if (dim.has_dim_value()) {
+        if (dim.dim_value() < 0) {
+          dim.clear_dim_value();
+        }
+      }
+    }
+  }
+}
+
+NodeArg::NodeArg(const std::string& name, const TypeProto* p_node_arg_type) {
   node_arg_info_.set_name(name);
   // If the name is empty, it means the arg does not exist.
   exists_ = !(name.empty());
   if (nullptr != p_node_arg_type) {
     (*node_arg_info_.mutable_type()) = *p_node_arg_type;
+    RemoveInvalidValues(*node_arg_info_.mutable_type());
     type_ = DataTypeUtils::ToType(node_arg_info_.type());
   } else {
     type_ = nullptr;

onnxruntime/core/providers/cpu/controlflow/loop.cc

@@ -412,19 +412,21 @@ Status LoopImpl::Execute(FeedsFetchesManager* ffm, const FeedsFetchesManager* ca
     auto& graph_outputs = subgraph_.GetOutputs();
 
     for (int i = num_loop_carried_vars_; i < num_outputs_; ++i) {
-      std::vector<int64_t> output_dims;
-      output_dims.push_back(0);  // num iterations is first dim
-
       // get shape from subgraph output if possible to attempt to have the correct rank
       auto* graph_output = graph_outputs.at(i + 1);  // + 1 as first subgraph output is condition value
       auto* graph_output_shape = graph_output->Shape();
 
-      if (graph_output_shape) {
-        output_dims.reserve(graph_output_shape->dim_size() + 1);
+      std::vector<int64_t> output_dims;
+      output_dims.reserve((graph_output_shape ? graph_output_shape->dim_size() : 0) + 1);
+      output_dims.push_back(0);  // num iterations is first dim
 
+      if (graph_output_shape) {
         const auto& tensor_shape = onnxruntime::utils::GetTensorShapeFromTensorShapeProto(*graph_output_shape);
         const auto& dims = tensor_shape.GetDims();
-        std::copy(dims.cbegin(), dims.cend(), std::back_inserter(output_dims));
+
+        // copy to output dims and use 0 for any symbolic dim
+        std::for_each(dims.cbegin(), dims.cend(),
+                      [&output_dims](const int64_t dim) { output_dims.push_back(dim < 0 ? 0 : dim); });
       } else {
         // TODO: We could try and call ExecuteGraph to get the output shape from fetches so the rank is correct,
         // however that could still fail as we would potentially be passing in invalid data.

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -21,7 +21,6 @@
 #define BACKEND_OPENMP ""
 #endif
 
-
 #if USE_MKLDNN
 #define BACKEND_MKLDNN "-MKL-DNN"
 #include "core/providers/mkldnn/mkldnn_execution_provider.h"
@@ -48,8 +47,6 @@
 #define BACKEND_OPENVINO ""
 #endif
 
-
-
 #if USE_OPENBLAS
 #define BACKEND_OPENBLAS "-OPENBLAS"
 #else
@@ -465,8 +462,11 @@ including arg name, arg type (contains both type and shape).)pbdoc")
                 if (shape->dim(i).has_dim_value()) {
                   res << shape->dim(i).dim_value();
                 } else if (shape->dim(i).has_dim_param()) {
+                  res << "'" << shape->dim(i).dim_param() << "'";
+                } else {
                   res << "None";
                 }
+
                 if (i < shape->dim_size() - 1) {
                   res << ", ";
                 }
@@ -491,6 +491,8 @@ including arg name, arg type (contains both type and shape).)pbdoc")
               if (shape->dim(i).has_dim_value()) {
                 arr[i] = py::cast(shape->dim(i).dim_value());
               } else if (shape->dim(i).has_dim_param()) {
+                arr[i] = py::cast(shape->dim(i).dim_param());
+              } else {
                 arr[i] = py::none();
               }
             }