comaniac commented on a change in pull request #5997:
URL: https://github.com/apache/incubator-tvm/pull/5997#discussion_r451838156
##########
File path: python/tvm/relay/analysis/analysis.py
##########
@@ -351,3 +353,42 @@ def search_fc_transpose(expr):
"""
ret = _ffi_api.search_fc_transpose(expr)
return ret
+
+
+def get_calibration_data(mod, data):
+ """Get the calibration data of a given relay graph
+
+ This pass use the graph runtime to get the calibration data of a module,
which
Review comment:
```suggestion
This pass uses the graph runtime to get the calibration data of a
module, which
```
Per offline discussion, please mention that this pass only works for the
graph without control flow.
##########
File path: src/relay/analysis/get_calibration_data.cc
##########
@@ -0,0 +1,217 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/analysis/get_calibration_data.cc
+ *
+ * \brief To get the calibration data, we need to perform two
+ * steps. First, we need to prepare the module that generate
+ * the tensor values (GetCalibrateModule). Second, we need to
+ * generate the mapping between the values and the functions
+ * (GetCalibrateOutputMap).
+ */
+
+#include <tvm/relay/analysis.h>
+#include <tvm/relay/expr.h>
+#include <tvm/relay/expr_functor.h>
+
+namespace tvm {
+namespace relay {
+
+/*!
+ * \brief This function returns a module that will be used by
+ * the relay graph runtime for collecting the calibration data.
+ * To do that, we first make all inputs and outputs of each
+ * function into the final output (i.e., the final output is a
+ * tuple of tensors). Then, we change the compiler attribute of
+ * each function. Finally, we mark all function to be inlined.
+ */
+
+class Collector : public ExprRewriter {
+ public:
+ explicit Collector(const IRModule& module) { glob_funcs_ =
module->functions; }
+
+ Expr Rewrite_(const CallNode* call, const Expr& post) final {
+ // check if the function implementation is available
+ // intrinsic functions are excluded for now
+ if (call->op->IsInstance<GlobalVarNode>()) {
+ auto var = Downcast<GlobalVar>(call->op);
+ CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not
defined";
+ // we only handle functions with Compiler attribute set
+ auto* fn = glob_funcs_[var].as<FunctionNode>();
+ auto func = GetRef<Function>(fn);
+ if (func->GetAttr<String>(attr::kCompiler)) {
+ // collect all the inputs and outputs
+ for (const auto& it : call->args) new_outputs_.push_back(it);
+ new_outputs_.push_back(post);
+ }
+ }
+ return post;
+ }
+
+ Array<Expr> GetNewOutputs() { return new_outputs_; }
+
+ private:
+ Map<GlobalVar, BaseFunc> glob_funcs_;
+ Array<Expr> new_outputs_;
+};
+
+Expr FlattenToTuple(const Array<Expr>& exprs, const IRModule& module) {
+ auto glob_funcs = module->functions;
+ Array<Expr> fields;
+ for (const auto& it : exprs) {
+ bool is_tuple = false;
+ if (auto* cn = it.as<CallNode>()) {
+ if (cn->op.as<GlobalVarNode>()) {
+ if (auto* fn =
glob_funcs[Downcast<GlobalVar>(cn->op)].as<FunctionNode>()) {
+ if (auto* tn = fn->body.as<TupleNode>()) {
+ is_tuple = true;
+ for (size_t i = 0; i < tn->fields.size(); i++) {
+ fields.push_back(TupleGetItem(it, i));
+ }
+ }
+ }
+ }
+ }
+ if (!is_tuple) {
+ fields.push_back(it);
+ }
+ }
+ return Tuple(fields);
+}
+
+IRModule GetCalibrateModule(IRModule module) {
+ auto glob_funcs = module->functions;
+ // module is mutable, hence, we make a copy of it.
+ module.CopyOnWrite();
+ for (const auto& pair : glob_funcs) {
+ if (auto* fn = pair.second.as<FunctionNode>()) {
+ auto func = GetRef<Function>(fn);
+ // we only collect the outputs for main function
+ if (pair.first->name_hint == "main") {
+ Collector collector(module);
+ PostOrderRewrite(func->body, &collector);
+ auto new_outputs = collector.GetNewOutputs();
+ if (!new_outputs.empty()) {
Review comment:
In which case `new_outputs` will be empty? If this is just a sanity
check, use `CHECK` instead.
##########
File path: src/relay/analysis/get_calibration_data.cc
##########
@@ -0,0 +1,217 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/analysis/get_calibration_data.cc
+ *
+ * \brief To get the calibration data, we need to perform two
+ * steps. First, we need to prepare the module that generate
+ * the tensor values (GetCalibrateModule). Second, we need to
+ * generate the mapping between the values and the functions
+ * (GetCalibrateOutputMap).
+ */
+
+#include <tvm/relay/analysis.h>
+#include <tvm/relay/expr.h>
+#include <tvm/relay/expr_functor.h>
+
+namespace tvm {
+namespace relay {
+
+/*!
+ * \brief This function returns a module that will be used by
+ * the relay graph runtime for collecting the calibration data.
+ * To do that, we first make all inputs and outputs of each
+ * function into the final output (i.e., the final output is a
+ * tuple of tensors). Then, we change the compiler attribute of
+ * each function. Finally, we mark all function to be inlined.
+ */
+
+class Collector : public ExprRewriter {
+ public:
+ explicit Collector(const IRModule& module) { glob_funcs_ =
module->functions; }
+
+ Expr Rewrite_(const CallNode* call, const Expr& post) final {
+ // check if the function implementation is available
+ // intrinsic functions are excluded for now
+ if (call->op->IsInstance<GlobalVarNode>()) {
+ auto var = Downcast<GlobalVar>(call->op);
+ CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not
defined";
+ // we only handle functions with Compiler attribute set
+ auto* fn = glob_funcs_[var].as<FunctionNode>();
+ auto func = GetRef<Function>(fn);
+ if (func->GetAttr<String>(attr::kCompiler)) {
+ // collect all the inputs and outputs
+ for (const auto& it : call->args) new_outputs_.push_back(it);
+ new_outputs_.push_back(post);
+ }
+ }
+ return post;
+ }
+
+ Array<Expr> GetNewOutputs() { return new_outputs_; }
+
+ private:
+ Map<GlobalVar, BaseFunc> glob_funcs_;
+ Array<Expr> new_outputs_;
+};
+
+Expr FlattenToTuple(const Array<Expr>& exprs, const IRModule& module) {
Review comment:
The naming is improper. IIUC, this function does the following things:
* Given an array of expressions, which could be a tensor or a call node.
* If it is a call node, getting the last op of the function body.
* If the last op of the function body is a tuple, add its fields to the new
tuple node fields.
* Otherwise, directly adding the expression to the new tuple node fields.
My concerns are:
* What if `it` itself is a tuple node? For example, the output of the main
function is already a tuple node in the original graph.
* It seems to me that you can directly check the type of `it` to see if it
returns a tuple? Something like (@zhiics please help checking this):
```
CHECK(it->checked_type_.defined());
if (auto tn = GetRef<TupleType>(it->checked_type_)) {
...
}
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
