*/
#include "Dialect/IR/TFLNodes.h"
+#include "Dialect/IR/TFLDialect.h"
#include "Dialect/Service/TFLShapeAnnot.h"
#include "Dialect/Service/TFLShapeInferenceRule.h"
#include <loco.h>
+#include <loco/IR/CanonicalDialect.h>
#include <loco/Service/ShapeInference.h>
+#include <loco/Service/CanonicalShapeInferenceRule.h>
+#include <loco/Service/MultiDialectShapeInferenceRule.h>
#include <stdex/Memory.h>
tfl_node->annot<locoex::TFLShapeAnnot>(std::move(shape_annot));
locoex::TFLShapeInferenceRule tfl_rule;
- loco::apply(&tfl_rule).to(g.get());
+ loco::CanonicalShapeInferenceRule canonical_rule;
+ loco::MultiDialectShapeInferenceRule rules;
+
+ rules.bind(loco::CanonicalDialect::get(), &canonical_rule)
+ .bind(locoex::TFLDialect::get(), &tfl_rule);
+
+ loco::apply(&rules).to(g.get());
// Verify
auto check_shape = [](locoex::TFLRelu *tfl_node) {
check_shape(tfl_node);
// step 2.
- loco::apply(&tfl_rule).to(g.get());
+ loco::apply(&rules).to(g.get());
check_shape(tfl_node);
}