[Relay][RFC] Relay IR Text Format (#1781)

.gitignore

@@ -209,3 +209,7 @@ tvm_t.*
 
 # patch sentinel
 patched.txt
+
+# Python type checking
+.mypy_cache/
+.pyre/

CMakeLists.txt

@@ -47,6 +47,7 @@ tvm_option(USE_ROCBLAS "Build with ROCM:RoCBLAS" OFF)
 tvm_option(USE_SORT "Build with sort support" OFF)
 tvm_option(USE_NNPACK "Build with nnpack support" OFF)
 tvm_option(USE_RANDOM "Build with random support" OFF)
+tvm_option(USE_ANTLR "Build with ANTLR for Relay parsing" OFF)
 
 # include directories
 include_directories("include")
@@ -183,6 +184,7 @@ include(cmake/modules/Metal.cmake)
 include(cmake/modules/ROCM.cmake)
 include(cmake/modules/SGX.cmake)
 include(cmake/modules/LLVM.cmake)
+include(cmake/modules/ANTLR.cmake)
 include(cmake/modules/contrib/BLAS.cmake)
 include(cmake/modules/contrib/Random.cmake)
 include(cmake/modules/contrib/Sort.cmake)

Jenkinsfile

@@ -98,6 +98,7 @@ stage('Build') {
            echo set\\(USE_GRAPH_RUNTIME ON\\) >> config.cmake
            echo set\\(USE_STACKVM_RUNTIME ON\\) >> config.cmake
            echo set\\(USE_GRAPH_RUNTIME_DEBUG ON\\) >> config.cmake
+           echo set\\(USE_ANTLR ON\\) >> config.cmake
            echo set\\(USE_BLAS openblas\\) >> config.cmake
            echo set\\(CMAKE_CXX_COMPILER g++\\) >> config.cmake
            echo set\\(CMAKE_CXX_FLAGS -Werror\\) >> config.cmake
@@ -133,6 +134,7 @@ stage('Build') {
            echo set\\(USE_LLVM llvm-config-4.0\\) >> config.cmake
            echo set\\(USE_NNPACK ON\\) >> config.cmake
            echo set\\(NNPACK_PATH /NNPACK/build/\\) >> config.cmake
+           echo set\\(USE_ANTLR ON\\) >> config.cmake
            echo set\\(CMAKE_CXX_COMPILER g++\\) >> config.cmake
            echo set\\(CMAKE_CXX_FLAGS -Werror\\) >> config.cmake
            """

cmake/config.cmake

@@ -128,3 +128,6 @@ set(USE_ROCBLAS OFF)
 
 # Whether use contrib sort
 set(USE_SORT OFF)
+
+# Build ANTLR parser for Relay text format
+set(USE_ANTLR OFF)

cmake/modules/ANTLR.cmake

@@ -0,0 +1,28 @@
+if(USE_ANTLR)
+  if(EXISTS /usr/local/lib/antlr-4.7.1-complete.jar)
+    set(ANTLR4 "/usr/local/lib/antlr-4.7.1-complete.jar")
+
+    set(RELAY_PARSER_DIR
+      ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm/relay/grammar)
+
+    set(RELAY_PARSER
+      ${RELAY_PARSER_DIR}/py2/RelayVisitor.py
+      ${RELAY_PARSER_DIR}/py2/RelayParser.py
+      ${RELAY_PARSER_DIR}/py2/RelayLexer.py
+
+      ${RELAY_PARSER_DIR}/py3/RelayVisitor.py
+      ${RELAY_PARSER_DIR}/py3/RelayParser.py
+      ${RELAY_PARSER_DIR}/py3/RelayLexer.py)
+
+    # Generate ANTLR grammar for parsing.
+    add_custom_command(OUTPUT ${RELAY_PARSER}
+      COMMAND $ENV{JAVA_HOME}/bin/java -jar ${ANTLR4} -visitor -no-listener -Dlanguage=Python2 ${RELAY_PARSER_DIR}/Relay.g4 -o ${RELAY_PARSER_DIR}/py2
+      COMMAND $ENV{JAVA_HOME}/bin/java -jar ${ANTLR4} -visitor -no-listener -Dlanguage=Python3 ${RELAY_PARSER_DIR}/Relay.g4 -o ${RELAY_PARSER_DIR}/py3
+      DEPENDS ${RELAY_PARSER_DIR}/Relay.g4
+      WORKING_DIRECTORY ${RELAY_PARSER_DIR})
+
+    add_custom_target(relay_parser ALL DEPENDS ${RELAY_PARSER})
+  else()
+    message(FATAL_ERROR "Can't find ANTLR4!")
+  endif()
+endif(USE_ANTLR)

docker/Dockerfile.ci_cpu

@@ -40,10 +40,3 @@ COPY install/ubuntu_install_nnpack.sh /install/ubuntu_install_nnpack.sh
 RUN bash /install/ubuntu_install_nnpack.sh
 
 ENV PATH $PATH:$CARGO_HOME/bin:/usr/lib/go-1.10/bin
-
-# ANTLR deps
-COPY install/ubuntu_install_java.sh /install/ubuntu_install_java.sh
-RUN bash /install/ubuntu_install_java.sh
-
-COPY install/ubuntu_install_antlr.sh /install/ubuntu_install_antlr.sh
-RUN bash /install/ubuntu_install_antlr.sh

docker/install/ubuntu_install_antlr.sh

@@ -1,5 +1,3 @@
 cd /usr/local/lib
 wget https://www.antlr.org/download/antlr-4.7.1-complete.jar
 cd -
-
-alias antlr4='java -jar /usr/local/lib/antlr-4.7.1-complete.jar'

python/tvm/relay/__init__.py

@@ -8,6 +8,7 @@
 from . import module
 from . import ir_pass
 from .build_module import build, build_config, create_executor
+from . import parser
 
 # Root operators
 from .op import Op
@@ -52,7 +53,6 @@
 If = expr.If
 TupleGetItem = expr.TupleGetItem
 
-
 # helper functions
 var = expr.var
 const = expr.const
@@ -63,3 +63,6 @@
 def _debug(*args):
     import pdb
     pdb.set_trace()
+
+# Parser
+fromtext = parser.fromtext