Propagate type annotations in the exact matcher

.gitmodules

@@ -6,10 +6,10 @@
 	url = https://github.com/dmlc/dlpack
 [submodule "3rdparty/rang"]
 	path = 3rdparty/rang
-	url = https://github.com/agauniyal/rang
-[submodule "3rdparty/vta-hw"]
-	path = 3rdparty/vta-hw
-	url = https://github.com/apache/incubator-tvm-vta
+  url = https://github.com/agauniyal/rang
 [submodule "3rdparty/libbacktrace"]
 	path = 3rdparty/libbacktrace
-	url = https://github.com/tlc-pack/libbacktrace.git
+  url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "3rdparty/vta-hw"]
+	path = 3rdparty/vta-hw
+	url = git@github.com:uwsampl/3la-vta.git

CMakeLists.txt

@@ -64,6 +64,8 @@ tvm_option(USE_BLAS "The blas library to be linked" none)
 tvm_option(USE_MKL "MKL root path when use MKL blas" OFF)
 tvm_option(USE_MKLDNN "Build with MKLDNN" OFF)
 tvm_option(USE_DNNL_CODEGEN "Enable MKLDNN (DNNL) codegen" OFF)
+tvm_option(USE_ILAVTA_CODEGEN "Enable ILA codegen for VTA" OFF)
+tvm_option(USE_ILAFLEX_CODEGEN "Enable ILA codegen for FlexNLP" OFF)
 tvm_option(USE_CUDNN "Build with cuDNN" OFF)
 tvm_option(USE_CUBLAS "Build with cuBLAS" OFF)
 tvm_option(USE_THRUST "Build with Thrust" OFF)
@@ -378,6 +380,9 @@ include(cmake/modules/contrib/EthosN.cmake)
 include(cmake/modules/contrib/BLAS.cmake)
 include(cmake/modules/contrib/CODEGENC.cmake)
 include(cmake/modules/contrib/DNNL.cmake)
+include(cmake/modules/contrib/ILAVTA.cmake)
+include(cmake/modules/contrib/ILAFlex.cmake)
+include(cmake/modules/contrib/ILACNN.cmake)
 include(cmake/modules/contrib/Random.cmake)
 include(cmake/modules/contrib/Posit.cmake)
 include(cmake/modules/contrib/MicroStandaloneRuntime.cmake)

README.md

@@ -15,8 +15,41 @@
 <!--- specific language governing permissions and limitations -->
 <!--- under the License. -->
 
