From 31aff686e0f73d12fbc6b5313e439597f9798f75 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Tue, 1 Oct 2019 16:40:26 -0700 Subject: [PATCH] Register opset-11 CPU kernel for 'If' op (#1948) * Initial commit * Update * Update * Update * Update * PR comments --- .../core/providers/cpu/controlflow/if.cc | 19 ++++-- .../providers/cpu/cpu_execution_provider.cc | 6 +- .../test/providers/cpu/controlflow/if_test.cc | 64 +++++++++++++++---- 3 files changed, 69 insertions(+), 20 deletions(-) diff --git a/onnxruntime/core/providers/cpu/controlflow/if.cc b/onnxruntime/core/providers/cpu/controlflow/if.cc index bad6fdb05b..a92f6e3b38 100644 --- a/onnxruntime/core/providers/cpu/controlflow/if.cc +++ b/onnxruntime/core/providers/cpu/controlflow/if.cc @@ -46,12 +46,21 @@ ONNX_OPERATOR_SET_SCHEMA( .TypeConstraint("B", {"tensor(bool)"}, "Only bool")); */ +ONNX_CPU_OPERATOR_VERSIONED_KERNEL(If, + 1, 10, + KernelDefBuilder() + .TypeConstraint("B", DataTypeImpl::GetTensorType()) + .TypeConstraint("V", DataTypeImpl::AllTensorTypes()), + If); + +// output shape rules requiring the output shapes of the 'THEN' and 'ELSE' +// branches to be the same were relaxed in opset-11 ONNX_CPU_OPERATOR_KERNEL(If, - 1, - KernelDefBuilder() - .TypeConstraint("B", DataTypeImpl::GetTensorType()) - .TypeConstraint("V", DataTypeImpl::AllTensorTypes()), - If); + 11, + KernelDefBuilder() + .TypeConstraint("B", DataTypeImpl::GetTensorType()) + .TypeConstraint("V", DataTypeImpl::AllTensorTypes()), + If); struct If::Info { Info(const onnxruntime::Node& node, const GraphViewer& subgraph_in) : subgraph(subgraph_in) { diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index 203385079b..f1a068c1f7 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -212,7 +212,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, bool, Expand); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, MLFloat16, Expand); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, 8, Scan); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, If); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, If); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, Loop); // Opset 9 @@ -357,6 +357,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Ma class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, LpPool); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Conv); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ConvTranspose); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, If); void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { @@ -554,7 +555,7 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, // Opset 9 @@ -699,6 +700,7 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/test/providers/cpu/controlflow/if_test.cc b/onnxruntime/test/providers/cpu/controlflow/if_test.cc index 59a23a3695..a3193c57a3 100644 --- a/onnxruntime/test/providers/cpu/controlflow/if_test.cc +++ b/onnxruntime/test/providers/cpu/controlflow/if_test.cc @@ -44,7 +44,8 @@ static const ONNX_NAMESPACE::GraphProto CreateSubgraph(bool then_branch, const R class IfOpTester : public OpTester { public: - IfOpTester(const RunOptions& options) : OpTester("If"), options_{options} { + IfOpTester(const RunOptions& options, int opset_version = 10) : + OpTester("If", opset_version), options_{options}, opset_version_(opset_version) { } protected: @@ -73,7 +74,18 @@ class IfOpTester : public OpTester { inputs = {split_input}; outputs = {&split_out_0, &split_out_1}; - graph.AddNode("split", "Split", "Split into 2", inputs, outputs); + auto& split_node = graph.AddNode("split", "Split", "Split into 2", inputs, outputs); + if (opset_version_ > 10) { + AttributeProto attr_proto; + attr_proto.set_name("split"); + attr_proto.set_type(AttributeProto_AttributeType_INTS); + + auto* split_attribute = attr_proto.mutable_ints(); + *split_attribute->Add() = 1; // split "unevenly" to create different shapes across the "then" and "else" branches + *split_attribute->Add() = 2; + + split_node.AddAttribute("split", attr_proto); + } } // add If node @@ -99,13 +111,14 @@ class IfOpTester : public OpTester { private: RunOptions options_; + int opset_version_; }; -/* Subgraphs looks like this. All inputs come from outer scope so we just + /* Subgraphs looks like this. All inputs come from outer scope so we just create a NodeArg with the input name. The numbers in [] are the values the tests are expected to produce as output from each node. -THEN branch +THEN branch (all opset versions) split_out_0 if_input_0 [1] \ | [1] \ | @@ -113,13 +126,21 @@ THEN branch | add_out_0 [2] -ELSE branch +ELSE branch (opset 10 and below) split_out_1 if_input_0 [1] \ | [10] \ | \------[Add] | add_out_1 [11] + +ELSE branch (opset 11 and above) + split_out_1 if_input_0 [1] + \ | + [10, 10] \ | + \------[Add] + | + add_out_1 [11, 11] */ static const ONNX_NAMESPACE::GraphProto CreateSubgraph(bool then_branch, const RunOptions& options) { @@ -185,8 +206,9 @@ void RunTest(bool condition_value, RunOptions options, bool is_tensorrt_supported = true, OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess, - const std::string& failure_message = "") { - IfOpTester test{options}; + const std::string& failure_message = "", + int opset_version = 10) { + IfOpTester test{options, opset_version}; test.AddShapeToTensorData(options.include_dim_values_in_main_graph, options.symbolic_dim_value_in_main_graph); @@ -196,18 +218,26 @@ void RunTest(bool condition_value, // it's outputs are 1:1 with the graph outputs. // simple tensor that we split into 2, and use one output for the 'then' and one for the 'else' branch in the If - test.AddInput("split_input", {2}, {1.f, 10.f}); + if (opset_version != 11) { + test.AddInput("split_input", {2}, {1.f, 10.f}); + } else { + // in the opset 11 test case, we are going to split "unevenly", to create different shapes in "then" and "else" branches + test.AddInput("split_input", {3}, {1.f, 10.f, 10.f}); + } // graph input to specify which branch to take test.AddInput("if_cond", {1}, {condition_value}); test.AddInput("if_graph_input_0", {1}, {1.f}); - std::vector output_shape{1}; - if (condition_value) { - test.AddOutput("if_out_0", output_shape, {2.f}); - } else { - test.AddOutput("if_out_0", output_shape, {11.f}); + if (opset_version != 11 && condition_value) { + test.AddOutput("if_out_0", {1}, {2.f}); + } else if (opset_version != 11) { + test.AddOutput("if_out_0", {1}, {11.f}); + } else if (opset_version == 11 && condition_value) { + test.AddOutput("if_out_0", {1}, {2.f}); + } else if (opset_version == 11) { + test.AddOutput("if_out_0", {2}, {11.f, 11.f}); } std::unordered_set excluded_providers; @@ -286,5 +316,13 @@ TEST(If, SymbolicShapeInMainGraph_NoShapeInSubgraph_False) { RunTest(false, options, false); } +TEST(If, Opset11ThenAndElseBranchesProduceDifferentOutputShapes) { + RunOptions options{}; + options.include_dim_values_in_main_graph = true; + options.include_dim_values_in_subgraph = false; + + RunTest(false, options, false, OpTester::ExpectResult::kExpectSuccess, "", 11); +} + } // namespace test } // namespace onnxruntime