Apply comment

Samsung · Oct 25, 2023 · b03103d · b03103d
1 parent 6225c61
commit b03103d
Show file tree

Hide file tree

Showing 12 changed files with 54 additions and 20 deletions.
diff --git a/compiler/circlechef/core/src/ModelChef.cpp b/compiler/circlechef/core/src/ModelChef.cpp
@@ -135,7 +135,7 @@ gather_builtincode_map(const ::circlechef::ModelRecipe &model_recipe)
 
   for (const auto &operation : model_recipe.operation())
   {
-    if (operation.type() == "Custom")
+    if (operation.has_extype() && operation.extype() == "Custom")
       continue;
 
     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
@@ -152,7 +152,7 @@ gather_builtincode_map(const ::circlechef::ModelRecipe &model_recipe)
     const auto &graph = model_recipe.graph(g);
     for (const auto &operation : graph.operation())
     {
-      if (operation.type() == "Custom")
+      if (operation.has_extype() && operation.extype() == "Custom")
         continue;
 
       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
@@ -172,7 +172,7 @@ std::set<std::string> gather_customcode_set(const ::circlechef::ModelRecipe &mod
   std::set<std::string> customcode_set;
   for (const auto &operation : model_recipe.operation())
   {
-    if (operation.type() == "Custom")
+    if (operation.has_extype() && operation.extype() == "Custom")
       customcode_set.insert(operation.custom_code());
   }
 
@@ -182,7 +182,7 @@ std::set<std::string> gather_customcode_set(const ::circlechef::ModelRecipe &mod
     const auto &graph = model_recipe.graph(g);
     for (const auto &operation : graph.operation())
     {
-      if (operation.type() == "Custom")
+      if (operation.has_extype() && operation.extype() == "Custom")
         customcode_set.insert(operation.custom_code());
     }
   }
@@ -418,7 +418,7 @@ template <typename T> void cook_graph(const T &graph, CookParams &cp)
     assert(operation.has_type());
 
     std::string op_type = operation.type();
-    if (op_type == "Custom")
+    if (operation.has_custom_code())
       op_type = operation.custom_code();
 
     auto op_chef = op_chef_registry().lookup(op_type).create(&operation);

diff --git a/compiler/tflchef/core/src/CustomOp/AddV2.cpp b/compiler/tflchef/core/src/CustomOp/AddV2.cpp
@@ -29,7 +29,11 @@ AddV2Chef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "AddV2");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "AddV2");
+  }
 
   /**
    * REGISTER_OP("AddV2")

diff --git a/compiler/tflchef/core/src/CustomOp/All.cpp b/compiler/tflchef/core/src/CustomOp/All.cpp
@@ -29,7 +29,11 @@ AllChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "All");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "All");
+  }
 
   /**
    * REGISTER_OP("All")

diff --git a/compiler/tflchef/core/src/CustomOp/BatchMatMulV2.cpp b/compiler/tflchef/core/src/CustomOp/BatchMatMulV2.cpp
@@ -29,7 +29,11 @@ BatchMatMulV2Chef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "BatchMatMulV2");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "BatchMatMulV2");
+  }
 
   /**
    * REGISTER_OP("BatchMatMulV2")

diff --git a/compiler/tflchef/core/src/CustomOp/BroadcastTo.cpp b/compiler/tflchef/core/src/CustomOp/BroadcastTo.cpp
@@ -29,7 +29,11 @@ BroadcastToChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "BroadcastTo");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "BroadcastTo");
+  }
 
   /**
    * REGISTER_OP("BroadcastTo")

diff --git a/compiler/tflchef/core/src/CustomOp/Erf.cpp b/compiler/tflchef/core/src/CustomOp/Erf.cpp
@@ -29,7 +29,11 @@ ErfChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "Erf");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "Erf");
+  }
 
   /**
    * REGISTER_OP("Erf")

diff --git a/compiler/tflchef/core/src/CustomOp/MatMul.cpp b/compiler/tflchef/core/src/CustomOp/MatMul.cpp
@@ -29,7 +29,11 @@ MatMulChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "MatMul");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "MatMul");
+  }
 
   /**
    * REGISTER_OP("MatMul")

diff --git a/compiler/tflchef/core/src/CustomOp/MatrixBandPart.cpp b/compiler/tflchef/core/src/CustomOp/MatrixBandPart.cpp
@@ -29,7 +29,11 @@ MatrixBandPartChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "MatrixBandPart");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "MatrixBandPart");
+  }
 
   /**
    * REGISTER_OP("MatrixBandPart")

diff --git a/compiler/tflchef/core/src/CustomOp/MaxPoolWithArgmax.cpp b/compiler/tflchef/core/src/CustomOp/MaxPoolWithArgmax.cpp
@@ -29,7 +29,11 @@ MaxPoolWithArgmaxChef::custom_value(flatbuffers::FlatBufferBuilder &fbb) const
 {
   auto &operation = (*_operation);
 
-  assert(operation.custom_code() == "MaxPoolWithArgmax");
+  if (operation.has_extype() && operation.extype() == "Custom")
+  {
+    assert(operation.has_custom_code());
+    assert(operation.custom_code() == "MaxPoolWithArgmax");
+  }
 
   /**
    * REGISTER_OP("MaxPoolWithArgmax")

diff --git a/compiler/tflchef/core/src/ModelChef.cpp b/compiler/tflchef/core/src/ModelChef.cpp
@@ -141,7 +141,7 @@ gather_builtincode_map(const ::tflchef::ModelRecipe &model_recipe)
 
   for (const auto &operation : model_recipe.operation())
   {
-    if (operation.type() == "Custom")
+    if (operation.has_extype() && operation.extype() == "Custom")
       continue;
 
     auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
@@ -157,7 +157,7 @@ gather_builtincode_map(const ::tflchef::ModelRecipe &model_recipe)
     const auto &graph = model_recipe.graph(g);
     for (const auto &operation : graph.operation())
     {
-      if (operation.type() == "Custom")
+      if (operation.has_extype() && operation.extype() == "Custom")
         continue;
 
       auto op_chef = op_chef_registry().lookup(operation.type()).create(&operation);
@@ -177,7 +177,7 @@ std::set<std::string> gather_customcode_set(const ::tflchef::ModelRecipe &model_
   std::set<std::string> customcode_set;
   for (const auto &operation : model_recipe.operation())
   {
-    if (operation.type() == "Custom")
+    if (operation.has_extype() && operation.extype() == "Custom")
       customcode_set.insert(operation.custom_code());
   }
 
@@ -187,7 +187,7 @@ std::set<std::string> gather_customcode_set(const ::tflchef::ModelRecipe &model_
     const auto &graph = model_recipe.graph(g);
     for (const auto &operation : graph.operation())
     {
-      if (operation.type() == "Custom")
+      if (operation.has_extype() && operation.extype() == "Custom")
         customcode_set.insert(operation.custom_code());
     }
   }
@@ -618,7 +618,7 @@ template <typename T> std::map<std::string, int32_t> cook_graph(const T &graph,
     assert(operation.has_type());
 
     std::string op_type = operation.type();
-    if (op_type == "Custom")
+    if (operation.has_custom_code())
       op_type = operation.custom_code();
 
     auto op_chef = op_chef_registry().lookup(op_type).create(&operation);

diff --git a/compiler/tflchef/tests/custom_erf/test.recipe b/compiler/tflchef/tests/custom_erf/test.recipe
@@ -9,10 +9,11 @@ operand {
   shape { dim: 1 dim: 3 dim: 3 dim: 2 }
 }
 operation {
-  type: "Custom"
+  type: "Erf"
   input: "ifm"
   output: "ofm"
   custom_code: "Erf"
+  extype: "Custom"
 }
 input: "ifm"
 output: "ofm"
diff --git a/res/TensorFlowLiteRecipes/All_000/test.recipe b/res/TensorFlowLiteRecipes/All_000/test.recipe
@@ -18,14 +18,15 @@ operand {
   }
 }
 operation {
-  type: "Custom"
+  type: "All"
   all_options {
     keep_dims: false
   }
   input: "ifm"
   input: "All/reduction_indices"
   output: "ofm"
   custom_code: "All"
+  extype: "Custom"
 }
 input: "ifm"
 output: "ofm"