[exo-tflite] type inferencing including custom op (#6320)

This commit enables exo-tflite to infer type canonical nodes as well as custop op.

Signed-off-by: Hyun Sik Yoon <hyunsik.yoon@samsung.com>

diff --git a/compiler/exo-tflite/src/TypeInference.cpp b/compiler/exo-tflite/src/TypeInference.cpp
index 18c56a3..ebfef90 100644 (file)
--- a/compiler/exo-tflite/src/TypeInference.cpp
+++ b/compiler/exo-tflite/src/TypeInference.cpp
@@ -20,8 +20,12 @@
 
 #include <loco/IR/CanonicalNode.h>
 #include <loco/IR/CanonicalNodeVisitor.h>
+#include <loco/IR/CanonicalDialect.h>
 #include <loco/Service/TypeInference.h>
 
+#include <locoex/COpDialect.h>
+#include <locoex/Service/COpTypeInference.h>
+
 #include <stdex/Memory.h>
 
 #include <type_traits>
@@ -60,8 +64,15 @@ tflite::TensorType translateLocoTypeToTFLite(loco::DataType dtype)
 
 void TypeInference::run(loco::Graph *g)
 {
-  loco::CanonicalTypeInferenceRule rule;
-  loco::apply(&rule).to(g);
+  loco::CanonicalTypeInferenceRule canonical_rule;
+  locoex::COpTypeInferenceRule cop_rule; // rule for custom op
+
+  loco::MultiDialectTypeInferenceRule rules;
+
+  rules.bind(loco::CanonicalDialect::get(), &canonical_rule)
+      .bind(locoex::COpDialect::get(), &cop_rule);
+
+  loco::apply(&rules).to(g);
 }
 
 tflite::TensorType TypeInference::get(loco::Node *node)