You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/06/25 11:38:37 UTC
[GitHub] [incubator-tvm] mbaret commented on a change in pull request #5919: [BYOC] JSON Runtime with DNNL End-to-End Flow
mbaret commented on a change in pull request #5919:
URL: https://github.com/apache/incubator-tvm/pull/5919#discussion_r445492584
##########
File path: src/relay/backend/graph_runtime_codegen.cc
##########
@@ -661,6 +660,8 @@ struct Handler<std::unordered_map<std::string, dmlc::any>> {
writer->WriteObjectKeyValue(k, dmlc::get<std::vector<std::vector<int64_t>>>(v));
} else if (SameType<std::vector<std::string>>(v)) {
writer->WriteObjectKeyValue(k, dmlc::get<std::vector<std::string>>(v));
+ } else if (SameType<std::vector<dmlc::any>>(v)) {
Review comment:
What does this change do?
##########
File path: src/relay/backend/contrib/codegen_json/codegen_json.h
##########
@@ -0,0 +1,353 @@
+/*
+ * 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/contrib/codegen_json.h
+ * \brief Utilities for json codegen and runtime
+ */
+#ifndef TVM_RELAY_BACKEND_CONTRIB_CODEGEN_JSON_CODEGEN_JSON_H_
+#define TVM_RELAY_BACKEND_CONTRIB_CODEGEN_JSON_CODEGEN_JSON_H_
+
+#include <dmlc/any.h>
+#include <dmlc/json.h>
+#include <tvm/node/container.h>
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/container.h>
+#include <tvm/tir/op.h>
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "../../../../runtime/contrib/json/json_node.h"
+#include "../../../../runtime/contrib/json/json_runtime.h"
+#include "../../utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+namespace contrib {
+
+using namespace tvm::runtime::json;
+
+using ShapeVector = std::vector<std::vector<int64_t>>;
+using TypeVector = std::vector<std::string>;
+using JSONGraphObjectPtr = std::shared_ptr<JSONGraphNode>;
+
+/*!
+ * \brief Helper class to extract all attributes of a certain op and save them
+ * into text format.
+ */
+class OpAttrExtractor : public AttrVisitor {
+ public:
+ explicit OpAttrExtractor(JSONGraphObjectPtr node) : node_(node) {}
+
+ template <typename T = double, typename = std::enable_if_t<std::is_floating_point<T>::value>>
+ std::string Fp2String(const T value) {
+ std::ostringstream out;
+ out.precision(std::numeric_limits<T>::max_digits10);
+ out << value;
+ return out.str();
+ }
+
+ void SetNodeAttr(const char* key, const std::vector<std::string>& value) {
+ std::vector<dmlc::any> attr;
+ attr.emplace_back(value);
+ node_->SetAttr(key, attr);
+ }
+
+ void Visit(const char* key, double* value) final { SetNodeAttr(key, {Fp2String(*value)}); }
+
+ void Visit(const char* key, int64_t* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, uint64_t* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, int* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, bool* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, std::string* value) final { SetNodeAttr(key, {*value}); }
+
+ void Visit(const char* key, DataType* value) final {
+ if (!value->is_void()) {
+ SetNodeAttr(key, {runtime::DLDataType2String(*value)});
+ } else {
+ SetNodeAttr(key, {""});
+ }
+ }
+
+ void Visit(const char* key, runtime::ObjectRef* value) final {
+ if (const auto* an = (*value).as<ArrayNode>()) {
+ std::vector<std::string> attr;
+ for (size_t i = 0; i < an->size(); ++i) {
+ if (const auto* im = (*an)[i].as<IntImmNode>()) {
+ attr.push_back(std::to_string(im->value));
+ } else if (const auto* fm = (*an)[i].as<FloatImmNode>()) {
+ attr.push_back(Fp2String(fm->value));
+ } else if (const auto* str = (*an)[i].as<StringObj>()) {
+ String s = GetRef<String>(str);
+ attr.push_back(s.operator std::string());
+ } else {
+ LOG(FATAL) << "Not supported type: " << (*an)[i]->GetTypeKey();
+ }
+ }
+ SetNodeAttr(key, attr);
+ } else if (!(*value).defined()) { // Skip NullValue
+ SetNodeAttr(key, std::vector<std::string>{""});
+ } else if (const auto* im = (*value).as<IntImmNode>()) {
+ SetNodeAttr(key, std::vector<std::string>{std::to_string(im->value)});
+ } else if (const auto* fm = (*value).as<FloatImmNode>()) {
+ SetNodeAttr(key, std::vector<std::string>{Fp2String(fm->value)});
+ } else if (const auto* str = (*value).as<StringObj>()) {
+ String s = GetRef<String>(str);
+ SetNodeAttr(key, std::vector<std::string>{s.operator std::string()});
+ } else {
+ LOG(FATAL) << "Not yet supported type: " << (*value)->GetTypeKey() << ": " << *value;
+ }
+ }
+
+ void Visit(const char* key, runtime::NDArray* value) final {
+ LOG(FATAL) << "NDArray is not allowed in op attribute";
+ }
+
+ void Visit(const char* key, void** value) final {
+ LOG(FATAL) << "void pointer is not allowed in op attribute";
+ }
+
+ void Extract(Object* node) {
+ if (node) {
+ reflection_->VisitAttrs(node, this);
+ }
+ }
+
+ private:
+ JSONGraphObjectPtr node_;
+ ReflectionVTable* reflection_ = ReflectionVTable::Global();
+};
+
+/*! \brief Serialize a Relay expression to JSON. */
+class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEntry>> {
+ public:
+ /*!
+ * \brief Constructor
+ *
+ * \param symbol The symbol that represents the graph being converted.
+ * \param expr The Relay expression to be converted to the JSON form.
+ */
+ JSONSerializer(const std::string& symbol, const Expr& expr) : symbol_(symbol), func_(expr) {}
+
+ void serialize() {
+ relay::Function func = Downcast<relay::Function>(func_);
+ // First we convert all the parameters into input nodes.
+ for (const auto& param : func->params) {
+ auto node_ptr = std::make_shared<JSONGraphNode>(param->name_hint(), "input" /* op_type_ */);
+ memo_[param] = AddNode(node_ptr, param);
+ }
+ heads_ = VisitExpr(func->body);
+ }
+
+ /*!\brief Return the required params. */
+ Array<String> GetParams() const { return params_; }
+
+ /*!\brief Return the generated json. */
+ std::string GetJSON() {
+ std::ostringstream os;
+ dmlc::JSONWriter writer(&os);
+ Save(&writer);
+ return os.str();
+ }
+
+ protected:
+ /*!
+ * \brief Add a node to graph.
+ *
+ * \param node A graph node. It is a shared pointer. Some attributes of it
+ * will be added, i.e. shape and type. These attributes are attached to
+ * the JSON graph in the end.
+ * \param expr The relay expression.
+ * \return A list of graph entry nodes. It the relay expr is a tuple type, we
+ * will flatten it.
+ */
+ std::vector<JSONGraphNodeEntry> AddNode(JSONGraphObjectPtr node, const Expr& expr) {
+ auto checked_type = expr->checked_type();
+ auto node_id = nodes_.size();
+ nodes_.push_back(node);
+ std::vector<JSONGraphNodeEntry> ret;
+ ShapeVector shape;
+ TypeVector dtype;
+ // Flatten tuple node.
+ if (const auto* tuple_type = checked_type.as<TupleTypeNode>()) {
+ for (size_t i = 0; i < tuple_type->fields.size(); ++i) {
+ const auto* tensor_type = tuple_type->fields[i].as<TensorTypeNode>();
+ CHECK(tensor_type) << "Expect TensorType, but received: ."
+ << tuple_type->fields[i]->GetTypeKey();
+ ret.push_back(JSONGraphNodeEntry(node_id, i));
+ shape.emplace_back(GetIntShape(tensor_type->shape));
+ dtype.emplace_back(DType2String(tensor_type->dtype));
+ }
+ node->SetNumOutput(tuple_type->fields.size());
+ } else {
+ const auto* tensor_type = checked_type.as<TensorTypeNode>();
+ CHECK(tensor_type) << "Expect TensorType, but received: ." << checked_type->GetTypeKey();
Review comment:
"." in the wrong place?
##########
File path: src/runtime/contrib/json/json_runtime.h
##########
@@ -0,0 +1,267 @@
+/*
+ * 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/runtime/json/json_runtime.h
+ * \brief Utilities for json runtime.
+ */
+
+#ifndef TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+
+#include <tvm/runtime/container.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/ndarray.h>
+
+#include <cstddef>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "json_node.h"
+
+namespace tvm {
+namespace runtime {
+namespace json {
+
+/*!
+ * \brief A json runtime that executes the serialized JSON format. This runtime
+ * can be extended by user defined runtime for execution.
+ */
+class JSONRuntimeBase : public ModuleNode {
+ public:
+ JSONRuntimeBase(const std::string& symbol_name, const std::string& graph_json,
+ const Array<String> const_names)
+ : symbol_name_(symbol_name), graph_json_(graph_json), const_names_(const_names) {
+ LoadGraph(graph_json_);
+ }
+
+ const char* type_key() const { return "json"; }
+
+ /*! \brief Initialize a specific json runtime. */
+ virtual void Init(const Array<NDArray>& consts) = 0;
+
+ /*! \brief Invoke the execution engine to inteprete a specific json runtime. */
+ virtual void Run() = 0;
+
+ /*!
+ * \brief Get a packed function.
+ * \param name The name/symbol of the function.
+ * \param sptr_to_self The pointer to the module node.
+ * \return The packed function.
+ */
+ virtual PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "get_symbol") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->symbol_name_; });
+ } else if (name == "get_const_vars") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->const_names_; });
+ } else if (this->symbol_name_ == name) {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK(this->initialized_) << "The module has not been initialized";
+
+ // Bind argument tensors to data entries.
+ this->SetInputOutputBuffers(args);
+ // Execute the subgraph.
+ this->Run();
+ });
+ } else if ("__init_" + this->symbol_name_ == name) {
+ // The function to initialize constant tensors.
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK_EQ(args.size(), 1U);
+ this->Init(args[0]);
+ this->initialized_ = true;
+ *rv = 0;
+ });
+ } else {
+ return PackedFunc(nullptr);
+ }
+ }
+
+ virtual void SaveToBinary(dmlc::Stream* stream) {
+ // Save the symbol
+ stream->Write(symbol_name_);
+ // Save the graph
+ stream->Write(graph_json_);
+ // Save the required const names
+ std::vector<std::string> consts;
+ for (const auto& it : const_names_) {
+ consts.push_back(it);
+ }
+ stream->Write(consts);
+ }
+
+ template <typename T,
+ typename = typename std::enable_if<std::is_base_of<JSONRuntimeBase, T>::value>::type>
+ static Module LoadFromBinary(void* strm) {
+ dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
+ std::string symbol;
+ std::string graph_json;
+ std::vector<std::string> consts;
+ // Load the symbol
+ CHECK(stream->Read(&symbol)) << "Loading symbol name failed";
+ CHECK(stream->Read(&graph_json)) << "Loading graph json failed";
+ CHECK(stream->Read(&consts)) << "Loading the const name list failed";
+ Array<String> const_names;
+ for (const auto& it : consts) {
+ const_names.push_back(it);
+ }
+ auto n = make_object<T>(symbol, graph_json, const_names);
+ return Module(n);
+ }
+
+ protected:
+ /*!
+ * \brief Set up the input and output buffers by binding their DLTensor pointers to the
+ * corresponding data entry.
+ *
+ * \param args The packed args.
+ */
+ void SetInputOutputBuffers(const TVMArgs& args) {
+ CHECK_EQ(args.size(), input_var_idx_.size() + outputs_.size())
+ << "Found mismatch in the number of provided data entryies and required.";
+
+ for (size_t i = 0; i < static_cast<size_t>(args.size()); i++) {
+ auto eid = i < input_var_idx_.size() ? EntryID(input_var_idx_[i], 0)
+ : EntryID(outputs_[i - input_var_idx_.size()]);
+ CHECK(args[i].type_code() == kTVMNDArrayHandle || args[i].type_code() == kTVMDLTensorHandle)
+ << "Expect NDArray or DLTensor as inputs";
+
+ const DLTensor* arg;
+ if (args[i].IsObjectRef<NDArray>()) {
+ NDArray arr = args[i];
+ arg = arr.operator->();
+ } else {
+ arg = args[i].operator DLTensor*();
+ }
+
+ // Assign input/output the NDArray pointers to data entry so that we can directly
+ // read/write host buffers.
+ data_entry_[eid] = arg;
+ }
+ }
+
+ /*!
+ * \brief Load the graph and record the entries for inputs and constants.
+ *
+ * \param graph_json The graph in the json format.
+ */
+ void LoadGraph(const std::string& graph_json) {
+ std::istringstream is(graph_json);
+ dmlc::JSONReader reader(&is);
+ this->Load(&reader);
+ std::vector<std::string> consts;
+ for (size_t i = 0; i < input_nodes_.size(); i++) {
+ uint32_t nid = input_nodes_[i];
+ std::string name = nodes_[nid].name_;
+ if (nodes_[nid].op_type_ == "input") {
+ input_var_idx_.push_back(nid);
+ } else {
+ CHECK_EQ(nodes_[nid].op_type_, "const");
+ auto pos = std::find(std::begin(const_names_), std::end(const_names_), name);
+ CHECK(pos != std::end(const_names_)) << "Found non-existent constant: " << name;
+ const_idx_.push_back(nid);
+ consts.push_back(name);
+ }
+ }
+ CHECK_EQ(consts.size(), const_names_.size())
+ << "Found mismatch for the number of constants in the graph and required.";
+
+ for (size_t i = 0; i < consts.size(); i++) {
+ CHECK_EQ(consts[i], const_names_[i])
+ << "The position of constant in the graph must be the same as the required.";
+ }
+
+ // Reserve data entries.
+ data_entry_.resize(NumEntries());
+ }
+
+ /*!
+ * \brief Set up the constants/weights for inference by binding their DLTensor pointer to
+ * the corresponding data entry.
+ *
+ * \param consts A list of constant NDArray to be used.
+ */
+ void SetupConstants(const Array<NDArray>& consts) {
+ for (size_t i = 0; i < consts.size(); ++i) {
+ data_entry_[const_idx_[i]] = consts[i].operator->();
+ }
+ }
+
+ // Load the graph.
+ void Load(dmlc::JSONReader* reader) {
+ reader->BeginObject();
+ std::string key;
+ while (reader->NextObjectItem(&key)) {
+ if (key == "nodes") {
+ reader->Read(&nodes_);
+ } else if (key == "arg_nodes") {
+ reader->Read(&input_nodes_);
+ } else if (key == "node_row_ptr") {
+ reader->Read(&node_row_ptr_);
+ } else if (key == "heads") {
+ reader->Read(&outputs_);
+ } else {
+ LOG(FATAL) << "Unknow key: " << key;
+ }
+ }
+ }
+
+ // Get the node entry index.
+ uint32_t EntryID(uint32_t nid, uint32_t index) const { return node_row_ptr_[nid] + index; }
+
+ // Get the node entry index.
+ uint32_t EntryID(const JSONGraphNodeEntry& e) const { return EntryID(e.id_, e.index_); }
+
+ // Number of node entries.
+ uint32_t NumEntries() const { return node_row_ptr_.back(); }
+
+ protected:
+ /* The only subgraph name for this module. */
+ std::string symbol_name_;
+ /* The graph. */
+ std::string graph_json_;
+ /* The required constant names. */
Review comment:
Should these have \brief too?
##########
File path: src/runtime/contrib/json/json_runtime.h
##########
@@ -0,0 +1,267 @@
+/*
+ * 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/runtime/json/json_runtime.h
+ * \brief Utilities for json runtime.
+ */
+
+#ifndef TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+
+#include <tvm/runtime/container.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/ndarray.h>
+
+#include <cstddef>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "json_node.h"
+
+namespace tvm {
+namespace runtime {
+namespace json {
+
+/*!
+ * \brief A json runtime that executes the serialized JSON format. This runtime
+ * can be extended by user defined runtime for execution.
+ */
+class JSONRuntimeBase : public ModuleNode {
+ public:
+ JSONRuntimeBase(const std::string& symbol_name, const std::string& graph_json,
+ const Array<String> const_names)
+ : symbol_name_(symbol_name), graph_json_(graph_json), const_names_(const_names) {
+ LoadGraph(graph_json_);
+ }
+
+ const char* type_key() const { return "json"; }
+
+ /*! \brief Initialize a specific json runtime. */
+ virtual void Init(const Array<NDArray>& consts) = 0;
+
+ /*! \brief Invoke the execution engine to inteprete a specific json runtime. */
+ virtual void Run() = 0;
+
+ /*!
+ * \brief Get a packed function.
+ * \param name The name/symbol of the function.
+ * \param sptr_to_self The pointer to the module node.
+ * \return The packed function.
+ */
+ virtual PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "get_symbol") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->symbol_name_; });
+ } else if (name == "get_const_vars") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->const_names_; });
+ } else if (this->symbol_name_ == name) {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK(this->initialized_) << "The module has not been initialized";
+
+ // Bind argument tensors to data entries.
+ this->SetInputOutputBuffers(args);
+ // Execute the subgraph.
+ this->Run();
+ });
+ } else if ("__init_" + this->symbol_name_ == name) {
+ // The function to initialize constant tensors.
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK_EQ(args.size(), 1U);
+ this->Init(args[0]);
+ this->initialized_ = true;
+ *rv = 0;
+ });
+ } else {
+ return PackedFunc(nullptr);
+ }
+ }
+
+ virtual void SaveToBinary(dmlc::Stream* stream) {
+ // Save the symbol
+ stream->Write(symbol_name_);
+ // Save the graph
+ stream->Write(graph_json_);
+ // Save the required const names
+ std::vector<std::string> consts;
+ for (const auto& it : const_names_) {
+ consts.push_back(it);
+ }
+ stream->Write(consts);
+ }
+
+ template <typename T,
+ typename = typename std::enable_if<std::is_base_of<JSONRuntimeBase, T>::value>::type>
+ static Module LoadFromBinary(void* strm) {
+ dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
+ std::string symbol;
+ std::string graph_json;
+ std::vector<std::string> consts;
+ // Load the symbol
+ CHECK(stream->Read(&symbol)) << "Loading symbol name failed";
+ CHECK(stream->Read(&graph_json)) << "Loading graph json failed";
+ CHECK(stream->Read(&consts)) << "Loading the const name list failed";
+ Array<String> const_names;
+ for (const auto& it : consts) {
+ const_names.push_back(it);
+ }
+ auto n = make_object<T>(symbol, graph_json, const_names);
+ return Module(n);
+ }
+
+ protected:
+ /*!
+ * \brief Set up the input and output buffers by binding their DLTensor pointers to the
+ * corresponding data entry.
+ *
+ * \param args The packed args.
+ */
+ void SetInputOutputBuffers(const TVMArgs& args) {
+ CHECK_EQ(args.size(), input_var_idx_.size() + outputs_.size())
+ << "Found mismatch in the number of provided data entryies and required.";
+
+ for (size_t i = 0; i < static_cast<size_t>(args.size()); i++) {
+ auto eid = i < input_var_idx_.size() ? EntryID(input_var_idx_[i], 0)
+ : EntryID(outputs_[i - input_var_idx_.size()]);
+ CHECK(args[i].type_code() == kTVMNDArrayHandle || args[i].type_code() == kTVMDLTensorHandle)
+ << "Expect NDArray or DLTensor as inputs";
+
+ const DLTensor* arg;
+ if (args[i].IsObjectRef<NDArray>()) {
+ NDArray arr = args[i];
+ arg = arr.operator->();
+ } else {
+ arg = args[i].operator DLTensor*();
+ }
+
+ // Assign input/output the NDArray pointers to data entry so that we can directly
+ // read/write host buffers.
+ data_entry_[eid] = arg;
+ }
+ }
+
+ /*!
+ * \brief Load the graph and record the entries for inputs and constants.
+ *
+ * \param graph_json The graph in the json format.
+ */
+ void LoadGraph(const std::string& graph_json) {
+ std::istringstream is(graph_json);
+ dmlc::JSONReader reader(&is);
+ this->Load(&reader);
+ std::vector<std::string> consts;
+ for (size_t i = 0; i < input_nodes_.size(); i++) {
+ uint32_t nid = input_nodes_[i];
+ std::string name = nodes_[nid].name_;
+ if (nodes_[nid].op_type_ == "input") {
+ input_var_idx_.push_back(nid);
+ } else {
+ CHECK_EQ(nodes_[nid].op_type_, "const");
+ auto pos = std::find(std::begin(const_names_), std::end(const_names_), name);
+ CHECK(pos != std::end(const_names_)) << "Found non-existent constant: " << name;
+ const_idx_.push_back(nid);
+ consts.push_back(name);
+ }
+ }
+ CHECK_EQ(consts.size(), const_names_.size())
+ << "Found mismatch for the number of constants in the graph and required.";
+
+ for (size_t i = 0; i < consts.size(); i++) {
+ CHECK_EQ(consts[i], const_names_[i])
+ << "The position of constant in the graph must be the same as the required.";
+ }
+
+ // Reserve data entries.
+ data_entry_.resize(NumEntries());
+ }
+
+ /*!
+ * \brief Set up the constants/weights for inference by binding their DLTensor pointer to
+ * the corresponding data entry.
+ *
+ * \param consts A list of constant NDArray to be used.
+ */
+ void SetupConstants(const Array<NDArray>& consts) {
+ for (size_t i = 0; i < consts.size(); ++i) {
+ data_entry_[const_idx_[i]] = consts[i].operator->();
+ }
+ }
+
+ // Load the graph.
+ void Load(dmlc::JSONReader* reader) {
+ reader->BeginObject();
+ std::string key;
+ while (reader->NextObjectItem(&key)) {
+ if (key == "nodes") {
+ reader->Read(&nodes_);
+ } else if (key == "arg_nodes") {
+ reader->Read(&input_nodes_);
+ } else if (key == "node_row_ptr") {
+ reader->Read(&node_row_ptr_);
+ } else if (key == "heads") {
+ reader->Read(&outputs_);
+ } else {
+ LOG(FATAL) << "Unknow key: " << key;
Review comment:
Unknown
##########
File path: src/relay/backend/contrib/codegen_json/codegen_json.h
##########
@@ -0,0 +1,353 @@
+/*
+ * 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/contrib/codegen_json.h
+ * \brief Utilities for json codegen and runtime
+ */
+#ifndef TVM_RELAY_BACKEND_CONTRIB_CODEGEN_JSON_CODEGEN_JSON_H_
+#define TVM_RELAY_BACKEND_CONTRIB_CODEGEN_JSON_CODEGEN_JSON_H_
+
+#include <dmlc/any.h>
+#include <dmlc/json.h>
+#include <tvm/node/container.h>
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/container.h>
+#include <tvm/tir/op.h>
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "../../../../runtime/contrib/json/json_node.h"
+#include "../../../../runtime/contrib/json/json_runtime.h"
+#include "../../utils.h"
+
+namespace tvm {
+namespace relay {
+namespace backend {
+namespace contrib {
+
+using namespace tvm::runtime::json;
+
+using ShapeVector = std::vector<std::vector<int64_t>>;
+using TypeVector = std::vector<std::string>;
+using JSONGraphObjectPtr = std::shared_ptr<JSONGraphNode>;
+
+/*!
+ * \brief Helper class to extract all attributes of a certain op and save them
+ * into text format.
+ */
+class OpAttrExtractor : public AttrVisitor {
+ public:
+ explicit OpAttrExtractor(JSONGraphObjectPtr node) : node_(node) {}
+
+ template <typename T = double, typename = std::enable_if_t<std::is_floating_point<T>::value>>
+ std::string Fp2String(const T value) {
+ std::ostringstream out;
+ out.precision(std::numeric_limits<T>::max_digits10);
+ out << value;
+ return out.str();
+ }
+
+ void SetNodeAttr(const char* key, const std::vector<std::string>& value) {
+ std::vector<dmlc::any> attr;
+ attr.emplace_back(value);
+ node_->SetAttr(key, attr);
+ }
+
+ void Visit(const char* key, double* value) final { SetNodeAttr(key, {Fp2String(*value)}); }
+
+ void Visit(const char* key, int64_t* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, uint64_t* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, int* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, bool* value) final { SetNodeAttr(key, {std::to_string(*value)}); }
+
+ void Visit(const char* key, std::string* value) final { SetNodeAttr(key, {*value}); }
+
+ void Visit(const char* key, DataType* value) final {
+ if (!value->is_void()) {
+ SetNodeAttr(key, {runtime::DLDataType2String(*value)});
+ } else {
+ SetNodeAttr(key, {""});
+ }
+ }
+
+ void Visit(const char* key, runtime::ObjectRef* value) final {
+ if (const auto* an = (*value).as<ArrayNode>()) {
+ std::vector<std::string> attr;
+ for (size_t i = 0; i < an->size(); ++i) {
+ if (const auto* im = (*an)[i].as<IntImmNode>()) {
+ attr.push_back(std::to_string(im->value));
+ } else if (const auto* fm = (*an)[i].as<FloatImmNode>()) {
+ attr.push_back(Fp2String(fm->value));
+ } else if (const auto* str = (*an)[i].as<StringObj>()) {
+ String s = GetRef<String>(str);
+ attr.push_back(s.operator std::string());
+ } else {
+ LOG(FATAL) << "Not supported type: " << (*an)[i]->GetTypeKey();
+ }
+ }
+ SetNodeAttr(key, attr);
+ } else if (!(*value).defined()) { // Skip NullValue
+ SetNodeAttr(key, std::vector<std::string>{""});
+ } else if (const auto* im = (*value).as<IntImmNode>()) {
+ SetNodeAttr(key, std::vector<std::string>{std::to_string(im->value)});
+ } else if (const auto* fm = (*value).as<FloatImmNode>()) {
+ SetNodeAttr(key, std::vector<std::string>{Fp2String(fm->value)});
+ } else if (const auto* str = (*value).as<StringObj>()) {
+ String s = GetRef<String>(str);
+ SetNodeAttr(key, std::vector<std::string>{s.operator std::string()});
+ } else {
+ LOG(FATAL) << "Not yet supported type: " << (*value)->GetTypeKey() << ": " << *value;
+ }
+ }
+
+ void Visit(const char* key, runtime::NDArray* value) final {
+ LOG(FATAL) << "NDArray is not allowed in op attribute";
+ }
+
+ void Visit(const char* key, void** value) final {
+ LOG(FATAL) << "void pointer is not allowed in op attribute";
+ }
+
+ void Extract(Object* node) {
+ if (node) {
+ reflection_->VisitAttrs(node, this);
+ }
+ }
+
+ private:
+ JSONGraphObjectPtr node_;
+ ReflectionVTable* reflection_ = ReflectionVTable::Global();
+};
+
+/*! \brief Serialize a Relay expression to JSON. */
+class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEntry>> {
+ public:
+ /*!
+ * \brief Constructor
+ *
+ * \param symbol The symbol that represents the graph being converted.
+ * \param expr The Relay expression to be converted to the JSON form.
+ */
+ JSONSerializer(const std::string& symbol, const Expr& expr) : symbol_(symbol), func_(expr) {}
+
+ void serialize() {
+ relay::Function func = Downcast<relay::Function>(func_);
+ // First we convert all the parameters into input nodes.
+ for (const auto& param : func->params) {
+ auto node_ptr = std::make_shared<JSONGraphNode>(param->name_hint(), "input" /* op_type_ */);
+ memo_[param] = AddNode(node_ptr, param);
+ }
+ heads_ = VisitExpr(func->body);
+ }
+
+ /*!\brief Return the required params. */
+ Array<String> GetParams() const { return params_; }
+
+ /*!\brief Return the generated json. */
+ std::string GetJSON() {
+ std::ostringstream os;
+ dmlc::JSONWriter writer(&os);
+ Save(&writer);
+ return os.str();
+ }
+
+ protected:
+ /*!
+ * \brief Add a node to graph.
+ *
+ * \param node A graph node. It is a shared pointer. Some attributes of it
+ * will be added, i.e. shape and type. These attributes are attached to
+ * the JSON graph in the end.
+ * \param expr The relay expression.
+ * \return A list of graph entry nodes. It the relay expr is a tuple type, we
+ * will flatten it.
+ */
+ std::vector<JSONGraphNodeEntry> AddNode(JSONGraphObjectPtr node, const Expr& expr) {
+ auto checked_type = expr->checked_type();
+ auto node_id = nodes_.size();
+ nodes_.push_back(node);
+ std::vector<JSONGraphNodeEntry> ret;
+ ShapeVector shape;
+ TypeVector dtype;
+ // Flatten tuple node.
+ if (const auto* tuple_type = checked_type.as<TupleTypeNode>()) {
+ for (size_t i = 0; i < tuple_type->fields.size(); ++i) {
+ const auto* tensor_type = tuple_type->fields[i].as<TensorTypeNode>();
+ CHECK(tensor_type) << "Expect TensorType, but received: ."
+ << tuple_type->fields[i]->GetTypeKey();
+ ret.push_back(JSONGraphNodeEntry(node_id, i));
+ shape.emplace_back(GetIntShape(tensor_type->shape));
+ dtype.emplace_back(DType2String(tensor_type->dtype));
+ }
+ node->SetNumOutput(tuple_type->fields.size());
+ } else {
+ const auto* tensor_type = checked_type.as<TensorTypeNode>();
+ CHECK(tensor_type) << "Expect TensorType, but received: ." << checked_type->GetTypeKey();
+ shape.emplace_back(GetIntShape(tensor_type->shape));
+ dtype.emplace_back(DType2String(tensor_type->dtype));
+ ret.push_back(JSONGraphNodeEntry(node_id, 0));
+ }
+ std::vector<dmlc::any> shape_attrs;
+ shape_attrs.emplace_back(shape);
+ node->SetAttr("shape", shape_attrs);
+
+ std::vector<dmlc::any> type_attrs;
+ type_attrs.emplace_back(dtype);
+ node->SetAttr("dtype", type_attrs);
+ return ret;
+ }
+
+ void SetCallNodeAttribute(JSONGraphObjectPtr node, const CallNode* cn) {
+ if (cn->op.as<OpNode>()) {
+ OpAttrExtractor extractor(node);
+ const Object* call_attr = cn->attrs.get();
+ extractor.Extract(const_cast<Object*>(call_attr));
+ } else if (const auto* fn = cn->op.as<FunctionNode>()) {
+ auto pattern = fn->GetAttr<String>(attr::kPartitionedFromPattern);
+ CHECK(pattern.defined());
+ std::vector<std::string> values;
+ values.push_back(pattern.value().operator std::string());
+ std::vector<dmlc::any> attr;
+ attr.emplace_back(values);
+ node->SetAttr("PartitionedFromPattern", attr);
Review comment:
Composite should also be set.
##########
File path: src/runtime/metadata_module.cc
##########
@@ -48,15 +48,22 @@ class MetadataModuleNode : public ModuleNode {
public:
MetadataModuleNode(const std::unordered_map<std::string, NDArray>& metadata,
const std::unordered_map<std::string, std::vector<std::string>>& sym_vars)
- : metadata_(metadata), sym_vars_(sym_vars) {}
+ : metadata_(metadata), sym_vars_(sym_vars) {
+ // Only the related submodules are cached to reduce the number of runtime
+ // symbol lookup for initialization. Otherwise, symbols/primitives in the
+ // DSO module will also be cached but they never need to be initialized.
+ for (const auto& it : sym_vars_) {
+ initialized_[it.first] = false;
+ }
+ }
Review comment:
These changes (and the others in this file) seem unrelated to this patch.
##########
File path: src/runtime/contrib/json/json_runtime.h
##########
@@ -0,0 +1,267 @@
+/*
+ * 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/runtime/json/json_runtime.h
Review comment:
src/runtime/contrib/json/json_runtime.h
##########
File path: tests/python/relay/test_json_runtime.py
##########
@@ -0,0 +1,625 @@
+# 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.
+"""Unit tests for JSON codegen and runtime."""
Review comment:
This seems more like it's testing the DNNL codegen and runtime rather than the underlying JSON implementation, to the point no tests will run without DNNL. I think we should aim to steer away from this sort of approach and see if there are ways we can more directly test the implementation.
##########
File path: src/runtime/contrib/json/json_runtime.h
##########
@@ -0,0 +1,267 @@
+/*
+ * 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/runtime/json/json_runtime.h
+ * \brief Utilities for json runtime.
+ */
+
+#ifndef TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_JSON_JSON_RUNTIME_H_
+
+#include <tvm/runtime/container.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/ndarray.h>
+
+#include <cstddef>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "json_node.h"
+
+namespace tvm {
+namespace runtime {
+namespace json {
+
+/*!
+ * \brief A json runtime that executes the serialized JSON format. This runtime
+ * can be extended by user defined runtime for execution.
+ */
+class JSONRuntimeBase : public ModuleNode {
+ public:
+ JSONRuntimeBase(const std::string& symbol_name, const std::string& graph_json,
+ const Array<String> const_names)
+ : symbol_name_(symbol_name), graph_json_(graph_json), const_names_(const_names) {
+ LoadGraph(graph_json_);
+ }
+
+ const char* type_key() const { return "json"; }
+
+ /*! \brief Initialize a specific json runtime. */
+ virtual void Init(const Array<NDArray>& consts) = 0;
+
+ /*! \brief Invoke the execution engine to inteprete a specific json runtime. */
+ virtual void Run() = 0;
+
+ /*!
+ * \brief Get a packed function.
+ * \param name The name/symbol of the function.
+ * \param sptr_to_self The pointer to the module node.
+ * \return The packed function.
+ */
+ virtual PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
+ if (name == "get_symbol") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->symbol_name_; });
+ } else if (name == "get_const_vars") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->const_names_; });
+ } else if (this->symbol_name_ == name) {
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK(this->initialized_) << "The module has not been initialized";
+
+ // Bind argument tensors to data entries.
+ this->SetInputOutputBuffers(args);
+ // Execute the subgraph.
+ this->Run();
+ });
+ } else if ("__init_" + this->symbol_name_ == name) {
+ // The function to initialize constant tensors.
+ return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+ CHECK_EQ(args.size(), 1U);
+ this->Init(args[0]);
+ this->initialized_ = true;
+ *rv = 0;
+ });
+ } else {
+ return PackedFunc(nullptr);
+ }
+ }
+
+ virtual void SaveToBinary(dmlc::Stream* stream) {
+ // Save the symbol
+ stream->Write(symbol_name_);
+ // Save the graph
+ stream->Write(graph_json_);
+ // Save the required const names
+ std::vector<std::string> consts;
+ for (const auto& it : const_names_) {
+ consts.push_back(it);
+ }
+ stream->Write(consts);
+ }
+
+ template <typename T,
+ typename = typename std::enable_if<std::is_base_of<JSONRuntimeBase, T>::value>::type>
+ static Module LoadFromBinary(void* strm) {
+ dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
+ std::string symbol;
+ std::string graph_json;
+ std::vector<std::string> consts;
+ // Load the symbol
+ CHECK(stream->Read(&symbol)) << "Loading symbol name failed";
Review comment:
This implies the module can only hold a single function, and I'm assuming this was the motivation behind the broader changes to the compiler flow to move from IRModule -> Function for external codegen. Are there any downsides/overheads to only allowing single function modules?
----------------------------------------------------------------
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