[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,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
+    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
+
+    data : Dict[str, NDArray]
+
+    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 = {}
+        value["inputs"] = ref_res[offset:offset+in_len]
+        value["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,192 @@
+/*
+ * 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 subgraphs
+ * (GetCalibrateOutputMap).
+ */
+
+#include <tvm/relay/expr.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/relay/analysis.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
+ * subgrpah into the final output (i.e., the final output is a
+ * tuple of tensors). Then, we change the compiler attribute of
+ * each subgraph. Finally, we mark all subgraph to be inlined.
+ */
+
+IRModule GetCalibrateModule(IRModule module) {
+  class OutputCollector : public ExprRewriter {
+   public:
+    OutputCollector(const Map<GlobalVar, BaseFunc>& glob_funcs)
+      : glob_funcs(glob_funcs) {}
+
+    Expr Rewrite_(const CallNode* call, const Expr& post) final {
+      if (call->op->IsInstance<GlobalVarNode>()) {
+        auto var = Downcast<GlobalVar>(call->op);
+        // check if it is a subgraph of the original graph
+        if (glob_funcs.count(var) > 0) {
+          for (size_t i = 0; i < call->args.size(); i++)
+            new_outputs.push_back(call->args[i]);
+          // need to flatten the output if it is a tuple
+          auto* fn = glob_funcs[var].as<FunctionNode>();
+          if (auto* tn = fn->body.as<TupleNode>()) {
+            for (size_t i = 0; i < tn->fields.size(); i++) {
+              new_outputs.push_back(TupleGetItem(post, i));
+            }
+          } else {
+            new_outputs.push_back(post);
+          }
+        }
+      }
+      return post;
+    }
+
+    Array<Expr> GetNewOutputs() {
+      return new_outputs;
+    }
+
+   private:
+    const Map<GlobalVar, BaseFunc>& glob_funcs;
+    Array<Expr> new_outputs;
+  };
+
+  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);
+      auto* gl_var = pair.first.as<GlobalVarNode>();
+      // we only collect the outputs for main function
+      if (gl_var->name_hint == "main") {
+        OutputCollector output_collector(glob_funcs);
+        PostOrderRewrite(func->body, &output_collector);
+        auto new_outputs = output_collector.GetNewOutputs();
+        if (!new_outputs.empty()) {
+          Array<Expr> fields;
+          for (const auto& output : new_outputs) {
+            fields.push_back(output);
+          }
+          auto tuple = Tuple(fields);
+          func = Function(func->params, tuple, tuple->checked_type_,
+                          func->type_params, func->attrs);
+        }
+      // inline the function if it is not main function
+      } else {
+        func = WithAttr(std::move(func), attr::kInline, tvm::Integer(1));
+      }
+      // reset the compiler attribute to null
+      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 subgraph. The key is a
+ * GlobalVar that corresponds to each subgraph 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.
+ */
+
+Map<GlobalVar, Array<Integer>> GetCalibrateOutputMap(const IRModule& module) {
+  class OutputMapper : public ExprRewriter {
+   public:
+    OutputMapper(Map<GlobalVar, Array<Integer>>* output_map,
+                 const Map<GlobalVar, BaseFunc>& glob_funcs,
+                 int* offset)
+      : output_map(output_map), glob_funcs(glob_funcs), offset(offset) {}
+
+    Expr Rewrite_(const CallNode* call, const Expr& post) final {
+      if (call->op->IsInstance<GlobalVarNode>()) {
+        auto var = Downcast<GlobalVar>(call->op);
+        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
+        int 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;
+    const Map<GlobalVar, BaseFunc>& glob_funcs;
+    int* offset;
+  };
+
+  Map<GlobalVar, Array<Integer>> output_map;
+  int offset = 0;
+  auto glob_funcs = module->functions;
+  for (const auto& pair : glob_funcs) {
+    if (auto* fn = pair.second.as<FunctionNode>()) {
+      auto* gl_var = pair.first.as<GlobalVarNode>();
+      if (gl_var->name_hint == "main") {
+        OutputMapper output_mapper(&output_map, glob_funcs, &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

tests/python/relay/test_analysis_get_calibration_data.py

@@ -0,0 +1,99 @@
+# 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.
+
+import numpy as np
+
+import tvm
+import tvm.relay.testing
+from tvm import relay
+from tvm.relay import transform
+from tvm.relay.analysis import get_calibration_data
+
+
+def check_data_size(mod, data):
+    assert len(data) == len(mod.functions) - 1
+    for key, value in mod.functions.items():
+        if key.name_hint != "main":
+            assert len(data[key]["inputs"]) == len(value.params)
+            if isinstance(value.body, relay.Tuple):
+                assert len(data[key]["outputs"]) == len(value.body.fields)
+            else:
+                assert len(data[key]["outputs"]) == 1
+
+def test_synthetic():
+    # A module with two subgraphs
+    mod = tvm.IRModule()
+
+    x0 = relay.var('x0', shape=(8, 8))
+    y0 = relay.var('y0', shape=(8, 8))
+    z0 = x0 + y0
+    z1 = x0 - y0
+    z2 = relay.Tuple((z0, z1))
+    f0 = relay.Function([x0, y0], z2)
+    g0 = relay.GlobalVar("g0")
+    mod[g0] = f0
+
+    x1 = relay.var('x1', shape=(8, 8))
+    y1 = relay.var('y1', shape=(8, 8))
+    z1 = x1 - y1
+    f1 = relay.Function([x1, y1], z1)
+    g1 = relay.GlobalVar("g1")
+    mod[g1] = f1
+
+    x = relay.var('x', shape=(8, 8))
+    y = relay.var('y', shape=(8, 8))
+    z = relay.var('z', shape=(8, 8))
+    c0 = relay.Call(g0, [x, y])
+    c1 = relay.Call(g1, [relay.TupleGetItem(c0, 0), z])
+    fm = relay.Function([x, y, z], c1)
+    mod["main"] = fm
+
+    x_data = np.random.rand(8, 8).astype('float32')
+    y_data = np.random.rand(8, 8).astype('float32')
+    z_data = np.random.rand(8, 8).astype('float32')
+    data = get_calibration_data(mod, {"x": x_data, "y": y_data, "z": z_data})
+
+    # Check the number and orders
+    check_data_size(mod, data)
+    tvm.testing.assert_allclose(data[g0]["inputs"][0].asnumpy(), x_data)
+    tvm.testing.assert_allclose(data[g0]["inputs"][1].asnumpy(), y_data)
+    tvm.testing.assert_allclose(data[g0]["outputs"][0].asnumpy(), x_data + y_data)
+    tvm.testing.assert_allclose(data[g0]["outputs"][1].asnumpy(), x_data - y_data)
+    tvm.testing.assert_allclose(data[g1]["inputs"][0].asnumpy(), x_data + y_data)
+    tvm.testing.assert_allclose(data[g1]["inputs"][1].asnumpy(), z_data)
+    tvm.testing.assert_allclose(data[g1]["outputs"][0].asnumpy(), x_data + y_data - z_data)
+
+def test_mobilenet_dnnl():
+    dtype = 'float32'
+    ishape = (1, 3, 224, 224)
+    mod, params = relay.testing.mobilenet.get_workload(
+        batch_size=1, dtype='float32')
+
+    mod = transform.AnnotateTarget(["dnnl"])(mod)
+    mod = transform.MergeCompilerRegions()(mod)
+    mod = transform.PartitionGraph()(mod)
+
+    i_data = np.random.uniform(0, 1, ishape).astype(dtype)
+    data = get_calibration_data(mod, {"data": i_data, **params})
+
+    # Check the number and orders
+    check_data_size(mod, data)
+
+if __name__ == "__main__":
+    test_synthetic()
+    test_mobilenet_dnnl()
+