[GitHub] [incubator-tvm] zhiics opened a new pull request #4258: [WIP][TVM] Bring Your Own Codegen to TVM

[GitBox <gi...@apache.org>]

zhiics opened a new pull request #4258: [WIP][TVM] Bring Your Own Codegen to TVM
URL: https://github.com/apache/incubator-tvm/pull/4258
 
 
   This is a WIP that enables different backends and/or hardware vendors to bring their own codegen tools to TVM. This is the collaboration between @comaniac and me. @jroesch also provided lots of suggestions in the initial design. The RFC can be found here:  https://discuss.tvm.ai/t/bring-your-own-codegen-to-tvm/4501/27
   
   Some high-level design and APIs involve the following parts:
   
   - Graph coloring/annotation
   Providing HW vendors an infra to customize where they want to execute an op.
   Two possible ways are allowed to annotate a graph:
     - Custom pass: users can write a Relay pass to decide how they want to partition graph using (`subgraph_begin` and `subgraph_end` annotations). For example, more sophisticated algorithm could be implemented to annotate the groups of operators.
     - A high-level API is used to help user/vendors to enable a convenient integration
          ```python
          @reg.register_extern_op("nn.conv2d")
          def conv2d(attrs, args, comp):
   	       return get_extern_op(comp, "conv2d")(attrs, args)
       ```
       Each codegen only needs provide the supported operators and they can invoke a separate build pipeline, e.g. `build_extern`, to invoke partitioning. On the completion of this pipeline, the operators that will be offloaded is wrapped with `subgraph_start` and `subgraph_end` annotation. The annotated program will then be sent to the normal `build` pipeline for code and artifacts generation.
       ```python
       @reg.register_extern_op("nn.conv2d")
       def conv2d(attrs, args, comp):
   	    return get_extern_op(comp, "conv2d")(attrs, args)
       ```
         
   - Graph partitioning
   It is a Relay pass that partitions a program into segments that could be executed on various hardware platforms based on the annotations. The current implementation has not fused consecutive subgraphs belonging to the same backend together yet. It will be handled by follow-up PRs.
   
   - Code generation
   Generate code for each segment of a partition Relay program. Each codegen tool is wrapped into a runtime module so that we can leverage the current TVM infra to do serialization and runtime invocation.
   
   FYI, we currently used GCC as an external codegen tool for easy prototyping and verification. It should be removed later when we land it.
   
   cc @tqchen @yzhliu @wweic @broune @soiferj @junrushao1994 @icemelon9 

----------------------------------------------------------------
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


With regards,
Apache Git Services