Skip to content

Commit

Permalink
[BYOC] JSON Runtime with DNNL End-to-End Flow (#5919)
Browse files Browse the repository at this point in the history 
* json runtime

* json dnnl WIP

* fix ArrayNode usages

* Support composite functions

* DNNL json runtime: conv2d/add/relu/dense/bn

* add a more complex example

* fix bias memory issue

* rebase to upstream

* merge to metadata module, remove the unused driver

* handle constant

* support composite functions

* support DNNL constant

* clean up

* Simplify dnnl user code

* GetDataSize

* fix dense bug

* improve cmake

* zero copy

* add unit test

* move json to contrib/json

* fix cmake

* lint

* max_digits10 for fp serialization

* only keep base getfunction

* fix lint

* zero copy for all data entries

* address comments

* enable ci

* address comment; fix bug

* address comment

Co-authored-by: Zhi Chen <chzhi@amazon.com>
  • Loading branch information
@comaniac @zhiics
comaniac and zhiics committed Jul 10, 2020
1 parent ba2b75b commit 8a0249c
Show file tree
Hide file tree
Showing 12 changed files with 2,179 additions and 14 deletions.
19 changes: 15 additions & 4 deletions cmake/modules/contrib/DNNL.cmake
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,25 @@
# specific language governing permissions and limitations
# under the License.

if(USE_DNNL_CODEGEN STREQUAL "ON")
file(GLOB DNNL_RELAY_CONTRIB_SRC src/relay/backend/contrib/dnnl/codegen.cc)
if((USE_DNNL_CODEGEN STREQUAL "ON") OR (USE_DNNL_CODEGEN STREQUAL "JSON"))
add_definitions(-DUSE_JSON_RUNTIME=1)
file(GLOB DNNL_RELAY_CONTRIB_SRC src/relay/backend/contrib/dnnl/*.cc)
list(APPEND COMPILER_SRCS ${DNNL_RELAY_CONTRIB_SRC})
list(APPEND COMPILER_SRCS ${JSON_RELAY_CONTRIB_SRC})

find_library(EXTERN_LIBRARY_DNNL dnnl)
list(APPEND TVM_RUNTIME_LINKER_LIBS ${EXTERN_LIBRARY_DNNL})
file(GLOB DNNL_CONTRIB_SRC src/runtime/contrib/dnnl/*)
file(GLOB DNNL_CONTRIB_SRC src/runtime/contrib/dnnl/dnnl_json_runtime.cc)
list(APPEND RUNTIME_SRCS ${DNNL_CONTRIB_SRC})
message(STATUS "Build with DNNL codegen: " ${EXTERN_LIBRARY_DNNL})
message(STATUS "Build with DNNL JSON runtime: " ${EXTERN_LIBRARY_DNNL})
elseif(USE_DNNL_CODEGEN STREQUAL "C_SRC")
file(GLOB DNNL_RELAY_CONTRIB_SRC src/relay/backend/contrib/dnnl/*.cc)
list(APPEND COMPILER_SRCS ${DNNL_RELAY_CONTRIB_SRC})

find_library(EXTERN_LIBRARY_DNNL dnnl)
list(APPEND TVM_RUNTIME_LINKER_LIBS ${EXTERN_LIBRARY_DNNL})
file(GLOB DNNL_CONTRIB_SRC src/runtime/contrib/dnnl/dnnl.cc)
list(APPEND RUNTIME_SRCS ${DNNL_CONTRIB_SRC})
message(STATUS "Build with DNNL C source module: " ${EXTERN_LIBRARY_DNNL})
endif()

353 changes: 353 additions & 0 deletions src/relay/backend/contrib/codegen_json/codegen_json.h
View file
Original file line number Diff line number Diff line change
@@ -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);
} 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});
} 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());
std::vector<dmlc::any> attr;
attr.emplace_back(values);
node->SetAttr("PartitionedFromPattern", attr);
}
}

std::vector<JSONGraphNodeEntry> VisitExprDefault_(const Object* op) {
LOG(FATAL) << "JSON runtime currently doesn't support " << op->GetTypeKey();
return {};
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const VarNode* vn) {
CHECK(memo_.count(GetRef<Expr>(vn)));
return memo_[GetRef<Expr>(vn)];
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* cn) {
std::string name = symbol_ + "_const_" + std::to_string(params_.size());
params_.push_back(name);
auto node = std::make_shared<JSONGraphNode>(name, "const" /* op_type_ */);
return AddNode(node, GetRef<Expr>(cn));
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const TupleNode* tn) {
std::vector<JSONGraphNodeEntry> fields;
for (const auto& field : tn->fields) {
auto ref = VisitExpr(field);
fields.insert(fields.end(), ref.begin(), ref.end());
}
return fields;
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const CallNode* cn) {
Expr expr = GetRef<Expr>(cn);
std::string name;
if (const auto* op_node = cn->op.as<OpNode>()) {
name = op_node->name;
} else if (const auto* fn = cn->op.as<FunctionNode>()) {
auto comp = fn->GetAttr<String>(attr::kComposite);
CHECK(comp.defined()) << "JSON runtime only supports composite functions.";
name = comp.value();
} else {
LOG(FATAL) << "JSON runtime does not support calls to " << cn->op->GetTypeKey();
}

std::vector<JSONGraphNodeEntry> inputs;
for (const auto& arg : cn->args) {
auto res = VisitExpr(arg);
inputs.insert(inputs.end(), res.begin(), res.end());
}
auto node = std::make_shared<JSONGraphNode>(name, /* name_ */
"kernel", /* op_type_ */
inputs, 1 /* num_outputs_ */);
SetCallNodeAttribute(node, cn);
return AddNode(node, GetRef<Expr>(cn));
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const LetNode* ln) {
CHECK_EQ(memo_.count(ln->var), 0);
memo_[ln->var] = VisitExpr(ln->value);
return VisitExpr(ln->body);
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const TupleGetItemNode* gtn) {
auto vtuple = VisitExpr(gtn->tuple);
return {vtuple[gtn->index]};
}

