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 2021/04/28 15:08:58 UTC
[GitHub] [tvm] manupa-arm opened a new pull request #7938: Improved MLF to contain workspace info
manupa-arm opened a new pull request #7938:
URL: https://github.com/apache/tvm/pull/7938
Improving the Model Library Format (MLF) 's metadata.json include workspace size information.
Added functionality to calculate workspace by each primfunc and main function. Additionally, the memory taken by constants and IO is also reported -- in case an executor want to copy them to workspace. Moreover, the workspace information required by each primfunc and main is reported in metadata.json in the Model Library Format(MLF).
--
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
[GitHub] [tvm] manupa-arm commented on pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#issuecomment-832805866
Hi @areusch ,
Thanks for the review.
I've addressed your comments and answered questions. PTAL when you have some time.
--
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
[GitHub] [tvm] areusch commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627645390
##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -27,6 +27,8 @@
from ..relay.backend import executor_factory
from ..relay import param_dict
+MAIN_FUNC_NAME_STR = "run"
Review comment:
okay sure, I agree "main" is more appropriate and the model_name becoming a prefix makes more sense in the long run. How about we align with `runtime::symbol::tvm_module_main`?
--
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
[GitHub] [tvm] areusch commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626847192
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
i sort of favor A3 (though not sure we need the "_func" part), since this function essentially replaced GraphExecutor#run.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626689505
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
I kept this to enable a developers from python to enable creation of FunctionInfo objects.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626849123
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
Cool! I like _tvm_<mod_name>_run.
@giuseros any objection?
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
Cool! I like \_tvm\_<mod_name>_run.
@giuseros any objection?
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626863808
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
So there is a subsequent PR from @giuseros that will introduce the mangling of names to include \_tvm\_<mod_name>_ prefix, therefore I will leave this at just "run" for now.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627535977
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
I think its good practice to have a constructor in python for FFI passable items.
The value I see it might be helpful to create unit tests for any future components that consume FunctionInfo.
Yes, I agree it is not obvious the value in this PR -- have a think, Its possible to take it out as its not mandatory to have this.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627530620
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
Review comment:
Alright; I moved it down
--
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
[GitHub] [tvm] manupa-arm commented on pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#issuecomment-833626761
Hi @areusch,
I addressed one comment but I still believe there is a value in retaining the constructor in the python for FFI passable object -- but not stressing that we absolutely need it though good to have. Let me know what you think.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626863808
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
So there is a subsequent PR @giuseros that will introduce the mangling of names to include \_tvm\_<mod_name>_ prefix, therefore I will leave this at just "run" for now.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626736200
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
I've had a discussion with @giuseros . So the AoT currently uses _tvm_<mod_name>_run_func for the main function.
Few options, we can consider :
A1 : run_func
A2 : main_func
A3 : _tvm_<mod_name>_run_func ( + all operators function names as well for consistency)
A4 : _tvm_<mod_name>_main_func ( + all operators function names as well for consistency)
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626688186
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
My perception is that all the functions are in the namespace of mod_name, thus my opinion it should be the main function of <mod_name>.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627646697
##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -27,6 +27,8 @@
from ..relay.backend import executor_factory
from ..relay import param_dict
+MAIN_FUNC_NAME_STR = "run"
Review comment:
Ack
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626736200
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
I've had a discussion with @giuseros . So the AoT currently uses \_tvm\_<mod_name>_run_func for the main function.
Few options, we can consider :
A1 : run_func
A2 : main_func
A3 : \_tvm\_<mod_name>_run_func ( + all operators function names as well for consistency)
A4 : \_tvm\_<mod_name>_main_func ( + all operators function names as well for consistency)
--
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
[GitHub] [tvm] manupa-arm commented on pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#issuecomment-833803491
I think I've addessed the comments, PTAL
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627614476
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
I see; I ll drop it for now until such a time its needed.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626686759
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
Review comment:
I did a different change -- kind of unified the inits. See if you agree.
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
+ std::unordered_map<int, int> device_io;
+ std::unordered_map<int, int> device_consts;
+
+ for (const auto& kv : storage_device_map_) {
+ auto sids = kv.second[0];
+ auto devices = kv.second[1];
+ CHECK_EQ(sids.size(), devices.size());
+ for (uint32_t i = 0; i < sids.size(); i++) {
+ sid_workspace[devices[i]][sids[i]] = 0;
+ device_io[devices[i]] = 0;
+ device_consts[devices[i]] = 0;
+ }
+ }
+
+ for (const auto& kv : storage_device_map_) {
+ auto size_bytes = CalculateRelayExprSizeBytes(kv.first->checked_type());
+ auto sids = kv.second[0];
+ auto devices = kv.second[1];
+ if (kv.first->IsInstance<ConstantNode>()) {
+ for (const auto& dev : devices) {
+ device_consts[dev] += size_bytes;
+ }
+ continue;
+ } else if (kv.first->IsInstance<VarNode>() || kv.first == func->body) {
+ for (const auto& dev : devices) {
+ device_io[dev] += size_bytes;
+ }
+ continue;
+ }
+ for (uint32_t i = 0; i < sids.size(); i++) {
+ if (size_bytes > sid_workspace[devices[i]][sids[i]]) {
Review comment:
Done
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626691300
##########
File path: tests/python/unittest/test_micro_model_library_format.py
##########
@@ -167,11 +199,68 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[
assert "p0" in params
+@tvm.testing.requires_micro
+def test_export_model_library_format_workspace():
+ with utils.TempDirectory.set_keep_for_debug(True):
Review comment:
Done
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627709495
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -577,10 +748,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
Review comment:
Done, actually no KMainFuncStr needed as we agreed to use runtime::symbol::tvm_module_main
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626685988
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
Review comment:
Done
--
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
[GitHub] [tvm] areusch commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627602572
##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -27,6 +27,8 @@
from ..relay.backend import executor_factory
from ..relay import param_dict
+MAIN_FUNC_NAME_STR = "run"
Review comment:
how come it's named MAIN_FUNC_STR? would be good to clarify this constant, i think, and explain our intentions in a comment. AOT_RUN_FUNCTION_NAME or something? would you guys be updating this after https://discuss.tvm.apache.org/t/mini-rfc-name-mangling-in-aot/9907 lands?
in the non-AOT case, there isn't a corresponding function to look to, so it's confusing why we name it "run" as opposed to model name. also, can we update `kMainFunc` to match?
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -577,10 +748,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
Review comment:
can you add a comment to keep in sync with MAIN_FUNC_STR in python/tvm/micro/model_library_format.py? and vice versa there?
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
so since Map is not mutable from Python, I wonder if it would make sense to also make FunctionInfo immutable in Python? I think the main question I have is: what should the API look like? it's easier to figure that out given an example of how it's used. i don't disagree with having a function like this, but since there isn't a way to use it now, someone could come and propose an API change and we wouldn't have any reasoning why to push back.
--
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
[GitHub] [tvm] manupa-arm commented on pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#issuecomment-834802842
Merge?
--
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
[GitHub] [tvm] areusch merged pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch merged pull request #7938:
URL: https://github.com/apache/tvm/pull/7938
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626736200
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
I've had a discussion with @giuseros . So the AoT currently uses _tvm_<mod_name>_run_func for the main function.
Few options, we can consider :
A1 : run_func
A2 : main_func
A3 : \_tvm\_<mod_name>_run_func ( + all operators function names as well for consistency)
A4 : \_tvm\_<mod_name>_main_func ( + all operators function names as well for consistency)
--
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
[GitHub] [tvm] areusch commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r625212228
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
+ std::unordered_map<int, int> device_io;
+ std::unordered_map<int, int> device_consts;
+
+ for (const auto& kv : storage_device_map_) {
+ auto sids = kv.second[0];
+ auto devices = kv.second[1];
+ CHECK_EQ(sids.size(), devices.size());
+ for (uint32_t i = 0; i < sids.size(); i++) {
+ sid_workspace[devices[i]][sids[i]] = 0;
+ device_io[devices[i]] = 0;
+ device_consts[devices[i]] = 0;
+ }
+ }
+
+ for (const auto& kv : storage_device_map_) {
+ auto size_bytes = CalculateRelayExprSizeBytes(kv.first->checked_type());
+ auto sids = kv.second[0];
+ auto devices = kv.second[1];
+ if (kv.first->IsInstance<ConstantNode>()) {
+ for (const auto& dev : devices) {
+ device_consts[dev] += size_bytes;
+ }
+ continue;
+ } else if (kv.first->IsInstance<VarNode>() || kv.first == func->body) {
+ for (const auto& dev : devices) {
+ device_io[dev] += size_bytes;
+ }
+ continue;
+ }
+ for (uint32_t i = 0; i < sids.size(); i++) {
+ if (size_bytes > sid_workspace[devices[i]][sids[i]]) {
Review comment:
can you add a brief note why you want the max here?
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
how come we need this?
##########
File path: tests/python/unittest/test_micro_model_library_format.py
##########
@@ -167,11 +199,68 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[
assert "p0" in params
+@tvm.testing.requires_micro
+def test_export_model_library_format_workspace():
+ with utils.TempDirectory.set_keep_for_debug(True):
Review comment:
nit: remove
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
Review comment:
can you add a one-liner comment to each of these explaining key and value?
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -551,10 +713,14 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, IRModule> lowered_funcs_;
+ /*! \brief lowered funcs */
+ Map<String, FunctionInfo> function_metadata_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
CompileEngine compile_engine_;
+ /*! \brief main function name */
+ const String kMainFuncStr = "main_func";
Review comment:
i wonder if this should be `mod_name`, as passed to GraphExecutorFactory?
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
Review comment:
would suggest to move this to line 268, where it's built up, so that it's clear it's not needed til then.
--
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
[GitHub] [tvm] manupa-arm commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
manupa-arm commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r627614016
##########
File path: python/tvm/micro/model_library_format.py
##########
@@ -27,6 +27,8 @@
from ..relay.backend import executor_factory
from ..relay import param_dict
+MAIN_FUNC_NAME_STR = "run"
Review comment:
Yes, we ll be updating the name to whatever gets agreed in the RFC for AoT case. I think this is not exactly the run function in second thoughts even in the AoT case. The actual run function is wrapping the the main function. Agree with establishing the connection.
For non-AoT case, this should be a name indicating relay "main" function. How about just "main" ? I still feel mod_name is parent namespace of functions (including this and operators)
--
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
[GitHub] [tvm] areusch commented on a change in pull request #7938: Improved MLF to contain workspace info
Posted by GitBox <gi...@apache.org>.
areusch commented on a change in pull request #7938:
URL: https://github.com/apache/tvm/pull/7938#discussion_r626867573
##########
File path: src/relay/backend/utils.cc
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 relay/backend/util.cc
+ * \brief Relay backend utilities.
+ */
+
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+void FunctionInfo::SetWorkspaceSize(Target tgt, tvm::Integer size) {
+ (*this)->workspace_sizes.Set(tgt, size);
+}
+
+TVM_REGISTER_NODE_TYPE(FunctionInfoNode);
+TVM_REGISTER_GLOBAL("relay.backend.FunctionInfo").set_body_typed([]() { return FunctionInfo(); });
Review comment:
why would someone need to do that, though?
##########
File path: src/relay/backend/graph_executor_codegen.cc
##########
@@ -189,9 +192,109 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
targets_ = targets;
}
+ /*!
+ * \brief Calculate the storage required to store the type of relay.Expr
+ *
+ * \param func The relay expr for which the storage is calculated
+ */
+ int64_t CalculateRelayExprSizeBytes(const Type& expr_type) {
+ if (expr_type->IsInstance<TupleTypeNode>()) {
+ auto tuple_type = Downcast<TupleType>(expr_type);
+ int64_t size = 0;
+ for (const auto& field : tuple_type->fields) {
+ size += CalculateRelayExprSizeBytes(field);
+ }
+ return size;
+ }
+ auto tensor_type = expr_type.as<TensorTypeNode>();
+ auto shape = tensor_type->shape;
+ int num_of_elements = 1;
+ for (const auto& dim_index_expr : shape) {
+ if (dim_index_expr->IsInstance<IntImmNode>()) {
+ num_of_elements *= dim_index_expr.as<IntImmNode>()->value;
+ } else {
+ // If shape is dynamic, we cannot calculate workspace in compile time.
+ num_of_elements = 0;
+ }
+ }
+ auto element_size = tensor_type->dtype.bytes();
+ return element_size * num_of_elements;
+ }
+
+ /*!
+ * \brief Update the "main" control function's metadata
+ *
+ * \param func The main function that contains calls to relay primitive functions
+ */
+ void UpdateMainWorkspaceSize(const Function& func) {
+ std::unordered_map<int, std::unordered_map<int, int>> sid_workspace;
+ std::unordered_map<int, int> device_workspace;
Review comment:
i think better is moving the variable declarations closer to where they're used. it makes it clear that device_workspace isn't of concern when analyzing the first code stanza.
--
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