/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ /*! * Copyright (c) 2018 by Contributors * * \file annotate.cc * * \brief Annotating the graph with simulated quantize operators. */ #include <tvm/relay/transform.h> #include <tvm/relay/analysis.h> #include "./quantize.h" namespace tvm { namespace relay { namespace quantize { using namespace relay::transform; class QAnnotateExpr; class QAnnotateExprNode : public TempExprNode { public: Expr expr; QAnnotateKind kind; void VisitAttrs(tvm::AttrVisitor* v) final { v->Visit("expr", &expr); v->Visit("kind", &kind); } TVM_DLL static QAnnotateExpr make(Expr expr, QAnnotateKind kind); Expr Realize() const final; static constexpr const char* _type_key = "relay.QAnnotateExpr"; TVM_DECLARE_NODE_TYPE_INFO(QAnnotateExprNode, TempExprNode); }; RELAY_DEFINE_NODE_REF(QAnnotateExpr, QAnnotateExprNode, TempExpr); Expr QAnnotateExprNode::Realize() const { return expr; } QAnnotateExpr QAnnotateExprNode::make(Expr expr, QAnnotateKind kind) { auto rnode = make_node<QAnnotateExprNode>(); rnode->expr = expr; rnode->kind = kind; return QAnnotateExpr(rnode); } TVM_REGISTER_API("relay._quantize.make_annotate_expr") .set_body([](TVMArgs args, TVMRetValue *ret) { *ret = QAnnotateExprNode::make(args[0], static_cast<QAnnotateKind>(args[1].operator int())); }); Pass QuantizeAnnotate() { // TODO(tvm-teams): since partition has added cast_hint in different // branches, try to remove this in the future. std::function<Expr(const Expr&)> fmulti_ref = [](const Expr& e) { if (e->derived_from<TempExprNode>()) { const auto* n = e.as<QAnnotateExprNode>(); CHECK(n); const PackedFunc* f = runtime::Registry::Get("relay.quantize.attach_simulated_quantize"); Expr ret = (*f)(n->expr, static_cast<int>(kQInput)); return static_cast<Expr>(QAnnotateExprNode::make(ret, kQInput)); } return e; }; runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func = [=](Function f, Module m, PassContext pc) { auto func = Downcast<Function>(ForwardRewrite(f, "FQAnnotateRewrite", nullptr, fmulti_ref)); auto new_params = func->params; for (const auto& x : FreeVars(func)) { new_params.push_back(x); } return FunctionNode::make(new_params, func->body, func->ret_type, func->type_params, func->attrs); }; return CreateFunctionPass(pass_func, 1, "QuantizeAnnotate", {}); } TVM_REGISTER_API("relay._quantize.QuantizeAnnotate") .set_body_typed(QuantizeAnnotate); } // namespace quantize } // namespace relay } // namespace tvm