[neurun] Add passing a param of operand::Set when creating the Node (#4867)

This commit adds passing a param of operand::Set when creating the Node in order to adjust operand information for each operation.

Signed-off-by: jiseob.jang <jiseob.jang@samsung.com>

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
index 5585873..4eb1a49 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
+++ b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.cc
@@ -28,7 +28,8 @@ OperationFactory::OperationFactory()
 {
   using namespace neurun::model;
 
-  _map[ANEURALNETWORKS_DEPTHWISE_CONV_2D] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_DEPTHWISE_CONV_2D] = [](const OperationFactory::Param &init_param,
+                                               neurun::model::operand::Set &) {
     assert((init_param.input_count == 8 || init_param.input_count == 11) &&
            init_param.output_count == 1);
 
@@ -88,7 +89,8 @@ OperationFactory::OperationFactory()
     return new operation::DepthwiseConv2DNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_MAX_POOL_2D] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_MAX_POOL_2D] = [](const OperationFactory::Param &init_param,
+                                         neurun::model::operand::Set &) {
     assert(init_param.input_count == 7 || init_param.input_count == 10);
     assert(init_param.output_count == 1);
 
@@ -149,7 +151,8 @@ OperationFactory::OperationFactory()
     return new operation::MaxPool2DNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_AVERAGE_POOL_2D] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_AVERAGE_POOL_2D] = [](const OperationFactory::Param &init_param,
+                                             neurun::model::operand::Set &) {
     // TODO We may reuse code here for MAX_POOL_2D. Seems like these two are identical
     assert(init_param.input_count == 7 || init_param.input_count == 10);
     assert(init_param.output_count == 1);
@@ -212,7 +215,8 @@ OperationFactory::OperationFactory()
     return new operation::AvgPool2DNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_CONCATENATION] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_CONCATENATION] = [](const OperationFactory::Param &init_param,
+                                           neurun::model::operand::Set &) {
     assert(init_param.input_count >= 2); // At least one one input tensor and axis
     assert(init_param.output_count == 1);
 
@@ -235,7 +239,8 @@ OperationFactory::OperationFactory()
     return new operation::ConcatNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_RESHAPE] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RESHAPE] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     // Each input should be interpreted as follows:
@@ -252,7 +257,8 @@ OperationFactory::OperationFactory()
     return new operation::ReshapeNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_FULLY_CONNECTED] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_FULLY_CONNECTED] = [](const OperationFactory::Param &init_param,
+                                             neurun::model::operand::Set &) {
     assert(init_param.input_count == 4 && init_param.output_count == 1);
 
     // Each input should be interpreted as follows:
@@ -271,7 +277,8 @@ OperationFactory::OperationFactory()
     return new operation::FullyConnectedNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_SOFTMAX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SOFTMAX] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     // Each input should be interpreted as follows:
@@ -288,7 +295,8 @@ OperationFactory::OperationFactory()
     return new operation::SoftmaxNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_CAST_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_CAST_EX] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -300,7 +308,8 @@ OperationFactory::OperationFactory()
     return new operation::CastNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_CONV_2D] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_CONV_2D] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     using neurun::model::operation::Conv2DNode;
 
     // inputCount is either 7 or 10 acccording to NN API specification.