-<img src=https://raw.githubusercontent.com/apache/tvm-site/main/images/logo/tvm-logo-small.png width=128/> Open Deep Learning Compiler Stack
-==============================================
+This is a fork of TVM for adding BYOC integrations for the 3LA project.
+
+Right now we have a VTA integration in `src/relay/backend/contrib/ilavta`. Note that you have to include the line `SET(USE_ILAVTA_CODEGEN ON)` in `build/config.cmake` before building TVM to support this (other flags that should be on: `USE_LLVM`, `USE_VTA_FSIM`). We have a test of this backend in `tests/python/relay/test_external_codegen.py` (see `test_extern_vta()`).
+
+This version also uses a fork of the VTA repo meant to dump logs.
+Try `vta/python/integration/matmul_tutorial.py` to use the dumping facility.
+VTA can be set into dumping mode by calling `vta.testing.simulator.dump_mode(True)`.
+You can specify the location at which the dump will be deposited using `vta.testing.simulator.dump_target(path)`; the default is `./vta_sim_dump.json`.
+See the readme at [the VTA fork](https://github.com/uwsampl/3la-vta) to see a description of the dumping mode and the dumping format.
+
+You can use `vta.testing.ila_converter.convert(dump_file, dest_file)` to convert a VTA simulator dump into an ILA program fragment.
+
+# 3LA environment setup
+
+## Docker setup
+Please follow the instruction [here](https://github.com/PrincetonUniversity/3la-integrate) to set up the 3LA integrated docker container.
+
+To attach to the container, run `sudo docker exec -it <name of the container> /bin/bash`
+
+Before running any 3LA related test, `source init.sh` under `/root` first.
+
+## 3LA tvm setup
+Please follow the steps [here](https://tvm.apache.org/docs/install/from_source.html#developers-get-source-from-github). Note to replace the github repo link to this repo. Then switch to `conv1d-codegen` or `3la-rebase-complete` branch.
+
+Before running `cmake`, please add the following lines to `config.cmake` 
+```cmake
+set(USE_ILAVTA_CODEGEN ON)
+set(USE_ILACNN_CODEGEN ON)
+set(USE_ILAFLEX_CODEGEN ON)
+```
+and then set `USE_LLVM` to `ON`.
+
+Before installing the python interface of this variant of tvm, you probably need to uninstall the tvm  that was installed when building the docker image (to do so, run `pip uninstall <package>`).
+
+<img src=https://raw.githubusercontent.com/apache/incubator-tvm-site/main/images/logo/tvm-logo-small.png width=128/> Open Deep Learning Compiler Stack
 [Documentation](https://tvm.apache.org/docs) |
 [Contributors](CONTRIBUTORS.md) |
 [Community](https://tvm.apache.org/community) |

cmake/modules/contrib/ILACNN.cmake

@@ -0,0 +1,9 @@
+if(USE_ILACNN_CODEGEN STREQUAL "ON")
+  add_definitions(-DUSE_ILACNN_RUNTIME=1)
+  file(GLOB ILACNN_RELAY_CONTRIB_SRC src/relay/backend/contrib/ilacnn/*.cc)
+  list(APPEND COMPILER_SRCS ${ILACNN_RELAY_CONTRIB_SRC})
+  list(APPEND COMPILER_SRCS ${JSON_RELAY_CONTRIB_SRC})
+
+  file(GLOB ILACNN_CONTRIB_SRC src/runtime/contrib/ilacnn/ilacnn_runtime.cc)
+  list(APPEND RUNTIME_SRCS ${ILACNN_CONTRIB_SRC})
+endif()

cmake/modules/contrib/ILAFlex.cmake

@@ -0,0 +1,9 @@
+if(USE_ILAFLEX_CODEGEN STREQUAL "ON") 
+  add_definitions(-DUSE_ILAFLEX_RUNTIME=1)
+  file(GLOB ILAFLEX_RELAY_CONTRIB_SRC src/relay/backend/contrib/ilaflex/*.cc)
+  list(APPEND COMPILER_SRCS ${ILAFLEX_RELAY_CONTRIB_SRC})
+  list(APPEND COMPILER_SRCS ${JSON_RELAY_CONTRIB_SRC})
+
+  file(GLOB ILAFLEX_CONTRIB_SRC src/runtime/contrib/ilaflex/ilaflex_runtime.cc)
+  list(APPEND RUNTIME_SRCS ${ILAFLEX_CONTRIB_SRC})
+endif()

cmake/modules/contrib/ILAVTA.cmake

@@ -0,0 +1,34 @@
+if(USE_ILAVTA_CODEGEN STREQUAL "ON") 
+  include_directories(BEFORE SYSTEM ${VTA_HW_PATH}/include)
+  add_definitions(-DUSE_ILAVTA_RUNTIME=1)
+  file(GLOB ILAVTA_RELAY_CONTRIB_SRC src/relay/backend/contrib/ilavta/*.cc)
+  list(APPEND COMPILER_SRCS ${ILAVTA_RELAY_CONTRIB_SRC})
+  list(APPEND COMPILER_SRCS ${JSON_RELAY_CONTRIB_SRC})
+
+  file(GLOB ILAVTA_CONTRIB_SRC src/runtime/contrib/ilavta/ilavta_runtime.cc)
+  list(APPEND ILAVTA_CONTRIB_SRC src/runtime/contrib/ilavta/ilavta_helpers.cc)
+  file(GLOB VTA_RUNTIME_SRCS ${VTA_HW_PATH}/src/*.cc)
+  list(APPEND VTA_RUNTIME_SRCS ${VTA_HW_PATH}/src/sim/sim_driver.cc)
+  list(APPEND VTA_RUNTIME_SRCS ${VTA_HW_PATH}/src/sim/sim_tlpp.cc)
+  list(APPEND VTA_RUNTIME_SRCS ${VTA_HW_PATH}/src/vmem/virtual_memory.cc)
+
+  list(APPEND RUNTIME_SRCS ${ILAVTA_CONTRIB_SRC})
+  list(APPEND RUNTIME_SRCS ${VTA_RUNTIME_SRCS})
+
+  set(VTA_CONFIG ${PYTHON} ${VTA_HW_PATH}/config/vta_config.py)
+
+  if(EXISTS ${CMAKE_CURRENT_BINARY_DIR}/vta_config.json)
+    message(STATUS "Use VTA config " ${CMAKE_CURRENT_BINARY_DIR}/vta_config.json)
+    set(VTA_CONFIG ${PYTHON} ${VTA_HW_PATH}/config/vta_config.py
+      --use-cfg=${CMAKE_CURRENT_BINARY_DIR}/vta_config.json)
+  endif()
+  execute_process(COMMAND ${VTA_CONFIG} --target OUTPUT_VARIABLE VTA_TARGET OUTPUT_STRIP_TRAILING_WHITESPACE)
+  message(STATUS "Build VTA runtime with target: " ${VTA_TARGET})
+  execute_process(COMMAND ${VTA_CONFIG} --defs OUTPUT_VARIABLE __vta_defs)
+  string(REGEX MATCHALL "(^| )-D[A-Za-z0-9_=.]*" VTA_DEFINITIONS "${__vta_defs}")
+
+  foreach(__def ${VTA_DEFINITIONS})
+    string(SUBSTRING ${__def} 3 -1 __strip_def)
+    add_definitions(-D${__strip_def})
+  endforeach()
+endif()

include/tvm/relay/attrs/nn.h

@@ -1097,7 +1097,7 @@ struct UpSampling3DAttrs : public tvm::AttrsNode<UpSampling3DAttrs> {
 /*! \brief Attributes used for the padding operator */
 struct PadAttrs : public tvm::AttrsNode<PadAttrs> {
   Array<Array<Integer>> pad_width;
-  std::string pad_mode;
+  tvm::String pad_mode;
 
   TVM_DECLARE_ATTRS(PadAttrs, "relay.attrs.PadAttrs") {
     TVM_ATTR_FIELD(pad_width).describe(

include/tvm/relay/attrs/transform.h

@@ -33,6 +33,20 @@
 namespace tvm {
 namespace relay {
 
+struct WindowsAttrs : public tvm::AttrsNode<WindowsAttrs> {
+  int axis;
+  Array<Integer> window_shape;
+  Array<Integer> strides;
+  TVM_DECLARE_ATTRS(WindowsAttrs, "relay.attrs.WindowsAttrs") {
+    TVM_ATTR_FIELD(axis).describe(
+        "What axis the windows begin forming over.");
+    TVM_ATTR_FIELD(window_shape).describe(
+        "The window shape to form over the input.");
+    TVM_ATTR_FIELD(strides).describe(
+        "How to stride the windows.");
+  }
+};
+
 /*! \brief data type cast */
 struct CastAttrs : public tvm::AttrsNode<CastAttrs> {
   DataType dtype;
@@ -154,7 +168,7 @@ struct GatherNDAttrs : public tvm::AttrsNode<GatherNDAttrs> {
 struct TakeAttrs : public tvm::AttrsNode<TakeAttrs> {
   Integer batch_dims;
   Integer axis;
-  std::string mode;
+  tvm::String mode;
 
   TVM_DECLARE_ATTRS(TakeAttrs, "relay.attrs.TakeAttrs") {
     TVM_ATTR_FIELD(batch_dims)
@@ -302,7 +316,7 @@ struct StridedSliceAttrs : public tvm::AttrsNode<StridedSliceAttrs> {
   Optional<Array<Integer>> begin;
   Optional<Array<Integer>> end;
   Optional<Array<Integer>> strides;
-  std::string slice_mode;
+  tvm::String slice_mode;
 
   TVM_DECLARE_ATTRS(StridedSliceAttrs, "relay.attrs.StridedSliceAttrs") {
     TVM_ATTR_FIELD(begin).describe("Indices for begin of slice, begin index is also inclusive");
@@ -478,6 +492,19 @@ struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
   }
 };  // struct UniqueAttrs
 
+/*! \brief Attributes for calling accelerators */
+struct AcceleratorCallAttrs : public tvm::AttrsNode<AcceleratorCallAttrs> {
+  std::string func_name;
+  Array<Integer> output_shape;
+  DataType output_dtype;
+  TVM_DECLARE_ATTRS(AcceleratorCallAttrs, "relay.attrs.AcceleratorCallAttrs") {
+    TVM_ATTR_FIELD(func_name).describe("The name of the accelerator function").set_default("unknown");
+    TVM_ATTR_FIELD(output_shape).describe("Inferred output shape for the accelerator call");
+    TVM_ATTR_FIELD(output_dtype).describe("Inferred / annotated output data type of the accelerator call")
+                                .set_default(NullValue<DataType>());
+  }
+};  // struct AcceleratorCallAttrs
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_TRANSFORM_H_