std::vector<JSONGraphNodeEntry> VisitExpr_(const FunctionNode* fn) {
CHECK(fn->GetAttr<String>(attr::kComposite).defined())
<< "JSON runtime only supports composite functions";
// FunctionNode should be handled by the caller.
return {};
}

/*!
* \brief Save to JSON graph
*
* \param writer A json writer
*/
void Save(dmlc::JSONWriter* writer) {
std::vector<size_t> arg_nodes;
for (size_t i = 0; i < nodes_.size(); ++i) {
auto node = nodes_[i];
if (node->IsLeaf()) {
arg_nodes.push_back(i);
}
}
size_t num_entry = 0;
std::vector<size_t> node_row_ptr{0};
for (auto node : nodes_) {
num_entry += node->GetNumOutput();
node_row_ptr.push_back(num_entry);
}
writer->BeginObject();
writer->WriteObjectKeyValue("nodes", nodes_);
writer->WriteObjectKeyValue("arg_nodes", arg_nodes);
writer->WriteObjectKeyValue("heads", heads_);
writer->WriteObjectKeyValue("node_row_ptr", node_row_ptr);
writer->EndObject();
}

private:
/*! \brief The symbol that represents the json graph. */
std::string symbol_;
/*! \brief The function to be serialized. */
const Expr func_;
/*! \brief JSON graph nodes. */
std::vector<JSONGraphObjectPtr> nodes_;
/*! \brief Output of the JSON graph. */
std::vector<JSONGraphNodeEntry> heads_;
/*! \brief The list of required constants. */
Array<String> params_;
};

} // namespace contrib
} // namespace backend
} // namespace relay
} // namespace tvm
#endif // TVM_RELAY_BACKEND_CONTRIB_CODEGEN_JSON_CODEGEN_JSON_H_

0 comments on commit 8a0249c

Please sign in to comment.