comaniac commented on a change in pull request #6343: URL: https://github.com/apache/incubator-tvm/pull/6343#discussion_r495280239
########## File path: python/tvm/contrib/target/vitis_ai.py ########## @@ -0,0 +1,138 @@ +# 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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel + +"""Utility to offload (sub-)models to Vitis-AI""" + +import warnings + +import pyxir +import pyxir.frontend.tvm + +from tvm.relay.expr import Tuple, Call, TupleGetItem +import tvm._ffi + + +class CodegenVitisAI: + + """Traverse Relay expression and convert into PyXIR XGraph format""" + + def __init__(self, model_name, function): + self.model_name = model_name + self.function = function + self.params = {} + + def convert_pyxir(self, target): + """Convert Relay expression to PyXIR XGraph""" + xgraph = pyxir.frontend.tvm.from_relay(self.function, + params=self.params, postprocessing=None) + xgraph = pyxir.partition(xgraph, targets=[target]) + return xgraph + + def get_output_names(self): + """Get output names from Relay expression""" + func = self.function + output_relay_ids = [] + expr = func.body + if isinstance(expr, Tuple): + for field in expr.fields: + output_relay_ids.append(hash(field)) + elif isinstance(expr, Call): + output_relay_ids.append(hash(expr)) + elif isinstance(expr, TupleGetItem): + output_relay_ids.append(hash(expr.tuple_value)) + else: + raise ValueError("Vitis-AI codegen does not support {} as output".format(type(expr))) + return output_relay_ids + + +@tvm._ffi.register_func("relay.ext.vitis_ai") +def vitis_ai_compiler(ref): + """Create a Vitis-AI runtime from the provided Relay expression""" + assert isinstance(ref, tvm.relay.function.Function) + + out_tensor_names = [] + name = str(ref.attrs.global_symbol) + + pass_context = tvm.get_global_func("transform.GetCurrentPassContext")() + + # The target Vitis-AI accelerator device + target = str(pass_context.config['relay.ext.vitis_ai.options.target']) \ + if 'relay.ext.vitis_ai.options.target' in pass_context.config else None + + # (Optional configs) The build and work directories to be used by Vitis-AI + vai_build_dir = str(pass_context.config['relay.ext.vitis_ai.options.build_dir']) \ + if 'relay.ext.vitis_ai.options.build_dir' in pass_context.config else \ + tvm.contrib.util.tempdir().relpath("") + vai_work_dir = str(pass_context.config['relay.ext.vitis_ai.options.work_dir']) \ + if 'relay.ext.vitis_ai.options.work_dir' in pass_context.config else \ + tvm.contrib.util.tempdir().relpath("") + + # (Optional configs) Export and load PyXIR runtime module to file if provided. This is used to + # compile and quantize a model on the host and deploy it at the edge + export_runtime_module = \ + str(pass_context.config['relay.ext.vitis_ai.options.export_runtime_module']) \ + if 'relay.ext.vitis_ai.options.export_runtime_module' in pass_context.config else "" + load_runtime_module = \ + str(pass_context.config['relay.ext.vitis_ai.options.load_runtime_module']) \ + if 'relay.ext.vitis_ai.options.load_runtime_module' in pass_context.config else "" + + # Config checks + if load_runtime_module != "" and target is not None: Review comment: ```suggestion if load_runtime_module and target is not None: ``` ditto to others. ########## File path: python/tvm/contrib/target/vitis_ai.py ########## @@ -0,0 +1,138 @@ +# 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. +# pylint: disable=invalid-name, unused-argument, import-outside-toplevel + +"""Utility to offload (sub-)models to Vitis-AI""" + +import warnings + +import pyxir +import pyxir.frontend.tvm + +from tvm.relay.expr import Tuple, Call, TupleGetItem +import tvm._ffi + + +class CodegenVitisAI: + + """Traverse Relay expression and convert into PyXIR XGraph format""" + + def __init__(self, model_name, function): + self.model_name = model_name + self.function = function + self.params = {} + + def convert_pyxir(self, target): + """Convert Relay expression to PyXIR XGraph""" + xgraph = pyxir.frontend.tvm.from_relay(self.function, + params=self.params, postprocessing=None) + xgraph = pyxir.partition(xgraph, targets=[target]) + return xgraph + + def get_output_names(self): + """Get output names from Relay expression""" + func = self.function + output_relay_ids = [] + expr = func.body + if isinstance(expr, Tuple): + for field in expr.fields: + output_relay_ids.append(hash(field)) + elif isinstance(expr, Call): + output_relay_ids.append(hash(expr)) + elif isinstance(expr, TupleGetItem): + output_relay_ids.append(hash(expr.tuple_value)) + else: + raise ValueError("Vitis-AI codegen does not support {} as output".format(type(expr))) + return output_relay_ids + + +@tvm._ffi.register_func("relay.ext.vitis_ai") +def vitis_ai_compiler(ref): + """Create a Vitis-AI runtime from the provided Relay expression""" + assert isinstance(ref, tvm.relay.function.Function) + + out_tensor_names = [] + name = str(ref.attrs.global_symbol) + + pass_context = tvm.get_global_func("transform.GetCurrentPassContext")() + + # The target Vitis-AI accelerator device + target = str(pass_context.config['relay.ext.vitis_ai.options.target']) \ + if 'relay.ext.vitis_ai.options.target' in pass_context.config else None + + # (Optional configs) The build and work directories to be used by Vitis-AI + vai_build_dir = str(pass_context.config['relay.ext.vitis_ai.options.build_dir']) \ + if 'relay.ext.vitis_ai.options.build_dir' in pass_context.config else \ + tvm.contrib.util.tempdir().relpath("") + vai_work_dir = str(pass_context.config['relay.ext.vitis_ai.options.work_dir']) \ + if 'relay.ext.vitis_ai.options.work_dir' in pass_context.config else \ + tvm.contrib.util.tempdir().relpath("") + + # (Optional configs) Export and load PyXIR runtime module to file if provided. This is used to + # compile and quantize a model on the host and deploy it at the edge + export_runtime_module = \ + str(pass_context.config['relay.ext.vitis_ai.options.export_runtime_module']) \ + if 'relay.ext.vitis_ai.options.export_runtime_module' in pass_context.config else "" + load_runtime_module = \ + str(pass_context.config['relay.ext.vitis_ai.options.load_runtime_module']) \ + if 'relay.ext.vitis_ai.options.load_runtime_module' in pass_context.config else "" + + # Config checks + if load_runtime_module != "" and target is not None: + warnings.warn("Both `load_runtime_module` and `target` configs were specified." + " The `load_runtime_module` points to a prebuilt runtime module with" + " an internal target so the `target` config will be ignored") + if load_runtime_module != "" and 'relay.ext.vitis_ai.options.build_dir' in pass_context.config: + warnings.warn("Both `load_runtime_module` and `build_dir` configs were specified." + " The `load_runtime_module` points to a prebuilt runtime module with" + " an internal build directory so the `build_dir` config will be ignored") + if load_runtime_module != "" and 'relay.ext.vitis_ai.options.work_dir' in pass_context.config: + warnings.warn("Both `load_runtime_module` and `work_dir` configs were specified." + " The `load_runtime_module` points to a prebuilt runtime module with" + " an internal work directory so the `work_dir` config will be ignored") + + + # If load_runtime_module is not set, we will build the PyXIR runtime module from scratch + if load_runtime_module == "": + # Convert Relay expression into XGraph and do partitioning inside PyXIR + builder = CodegenVitisAI(name, ref) + xgraph = builder.convert_pyxir(target) + output_relay_ids = builder.get_output_names() + layers = xgraph.get_layers() + + # Get the output tensor names using XGraph and output Relay ids + out_tensor_names = [] + for layer in layers: + if not layer.internal: + for relay_id in layer.attrs['relay_id']: + if relay_id in output_relay_ids: + out_tensor_names.append(layer.name) + break + if len(out_tensor_names) == 0: Review comment: ```suggestion if not out_tensor_names: ``` ########## File path: src/runtime/contrib/vitis_ai/vitis_ai_runtime.cc ########## @@ -0,0 +1,204 @@ +/* + * 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 vitis_ai_runtime.cc + */ + +#include "vitis_ai_runtime.h" + +#include <tvm/ir/transform.h> +#include <tvm/runtime/registry.h> + +#include <fstream> +#include <streambuf> + +using namespace pyxir::runtime; + +namespace tvm { +namespace runtime { + +/*! \brief The target Vitis-AI accelerator device */ +TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.vitis_ai.options.target", String); +/*! \brief (Optional config) The build directory to be used by Vitis-AI */ +TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.vitis_ai.options.build_dir", String); +/*! \brief (Optional config) The work directory to be used by Vitis-AI */ +TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.vitis_ai.options.work_dir", String); +/*! \brief (Optional config) Export PyXIR runtime module to disk during serialization if provided */ +TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.vitis_ai.options.export_runtime_module", String); +/*! \brief (Optional config) Load PyXIR runtime module from disk */ +TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.vitis_ai.options.load_runtime_module", String); + +VitisAIRuntime::VitisAIRuntime(const std::string& symbol_name, const Array<String> const_names, + const std::string& serialized_rt_mod, + const std::string& export_rt_mod_path) + : symbol_name_(symbol_name), + const_names_(const_names), + export_rt_mod_path_(export_rt_mod_path) { + std::istringstream sstream(serialized_rt_mod); + rt_mod_.reset(new RuntimeModule()); + rt_mod_->deserialize(sstream); + in_tensor_names_ = rt_mod_->get_in_tensor_names(); + out_tensor_names_ = rt_mod_->get_out_tensor_names(); +} + +VitisAIRuntime::VitisAIRuntime(const std::string& symbol_name, const std::string& xgraph_str, + const Array<String> const_names, const std::string& target, + const std::string& build_dir, const std::string& work_dir, + const std::string& export_rt_mod_path) + : symbol_name_(symbol_name), + const_names_(const_names), + export_rt_mod_path_(export_rt_mod_path) { + std::istringstream xgraph_sstream(xgraph_str); + pyxir::XGraphHolder xgraph = std::make_shared<pyxir::graph::XGraph>(""); + pyxir::read(xgraph, xgraph_sstream); + in_tensor_names_ = xgraph->get_input_names(); + out_tensor_names_ = xgraph->get_meta_attr("tvm_out_tensors").get_strings(); + + pyxir::partition(xgraph, std::vector<std::string>{target}, ""); + + pyxir::RunOptionsHolder run_options(new pyxir::runtime::RunOptions()); + run_options->on_the_fly_quantization = true; + run_options->build_dir = build_dir; + if (!work_dir.empty()) run_options->work_dir = work_dir; + rt_mod_ = + pyxir::build_rt(xgraph, target, in_tensor_names_, out_tensor_names_, "vai", run_options); +} + +Module VitisAIRuntimeCreate(const std::string& name, const std::string& xgraph_str, + const std::string& target, const std::string& build_dir, + const std::string& work_dir, const std::string& export_rt_mod_path) { + Array<String> const_vars; + auto exec = make_object<VitisAIRuntime>(name, xgraph_str, const_vars, target, build_dir, work_dir, + export_rt_mod_path); + return Module(exec); +} + +TVM_REGISTER_GLOBAL("tvm.vitis_ai_runtime.from_xgraph").set_body([](TVMArgs args, TVMRetValue* rv) { + *rv = VitisAIRuntimeCreate(args[0], args[1], args[2], args[3], args[4], args[5]); +}); + +Module VitisAIRuntimeCreate(const std::string& name, const std::string& serialized_rt_mod, + const std::string& export_rt_mod_path) { + Array<String> const_vars; + auto exec = make_object<VitisAIRuntime>(name, const_vars, serialized_rt_mod, export_rt_mod_path); + return Module(exec); +} + +TVM_REGISTER_GLOBAL("tvm.vitis_ai_runtime.from_rt_mod").set_body([](TVMArgs args, TVMRetValue* rv) { + std::string load_rt_mod_path = args[1]; + assert(!load_rt_mod_path.empty()); + std::ifstream in_file(load_rt_mod_path); + std::stringstream buffer; + buffer << in_file.rdbuf(); + std::string serialized_rt_mod = buffer.str(); + in_file.close(); + *rv = VitisAIRuntimeCreate(args[0], serialized_rt_mod, args[2]); +}); + +Module VitisAIRuntimeLoadFromBinary(void* strm ) { + dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm); + std::string symbol_name; + std::vector<std::string> const_vars; + std::string serialized_rt_mod; + std::string export_rt_mod_path; + stream->Read(&serialized_rt_mod); + stream->Read(&export_rt_mod_path); + stream->Read(&symbol_name); + stream->Read(&const_vars); + Array<String> const_names; + for (const auto& it : const_vars) { + const_names.push_back(it); + } + auto exec = + make_object<VitisAIRuntime>(symbol_name, const_names, serialized_rt_mod, export_rt_mod_path); + return Module(exec); +} + +TVM_REGISTER_GLOBAL("runtime.module.loadbinary_VitisAIRuntime") + .set_body_typed(VitisAIRuntimeLoadFromBinary); + +void VitisAIRuntime::SaveToBinary(dmlc::Stream* stream) { + std::ostringstream sstream; + rt_mod_->serialize(sstream); + stream->Write(sstream.str()); + stream->Write(export_rt_mod_path_); + stream->Write(symbol_name_); + std::vector<std::string> consts; + for (const auto& it : const_names_) { + consts.push_back(it); + } + stream->Write(consts); + + // If export_runtime_module_ member variable is set, we will additionally export the PyXIR Review comment: ```suggestion // If export_rt_mod_path_ member variable is set, we will additionally export the PyXIR ``` ---------------------------------------------------------------- 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]
