From f22d1df04f582402431ba4e65127af16688fead1 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Tue, 18 Jun 2019 07:45:09 -0700 Subject: [PATCH] Fix a bug that constant folding doesn't check the returned value of OpKernel::Compute (#1243) --- include/onnxruntime/core/framework/op_kernel.h | 5 ++--- onnxruntime/core/optimizer/constant_folding.cc | 2 +- onnxruntime/test/optimizer/optimizer_test.cc | 3 ++- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/include/onnxruntime/core/framework/op_kernel.h b/include/onnxruntime/core/framework/op_kernel.h index 3afd667438..77ec9ed292 100644 --- a/include/onnxruntime/core/framework/op_kernel.h +++ b/include/onnxruntime/core/framework/op_kernel.h @@ -40,10 +40,9 @@ class OpKernel { return op_kernel_info_.GetKernelDef(); } - virtual Status Compute(OpKernelContext* context) const = 0; + virtual Status Compute(OpKernelContext* context) const ORT_MUST_USE_RESULT = 0; - virtual Status ComputeAsync(OpKernelContext*, - DoneCallback) const { + virtual Status ComputeAsync(OpKernelContext*, DoneCallback) const ORT_MUST_USE_RESULT { ORT_NOT_IMPLEMENTED(__FUNCTION__, " is not implemented"); } diff --git a/onnxruntime/core/optimizer/constant_folding.cc b/onnxruntime/core/optimizer/constant_folding.cc index db5cfc2d76..b6273a870b 100644 --- a/onnxruntime/core/optimizer/constant_folding.cc +++ b/onnxruntime/core/optimizer/constant_folding.cc @@ -49,7 +49,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level) auto* kernel = info.GetKernel(node->Index()); OpKernelContext op_kernel_context(&frame, kernel, ::onnxruntime::logging::LoggingManager::DefaultLogger()); - kernel->Compute(&op_kernel_context); + ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context)); std::vector fetches; frame.GetOutputs(fetches); diff --git a/onnxruntime/test/optimizer/optimizer_test.cc b/onnxruntime/test/optimizer/optimizer_test.cc index 020e322051..fce027650f 100644 --- a/onnxruntime/test/optimizer/optimizer_test.cc +++ b/onnxruntime/test/optimizer/optimizer_test.cc @@ -76,7 +76,8 @@ TEST(OptimizerTest, Basic) { OpKernelContext op_kernel_context(&frame, kernel, logger); - kernel->Compute(&op_kernel_context); + auto st = kernel->Compute(&op_kernel_context); + ASSERT_TRUE(st.IsOK()) << st.ErrorMessage(); std::vector fetches; frame.GetOutputs(fetches);