Output Tensor Shape Validation b/w ONNX inference and ORT (#7252)

* Adding Output Shape Validation for ORT-CPU execution flow

* Skipping validation check in-case output is not a tensor. Fixed conv_transpose test. Ignoring pad and reduction test

* Comparison b/w signed and un-signed int. Removed const for a primitive variable

* Commented the un-used test function signature

* Removed exception instead logging warning. Because there are lots of ORT tests which are failing because of this validation

* Fixed warning condition and test

* Fixed test and addressed comment on the PR

* Output shape verification will happen only for final output nodes of the model

* Changed variable name from camel case to underscore style

* Enable the tests as the validation failure will now logs warning instead of throwing an exception

* Adding Output Shape Validation for ORT-CPU execution flow

* Resolve merge conflict

* Comparison b/w signed and un-signed int. Removed const for a primitive variable

* Commented the un-used test function signature

* Removed exception instead logging warning. Because there are lots of ORT tests which are failing because of this validation

* Fixed warning condition and test

* Fixed test and addressed comment on the PR

* Output shape verification will happen only for final output nodes of the model

* Changed variable name from camel case to underscore style

* Enable the tests as the validation failure will now logs warning instead of throwing an exception

* Remove duplicate function "GetLogger()"

Remove duplicate function "GetLogger()"

* Fixed typo in method name TestConvTransposeOpInitializer

Fixed typo in method name "TestConvTransposeOpInitializer"

Co-authored-by: Sumit Agarwal <sumitagarwal@microsoft.com>

onnxruntime/core/framework/execution_frame.cc

@@ -138,10 +138,10 @@ OrtValue* IExecutionFrame::GetMutableNodeInputOrOutputMLValue(int index) {
 // This method is not thread safe!
 // Return S_OK and nullptr if index map to an value that is an unused optional input/output
 
-Status IExecutionFrame::GetOrCreateNodeOutputMLValue(int index, const TensorShape* shape, OrtValue*& p_ort_value,
-                                                     size_t nnz) {
+Status IExecutionFrame::GetOrCreateNodeOutputMLValue(const int output_index, int output_arg_index, const TensorShape* shape, OrtValue*& p_ort_value,
+                                                     const Node& node, size_t nnz) {
   auto status = Status::OK();
-  int ort_value_idx = GetNodeIdxToMLValueIdx(index);
+  int ort_value_idx = GetNodeIdxToMLValueIdx(output_arg_index);
 
   // return nullptr if it is optional
   if (ort_value_idx == NodeIndexInfo::kInvalidEntry) {
@@ -158,6 +158,9 @@ Status IExecutionFrame::GetOrCreateNodeOutputMLValue(int index, const TensorShap
                     " Requested shape:", shape ? shape->ToString() : "null");
       }
     } else {
+      if (this->IsOutput(ort_value_idx)) {
+        VerifyOutputSizes(output_index, node, shape);
+      }
       status = CreateNodeOutputMLValueImpl(*p_ort_value, ort_value_idx, shape, nnz);
     }
   }
@@ -165,6 +168,51 @@ Status IExecutionFrame::GetOrCreateNodeOutputMLValue(int index, const TensorShap
   return status;
 }
 
+void IExecutionFrame::VerifyOutputSizes(int output_index, const Node& node,
+                                       const TensorShape* output_shape) {
+  ProtoHelperNodeContext protoContext(node);
+  OpNodeProtoHelper<ProtoHelperNodeContext> info(&protoContext);
+
+  const onnx::TypeProto* outputproto = info.GetOutputType(output_index);
+
+  if (outputproto == nullptr) {
+    return;
+  }
+
+  if (outputproto->value_case() != onnx::TypeProto::kTensorType) {
+    return;
+  }
+  const auto& tensortype = outputproto->tensor_type();
+  if (tensortype.has_shape()) {
+    const auto& shape = tensortype.shape();
+
+    int actual_num_dims = static_cast<int>(output_shape->NumDimensions());
+    int expected_num_dims = shape.dim_size();
+    if (actual_num_dims != expected_num_dims) {
+      if (GetLogger() != nullptr) {
+        LOGS(*GetLogger(), WARNING) << "Number of dimension(" << actual_num_dims << ") of output shape didn't match "
+                                               << "with model's expected number("<< expected_num_dims <<") of dimension of output shape";
+      }
+      return;
+    }
+
+    for (uint32_t output_dim = 0, end = actual_num_dims; output_dim < end; ++output_dim) {
+      const auto dim = shape.dim(output_dim);
+      
+      if (dim.has_dim_value()) {
+        int64_t expected_size = dim.dim_value();
+        int64_t actual_size = (*output_shape)[output_dim];
+        if (expected_size != actual_size) {
+          if (GetLogger() != NULL) {
+            LOGS(*GetLogger(), WARNING) << "Actual dimension(" << actual_size << ") didn't match with expected dimension ("
+                                                   << expected_size << ") for outputProtoIndex: " << output_index;
+          }
+          return;
+        }
+      }
+    }
+  }
+}
 bool IExecutionFrame::TryGetInferredShape(int /*index*/, TensorShape& /*shape*/) const {
   // By default, there is not information about inferred shape, so this default
   // implementation always returns false. The derived class of IExecutionFrame
@@ -687,6 +735,10 @@ Status ExecutionFrame::AllocateAsPerAllocationPlan(OrtValue& ort_value, int ort_
   }
 }
 
+const logging::Logger* ExecutionFrame::GetLogger() {
+  return &(session_state_.Logger());
+}
+
 AllocatorPtr ExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {
   return session_state_.GetAllocator(info);
 }

onnxruntime/core/framework/execution_frame.h

@@ -61,7 +61,8 @@ class IExecutionFrame {
   // This method is not thread safe!
   // Return S_OK and nullptr if index map to an value that is an unused optional input/output
   // Shape is required for tensors but not traditional ML values.
-  Status GetOrCreateNodeOutputMLValue(int index, const TensorShape* shape, OrtValue*& p_ort_value, size_t nnz = 0);
+  Status GetOrCreateNodeOutputMLValue(const int index, int output_arg_index, const TensorShape* shape, 
+                                      OrtValue*& p_ort_value, const Node& node, size_t nnz = 0);
 
   // This function try retrieve the inferred shapes for the given NodeArg index.
   // If the retrival is sucessful, this function returns true and false otherwise.
@@ -96,6 +97,11 @@ class IExecutionFrame {
     return all_values_[ort_value_index];
   }
 
+  void VerifyOutputSizes(int output_index, const onnxruntime::Node& node,
+                         const onnxruntime::TensorShape* output_shape);
+
+  virtual const logging::Logger* GetLogger() = 0;
+
   virtual AllocatorPtr GetAllocatorImpl(const OrtMemoryInfo& info) const = 0;
 
   virtual Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape,
@@ -174,6 +180,7 @@ class ExecutionFrame final : public IExecutionFrame {
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(ExecutionFrame);
 
+  const logging::Logger* GetLogger() override;
   AllocatorPtr GetAllocatorImpl(const OrtMemoryInfo& info) const override;
   Status ReleaseMLValueImpl(int ort_value_idx) override;
   Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) override;

onnxruntime/core/framework/op_kernel.cc

@@ -72,7 +72,7 @@ OrtValue* OpKernelContext::OutputMLValue(int index, const TensorShape& shape, si
   //I believe it's a false alarm.
 
   OrtValue* p_ml_value = nullptr;
-  Status status = execution_frame_->GetOrCreateNodeOutputMLValue(GetOutputArgIndex(index), &shape, p_ml_value, nnz);
+  Status status = execution_frame_->GetOrCreateNodeOutputMLValue(index, GetOutputArgIndex(index), &shape, p_ml_value, kernel_->Node(), nnz);
   ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
   return p_ml_value;
 }
@@ -134,7 +134,7 @@ Fence_t OpKernelContext::OutputFence(int index) const {
 OrtValue* OpKernelContext::GetOrCreateOutputMLValue(int index) {
   auto output_arg_index = GetOutputArgIndex(index);
   OrtValue* value = nullptr;
-  auto status = execution_frame_->GetOrCreateNodeOutputMLValue(output_arg_index, nullptr, value);
+  auto status = execution_frame_->GetOrCreateNodeOutputMLValue(index, output_arg_index, nullptr, value, kernel_->Node());
   ORT_ENFORCE(status.IsOK(), status.ErrorMessage());
   return value;
 }

onnxruntime/core/optimizer/optimizer_execution_frame.cc

@@ -120,6 +120,10 @@ OptimizerExecutionFrame::OptimizerExecutionFrame(const Info& info,
   Init(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), fetches);
 }
 
+const logging::Logger* OptimizerExecutionFrame::GetLogger() {
+  return nullptr;
+}
+
 AllocatorPtr OptimizerExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {
   return info_.GetAllocator(info);
 }

onnxruntime/core/optimizer/optimizer_execution_frame.h

@@ -88,6 +88,8 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OptimizerExecutionFrame);
 
+  const logging::Logger* GetLogger() override;
+
   AllocatorPtr GetAllocatorImpl(const OrtMemoryInfo& info) const override;
 
   Status CreateNodeOutputMLValueImpl(OrtValue& ort_value, int ort_value_idx, const TensorShape* shape, size_t nnz) override;

onnxruntime/test/framework/execution_frame_test.cc

@@ -105,6 +105,58 @@ TEST_F(ExecutionFrameTest, TensorAllocationTest) {
   ASSERT_EQ(tensor2->template Data<float>(), p_tensor->template Data<float>());
 }
 
+TEST_F(ExecutionFrameTest, OutputShapeValidationTest) {
+  onnxruntime::Model model("test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), 
+      {{kOnnxDomain, 12}}, {}, DefaultLoggingManager().DefaultLogger());
+  onnxruntime::Graph& graph = model.MainGraph();
+  TypeProto tensor_float;
+  tensor_float.mutable_tensor_type()->set_elem_type(TensorProto_DataType_FLOAT);
+  onnxruntime::NodeArg input_def("X", &tensor_float), output_def("Y", &tensor_float);
+
+  onnx::TensorShapeProto new_shape;
+  new_shape.add_dim()->set_dim_value(2);
+  new_shape.add_dim()->set_dim_value(3);
+  output_def.SetShape(new_shape);
+
+  onnxruntime::Node* node = &graph.AddNode("node1", "Relu", "Relu operator", ArgMap{&input_def}, ArgMap{&output_def});
+  node->SetExecutionProviderType(kCpuExecutionProvider);
+  ASSERT_STATUS_OK(graph.Resolve());
+
+  auto cpu_xp = CreateCPUExecutionProvider();
+  auto xp_typ = cpu_xp->Type();
+  ExecutionProviders execution_providers;
+  execution_providers.Add(xp_typ, std::move(cpu_xp));
+  KernelRegistryManager kernel_registry_manager;
+  ASSERT_STATUS_OK(kernel_registry_manager.RegisterKernels(execution_providers));
+
+  DataTransferManager dtm;
+  profiling::Profiler profiler;
+  SessionState state(graph, execution_providers, true, &tp_, nullptr, dtm,
+                     DefaultLoggingManager().DefaultLogger(), profiler);
+
+  node->SetExecutionProviderType(xp_typ);
+
+  ASSERT_STATUS_OK(state.FinalizeSessionState(ORT_TSTR(""), kernel_registry_manager));
+
+  vector<OrtValue> outputs;
+  ExecutionFrame frame({}, {}, {}, outputs, {}, state);
+
+  int start_index = frame.GetNodeOffset(node->Index());
+  ASSERT_EQ(start_index, 0);
+  TensorShape actual_shape_same_as_input(std::vector<int64_t>{2, 3});
+  TensorShape actual_shape_diff_from_input(std::vector<int64_t>{2, 9});
+  
+  OrtValue* p_ml_value = frame.GetMutableNodeInputOrOutputMLValue(0);
+  ASSERT_TRUE(p_ml_value != nullptr);
+
+  // Calling the method with correct shape. It should work without any warnings.
+  ASSERT_STATUS_OK(frame.GetOrCreateNodeOutputMLValue(int(node->Index()), 1, &actual_shape_same_as_input, p_ml_value, *node, size_t(0)));
+
+  frame.ReleaseMLValue(1);
+  // Calling the method with in-correct shape. It should work but this time it should display a warning message.
+  ASSERT_STATUS_OK(frame.GetOrCreateNodeOutputMLValue(int(node->Index()), 1, &actual_shape_diff_from_input, p_ml_value, *node, size_t(0)));
+}
+
 TEST_F(ExecutionFrameTest, FeedInDataTest) {
   onnxruntime::Model model("test", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
                            std::unordered_map<std::string, int>{{"", 10}}, {},

onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc

@@ -29,7 +29,7 @@ void TestConvTransposeOpInitializer(const ConvTransposeOpAttributes& attributes,
                                     OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess,
                                     const std::string& err_str = "",
                                     const std::unordered_set<std::string>& excluded_provider_types = {kTensorrtExecutionProvider}) {
-  OpTester test("ConvTranspose");
+  OpTester test("ConvTranspose", 11);
   test.AddAttribute("kernel_shape", attributes.kernel_shape);
   test.AddAttribute("group", attributes.group);
 

onnxruntime/test/providers/cpu/tensor/pad_test.cc

@@ -678,6 +678,22 @@ TYPED_TEST(PadOpTest, Pad_Constant_DimWithZeroInput) {
                                   {2, 2, 2},
                                   {T(1), T(1), T(1), T(1), T(1), T(1), T(1), T(1)});
 }
+// Added output shape verification b/w the output shape generated by operator specific ONNX inference and
+// the output shape generated by operator specific ORT implementation. After adding this verification, 
+// this test logs warning as validation fails for 2 data types out of 8 data types i.e. Float and Double.
+// Reason:
+//  Pad ORT implementation output shape does not match with Pad ONNX inference function output shape.
+//  
+// For Float and Double this test gets executed for 2 different opset version, 10 and 11. Specifically this 
+// test is failing for opset version 10. 
+//  Investigation Analysis: Different ONNX inference class/method gets executed per opset version. Main difference b/w the 2 
+//          pad operator ONNX inference class/method is:
+//              Older Pad operator ONNX inference: Accepts "pads and values" as attribute.
+//              Newer Pad operator ONNX inference: Accetps "pads and values" as input.
+//          For newer version, "pads & values" fields have not been added as initializer, thus instead of shape 
+//          inference, rank inference gets triggered. Whereas, in older version shape inference gets executed 
+//          as "pads & values" fields have been added as attribute.
+//      In order to remove the warning, shape inference methods needs to be fixed.
 
 TYPED_TEST(PadOpTest, Pad_Edge_DimWithZeroInput) {
   using T = TypeParam;