manupa-arm commented on a change in pull request #7785:
URL: https://github.com/apache/tvm/pull/7785#discussion_r615725747
##########
File path: src/relay/backend/build_module.cc
##########
@@ -59,17 +64,50 @@ struct BuildOutput {
*/
struct GraphCodegen {
public:
- GraphCodegen() {
- auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
- mod = (*pf)();
+ explicit GraphCodegen(Target target_host) : target_host_(target_host) {
+ const String executor_str =
+ target_host->GetAttr<String>("executor").value_or(kTvmExecutorGraph);
+ if (executor_str == kTvmExecutorGraph) {
+ executor_ = ExecutorType::Graph;
+ auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
+ mod = (*pf)();
+ } else if (executor_str == kTvmExecutorAot) {
Review comment:
I think we'd need a ExecutorCodegen bass class and specialization of
GraphCodegen and AOTCodegen. The switch between the which functionality to run
should be done in a higher-level in the build_module.cc. Therefore we can avoid
the having to check which executor inside each of these functions. WDYT ?
##########
File path: src/relay/backend/aot_codegen.cc
##########
@@ -0,0 +1,704 @@
+/*
+ * 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/graph_codegen.cc
+ * \brief Graph runtime codegen
+ */
+
+#include <dmlc/any.h>
+#include <tvm/ir/module.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/tir/builtin.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt.h>
+
+#include <algorithm>
+#include <list>
+#include <string>
+#include <vector>
+
+#include "../../runtime/meta_data.h"
+#include "compile_engine.h"
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+using IntegerArray = Array<Integer>;
+using ShapeVector = std::vector<std::vector<int64_t>>;
+using GraphAttrs = std::unordered_map<std::string, dmlc::any>;
+using TargetsMap = std::unordered_map<int, Target>;
+
+/*! \brief Lowered outputs */
+struct AOTLoweredOutput {
Review comment:
We should make this a sub-class of LoweredOutput that could be made
common between Graph and AoT codegens. This is essentially the output code each
of the executor codegens. This will help in having generic interfaces in
creation of metadata module and passing around special attributes (e.g.,
runner_func, json) that are specific to each executor.
##########
File path: src/relay/backend/build_module.cc
##########
@@ -101,7 +139,18 @@ struct GraphCodegen {
return ret;
}
+ runtime::AOTMetadata GetAOTMetadata() {
Review comment:
If the relevant information is available in LoweredOutput of
AoTExecutorCodegen, we should not need to get this one seperately -- also look
at my comment where this information is passed onto CreateMetadataModule(...)
##########
File path: src/relay/backend/build_module.cc
##########
@@ -43,12 +43,17 @@ namespace backend {
using TargetsMap = Map<tvm::Integer, tvm::Target>;
using namespace tvm::relay::transform;
+/*!
+ * Type of supported executors
+ */
+enum class ExecutorType { Graph, Aot };
+
/*!
* \brief Output of building module
- *
*/
struct BuildOutput {
std::string graph_json;
+ tir::PrimFunc runner_function;
Review comment:
I am not exactly sure why we need the runner_function in the BuildOutput
because I thought this only needed to be visible to MetadataModule as it will
be pointed to by tvm_model_t.
If we really need to have it, I think its better if we can have a parent
class BuildOutput that contains just runtime.Module and params and specialize
to contain either graph json or runner function. Its confusing to me to have a
struct that has both graph_json and runner_function. WDYT ?
##########
File path: src/relay/backend/build_module.cc
##########
@@ -59,17 +64,50 @@ struct BuildOutput {
*/
struct GraphCodegen {
public:
- GraphCodegen() {
- auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
- mod = (*pf)();
+ explicit GraphCodegen(Target target_host) : target_host_(target_host) {
+ const String executor_str =
+ target_host->GetAttr<String>("executor").value_or(kTvmExecutorGraph);
+ if (executor_str == kTvmExecutorGraph) {
+ executor_ = ExecutorType::Graph;
+ auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
+ mod = (*pf)();
+ } else if (executor_str == kTvmExecutorAot) {
+ executor_ = ExecutorType::Aot;
+ auto pf = GetPackedFunc("relay.build_module._GraphAOTCodegen");
+ mod = (*pf)();
+ } else {
+ LOG(FATAL) << "Executor " << executor_str << " not supported";
+ }
}
~GraphCodegen() {}
- void Init(runtime::Module* m, TargetsMap targets) { CallFunc("init", m,
targets); }
+ void Init(runtime::Module* m, TargetsMap targets) {
+ if (executor_ == ExecutorType::Graph) {
+ CallFunc("init", m, targets);
+ } else if (executor_ == ExecutorType::Aot) {
+ CallFunc("init", m, targets, target_host_);
+ } else {
+ LOG(FATAL) << "Executor not supported";
+ }
+ }
void Codegen(const Function& func) { CallFunc("codegen", func); }
- std::string GetJSON() { return CallFunc<std::string>("get_graph_json",
nullptr); }
+ std::string GetJSON() {
+ if (executor_ == ExecutorType::Graph) {
+ return CallFunc<std::string>("get_graph_json", nullptr);
+ } else {
+ return "";
+ }
+ }
+
+ tir::PrimFunc GetRunnerFunction() {
Review comment:
This should only be implemented in AOTCodegen.
##########
File path: src/relay/backend/build_module.cc
##########
@@ -59,17 +64,50 @@ struct BuildOutput {
*/
struct GraphCodegen {
public:
- GraphCodegen() {
- auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
- mod = (*pf)();
+ explicit GraphCodegen(Target target_host) : target_host_(target_host) {
+ const String executor_str =
+ target_host->GetAttr<String>("executor").value_or(kTvmExecutorGraph);
+ if (executor_str == kTvmExecutorGraph) {
+ executor_ = ExecutorType::Graph;
+ auto pf = GetPackedFunc("relay.build_module._GraphExecutorCodegen");
+ mod = (*pf)();
+ } else if (executor_str == kTvmExecutorAot) {
+ executor_ = ExecutorType::Aot;
+ auto pf = GetPackedFunc("relay.build_module._GraphAOTCodegen");
+ mod = (*pf)();
+ } else {
+ LOG(FATAL) << "Executor " << executor_str << " not supported";
+ }
}
~GraphCodegen() {}
- void Init(runtime::Module* m, TargetsMap targets) { CallFunc("init", m,
targets); }
+ void Init(runtime::Module* m, TargetsMap targets) {
+ if (executor_ == ExecutorType::Graph) {
+ CallFunc("init", m, targets);
+ } else if (executor_ == ExecutorType::Aot) {
+ CallFunc("init", m, targets, target_host_);
+ } else {
+ LOG(FATAL) << "Executor not supported";
+ }
+ }
void Codegen(const Function& func) { CallFunc("codegen", func); }
- std::string GetJSON() { return CallFunc<std::string>("get_graph_json",
nullptr); }
+ std::string GetJSON() {
+ if (executor_ == ExecutorType::Graph) {
+ return CallFunc<std::string>("get_graph_json", nullptr);
+ } else {
Review comment:
I think we can avoid these, if you agree to above comment
##########
File path: src/relay/backend/build_module.cc
##########
@@ -473,15 +540,17 @@ class RelayBuildModule : public runtime::ModuleNode {
// Relay IRModule -> IRModule optimizations.
relay_module = Optimize(relay_module, targets_, params);
+
// Get the updated function.
auto func = Downcast<Function>(relay_module->Lookup("main"));
// Generate code for the updated function.
- graph_codegen_ = std::unique_ptr<GraphCodegen>(new GraphCodegen());
+ graph_codegen_ = std::unique_ptr<GraphCodegen>(new
GraphCodegen(target_host));
Review comment:
We should not overload aot_codegen as graph_codegen -- IMHO, the switch
to use which executor should be reflected somewhere at this source level.
##########
File path: src/relay/backend/aot_codegen.cc
##########
@@ -0,0 +1,704 @@
+/*
+ * 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/graph_codegen.cc
+ * \brief Graph runtime codegen
+ */
+
+#include <dmlc/any.h>
+#include <tvm/ir/module.h>
+#include <tvm/relay/expr_functor.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/tir/builtin.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt.h>
+
+#include <algorithm>
+#include <list>
+#include <string>
+#include <vector>
+
+#include "../../runtime/meta_data.h"
+#include "compile_engine.h"
+#include "utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+
+using IntegerArray = Array<Integer>;
+using ShapeVector = std::vector<std::vector<int64_t>>;
+using GraphAttrs = std::unordered_map<std::string, dmlc::any>;
+using TargetsMap = std::unordered_map<int, Target>;
+
+/*! \brief Lowered outputs */
+struct AOTLoweredOutput {
+ tir::PrimFunc runner_func;
+ Map<String, IRModule> lowered_funcs;
+ Array<tvm::runtime::Module> external_mods;
+ std::unordered_map<std::string, std::pair<int, const tvm::runtime::NDArray>>
params;
+ runtime::AOTMetadata aot_metadata;
+};
+
+class AotReturnSidVisitor : public ExprVisitor {
+ public:
+ explicit AotReturnSidVisitor(Map<Expr, Array<IntegerArray>>
storage_device_map)
+ : storage_device_map_{storage_device_map}, return_sid_{-1} {}
+
+ IntegerArray FindReturnSid(Function func) {
+ VisitExpr(func->body);
+ return return_sid_;
+ }
+
+ protected:
+ void AssignReturnSid(Expr e) {
+ auto iter = storage_device_map_.find(e);
+ if (iter != storage_device_map_.end()) {
+ return_sid_ = (*iter).second[0];
+ }
+ }
+
+ void VisitExpr_(const ConstantNode* cn) override {
+ ExprVisitor::VisitExpr_(cn);
+ AssignReturnSid(GetRef<Expr>(cn));
+ }
+
+ void VisitExpr_(const VarNode* vn) override {
+ ExprVisitor::VisitExpr_(vn);
+ AssignReturnSid(GetRef<Expr>(vn));
+ }
+
+ void VisitExpr_(const CallNode* cn) override {
+ ExprVisitor::VisitExpr_(cn);
+ AssignReturnSid(GetRef<Expr>(cn));
+ }
+
+ void VisitExpr_(const LetNode* op) override { VisitExpr(op->body); }
+
+ void VisitExpr_(const TupleNode* tn) override {
+ ExprVisitor::VisitExpr_(tn);
+ AssignReturnSid(GetRef<Expr>(tn));
+ }
+
+ private:
+ Map<Expr, Array<IntegerArray>> storage_device_map_;
+ IntegerArray return_sid_;
+};
+
+/*! \brief Code generator for AOT executor */
+class AOTCodegen : public ExprVisitor {
+ protected:
+ /*!
+ * \brief Utility function to allocate a DLTensor or TVMValue
+ * \param type the type of allocation
+ * \param num the number of variable to allocate on the stack
+ * \return PrimExpr representing the allocated object
+ */
+ PrimExpr StackAlloca(std::string type, size_t num) {
+ Array<PrimExpr> args = {tir::StringImm(type), ConstInt32(num)};
+ return tir::Call(DataType::Handle(), tir::builtin::tvm_stack_alloca(),
args);
+ }
+
+ /*!
+ * \brief Utility function to allocate memory for storage identifiers
+ * \param memory_size_byte size in bytes of the allocation
+ * \return PrimExpr representing the allocated memory
+ */
+ PrimExpr AllocateBackendMemory(int memory_size_byte) {
+ // TODO(giuseros): use tir::Allocate instead of TVMBackendAllocWorkspace
Review comment:
@giuseros , if its not too much trouble lets move to tir.allocates -- I
think its a very small change to require its own PR (if I am not mistaken)
##########
File path: src/relay/backend/build_module.cc
##########
@@ -524,7 +593,8 @@ class RelayBuildModule : public runtime::ModuleNode {
}
auto ext_mods = graph_codegen_->GetExternalModules();
- ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod,
ext_mods, GetTargetHost());
+ ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod,
ext_mods, GetTargetHost(),
+
graph_codegen_->GetAOTMetadata());
Review comment:
It seems a bit ad-hoc to include just the AoTMetadata here. I'd suggest
to include LoweredOutput here. Thus, that way inside CreateMetadataModule -->
it will know how to consume the output of ExecutorCodegen to create correct
metadata module. WDYT ?
--
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:
[email protected]