@@ -362,7 +371,8 @@ OperationFactory::OperationFactory()
     return new Conv2DNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_ADD] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_ADD] = [](const OperationFactory::Param &init_param,
+                                 neurun::model::operand::Set &) {
     assert(init_param.input_count == 3);
     assert(init_param.output_count == 1);
 
@@ -381,7 +391,8 @@ OperationFactory::OperationFactory()
     return new operation::AddNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_REDUCE_SUM_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_REDUCE_SUM_EX] = [](const OperationFactory::Param &init_param,
+                                           neurun::model::operand::Set &) {
     assert(init_param.input_count == 2);
     assert(init_param.output_count == 1);
 
@@ -400,7 +411,8 @@ OperationFactory::OperationFactory()
     return new operation::ReduceSumNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_SUB] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SUB] = [](const OperationFactory::Param &init_param,
+                                 neurun::model::operand::Set &) {
     assert(init_param.input_count == 3);
     assert(init_param.output_count == 1);
 
@@ -419,7 +431,8 @@ OperationFactory::OperationFactory()
     return new operation::SubNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_STRIDED_SLICE] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_STRIDED_SLICE] = [](const OperationFactory::Param &init_param,
+                                           neurun::model::operand::Set &) {
     assert(init_param.input_count == 7 && init_param.output_count == 1);
 
     operand::IndexSet inputs{init_param.inputs[0]};
@@ -459,7 +472,8 @@ OperationFactory::OperationFactory()
     return new operation::StridedSliceNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_TRANSPOSE] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_TRANSPOSE] = [](const OperationFactory::Param &init_param,
+                                       neurun::model::operand::Set &) {
     // TODO make this work with init_param.input_count == 1 (when permutation vector is optional)
 
     // Inputs
@@ -482,7 +496,8 @@ OperationFactory::OperationFactory()
     return new operation::TransposeNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_MUL] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_MUL] = [](const OperationFactory::Param &init_param,
+                                 neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -501,7 +516,8 @@ OperationFactory::OperationFactory()
     return new operation::MulNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_SQUEEZE] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SQUEEZE] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 || init_param.input_count == 2);
     assert(init_param.output_count == 1);
 
@@ -528,7 +544,8 @@ OperationFactory::OperationFactory()
     return new operation::SqueezeNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_TANH] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_TANH] = [](const OperationFactory::Param &init_param,
+                                  neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -541,7 +558,8 @@ OperationFactory::OperationFactory()
     return new operation::TanhNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_LOGISTIC] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_LOGISTIC] = [](const OperationFactory::Param &init_param,
+                                      neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -554,7 +572,8 @@ OperationFactory::OperationFactory()
     return new operation::LogisticNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_DIV] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_DIV] = [](const OperationFactory::Param &init_param,
+                                 neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -572,7 +591,8 @@ OperationFactory::OperationFactory()
     return new operation::DivNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_EXP_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_EXP_EX] = [](const OperationFactory::Param &init_param,
+                                    neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -585,7 +605,8 @@ OperationFactory::OperationFactory()
     return new operation::ExpNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_REDUCE_MAX_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_REDUCE_MAX_EX] = [](const OperationFactory::Param &init_param,
+                                           neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -602,7 +623,8 @@ OperationFactory::OperationFactory()
     return new operation::ReduceMaxNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_NOT_EQUAL_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_NOT_EQUAL_EX] = [](const OperationFactory::Param &init_param,
+                                          neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -619,7 +641,8 @@ OperationFactory::OperationFactory()
     return new operation::ComparisonNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_LOGICAL_AND_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_LOGICAL_AND_EX] = [](const OperationFactory::Param &init_param,
+                                            neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -633,7 +656,8 @@ OperationFactory::OperationFactory()
     return new operation::LogicalAndNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_RSQRT_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RSQRT_EX] = [](const OperationFactory::Param &init_param,
+                                      neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -646,7 +670,8 @@ OperationFactory::OperationFactory()
     return new operation::RSQRTNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_RELU] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RELU] = [](const OperationFactory::Param &init_param,
+                                  neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -659,7 +684,8 @@ OperationFactory::OperationFactory()
     return new operation::ReLUNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_RESIZE_BILINEAR] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RESIZE_BILINEAR] = [](const OperationFactory::Param &init_param,
+                                             neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -678,7 +704,8 @@ OperationFactory::OperationFactory()
     return new operation::ResizeBilinearNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_RELU1] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RELU1] = [](const OperationFactory::Param &init_param,
+                                   neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -691,7 +718,8 @@ OperationFactory::OperationFactory()
     return new operation::ReLU1Node{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_RELU6] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_RELU6] = [](const OperationFactory::Param &init_param,
+                                   neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -704,7 +732,8 @@ OperationFactory::OperationFactory()
     return new operation::ReLU6Node{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_FLOOR] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_FLOOR] = [](const OperationFactory::Param &init_param,
+                                   neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -716,7 +745,8 @@ OperationFactory::OperationFactory()
     return new operation::FloorNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_SPACE_TO_DEPTH] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SPACE_TO_DEPTH] = [](const OperationFactory::Param &init_param,
+                                            neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -733,7 +763,8 @@ OperationFactory::OperationFactory()
     return new operation::SpaceToDepthNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_L2_POOL_2D] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_L2_POOL_2D] = [](const OperationFactory::Param &init_param,
+                                        neurun::model::operand::Set &) {
     assert(init_param.input_count == 10 || init_param.input_count == 7);
     assert(init_param.output_count == 1);
 
@@ -791,7 +822,8 @@ OperationFactory::OperationFactory()
     return new operation::L2Pool2DNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_EMBEDDING_LOOKUP] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_EMBEDDING_LOOKUP] = [](const OperationFactory::Param &init_param,
+                                              neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -805,7 +837,8 @@ OperationFactory::OperationFactory()
     return new operation::EmbeddingLookupNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_L2_NORMALIZATION] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_L2_NORMALIZATION] = [](const OperationFactory::Param &init_param,
+                                              neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -817,7 +850,8 @@ OperationFactory::OperationFactory()
     return new operation::L2NormalizationNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_HASHTABLE_LOOKUP] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_HASHTABLE_LOOKUP] = [](const OperationFactory::Param &init_param,
+                                              neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 2);
 
     // Each output should be interpreted as follows:
@@ -836,7 +870,8 @@ OperationFactory::OperationFactory()
     return new operation::HashtableLookupNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_PRELU_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_PRELU_EX] = [](const OperationFactory::Param &init_param,
+                                      neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -850,7 +885,8 @@ OperationFactory::OperationFactory()
     return new operation::PReLUNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_TRANSPOSE_CONV_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_TRANSPOSE_CONV_EX] = [](const OperationFactory::Param &init_param,
+                                               neurun::model::operand::Set &) {
     assert(init_param.input_count == 6 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -874,7 +910,8 @@ OperationFactory::OperationFactory()
     return new operation::TransposeConvNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_SQRT_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SQRT_EX] = [](const OperationFactory::Param &init_param,
+                                     neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -886,7 +923,8 @@ OperationFactory::OperationFactory()
     return new operation::SQRTNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_LOGICAL_OR_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_LOGICAL_OR_EX] = [](const OperationFactory::Param &init_param,
+                                           neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -900,7 +938,8 @@ OperationFactory::OperationFactory()
     return new operation::LogicalOrNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_LOGICAL_NOT_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_LOGICAL_NOT_EX] = [](const OperationFactory::Param &init_param,
+                                            neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -913,7 +952,8 @@ OperationFactory::OperationFactory()
     return new operation::LogicalNotNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_EQUAL_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_EQUAL_EX] = [](const OperationFactory::Param &init_param,
+                                      neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -930,7 +970,8 @@ OperationFactory::OperationFactory()
     return new operation::ComparisonNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_SQUARED_DIFFERENCE_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_SQUARED_DIFFERENCE_EX] = [](const OperationFactory::Param &init_param,
+                                                   neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -944,7 +985,8 @@ OperationFactory::OperationFactory()
     return new operation::SquaredDifferenceNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_TOPK_V2_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_TOPK_V2_EX] = [](const OperationFactory::Param &init_param,
+                                        neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 2);
 
     // Each output should be interpreted as follows:
@@ -965,7 +1007,8 @@ OperationFactory::OperationFactory()
     return new operation::TopKV2Node{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_GATHER_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_GATHER_EX] = [](const OperationFactory::Param &init_param,
+                                       neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -983,7 +1026,8 @@ OperationFactory::OperationFactory()
     return new operation::GatherNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_NEG_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_NEG_EX] = [](const OperationFactory::Param &init_param,
+                                    neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -996,7 +1040,8 @@ OperationFactory::OperationFactory()
     return new operation::NegNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_ABS_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_ABS_EX] = [](const OperationFactory::Param &init_param,
+                                    neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -1009,7 +1054,8 @@ OperationFactory::OperationFactory()
     return new operation::AbsNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_ARGMAX_EX] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_ARGMAX_EX] = [](const OperationFactory::Param &init_param,
+                                       neurun::model::operand::Set &) {
     assert(init_param.input_count == 2 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -1026,7 +1072,8 @@ OperationFactory::OperationFactory()
     return new operation::ArgMaxNode{inputs, outputs, param};
   };
 
-  _map[ANEURALNETWORKS_DEQUANTIZE] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_DEQUANTIZE] = [](const OperationFactory::Param &init_param,
+                                        neurun::model::operand::Set &) {
     assert(init_param.input_count == 1 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -1039,7 +1086,8 @@ OperationFactory::OperationFactory()
     return new operation::DequantizeNode{inputs, outputs};
   };
 
-  _map[ANEURALNETWORKS_MEAN] = [](const OperationFactory::Param &init_param) {
+  _map[ANEURALNETWORKS_MEAN] = [](const OperationFactory::Param &init_param,
+                                  neurun::model::operand::Set &) {
     assert(init_param.input_count == 3 && init_param.output_count == 1);
 
     operand::IndexSet outputs{init_param.outputs[0]};
@@ -1060,7 +1108,8 @@ OperationFactory::OperationFactory()
 }
 
 neurun::model::operation::Node *OperationFactory::create(ANeuralNetworksOperationType type,
-                                                         const OperationFactory::Param &param)
+                                                         const OperationFactory::Param &param,
+                                                         neurun::model::operand::Set &operands)
 {
-  return _map.at(type)(param);
+  return _map.at(type)(param, operands);
 }

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.h b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.h
index f9735f6..19bb6a8 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.h
+++ b/runtimes/neurun/frontend/nnapi/wrapper/OperationFactory.h
@@ -19,6 +19,7 @@
 
 #include <unordered_map>
 
+#include "model/operand/Set.h"
 #include "model/operation/Node.h"
 #include "NeuralNetworks.h"
 #include "NeuralNetworksEx.h"
@@ -38,8 +39,8 @@ public:
   };
 
 public:
