tqchen commented on a change in pull request #4718: [Docs] Bring Your Own
Codegen Guide -- Part 2
URL: https://github.com/apache/incubator-tvm/pull/4718#discussion_r369160146
##########
File path: docs/dev/relay_bring_your_own_codegen.rst
##########
@@ -516,10 +516,417 @@ So that users can configure whether to include your
compiler when configuring TV
Implement a Codegen for Your Representation
*******************************************
-Although we have demonstrated how to implement a C codegen, your hardware may
require other forms of graph representation, such as JSON. In this case, you
can slightly modify ``CodegenC`` class we have implemented to generate your own
graph representation, and implement a customized runtime module to let TVM
runtime know how this graph representation should be executed. **(TBA)**
+Although we have demonstrated how to implement a C codegen, your hardware may
require other forms of graph representation, such as JSON. In this case, you
could modify ``CodegenC`` class we have implemented to generate your own graph
representation and implement a customized runtime module to let TVM runtime
know how this graph representation should be executed.
-Implement CodegenJSON
-=====================
+To simplify, we define a graph representation named "ExampleJSON" in this
guide. ExampleJSON does not mean the real JSON but just a simple representation
for graphs without a control flow. For example, assuming we have the following
subgraph named ``subgraph_0``:
+
+::
+
+ input0
+ |
+ add <-- input1
+ |
+ subtract <-- input2
+ |
+ multiply <-- input3
+ |
+ out
+
+Then the ExampleJON of this subgraph looks like:
+
+.. code-block:: json
+
+ subgraph_0
+ input 0 10 10
+ input 1 10 10
+ input 2 10 10
+ input 3 10 10
+ add 4 inputs: 0 1 shape: 10 10
+ sub 5 inputs: 4 2 shape: 10 10
+ add 6 inputs: 5 3 shape: 10 10
+
+The ``input`` keyword declares an input tensor with its ID and shape; while
the other statements describes computations in ``<op> <output ID> inputs:
[input ID] shape: [shape]`` syntax.
+
+In this section, our goal is to implement the following customized TVM runtime
module to execute ExampleJSON graphs.
+
+.. code-block:: c++
+
+ runtime::Module ExampleJsonCompiler(const NodeRef& ref) {
+ ExampleJsonCodeGen codegen(ref);
+ std::string code = codegen.gen(); // Note 1
+ const auto* pf =
runtime::Registry::Get("module.examplejson_module_create"); // Note 2
+ CHECK(pf != nullptr) << "Cannot find ExampleJson module to create the
external runtime module";
+ return (*pf)(code);
+ }
+
TVM_REGISTER_GLOBAL("relay.ext.examplejsoncompiler").set_body_typed(ExampleJsonCompiler);
+
+**Note 1**: We will implement a customized codegen later to generate a
ExampleJSON code string by taking a subgraph.
+
+**Note 2**: This line obtains a pointer to a function for creating the
customized runtime module. You can see that it takes subgraph code in
ExampleJSON format we just generated and initializes a runtime module.
+
+In the following sections, we are going to introduce 1) how to implement
``ExampleJsonCodeGen`` and 2) how to implement and register
``examplejson_module_create``.
+
+Implement ExampleJsonCodeGen
+============================
+
+Similar to the C codegen, we also derive ``ExampleJsonCodeGen`` from
``ExprVisitor`` to make use of visitor patterns for subgraph traversing. On the
other hand, we do not have to inherit ``CodegenCBase`` because we do not need
TVM C++ wrappers. The codegen class is implemented as follows:
+
+.. code-block:: c++
+
+ #include <tvm/relay/expr_functor.h>
+ #include <tvm/relay/transform.h>
+ #include <tvm/relay/type.h>
+ #include <tvm/runtime/module.h>
+ #include <tvm/runtime/object.h>
+
+ #include <fstream>
+ #include <sstream>
+
+ namespace tvm {
+ namespace relay {
+ namespace contrib {
+
+ class ExampleJsonCodeGen : public ExprVisitor {
+ public:
+ explicit ExampleJsonCodeGen();
+
+ // Note 1
+ void VisitExpr_(const VarNode* node) { /* Skip in this example. */ }
+ void VisitExpr_(const CallNode* call) final { /* Skip in this example.
*/ }
+
+ // Note 2
+ std::string gen(NodeRef& ref) {
+ this->code = "";
+ if (ref->IsInstance<FunctionNode>()) {
+ this->visit(Downcast<Function>(ref));
+ } else if (ref->IsInstance<relay::ModuleNode>()) {
+ relay::Module mod = Downcast<relay::Module>(ref);
+ for (const auto& it : mod->functions) {
+ this->visit(Downcast<Function>(it.second));
+ }
+ } else {
+ LOG(FATAL) << "The input ref is expected to be a Relay
function or module";
+ }
+ return this->code;
+ }
+
+ private:
+ /*! \brief The function id that represents a C source function. */
+ std::string code;
+ }
+
+**Note 1**: We again implement corresponding visitor functions to generate
ExampleJSON code and store it to a class variable ``code`` (we skip the visitor
function implementation in this example as their concepts are basically the
same as C codegen). After finished the graph visiting, we should have an
ExampleJSON graph in ``code``.
+
+**Note 2**: We define an internal API ``gen`` to take a subgraph and generate
a ExampleJSON code. This API can be in an arbitrary name you prefer.
+
+The next step is to implement a customized runtime to make use of the output
of ``ExampleJsonCodeGen``.
+
+Implement a Customized Runtime
+==============================
+
+In this section, we will implement a customized TVM runtime step-by-step and
register it to TVM runtime modules. The customized runtime should be located at
``src/runtime/contrib/<your-runtime-name>/``. In our example, we name our
runtime "example_ext_runtime" and put it under
`here<src/runtime/contrib/example_ext_runtime/example_ext_runtime.cc>`_. Feel
free to check this file for a complete implementation.
+
+Again, we first define a customized runtime class as follows. The class has to
be derived from TVM ``ModuleNode`` in order to be compatible with other TVM
runtime modules.
+
+.. code-block:: c++
+
+ #include <dmlc/logging.h>
+ #include <tvm/runtime/c_runtime_api.h>
+ #include <tvm/runtime/memory.h>
+ #include <tvm/runtime/module.h>
+ #include <tvm/runtime/ndarray.h>
+ #include <tvm/runtime/object.h>
+ #include <tvm/runtime/packed_func.h>
+ #include <tvm/runtime/registry.h>
+
+ #include <fstream>
+ #include <cmath>
+ #include <map>
+ #include <sstream>
+ #include <string>
+ #include <vector>
+
+ namespace tvm {
+ namespace runtime {
+ class ExampleJsonModule : public ModuleNode {
+ public:
+ explicit ExampleJsonModule(std::string graph_json);
+
+ PackedFunc GetFunction(const std::string& name,
+ const ObjectPtr<Object>& sptr_to_self) final;
+
+ const char* type_key() const { return "examplejson"; }
+
+ void SaveToBinary(dmlc::Stream* stream) final;
+
+ static Module LoadFromBinary(void* strm);
+
+ static Module Create(const std::string& path);
+
+ std::string GetSource(const std::string& format = "");
+
+ void Run(int id, const std::vector<int>& inputs, int output);
+
+ void ParseJson(const std::string& json);
+
+ private:
+ /* \brief The json string that represents a computational graph. */
+ std::string graph_json_;
+ /* \brief The subgraph that being processed. */
+ std::string curr_subgraph_;
+ /*! \brief A simple graph from subgraph id to node entries. */
+ std::map<std::string, std::vector<NodeEntry> > graph_;
+ /* \brief A simple pool to contain the tensor for each node in the
graph. */
+ std::vector<NDArray> data_entry_;
+ /* \brief A mapping from node id to op name. */
+ std::vector<std::string> op_id_;
+ };
+
+In particular, there are some functions derived from ``ModuleNode`` that we
must implement in ``ExampleJsonModule``:
+
+* Constructor: The constructor of this class should accept a subgraph (in your
representation), process and store it in any format you like. The saved
subgraph could be used by the following two functions.
+
+* ``GetFunction``: This is the most important function in this class. When TVM
runtime wants to execute a subgraph with your compiler tag, TVM runtime invokes
this function from your customized runtime module. It provides the function
name as well as runtime arguments, and ``GetFunction`` should return a packed
function implementation for TVM runtime to execute.
+
+* ``SaveToBinary`` and ``LoadFromBinary``: ``SaveToBinary`` serialize the
runtime module to a binary format for later deployment. This function will be
called by TVM when users use ``export_library`` API. On the other hand, since
we are now using our own graph representation, we have to make sure that
``LoadFromBinary`` is able to construct the same runtime module by taking the
serialized binary generated by ``SaveToBinary``.
+
+* ``GetSource`` (optional): If you would like to see the generated ExampleJSON
code, you can implement this function to dump it; otherwise you can skip the
implementation.
+
+Other functions and class variables will be introduced along with the
implementation of above must-have functions.
+
+Implement Constructor
+---------------------
+
+.. code-block:: c++
+
+ explicit ExampleJsonModule(std::string graph_json) {
+ this->graph_json_ = graph_json;
+ ParseJson(this->graph_json_);
+ }
+
+Then, we implement ``ParseJson`` to parse a subgraph in ExampleJSON format and
construct a graph in memory for later usage. Since we do not support subgraph
with branches in this example, we simply use an array to store every nodes in a
subgraph in order.
Review comment:
Add a summary section to show a checklist of things todo when adding a
customized runtime and compiler
----------------------------------------------------------------
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]
With regards,
Apache Git Services
