mbaret commented on a change in pull request #5919: URL: https://github.com/apache/incubator-tvm/pull/5919#discussion_r445662045
########## 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: Ah yes :) ---------------------------------------------------------------- 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]