-  using Generator =
-      std::function<neurun::model::operation::Node *(const OperationFactory::Param &)>;
+  using Generator = std::function<neurun::model::operation::Node *(const OperationFactory::Param &,
+                                                                   neurun::model::operand::Set &)>;
 
 public:
   static OperationFactory &instance();
@@ -49,7 +50,8 @@ private:
 
 public:
   neurun::model::operation::Node *create(ANeuralNetworksOperationType,
-                                         const OperationFactory::Param &param);
+                                         const OperationFactory::Param &param,
+                                         neurun::model::operand::Set &operands);
   // TODO add "register" method for separating registration, possibly supporting custom-ops
 
 private:

diff --git a/runtimes/neurun/frontend/nnapi/wrapper/model.cc b/runtimes/neurun/frontend/nnapi/wrapper/model.cc
index 77e883f..a7ed51a 100644 (file)
--- a/runtimes/neurun/frontend/nnapi/wrapper/model.cc
+++ b/runtimes/neurun/frontend/nnapi/wrapper/model.cc
@@ -104,7 +104,7 @@ bool ANeuralNetworksModel::addOperation(ANeuralNetworksOperationType type, uint3
     auto &factory = OperationFactory::instance();
     OperationFactory::Param param{inputCount, inputs, outputCount, outputs};
 
-    auto node = factory.create(type, param);
+    auto node = factory.create(type, param, _model->operands);
     _model->operations.append(std::unique_ptr<neurun::model::operation::Node>{node});
   }
   catch (const std::exception &e)
@@ -132,7 +132,7 @@ bool ANeuralNetworksModel::addOperationEx(ANeuralNetworksOperationTypeEx type, u
     auto &factory = OperationFactory::instance();
     OperationFactory::Param param{inputCount, inputs, outputCount, outputs};
 
-    auto node = factory.create(type, param);
+    auto node = factory.create(type, param, _model->operands);
     _model->operations.append(std::unique_ptr<neurun::model::operation::Node>{node});
   }
   catch (const std::exception &e)