#include <loco/IR/CanonicalNode.h>
#include <loco/IR/CanonicalNodeVisitor.h>
+#include <loco/Service/TypeInference.h>
#include <stdex/Memory.h>
#include <type_traits>
+// TODO Remove unused code
namespace
{
void TypeInference::run(loco::Graph *g)
{
- TypeContext ctx;
- TypeGetter type_getter{ctx};
-
- for (auto node : loco::postorder_traversal(loco::output_nodes(g)))
- {
- if (auto canonical_node = dynamic_cast<loco::CanonicalNode *>(node))
- {
- auto tflite_type = canonical_node->accept(&type_getter);
- node->annot(stdex::make_unique<TypeAnnotation>(tflite_type));
-
- // Update type context to allow type getter to retrieve the type of node's arguments
- ctx._node_to_type[node] = tflite_type;
- }
- }
+ loco::CanonicalTypeInferenceRule rule;
+ loco::apply(&rule).to(g);
}
tflite::TensorType TypeInference::get(loco::Node *node)
{
- assert(node->annot<TypeAnnotation>() != nullptr);
- return node->annot<TypeAnnotation>()->type();
+ assert(loco::dtype_known(node));
+ return translateLocoTypeToTFLite(loco::dtype_get(node));
}