[Relay] Add pass for getting calibration data from a relay module

include/tvm/relay/analysis.h

@@ -236,6 +236,24 @@ TVM_DLL Array<Pattern> UnmatchedCases(const Match& match, const IRModule& mod);
  */
 TVM_DLL std::unordered_map<const Object*, size_t> GetExprRefCount(const Expr& body);
 
+/*!
+ * \brief Get the updated module for collecting calibration data.
+ *
+ * \param mod The module to be updated.
+ *
+ * \return The updated module.
+ */
+TVM_DLL IRModule GetCalibrateModule(IRModule mod);
+
+/*!
+ * \brief Get the output map between subgrpahs and its inputs/output.
+ *
+ * \param mod The module for running calibration.
+ *
+ * \return The mapping between a subgraph name and its postition in the output tuple.
+ */
+TVM_DLL Map<GlobalVar, Array<Integer>> GetCalibrateOutputMap(const IRModule& mod);
+
 }  // namespace relay
 }  // namespace tvm
 

python/tvm/relay/analysis/analysis.py

@@ -21,6 +21,8 @@
 configuring the passes and scripting them in Python.
 """
 from tvm.ir import IRModule
+from tvm.relay import transform, build_module
+from tvm.runtime.ndarray import cpu
 
 from . import _ffi_api
 from .feature import Feature
@@ -351,3 +353,43 @@ 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
+    includes the input and output values of each subgraph. The returned data uses
+    the GlobalVar of each subgraph as a key. Users can further access the inputs and
+    outputs by using `inputs` or  `outputs` as the key.
+
+    Parameters
+    ----------
+    mod : tvm.IRModule
+        The input module for collecting the calibration data
+
+    data : Dict[str, NDArray]
+        The input data for running the module
+
+    Returns
+    -------
+    data : Dict[tvm.relay.GlobalVar, Dict[str, NDArray]]
+    """
+    output_map = _ffi_api.get_calibrate_output_map(mod)
+
+    mod = _ffi_api.get_calibrate_module(mod)
+    mod = transform.Inline()(mod)
+
+    ref_ex = build_module.create_executor("graph", mod=mod, ctx=cpu(0))
+    ref_res = ref_ex.evaluate()(**data)
+
+    calib_data = {}
+    for gvar, indices in output_map.items():
+        offset = int(indices[0])
+        in_len = int(indices[1])
+        out_len = int(indices[2])
+        value = {"inputs": ref_res[offset:offset + in_len],
+                 "outputs": ref_res[offset + in_len:offset + in_len + out_len]}
+        calib_data[gvar] = value
+
+    return calib_data

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()) {
+          Expr tuple = FlattenToTuple(new_outputs, module);
+          func =
+              Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs);
+          module->Update(pair.first, func);
+        }
+      }
+    }
+  }
+  // reset the attribute of functions for running graph runtime
+  for (const auto& pair : glob_funcs) {
+    if (auto* fn = pair.second.as<FunctionNode>()) {
+      auto func = GetRef<Function>(fn);
+      if (func->GetAttr<String>(attr::kCompiler)) {
+        // we need to inline the functions in order to run grpah runtime
+        func = WithAttr(std::move(func), attr::kInline, tvm::Integer(1));
+        // reset the compiler attribute to null for llvm execution
+        func = WithAttr(std::move(func), attr::kCompiler, NullValue<ObjectRef>());
+        module->Update(pair.first, func);
+      }
+    }
+  }
+  return module;
+}
+
+/*!
+ * \brief This function generates the output mapping between
+ * the calibration data and each function. The key is a
+ * GlobalVar that corresponds to each function and the value
+ * is an array of integers. The size of the array is always
+ * three. The first value is the offset the points to the start.
+ * The second value is the number of inputs. The third value
+ * is the number of outputs.
+ */
+
+class OutputMapper : public ExprRewriter {
+ public:
+  OutputMapper(Map<GlobalVar, Array<Integer>>* output_map, const IRModule& module, size_t* offset)
+      : output_map_(output_map), offset_(offset) {
+    glob_funcs_ = module->functions;
+  }
+
+  Expr Rewrite_(const CallNode* call, const Expr& post) final {
+    if (call->op->IsInstance<GlobalVarNode>()) {
+      auto var = Downcast<GlobalVar>(call->op);
+      CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not defined";
+      CHECK_EQ(output_map_->count(var), 0)
+          << "Repeated function call " << var << " is not supported.";
+      // 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)) {
+        Array<Integer> info;
+        // the first value is the offset
+        info.push_back(Integer(*offset_));
+        // the second value is the number of inputs
+        info.push_back(Integer(call->args.size()));
+        // the third value is the number of outputs
+        // we need to check if the output is a tuple
+        size_t out_size = 1;
+        auto* fn = glob_funcs_[var].as<FunctionNode>();
+        if (auto* tn = fn->body.as<TupleNode>()) {
+          info.push_back(Integer(tn->fields.size()));
+          out_size = tn->fields.size();
+        } else {
+          info.push_back(Integer(1));
+        }
+        output_map_->Set(var, info);
+        // calculate the offset for the next function
+        *offset_ = *offset_ + call->args.size() + out_size;
+      }
+    }
+    return post;
+  }
+
+ private:
+  Map<GlobalVar, Array<Integer>>* output_map_;
+  Map<GlobalVar, BaseFunc> glob_funcs_;
+  size_t* offset_;
+};
+
+Map<GlobalVar, Array<Integer>> GetCalibrateOutputMap(const IRModule& module) {
+  Map<GlobalVar, Array<Integer>> output_map;
+  size_t offset = 0;
+  auto glob_funcs = module->functions;
+  for (const auto& pair : glob_funcs) {
+    if (auto* fn = pair.second.as<FunctionNode>()) {
+      if (pair.first->name_hint == "main") {
+        OutputMapper output_mapper(&output_map, module, &offset);
+        auto func = GetRef<Function>(fn);
+        PostOrderRewrite(func->body, &output_mapper);
+      }
+    }
+  }
+
+  return output_map;
+}
+
+TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_module").set_body_typed([](IRModule mod) {
+  return GetCalibrateModule(mod);
+});
+
+TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_output_map")
+    .set_body_typed([](const IRModule& mod) { return GetCalibrateOutputMap(mod); });
+
+}  // namespace relay
+}  // namespace tvm