[GitHub] [incubator-tvm] seanlatias commented on a change in pull request #5997: [Relay] Add pass for getting calibration data from a relay module

[GitBox <gi...@apache.org>]

seanlatias commented on a change in pull request #5997:
URL: https://github.com/apache/incubator-tvm/pull/5997#discussion_r452370382



##########
File path: src/relay/analysis/get_calibration_data.cc
##########
@@ -0,0 +1,207 @@
+/*
+ * 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 FlattenOutputTuple(const Array<Expr>& exprs) {
+  Array<Expr> fields;
+  for (const auto& it : exprs) {
+    CHECK(it->checked_type_.defined());
+    if (auto* tn = it->checked_type_.as<TupleTypeNode>()) {
+      // TODO: for now input argument cannot be a tuple
+      CHECK(it->IsInstance<CallNode>());

Review comment:
       Because we only allow output to be a tuple. And output must be a CallNode.




----------------------------------------------------------------
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:
users@infra.apache.org