You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/07/06 17:45:28 UTC
[GitHub] [incubator-tvm] comaniac commented on a change in pull request #5997: [Relay] Add pass for getting calibration data from a relay module
comaniac commented on a change in pull request #5997:
URL: https://github.com/apache/incubator-tvm/pull/5997#discussion_r450365939
##########
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
+ 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
Review comment:
Add description.
##########
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
+ 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]
Review comment:
ditto.
##########
File path: 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
Review comment:
Add the reason about why we do this.
##########
File path: 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():
Review comment:
Make a better name of this test.
##########
File path: 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) {
Review comment:
1. Better to be "check if its function implementation is available".
2. It seems to me that we can still collect its data even if its implementation is unavailable. Although it looks useless for BYOC, it might be useful for other use cases. Maybe we can add an option like `include_intransic_func=False` to the API.
##########
File path: 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
Review comment:
1. Move this comment to the else branch.
2. Add the reason about why we do this.
##########
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:
We may need to think of the semantic for the module with control flows, which has to use VM instead of graph runtime.
##########
File path: 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():
Review comment:
Check if DNNL codegen is enabled before running this test.
##########
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
+ 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]
Review comment:
```suggestion
value = {"inputs": ref_res[offset:offset + in_len],
"outputs": ref_res[offset + in_len:offset + in_len + out_len]}
```
##########
File path: 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;
Review comment:
Use `varName_` for all class variables.
##########
File path: 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;
Review comment:
1. s/int/size_t/
2. Ditto to `offset`.
----------------------------------------------------------------
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