From 373a46207f8ad63c61440751dba9c043b0e4e6d6 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Wed, 17 Jun 2020 18:46:57 +0000 Subject: [PATCH 01/25] add simple pass to extract outputs --- 3rdparty/dlpack | 2 +- include/tvm/relay/transform.h | 2 + python/tvm/relay/transform/transform.py | 4 + .../transforms/calibrate_partition_graph.cc | 104 ++++++++++++++++++ 4 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 src/relay/transforms/calibrate_partition_graph.cc diff --git a/3rdparty/dlpack b/3rdparty/dlpack index 3ec04430e89a..0acb731e0e43 160000 --- a/3rdparty/dlpack +++ b/3rdparty/dlpack @@ -1 +1 @@ -Subproject commit 3ec04430e89a6834e5a1b99471f415fa939bf642 +Subproject commit 0acb731e0e43d15deee27b66f10e4c5b4e667913 diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h index 39212d8fd1e5..9311458e4edd 100644 --- a/include/tvm/relay/transform.h +++ b/include/tvm/relay/transform.h @@ -349,6 +349,8 @@ TVM_DLL Pass Inline(); */ TVM_DLL Pass RemoveUnusedFunctions(Array entry_functions); +TVM_DLL Pass CalibratePartitionGraph(); + } // namespace transform /*! diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py index 8f4ec1046500..8b5663d0d736 100644 --- a/python/tvm/relay/transform/transform.py +++ b/python/tvm/relay/transform/transform.py @@ -604,6 +604,10 @@ def Inline(): return _ffi_api.Inline() +def CalibratePartitionGraph(): + return _ffi_api.CalibratePartitionGraph() + + def gradient(expr, mod=None, mode='higher_order'): """ Transform the input function, diff --git a/src/relay/transforms/calibrate_partition_graph.cc b/src/relay/transforms/calibrate_partition_graph.cc new file mode 100644 index 000000000000..57568a3c08ca --- /dev/null +++ b/src/relay/transforms/calibrate_partition_graph.cc @@ -0,0 +1,104 @@ +/* + * 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/transforms/calibrate_partition_graph.cc + * + * \brief Partition an input function into multiple functions according based + */ + +#include +#include +#include + +#include + +namespace tvm { +namespace relay { + +namespace calibrate_partition { + +/* +*/ + +IRModule CalibratePartition(IRModule module) { + class OutputCollector : public ExprRewriter { + public: + OutputCollector() = default; + + Expr Rewrite_(const CallNode* call, const Expr& post) final { + if (call->op->IsInstance()) { + LOG(INFO) << call->op; + auto var = Downcast(call->op); + output_map.Set(var, post); + } + return post; + } + + Map GetOutputMap() { + return output_map; + } + + private: + Map output_map; + + }; + + 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()) { + auto func = GetRef(fn); + // Collect the output + OutputCollector output_collector; + auto body = PostOrderRewrite(func->body, &output_collector); + auto output_map = output_collector.GetOutputMap(); + for (const auto& kv : output_map) { + LOG(INFO) << kv.first << " " << kv.second; + } + func = Function(func->params, body, func->ret_type, func->type_params, func->attrs); + // Reset the compiler attribute to null + func = WithAttr(std::move(func), attr::kCompiler, NullValue()); + module->Update(pair.first, func); + } + } + return module; +} + +} // namespace calibrate_partition + +namespace transform { + +Pass CalibratePartitionGraph() { + runtime::TypedPackedFunc calib_func = [=](IRModule m, + PassContext pc) { + return calibrate_partition::CalibratePartition(m); + }; + + auto partition_pass = CreateModulePass(calib_func, 0, "CalibratePartitionGraph", {}); + return Sequential({partition_pass, InferType()}); +} + +TVM_REGISTER_GLOBAL("relay._transform.CalibratePartitionGraph").set_body_typed(transform::CalibratePartitionGraph); + +} // namespace transform + +} // namespace relay +} // namespace tvm From 9da5b2e9dea8b6844310557c06c45268a6a36ddf Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 22 Jun 2020 17:34:19 +0000 Subject: [PATCH 02/25] complete pass that collects all function inputs/outputs --- .../transforms/calibrate_partition_graph.cc | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/src/relay/transforms/calibrate_partition_graph.cc b/src/relay/transforms/calibrate_partition_graph.cc index 57568a3c08ca..0d017b146f83 100644 --- a/src/relay/transforms/calibrate_partition_graph.cc +++ b/src/relay/transforms/calibrate_partition_graph.cc @@ -44,19 +44,20 @@ IRModule CalibratePartition(IRModule module) { Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { - LOG(INFO) << call->op; auto var = Downcast(call->op); - output_map.Set(var, post); + for (size_t i = 0; i < call->args.size(); i++) + new_outputs.push_back(call->args[i]); + new_outputs.push_back(post); } return post; } - Map GetOutputMap() { - return output_map; + Array GetNewOutputs() { + return new_outputs; } private: - Map output_map; + Array new_outputs; }; @@ -69,11 +70,16 @@ IRModule CalibratePartition(IRModule module) { // Collect the output OutputCollector output_collector; auto body = PostOrderRewrite(func->body, &output_collector); - auto output_map = output_collector.GetOutputMap(); - for (const auto& kv : output_map) { - LOG(INFO) << kv.first << " " << kv.second; + auto new_outputs = output_collector.GetNewOutputs(); + if (!new_outputs.empty()) { + Array fields; + fields.push_back(body); + 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); } - func = Function(func->params, body, func->ret_type, func->type_params, func->attrs); // Reset the compiler attribute to null func = WithAttr(std::move(func), attr::kCompiler, NullValue()); module->Update(pair.first, func); From 91d0f09b34f1656c1355b50682c2bcfc8e824cd9 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Wed, 24 Jun 2020 18:28:59 +0000 Subject: [PATCH 03/25] add analysis pass for collecting outputs --- include/tvm/relay/analysis.h | 2 + python/tvm/relay/analysis/analysis.py | 3 ++ .../transforms/calibrate_partition_graph.cc | 41 +++++++++++++++++++ 3 files changed, 46 insertions(+) diff --git a/include/tvm/relay/analysis.h b/include/tvm/relay/analysis.h index b4b1b9dcc4e8..ce042e426e85 100644 --- a/include/tvm/relay/analysis.h +++ b/include/tvm/relay/analysis.h @@ -236,6 +236,8 @@ TVM_DLL Array UnmatchedCases(const Match& match, const IRModule& mod); */ TVM_DLL std::unordered_map GetExprRefCount(const Expr& body); +TVM_DLL Map GetCalibrateOutputMap(const IRModule& mod); + } // namespace relay } // namespace tvm diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index c237859eb987..85381f51ec5d 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -351,3 +351,6 @@ def search_fc_transpose(expr): """ ret = _ffi_api.search_fc_transpose(expr) return ret + +def get_calibrate_output_map(mod): + return _ffi_api.get_calibrate_output_map(mod) diff --git a/src/relay/transforms/calibrate_partition_graph.cc b/src/relay/transforms/calibrate_partition_graph.cc index 0d017b146f83..c2bccb7bddf0 100644 --- a/src/relay/transforms/calibrate_partition_graph.cc +++ b/src/relay/transforms/calibrate_partition_graph.cc @@ -26,6 +26,7 @@ #include #include #include +#include #include @@ -106,5 +107,45 @@ TVM_REGISTER_GLOBAL("relay._transform.CalibratePartitionGraph").set_body_typed(t } // namespace transform +Map GetCalibrateOutputMap(const IRModule& module) { + class OutputMapper : public ExprRewriter { + public: + OutputMapper(Map& output_map, int& offset) : output_map(output_map), offset(offset) {} + + Expr Rewrite_(const CallNode* call, const Expr& post) final { + if (call->op->IsInstance()) { + auto var = Downcast(call->op); + output_map.Set(var, Integer(offset)); + offset = offset + call->args.size() + 1; + } + return post; + } + + private: + Map& output_map; + int& offset; + + }; + + Map output_map; + int offset = 1; + auto glob_funcs = module->functions; + for (const auto& pair : glob_funcs) { + if (auto* fn = pair.second.as()) { + auto func = GetRef(fn); + // Collect the output + OutputMapper output_mapper(output_map, offset); + auto body = PostOrderRewrite(func->body, &output_mapper); + } + } + + return output_map; +} + +TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_output_map") + .set_body_typed([](const IRModule& mod) { + return GetCalibrateOutputMap(mod); +}); + } // namespace relay } // namespace tvm From 7c1e0cac96e0b3efdf106594c210ef7d46fc884c Mon Sep 17 00:00:00 2001 From: seanlatias Date: Fri, 26 Jun 2020 14:30:10 +0000 Subject: [PATCH 04/25] reorganize the files --- include/tvm/relay/analysis.h | 2 ++ include/tvm/relay/transform.h | 2 -- python/tvm/relay/analysis/analysis.py | 18 +++++++++++++ .../calibrate_partition_graph.cc | 25 ++++--------------- 4 files changed, 25 insertions(+), 22 deletions(-) rename src/relay/{transforms => analysis}/calibrate_partition_graph.cc (84%) diff --git a/include/tvm/relay/analysis.h b/include/tvm/relay/analysis.h index ce042e426e85..48d52e31666e 100644 --- a/include/tvm/relay/analysis.h +++ b/include/tvm/relay/analysis.h @@ -236,6 +236,8 @@ TVM_DLL Array UnmatchedCases(const Match& match, const IRModule& mod); */ TVM_DLL std::unordered_map GetExprRefCount(const Expr& body); +TVM_DLL IRModule GetCalibrateModule(IRModule mod); + TVM_DLL Map GetCalibrateOutputMap(const IRModule& mod); } // namespace relay diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h index 9311458e4edd..39212d8fd1e5 100644 --- a/include/tvm/relay/transform.h +++ b/include/tvm/relay/transform.h @@ -349,8 +349,6 @@ TVM_DLL Pass Inline(); */ TVM_DLL Pass RemoveUnusedFunctions(Array entry_functions); -TVM_DLL Pass CalibratePartitionGraph(); - } // namespace transform /*! diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index 85381f51ec5d..ea50e6a2b312 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/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 @@ -352,5 +354,21 @@ def search_fc_transpose(expr): ret = _ffi_api.search_fc_transpose(expr) return ret +def get_calibrate_module(mod): + return _ffi_api.get_calibrate_module(mod) + def get_calibrate_output_map(mod): return _ffi_api.get_calibrate_output_map(mod) + +def get_calibrate_results(mod, i_data, params): + output_map = _ffi_api.get_calibrate_output_map(mod) + + mod = get_calibrate_module(mod) + print(mod) + + mod = transform.Inline()(mod) + + ref_ex = build_module.create_executor("graph", mod=mod, ctx=cpu(0)) + ref_res = ref_ex.evaluate()(i_data, **params) + print(ref_res) + diff --git a/src/relay/transforms/calibrate_partition_graph.cc b/src/relay/analysis/calibrate_partition_graph.cc similarity index 84% rename from src/relay/transforms/calibrate_partition_graph.cc rename to src/relay/analysis/calibrate_partition_graph.cc index c2bccb7bddf0..08fbf681a2c0 100644 --- a/src/relay/transforms/calibrate_partition_graph.cc +++ b/src/relay/analysis/calibrate_partition_graph.cc @@ -33,12 +33,10 @@ namespace tvm { namespace relay { -namespace calibrate_partition { - /* */ -IRModule CalibratePartition(IRModule module) { +IRModule GetCalibrateModule(IRModule module) { class OutputCollector : public ExprRewriter { public: OutputCollector() = default; @@ -89,23 +87,10 @@ IRModule CalibratePartition(IRModule module) { return module; } -} // namespace calibrate_partition - -namespace transform { - -Pass CalibratePartitionGraph() { - runtime::TypedPackedFunc calib_func = [=](IRModule m, - PassContext pc) { - return calibrate_partition::CalibratePartition(m); - }; - - auto partition_pass = CreateModulePass(calib_func, 0, "CalibratePartitionGraph", {}); - return Sequential({partition_pass, InferType()}); -} - -TVM_REGISTER_GLOBAL("relay._transform.CalibratePartitionGraph").set_body_typed(transform::CalibratePartitionGraph); - -} // namespace transform +TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_module") + .set_body_typed([](IRModule mod) { + return GetCalibrateModule(mod); +}); Map GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { From 08c86643c0b721be5831da4d3b0bff4a59a46ffb Mon Sep 17 00:00:00 2001 From: seanlatias Date: Thu, 2 Jul 2020 19:29:38 +0000 Subject: [PATCH 05/25] add the first test --- include/tvm/relay/analysis.h | 16 ++++- python/tvm/relay/analysis/analysis.py | 25 +++++--- .../analysis/calibrate_partition_graph.cc | 63 +++++++++++-------- .../test_analysis_get_calibration_data.py | 58 +++++++++++++++++ 4 files changed, 124 insertions(+), 38 deletions(-) create mode 100644 tests/python/relay/test_analysis_get_calibration_data.py diff --git a/include/tvm/relay/analysis.h b/include/tvm/relay/analysis.h index 48d52e31666e..8eda7dd824ca 100644 --- a/include/tvm/relay/analysis.h +++ b/include/tvm/relay/analysis.h @@ -236,9 +236,23 @@ TVM_DLL Array UnmatchedCases(const Match& match, const IRModule& mod); */ TVM_DLL std::unordered_map 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); -TVM_DLL Map GetCalibrateOutputMap(const 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> GetCalibrateOutputMap(const IRModule& mod); } // namespace relay } // namespace tvm diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index ea50e6a2b312..b13049f0af30 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -354,21 +354,26 @@ def search_fc_transpose(expr): ret = _ffi_api.search_fc_transpose(expr) return ret -def get_calibrate_module(mod): - return _ffi_api.get_calibrate_module(mod) -def get_calibrate_output_map(mod): - return _ffi_api.get_calibrate_output_map(mod) - -def get_calibrate_results(mod, i_data, params): +def get_calibration_data(mod, data): output_map = _ffi_api.get_calibrate_output_map(mod) + print(output_map) - mod = get_calibrate_module(mod) + mod = _ffi_api.get_calibrate_module(mod) print(mod) mod = transform.Inline()(mod) ref_ex = build_module.create_executor("graph", mod=mod, ctx=cpu(0)) - ref_res = ref_ex.evaluate()(i_data, **params) - print(ref_res) - + ref_res = ref_ex.evaluate()(**data) + + calib_data = {} + for gvar, indices in output_map.items(): + offset = int(indices[0]) + length = int(indices[1]) + value = {} + value["intputs"] = ref_res[offset:offset+length] + value["outputs"] = [ref_res[offset+length]] + calib_data[gvar] = value + + return calib_data diff --git a/src/relay/analysis/calibrate_partition_graph.cc b/src/relay/analysis/calibrate_partition_graph.cc index 08fbf681a2c0..e845900c1cd6 100644 --- a/src/relay/analysis/calibrate_partition_graph.cc +++ b/src/relay/analysis/calibrate_partition_graph.cc @@ -18,18 +18,15 @@ */ /* - * \file src/relay/transforms/calibrate_partition_graph.cc + * \file src/relay/analysis/get_calibration_data.cc * - * \brief Partition an input function into multiple functions according based + * \brief */ #include #include -#include #include -#include - namespace tvm { namespace relay { @@ -39,14 +36,18 @@ namespace relay { IRModule GetCalibrateModule(IRModule module) { class OutputCollector : public ExprRewriter { public: - OutputCollector() = default; + OutputCollector(const Map& glob_funcs) + : glob_funcs(glob_funcs) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); - for (size_t i = 0; i < call->args.size(); i++) - new_outputs.push_back(call->args[i]); - new_outputs.push_back(post); + // 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]); + new_outputs.push_back(post); + } } return post; } @@ -56,6 +57,7 @@ IRModule GetCalibrateModule(IRModule module) { } private: + const Map& glob_funcs; Array new_outputs; }; @@ -66,18 +68,23 @@ IRModule GetCalibrateModule(IRModule module) { for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { auto func = GetRef(fn); - // Collect the output - OutputCollector output_collector; - auto body = PostOrderRewrite(func->body, &output_collector); - auto new_outputs = output_collector.GetNewOutputs(); - if (!new_outputs.empty()) { - Array fields; - fields.push_back(body); - for (const auto& output : new_outputs) { - fields.push_back(output); + auto* gl_var = pair.first.as(); + if (gl_var->name_hint == "main") { + // Collect the output + OutputCollector output_collector(glob_funcs); + auto body = PostOrderRewrite(func->body, &output_collector); + auto new_outputs = output_collector.GetNewOutputs(); + if (!new_outputs.empty()) { + Array 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); } - auto tuple = Tuple(fields); - func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); + } else { + // Inline the function if it is not main function + func = WithAttr(std::move(func), attr::kInline, tvm::Integer(1)); } // Reset the compiler attribute to null func = WithAttr(std::move(func), attr::kCompiler, NullValue()); @@ -92,28 +99,30 @@ TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_module") return GetCalibrateModule(mod); }); -Map GetCalibrateOutputMap(const IRModule& module) { +Map> GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { public: - OutputMapper(Map& output_map, int& offset) : output_map(output_map), offset(offset) {} + OutputMapper(Map>& output_map, int& offset) : output_map(output_map), offset(offset) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); - output_map.Set(var, Integer(offset)); + Array info; + info.push_back(Integer(offset)); + info.push_back(Integer(call->args.size())); + output_map.Set(var, info); offset = offset + call->args.size() + 1; } return post; } private: - Map& output_map; + Map>& output_map; int& offset; - }; - Map output_map; - int offset = 1; + Map> output_map; + int offset = 0; auto glob_funcs = module->functions; for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { diff --git a/tests/python/relay/test_analysis_get_calibration_data.py b/tests/python/relay/test_analysis_get_calibration_data.py new file mode 100644 index 000000000000..c1b477307e86 --- /dev/null +++ b/tests/python/relay/test_analysis_get_calibration_data.py @@ -0,0 +1,58 @@ +# 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 tvm +from tvm import relay +from tvm.relay.analysis import get_calibration_data + +import numpy as np + +def test_basic(): + # 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 + f0 = relay.Function([x0, y0], z0) + 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, [c0, 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}) + + print(data) + +test_basic() + From e2d128133fb78e8cd14c4219fb00755f9c9852ac Mon Sep 17 00:00:00 2001 From: seanlatias Date: Fri, 3 Jul 2020 13:43:25 +0000 Subject: [PATCH 06/25] update test with tuples --- python/tvm/relay/analysis/analysis.py | 12 +++- .../analysis/calibrate_partition_graph.cc | 47 ++++++++++++---- .../test_analysis_get_calibration_data.py | 55 +++++++++++++++++-- 3 files changed, 94 insertions(+), 20 deletions(-) diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index b13049f0af30..4b34afbd6547 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -356,6 +356,8 @@ def search_fc_transpose(expr): def get_calibration_data(mod, data): + """Get the calibration data of a given relay graph + """ output_map = _ffi_api.get_calibrate_output_map(mod) print(output_map) @@ -370,10 +372,14 @@ def get_calibration_data(mod, data): calib_data = {} for gvar, indices in output_map.items(): offset = int(indices[0]) - length = int(indices[1]) + in_len = int(indices[1]) + out_len = int(indices[2]) value = {} - value["intputs"] = ref_res[offset:offset+length] - value["outputs"] = [ref_res[offset+length]] + value["inputs"] = ref_res[offset:offset+in_len] + if out_len == 1: + value["outputs"] = [ref_res[offset+in_len]] + else: + value["outputs"] = ref_res[offset+in_len:offset+in_len+out_len] calib_data[gvar] = value return calib_data diff --git a/src/relay/analysis/calibrate_partition_graph.cc b/src/relay/analysis/calibrate_partition_graph.cc index e845900c1cd6..f81c05458bd6 100644 --- a/src/relay/analysis/calibrate_partition_graph.cc +++ b/src/relay/analysis/calibrate_partition_graph.cc @@ -46,7 +46,15 @@ IRModule GetCalibrateModule(IRModule module) { if (glob_funcs.count(var) > 0) { for (size_t i = 0; i < call->args.size(); i++) new_outputs.push_back(call->args[i]); - new_outputs.push_back(post); + // need to flatten the output if it is a tuple + auto* fn = glob_funcs[var].as(); + if (auto* tn = fn->body.as()) { + 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; @@ -94,15 +102,13 @@ IRModule GetCalibrateModule(IRModule module) { return module; } -TVM_REGISTER_GLOBAL("relay.analysis.get_calibrate_module") - .set_body_typed([](IRModule mod) { - return GetCalibrateModule(mod); -}); - Map> GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { public: - OutputMapper(Map>& output_map, int& offset) : output_map(output_map), offset(offset) {} + OutputMapper(Map>& output_map, + const Map& 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()) { @@ -110,14 +116,23 @@ Map> GetCalibrateOutputMap(const IRModule& module) { Array info; info.push_back(Integer(offset)); info.push_back(Integer(call->args.size())); + int out_size = 1; + auto* fn = glob_funcs[var].as(); + if (auto* tn = fn->body.as()) { + info.push_back(Integer(tn->fields.size())); + out_size = tn->fields.size(); + } else { + info.push_back(Integer(1)); + } output_map.Set(var, info); - offset = offset + call->args.size() + 1; + offset = offset + call->args.size() + out_size; } return post; } private: Map>& output_map; + const Map& glob_funcs; int& offset; }; @@ -126,16 +141,24 @@ Map> GetCalibrateOutputMap(const IRModule& module) { auto glob_funcs = module->functions; for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { - auto func = GetRef(fn); - // Collect the output - OutputMapper output_mapper(output_map, offset); - auto body = PostOrderRewrite(func->body, &output_mapper); + auto* gl_var = pair.first.as(); + if (gl_var->name_hint == "main") { + // Collect the output + OutputMapper output_mapper(output_map, glob_funcs, offset); + auto func = GetRef(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); diff --git a/tests/python/relay/test_analysis_get_calibration_data.py b/tests/python/relay/test_analysis_get_calibration_data.py index c1b477307e86..bf9efc3394ca 100644 --- a/tests/python/relay/test_analysis_get_calibration_data.py +++ b/tests/python/relay/test_analysis_get_calibration_data.py @@ -16,19 +16,23 @@ # under the License. import tvm +import tvm.relay.testing from tvm import relay +from tvm.relay import transform from tvm.relay.analysis import get_calibration_data import numpy as np -def test_basic(): +def make_basic_graph(): # 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 - f0 = relay.Function([x0, y0], z0) + z1 = x0 - y0 + z2 = relay.Tuple((z0, z1)) + f0 = relay.Function([x0, y0], z2) g0 = relay.GlobalVar("g0") mod[g0] = f0 @@ -43,16 +47,57 @@ def test_basic(): y = relay.var('y', shape=(8, 8)) z = relay.var('z', shape=(8, 8)) c0 = relay.Call(g0, [x, y]) - c1 = relay.Call(g1, [c0, z]) + c1 = relay.Call(g1, [relay.TupleGetItem(c0, 0), z]) fm = relay.Function([x, y, z], c1) mod["main"] = fm + return mod, g0, g1 + +def test_basic_single_data(): + mod, g0, g1 = make_basic_graph() + 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}) - print(data) + # Check the number and orders + assert len(data) == 2 + assert len(data[g0]["inputs"]) == 2 + assert len(data[g0]["outputs"]) == 2 + assert len(data[g1]["inputs"]) == 2 + assert len(data[g1]["outputs"]) == 1 + 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 + assert len(data) == 2 + #assert len(data[g0]["inputs"]) == 2 + """ + 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[g1]["inputs"][1].asnumpy(), z_data) + """ -test_basic() +test_basic_single_data() From db0bf7dd39ce5ce9dae696f959cfe6bd7e0e8d62 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 12:48:30 +0000 Subject: [PATCH 07/25] clean up Python code --- python/tvm/relay/analysis/analysis.py | 23 +++++++---- .../test_analysis_get_calibration_data.py | 40 +++++++++---------- 2 files changed, 34 insertions(+), 29 deletions(-) diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index 4b34afbd6547..1b8584347058 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -357,13 +357,25 @@ def search_fc_transpose(expr): 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) - print(output_map) mod = _ffi_api.get_calibrate_module(mod) - print(mod) - mod = transform.Inline()(mod) ref_ex = build_module.create_executor("graph", mod=mod, ctx=cpu(0)) @@ -376,10 +388,7 @@ def get_calibration_data(mod, data): out_len = int(indices[2]) value = {} value["inputs"] = ref_res[offset:offset+in_len] - if out_len == 1: - value["outputs"] = [ref_res[offset+in_len]] - else: - value["outputs"] = ref_res[offset+in_len:offset+in_len+out_len] + value["outputs"] = ref_res[offset+in_len:offset+in_len+out_len] calib_data[gvar] = value return calib_data diff --git a/tests/python/relay/test_analysis_get_calibration_data.py b/tests/python/relay/test_analysis_get_calibration_data.py index bf9efc3394ca..fe23bc35be84 100644 --- a/tests/python/relay/test_analysis_get_calibration_data.py +++ b/tests/python/relay/test_analysis_get_calibration_data.py @@ -15,15 +15,26 @@ # 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 -import numpy as np -def make_basic_graph(): +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() @@ -51,22 +62,13 @@ def make_basic_graph(): fm = relay.Function([x, y, z], c1) mod["main"] = fm - return mod, g0, g1 - -def test_basic_single_data(): - mod, g0, g1 = make_basic_graph() - 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 - assert len(data) == 2 - assert len(data[g0]["inputs"]) == 2 - assert len(data[g0]["outputs"]) == 2 - assert len(data[g1]["inputs"]) == 2 - assert len(data[g1]["outputs"]) == 1 + 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) @@ -81,23 +83,17 @@ def test_mobilenet_dnnl(): 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 - assert len(data) == 2 - #assert len(data[g0]["inputs"]) == 2 - """ - 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[g1]["inputs"][1].asnumpy(), z_data) - """ + check_data_size(mod, data) -test_basic_single_data() +if __name__ == "__main__": + test_synthetic() + test_mobilenet_dnnl() From 84ae166e4b644ac2d6f47eccb6dbecfca48e70f8 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 13:03:29 +0000 Subject: [PATCH 08/25] merge with upstream --- 3rdparty/dlpack | 2 +- 3rdparty/dmlc-core | 2 +- 3rdparty/vta-hw | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/3rdparty/dlpack b/3rdparty/dlpack index 0acb731e0e43..3ec04430e89a 160000 --- a/3rdparty/dlpack +++ b/3rdparty/dlpack @@ -1 +1 @@ -Subproject commit 0acb731e0e43d15deee27b66f10e4c5b4e667913 +Subproject commit 3ec04430e89a6834e5a1b99471f415fa939bf642 diff --git a/3rdparty/dmlc-core b/3rdparty/dmlc-core index 6c401e242c59..ebcaa42d0b72 160000 --- a/3rdparty/dmlc-core +++ b/3rdparty/dmlc-core @@ -1 +1 @@ -Subproject commit 6c401e242c59a1f4c913918246591bb13fd714e7 +Subproject commit ebcaa42d0b7271265e6a02ff468c6753a8fbecdb diff --git a/3rdparty/vta-hw b/3rdparty/vta-hw index db65157208ec..f1c3385e2042 160000 --- a/3rdparty/vta-hw +++ b/3rdparty/vta-hw @@ -1 +1 @@ -Subproject commit db65157208ec8fabb7b548c94596211b9db04190 +Subproject commit f1c3385e2042e876c719d1b30dd90ed533c51ed7 From 0c76cf45338c23bb3d1483ec915008c33ffacac4 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 13:06:26 +0000 Subject: [PATCH 09/25] clean up transform.py --- python/tvm/relay/transform/transform.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py index 2ec14f4427ae..ede63808d4fd 100644 --- a/python/tvm/relay/transform/transform.py +++ b/python/tvm/relay/transform/transform.py @@ -649,10 +649,6 @@ def Inline(): return _ffi_api.Inline() -def CalibratePartitionGraph(): - return _ffi_api.CalibratePartitionGraph() - - def gradient(expr, mod=None, mode='higher_order'): """ Transform the input function, From 4f13b0c0fbdbb7e26013a85326eebe73c045ffb7 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 14:32:16 +0000 Subject: [PATCH 10/25] add comments for cpp files --- ...ition_graph.cc => get_calibration_data.cc} | 46 ++++++++++++++----- 1 file changed, 35 insertions(+), 11 deletions(-) rename src/relay/analysis/{calibrate_partition_graph.cc => get_calibration_data.cc} (76%) diff --git a/src/relay/analysis/calibrate_partition_graph.cc b/src/relay/analysis/get_calibration_data.cc similarity index 76% rename from src/relay/analysis/calibrate_partition_graph.cc rename to src/relay/analysis/get_calibration_data.cc index f81c05458bd6..d092e93b0528 100644 --- a/src/relay/analysis/calibrate_partition_graph.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -17,10 +17,14 @@ * under the License. */ -/* +/*! * \file src/relay/analysis/get_calibration_data.cc * - * \brief + * \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 @@ -30,8 +34,14 @@ 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 { @@ -67,7 +77,6 @@ IRModule GetCalibrateModule(IRModule module) { private: const Map& glob_funcs; Array new_outputs; - }; auto glob_funcs = module->functions; @@ -77,10 +86,10 @@ IRModule GetCalibrateModule(IRModule module) { if (auto* fn = pair.second.as()) { auto func = GetRef(fn); auto* gl_var = pair.first.as(); + // we only collect the outputs for main function if (gl_var->name_hint == "main") { - // Collect the output OutputCollector output_collector(glob_funcs); - auto body = PostOrderRewrite(func->body, &output_collector); + PostOrderRewrite(func->body, &output_collector); auto new_outputs = output_collector.GetNewOutputs(); if (!new_outputs.empty()) { Array fields; @@ -88,13 +97,14 @@ IRModule GetCalibrateModule(IRModule module) { fields.push_back(output); } auto tuple = Tuple(fields); - func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); + func = Function(func->params, tuple, tuple->checked_type_, + func->type_params, func->attrs); } + // inline the function if it is not main function } else { - // Inline the function if it is not main function func = WithAttr(std::move(func), attr::kInline, tvm::Integer(1)); } - // Reset the compiler attribute to null + // reset the compiler attribute to null func = WithAttr(std::move(func), attr::kCompiler, NullValue()); module->Update(pair.first, func); } @@ -102,6 +112,16 @@ IRModule GetCalibrateModule(IRModule module) { 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> GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { public: @@ -114,8 +134,12 @@ Map> GetCalibrateOutputMap(const IRModule& module) { if (call->op->IsInstance()) { auto var = Downcast(call->op); Array 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(); if (auto* tn = fn->body.as()) { @@ -125,6 +149,7 @@ Map> GetCalibrateOutputMap(const IRModule& module) { 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; @@ -143,7 +168,6 @@ Map> GetCalibrateOutputMap(const IRModule& module) { if (auto* fn = pair.second.as()) { auto* gl_var = pair.first.as(); if (gl_var->name_hint == "main") { - // Collect the output OutputMapper output_mapper(output_map, glob_funcs, offset); auto func = GetRef(fn); PostOrderRewrite(func->body, &output_mapper); From 22b10db0d50c540f00e69cf67290d9c41c48a135 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 17:34:43 +0000 Subject: [PATCH 11/25] fix lint issues --- src/relay/analysis/get_calibration_data.cc | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index d092e93b0528..b789a5aea9a6 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -97,7 +97,7 @@ IRModule GetCalibrateModule(IRModule module) { fields.push_back(output); } auto tuple = Tuple(fields); - func = Function(func->params, tuple, tuple->checked_type_, + func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); } // inline the function if it is not main function @@ -125,9 +125,9 @@ IRModule GetCalibrateModule(IRModule module) { Map> GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { public: - OutputMapper(Map>& output_map, + OutputMapper(Map>* output_map, const Map& glob_funcs, - int& offset) + int* offset) : output_map(output_map), glob_funcs(glob_funcs), offset(offset) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { @@ -135,7 +135,7 @@ Map> GetCalibrateOutputMap(const IRModule& module) { auto var = Downcast(call->op); Array info; // the first value is the offset - info.push_back(Integer(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 @@ -148,19 +148,19 @@ Map> GetCalibrateOutputMap(const IRModule& module) { } else { info.push_back(Integer(1)); } - output_map.Set(var, info); + output_map->Set(var, info); // calculate the offset for the next function - offset = offset + call->args.size() + out_size; + *offset = *offset + call->args.size() + out_size; } return post; } private: - Map>& output_map; + Map>* output_map; const Map& glob_funcs; - int& offset; + int* offset; }; - + Map> output_map; int offset = 0; auto glob_funcs = module->functions; @@ -168,7 +168,7 @@ Map> GetCalibrateOutputMap(const IRModule& module) { if (auto* fn = pair.second.as()) { auto* gl_var = pair.first.as(); if (gl_var->name_hint == "main") { - OutputMapper output_mapper(output_map, glob_funcs, offset); + OutputMapper output_mapper(&output_map, glob_funcs, &offset); auto func = GetRef(fn); PostOrderRewrite(func->body, &output_mapper); } From d188d4c742e5eda34085323bc57a27c099bddcaa Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 6 Jul 2020 17:44:13 +0000 Subject: [PATCH 12/25] update submodules --- 3rdparty/dmlc-core | 2 +- 3rdparty/vta-hw | 2 +- src/relay/analysis/get_calibration_data.cc | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/3rdparty/dmlc-core b/3rdparty/dmlc-core index ebcaa42d0b72..6c401e242c59 160000 --- a/3rdparty/dmlc-core +++ b/3rdparty/dmlc-core @@ -1 +1 @@ -Subproject commit ebcaa42d0b7271265e6a02ff468c6753a8fbecdb +Subproject commit 6c401e242c59a1f4c913918246591bb13fd714e7 diff --git a/3rdparty/vta-hw b/3rdparty/vta-hw index f1c3385e2042..db65157208ec 160000 --- a/3rdparty/vta-hw +++ b/3rdparty/vta-hw @@ -1 +1 @@ -Subproject commit f1c3385e2042e876c719d1b30dd90ed533c51ed7 +Subproject commit db65157208ec8fabb7b548c94596211b9db04190 diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index b789a5aea9a6..5cccea4f393c 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -46,7 +46,7 @@ namespace relay { IRModule GetCalibrateModule(IRModule module) { class OutputCollector : public ExprRewriter { public: - OutputCollector(const Map& glob_funcs) + explicit OutputCollector(const Map& glob_funcs) : glob_funcs(glob_funcs) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { From 81ac706a0d3a74e84627231977186befef79c4c8 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Tue, 7 Jul 2020 14:37:51 +0000 Subject: [PATCH 13/25] modify files according to the review --- python/tvm/relay/analysis/analysis.py | 7 ++- src/relay/analysis/get_calibration_data.cc | 56 +++++++++---------- .../test_analysis_get_calibration_data.py | 10 +++- 3 files changed, 39 insertions(+), 34 deletions(-) diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index 1b8584347058..e0bae799afe8 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -366,8 +366,10 @@ def get_calibration_data(mod, data): Parameters ---------- mod : tvm.IRModule + The input module for collecting the calibration data data : Dict[str, NDArray] + The input data for running the module Returns ------- @@ -386,9 +388,8 @@ def get_calibration_data(mod, data): 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] + 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 diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 5cccea4f393c..7ff340ad89c0 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -23,7 +23,7 @@ * \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 + * generate the mapping between the values and the functions * (GetCalibrateOutputMap). */ @@ -38,32 +38,32 @@ 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 + * function 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. + * each function. Finally, we mark all function to be inlined. */ IRModule GetCalibrateModule(IRModule module) { class OutputCollector : public ExprRewriter { public: - explicit OutputCollector(const Map& glob_funcs) - : glob_funcs(glob_funcs) {} + explicit OutputCollector(const Map& glob_funcs) : glob_funcs_(glob_funcs) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); - // check if it is a subgraph of the original graph + // check if the function implementation is available + // intrinsic functions are excluded for now if (glob_funcs.count(var) > 0) { for (size_t i = 0; i < call->args.size(); i++) - new_outputs.push_back(call->args[i]); + new_outputs_.push_back(call->args[i]); // need to flatten the output if it is a tuple - auto* fn = glob_funcs[var].as(); + auto* fn = glob_funcs_[var].as(); if (auto* tn = fn->body.as()) { for (size_t i = 0; i < tn->fields.size(); i++) { - new_outputs.push_back(TupleGetItem(post, i)); + new_outputs_.push_back(TupleGetItem(post, i)); } } else { - new_outputs.push_back(post); + new_outputs_.push_back(post); } } } @@ -71,12 +71,12 @@ IRModule GetCalibrateModule(IRModule module) { } Array GetNewOutputs() { - return new_outputs; + return new_outputs_; } private: - const Map& glob_funcs; - Array new_outputs; + const Map& glob_funcs_; + Array new_outputs_; }; auto glob_funcs = module->functions; @@ -100,11 +100,11 @@ IRModule GetCalibrateModule(IRModule module) { func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); } - // inline the function if it is not main function } else { + // 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 + // reset the compiler attribute to null for llvm execution func = WithAttr(std::move(func), attr::kCompiler, NullValue()); module->Update(pair.first, func); } @@ -114,8 +114,8 @@ IRModule GetCalibrateModule(IRModule 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 + * 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 @@ -127,42 +127,42 @@ Map> GetCalibrateOutputMap(const IRModule& module) { public: OutputMapper(Map>* output_map, const Map& glob_funcs, - int* offset) - : output_map(output_map), glob_funcs(glob_funcs), offset(offset) {} + size_t* offset) + : output_map_(output_map), glob_funcs_(glob_funcs), offset_(offset) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); Array info; // the first value is the offset - info.push_back(Integer(*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(); + size_t out_size = 1; + auto* fn = glob_funcs_[var].as(); if (auto* tn = fn->body.as()) { info.push_back(Integer(tn->fields.size())); out_size = tn->fields.size(); } else { info.push_back(Integer(1)); } - output_map->Set(var, info); + output_map_->Set(var, info); // calculate the offset for the next function - *offset = *offset + call->args.size() + out_size; + *offset_ = *offset_ + call->args.size() + out_size; } return post; } private: - Map>* output_map; - const Map& glob_funcs; - int* offset; + Map>* output_map_; + const Map& glob_funcs_; + size_t* offset_; }; Map> output_map; - int offset = 0; + size_t offset = 0; auto glob_funcs = module->functions; for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { diff --git a/tests/python/relay/test_analysis_get_calibration_data.py b/tests/python/relay/test_analysis_get_calibration_data.py index fe23bc35be84..8b4a1aaadab0 100644 --- a/tests/python/relay/test_analysis_get_calibration_data.py +++ b/tests/python/relay/test_analysis_get_calibration_data.py @@ -34,7 +34,7 @@ def check_data_size(mod, data): else: assert len(data[key]["outputs"]) == 1 -def test_synthetic(): +def test_simple_graph(): # A module with two subgraphs mod = tvm.IRModule() @@ -54,6 +54,7 @@ def test_synthetic(): 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)) @@ -78,6 +79,10 @@ def test_synthetic(): tvm.testing.assert_allclose(data[g1]["outputs"][0].asnumpy(), x_data + y_data - z_data) def test_mobilenet_dnnl(): + if not tvm.get_global_func("relay.ext.dnnl", True): + print("skip because DNNL codegen is not available") + return + dtype = 'float32' ishape = (1, 3, 224, 224) mod, params = relay.testing.mobilenet.get_workload( @@ -94,6 +99,5 @@ def test_mobilenet_dnnl(): check_data_size(mod, data) if __name__ == "__main__": - test_synthetic() + test_simple_graph() test_mobilenet_dnnl() - From 1bf93f97660a33727661bd3bfcd7ef423dd7a553 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Tue, 7 Jul 2020 14:47:04 +0000 Subject: [PATCH 14/25] fix style and typo --- src/relay/analysis/get_calibration_data.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 7ff340ad89c0..026cadf5d4bf 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -44,16 +44,16 @@ namespace relay { */ IRModule GetCalibrateModule(IRModule module) { - class OutputCollector : public ExprRewriter { + class Collector : public ExprRewriter { public: - explicit OutputCollector(const Map& glob_funcs) : glob_funcs_(glob_funcs) {} + explicit Collector(const Map& glob_funcs) : glob_funcs_(glob_funcs) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); // check if the function implementation is available // intrinsic functions are excluded for now - if (glob_funcs.count(var) > 0) { + 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 @@ -88,9 +88,9 @@ IRModule GetCalibrateModule(IRModule module) { auto* gl_var = pair.first.as(); // 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(); + Collector collector(glob_funcs); + PostOrderRewrite(func->body, &collector); + auto new_outputs = collector.GetNewOutputs(); if (!new_outputs.empty()) { Array fields; for (const auto& output : new_outputs) { From 02b929449a4502e6df7a2e35dcf30bf0be8a4fad Mon Sep 17 00:00:00 2001 From: seanlatias Date: Tue, 7 Jul 2020 14:54:51 +0000 Subject: [PATCH 15/25] fix lint error --- src/relay/analysis/get_calibration_data.cc | 27 ++++++++-------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 026cadf5d4bf..d8ed5cef2474 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -27,9 +27,9 @@ * (GetCalibrateOutputMap). */ +#include #include #include -#include namespace tvm { namespace relay { @@ -54,8 +54,7 @@ IRModule GetCalibrateModule(IRModule module) { // check if the function implementation is available // intrinsic functions are excluded for now if (glob_funcs_.count(var) > 0) { - for (size_t i = 0; i < call->args.size(); i++) - new_outputs_.push_back(call->args[i]); + 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(); if (auto* tn = fn->body.as()) { @@ -70,9 +69,7 @@ IRModule GetCalibrateModule(IRModule module) { return post; } - Array GetNewOutputs() { - return new_outputs_; - } + Array GetNewOutputs() { return new_outputs_; } private: const Map& glob_funcs_; @@ -97,8 +94,8 @@ IRModule GetCalibrateModule(IRModule module) { fields.push_back(output); } auto tuple = Tuple(fields); - func = Function(func->params, tuple, tuple->checked_type_, - func->type_params, func->attrs); + func = + Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); } } else { // we need to inline the functions in order to run grpah runtime @@ -126,9 +123,8 @@ Map> GetCalibrateOutputMap(const IRModule& module) { class OutputMapper : public ExprRewriter { public: OutputMapper(Map>* output_map, - const Map& glob_funcs, - size_t* offset) - : output_map_(output_map), glob_funcs_(glob_funcs), offset_(offset) {} + const Map& glob_funcs, size_t* offset) + : output_map_(output_map), glob_funcs_(glob_funcs), offset_(offset) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { @@ -178,15 +174,12 @@ Map> GetCalibrateOutputMap(const IRModule& module) { 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_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); -}); + .set_body_typed([](const IRModule& mod) { return GetCalibrateOutputMap(mod); }); } // namespace relay } // namespace tvm From 52b6116b9b6cb2a63643f6706ef55d89125fb079 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Tue, 7 Jul 2020 20:28:58 +0000 Subject: [PATCH 16/25] add checks for repeated function calls --- src/relay/analysis/get_calibration_data.cc | 25 +++++++++++----------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index d8ed5cef2474..d8e2a2d97728 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -49,21 +49,20 @@ IRModule GetCalibrateModule(IRModule module) { explicit Collector(const Map& glob_funcs) : glob_funcs_(glob_funcs) {} 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()) { auto var = Downcast(call->op); - // check if the function implementation is available - // intrinsic functions are excluded for now - 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(); - if (auto* tn = fn->body.as()) { - for (size_t i = 0; i < tn->fields.size(); i++) { - new_outputs_.push_back(TupleGetItem(post, i)); - } - } else { - new_outputs_.push_back(post); + CHECK(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(); + if (auto* tn = fn->body.as()) { + 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; @@ -129,6 +128,8 @@ Map> GetCalibrateOutputMap(const IRModule& module) { Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); + CHECK(glob_funcs_.count(var) > 0); + CHECK(output_map_->count(var) == 0) << "Repeated function call is not supported."; Array info; // the first value is the offset info.push_back(Integer(*offset_)); From 604c9075dfd095d63a46b611cdcde4ecd41e1a58 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Tue, 7 Jul 2020 20:48:38 +0000 Subject: [PATCH 17/25] fix lint error --- src/relay/analysis/get_calibration_data.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index d8e2a2d97728..ac9765df68e3 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -53,7 +53,7 @@ IRModule GetCalibrateModule(IRModule module) { // intrinsic functions are excluded for now if (call->op->IsInstance()) { auto var = Downcast(call->op); - CHECK(glob_funcs_.count(var) > 0); + CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not defined"; 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(); @@ -128,8 +128,9 @@ Map> GetCalibrateOutputMap(const IRModule& module) { Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); - CHECK(glob_funcs_.count(var) > 0); - CHECK(output_map_->count(var) == 0) << "Repeated function call is not supported."; + 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."; Array info; // the first value is the offset info.push_back(Integer(*offset_)); From 031d79e2e39f8f3b1f5d686950a49513c77f7c86 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Wed, 8 Jul 2020 20:47:31 +0000 Subject: [PATCH 18/25] merge review comments --- src/relay/analysis/get_calibration_data.cc | 149 +++++++++++------- .../test_analysis_get_calibration_data.py | 2 + 2 files changed, 92 insertions(+), 59 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index ac9765df68e3..9760996d2b04 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -43,66 +43,91 @@ namespace relay { * each function. Finally, we mark all function to be inlined. */ -IRModule GetCalibrateModule(IRModule module) { - class Collector : public ExprRewriter { - public: - explicit Collector(const Map& glob_funcs) : glob_funcs_(glob_funcs) {} - - 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()) { - auto var = Downcast(call->op); - CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not defined"; - 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(); - if (auto* tn = fn->body.as()) { - for (size_t i = 0; i < tn->fields.size(); i++) { - new_outputs_.push_back(TupleGetItem(post, i)); - } - } else { - new_outputs_.push_back(post); - } +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()) { + auto var = Downcast(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(); + auto func = GetRef(fn); + if (func->GetAttr(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; } + return post; + } - Array GetNewOutputs() { return new_outputs_; } + Array GetNewOutputs() { return new_outputs_; } - private: - const Map& glob_funcs_; - Array new_outputs_; - }; + private: + Map glob_funcs_; + Array new_outputs_; +}; +Expr FlattenToTuple(const Array& exprs, const IRModule& module) { + auto glob_funcs = module->functions; + Array fields; + for (const auto& it : exprs) { + bool is_tuple = false; + if (auto* cn = it.as()) { + if (cn->op.as()) { + if (auto* fn = glob_funcs[Downcast(cn->op)].as()) { + if (auto* tn = fn->body.as()) { + 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()) { auto func = GetRef(fn); - auto* gl_var = pair.first.as(); // we only collect the outputs for main function - if (gl_var->name_hint == "main") { - Collector collector(glob_funcs); + if (pair.first->name_hint == "main") { + Collector collector(module); PostOrderRewrite(func->body, &collector); auto new_outputs = collector.GetNewOutputs(); if (!new_outputs.empty()) { - Array fields; - for (const auto& output : new_outputs) { - fields.push_back(output); - } - auto tuple = Tuple(fields); + Expr tuple = FlattenToTuple(new_outputs, module); func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); + module->Update(pair.first, func); } - } else { + } + } + } + // reset the attribute of functions for running graph runtime + for (const auto& pair : glob_funcs) { + if (auto* fn = pair.second.as()) { + auto func = GetRef(fn); + if (func->GetAttr(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()); + module->Update(pair.first, func); } - // reset the compiler attribute to null for llvm execution - func = WithAttr(std::move(func), attr::kCompiler, NullValue()); - module->Update(pair.first, func); } } return module; @@ -118,19 +143,23 @@ IRModule GetCalibrateModule(IRModule module) { * is the number of outputs. */ -Map> GetCalibrateOutputMap(const IRModule& module) { - class OutputMapper : public ExprRewriter { - public: - OutputMapper(Map>* output_map, - const Map& glob_funcs, size_t* offset) - : output_map_(output_map), glob_funcs_(glob_funcs), offset_(offset) {} - - Expr Rewrite_(const CallNode* call, const Expr& post) final { - if (call->op->IsInstance()) { - auto var = Downcast(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."; +class OutputMapper : public ExprRewriter { + public: + OutputMapper(Map>* 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()) { + auto var = Downcast(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(); + auto func = GetRef(fn); + if (func->GetAttr(attr::kCompiler)) { Array info; // the first value is the offset info.push_back(Integer(*offset_)); @@ -150,15 +179,17 @@ Map> GetCalibrateOutputMap(const IRModule& module) { // calculate the offset for the next function *offset_ = *offset_ + call->args.size() + out_size; } - return post; } + return post; + } - private: - Map>* output_map_; - const Map& glob_funcs_; - size_t* offset_; - }; + private: + Map>* output_map_; + Map glob_funcs_; + size_t* offset_; +}; +Map> GetCalibrateOutputMap(const IRModule& module) { Map> output_map; size_t offset = 0; auto glob_funcs = module->functions; @@ -166,7 +197,7 @@ Map> GetCalibrateOutputMap(const IRModule& module) { if (auto* fn = pair.second.as()) { auto* gl_var = pair.first.as(); if (gl_var->name_hint == "main") { - OutputMapper output_mapper(&output_map, glob_funcs, &offset); + OutputMapper output_mapper(&output_map, module, &offset); auto func = GetRef(fn); PostOrderRewrite(func->body, &output_mapper); } diff --git a/tests/python/relay/test_analysis_get_calibration_data.py b/tests/python/relay/test_analysis_get_calibration_data.py index 8b4a1aaadab0..9a29f2e73e23 100644 --- a/tests/python/relay/test_analysis_get_calibration_data.py +++ b/tests/python/relay/test_analysis_get_calibration_data.py @@ -44,6 +44,7 @@ def test_simple_graph(): z1 = x0 - y0 z2 = relay.Tuple((z0, z1)) f0 = relay.Function([x0, y0], z2) + f0 = f0.with_attr("Compiler", "test_graph") g0 = relay.GlobalVar("g0") mod[g0] = f0 @@ -51,6 +52,7 @@ def test_simple_graph(): y1 = relay.var('y1', shape=(8, 8)) z1 = x1 - y1 f1 = relay.Function([x1, y1], z1) + f1 = f1.with_attr("Compiler", "test_graph") g1 = relay.GlobalVar("g1") mod[g1] = f1 From f68ecba820b72c44217b9acdfc7dd7eb33922d95 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Wed, 8 Jul 2020 20:51:37 +0000 Subject: [PATCH 19/25] small simplification --- src/relay/analysis/get_calibration_data.cc | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 9760996d2b04..1a3b69df21ef 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -195,8 +195,7 @@ Map> GetCalibrateOutputMap(const IRModule& module) { auto glob_funcs = module->functions; for (const auto& pair : glob_funcs) { if (auto* fn = pair.second.as()) { - auto* gl_var = pair.first.as(); - if (gl_var->name_hint == "main") { + if (pair.first->name_hint == "main") { OutputMapper output_mapper(&output_map, module, &offset); auto func = GetRef(fn); PostOrderRewrite(func->body, &output_mapper); From bf242186521a496dddde5e3fe6ff4a1ebe303528 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Thu, 9 Jul 2020 16:49:37 +0000 Subject: [PATCH 20/25] revise the code according to the review comments --- python/tvm/relay/analysis/analysis.py | 12 ++++++-- src/relay/analysis/get_calibration_data.cc | 32 ++++++++-------------- 2 files changed, 20 insertions(+), 24 deletions(-) diff --git a/python/tvm/relay/analysis/analysis.py b/python/tvm/relay/analysis/analysis.py index e0bae799afe8..632af460ce96 100644 --- a/python/tvm/relay/analysis/analysis.py +++ b/python/tvm/relay/analysis/analysis.py @@ -358,11 +358,17 @@ def search_fc_transpose(expr): 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 + This pass uses the graph runtime to get the calibration data of a module, which + includes the input and output values of each function. The returned data uses + the GlobalVar of each function as a key. Users can further access the inputs and outputs by using `inputs` or `outputs` as the key. + Following are some limitations: + 1. The input module (graph) cannot have control flows. + 2. The input arguments of each function cannot be tuples (outputs can be tuples). + 3. We only handle top-level functions (i.e., nested function is not handled). + 4. We only handle functions with `Compiler` attribute being set. + Parameters ---------- mod : tvm.IRModule diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 1a3b69df21ef..56767a2ee273 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -72,24 +72,17 @@ class Collector : public ExprRewriter { Array new_outputs_; }; -Expr FlattenToTuple(const Array& exprs, const IRModule& module) { - auto glob_funcs = module->functions; +Expr FlattenOutputTuple(const Array& exprs) { Array fields; for (const auto& it : exprs) { - bool is_tuple = false; - if (auto* cn = it.as()) { - if (cn->op.as()) { - if (auto* fn = glob_funcs[Downcast(cn->op)].as()) { - if (auto* tn = fn->body.as()) { - is_tuple = true; - for (size_t i = 0; i < tn->fields.size(); i++) { - fields.push_back(TupleGetItem(it, i)); - } - } - } + CHECK(it->checked_type_.defined()); + if (auto* tn = it->checked_type_.as()) { + // TODO: for now input argument cannot be a tuple + CHECK(it->IsInstance()); + for (size_t i = 0; i < tn->fields.size(); i++) { + fields.push_back(TupleGetItem(it, i)); } - } - if (!is_tuple) { + } else { fields.push_back(it); } } @@ -108,12 +101,9 @@ IRModule GetCalibrateModule(IRModule module) { 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); - } + Expr tuple = FlattenOutputTuple(new_outputs); + func = Function(func->params, tuple, tuple->checked_type_, func->type_params, func->attrs); + module->Update(pair.first, func); } } } From 9ccd88f6cd28bfd41df44707a287f28228376859 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Thu, 9 Jul 2020 16:57:44 +0000 Subject: [PATCH 21/25] add username in TODO --- src/relay/analysis/get_calibration_data.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 56767a2ee273..47e060b895e6 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -77,7 +77,7 @@ Expr FlattenOutputTuple(const Array& exprs) { for (const auto& it : exprs) { CHECK(it->checked_type_.defined()); if (auto* tn = it->checked_type_.as()) { - // TODO: for now input argument cannot be a tuple + // TODO(seanlatias): for now input argument cannot be a tuple CHECK(it->IsInstance()); for (size_t i = 0; i < tn->fields.size(); i++) { fields.push_back(TupleGetItem(it, i)); From cfd1e208bdce889d8d0073c6c5b208dc45e1b7bb Mon Sep 17 00:00:00 2001 From: seanlatias Date: Thu, 9 Jul 2020 17:58:37 +0000 Subject: [PATCH 22/25] use IRModule directly --- src/relay/analysis/get_calibration_data.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index 47e060b895e6..adc0fb1e92ff 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -45,13 +45,14 @@ namespace relay { class Collector : public ExprRewriter { public: - explicit Collector(const IRModule& module) { glob_funcs_ = module->functions; } + explicit Collector(const IRModule& module) : module_(module) {} 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()) { auto var = Downcast(call->op); + auto glob_funcs_ = module_->functions; 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(); @@ -68,7 +69,7 @@ class Collector : public ExprRewriter { Array GetNewOutputs() { return new_outputs_; } private: - Map glob_funcs_; + const IRModule& module_; Array new_outputs_; }; @@ -136,13 +137,12 @@ IRModule GetCalibrateModule(IRModule module) { class OutputMapper : public ExprRewriter { public: OutputMapper(Map>* output_map, const IRModule& module, size_t* offset) - : output_map_(output_map), offset_(offset) { - glob_funcs_ = module->functions; - } + : output_map_(output_map), module_(module), offset_(offset) {} Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); + auto glob_funcs_ = module_->functions; 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."; @@ -175,7 +175,7 @@ class OutputMapper : public ExprRewriter { private: Map>* output_map_; - Map glob_funcs_; + const IRModule& module_; size_t* offset_; }; From c1eb082e5dc099a4c89f9e9596449c92a3e75e6c Mon Sep 17 00:00:00 2001 From: seanlatias Date: Thu, 9 Jul 2020 18:12:21 +0000 Subject: [PATCH 23/25] use better APIs according to the review --- src/relay/analysis/get_calibration_data.cc | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index adc0fb1e92ff..c8bbf948240e 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -52,10 +52,9 @@ class Collector : public ExprRewriter { // intrinsic functions are excluded for now if (call->op->IsInstance()) { auto var = Downcast(call->op); - auto glob_funcs_ = module_->functions; - CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not defined"; + CHECK(module_->ContainGlobalVar(var->name_hint)) << "Function " << var << " is not defined"; // we only handle functions with Compiler attribute set - auto* fn = glob_funcs_[var].as(); + auto* fn = module_->Lookup(var).as(); auto func = GetRef(fn); if (func->GetAttr(attr::kCompiler)) { // collect all the inputs and outputs @@ -142,12 +141,11 @@ class OutputMapper : public ExprRewriter { Expr Rewrite_(const CallNode* call, const Expr& post) final { if (call->op->IsInstance()) { auto var = Downcast(call->op); - auto glob_funcs_ = module_->functions; - CHECK_GT(glob_funcs_.count(var), 0) << "Function " << var << " is not defined"; + CHECK(module_->ContainGlobalVar(var->name_hint)) << "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(); + auto* fn = module_->Lookup(var).as(); auto func = GetRef(fn); if (func->GetAttr(attr::kCompiler)) { Array info; @@ -158,7 +156,6 @@ class OutputMapper : public ExprRewriter { // 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(); if (auto* tn = fn->body.as()) { info.push_back(Integer(tn->fields.size())); out_size = tn->fields.size(); From 1cc0c451557c9db7cc7c20fa9f57ef404d787bae Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 13 Jul 2020 11:32:06 +0000 Subject: [PATCH 24/25] apply comments from the reviewer --- src/relay/analysis/get_calibration_data.cc | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/relay/analysis/get_calibration_data.cc b/src/relay/analysis/get_calibration_data.cc index c8bbf948240e..34d0d0002b6a 100644 --- a/src/relay/analysis/get_calibration_data.cc +++ b/src/relay/analysis/get_calibration_data.cc @@ -21,7 +21,7 @@ * \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 + * steps. First, we need to prepare the module that generates * the tensor values (GetCalibrateModule). Second, we need to * generate the mapping between the values and the functions * (GetCalibrateOutputMap). @@ -54,8 +54,7 @@ class Collector : public ExprRewriter { auto var = Downcast(call->op); CHECK(module_->ContainGlobalVar(var->name_hint)) << "Function " << var << " is not defined"; // we only handle functions with Compiler attribute set - auto* fn = module_->Lookup(var).as(); - auto func = GetRef(fn); + auto func = Downcast(module_->Lookup(var)); if (func->GetAttr(attr::kCompiler)) { // collect all the inputs and outputs for (const auto& it : call->args) new_outputs_.push_back(it); @@ -144,9 +143,8 @@ class OutputMapper : public ExprRewriter { CHECK(module_->ContainGlobalVar(var->name_hint)) << "Function " << var << " is not defined"; CHECK_EQ(output_map_->count(var), 0) << "Repeated function call " << var << " is not supported."; + auto func = Downcast(module_->Lookup(var)); // we only handle functions with Compiler attribute set - auto* fn = module_->Lookup(var).as(); - auto func = GetRef(fn); if (func->GetAttr(attr::kCompiler)) { Array info; // the first value is the offset @@ -156,7 +154,7 @@ class OutputMapper : public ExprRewriter { // the third value is the number of outputs // we need to check if the output is a tuple size_t out_size = 1; - if (auto* tn = fn->body.as()) { + if (auto* tn = func->body.as()) { info.push_back(Integer(tn->fields.size())); out_size = tn->fields.size(); } else { From 79fd577877bed6e32b54119e89e21e038620b336 Mon Sep 17 00:00:00 2001 From: seanlatias Date: Mon, 13 Jul 2020 15:37:44 +0000 Subject: [PATCH 25/25] retrigger ci