mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Register opset-11 CPU kernel for 'If' op (#1948)
* Initial commit * Update * Update * Update * Update * PR comments
This commit is contained in:
Hariharan Seshadri
GitHub
parent
d1b1cdc5c4
commit
31aff686e0
3 changed files with 69 additions and 20 deletions
|
|
@ -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<bool>())
|
||||
.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<bool>())
|
||||
.TypeConstraint("V", DataTypeImpl::AllTensorTypes()),
|
||||
If);
|
||||
11,
|
||||
KernelDefBuilder()
|
||||
.TypeConstraint("B", DataTypeImpl::GetTensorType<bool>())
|
||||
.TypeConstraint("V", DataTypeImpl::AllTensorTypes()),
|
||||
If);
|
||||
|
||||
struct If::Info {
|
||||
Info(const onnxruntime::Node& node, const GraphViewer& subgraph_in) : subgraph(subgraph_in) {
|
||||
|
|
|
|||
|
|
@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, bool, Expand)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, MLFloat16, Expand)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 8, 8, Scan)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, If)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, If)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, 10, Loop)>,
|
||||
|
||||
// Opset 9
|
||||
|
|
@ -699,6 +700,7 @@ void RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
|
|||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, LpPool)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Conv)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, ConvTranspose)>,
|
||||
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, If)>,
|
||||
};
|
||||
|
||||
for (auto& function_table_entry : function_table) {
|
||||
|
|
|
|||
|
|
@ -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<float>("split_input", {2}, {1.f, 10.f});
|
||||
if (opset_version != 11) {
|
||||
test.AddInput<float>("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<float>("split_input", {3}, {1.f, 10.f, 10.f});
|
||||
}
|
||||
|
||||
// graph input to specify which branch to take
|
||||
test.AddInput<bool>("if_cond", {1}, {condition_value});
|
||||
|
||||
test.AddInput<float>("if_graph_input_0", {1}, {1.f});
|
||||
|
||||
std::vector<int64_t> output_shape{1};
|
||||
if (condition_value) {
|
||||
test.AddOutput<float>("if_out_0", output_shape, {2.f});
|
||||
} else {
|
||||
test.AddOutput<float>("if_out_0", output_shape, {11.f});
|
||||
if (opset_version != 11 && condition_value) {
|
||||
test.AddOutput<float>("if_out_0", {1}, {2.f});
|
||||
} else if (opset_version != 11) {
|
||||
test.AddOutput<float>("if_out_0", {1}, {11.f});
|
||||
} else if (opset_version == 11 && condition_value) {
|
||||
test.AddOutput<float>("if_out_0", {1}, {2.f});
|
||||
} else if (opset_version == 11) {
|
||||
test.AddOutput<float>("if_out_0", {2}, {11.f, 11.f});
|
||||
}
|
||||
|
||||
std::unordered_set<std::string> 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
|
||||
|
|
|
|||
Loading…
Reference in a new issue