diff --git a/.github/ISSUE_TEMPLATE/flaky-test.md b/.github/ISSUE_TEMPLATE/flaky-test.md
new file mode 100644
index 000000000000..1d61bbb632a4
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/flaky-test.md
@@ -0,0 +1,18 @@
+---
+name: "\U00002744 Flaky Test"
+about: Report a flaky test, make sure to include links to CI runs, a sample failure log, and the name of the test(s)
+title: "[Flaky Test] "
+labels: "test: flaky"
+---
+
+Thanks for participating in the TVM community! We use https://discuss.tvm.ai for any general usage questions and discussions. The issue tracker is used for actionable items such as feature proposals discussion, roadmaps, and bug tracking. You are always welcomed to post on the forum first :smile_cat:
+
+These tests were found to be flaky (intermittently failing on `main` or failed in a PR with unrelated changes). As per [the docs](https://github.com/apache/tvm/blob/main/docs/contribute/ci.rst#handling-flaky-failures, these failures will be disabled in a PR that references this issue until the test owners can fix the source of the flakiness.
+
+### Test(s)
+
+- `tests/python/some_file.py::the_test_name`
+
+### Jenkins Links
+
+- Please provide link(s) to failed CI runs. If runs are for a PR, explain why your PR did not break the test (e.g. did not touch that part of the codebase)
diff --git a/.github/workflows/update_last_successful_branch.yml b/.github/workflows/update_last_successful_branch.yml
new file mode 100644
index 000000000000..1e8def4040ae
--- /dev/null
+++ b/.github/workflows/update_last_successful_branch.yml
@@ -0,0 +1,43 @@
+# 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.
+
+# GH actions.
+# We use it to cover windows and mac builds
+# Jenkins is still the primary CI
+
+name: Update last-successful branch
+
+on:
+  schedule:
+    - cron: "0/15 * * * *"
+  workflow_dispatch:
+
+concurrency:
+  group: update-last-successful-branch
+  cancel-in-progress: true
+
+jobs:
+  update-last-successful-branch:
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v2
+      - name: Update last-successful branch
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          set -eux
+          python tests/scripts/update_branch.py || echo step failed
diff --git a/.gitignore b/.gitignore
index b5ff0f77581b..4e2c9824d307 100644
--- a/.gitignore
+++ b/.gitignore
@@ -108,6 +108,8 @@ DerivedData/
 *.class
 jvm/*/target/
 jvm/*/*/target/
+jvm/native/*/generated
+jvm/native/src/main/native/org_apache_tvm_native_c_api.h
 *.worksheet
 *.idea
 *.iml
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f8d0ab95b80a..9c6a7dddfdf6 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -38,6 +38,7 @@ tvm_option(USE_LLVM "Build with LLVM, can be set to specific llvm-config path" O
 tvm_option(USE_STACKVM_RUNTIME "Include stackvm into the runtime" OFF)
 tvm_option(USE_GRAPH_EXECUTOR "Build with tiny graph executor" ON)
 tvm_option(USE_GRAPH_EXECUTOR_CUDA_GRAPH "Build with tiny graph executor with CUDA Graph for GPUs" OFF)
+tvm_option(USE_AOT_EXECUTOR "Build with AOT executor" ON)
 tvm_option(USE_PROFILER "Build profiler for the VM and graph executor" ON)
 tvm_option(USE_OPENMP "Build with OpenMP thread pool implementation" OFF)
 tvm_option(USE_RELAY_DEBUG "Building Relay in debug mode..." OFF)
@@ -395,6 +396,13 @@ if(USE_PROFILER)
   list(APPEND RUNTIME_SRCS ${RUNTIME_VM_PROFILER_SRCS})
 endif(USE_PROFILER)
 
+if(USE_AOT_EXECUTOR)
+  message(STATUS "Build with AOT Executor support...")
+  file(GLOB RUNTIME_AOT_EXECUTOR_SRCS src/runtime/aot_executor/*.cc)
+  list(APPEND RUNTIME_SRCS ${RUNTIME_AOT_EXECUTOR_SRCS})
+
+endif(USE_AOT_EXECUTOR)
+
 # Enable ctest if gtest is available
 if(USE_GTEST)
   # Check env var for backward compatibility. A better way to specify package
@@ -482,6 +490,7 @@ add_library(tvm_runtime_objs OBJECT ${RUNTIME_SRCS})
 add_library(tvm_libinfo_objs OBJECT ${LIBINFO_FILE})
 
 add_library(tvm SHARED $<TARGET_OBJECTS:tvm_objs> $<TARGET_OBJECTS:tvm_runtime_objs> $<TARGET_OBJECTS:tvm_libinfo_objs>)
+target_include_directories(tvm PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
 set_property(TARGET tvm APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
 set_property(TARGET tvm APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
 if(BUILD_STATIC_RUNTIME)
@@ -496,6 +505,7 @@ else()
   add_library(tvm_runtime SHARED $<TARGET_OBJECTS:tvm_runtime_objs> $<TARGET_OBJECTS:tvm_libinfo_objs>)
   set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_NO_UNDEFINED_SYMBOLS}")
 endif()
+target_include_directories(tvm_runtime PUBLIC "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
 set_property(TARGET tvm_runtime APPEND PROPERTY LINK_OPTIONS "${TVM_VISIBILITY_FLAG}")
 
 target_compile_definitions(tvm_objs PUBLIC DMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>)
@@ -566,17 +576,6 @@ target_link_libraries(tvm_runtime PRIVATE ${TVM_RUNTIME_LINKER_LIBS})
 
 # Set flags for clang
 include(cmake/modules/ClangFlags.cmake)
-
-# Related headers
-target_include_directories(
-  tvm
-  PUBLIC "topi/include")
-target_include_directories(
-  tvm_objs
-  PUBLIC "topi/include")
-target_include_directories(
-  tvm_libinfo_objs
-  PUBLIC "topi/include")
 set(CRC16_INCLUDE_PATH "3rdparty/libcrc/include")
 target_include_directorieS(
   tvm_objs
@@ -621,33 +620,56 @@ endif()
 add_custom_target(runtime DEPENDS tvm_runtime)
 
 # Installation rules
-install(TARGETS tvm DESTINATION lib${LIB_SUFFIX})
-install(TARGETS tvm_runtime DESTINATION lib${LIB_SUFFIX})
+install(TARGETS tvm EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
+install(TARGETS tvm_runtime EXPORT ${PROJECT_NAME}Targets DESTINATION lib${LIB_SUFFIX})
 
 if (INSTALL_DEV)
   install(
-    DIRECTORY "include/." DESTINATION "include"
+    DIRECTORY "include/" DESTINATION "include"
     FILES_MATCHING
     PATTERN "*.h"
   )
   install(
-    DIRECTORY "3rdparty/dlpack/include/." DESTINATION "include"
+    DIRECTORY "3rdparty/dlpack/include/" DESTINATION "include"
     FILES_MATCHING
     PATTERN "*.h"
     )
   install(
-    DIRECTORY "3rdparty/dmlc-core/include/." DESTINATION "include"
+    DIRECTORY "3rdparty/dmlc-core/include/" DESTINATION "include"
     FILES_MATCHING
     PATTERN "*.h"
     )
 else(INSTALL_DEV)
   install(
-    DIRECTORY "include/tvm/runtime/." DESTINATION "include/tvm/runtime"
+    DIRECTORY "include/tvm/runtime/" DESTINATION "include/tvm/runtime"
     FILES_MATCHING
     PATTERN "*.h"
     )
 endif(INSTALL_DEV)
 
+include(GNUInstallDirs)
+include(CMakePackageConfigHelpers)
+set(PROJECT_CONFIG_CONTENT "@PACKAGE_INIT@\n")
+string(APPEND PROJECT_CONFIG_CONTENT "include(CMakeFindDependencyMacro)\n")
+string(APPEND PROJECT_CONFIG_CONTENT "find_dependency(Threads REQUIRED)\n")
+string(APPEND PROJECT_CONFIG_CONTENT
+       "include(\"\${CMAKE_CURRENT_LIST_DIR}/${PROJECT_NAME}Targets.cmake\")")
+file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_CONFIG_CONTENT})
+
+install(EXPORT ${PROJECT_NAME}Targets
+  NAMESPACE ${PROJECT_NAME}::
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME})
+
+# Create config for find_package()
+configure_package_config_file(
+  "${CMAKE_CURRENT_BINARY_DIR}/temp_config_file.cmake" ${PROJECT_NAME}Config.cmake
+  INSTALL_DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
+
+install(
+  FILES
+  "${PROJECT_BINARY_DIR}/${PROJECT_NAME}Config.cmake"
+  DESTINATION "${CMAKE_INSTALL_LIBDIR}/cmake/${PROJECT_NAME}")
+
 # More target definitions
 if(MSVC)
   target_compile_definitions(tvm_objs PRIVATE -DTVM_EXPORTS)
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
index 5e37f09edcf2..03fd4abbb826 100644
--- a/CONTRIBUTORS.md
+++ b/CONTRIBUTORS.md
@@ -64,6 +64,7 @@ We do encourage everyone to work anything they are interested in.
 - [Christopher Sidebottom](https://github.com/Mousius): @Mousius - arm, ethos-u, relay
 - [Junru Shao](https://github.com/junrushao1994) (PMC): @junrushao1994 - relay, compiler
 - [Haichen Shen](https://github.com/icemelon) (PMC): @icemelon - relay, topi
+- [Chris Sullivan](https://github.com/csullivan): @csullivan - amd backend
 - [Siva Rama Krishna Reddy](https://github.com/srkreddy1238): @srkreddy1238 - frontends, golang
 - [Zhixun Tan](https://github.com/phisiart): @phisiart - opengl, web
 - [Andrew Tulloch](https://github.com/ajtulloch): @ajtulloch - topi, compiler, runtime
diff --git a/Jenkinsfile b/Jenkinsfile
index b0e5f77612f5..840f72640900 100755
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -47,11 +47,12 @@ import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
 // NOTE: these lines are scanned by docker/dev_common.sh. Please update the regex as needed. -->
 ci_lint = "tlcpack/ci-lint:v0.68"
 ci_gpu = "tlcpack/ci-gpu:v0.81"
-ci_cpu = "tlcpack/ci-cpu:v0.80"
+ci_cpu = "tlcpack/ci-cpu:v0.81"
 ci_wasm = "tlcpack/ci-wasm:v0.71"
 ci_i386 = "tlcpack/ci-i386:v0.74"
 ci_qemu = "tlcpack/ci-qemu:v0.10"
 ci_arm = "tlcpack/ci-arm:v0.07"
+ci_hexagon = "tlcpack/ci-hexagon:v0.01"
 // <--- End of regex-scanned config.
 
 // Parameters to allow overriding (in Jenkins UI), the images
@@ -256,20 +257,20 @@ def cpp_unittest(image) {
 }
 
 stage('Build') {
-    parallel 'BUILD: GPU': {
-      if (!skip_ci) {
-        node('GPUBUILD') {
-          ws(per_exec_ws('tvm/build-gpu')) {
-            init_git()
-            sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh"
-            make(ci_gpu, 'build', '-j2')
-            pack_lib('gpu', tvm_multilib)
-            // compiler test
-            sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu_other.sh"
-            make(ci_gpu, 'build2', '-j2')
-          }
+  parallel 'BUILD: GPU': {
+    if (!skip_ci) {
+      node('GPUBUILD') {
+        ws(per_exec_ws('tvm/build-gpu')) {
+          init_git()
+          sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu.sh"
+          make(ci_gpu, 'build', '-j2')
+          pack_lib('gpu', tvm_multilib)
+          // compiler test
+          sh "${docker_run} ${ci_gpu} ./tests/scripts/task_config_build_gpu_other.sh"
+          make(ci_gpu, 'build2', '-j2')
         }
       }
+    }
   },
   'BUILD: CPU': {
     if (!skip_ci && is_docs_only_build != 1) {
@@ -360,32 +361,64 @@ stage('Build') {
             script: "${docker_run} ${ci_qemu} ./tests/scripts/task_config_build_qemu.sh",
             label: 'Create QEMU cmake config',
           )
-          make(ci_qemu, 'build', '-j2')
-          timeout(time: max_time, unit: 'MINUTES') {
-            ci_setup(ci_qemu)
+          try {
+            make(ci_qemu, 'build', '-j2')
+            timeout(time: max_time, unit: 'MINUTES') {
+              ci_setup(ci_qemu)
+              sh (
+                script: "${docker_run} ${ci_qemu} ./tests/scripts/task_python_microtvm.sh",
+                label: 'Run microTVM tests',
+              )
+              sh (
+                script: "${docker_run} ${ci_qemu} ./tests/scripts/task_demo_microtvm.sh",
+                label: 'Run microTVM demos',
+              )
+            }
+          } finally {
+            junit 'build/pytest-results/*.xml'
+          }
+        }
+      }
+     } else {
+      Utils.markStageSkippedForConditional('BUILD: QEMU')
+    }
+  },
+  'BUILD: Hexagon': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('CPU') {
+        ws(per_exec_ws('tvm/build-hexagon')) {
+          init_git()
+          sh (
+            script: "${docker_run} ${ci_hexagon} ./tests/scripts/task_config_build_hexagon.sh",
+            label: 'Create Hexagon cmake config',
+          )
+          try {
+            make(ci_hexagon, 'build', '-j2')
             sh (
-              script: "${docker_run} ${ci_qemu} ./tests/scripts/task_python_microtvm.sh",
-              label: 'Run microTVM tests',
+              script: "${docker_run} ${ci_hexagon} ./tests/scripts/task_build_hexagon_api.sh",
+              label: 'Build Hexagon API',
             )
             sh (
-              script: "${docker_run} ${ci_qemu} ./tests/scripts/task_demo_microtvm.sh",
-              label: 'Run microTVM demos',
+              script: "${docker_run} ${ci_hexagon} ./tests/scripts/task_python_hexagon.sh",
+              label: 'Run Hexagon tests',
             )
+          } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
       }
      } else {
-      Utils.markStageSkippedForConditional('BUILD: QEMU')
+      Utils.markStageSkippedForConditional('BUILD: Hexagon')
     }
   }
 }
 
 stage('Test') {
-    parallel 'unittest: GPU': {
-      if (!skip_ci && is_docs_only_build != 1) {
-        node('TensorCore') {
-          ws(per_exec_ws('tvm/ut-python-gpu')) {
+  parallel 'unittest: GPU': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('TensorCore') {
+        ws(per_exec_ws('tvm/ut-python-gpu')) {
+          try {
             init_git()
             unpack_lib('gpu', tvm_multilib)
             timeout(time: max_time, unit: 'MINUTES') {
@@ -402,18 +435,21 @@ stage('Test') {
                 script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_integration_gpuonly.sh",
                 label: 'Run Python GPU integration tests',
               )
-              junit 'build/pytest-results/*.xml'
             }
+          } finally {
+            junit 'build/pytest-results/*.xml'
           }
         }
-      } else {
-        Utils.markStageSkippedForConditional('unittest: GPU')
       }
-    },
-    'integration: CPU': {
-      if (!skip_ci && is_docs_only_build != 1) {
-        node('CPU') {
-          ws(per_exec_ws('tvm/ut-python-cpu')) {
+    } else {
+      Utils.markStageSkippedForConditional('unittest: GPU')
+    }
+  },
+  'integration: CPU': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('CPU') {
+        ws(per_exec_ws('tvm/ut-python-cpu')) {
+          try {
             init_git()
             unpack_lib('cpu', tvm_multilib_tsim)
             timeout(time: max_time, unit: 'MINUTES') {
@@ -422,18 +458,21 @@ stage('Test') {
                 script: "${docker_run} ${ci_cpu} ./tests/scripts/task_python_integration.sh",
                 label: 'Run CPU integration tests',
               )
-              junit 'build/pytest-results/*.xml'
             }
+          } finally {
+            junit 'build/pytest-results/*.xml'
           }
         }
-      } else {
-        Utils.markStageSkippedForConditional('integration: CPU')
       }
-    },
-    'unittest: CPU': {
-      if (!skip_ci && is_docs_only_build != 1) {
-        node('CPU') {
-          ws(per_exec_ws("tvm/ut-python-cpu")) {
+    } else {
+      Utils.markStageSkippedForConditional('integration: CPU')
+    }
+  },
+  'unittest: CPU': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('CPU') {
+        ws(per_exec_ws("tvm/ut-python-cpu")) {
+          try {
             init_git()
             unpack_lib('cpu', tvm_multilib_tsim)
             timeout(time: max_time, unit: 'MINUTES') {
@@ -444,18 +483,21 @@ stage('Test') {
                 script: "${docker_run} ${ci_cpu} ./tests/scripts/task_python_vta_tsim.sh",
                 label: "Run VTA tests in TSIM",
               )
-              junit "build/pytest-results/*.xml"
             }
+          } finally {
+            junit 'build/pytest-results/*.xml'
           }
         }
-      } else {
-        Utils.markStageSkippedForConditional('unittest: CPU')
       }
-    },
-    'python3: i386': {
-      if (!skip_ci && is_docs_only_build != 1) {
-        node('CPU') {
-          ws(per_exec_ws('tvm/ut-python-i386')) {
+    } else {
+      Utils.markStageSkippedForConditional('unittest: CPU')
+    }
+  },
+  'python3: i386': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('CPU') {
+        ws(per_exec_ws('tvm/ut-python-i386')) {
+          try {
             init_git()
             unpack_lib('i386', tvm_multilib)
             timeout(time: max_time, unit: 'MINUTES') {
@@ -466,18 +508,21 @@ stage('Test') {
                 label: 'Run i386 integration tests',
               )
               fsim_test(ci_i386)
-              junit 'build/pytest-results/*.xml'
             }
+          } finally {
+            junit 'build/pytest-results/*.xml'
           }
         }
-     } else {
-        Utils.markStageSkippedForConditional('python3: i386')
       }
-    },
-    'python3: arm': {
-      if (!skip_ci && is_docs_only_build != 1) {
-        node('ARM') {
-          ws(per_exec_ws('tvm/ut-python-arm')) {
+    } else {
+      Utils.markStageSkippedForConditional('python3: i386')
+    }
+  },
+  'python3: arm': {
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('ARM') {
+        ws(per_exec_ws('tvm/ut-python-arm')) {
+          try {
             init_git()
             unpack_lib('arm', tvm_multilib)
             timeout(time: max_time, unit: 'MINUTES') {
@@ -487,47 +532,55 @@ stage('Test') {
                 script: "${docker_run} ${ci_arm} ./tests/scripts/task_python_arm_compute_library.sh",
                 label: 'Run test_arm_compute_lib test',
               )
-              junit 'build/pytest-results/*.xml'
             // sh "${docker_run} ${ci_arm} ./tests/scripts/task_python_integration.sh"
             }
+          } finally {
+            junit 'build/pytest-results/*.xml'
           }
         }
-      } else {
-         Utils.markStageSkippedForConditional('python3: arm')
       }
-    },
+    } else {
+      Utils.markStageSkippedForConditional('python3: arm')
+    }
+  },
   'topi: GPU': {
-  if (!skip_ci && is_docs_only_build != 1) {
-    node('GPU') {
-      ws(per_exec_ws('tvm/topi-python-gpu')) {
-        init_git()
-        unpack_lib('gpu', tvm_multilib)
-        timeout(time: max_time, unit: 'MINUTES') {
-          ci_setup(ci_gpu)
-          sh (
-            script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_topi.sh",
-            label: 'Run TOPI tests',
-          )
-          junit 'build/pytest-results/*.xml'
+    if (!skip_ci && is_docs_only_build != 1) {
+      node('GPU') {
+        ws(per_exec_ws('tvm/topi-python-gpu')) {
+          try {
+            init_git()
+            unpack_lib('gpu', tvm_multilib)
+            timeout(time: max_time, unit: 'MINUTES') {
+              ci_setup(ci_gpu)
+              sh (
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_topi.sh",
+                label: 'Run TOPI tests',
+              )
+            }
+          } finally {
+            junit 'build/pytest-results/*.xml'
+          }
         }
       }
-    }
     } else {
       Utils.markStageSkippedForConditional('topi: GPU')
-  }
+    }
   },
   'frontend: GPU': {
     if (!skip_ci && is_docs_only_build != 1) {
       node('GPU') {
         ws(per_exec_ws('tvm/frontend-python-gpu')) {
-          init_git()
-          unpack_lib('gpu', tvm_multilib)
-          timeout(time: max_time, unit: 'MINUTES') {
-            ci_setup(ci_gpu)
-            sh (
-              script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh",
-              label: 'Run Python frontend tests',
-            )
+          try {
+            init_git()
+            unpack_lib('gpu', tvm_multilib)
+            timeout(time: max_time, unit: 'MINUTES') {
+              ci_setup(ci_gpu)
+              sh (
+                script: "${docker_run} ${ci_gpu} ./tests/scripts/task_python_frontend.sh",
+                label: 'Run Python frontend tests',
+              )
+            }
+          } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
@@ -540,14 +593,17 @@ stage('Test') {
     if (!skip_ci && is_docs_only_build != 1) {
       node('CPU') {
         ws(per_exec_ws('tvm/frontend-python-cpu')) {
-          init_git()
-          unpack_lib('cpu', tvm_multilib)
-          timeout(time: max_time, unit: 'MINUTES') {
-            ci_setup(ci_cpu)
-            sh (
-              script: "${docker_run} ${ci_cpu} ./tests/scripts/task_python_frontend_cpu.sh",
-              label: 'Run Python frontend tests',
-            )
+          try {
+            init_git()
+            unpack_lib('cpu', tvm_multilib)
+            timeout(time: max_time, unit: 'MINUTES') {
+              ci_setup(ci_cpu)
+              sh (
+                script: "${docker_run} ${ci_cpu} ./tests/scripts/task_python_frontend_cpu.sh",
+                label: 'Run Python frontend tests',
+              )
+            }
+          } finally {
             junit 'build/pytest-results/*.xml'
           }
         }
@@ -595,13 +651,13 @@ stage('Build packages') {
 */
 
 stage('Deploy') {
-    node('doc') {
-      ws(per_exec_ws('tvm/deploy-docs')) {
-        if (env.BRANCH_NAME == 'main') {
+  node('doc') {
+    ws(per_exec_ws('tvm/deploy-docs')) {
+      if (env.BRANCH_NAME == 'main') {
         unpack_lib('docs', 'docs.tgz')
         sh 'cp docs.tgz /var/docs/docs.tgz'
         sh 'tar xf docs.tgz -C /var/docs'
-        }
       }
     }
+  }
 }
diff --git a/apps/android_rpc/.gitignore b/apps/android_rpc/.gitignore
index 39fb081a42a8..1f8bc3b3e5e6 100644
--- a/apps/android_rpc/.gitignore
+++ b/apps/android_rpc/.gitignore
@@ -7,3 +7,9 @@
 /build
 /captures
 .externalNativeBuild
+
+/app/src/main/jni/jni_helper_func.h
+/app/src/main/jni/org_apache_tvm_native_c_api.cc
+/app/src/main/jni/org_apache_tvm_native_c_api.h
+/app/src/main/obj/
+dev_tools/tvmrpc.keystore
diff --git a/apps/cpp_rpc/CMakeLists.txt b/apps/cpp_rpc/CMakeLists.txt
index 966448929a0a..657d486500a9 100644
--- a/apps/cpp_rpc/CMakeLists.txt
+++ b/apps/cpp_rpc/CMakeLists.txt
@@ -51,6 +51,10 @@ if (BUILD_FOR_ANDROID AND USE_HEXAGON_SDK)
   list(APPEND TVM_RPC_LINKER_LIBS cdsprpc log)
 endif()
 
-list(APPEND TVM_RPC_LINKER_LIBS tvm_runtime)
+if(BUILD_STATIC_RUNTIME)
+  list(APPEND TVM_RPC_LINKER_LIBS -Wl,--whole-archive tvm_runtime -Wl,--no-whole-archive)
+else()
+  list(APPEND TVM_RPC_LINKER_LIBS tvm_runtime)
+endif()
  
 target_link_libraries(tvm_rpc ${TVM_RPC_LINKER_LIBS})
diff --git a/apps/hexagon_api/CMakeLists.txt b/apps/hexagon_api/CMakeLists.txt
index 557dcfb85045..052fd00020cb 100644
--- a/apps/hexagon_api/CMakeLists.txt
+++ b/apps/hexagon_api/CMakeLists.txt
@@ -23,6 +23,31 @@ else()
 endif()
 file(MAKE_DIRECTORY ${HEXAGON_API_BINARY_DIR})
 
+# Build X86 binaries:
+# - tvm_rpc_x86
+
+ExternalProject_Add(x86_tvm_runtime_rpc
+  SOURCE_DIR "${TVM_SOURCE_DIR}"
+  BUILD_COMMAND $(MAKE) runtime tvm_rpc
+  CMAKE_ARGS
+    "-DUSE_HEXAGON_TOOLCHAIN=${USE_HEXAGON_TOOLCHAIN}"
+    "-DCMAKE_CXX_STANDARD=14"
+    "-DUSE_LIBBACKTRACE=OFF"
+    "-DUSE_RPC=ON"
+    "-DUSE_CPP_RPC=ON"
+    "-DUSE_HEXAGON_RPC=ON"
+    "-DBUILD_STATIC_RUNTIME=ON"
+  INSTALL_COMMAND ""
+  BUILD_ALWAYS ON
+)
+ExternalProject_Get_Property(x86_tvm_runtime_rpc BINARY_DIR)
+ExternalProject_Add_Step(x86_tvm_runtime_rpc copy_rpc_server
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different
+    ${BINARY_DIR}/tvm_rpc
+    ${HEXAGON_API_BINARY_DIR}/tvm_rpc_x86
+  DEPENDEES install
+)
+
 # Build Android binaries:
 # - libtvm_runtime.so
 # - tvm_rpc_android
@@ -38,7 +63,6 @@ ExternalProject_Add(android_tvm_runtime_rpc
     "-DUSE_HEXAGON_ARCH=${USE_HEXAGON_ARCH}"
     "-DCMAKE_CXX_STANDARD=14"
     "-DUSE_LIBBACKTRACE=OFF"
-    "-DUSE_LLVM=OFF"
     "-DUSE_RPC=ON"
     "-DUSE_CPP_RPC=ON"
     "-DUSE_HEXAGON_RPC=ON"
@@ -66,7 +90,7 @@ ExternalProject_Add_Step(android_tvm_runtime_rpc copy_rpc_server
 
 ExternalProject_Add(hexagon_tvm_runtime_rpc
   SOURCE_DIR "${TVM_SOURCE_DIR}"
-  BUILD_COMMAND $(MAKE) runtime
+  BUILD_COMMAND $(MAKE) runtime hexagon_rpc_sim
   CMAKE_ARGS
     "-DCMAKE_C_COMPILER=${USE_HEXAGON_TOOLCHAIN}/bin/hexagon-clang"
     "-DCMAKE_CXX_COMPILER=${USE_HEXAGON_TOOLCHAIN}/bin/hexagon-clang++"
@@ -84,6 +108,7 @@ ExternalProject_Add_Step(hexagon_tvm_runtime_rpc copy_binaries
   COMMAND ${CMAKE_COMMAND} -E copy_if_different
     ${BINARY_DIR}/libtvm_runtime.a
     ${BINARY_DIR}/libhexagon_rpc_skel.so
+    ${BINARY_DIR}/libhexagon_rpc_sim.so
     ${HEXAGON_API_BINARY_DIR}
   DEPENDEES install
 )
diff --git a/apps/microtvm/arduino/template_project/src/example_project/model.c b/apps/microtvm/arduino/template_project/src/example_project/model.c
index 553665191b14..25d609dacce1 100644
--- a/apps/microtvm/arduino/template_project/src/example_project/model.c
+++ b/apps/microtvm/arduino/template_project/src/example_project/model.c
@@ -86,7 +86,7 @@ tvm_crt_error_t TVMPlatformGenerateRandom(uint8_t* buffer, size_t num_bytes) {
 void TVMInitialize() { StackMemoryManager_Init(&app_workspace, g_aot_memory, WORKSPACE_SIZE); }
 
 void TVMExecute(void* input_data, void* output_data) {
-  int ret_val = tvmgen_default_run_model(input_data, output_data);
+  int ret_val = tvmgen_default___tvm_main__(input_data, output_data);
   if (ret_val != 0) {
     TVMPlatformAbort(kTvmErrorPlatformCheckFailure);
   }
diff --git a/apps/microtvm/ethosu/src/demo_bare_metal.c b/apps/microtvm/ethosu/src/demo_bare_metal.c
index 753dd0712852..78d79dc23e4d 100644
--- a/apps/microtvm/ethosu/src/demo_bare_metal.c
+++ b/apps/microtvm/ethosu/src/demo_bare_metal.c
@@ -40,7 +40,7 @@ int main(int argc, char** argv) {
 
   printf("Running inference\n");
   struct tvmgen_default_outputs outputs = {
-      .output = output,
+      .MobilenetV2_Predictions_Reshape_11 = output,
   };
   struct tvmgen_default_inputs inputs = {
       .tfl_quantize = input,
diff --git a/apps/microtvm/ethosu/src/demo_freertos.c b/apps/microtvm/ethosu/src/demo_freertos.c
index cd5b76b15f49..2be8f1148e28 100644
--- a/apps/microtvm/ethosu/src/demo_freertos.c
+++ b/apps/microtvm/ethosu/src/demo_freertos.c
@@ -106,7 +106,7 @@ static void prvInferenceTask(void* pvParameters) {
       .tfl_quantize = pucReceivedData,
   };
   struct tvmgen_default_outputs xOutputs = {
-      .output = output,
+      .MobilenetV2_Predictions_Reshape_11 = output,
   };
   struct ethosu_driver* xDriver = ethosu_reserve_driver();
   struct tvmgen_default_devices xDevices = {
diff --git a/apps/microtvm/zephyr/template_project/microtvm_api_server.py b/apps/microtvm/zephyr/template_project/microtvm_api_server.py
index a7774899acb8..de95be949ad9 100644
--- a/apps/microtvm/zephyr/template_project/microtvm_api_server.py
+++ b/apps/microtvm/zephyr/template_project/microtvm_api_server.py
@@ -20,6 +20,7 @@
 import collections.abc
 import enum
 import fcntl
+import json
 import logging
 import os
 import os.path
@@ -35,7 +36,7 @@
 import tempfile
 import threading
 import time
-import json
+import usb
 
 import serial
 import serial.tools.list_ports
@@ -179,8 +180,6 @@ def openocd_serial(options):
     if "openocd_serial" in options:
         return options["openocd_serial"]
 
-    import usb  # pylint: disable=import-outside-toplevel
-
     find_kw = BOARD_USB_FIND_KW[CMAKE_CACHE["BOARD"]]
     boards = usb.core.find(find_all=True, **find_kw)
     serials = []
@@ -588,6 +587,13 @@ def _set_nonblock(fd):
 
 
 class ZephyrSerialTransport:
+
+    NRF5340_VENDOR_ID = 0x1366
+
+    # NRF5340_DK v1.0.0 uses VCOM2
+    # NRF5340_DK v2.0.0 uses VCOM1
+    NRF5340_DK_BOARD_VCOM_BY_PRODUCT_ID = {0x1055: "VCOM2", 0x1051: "VCOM1"}
+
     @classmethod
     def _lookup_baud_rate(cls, options):
         # TODO(mehrdadh): remove this hack once dtlib.py is a standalone project
@@ -625,7 +631,17 @@ def _find_nrf_serial_port(cls, options):
             parts = line.split()
             ports_by_vcom[parts[2]] = parts[1]
 
-        return ports_by_vcom["VCOM2"]
+        nrf_board = usb.core.find(idVendor=cls.NRF5340_VENDOR_ID)
+
+        if nrf_board == None:
+            raise Exception("_find_nrf_serial_port: unable to find NRF5340DK")
+
+        if nrf_board.idProduct in cls.NRF5340_DK_BOARD_VCOM_BY_PRODUCT_ID:
+            vcom_port = cls.NRF5340_DK_BOARD_VCOM_BY_PRODUCT_ID[nrf_board.idProduct]
+        else:
+            raise Exception("_find_nrf_serial_port: unable to find known NRF5340DK product ID")
+
+        return ports_by_vcom[vcom_port]
 
     @classmethod
     def _find_openocd_serial_port(cls, options):
diff --git a/apps/microtvm/zephyr/template_project/src/aot_demo/main.c b/apps/microtvm/zephyr/template_project/src/aot_demo/main.c
index 3946727b26a8..27d74d7ba76b 100644
--- a/apps/microtvm/zephyr/template_project/src/aot_demo/main.c
+++ b/apps/microtvm/zephyr/template_project/src/aot_demo/main.c
@@ -178,7 +178,7 @@ void TVMInfer() {
       .input_1 = input_data,
   };
   struct tvmgen_default_outputs outputs = {
-      .output = output_data,
+      .Identity = output_data,
   };
 
   StackMemoryManager_Init(&app_workspace, g_aot_memory, WORKSPACE_SIZE);
diff --git a/apps/microtvm/zephyr_cmsisnn/src/main.c b/apps/microtvm/zephyr_cmsisnn/src/main.c
index 7696915f7aa2..31f6cd0cc1d0 100644
--- a/apps/microtvm/zephyr_cmsisnn/src/main.c
+++ b/apps/microtvm/zephyr_cmsisnn/src/main.c
@@ -34,7 +34,7 @@ extern float output_storage[12];
 
 extern const size_t output_len;
 
-static uint8_t g_crt_workspace[TVMGEN_DEFAULT_WORKSPACE_SIZE + 256];
+static uint8_t g_crt_workspace[TVMGEN_DEFAULT_WORKSPACE_SIZE];
 tvm_workspace_t app_workspace;
 
 void TVMLogf(const char* msg, ...) {
@@ -62,7 +62,7 @@ void main(void) {
   StackMemoryManager_Init(&app_workspace, g_crt_workspace, TVMGEN_DEFAULT_WORKSPACE_SIZE);
 
   struct tvmgen_default_inputs inputs = {.input = input_storage};
-  struct tvmgen_default_outputs outputs = {.output = output_storage};
+  struct tvmgen_default_outputs outputs = {.Identity = output_storage};
 
   if (tvmgen_default_run(&inputs, &outputs) != 0) {
     printk("Model run failed\n");
diff --git a/cmake/modules/Arduino.cmake b/cmake/modules/Arduino.cmake
index 54c144081efa..f603e5d5df2b 100644
--- a/cmake/modules/Arduino.cmake
+++ b/cmake/modules/Arduino.cmake
@@ -43,7 +43,7 @@ if(USE_MICRO)
       math(EXPR job_spec_stop "${job_spec_length} - 3")
 
       list(GET job_spec 0 job_src_base)
-      set(job_src_base "${CMAKE_SOURCE_DIR}/${job_src_base}")
+      set(job_src_base "${CMAKE_CURRENT_SOURCE_DIR}/${job_src_base}")
       foreach(copy_pattern_index RANGE 1 "${job_spec_stop}" 3)
         list(GET job_spec ${copy_pattern_index} copy_pattern)
         math(EXPR copy_dest_index "${copy_pattern_index} + 2")
diff --git a/cmake/modules/Hexagon.cmake b/cmake/modules/Hexagon.cmake
index b90528d0767d..a8844a22e164 100644
--- a/cmake/modules/Hexagon.cmake
+++ b/cmake/modules/Hexagon.cmake
@@ -47,92 +47,200 @@ function(find_hexagon_toolchain)
   endif()
 endfunction()
 
-if(BUILD_FOR_HEXAGON)
-  find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
-  # Add SDK and QuRT includes when building for Hexagon.
-  include_directories(SYSTEM ${HEXAGON_SDK_INCLUDES} ${HEXAGON_QURT_INCLUDES})
+macro(file_glob_append _output_list)
+  tvm_file_glob(GLOB _tmp0 ${ARGN})
+  set(_tmp1 ${${_output_list}})
+  list(APPEND _tmp1 ${_tmp0})
+  set(${_output_list} ${_tmp1})
+endmacro()
+
+set(TVMRT_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime")
+
+# First, verify that USE_HEXAGON_DEVICE has a valid value.
+if(DEFINED USE_HEXAGON_DEVICE)
+  if(NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_SIM}" AND
+     NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_HW}" AND
+     NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_NONE}")
+    message(SEND_ERROR "USE_HEXAGON_DEVICE must be one of "
+            "[${PICK_NONE}|${PICK_SIM}|${PICK_HW}]")
+    set(USE_HEXAGON_DEVICE OFF)
+  endif()
 endif()
 
-if (NOT USE_HEXAGON_SDK STREQUAL "" AND
-    NOT USE_HEXAGON_SDK STREQUAL "/path/to/sdk")
-  set(HEXAGON_SDK_PATH_DEFINED ${USE_HEXAGON_SDK})
+# This .cmake file is included when building any part of TVM for any
+# architecture. It shouldn't require any Hexagon-specific parameters
+# (like the path to the SDK), unless it's needed.
+#
+# Aside from building the code for Hexagon, two flags can enable some
+# Hexagon-related functionality:
+# - USE_HEXAGON_DEVICE
+# - USE_HEXAGON_RPC
+#
+# USE_HEXAGON_RPC:
+# - When building for Hexagon, this will build the Hexagon endpoint of the
+#   RPC server: the FastRPC skel library (with TVM runtime built into it),
+#   and the standalone RPC server for simulator.
+# - When building for Android, this will build the (intermediary) RPC server,
+#   including the "stub" code for the FastRPC implementation of the RPC
+#   channel.
+# - When building for x86, this will build the host-side code that instan-
+#   tiates the simulator.
+
+if(NOT BUILD_FOR_HEXAGON AND NOT BUILD_FOR_ANDROID)
+  set(BUILD_FOR_HOST TRUE)
 endif()
 
-if (BUILD_FOR_ANDROID AND HEXAGON_SDK_PATH_DEFINED)
-  find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
+
+if(NOT USE_HEXAGON_DEVICE AND NOT USE_HEXAGON_RPC AND NOT BUILD_FOR_HEXAGON)
+  # If nothing related to Hexagon is enabled, add phony Hexagon codegen,
+  # and some stuff needed by cpptests (this part is a temporary workaround
+  # until e2e support for Hexagon is enabled).
+  if(BUILD_FOR_HOST)
+    list(APPEND COMPILER_SRCS src/target/opt/build_hexagon_off.cc)
+  endif()
+  list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_buffer.cc)
+  list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_common.cc)
+  list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_user_dma.cc)
+  return()
+endif()
+
+
+function(add_android_paths)
+  if(NOT DEFINED HEXAGON_SDK_INCLUDES OR
+     NOT DEFINED HEXAGON_RPCMEM_ROOT OR
+     NOT DEFINED HEXAGON_REMOTE_ROOT)
+    message(FATAL_ERROR "This function must be called after find_hexagon_sdk_root")
+  endif()
   include_directories(SYSTEM
     ${HEXAGON_SDK_INCLUDES}
     ${HEXAGON_RPCMEM_ROOT}/inc
-    ${HEXAGON_REMOTE_ROOT})
+    ${HEXAGON_REMOTE_ROOT}
+  )
   link_directories(${HEXAGON_REMOTE_ROOT})
-  list(APPEND TVM_RUNTIME_LINKER_LIBS cdsprpc)
-endif()
+endfunction()
 
-# Don't run these checks when compiling Hexagon device code,
-# e.g. when compiling the TVM runtime for Hexagon.
-if (NOT BUILD_FOR_HEXAGON AND NOT BUILD_FOR_ANDROID)
-  if(USE_HEXAGON_DEVICE STREQUAL "OFF")
-    list(APPEND COMPILER_SRCS src/target/opt/build_hexagon_off.cc)
-    # append select runtime sources for unit testing
-    list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_buffer.cc)
-    list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_common.cc)
-    list(APPEND RUNTIME_SRCS src/runtime/hexagon/hexagon/hexagon_user_dma.cc)
-    return()
-  elseif(NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_SIM}" AND
-         NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_HW}")
-    message(SEND_ERROR "USE_HEXAGON_DEVICE must be one of "
-            "[${PICK_NONE}|${PICK_SIM}|${PICK_HW}]")
-    return()
+function(add_hexagon_wrapper_paths)
+  if(NOT DEFINED HEXAGON_TOOLCHAIN)
+    message(FATAL_ERROR "This function must be called after find_hexagon_toolchain")
   endif()
-endif()
+  include_directories(SYSTEM
+    "${HEXAGON_TOOLCHAIN}/include/iss"
+  )
+  link_directories("${HEXAGON_TOOLCHAIN}/lib/iss")
+endfunction()
 
-# If no Hexagon support is enabled (other than some stub code), cmake
-# execution should stop before reaching this point.
+# Common sources for TVM runtime with Hexagon support
+file_glob_append(RUNTIME_HEXAGON_COMMON_SRCS
+  "${TVMRT_SOURCE_DIR}/hexagon/hexagon_module.cc"
+  "${TVMRT_SOURCE_DIR}/hexagon/hexagon/*.cc"
+)
 
-if(NOT USE_HEXAGON_SDK OR NOT USE_HEXAGON_ARCH)
-  message(SEND_ERROR "Please set USE_HEXAGON_SDK to the Hexagon SDK root, "
-          "and USE_HEXAGON_ARCH to the Hexagon architecture version")
-  return()
-endif()
 
-if(USE_HEXAGON_LAUNCHER STREQUAL "ON")
-  message(SEND_ERROR "USE_HEXAGON_LAUNCHER is deprecated, please build apps separately")
+if(BUILD_FOR_HEXAGON)
+  find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
+  # Add SDK and QuRT includes when building for Hexagon.
+  include_directories(SYSTEM ${HEXAGON_SDK_INCLUDES} ${HEXAGON_QURT_INCLUDES})
+
+  list(APPEND RUNTIME_HEXAGON_SRCS ${RUNTIME_HEXAGON_COMMON_SRCS})
 endif()
 
-# find_hexagon_sdk_root has been called at this point.
+
+if(USE_HEXAGON_DEVICE)
+  function(invalid_device_value_for BUILD_TARGET)
+    message(SEND_ERROR
+      "USE_HEXAGON_DEVICE=${USE_HEXAGON_DEVICE} is not supported when "
+      "building for ${BUILD_TARGET}"
+    )
+  endfunction()
+
+  list(APPEND RUNTIME_HEXAGON_SRCS ${RUNTIME_HEXAGON_COMMON_SRCS})
+
+  if(BUILD_FOR_HOST)
+    if(NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_SIM}")
+      invalid_device_value_for("host")
+    endif()
+    find_hexagon_toolchain()
+    add_hexagon_wrapper_paths()
+    file_glob_append(RUNTIME_HEXAGON_SRCS
+      "${TVMRT_SOURCE_DIR}/hexagon/android/*.cc"
+      "${TVMRT_SOURCE_DIR}/hexagon/android/sim/*.cc"
+    )
+    list(APPEND TVM_RUNTIME_LINKER_LIBS "-lwrapper")
+
+    ExternalProject_Add(sim_dev
+      SOURCE_DIR "${TVMRT_SOURCE_DIR}/hexagon/android/sim/driver"
+      CMAKE_ARGS
+        "-DCMAKE_C_COMPILER=${HEXAGON_TOOLCHAIN}/bin/hexagon-clang"
+        "-DCMAKE_CXX_COMPILER=${HEXAGON_TOOLCHAIN}/bin/hexagon-clang++"
+        "-DHEXAGON_ARCH=${USE_HEXAGON_ARCH}"
+      INSTALL_COMMAND "true"
+    )
+
+  elseif(BUILD_FOR_ANDROID)
+    if(NOT USE_HEXAGON_DEVICE STREQUAL "${PICK_HW}")
+      invalid_device_value_for("Android")
+    endif()
+    find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
+    find_hexagon_toolchain()
+    add_android_paths()
+    file_glob_append(RUNTIME_HEXAGON_SRCS
+      "${TVMRT_SOURCE_DIR}/hexagon/android/*.cc"
+      "${TVMRT_SOURCE_DIR}/hexagon/android/target/*.cc"
+    )
+    # Hexagon runtime uses __android_log_print, which is in liblog.
+    list(APPEND TVM_RUNTIME_LINKER_LIBS dl log cdsprpc)
+
+  elseif(BUILD_FOR_HEXAGON)
+    invalid_device_value_for("Hexagon")
+  endif()
+endif()   # USE_HEXAGON_DEVICE
+
 
 if(USE_HEXAGON_RPC)
-  set(TVMRT_SOURCE_DIR "${CMAKE_SOURCE_DIR}/src/runtime")
-  set(QAIC_EXE "${HEXAGON_QAIC_EXE}")
-  foreach(INCDIR IN LISTS HEXAGON_SDK_INCLUDES HEXAGON_REMOTE_ROOT)
-    list(APPEND QAIC_FLAGS "-I${INCDIR}")
-  endforeach()
-
-  add_custom_command(
-    OUTPUT
-      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.h"
-      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
-      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
-    COMMAND
-      ${QAIC_EXE} ${QAIC_FLAGS} "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
-        -o "${TVMRT_SOURCE_DIR}/hexagon/rpc"
-    MAIN_DEPENDENCY "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
-  )
+  function(build_rpc_idl)
+    set(QAIC_EXE "${HEXAGON_QAIC_EXE}")
+    foreach(INCDIR IN LISTS HEXAGON_SDK_INCLUDES HEXAGON_REMOTE_ROOT)
+      list(APPEND QAIC_FLAGS "-I${INCDIR}")
+    endforeach()
+
+    add_custom_command(
+      OUTPUT
+        "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.h"
+        "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
+        "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
+      COMMAND
+        ${QAIC_EXE} ${QAIC_FLAGS} "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
+          -o "${TVMRT_SOURCE_DIR}/hexagon/rpc"
+      MAIN_DEPENDENCY "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc.idl"
+    )
+  endfunction()
+
+  list(APPEND RUNTIME_HEXAGON_SRCS ${RUNTIME_HEXAGON_COMMON_SRCS})
 
   if(BUILD_FOR_ANDROID)
     # Android part
-    tvm_file_glob(GLOB RUNTIME_HEXAGON_SRCS src/runtime/hexagon/host/*.cc)
-    tvm_file_glob(GLOB RUNTIME_HEXAGON_SRCS "${TVMRT_SOURCE_DIR}/hexagon/rpc/android/*.cc")
-    list(APPEND RUNTIME_HEXAGON_SRCS "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c")
+    find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
+    add_android_paths()
+    build_rpc_idl()
+    file_glob_append(RUNTIME_HEXAGON_SRCS
+      "${TVMRT_SOURCE_DIR}/hexagon/host/*.cc"
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/android/*.cc"
+    )
+    # Add this file separately, because it's auto-generated, and glob won't
+    # find it during cmake-time.
+    list(APPEND RUNTIME_HEXAGON_SRCS
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_stub.c"
+    )
+    list(APPEND TVM_RUNTIME_LINKER_LIBS cdsprpc)
 
   elseif(BUILD_FOR_HEXAGON)
     # Hexagon part
+    find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
     find_hexagon_toolchain()
-    message(STATUS "HEXAGON_TOOLCHAIN: ${HEXAGON_TOOLCHAIN}")
+    build_rpc_idl()
 
-    add_library(hexagon_rpc_skel SHARED
-      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
-      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon/rpc_server.cc"
+    # Include the generic RPC code into the TVM runtime.
+    list(APPEND RUNTIME_HEXAGON_SRCS
       "${TVMRT_SOURCE_DIR}/minrpc/minrpc_server.h"
       "${TVMRT_SOURCE_DIR}/minrpc/rpc_reference.h"
       "${TVMRT_SOURCE_DIR}/rpc/rpc_module.cc"
@@ -140,60 +248,33 @@ if(USE_HEXAGON_RPC)
       "${TVMRT_SOURCE_DIR}/rpc/rpc_session.cc"
       "${TVMRT_SOURCE_DIR}/rpc/rpc_local_session.cc"
     )
+    # Add the hardware-specific RPC code into the skel library.
+    add_library(hexagon_rpc_skel SHARED
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon/rpc_server.cc"
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/hexagon_rpc_skel.c"
+    )
     target_include_directories(hexagon_rpc_skel
       SYSTEM PRIVATE "${TVMRT_SOURCE_DIR}/hexagon/rpc"
     )
-  endif()
-endif()
-
-if(USE_HEXAGON_DEVICE STREQUAL "${PICK_SIM}")
-  find_hexagon_toolchain()
-  message(STATUS "Hexagon toolchain: ${HEXAGON_TOOLCHAIN}")
-  tvm_file_glob(GLOB RUNTIME_HEXAGON_SIM_SRCS src/runtime/hexagon/android/sim/*.cc)
-  include_directories(SYSTEM "${HEXAGON_TOOLCHAIN}/include/iss")
-  link_directories("${HEXAGON_TOOLCHAIN}/lib/iss")
-  list(APPEND TVM_RUNTIME_LINKER_LIBS "-lwrapper")
-  ExternalProject_Add(sim_dev
-    SOURCE_DIR "${CMAKE_SOURCE_DIR}/src/runtime/hexagon/android/sim/driver"
-    CMAKE_ARGS
-      "-DCMAKE_C_COMPILER=${HEXAGON_TOOLCHAIN}/bin/hexagon-clang"
-      "-DCMAKE_CXX_COMPILER=${HEXAGON_TOOLCHAIN}/bin/hexagon-clang++"
-      "-DHEXAGON_ARCH=${USE_HEXAGON_ARCH}"
-    INSTALL_COMMAND "true"
-  )
-elseif(USE_HEXAGON_DEVICE STREQUAL "${PICK_HW}")
-  find_hexagon_sdk_root("${USE_HEXAGON_SDK}" "${USE_HEXAGON_ARCH}")
-  find_hexagon_toolchain()
-  tvm_file_glob(GLOB RUNTIME_HEXAGON_DEVICE_SRCS src/runtime/hexagon/android/target/*.cc)
-
-  include_directories(SYSTEM
-    ${HEXAGON_SDK_INCLUDES}
-    ${HEXAGON_RPCMEM_ROOT}/inc
-    ${HEXAGON_REMOTE_ROOT}
-  )
+    # Add the simulator-specific RPC code into a shared library to be
+    # executed via run_main_on_sim.
+    add_library(hexagon_rpc_sim SHARED
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/simulator/rpc_server.cc"
+    )
+    target_link_libraries(hexagon_rpc_sim
+      -Wl,--whole-archive tvm_runtime -Wl,--no-whole-archive
+    )
 
-  list(APPEND TVM_RUNTIME_LINKER_LIBS "dl")
-  if(BUILD_FOR_ANDROID)
-    # Hexagon runtime uses __android_log_print, which is in liblog.
-    list(APPEND TVM_RUNTIME_LINKER_LIBS "log")
+  elseif(BUILD_FOR_HOST)
+    find_hexagon_toolchain()
+    add_hexagon_wrapper_paths()
+    file_glob_append(RUNTIME_HEXAGON_SRCS
+      "${TVMRT_SOURCE_DIR}/hexagon/host/*.cc"
+      "${TVMRT_SOURCE_DIR}/hexagon/rpc/simulator/session.cc"
+    )
+    list(APPEND TVM_RUNTIME_LINKER_LIBS "-lwrapper")
   endif()
-endif()
+endif()   # USE_HEXAGON_RPC
 
-set(RUNTIME_HEXAGON_COMMON_SRCS src/runtime/hexagon/hexagon_module.cc)
-if (USE_HEXAGON_DEVICE STREQUAL "${PICK_NONE}")
-  if(BUILD_FOR_HEXAGON)
-    tvm_file_glob(GLOB RUNTIME_HEXAGON_SRCS src/runtime/hexagon/hexagon/*.cc)
-  elseif(BUILD_FOR_ANDROID AND HEXAGON_SDK_PATH_DEFINED)
-  else()
-    tvm_file_glob(GLOB RUNTIME_HEXAGON_SRCS src/runtime/hexagon/host/*.cc)
-  endif()
-else()
-  tvm_file_glob(GLOB RUNTIME_HEXAGON_SRCS src/runtime/hexagon/android/*.cc)
-endif()
 
-list(APPEND RUNTIME_SRCS
-  ${RUNTIME_HEXAGON_SRCS}
-  ${RUNTIME_HEXAGON_SIM_SRCS}
-  ${RUNTIME_HEXAGON_DEVICE_SRCS}
-  ${RUNTIME_HEXAGON_COMMON_SRCS}
-)
+list(APPEND RUNTIME_SRCS ${RUNTIME_HEXAGON_SRCS})
diff --git a/cmake/modules/RustExt.cmake b/cmake/modules/RustExt.cmake
index 2922bc48dee2..e30caf0a0b04 100644
--- a/cmake/modules/RustExt.cmake
+++ b/cmake/modules/RustExt.cmake
@@ -16,8 +16,8 @@
 # under the License.
 
 if(USE_RUST_EXT)
-    set(RUST_SRC_DIR "${CMAKE_SOURCE_DIR}/rust")
-    set(CARGO_OUT_DIR "${CMAKE_SOURCE_DIR}/rust/target")
+    set(RUST_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/rust")
+    set(CARGO_OUT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/rust/target")
 
     if(USE_RUST_EXT STREQUAL "STATIC")
         set(COMPILER_EXT_PATH "${CARGO_OUT_DIR}/release/libcompiler_ext.a")
diff --git a/cmake/modules/StandaloneCrt.cmake b/cmake/modules/StandaloneCrt.cmake
index e63521e26dbf..c6e6dc77b442 100644
--- a/cmake/modules/StandaloneCrt.cmake
+++ b/cmake/modules/StandaloneCrt.cmake
@@ -57,7 +57,7 @@ if(USE_MICRO)
       math(EXPR job_spec_stop "${job_spec_length} - 3")
 
       list(GET job_spec 0 job_src_base)
-      set(job_src_base "${CMAKE_SOURCE_DIR}/${job_src_base}")
+      set(job_src_base "${CMAKE_CURRENT_SOURCE_DIR}/${job_src_base}")
       foreach(copy_pattern_index RANGE 1 "${job_spec_stop}" 3)
         list(GET job_spec ${copy_pattern_index} copy_pattern)
         math(EXPR copy_dest_index "${copy_pattern_index} + 2")
@@ -93,7 +93,7 @@ if(USE_MICRO)
     endforeach()
 
     set(make_common_args
-        "CRT_CONFIG=${CMAKE_SOURCE_DIR}/src/runtime/micro/crt_config.h"
+        "CRT_CONFIG=${CMAKE_CURRENT_SOURCE_DIR}/src/runtime/micro/crt_config.h"
         "BUILD_DIR=${host_build_dir_abspath}"
         "EXTRA_CFLAGS=-fPIC"
         "EXTRA_CXXFLAGS=-fPIC"
@@ -124,9 +124,9 @@ if(USE_MICRO)
     # Create the `crttest` target if we can find GTest.  If not, we create dummy
     # targets that give the user an informative error message.
     if(GTEST_FOUND)
-      tvm_file_glob(GLOB TEST_SRCS ${CMAKE_SOURCE_DIR}/tests/crt/*.cc)
+      tvm_file_glob(GLOB TEST_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/tests/crt/*.cc)
       add_executable(crttest ${TEST_SRCS})
-      target_include_directories(crttest SYSTEM PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/standalone_crt/include ${CMAKE_SOURCE_DIR}/src/runtime/micro)
+      target_include_directories(crttest SYSTEM PUBLIC ${CMAKE_CURRENT_BINARY_DIR}/standalone_crt/include ${CMAKE_CURRENT_SOURCE_DIR}/src/runtime/micro)
       target_link_libraries(crttest PRIVATE ${cmake_crt_libraries} GTest::GTest GTest::Main pthread dl)
       set_target_properties(crttest PROPERTIES EXCLUDE_FROM_ALL 1)
       set_target_properties(crttest PROPERTIES EXCLUDE_FROM_DEFAULT_BUILD 1)
diff --git a/cmake/modules/Zephyr.cmake b/cmake/modules/Zephyr.cmake
index 048240375cd6..a1cafce29d6e 100644
--- a/cmake/modules/Zephyr.cmake
+++ b/cmake/modules/Zephyr.cmake
@@ -43,7 +43,7 @@ if(USE_MICRO)
       math(EXPR job_spec_stop "${job_spec_length} - 3")
 
       list(GET job_spec 0 job_src_base)
-      set(job_src_base "${CMAKE_SOURCE_DIR}/${job_src_base}")
+      set(job_src_base "${CMAKE_CURRENT_SOURCE_DIR}/${job_src_base}")
       foreach(copy_pattern_index RANGE 1 "${job_spec_stop}" 3)
         list(GET job_spec ${copy_pattern_index} copy_pattern)
         math(EXPR copy_dest_index "${copy_pattern_index} + 2")
diff --git a/cmake/modules/contrib/EthosN.cmake b/cmake/modules/contrib/EthosN.cmake
index 244f461b251d..dbf5549180aa 100644
--- a/cmake/modules/contrib/EthosN.cmake
+++ b/cmake/modules/contrib/EthosN.cmake
@@ -40,8 +40,7 @@ if(NOT USE_ETHOSN STREQUAL "OFF")
 
     list(APPEND TVM_LINKER_LIBS ${ETHOSN_COMPILER_LIBRARY}
       ${ETHOSN_RUNTIME_LIBRARY})
-    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ETHOSN_COMPILER_LIBRARY}
-      ${ETHOSN_RUNTIME_LIBRARY})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ETHOSN_RUNTIME_LIBRARY})
 
     if(NOT MSVC)
       set_source_files_properties(${COMPILER_ETHOSN_SRCS}
diff --git a/cmake/utils/FindEthosN.cmake b/cmake/utils/FindEthosN.cmake
index 3b01374e1959..07eaf5ddeee1 100644
--- a/cmake/utils/FindEthosN.cmake
+++ b/cmake/utils/FindEthosN.cmake
@@ -58,7 +58,16 @@ macro(find_ethosn use_ethosn)
       PATHS ${__ethosn_stack}/lib)
     find_library(ETHOSN_COMPILER_LIBRARY NAMES EthosNSupport)
 
-    set(ETHOSN_PACKAGE_VERSION "0.1.1")
+    list(GET ETHOSN_INCLUDE_DIRS 0 filename)
+    set(filename "${filename}/ethosn_support_library/Support.hpp")
+    file(READ ${filename} ETHOSN_SUPPORT_H)
+    string(REGEX MATCH "VERSION_MAJOR ([0-9]*)" _ ${ETHOSN_SUPPORT_H})
+    set(ver_major ${CMAKE_MATCH_1})
+    string(REGEX MATCH "VERSION_MINOR ([0-9]*)" _ ${ETHOSN_SUPPORT_H})
+    set(ver_minor ${CMAKE_MATCH_1})
+    string(REGEX MATCH "VERSION_PATCH ([0-9]*)" _ ${ETHOSN_SUPPORT_H})
+    set(ver_patch ${CMAKE_MATCH_1})
+    set(ETHOSN_PACKAGE_VERSION "${ver_major}.${ver_minor}.${ver_patch}")
     set(ETHOSN_DEFINITIONS -DETHOSN_API_VERSION=${USE_ETHOSN_API_VERSION})
 
     # Runtime hardware support. Driver library also needed for
diff --git a/docker/Dockerfile.ci_cpu b/docker/Dockerfile.ci_cpu
index 962d738a9fc2..9f9f9a0b07c1 100644
--- a/docker/Dockerfile.ci_cpu
+++ b/docker/Dockerfile.ci_cpu
@@ -16,7 +16,6 @@
 # under the License.
 
 # CI docker CPU env
-# tag: v0.62
 FROM ubuntu:18.04
 
 RUN apt-get update --fix-missing
@@ -33,8 +32,8 @@ RUN pip config set global.no-cache-dir false
 COPY install/ubuntu_install_python_package.sh /install/ubuntu_install_python_package.sh
 RUN bash /install/ubuntu_install_python_package.sh
 
-COPY install/ubuntu1804_install_llvm.sh /install/ubuntu1804_install_llvm.sh
-RUN bash /install/ubuntu1804_install_llvm.sh
+COPY install/ubuntu_install_llvm_source.sh /install/ubuntu1804_install_llvm_source.sh
+RUN bash /install/ubuntu1804_install_llvm_source.sh
 
 COPY install/ubuntu_install_dnnl.sh /install/ubuntu_install_dnnl.sh
 RUN bash /install/ubuntu_install_dnnl.sh
@@ -99,6 +98,10 @@ COPY install/ubuntu_download_arm_compute_lib_binaries.sh /install/ubuntu_downloa
 RUN bash /install/ubuntu_download_arm_compute_lib_binaries.sh
 
 # Caffe deps
+COPY install/ubuntu_install_boost.sh /install/ubuntu_install_boost.sh
+RUN bash /install/ubuntu_install_boost.sh
+
+# Caffe
 COPY install/ubuntu_install_caffe.sh /install/ubuntu_install_caffe.sh
 RUN bash /install/ubuntu_install_caffe.sh
 
diff --git a/docker/Dockerfile.ci_hexagon b/docker/Dockerfile.ci_hexagon
new file mode 100644
index 000000000000..f1293019bbba
--- /dev/null
+++ b/docker/Dockerfile.ci_hexagon
@@ -0,0 +1,54 @@
+# 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.
+
+# CI docker Hexagon env
+# tag: v0.01
+FROM tvmcihexagon/ci-hexagon-base:v0.01_SDK4.2.0.2
+
+RUN apt-get update --fix-missing
+RUN apt-get install -y ca-certificates gnupg2 libxml2-dev
+
+COPY install/ubuntu_install_core.sh /install/ubuntu_install_core.sh
+RUN bash /install/ubuntu_install_core.sh
+
+COPY install/ubuntu1804_install_python.sh /install/ubuntu1804_install_python.sh
+RUN bash /install/ubuntu1804_install_python.sh
+
+# Globally disable pip cache
+RUN pip config set global.cache-dir false
+
+COPY install/ubuntu_install_python_package.sh /install/ubuntu_install_python_package.sh
+RUN bash /install/ubuntu_install_python_package.sh
+
+# ANTLR deps
+COPY install/ubuntu_install_java.sh /install/ubuntu_install_java.sh
+RUN bash /install/ubuntu_install_java.sh
+
+# Android SDK
+COPY install/ubuntu_install_androidsdk.sh /install/ubuntu_install_androidsdk.sh
+RUN bash /install/ubuntu_install_androidsdk.sh
+ENV ANDROID_HOME=/opt/android-sdk-linux
+ENV ANDROID_NDK_HOME=/opt/android-sdk-linux/ndk/21.3.6528147
+ENV PATH /opt/android-sdk-linux/platform-tools:$PATH
+
+# Hexagon
+# HEXAGON_SDK_PATH is and env variable in tvmcihexagon/ci-hexagon-base
+COPY install/ubuntu_install_hexagon.sh /install/ubuntu_install_hexagon.sh
+RUN bash /install/ubuntu_install_hexagon.sh
+ENV CLANG_LLVM_HOME /opt/clang-llvm
+ENV LD_LIBRARY_PATH $LD_LIBRARY_PATH:/opt/clang-llvm/lib
+ENV HEXAGON_TOOLCHAIN "${HEXAGON_SDK_PATH}/tools/HEXAGON_Tools/8.4.09/Tools"
diff --git a/docker/Dockerfile.ci_i386 b/docker/Dockerfile.ci_i386
index d4ce370c4205..e5e02b9755e4 100644
--- a/docker/Dockerfile.ci_i386
+++ b/docker/Dockerfile.ci_i386
@@ -18,7 +18,7 @@
 # CI docker i386 env
 # tag: v0.53
 
-FROM ioft/i386-ubuntu:16.04
+FROM i386/ubuntu:18.04
 
 RUN apt-get update --fix-missing && apt-get install -y ca-certificates
 
@@ -28,8 +28,8 @@ RUN bash /install/ubuntu_install_core.sh
 COPY install/ubuntu_install_llvm.sh /install/ubuntu_install_llvm.sh
 RUN bash /install/ubuntu_install_llvm.sh
 
-COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
-RUN bash /install/ubuntu_install_python.sh
+COPY install/ubuntu1804_install_python.sh /install/ubuntu1804_install_python.sh
+RUN bash /install/ubuntu1804_install_python.sh
 
 # Globally disable pip cache
 RUN pip config set global.no-cache-dir false
diff --git a/docker/bash.sh b/docker/bash.sh
index 6f31aa7a5180..4a695c96f0ce 100755
--- a/docker/bash.sh
+++ b/docker/bash.sh
@@ -314,7 +314,7 @@ fi
 # PID namespace (--pid=host).  The process inside does not have pid 1
 # and SIGKILL is propagated to the process inside, allowing jenkins to
 # kill it if needed.
-DOCKER_FLAGS+=( --rm --pid=host)
+DOCKER_FLAGS+=( --pid=host)
 
 # Expose services running in container to the host.
 if $USE_NET_HOST; then
diff --git a/docker/build.sh b/docker/build.sh
index 39fe7a024246..de320115e2cd 100755
--- a/docker/build.sh
+++ b/docker/build.sh
@@ -79,7 +79,7 @@ if [[ "$1" == "--net=host" ]]; then
     shift 1
 fi
 
-DOCKER_NO_CACHE_ARG=--no-cache
+#DOCKER_NO_CACHE_ARG=--no-cache
 
 if [[ "$1" == "--cache-from" ]]; then
     shift 1
diff --git a/docker/dev_common.sh b/docker/dev_common.sh
index 6823f754ae16..11ba34eb06e2 100644
--- a/docker/dev_common.sh
+++ b/docker/dev_common.sh
@@ -56,6 +56,11 @@ function lookup_image_spec() {
 }
 
 function run_docker() {
+    docker_bash_args=( )
+    while [ "x${1:0:1}" == "x-" ]; do
+        docker_bash_args=( "${docker_bash_args[@]}" "$1" )
+        shift
+    done
     image_name="$1"  # Name of the Jenkinsfile var to find
     shift
 
@@ -65,5 +70,6 @@ function run_docker() {
         exit 2
     fi
 
-    "${GIT_TOPLEVEL}/docker/bash.sh" "${image_spec}" "$@"
+    docker_bash_args=( "${docker_bash_args[@]}" "${image_spec}" "$@" )
+    "${GIT_TOPLEVEL}/docker/bash.sh" "${docker_bash_args[@]}"
 }
diff --git a/docker/install/ubuntu1804_install_python.sh b/docker/install/ubuntu1804_install_python.sh
index 5d39918f918c..4676ffb8b20d 100755
--- a/docker/install/ubuntu1804_install_python.sh
+++ b/docker/install/ubuntu1804_install_python.sh
@@ -20,13 +20,27 @@ set -e
 set -u
 set -o pipefail
 
-# install python and pip, don't modify this, modify install_python_package.sh
+
+cleanup() {
+  rm -rf base-requirements.txt
+}
+
+trap cleanup 0
+
+
+# Install python and pip. Don't modify this to add Python package dependencies,
+# instead modify install_python_package.sh
 apt-get update
 apt-get install -y software-properties-common
-apt-get install -y python3-dev python3-setuptools
-
-# Install pip
-cd /tmp && wget -q https://bootstrap.pypa.io/pip/3.6/get-pip.py && python3 get-pip.py
+apt-get install -y python3.7 python3.7-dev python3-pip
+update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.7 1
 
 # Pin pip and setuptools versions
-pip3 install pip==19.3.1 setuptools==58.4.0
+# Hashes generated via:
+#   $ pip download <package>==<version>
+#   $ pip hash --algorithm sha512 <package>.whl
+cat <<EOF > base-requirements.txt
+pip==22.0.3 --hash=sha512:12ca75130a1ce9807060a66dd2341afc7c7e663357ca4b937868dbc733634e11bae49ffff96acb0f5f3fb16cb14f680b9b6d185155a711c6098eda5cfbf2f8f5
+setuptools==60.9.1 --hash=sha512:24c21006f0650209e6a934a5366614e32a98fbdf11bc0e941419731aa3c8cb6d217e486608b4834b8f89c14dc36b3bafa30e795c87b58705be2b156f78acc3c5
+EOF
+pip3 install -r base-requirements.txt
diff --git a/docker/install/ubuntu_install_boost.sh b/docker/install/ubuntu_install_boost.sh
new file mode 100644
index 000000000000..e226bbc5d96e
--- /dev/null
+++ b/docker/install/ubuntu_install_boost.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+# 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.
+
+set -euxo pipefail
+
+cleanup() {
+  rm -rf /boost_1_67_0 /boost_1_67_0.tar.gz
+}
+
+trap cleanup 0
+
+curl -LO https://boostorg.jfrog.io/artifactory/main/release/1.67.0/source/boost_1_67_0.tar.gz
+BOOST_HASH=8c247e040303a97895cee9c9407ef205e2c3ab09f0b8320997835ad6221dff23a87231629498ccfd0acca473f74e9ec27b8bd774707b062228df1e5f72d44c92
+echo "$BOOST_HASH" boost_1_67_0.tar.gz | sha512sum -c
+tar -xf boost_1_67_0.tar.gz
+
+pushd boost_1_67_0
+./bootstrap.sh --with-python="$(which python3.7)"
+./b2 install --with-python --with-system --with-filesystem --with-thread --with-regex
+popd
+
+ln -s /usr/local/lib/libboost_python37.so.1.67.0 /usr/local/lib/libboost_python.so
+ln -s /usr/local/lib/libboost_python37.a /usr/local/lib/libboost_python.a
diff --git a/docker/install/ubuntu_install_caffe.sh b/docker/install/ubuntu_install_caffe.sh
index b26f339c3fa1..ab71eab54a2d 100644
--- a/docker/install/ubuntu_install_caffe.sh
+++ b/docker/install/ubuntu_install_caffe.sh
@@ -16,17 +16,15 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set -e
-set -u
-set -o pipefail
+set -euxo pipefail
 
 apt-get update --fix-missing
 
-# Install dependencies
-apt-get install -y --no-install-recommends libboost-filesystem-dev libboost-python-dev \
-    libboost-system-dev libboost-thread-dev libboost-regex-dev protobuf-compiler \
+# # Install dependencies
+apt-get install -y --no-install-recommends protobuf-compiler \
     libprotobuf-dev libhdf5-serial-dev libopenblas-dev libgflags-dev libgoogle-glog-dev
 
+
 # install python packages
 pip install "numpy" "protobuf" "scikit-image" "six"
 
@@ -36,15 +34,23 @@ CAFFE_HOME="/opt/caffe"
 git clone --branch=ssd --depth 1 https://github.com/weiliu89/caffe /caffe_src
 cd /caffe_src
 
+
 echo "Building Caffe"
 mkdir /caffe_src/build && cd /caffe_src/build
-cmake .. -DCMAKE_INSTALL_PREFIX=${CAFFE_HOME} -DCMAKE_BUILD_TYPE=Release -DCPU_ONLY=1 \
-    -Dpython_version=3 -DUSE_OPENCV=OFF -DUSE_LEVELDB=OFF -DUSE_LMDB=OFF -DBUILD_docs=OFF -DBLAS=open
+cmake -DCMAKE_INSTALL_PREFIX=${CAFFE_HOME}\
+    -DCMAKE_BUILD_TYPE=Release \
+    -DCPU_ONLY=1 \
+    -Dpython_version=3 \
+    -DUSE_OPENCV=OFF \
+    -DUSE_LEVELDB=OFF \
+    -DUSE_LMDB=OFF \
+    -DBUILD_docs=OFF \
+    -DBLAS=open \
+    ..
+
 make all -j$(expr $(nproc) - 1)
 make pycaffe -j$(expr $(nproc) - 1)
 make test -j$(expr $(nproc) - 1)
-make runtest -j$(expr $(nproc) - 1)
-make pytest -j$(expr $(nproc) - 1)
 
 echo "Installing Caffe to /opt/caffe"
 make install
@@ -54,4 +60,4 @@ cd / && rm -rf /caffe_src
 
 PYCAFFE_ROOT=${CAFFE_HOME}/python
 echo "${CAFFE_HOME}/lib" >> /etc/ld.so.conf.d/caffe.conf && ldconfig
-ln -s ${PYCAFFE_ROOT}/caffe /usr/local/lib/python3.6/dist-packages/caffe
+ln -s ${PYCAFFE_ROOT}/caffe /usr/local/lib/python3.7/dist-packages/caffe
diff --git a/docker/install/ubuntu_install_dnnl.sh b/docker/install/ubuntu_install_dnnl.sh
index fb0efd52a2c1..34f917ce6606 100755
--- a/docker/install/ubuntu_install_dnnl.sh
+++ b/docker/install/ubuntu_install_dnnl.sh
@@ -21,9 +21,9 @@ set -u
 set -o pipefail
 
 cd /usr/local/
-wget -q https://github.com/oneapi-src/oneDNN/releases/download/v1.5/dnnl_lnx_1.5.0_cpu_gomp.tgz
-tar -xzf dnnl_lnx_1.5.0_cpu_gomp.tgz
-mv dnnl_lnx_1.5.0_cpu_gomp/include/* /usr/local/include/
-mv dnnl_lnx_1.5.0_cpu_gomp/lib/libdnnl* /usr/local/lib/
-rm -rf dnnl_lnx_1.5.0_cpu_gomp.tgz dnnl_lnx_1.5.0_cpu_gomp
+wget -q https://github.com/oneapi-src/oneDNN/releases/download/v2.2/dnnl_lnx_2.2.0_cpu_gomp.tgz
+tar -xzf dnnl_lnx_2.2.0_cpu_gomp.tgz
+mv dnnl_lnx_2.2.0_cpu_gomp/include/* /usr/local/include/
+mv dnnl_lnx_2.2.0_cpu_gomp/lib/libdnnl* /usr/local/lib/
+rm -rf dnnl_lnx_2.2.0_cpu_gomp.tgz dnnl_lnx_2.2.0_cpu_gomp
 cd -
diff --git a/docker/install/ubuntu_install_hexagon.sh b/docker/install/ubuntu_install_hexagon.sh
new file mode 100644
index 000000000000..46d2a44cfaa5
--- /dev/null
+++ b/docker/install/ubuntu_install_hexagon.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+# 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.
+
+set -o errexit -o nounset
+set -o pipefail
+
+# Install LLVM/clang
+CLANG_LLVM_HOME=/opt/clang-llvm
+CLANG_LLVM_VERSION=13.0.0
+CLANG_LLVM_FILENAME=clang_llvm.tar.xz
+wget -q https://github.com/llvm/llvm-project/releases/download/llvmorg-${CLANG_LLVM_VERSION}/clang+llvm-${CLANG_LLVM_VERSION}-x86_64-linux-gnu-ubuntu-16.04.tar.xz -O ${CLANG_LLVM_FILENAME}
+mkdir ${CLANG_LLVM_HOME}
+tar -xvf ${CLANG_LLVM_FILENAME} -C ${CLANG_LLVM_HOME} --strip-components=1
+rm ${CLANG_LLVM_FILENAME}
diff --git a/docker/install/ubuntu_install_llvm.sh b/docker/install/ubuntu_install_llvm.sh
index 1f4eec6ee6ef..6616bfc5eb2b 100755
--- a/docker/install/ubuntu_install_llvm.sh
+++ b/docker/install/ubuntu_install_llvm.sh
@@ -16,9 +16,10 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set -e
-set -u
-set -o pipefail
+set -euxo pipefail
+
+apt-get update
+apt-get install -y gnupg
 
 echo deb http://apt.llvm.org/xenial/ llvm-toolchain-xenial-4.0 main\
      >> /etc/apt/sources.list.d/llvm.list
@@ -46,4 +47,4 @@ echo deb-src http://apt.llvm.org/xenial/ llvm-toolchain-xenial main\
      >> /etc/apt/sources.list.d/llvm.list
 
 wget -q -O - http://apt.llvm.org/llvm-snapshot.gpg.key|sudo apt-key add -
-apt-get update && apt-get install -y llvm-4.0 llvm-9 llvm-8 llvm-7 clang-9 libclang-9-dev clang-8 libclang-8-dev clang-7 libclang-7-dev
+apt-get install -y llvm-4.0 llvm-9 llvm-8 llvm-7 clang-9 libclang-9-dev clang-8 libclang-8-dev clang-7 libclang-7-dev
diff --git a/docker/lint.sh b/docker/lint.sh
index 4bad5ea3b923..a968bc1e6421 100755
--- a/docker/lint.sh
+++ b/docker/lint.sh
@@ -37,6 +37,9 @@ function run_lint_step() {
             ;;
         asf)
             cmd=( tests/lint/check_asf_header.sh --local )
+            if [ $inplace_fix -eq 1 ]; then
+                cmd=( "${cmd[@]}" --fix )
+            fi
             ;;
         clang_format)
             if [ $inplace_fix -eq 0 ]; then
@@ -77,6 +80,9 @@ function run_lint_step() {
             echo "error: don't know how to run lint step: $1" >&2
             echo "usage: ${SCRIPT_NAME} [-i] <lint_step>" >&2
             echo >&2
+            echo "options:" >&2
+            echo " -i    Fix lint errors in-place where possible (modifies non-compliant files)" >&2
+            echo >&2
             echo "available lint_step: ${DEFAULT_STEPS[@]}" >&2
             exit 2
             ;;
@@ -84,7 +90,7 @@ function run_lint_step() {
     shift
 
     if [ $validate_only -eq 0 ]; then
-        run_docker "ci_lint" "${cmd[@]}"
+        run_docker -it "ci_lint" "${cmd[@]}"
     fi
 }
 
diff --git a/docker/with_the_same_user b/docker/with_the_same_user
index 89ae8f5044da..f8ebc1482d1e 100644
--- a/docker/with_the_same_user
+++ b/docker/with_the_same_user
@@ -1,4 +1,5 @@
 #!/usr/bin/env bash
+set -x
 
 # Licensed to the Apache Software Foundation (ASF) under one
 # or more contributor license agreements.  See the NOTICE file
diff --git a/docs/reference/api/python/contrib.rst b/docs/reference/api/python/contrib.rst
index 0eb3024c2d08..52d3faff0fc4 100644
--- a/docs/reference/api/python/contrib.rst
+++ b/docs/reference/api/python/contrib.rst
@@ -93,6 +93,16 @@ tvm.contrib.random
     :members:
 
 
+tvm.contrib.relay_viz
+~~~~~~~~~~~~~~~~~~~~~
+.. automodule:: tvm.contrib.relay_viz
+    :members:
+.. automodule:: tvm.contrib.relay_viz.interface
+    :members:
+.. automodule:: tvm.contrib.relay_viz.terminal
+    :members:
+
+
 tvm.contrib.rocblas
 ~~~~~~~~~~~~~~~~~~~
 .. automodule:: tvm.contrib.rocblas
diff --git a/gallery/how_to/work_with_relay/using_relay_viz.py b/gallery/how_to/work_with_relay/using_relay_viz.py
new file mode 100644
index 000000000000..f61fc41f4f14
--- /dev/null
+++ b/gallery/how_to/work_with_relay/using_relay_viz.py
@@ -0,0 +1,159 @@
+# 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.
+# pylint: disable=line-too-long
+"""
+Use Relay Visualizer to Visualize Relay
+============================================================
+**Author**: `Chi-Wei Wang <https://github.com/chiwwang>`_
+
+Relay IR module can contain lots of operations. Although an individual
+operation is usually easy to understand, putting them together can cause
+a complicated, hard-to-read graph. Things can get even worse with optimiztion-passes
+coming into play.
+
+This utility visualizes an IR module as nodes and edges. It defines a set of interfaces including
+parser, plotter(renderer), graph, node, and edges.
+A default parser is provided. Users can implement their own renderers to render the graph.
+
+Here we use a renderer rendering graph in the text-form.
+It is a lightweight, AST-like visualizer, inspired by `clang ast-dump <https://clang.llvm.org/docs/IntroductionToTheClangAST.html>`_.
+We will introduce how to implement customized parsers and renderers through interface classes.
+"""
+from typing import (
+    Dict,
+    Union,
+    Tuple,
+    List,
+)
+import tvm
+from tvm import relay
+from tvm.contrib import relay_viz
+from tvm.contrib.relay_viz.interface import (
+    VizEdge,
+    VizNode,
+    VizParser,
+)
+from tvm.contrib.relay_viz.terminal import (
+    TermGraph,
+    TermPlotter,
+    TermVizParser,
+)
+
+######################################################################
+# Define a Relay IR Module with multiple GlobalVar
+# ------------------------------------------------
+# Let's build an example Relay IR Module containing multiple ``GlobalVar``.
+# We define an ``add`` function and call it in the main function.
+data = relay.var("data")
+bias = relay.var("bias")
+add_op = relay.add(data, bias)
+add_func = relay.Function([data, bias], add_op)
+add_gvar = relay.GlobalVar("AddFunc")
+
+input0 = relay.var("input0")
+input1 = relay.var("input1")
+input2 = relay.var("input2")
+add_01 = relay.Call(add_gvar, [input0, input1])
+add_012 = relay.Call(add_gvar, [input2, add_01])
+main_func = relay.Function([input0, input1, input2], add_012)
+main_gvar = relay.GlobalVar("main")
+
+mod = tvm.IRModule({main_gvar: main_func, add_gvar: add_func})
+
+######################################################################
+# Render the graph with Relay Visualizer on the terminal
+# ------------------------------------------------------
+# The terminal can show a Relay IR module in text similar to clang AST-dump.
+# We should see ``main`` and ``AddFunc`` function. ``AddFunc`` is called twice in the ``main`` function.
+viz = relay_viz.RelayVisualizer(mod)
+viz.render()
+
+######################################################################
+# Customize Parser for Interested Relay Types
+# -------------------------------------------
+# Sometimes we want to emphasize interested information, or parse things differently for a specific usage.
+# It is possible to provide customized parsers as long as it obeys the interface.
+# Here demostrate how to customize parsers for ``relay.var``.
+# We need to implement abstract interface :py:class:`tvm.contrib.relay_viz.interface.VizParser`.
+class YourAwesomeParser(VizParser):
+    def __init__(self):
+        self._delegate = TermVizParser()
+
+    def get_node_edges(
+        self,
+        node: relay.Expr,
+        relay_param: Dict[str, tvm.runtime.NDArray],
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+
+        if isinstance(node, relay.Var):
+            node = VizNode(node_to_id[node], "AwesomeVar", f"name_hint {node.name_hint}")
+            # no edge is introduced. So return an empty list.
+            return node, []
+
+        # delegate other types to the other parser.
+        return self._delegate.get_node_edges(node, relay_param, node_to_id)
+
+
+######################################################################
+# Pass the parser and an interested renderer to visualizer.
+# Here we just the terminal renderer.
+viz = relay_viz.RelayVisualizer(mod, {}, TermPlotter(), YourAwesomeParser())
+viz.render()
+
+######################################################################
+# Customization around Graph and Plotter
+# -------------------------------------------
+# Besides parsers, we can also customize graph and renderers by implementing
+# abstract class :py:class:`tvm.contrib.relay_viz.interface.VizGraph` and
+# :py:class:`tvm.contrib.relay_viz.interface.Plotter`.
+# Here we override the ``TermGraph`` defined in ``terminal.py`` for easier demo.
+# We add a hook duplicating above ``AwesomeVar``, and make ``TermPlotter`` use the new class.
+class AwesomeGraph(TermGraph):
+    def node(self, viz_node):
+        # add the node first
+        super().node(viz_node)
+        # if it's AwesomeVar, duplicate it.
+        if viz_node.type_name == "AwesomeVar":
+            duplicated_id = f"duplciated_{viz_node.identity}"
+            duplicated_type = "double AwesomeVar"
+            super().node(VizNode(duplicated_id, duplicated_type, ""))
+            # connect the duplicated var to the original one
+            super().edge(VizEdge(duplicated_id, viz_node.identity))
+
+
+# override TermPlotter to use `AwesomeGraph` instead
+class AwesomePlotter(TermPlotter):
+    def create_graph(self, name):
+        self._name_to_graph[name] = AwesomeGraph(name)
+        return self._name_to_graph[name]
+
+
+viz = relay_viz.RelayVisualizer(mod, {}, AwesomePlotter(), YourAwesomeParser())
+viz.render()
+
+######################################################################
+# Summary
+# -------
+# This tutorial demonstrates the usage of Relay Visualizer and customization.
+# The class :py:class:`tvm.contrib.relay_viz.RelayVisualizer` is composed of interfaces
+# defined in ``interface.py``.
+#
+# It is aimed for quick look-then-fix iterations.
+# The constructor arguments are intended to be simple, while the customization is still
+# possible through a set of interface classes.
+#
diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h
index 0e43abb54b93..4a00de802c61 100644
--- a/include/tvm/ir/expr.h
+++ b/include/tvm/ir/expr.h
@@ -178,6 +178,11 @@ class RelayExprNode : public BaseExprNode {
    *
    * For expressions that have the function type, the virtual device describes where the result of
    * the call to the function or closure is stored (instead of where the function itself is stored).
+   * For example, the virtual device of f = fn(x) { body } is the virtual device of f(y), not where
+   * the function itself is stored. Note that f(y)'s virtual device will be the same as the virtual
+   * device of body. For more details, see the documentation in
+   * src/relay/transforms/device_planner.cc.
+   *
    * The VirtualDevice's Target field describes how the body of the function should be compiled.
    *
    * Set to VirtualDevice::FullyUnconstrained by default.
@@ -190,6 +195,13 @@ class RelayExprNode : public BaseExprNode {
   /*!
    * \return The virtual device (VirtualDevice).
    * If the virtual device is not defined, returns VirtualDevice::FullyUnconstrained().
+   * Note that for function types, the virtual device is the device where the result of a
+   * call to the function is stored, not where the function itself lives.
+   * For example, the virtual device of f = fn(x) { body } is the virtual device of f(y), not where
+   * the function itself is stored. Note that f(y)'s virtual device will be the same as the virtual
+   * device of body.
+   *
+   * See the documentation of the virtual_device_ field (above) for more details.
    */
   VirtualDevice virtual_device() const;
 
diff --git a/include/tvm/ir/function.h b/include/tvm/ir/function.h
index 051c05dd3d01..1493544e7324 100644
--- a/include/tvm/ir/function.h
+++ b/include/tvm/ir/function.h
@@ -190,26 +190,6 @@ constexpr const char* kTarget = "target";
  */
 constexpr const char* kGlobalSymbol = "global_symbol";
 
-/*!
- * \brief The \p VirtualDevice which will hold each of the functions parameters.
- *
- * Only supported on Relay \p Functions. Generally added by the \p PlanDevices pass, but
- * may be included as an annotation on user programs.
- *
- * Type: Array<VirtualDevice>
- */
-constexpr const char* kParamVirtualDevice = "param_virtual_devices";
-
-/*!
- * \brief The \p VirtualDevice which will hold the function result.
- *
- * Only supported on Relay \p Functions. Generally added by the \p PlanDevices pass, but
- * may be included as an annotation on user programs.
- *
- * Type: VirtualDevice
- */
-constexpr const char* kResultVirtualDevice = "result_virtual_device";
-
 }  // namespace attr
 }  // namespace tvm
 #endif  // TVM_IR_FUNCTION_H_
diff --git a/include/tvm/ir/memory_pools.h b/include/tvm/ir/memory_pools.h
index c6e52648ebd4..d4eefdc910b0 100644
--- a/include/tvm/ir/memory_pools.h
+++ b/include/tvm/ir/memory_pools.h
@@ -103,29 +103,34 @@ struct PoolInfoNode : public Object {
   TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoNode, Object);
 };
 
+/*!
+ * \brief The string parameter to indicate read and write access to a pool
+ * This needs to be kept in sync with PoolInfo.READ_WRITE_ACCESS in
+ * python/tvm/ir/memory_pools.py
+ */
+static constexpr const char* kTargetPoolReadWriteAccess = "rw";
+
+/*!
+ * \brief The string parameter to indicate read only access to a pool
+ * This needs to be kept in sync with PoolInfo.READ_ONLY_ACCESS in
+ * python/tvm/ir/memory_pools.py
+ */
+static constexpr const char* kTargetPoolReadOnlyAccess = "ro";
+
+/*! \brief The PoolSize is unrestricted for the memory planner */
+static const int kUnrestrictedPoolSizeHint = -1;
+
+/*! \brief The clock frequency is not known */
+static const int kUnknownClockFrequency = -1;
+
+/*! \brief The read bandwidth is not known */
+static const int kUnknownReadBandwidth = -1;
+
+/*! \brief The write bandwidth is not known */
+static const int kUnknownWriteBandwidth = -1;
+
 class PoolInfo : public ObjectRef {
  public:
-  /*!
-   * \brief The string parameter to indicate read and write access to a pool
-   * This needs to be kept in sync with PoolInfo.READ_WRITE_ACCESS in
-   * python/tvm/ir/memory_pools.py
-   */
-  static constexpr const char* kTargetPoolReadWriteAccess = "rw";
-  /*!
-   * \brief The string parameter to indicate read only access to a pool
-   * This needs to be kept in sync with PoolInfo.READ_ONLY_ACCESS in
-   * python/tvm/ir/memory_pools.py
-   */
-  static constexpr const char* kTargetPoolReadOnlyAccess = "ro";
-  /*! \brief The PoolSize is unrestricted for the memory planner */
-  static const int kUnrestrictedPoolSizeHint = -1;
-  /*! \brief The clock frequency is not known */
-  static const int kUnknownClockFrequency = -1;
-  /*! \brief The read bandwidth is not known */
-  static const int kUnknownReadBandwidth = -1;
-  /*! \brief The write bandwidth is not known */
-  static const int kUnknownWriteBandwidth = -1;
-
   TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
                    Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
                    Integer clock_frequency_hz = kUnknownClockFrequency,
diff --git a/include/tvm/ir/module.h b/include/tvm/ir/module.h
index d308db87af8b..e2b47ef324df 100644
--- a/include/tvm/ir/module.h
+++ b/include/tvm/ir/module.h
@@ -40,7 +40,44 @@
 #include <vector>
 
 namespace tvm {
+/*!
+ * \brief Describes one parameter that should be linked into the generated module.
+ *
+ * When parameters are to be linked in with generated code (i.e. on target_host-compatible
+ * backends), Relay attaches instances of this object to a global TIR function. Code-generators
+ * use the information contained in this node to include the parameter data in the generated
+ * module.
+ */
+class LinkedParamNode : public Object {
+ public:
+  /*! \brief Unique numeric identifier used by runtimes to lookup this parameter. */
+  int64_t id;
+
+  /*! \brief Parameter data which should get linked into the final module. */
+  ::tvm::runtime::NDArray param;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("id", &id);
+    v->Visit("param", &param);
+  }
+
+  static constexpr const char* _type_key = "tir.LinkedParam";
+  TVM_DECLARE_FINAL_OBJECT_INFO(LinkedParamNode, Object);
+};
+
+/*!
+ * \brief Managed reference to LinkedParamNode.
+ */
+class LinkedParam : public ObjectRef {
+ public:
+  TVM_DLL LinkedParam(int64_t id, tvm::runtime::NDArray param);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(LinkedParam, ObjectRef, LinkedParamNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(LinkedParamNode);
+};
+
 class IRModule;
+
 /*!
  * \brief IRModule that holds functions and type definitions.
  *
@@ -504,6 +541,11 @@ constexpr const char* kRuntime = "runtime";
  */
 constexpr const char* kWorkspaceMemoryPools = "workspace_memory_pools";
 
+/*
+ * \brief Module attribute for tir constants
+ */
+constexpr const char* kConstantsArray = "Constants";
+
 }  // namespace attr
 }  // namespace tvm
 #endif  // TVM_IR_MODULE_H_
diff --git a/include/tvm/ir/transform.h b/include/tvm/ir/transform.h
index cb556fc13de7..d8f6632a66ca 100644
--- a/include/tvm/ir/transform.h
+++ b/include/tvm/ir/transform.h
@@ -392,7 +392,62 @@ class Pass : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(Pass, ObjectRef, PassNode);
 };
 
-class SequentialNode;
+/*!
+ * \brief The SequentialNode contains a set of passes that transform Relay
+ * programs from one AST to another semantically equivalent one.
+ *
+ * One example of this level of pass is that the pass manager needs to correctly
+ * perform a host of optimizations with a given optimization level and disabled
+ * passes.
+ */
+class SequentialNode : public PassNode {
+ public:
+  /* \brief The pass meta data.*/
+  PassInfo pass_info;
+
+  /*! \brief A list of passes that used to compose a sequential pass. */
+  tvm::Array<Pass> passes;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("pass_info", &pass_info);
+    v->Visit("passes", &passes);
+  }
+
+  /*!
+   * \brief Get the pass information/meta data.
+   */
+  PassInfo Info() const override { return pass_info; }
+
+  /*!
+   * \brief Resolve the pass dependency. It globs all required passes by
+   *        a given pass and executes them.
+   *
+   * \param mod The module that an optimization pass runs on.
+   *
+   * \return The updated module after resolving pass dependencies.
+   *
+   * TODO(zhiics) Build a dependency graph among the passes using provided
+   * metadata, i.e. required_passes. Likely, we can have a data structure, i.e.
+   * PassInfo, to store the relevant information including the parent passes.
+   */
+  void ResolveDependency(const IRModule& mod);
+
+  /*!
+   * \brief Perform optimizations on a series of passes. The aforementioned
+   *        typical pass manager jobs could be done by it. This function could
+   *        be overloaded to focus on different metrics, i.e. performance,
+   *        memory footprint, etc.
+   *
+   * \param mod The module that these passes are applied on.
+   * \param pass_ctx The context that these passes execute on.
+   *
+   * \return Return the updated module.
+   */
+  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
+
+  static constexpr const char* _type_key = "transform.Sequential";
+  TVM_DECLARE_FINAL_OBJECT_INFO(SequentialNode, PassNode);
+};
 
 class Sequential : public Pass {
  public:
@@ -418,7 +473,7 @@ class Sequential : public Pass {
   explicit Sequential(ObjectPtr<Object> n) : Pass(n) {}
 
   const SequentialNode* operator->() const;
-  using ContainerType = Sequential;
+  using ContainerType = SequentialNode;
 };
 
 /*
diff --git a/include/tvm/node/structural_hash.h b/include/tvm/node/structural_hash.h
index 887a012cfc93..a30a2c59d0d1 100644
--- a/include/tvm/node/structural_hash.h
+++ b/include/tvm/node/structural_hash.h
@@ -25,6 +25,7 @@
 
 #include <tvm/node/functor.h>
 #include <tvm/runtime/data_type.h>
+#include <tvm/runtime/ndarray.h>
 
 #include <functional>
 #include <string>
@@ -199,5 +200,13 @@ class SHashReducer {
   bool map_free_vars_;
 };
 
+class SEqualReducer;
+struct NDArrayContainerTrait {
+  static constexpr const std::nullptr_t VisitAttrs = nullptr;
+  static void SHashReduce(const runtime::NDArray::Container* key, SHashReducer hash_reduce);
+  static bool SEqualReduce(const runtime::NDArray::Container* lhs,
+                           const runtime::NDArray::Container* rhs, SEqualReducer equal);
+};
+
 }  // namespace tvm
 #endif  // TVM_NODE_STRUCTURAL_HASH_H_
diff --git a/include/tvm/relay/dataflow_matcher.h b/include/tvm/relay/dataflow_matcher.h
index 10e461645c8b..8dd5fbdd5eac 100644
--- a/include/tvm/relay/dataflow_matcher.h
+++ b/include/tvm/relay/dataflow_matcher.h
@@ -106,6 +106,17 @@ Expr RewritePatterns(Array<DFPatternCallback> callbacks, Expr expr, IRModule mod
  */
 Expr PartitionPattern(DFPattern pattern, Expr expr, Map<String, ObjectRef> attrs, PackedFunc check);
 
+/*!
+ * \brief Infer the type of an expression.
+ *
+ * \param expr The expression to rewrite
+ *
+ * \return Return An Expr with unambiguous type information filled in, as well as it's
+ * checked type field populated with the result type.
+ *
+ */
+Expr InferType(const Expr& expr);
+
 }  // namespace relay
 }  // namespace tvm
 
diff --git a/include/tvm/relay/executor.h b/include/tvm/relay/executor.h
index 6d1bd2de7f57..858ba5cfe198 100644
--- a/include/tvm/relay/executor.h
+++ b/include/tvm/relay/executor.h
@@ -113,6 +113,8 @@ class ExecutorNode : public Object {
   }
 
   static constexpr const char* _type_key = "Executor";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(ExecutorNode, Object);
 };
 
@@ -122,8 +124,6 @@ class ExecutorNode : public Object {
  */
 class Executor : public ObjectRef {
  public:
-  Executor() = default;
-
   /*!
    * \brief Create a new Executor object using the registry
    * \throws Error if name is not registered
@@ -147,7 +147,8 @@ class Executor : public ObjectRef {
   TVM_DLL static Map<String, String> ListExecutorOptions(const String& name);
 
   /*! \brief specify container node */
-  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Executor, ObjectRef, ExecutorNode);
+  TVM_DEFINE_OBJECT_REF_METHODS(Executor, ObjectRef, ExecutorNode);
+  TVM_DEFINE_OBJECT_REF_COW_METHOD(ExecutorNode)
 
  private:
   /*!
diff --git a/include/tvm/relay/expr_functor.h b/include/tvm/relay/expr_functor.h
index f96faffb24f4..fab579d0edd9 100644
--- a/include/tvm/relay/expr_functor.h
+++ b/include/tvm/relay/expr_functor.h
@@ -476,8 +476,19 @@ void ExpandDataflow(Expr expr, FCheckVisited fcheck_visited, FVisitLeaf fvisit_l
   auto fexpand_expr = [](const Expr& expr) {
     std::vector<Expr> result;
     if (const CallNode* op = expr.as<CallNode>()) {
-      for (auto it = op->args.rbegin(); it != op->args.rend(); ++it) {
-        result.push_back(*it);
+      if (op->op == Op::Get("call_lowered")) {
+        // Ignore the intermediate tuple since this is purely a calling-convention detail
+        const auto* tuple_args = op->args[1].as<TupleNode>();
+        ICHECK(tuple_args)
+            << "Expected second arg to call_lowered to be a Tuple of input arguments.";
+        for (auto it = tuple_args->fields.rbegin(); it != tuple_args->fields.rend(); ++it) {
+          result.push_back(*it);
+        }
+        result.push_back(op->args[0]);
+      } else {
+        for (auto it = op->args.rbegin(); it != op->args.rend(); ++it) {
+          result.push_back(*it);
+        }
       }
       result.push_back(op->op);
     } else if (const TupleNode* op = expr.as<TupleNode>()) {
diff --git a/include/tvm/relay/interpreter.h b/include/tvm/relay/interpreter.h
index eed6d0ffc1e4..f71107258d9a 100644
--- a/include/tvm/relay/interpreter.h
+++ b/include/tvm/relay/interpreter.h
@@ -184,10 +184,12 @@ TypedPackedFunc<ObjectRef(Array<Expr>)> EvalFunction(IRModule mod, Expr expr, De
  * \param import_set Already imported external modules.
  * \param device The device on which all primitives will be executed.
  * \param target The compiler target flag for compiling primitives.
+ * \param attrs Attributes for the expression to be evaluated with
  * @return The object representing the result.
  */
 ObjectRef Eval(Expr expr, Map<GlobalTypeVar, TypeData> type_definitions,
-               std::unordered_set<String> import_set, Device device, Target target);
+               std::unordered_set<String> import_set, Device device, Target target,
+               Map<String, ObjectRef> attrs = {});
 
 }  // namespace relay
 }  // namespace tvm
diff --git a/include/tvm/relay/runtime.h b/include/tvm/relay/runtime.h
index c4cabf5a5548..10e124bc339b 100644
--- a/include/tvm/relay/runtime.h
+++ b/include/tvm/relay/runtime.h
@@ -44,6 +44,12 @@ class AttrRegistry;
 
 namespace relay {
 
+/*! \brief Value used with Runtime::name to indicate the C++ runtime. */
+static constexpr const char* kTvmRuntimeCpp = "cpp";
+
+/*! \brief Value used with Runtime::name to indicate the C runtime. */
+static constexpr const char* kTvmRuntimeCrt = "crt";
+
 /*!
  * \brief Runtime information.
  *
@@ -105,6 +111,8 @@ class RuntimeNode : public Object {
   }
 
   static constexpr const char* _type_key = "Runtime";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(RuntimeNode, Object);
 };
 
diff --git a/include/tvm/relay/transform.h b/include/tvm/relay/transform.h
index 4bbc2df0ae48..ea3a5dba6bf7 100644
--- a/include/tvm/relay/transform.h
+++ b/include/tvm/relay/transform.h
@@ -499,6 +499,10 @@ TVM_DLL Pass PlanDevices(CompilationConfig config);
 /*!
  * \brief Bind the free variables to a Relay expression. This is a helper
  * function usually called by other pass functions to help optimizations.
+ * If any free variables are introduced into a function, those are added
+ * to the functoin parameters.
+ * Additionally this may change the order of parameters if you map a variable
+ * to a variable.
  *
  * \param expr The input expression.
  * \param binds The variable to expression map that will be used to help the
@@ -508,6 +512,19 @@ TVM_DLL Pass PlanDevices(CompilationConfig config);
  */
 TVM_DLL Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& binds);
 
+/*!
+ * \brief Substitute variables with new variables (including function parameters) in a function.
+ * This is a helper function usually called by other pass functions to help optimizations.
+ * Expects all values in the bind map to be Vars.
+ *
+ * \param func The input function.
+ * \param binds The variable to expression map that will be used to help the
+ *        binding.
+ *
+ * \return The updated expression.
+ */
+TVM_DLL Function SubstituteBoundVars(const Function& func, const tvm::Map<Var, Expr>& binds);
+
 /*!
  * \brief Apply rewrite rules to rewrite the expr in post DFS order. This
  * function is used as a helper function to rewrtie an expression in a pass.
diff --git a/include/tvm/runtime/container/optional.h b/include/tvm/runtime/container/optional.h
index bea4228c48b8..9961d5eeba0b 100644
--- a/include/tvm/runtime/container/optional.h
+++ b/include/tvm/runtime/container/optional.h
@@ -93,6 +93,11 @@ class Optional : public ObjectRef {
     ICHECK(data_ != nullptr);
     return T(data_);
   }
+  /*!
+   * \return The internal object pointer with container type of T.
+   * \note This function do not perform not-null checking.
+   */
+  const ContainerType* get() const { return static_cast<ContainerType*>(data_.get()); }
   /*!
    * \return The contained value if the Optional is not null
    *         otherwise return the default_value.
diff --git a/include/tvm/runtime/metadata.h b/include/tvm/runtime/metadata.h
new file mode 100644
index 000000000000..4f02efb8cee4
--- /dev/null
+++ b/include/tvm/runtime/metadata.h
@@ -0,0 +1,171 @@
+/*
+ * 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 tvm/runtime/metadata.h
+ * \brief Defines types which can be used in Metadata.
+ */
+#ifndef TVM_RUNTIME_METADATA_H_
+#define TVM_RUNTIME_METADATA_H_
+
+#include <inttypes.h>
+#ifdef __cplusplus
+#include <memory>
+#include <string>
+#include <vector>
+#endif
+#include <tvm/runtime/c_runtime_api.h>
+#ifdef __cplusplus
+#include <tvm/runtime/metadata_base.h>
+#include <tvm/support/span.h>
+#endif
+
+// Version number recorded in emitted artifacts for runtime checking.
+#define TVM_METADATA_VERSION 1
+
+#ifdef __cplusplus
+namespace tvm {
+namespace runtime {
+namespace metadata {
+/*!
+ * \brief Version of metadata emitted and understood by this compiler/runtime.
+ * Should be populated into the `version` field of all TVMMetadata.
+ */
+static const constexpr int64_t kMetadataVersion = TVM_METADATA_VERSION;
+}  // namespace metadata
+}  // namespace runtime
+}  // namespace tvm
+
+extern "C" {
+#endif
+
+/*!
+ * \brief Top-level metadata structure. Holds all other metadata types.
+ */
+struct TVMMetadata {
+  /*! \brief Version identifier for this metadata. */
+  int64_t version;
+  /*! \brief Inputs to the AOT run_model function.
+   * The order of the elements is the same as in the arguments to run_model. That is to say,
+   * this array specifies the first `num_inputs` arguments to run_model.
+   */
+  const struct TVMTensorInfo* inputs;
+  /*! \brief Number of elements in `inputs` array. */
+  int64_t num_inputs;
+  /*! \brief Outputs of the AOT run_model function.
+   * The order of the elements is the same as in the arguments to run_model. That is to say,
+   * this array specifies the last `num_outputs` arguments to run_model.
+   */
+  const struct TVMTensorInfo* outputs;
+  /*! \brief Number of elements in `outputs` array. */
+  int64_t num_outputs;
+  /*! \brief Memory Pools needed by the AOT run_model function.
+   * The order of the elements is the same as in the arguments to run_model. That is to say,
+   * this array specifies the last `num_pools` arguments to run_model.
+   */
+  const struct TVMTensorInfo* pools;
+  /*! \brief Number of elements in `pools` array. */
+  int64_t num_pools;
+  /*! \brief Name of the model, as passed to tvm.relay.build. */
+  const char* mod_name;
+};
+
+/*!
+ * \brief Describes one tensor argument to `run_model`.
+ * NOTE: while TIR allows for other types of arguments, such as scalars, the AOT run_model
+ * function does not currently accept these. Therefore it's not possible to express those
+ * in this metadata. A future patch may modify this.
+ */
+struct TVMTensorInfo {
+  /*! \brief Name of the tensor, as specified in the Relay program. */
+  const char* name;
+  /*! \brief Shape of the tensor. */
+  const int64_t* shape;
+  /*! \brief Rank of this tensor. */
+  int64_t num_shape;
+  /*! \brief Data type of one element of this tensor. */
+  DLDataType dtype;
+};
+#ifdef __cplusplus
+}  // extern "C"
+#include <tvm/runtime/object.h>
+namespace tvm {
+namespace runtime {
+namespace metadata {
+
+class Metadata;
+class TensorInfo;
+
+class MetadataNode : public MetadataBaseNode {
+ public:
+  explicit MetadataNode(const struct ::TVMMetadata* data) : data_{data} {}
+  static constexpr const char* _type_key = "metadata.MetadataNode";
+  const char* get_c_struct_name() const override;
+  inline int64_t version() const { return int64_t(data_->version); }
+  inline int64_t num_inputs() const { return data_->num_inputs; }
+  ArrayAccessor<struct TVMTensorInfo, TensorInfo> inputs();
+  inline int64_t num_outputs() const { return data_->num_outputs; }
+  ArrayAccessor<struct TVMTensorInfo, TensorInfo> outputs();
+  inline int64_t num_pools() const { return data_->num_pools; }
+  ArrayAccessor<struct TVMTensorInfo, TensorInfo> pools();
+  inline ::tvm::runtime::String mod_name() const { return ::tvm::runtime::String(data_->mod_name); }
+  const struct ::TVMMetadata* data() const { return data_; }
+  TVM_DECLARE_FINAL_OBJECT_INFO(MetadataNode, MetadataBaseNode);
+
+ private:
+  const struct ::TVMMetadata* data_;
+};
+
+class Metadata : public MetadataBase {
+ public:
+  explicit Metadata(const struct ::TVMMetadata* data);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Metadata, MetadataBase, MetadataNode);
+};
+
+class TensorInfoNode : public MetadataBaseNode {
+ public:
+  explicit TensorInfoNode(const struct ::TVMTensorInfo* data) : data_{data} {}
+  static constexpr const char* _type_key = "metadata.TensorInfoNode";
+  const char* get_c_struct_name() const override;
+  inline ::tvm::runtime::String name() const { return ::tvm::runtime::String(data_->name); }
+  inline int64_t num_shape() const { return data_->num_shape; }
+  inline ::tvm::support::Span<const int64_t, int64_t> shape() const {
+    return ::tvm::support::Span<const int64_t, int64_t>(data_->shape,
+                                                        data_->shape + data_->num_shape);
+  }
+  inline ::tvm::runtime::DataType dtype() const { return ::tvm::runtime::DataType(data_->dtype); }
+  const struct ::TVMTensorInfo* data() const { return data_; }
+  TVM_DECLARE_FINAL_OBJECT_INFO(TensorInfoNode, MetadataBaseNode);
+
+ private:
+  const struct ::TVMTensorInfo* data_;
+};
+
+class TensorInfo : public MetadataBase {
+ public:
+  explicit TensorInfo(const struct ::TVMTensorInfo* data);
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(TensorInfo, MetadataBase, TensorInfoNode);
+};
+
+}  // namespace metadata
+}  // namespace runtime
+}  // namespace tvm
+#endif  // defined(__cplusplus)
+
+#endif  // TVM_RUNTIME_METADATA_H_
diff --git a/include/tvm/runtime/metadata_base.h b/include/tvm/runtime/metadata_base.h
new file mode 100644
index 000000000000..442c42c303b1
--- /dev/null
+++ b/include/tvm/runtime/metadata_base.h
@@ -0,0 +1,216 @@
+/*
+ * 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 tvm/runtime/metadata_base.h
+ * \brief Defines types which can be used in Metadata.
+ */
+#ifndef TVM_RUNTIME_METADATA_BASE_H_
+#define TVM_RUNTIME_METADATA_BASE_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/runtime/object.h>
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+namespace metadata {
+
+/*!
+ * \brief Common base class for all Metadata.
+ *
+ * This class is used in the visitor classes as a internal check to ensure that verify that all
+ * parts of the Metadata struct used in codegen are Metadata objects.
+ */
+class MetadataBaseNode : public ::tvm::runtime::Object {
+ public:
+  virtual const char* get_c_struct_name() const = 0;
+
+  static constexpr const char* _type_key = "metadata.MetadataBaseNode";
+  TVM_DECLARE_BASE_OBJECT_INFO(MetadataBaseNode, ::tvm::runtime::Object);
+};
+
+/*! \brief Reference class for the common MetadataBaseNode class. */
+class MetadataBase : public ::tvm::runtime::ObjectRef {
+ public:
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(MetadataBase, ::tvm::runtime::ObjectRef, MetadataBaseNode);
+};
+
+template <typename C, class Ref>
+class ArrayAccessor;
+
+/*! \brief An iterator implementation that lazily instantiates the C++ wrapping Metadata class. */
+template <typename C, class Ref>
+class ArrayIterator {
+ public:
+  ArrayIterator(size_t index, const ArrayAccessor<C, Ref>* parent)
+      : index_{index}, parent_{parent} {}
+
+  inline Ref operator*() { return (*parent_)[index_]; }
+
+  inline ArrayIterator<C, Ref>& operator++() {
+    if (index_ < parent_->size()) {
+      index_++;
+    }
+
+    return *this;
+  }
+
+  inline bool operator==(const ArrayIterator<C, Ref>& other) const {
+    return parent_ == other.parent_ && index_ == other.index_;
+  }
+
+  inline bool operator!=(const ArrayIterator<C, Ref>& other) const { return !operator==(other); }
+
+ private:
+  size_t index_;
+  const ArrayAccessor<C, Ref>* parent_;
+};
+
+/*! \brief A span-like class which permits access to Array fields with complex elements.
+ * These array fields should be accessed from C++ using the Metadata wrapper classes. This class
+ * lazily instantiates those wrappers as they are accessed.
+ */
+template <typename C, class Ref>
+class ArrayAccessor {
+ public:
+  using value_type = Ref;
+  using iterator = ArrayIterator<C, Ref>;
+  using const_iterator = iterator;
+
+  template <typename T = typename std::enable_if<std::is_base_of<ObjectRef, Ref>::value>::type>
+  ArrayAccessor(const C* data, size_t num_data) : data_{data}, num_data_{num_data} {}
+
+  inline size_t size() const { return num_data_; }
+
+  inline Ref operator[](size_t index) const {
+    if (index >= num_data_) {
+      throw std::runtime_error("Index out of range");
+    }
+
+    return Ref(&data_[index]);
+  }
+
+  inline ArrayIterator<C, Ref> begin() const { return ArrayIterator<C, Ref>{0, this}; }
+
+  inline ArrayIterator<C, Ref> end() const { return ArrayIterator<C, Ref>{num_data_, this}; }
+
+ private:
+  const C* data_;
+  size_t num_data_;
+};
+
+/*! \brief A specialization of ArrayAccessor for String.
+ * This class is needed because the String constructor signature is different from the typical
+ * Metadata subclass.
+ */
+template <>
+class ArrayAccessor<const char*, ::tvm::runtime::String> {
+ public:
+  using value_type = ::tvm::runtime::String;
+  using iterator = ArrayIterator<const char*, ::tvm::runtime::String>;
+  using const_iterator = iterator;
+
+  ArrayAccessor(const char** data, size_t num_data) : data_{data}, num_data_{num_data} {}
+
+  inline size_t size() const { return num_data_; }
+
+  inline ::tvm::runtime::String operator[](size_t index) const {
+    if (index >= num_data_) {
+      throw std::runtime_error("Index out of range");
+    }
+    return ::tvm::runtime::String(data_[index]);
+  }
+
+  inline ArrayIterator<const char*, ::tvm::runtime::String> begin() const {
+    return ArrayIterator<const char*, ::tvm::runtime::String>{0, this};
+  }
+
+  inline ArrayIterator<const char*, ::tvm::runtime::String> end() const {
+    return ArrayIterator<const char*, ::tvm::runtime::String>{num_data_, this};
+  }
+
+ private:
+  const char** data_;
+  size_t num_data_;
+};
+
+/*! \brief Enumerates the primitive types which can be part of a Metadata instance.
+ *
+ * These are separate from TIR DataType because TIR does not model structs.
+ */
+enum MetadataKind : uint8_t {
+  kUint64 = 0,
+  kInt64 = 1,
+  kBool = 2,
+  kString = 3,
+  kHandle = 4,
+  kMetadata = 5,
+};
+
+/*! \brief Container for arrays in the metadata.
+ *
+ * Type information is needed when emitting arrays. This container augments the data field with
+ * the necessary typing information.
+ */
+class MetadataArrayNode : public MetadataBaseNode {
+ public:
+  MetadataArrayNode(Array<ObjectRef> array, MetadataKind kind, const char* type_key)
+      : array(::std::move(array)), kind{kind}, type_key{type_key} {}
+
+  const char* get_c_struct_name() const final;
+
+  std::string get_element_c_struct_name() const {
+    CHECK(kind == MetadataKind::kMetadata)
+      << "cannot get struct name for MetadataArray with kind=" << kind;
+    constexpr int prefix_size = sizeof("metadata.") - 1;
+    constexpr int suffix_size = sizeof("Node") - 1;
+    std::string type_key_str(type_key);
+    return std::string{"TVM"} + type_key_str.substr(prefix_size, type_key_str.size() - prefix_size - suffix_size);
+  }
+
+  Array<ObjectRef> array;
+
+  /*! \brief Describes the storage class of the emitted struct member. */
+  MetadataKind kind;
+
+  /*! \brief When `kind` is Metadata, type_key of the MetadataBaseNode used with this array. */
+  const char* type_key;
+
+  static constexpr const char* _type_key = "metadata.MetadataArrayNode";
+  TVM_DECLARE_BASE_OBJECT_INFO(MetadataArrayNode, MetadataBaseNode);
+};
+
+/*! \brief Reference class for MetadataArray. */
+class MetadataArray : public MetadataBase {
+ public:
+  MetadataArray(Array<ObjectRef> array, MetadataKind kind, const char* struct_name);
+
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(MetadataArray, MetadataBase, MetadataArrayNode);
+};
+
+}  // namespace metadata
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_METADATA_BASE_H_
diff --git a/include/tvm/runtime/module.h b/include/tvm/runtime/module.h
index 7b5326a44921..2e2a79b1ca53 100644
--- a/include/tvm/runtime/module.h
+++ b/include/tvm/runtime/module.h
@@ -235,8 +235,6 @@ constexpr const char* tvm_module_main = "__tvm_main__";
 constexpr const char* tvm_param_prefix = "__tvm_param__";
 /*! \brief A PackedFunc that looks up linked parameters by storage_id. */
 constexpr const char* tvm_lookup_linked_param = "_lookup_linked_param";
-/*! \brief The main AOT executor function generated from TIR */
-constexpr const char* tvm_run_func_suffix = "run_model";
 /*! \brief Model entrypoint generated as an interface to the AOT function outside of TIR */
 constexpr const char* tvm_entrypoint_suffix = "run";
 }  // namespace symbol
diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h
index a3bb569045eb..232ccd11f3c3 100644
--- a/include/tvm/runtime/packed_func.h
+++ b/include/tvm/runtime/packed_func.h
@@ -56,6 +56,10 @@ class TVMArgValue;
 class TVMMovableArgValueWithContext_;
 class TVMRetValue;
 class TVMArgsSetter;
+template <typename FType>
+class TypedPackedFunc;
+template <typename TSignature>
+struct SignaturePrinter;
 
 /*!
  * \brief Object container class that backs PackedFunc.
@@ -178,6 +182,9 @@ class PackedFunc : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(PackedFunc, ObjectRef, PackedFuncObj);
 };
 
+/*! \brief Using static function to output TypedPackedFunc signature */
+using FSig = std::string();
+
 /*!
  * \brief Please refer to \ref TypedPackedFuncAnchor "TypedPackedFunc<R(Args..)>"
  */
@@ -751,12 +758,16 @@ class TVMMovableArgValueWithContext_ {
    * \param value The other return value.
    * \param type_code The code associated with the type of the value.
    * \param arg_index In a function call, this argument is at index arg_index (0-indexed).
-   * \param optional_name Name of the function being called. Can be nullptr if the function is not
+   * \param optional_name Name of the function being called. Can be nullptr if the function is not.
+   * \param f_sig Pointer to static function outputting signature of the function being called.
    * named.
    */
   TVMMovableArgValueWithContext_(TVMValue value, int type_code, int arg_index,
-                                 const std::string* optional_name)
-      : value_(value, type_code), arg_index_(arg_index), optional_name_(optional_name) {}
+                                 const std::string* optional_name, FSig* f_sig)
+      : value_(value, type_code),
+        arg_index_(arg_index),
+        optional_name_(optional_name),
+        f_sig_(f_sig) {}
 
   template <typename T>
   operator T() const {
@@ -764,7 +775,8 @@ class TVMMovableArgValueWithContext_ {
       return value_;  // implicit conversion happens here
     } catch (dmlc::Error& e) {
       LOG(FATAL) << "In function " << (optional_name_ == nullptr ? "<anonymous>" : *optional_name_)
-                 << ": error while converting argument " << arg_index_ << ": " << e.what();
+                 << (f_sig_ == nullptr ? "" : (*f_sig_)()) << ": error while converting argument "
+                 << arg_index_ << ": " << e.what();
       throw;  // never reached, LOG(FATAL) throws, but this silences a warning.
     }
   }
@@ -773,6 +785,7 @@ class TVMMovableArgValueWithContext_ {
   TVMMovableArgValue_ value_;
   int arg_index_;
   const std::string* optional_name_;
+  FSig* f_sig_;
 };
 
 /*!
@@ -1268,6 +1281,61 @@ inline void for_each(const F& f, Args&&... args) {  // NOLINT(*)
   for_each_dispatcher<sizeof...(Args) == 0, 0, F>::run(f, std::forward<Args>(args)...);
 }
 
+namespace parameter_pack {
+
+template <typename... EnumArgs>
+struct EnumeratedParamPack {
+  struct Invoke {
+    template <template <size_t i, typename TArgument> class Functor, typename... ExtraParams>
+    static void F(ExtraParams&&... extra_params) {
+      using TExpander = int[];
+      (void)TExpander{
+          0,
+          (Functor<EnumArgs::i, typename EnumArgs::T>::F(extra_params...), 0)...,
+      };
+    }
+  };
+};
+
+template <typename... Args>
+struct EnumerateImpl {
+ private:
+  template <size_t _i, typename _T>
+  struct Item {
+    static const constexpr size_t i = _i;
+    using T = _T;
+  };
+
+  template <typename...>
+  struct Zipper;
+
+  template <std::size_t... id>
+  struct Zipper<std::integer_sequence<std::size_t, id...>> {
+    using T = EnumeratedParamPack<Item<id, Args>...>;
+  };
+
+ public:
+  using T = typename Zipper<std::index_sequence_for<Args...>>::T;
+};
+
+template <typename... Args>
+using Enumerate = typename EnumerateImpl<Args...>::T;
+
+template <typename... Args>
+struct ParamPack {
+  template <template <size_t i, typename TArgument> class Functor, typename... ExtraParams>
+  static void InvokeWithoutArg(ExtraParams&&... extra_params) {
+    Enumerate<Args...>::Invoke::template F<Functor, ExtraParams...>(
+        std::forward<ExtraParams>(extra_params)...);
+  }
+};
+
+}  // namespace parameter_pack
+
+/*!
+ * \brief Template class to get function signature of a function or functor.
+ * \tparam T The function/functor type.
+ */
 template <typename T>
 struct func_signature_helper {
   using FType = void;
@@ -1276,28 +1344,36 @@ struct func_signature_helper {
 template <typename T, typename R, typename... Args>
 struct func_signature_helper<R (T::*)(Args...)> {
   using FType = R(Args...);
+  using ParamType = parameter_pack::ParamPack<Args...>;
+  using RetType = R;
   static_assert(!std::is_reference<R>::value, "TypedPackedFunc return reference");
 };
 
 template <typename T, typename R, typename... Args>
 struct func_signature_helper<R (T::*)(Args...) const> {
   using FType = R(Args...);
+  using ParamType = parameter_pack::ParamPack<Args...>;
+  using RetType = R;
   static_assert(!std::is_reference<R>::value, "TypedPackedFunc return reference");
 };
 
 /*!
- * \brief template class to get function signature of a function or functor.
+ * \brief Template class to get function signature of a function or functor.
  * \tparam T The function/functor type.
  */
 template <typename T>
 struct function_signature {
   using FType = typename func_signature_helper<decltype(&T::operator())>::FType;
+  using ParamType = typename func_signature_helper<decltype(&T::operator())>::ParamType;
+  using RetType = typename func_signature_helper<decltype(&T::operator())>::RetType;
 };
 
 // handle case of function.
 template <typename R, typename... Args>
 struct function_signature<R(Args...)> {
   using FType = R(Args...);
+  using ParamType = parameter_pack::ParamPack<Args...>;
+  using RetType = R;
   static_assert(!std::is_reference<R>::value, "TypedPackedFunc return reference");
 };
 
@@ -1305,8 +1381,129 @@ struct function_signature<R(Args...)> {
 template <typename R, typename... Args>
 struct function_signature<R (*)(Args...)> {
   using FType = R(Args...);
+  using ParamType = detail::parameter_pack::ParamPack<Args...>;
+  using RetType = R;
   static_assert(!std::is_reference<R>::value, "TypedPackedFunc return reference");
 };
+
+template <typename TSignature>
+struct SignaturePrinter;
+
+namespace type2str {
+
+template <typename T>
+struct TypeSimplifier;
+
+template <typename T>
+struct Type2Str {
+  template <typename = std::enable_if_t<std::is_base_of<ObjectRef, T>::value>>
+  static std::string v() {
+    return T::ContainerType::_type_key;
+  }
+};
+template <>
+struct Type2Str<int> {
+  static std::string v() { return "int"; }
+};
+template <>
+struct Type2Str<double> {
+  static std::string v() { return "double"; }
+};
+template <>
+struct Type2Str<int64_t> {
+  static std::string v() { return "int64_t"; }
+};
+template <>
+struct Type2Str<uint64_t> {
+  static std::string v() { return "uint64_t"; }
+};
+template <>
+struct Type2Str<bool> {
+  static std::string v() { return "bool"; }
+};
+template <>
+struct Type2Str<void> {
+  static std::string v() { return "void"; }
+};
+template <>
+struct Type2Str<std::basic_string<char>> {
+  static std::string v() { return "basic_string<char>"; }
+};
+template <typename K, typename V>
+struct Type2Str<Map<K, V>> {
+  static std::string v() {
+    return "Map<" + TypeSimplifier<K>::v() + ", " + TypeSimplifier<V>::v() + ">";
+  }
+};
+template <>
+struct Type2Str<DLDevice> {
+  static std::string v() { return "DLDevice"; }
+};
+template <>
+struct Type2Str<DLTensor> {
+  static std::string v() { return "DLTensor"; }
+};
+template <>
+struct Type2Str<DataType> {
+  static std::string v() { return "DataType"; }
+};
+template <>
+struct Type2Str<TVMRetValue> {
+  static std::string v() { return "TVMRetValue"; }
+};
+template <>
+struct Type2Str<TVMArgValue> {
+  static std::string v() { return "TVMArgValue"; }
+};
+template <typename FType>
+struct Type2Str<TypedPackedFunc<FType>> {
+  static std::string v() { return SignaturePrinter<function_signature<FType>>::F(); }
+};
+template <typename T>
+struct Type2Str<Array<T>> {
+  static std::string v() { return "Array<" + TypeSimplifier<T>::v() + ">"; }
+};
+
+/*!
+ * \brief Template class to remove const, pointer and reference of original type.
+ * \tparam T The original type.
+ */
+template <typename T>
+struct TypeSimplifier {
+  static std::string v() {
+    using U = typename std::remove_cv<
+        typename std::remove_reference<typename std::remove_pointer<T>::type>::type>::type;
+    return (std::is_const<T>::value ? "const " : "") + Type2Str<U>::v() +
+           (std::is_pointer<T>::value ? "*" : "") + (std::is_reference<T>::value ? "&" : "");
+  }
+};
+
+}  // namespace type2str
+
+/*!
+ * \brief Template class to generate static function outputting signature of a function or functor.
+ * \tparam TSignature The function/functor signature type generated by `function_signature`.
+ */
+template <typename TSignature>
+struct SignaturePrinter {
+  using ParamType = typename TSignature::ParamType;
+  using RetType = typename TSignature::RetType;
+
+  template <size_t i, typename TArgument>
+  struct PrintParamType {
+    static void F(std::ostream& os) {
+      os << (i == 0 ? "" : ", ") << i << ": " << type2str::TypeSimplifier<TArgument>::v();
+    }
+  };
+
+  static std::string F() {
+    std::ostringstream oss;
+    oss << "(";
+    ParamType::template InvokeWithoutArg<PrintParamType>(oss);
+    oss << ") -> " << type2str::TypeSimplifier<RetType>::v();
+    return oss.str();
+  }
+};
 }  // namespace detail
 
 /* \brief argument settter to PackedFunc */
@@ -1422,22 +1619,22 @@ namespace detail {
 template <typename R, int nleft, int index, typename F>
 struct unpack_call_dispatcher {
   template <typename... Args>
-  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, const F& f,
+  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, FSig* f_sig, const F& f,
                                     const TVMArgs& args_pack, TVMRetValue* rv,
                                     Args&&... unpacked_args) {
     // construct a movable argument value
     // which allows potential move of argument to the input of F.
     unpack_call_dispatcher<R, nleft - 1, index + 1, F>::run(
-        optional_name, f, args_pack, rv, std::forward<Args>(unpacked_args)...,
+        optional_name, f_sig, f, args_pack, rv, std::forward<Args>(unpacked_args)...,
         TVMMovableArgValueWithContext_(args_pack.values[index], args_pack.type_codes[index], index,
-                                       optional_name));
+                                       optional_name, f_sig));
   }
 };
 
 template <typename R, int index, typename F>
 struct unpack_call_dispatcher<R, 0, index, F> {
   template <typename... Args>
-  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, const F& f,
+  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, FSig* f_sig, const F& f,
                                     const TVMArgs& args_pack, TVMRetValue* rv,
                                     Args&&... unpacked_args) {
     using RetType = decltype(f(std::forward<Args>(unpacked_args)...));
@@ -1452,7 +1649,7 @@ struct unpack_call_dispatcher<R, 0, index, F> {
 template <int index, typename F>
 struct unpack_call_dispatcher<void, 0, index, F> {
   template <typename... Args>
-  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, const F& f,
+  TVM_ALWAYS_INLINE static void run(const std::string* optional_name, FSig* f_sig, const F& f,
                                     const TVMArgs& args_pack, TVMRetValue* rv,
                                     Args&&... unpacked_args) {
     f(std::forward<Args>(unpacked_args)...);
@@ -1462,11 +1659,12 @@ struct unpack_call_dispatcher<void, 0, index, F> {
 template <typename R, int nargs, typename F>
 TVM_ALWAYS_INLINE void unpack_call(const std::string* optional_name, const F& f,
                                    const TVMArgs& args, TVMRetValue* rv) {
+  FSig* f_sig = detail::SignaturePrinter<detail::function_signature<F>>::F;
   CHECK_EQ(nargs, args.size()) << "Function "
                                << (optional_name == nullptr ? "<anonymous>" : *optional_name)
-                               << " expects " << nargs << " arguments but " << args.size()
-                               << " were provided";
-  unpack_call_dispatcher<R, nargs, 0, F>::run(optional_name, f, args, rv);
+                               << (f_sig == nullptr ? "" : (*f_sig)()) << " expects " << nargs
+                               << " arguments but " << args.size() << " were provided";
+  unpack_call_dispatcher<R, nargs, 0, F>::run(optional_name, f_sig, f, args, rv);
 }
 
 template <typename FType>
@@ -1520,10 +1718,11 @@ TypedPackedFunc<R(Args...)>::TypedPackedFunc(TVMMovableArgValueWithContext_&& va
 template <typename R, typename... Args>
 template <typename FType>
 inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda, std::string name) {
-  packed_ = PackedFunc([flambda, name](const TVMArgs& args, TVMRetValue* rv) {
+  FSig* f_sig = detail::SignaturePrinter<detail::function_signature<FType>>::F;
+  packed_ = PackedFunc([flambda, name, f_sig](const TVMArgs& args, TVMRetValue* rv) {
     if (args.size() != sizeof...(Args)) {
-      LOG(FATAL) << "Function " << name << " expects " << sizeof...(Args) << " arguments, but "
-                 << args.size() << " were provided.";
+      LOG(FATAL) << "Function " << name << (f_sig == nullptr ? "" : (*f_sig)()) << " expects "
+                 << sizeof...(Args) << " arguments, but " << args.size() << " were provided.";
     }
     detail::unpack_call<R, sizeof...(Args)>(&name, flambda, args, rv);
   });
@@ -1532,10 +1731,11 @@ inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda, std::s
 template <typename R, typename... Args>
 template <typename FType>
 inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
-  packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {
+  FSig* f_sig = detail::SignaturePrinter<detail::function_signature<FType>>::F;
+  packed_ = PackedFunc([flambda, f_sig](const TVMArgs& args, TVMRetValue* rv) {
     if (args.size() != sizeof...(Args)) {
-      LOG(FATAL) << "Function <anonymous> expects " << sizeof...(Args) << " arguments, but "
-                 << args.size() << " were provided.";
+      LOG(FATAL) << "Function <anonymous> " << (*f_sig)() << " expects " << sizeof...(Args)
+                 << " arguments, but " << args.size() << " were provided.";
     }
     detail::unpack_call<R, sizeof...(Args)>(nullptr, flambda, args, rv);
   });
diff --git a/include/tvm/runtime/registry.h b/include/tvm/runtime/registry.h
index 316fbeb7e891..7d781461c569 100644
--- a/include/tvm/runtime/registry.h
+++ b/include/tvm/runtime/registry.h
@@ -167,11 +167,12 @@ class Registry {
    */
   template <typename T, typename R, typename... Args>
   Registry& set_body_method(R (T::*f)(Args...)) {
-    auto fwrap = [f](T target, Args... params) -> R {
+    using R_ = typename std::remove_reference<R>::type;
+    auto fwrap = [f](T target, Args... params) -> R_ {
       // call method pointer
       return (target.*f)(params...);
     };
-    return set_body(TypedPackedFunc<R(T, Args...)>(fwrap, name_));
+    return set_body(TypedPackedFunc<R_(T, Args...)>(fwrap, name_));
   }
 
   /*!
diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 6359da0a5375..774bca1e2d28 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -218,6 +218,13 @@ class Executable : public ModuleNode {
    */
   void SetLib(const runtime::Module& lib);
 
+  /*!
+   * \brief Get VMFunction.
+   * \param func_name The function's name.
+   * \return VMFunction.
+   */
+  const VMFunction& GetVMFunctionWithName(const std::string& func_name) const;
+
   /*!
    * \brief Get the arity of the VMFunction.
    * \param func Function name.
diff --git a/include/tvm/runtime/vm/vm.h b/include/tvm/runtime/vm/vm.h
index d7311951b702..139c8ba5fcc8 100644
--- a/include/tvm/runtime/vm/vm.h
+++ b/include/tvm/runtime/vm/vm.h
@@ -93,7 +93,9 @@ struct VMFunction {
         params(std::move(params)),
         instructions(std::move(instructions)),
         register_file_size(register_file_size),
-        param_device_indexes(std::move(param_device_indexes)) {}
+        param_device_indexes(std::move(param_device_indexes)) {
+    ICHECK_EQ(params.size(), param_device_indexes.size());
+  }
 
   VMFunction() = default;
 
@@ -270,6 +272,15 @@ class VirtualMachine : public runtime::ModuleNode {
    */
   void SetInput(std::string name, TVMArgs args, int offset);
 
+  /*!
+   * \brief Set one input tensor with index or name to a function.
+   * \param name The function name.
+   * \param tag index or name of the input tensor .
+   * \param tensor the input tensor. If the tensor is not of the correct device for the function,
+   * they will be copied to the device.
+   */
+  void SetOneInput(std::string name, const TVMArgValue& tag, const TVMArgValue& tensor);
+
   /*!
    * \brief Internal hook for profiling the start of an op.
    *
@@ -286,6 +297,48 @@ class VirtualMachine : public runtime::ModuleNode {
    */
   virtual void OpStopHook();
 
+ private:
+  /*!
+   * \brief Get index of input tensor from its name.
+   * \param func_name The function's name.
+   * \param input_name The input tensor name.
+   * \return The input tensor index.
+   */
+  int64_t GetInputIndexFromVMFunction(const std::string& func_name,
+                                      const std::string& input_name) const;
+
+  /*!
+   * \brief Get index of input tensor from its name.
+   * \param params parameter names.
+   * \param input_name The input tensor name.
+   * \return The input tensor index.
+   */
+  int64_t GetInputIndexFromName(const std::vector<std::string>& params,
+                                const std::string& input_name) const;
+
+  /*!
+   * \brief Check executable exists and get VM function from it.
+   * \param func_name The function's name.
+   * \return VM function.
+   */
+  const VMFunction& CheckAndGetVMFunction(const std::string& func_name) const;
+
+  /*!
+   * \brief Creats inputs_ field, if it exists check its size.
+   * \param func_name The function's name.
+   * \param size inputs_ field size.
+   * \return VM function.
+   */
+  void CreateInputsOrCheckSize(const std::string& func_name, size_t size);
+
+  /*!
+   * \brief Set one input tensor with given index to set of input tensors if need copy to given
+   * device. \param tensors the input tensors set (destination) \param tensor some tensor (not
+   * neccessary DLTensor). \param index The input tensor index. \param dev device to copy if need.
+   */
+  void SetInputTensorWithIndex(std::vector<ObjectRef>& tensors,  // NOLINT(*)
+                               const TVMArgValue& tensor, int index, Device dev);
+
  protected:
   /*! \brief The virtual machine's packed function table. */
   std::vector<PackedFunc> packed_funcs_;
diff --git a/include/tvm/support/span.h b/include/tvm/support/span.h
new file mode 100644
index 000000000000..faa849c4a6f0
--- /dev/null
+++ b/include/tvm/support/span.h
@@ -0,0 +1,103 @@
+/*
+ * 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 tvm/support/span.h
+ * \brief Reimplementation of part of C++-20 style span.
+ */
+#ifndef TVM_SUPPORT_SPAN_H_
+#define TVM_SUPPORT_SPAN_H_
+
+#include <cstddef>
+#include <iterator>
+#include <type_traits>
+#include <vector>
+
+namespace tvm {
+namespace support {
+
+/*!
+ * \brief A partial implementation of the C++20 std::span.
+ *
+ * At the time of writing, TVM must compile against C++14.
+ */
+template <class T, class W>
+class Span {
+ public:
+  using value_type = W;
+  using const_W = typename ::std::add_const<W>::type;
+
+  template <class W1>
+  class iterator_base : public std::iterator<std::input_iterator_tag, W> {
+   public:
+    inline iterator_base(T* ptr, T* end) : ptr_{ptr}, end_{end} { CHECK_GE(end, ptr); }
+
+    inline W1 operator*() { return W1(*ptr_); }
+
+    inline iterator_base<W1>& operator++() {
+      if (ptr_ != end_) ptr_++;
+      return *this;
+    }
+
+    inline bool operator==(iterator_base<W1> other) {
+      return ptr_ == other.ptr_ && end_ == other.end_;
+    }
+
+    inline bool operator!=(iterator_base<W1> other) { return !(*this == other); }
+
+    template <class X = W1, typename = ::std::enable_if_t<!::std::is_const<X>::value> >
+    inline operator iterator_base<const_W>() const {
+      return iterator_base<const_W>(ptr_, end_);
+    }
+
+   private:
+    T* ptr_;
+    T* end_;
+  };
+
+  using iterator = iterator_base<W>;
+  using const_iterator = iterator_base<const_W>;
+
+  inline Span(T* begin, int num_elements) : begin_{begin}, end_{begin + num_elements} {}
+  inline Span(T* begin, T* end) : begin_{begin}, end_{end} {}
+
+  inline iterator begin() const { return iterator(begin_, end_); }
+
+  inline iterator end() const { return iterator(end_, end_); }
+
+  size_t size() const { return end_ - begin_; }
+
+  inline W operator[](int i) {
+    T* to_return = begin_ + i;
+    ICHECK_LT(to_return, end_) << "Span access out of bounds: " << i;
+    return W(*to_return);
+  }
+
+  inline operator std::vector<W>() { return std::vector<W>(begin(), end()); }
+
+ protected:
+  T* begin_;
+  T* end_;
+};
+
+}  // namespace support
+}  // namespace tvm
+
+#endif  // TVM_SUPPORT_SPAN_H_
diff --git a/include/tvm/target/virtual_device.h b/include/tvm/target/virtual_device.h
index 07011ea412a3..37f4b23b12c2 100644
--- a/include/tvm/target/virtual_device.h
+++ b/include/tvm/target/virtual_device.h
@@ -162,6 +162,7 @@ using MemoryScope = String;
  *
  * These operations are needed during device planning.
  */
+
 class VirtualDeviceNode : public AttrsNode<VirtualDeviceNode> {
  private:
   /*!
@@ -361,6 +362,13 @@ class VirtualDeviceCache {
   std::unordered_set<VirtualDevice, StructuralHash, StructuralEqual> cache_;
 };
 
+/*! brief The attribute key for the virtual device. This key will be promoted to first class on
+ * functions. For use in the parser and printer only.
+ *
+ * Type: VirtualDevice
+ */
+constexpr const char* kVirtualDevice = "virtual_device";
+
 }  // namespace tvm
 
 #endif  //  TVM_TARGET_VIRTUAL_DEVICE_H_
diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h
index 13f39317dbe4..89074d83e1d6 100644
--- a/include/tvm/te/operation.h
+++ b/include/tvm/te/operation.h
@@ -59,8 +59,11 @@ class TVM_DLL OperationNode : public Object {
   std::string name;
   /*! \brief optional tag of the operation */
   std::string tag;
-  /*! \brief additional attributes of the operation*/
+  /*! \brief additional attributes of the operation */
   Map<String, ObjectRef> attrs;
+  /*! \brief output tensors */
+  Array<Tensor> outputs;
+
   // virtual destructor.
   virtual ~OperationNode() {}
   /*! \return number of outputs */
@@ -473,7 +476,7 @@ class HybridOpNode : public OperationNode {
   /*! \brief The input tensors */
   Array<Tensor> inputs;
   /*! \brief Symbolic placeholder representation of outputs */
-  Array<Tensor> outputs;
+  Array<Tensor> symbolic_outputs;
   /*! \brief The axis of iterations */
   Array<IterVar> axis;
   /*! \brief the statement that generates the computation. This is
@@ -509,6 +512,7 @@ class HybridOpNode : public OperationNode {
     v->Visit("attrs", &attrs);
     v->Visit("inputs", &inputs);
     v->Visit("outputs", &outputs);
+    v->Visit("symbolic_outputs", &symbolic_outputs);
     v->Visit("axis", &axis);
     v->Visit("body", &body);
   }
diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h
index e6665f965b5b..d8a5ea67d844 100644
--- a/include/tvm/tir/builtin.h
+++ b/include/tvm/tir/builtin.h
@@ -626,6 +626,13 @@ TVM_DLL const Op& texture2d_store();
  */
 TVM_DLL const Op& texture2d_load();
 
+/*!
+ * \brief Copy 1d memory from source to destination
+ * Same semantics as memcpy(destination, source, size)
+ * Allows for device specific implementations e.g. direct memory access (DMA)
+ */
+TVM_DLL const Op& mem_copy();
+
 /*! \brief The kind of structure field info used in intrinsic */
 enum TVMStructFieldKind : int {
   // array head address
diff --git a/include/tvm/tir/data_layout.h b/include/tvm/tir/data_layout.h
index 73da05c549e2..81c3e98e663d 100644
--- a/include/tvm/tir/data_layout.h
+++ b/include/tvm/tir/data_layout.h
@@ -295,9 +295,13 @@ class BijectiveLayoutNode : public Object {
   /*! \brief Describes how source axes can be mapped to the destination axes,
    *   e.g., [i0 / 16, i1, i0 % 16] can describe NC -> NC16n
    */
-  Array<PrimExpr> forward_rule;
+  Array<PrimExpr> index_forward_rule;
   /*! \brief Describes how destination axes can be mapped to the source axes */
-  Array<PrimExpr> backward_rule;
+  Array<PrimExpr> index_backward_rule;
+  /*! \brief Describes how source shapes can be mapped to the destination shapes */
+  Array<PrimExpr> shape_forward_rule;
+  /*! \brief Describes how destination shapes can be mapped to the source shapes */
+  Array<PrimExpr> shape_backward_rule;
 
   /*! \brief The source layout */
   Layout src_layout;
@@ -307,8 +311,10 @@ class BijectiveLayoutNode : public Object {
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("src_layout", &src_layout);
     v->Visit("dst_layout", &dst_layout);
-    v->Visit("forward_rule", &forward_rule);
-    v->Visit("backward_rule", &backward_rule);
+    v->Visit("index_forward_rule", &index_forward_rule);
+    v->Visit("index_backward_rule", &index_backward_rule);
+    v->Visit("shape_forward_rule", &shape_forward_rule);
+    v->Visit("shape_backward_rule", &shape_backward_rule);
   }
 
   static constexpr const char* _type_key = "tir.BijectiveLayout";
diff --git a/include/tvm/tir/function.h b/include/tvm/tir/function.h
index 1ab911b756df..2b3c4b5fe003 100644
--- a/include/tvm/tir/function.h
+++ b/include/tvm/tir/function.h
@@ -151,42 +151,6 @@ class PrimFunc : public BaseFunc {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(PrimFuncNode);
 };
 
-/*!
- * \brief Describes one parameter that should be linked into the generated module.
- *
- * When parameters are to be linked in with generated code (i.e. on target_host-compatible
- * backends), Relay attaches instances of this object to a global TIR function. Code-generators
- * use the information contained in this node to include the parameter data in the generated
- * module.
- */
-class LinkedParamNode : public Object {
- public:
-  /*! \brief Unique numeric identifier used by runtimes to lookup this parameter. */
-  int64_t id;
-
-  /*! \brief Parameter data which should get linked into the final module. */
-  ::tvm::runtime::NDArray param;
-
-  void VisitAttrs(tvm::AttrVisitor* v) {
-    v->Visit("id", &id);
-    v->Visit("param", &param);
-  }
-
-  static constexpr const char* _type_key = "tir.LinkedParam";
-  TVM_DECLARE_FINAL_OBJECT_INFO(LinkedParamNode, Object);
-};
-
-/*!
- * \brief Managed reference to LinkedParamNode.
- */
-class LinkedParam : public ObjectRef {
- public:
-  TVM_DLL LinkedParam(int64_t id, ::tvm::runtime::NDArray param);
-
-  TVM_DEFINE_OBJECT_REF_METHODS(LinkedParam, ObjectRef, LinkedParamNode);
-  TVM_DEFINE_OBJECT_REF_COW_METHOD(LinkedParamNode);
-};
-
 /*!
  * \brief Tensor intrinsics for tensorization
  */
@@ -239,7 +203,7 @@ class TensorIntrin : public ObjectRef {
   TVM_DEFINE_OBJECT_REF_METHODS(TensorIntrin, ObjectRef, TensorIntrinNode)
 };
 
-/*!
+/*
  * \brief Specialize parameters of PrimFunc.
  * \param func The PrimFunc to be specialized.
  * \param param_map The mapping from function params to the instance.
diff --git a/include/tvm/tir/op.h b/include/tvm/tir/op.h
index 9cf7d0a3cd1f..9c3ea135c68d 100644
--- a/include/tvm/tir/op.h
+++ b/include/tvm/tir/op.h
@@ -495,6 +495,22 @@ TVM_DLL PrimExpr truncmod(PrimExpr a, PrimExpr b, Span span = Span());
  *       index types(int32, int64) when possible.
  */
 TVM_DLL PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span = Span());
+/*!
+ * \brief compute ceil(a / b) where a and b are non-negative.
+ *
+ * Use this function for shape split calculation.
+ *
+ * This function might take advantage of the fact
+ * that a and b are non-negative.
+ *
+ * \param a left operand
+ * \param b right operand
+ * \param span The location of this operation in the source.
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       shape types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span());
 /*!
  * \brief compute the remainder floor(a / b) where a and b are non-negative.
  *
@@ -521,6 +537,17 @@ TVM_DLL PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span = Span());
  *       index types(int32, int64) when possible.
  */
 TVM_DLL PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span = Span());
+/*!
+ * \brief compute ceil(a / b)
+ *
+ * \param a left operand
+ * \param b right operand
+ * \param span The location of this operation in the source.
+ * \return The result expression.
+ * \note this function does eager constant folding for
+ *       index types(int32, int64) when possible.
+ */
+TVM_DLL PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span = Span());
 /*!
  * \brief compute the remainder of floordiv
  *
@@ -835,10 +862,18 @@ TVM_DLL PrimExpr q_multiply_shift(PrimExpr x, PrimExpr y, PrimExpr q, PrimExpr s
                                   Span span = Span());
 
 // Intrinsic operators
-#define TVM_DECLARE_INTRIN_UNARY(OpName)                   \
-  inline PrimExpr OpName(PrimExpr x, Span span = Span()) { \
-    static const Op& op = Op::Get("tir." #OpName);         \
-    return tir::Call(x.dtype(), op, {x}, span);            \
+#define TVM_DECLARE_INTRIN_UNARY(OpName)                       \
+  inline PrimExpr OpName(PrimExpr x, Span span = Span()) {     \
+    static const Op& op = Op::Get("tir." #OpName);             \
+    if (x.dtype().is_bfloat16()) {                             \
+      DataType srcType = x.dtype();                            \
+      DataType dstType(kDLFloat, 32, srcType.lanes());         \
+      PrimExpr castX = tir::Cast(dstType, {x}, span);          \
+      PrimExpr result = tir::Call(dstType, op, {castX}, span); \
+      return tir::Cast(srcType, {result}, span);               \
+    } else {                                                   \
+      return tir::Call(x.dtype(), op, {x}, span);              \
+    }                                                          \
   }
 
 TVM_DECLARE_INTRIN_UNARY(exp);
diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h
index edb789b0bd7f..972f78171569 100644
--- a/include/tvm/tir/stmt.h
+++ b/include/tvm/tir/stmt.h
@@ -559,16 +559,18 @@ class AllocateNode : public StmtNode {
    *        Otherwise return 0.
    * \return The result.
    */
-  int32_t constant_allocation_size() const { return constant_allocation_size(extents); }
+  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
   /*!
    * \brief If the buffer size is constant, return the size.
    *        Otherwise return 0.
    * \param extents The extents of the buffer.
    * \return The result.
    */
-  TVM_DLL static int32_t constant_allocation_size(const Array<PrimExpr>& extents);
+  TVM_DLL static int64_t ConstantAllocationSize(const Array<PrimExpr>& extents);
 
   static constexpr const char* _type_key = "tir.Allocate";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
   TVM_DECLARE_FINAL_OBJECT_INFO(AllocateNode, StmtNode);
 };
 
@@ -585,6 +587,96 @@ class Allocate : public Stmt {
   TVM_DEFINE_OBJECT_REF_METHODS(Allocate, Stmt, AllocateNode);
 };
 
+/*!
+ * \brief Allocate a buffer that can be used in body.
+ */
+class AllocateConstNode : public StmtNode {
+ public:
+  /*! \brief The buffer variable. */
+  Var buffer_var;
+  /*! \brief The optional data associated to the constant.
+   */
+  Optional<runtime::NDArray> data;
+  /*! \brief If the PrimFunc containing the Stmt is added to IRModule,
+       this is an optional index to indicate the index within
+       "Constants" attribute, that is a Array<NDArray> of IRModule.
+   */
+  Optional<Integer> irmod_storage_idx;
+  /*! \brief The type of the buffer. */
+  DataType dtype;
+  /*! \brief The extents of the buffer. */
+  Array<PrimExpr> extents;
+  /*! \brief The body to be executed. */
+  Stmt body;
+  /*!
+   * \brief Additional annotations about the allocation.
+   *
+   *  These annotations can be used as auxiliary hint
+   *  to future transformations.
+   */
+  Map<String, ObjectRef> annotations;
+
+  void VisitAttrs(AttrVisitor* v) {
+    v->Visit("buffer_var", &buffer_var);
+    v->Visit("data", &data);
+    v->Visit("irmod_storage_idx", &irmod_storage_idx);
+    v->Visit("dtype", &dtype);
+    v->Visit("extents", &extents);
+    v->Visit("body", &body);
+    v->Visit("annotations", &annotations);
+    v->Visit("span", &span);
+  }
+
+  bool SEqualReduce(const AllocateConstNode* other, SEqualReducer equal) const {
+    return equal.DefEqual(buffer_var, other->buffer_var) && equal(dtype, other->dtype) &&
+           equal(extents, other->extents) && equal(data, other->data) && equal(body, other->body) &&
+           equal(annotations, other->annotations);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce.DefHash(buffer_var);
+    hash_reduce(dtype);
+    hash_reduce(extents);
+    hash_reduce(body);
+    hash_reduce(annotations);
+    hash_reduce(data);
+  }
+
+  /*!
+   * \brief If the buffer size is constant, return the size.
+   *        Otherwise return 0.
+   * \return The result.
+   */
+  int64_t ConstantAllocationSize() const { return ConstantAllocationSize(extents); }
+  /*!
+   * \brief If the buffer size is constant, return the size.
+   *        Otherwise return 0.
+   * \param extents The extents of the buffer.
+   * \return The result.
+   */
+  TVM_DLL static int64_t ConstantAllocationSize(const Array<PrimExpr>& extents);
+
+  static constexpr const char* _type_key = "tir.AllocateConst";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
+  TVM_DECLARE_FINAL_OBJECT_INFO(AllocateConstNode, StmtNode);
+};
+
+/*!
+ * \brief Managed reference to AllocateConstNode.
+ * \sa AllocateConstNode
+ */
+class AllocateConst : public Stmt {
+ public:
+  /* The constructor to create a IRNode with constant data
+   * depending on the type of ObjectRef, it will either
+   * create AllocateConstNode with irmod_storage_idx or data
+   */
+  TVM_DLL AllocateConst(Var buffer_var, DataType dtype, Array<PrimExpr> extents,
+                        ObjectRef data_or_idx, Stmt body, Span span = Span());
+  TVM_DEFINE_OBJECT_REF_METHODS(AllocateConst, Stmt, AllocateConstNode);
+};
+
 /*!
  * \brief The container of seq statement.
  *        Represent a sequence of statements.
@@ -1224,7 +1316,7 @@ class BlockRealize : public Stmt {
   TVM_DEFINE_OBJECT_REF_COW_METHOD(BlockRealizeNode);
 };
 
-/*! \brief namespace of possible attribute sin AttrStmt.attr_key */
+/*! \brief namespace of possible attributes in AttrStmt.attr_key */
 namespace attr {
 // The above attr does not pass to ir stage.
 /*! \brief Mark launching extent of thread, used by device API. */
@@ -1361,6 +1453,12 @@ constexpr const char* script_parsing_detect_access = "tir.script_parsing_detect_
  */
 constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
 
+/*! \brief Mark the stage of a statement in the software pipeline */
+constexpr const char* software_pipeline_stage = "software_pipeline_stage";
+
+/*! \brief Mark the order of a statement in the software pipeline */
+constexpr const char* software_pipeline_order = "software_pipeline_order";
+
 /*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */
 constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure";
 
diff --git a/include/tvm/tir/stmt_functor.h b/include/tvm/tir/stmt_functor.h
index 0b4ace20078c..16da91c2a2a3 100644
--- a/include/tvm/tir/stmt_functor.h
+++ b/include/tvm/tir/stmt_functor.h
@@ -87,6 +87,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
   virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const AllocateNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
+  virtual R VisitStmt_(const AllocateConstNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const StoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
   virtual R VisitStmt_(const BufferRealizeNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
@@ -113,6 +114,7 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
     IR_STMT_FUNCTOR_DISPATCH(ForNode);
     IR_STMT_FUNCTOR_DISPATCH(WhileNode);
     IR_STMT_FUNCTOR_DISPATCH(AllocateNode);
+    IR_STMT_FUNCTOR_DISPATCH(AllocateConstNode);
     IR_STMT_FUNCTOR_DISPATCH(StoreNode);
     IR_STMT_FUNCTOR_DISPATCH(AssertStmtNode);
     IR_STMT_FUNCTOR_DISPATCH(ProducerStoreNode);
@@ -155,6 +157,7 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
   void VisitStmt_(const ForNode* op) override;
   void VisitStmt_(const WhileNode* op) override;
   void VisitStmt_(const AllocateNode* op) override;
+  void VisitStmt_(const AllocateConstNode* op) override;
   void VisitStmt_(const StoreNode* op) override;
   void VisitStmt_(const BufferStoreNode* op) override;
   void VisitStmt_(const BufferRealizeNode* op) override;
@@ -255,6 +258,7 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
   Stmt VisitStmt_(const ForNode* op) override;
   Stmt VisitStmt_(const WhileNode* op) override;
   Stmt VisitStmt_(const AllocateNode* op) override;
+  Stmt VisitStmt_(const AllocateConstNode* op) override;
   Stmt VisitStmt_(const StoreNode* op) override;
   Stmt VisitStmt_(const BufferStoreNode* op) override;
   Stmt VisitStmt_(const BufferRealizeNode* op) override;
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
index 6d17c396c12f..3bb5491affdf 100644
--- a/include/tvm/tir/transform.h
+++ b/include/tvm/tir/transform.h
@@ -29,6 +29,7 @@
 #include <tvm/tir/function.h>
 
 #include <string>
+#include <vector>
 
 namespace tvm {
 namespace tir {
@@ -500,6 +501,116 @@ TVM_DLL Pass ConvertForLoopsToSerial();
  */
 TVM_DLL Pass UnifiedStaticMemoryPlanner();
 
+/*!
+ * \brief This pass transforms annotated loops into pipelined ones where producers and consumers
+ * are overlapped with the information provided in loop annotations, which enables optimization
+ * techniques like prefetching and pipeline parallelism.
+ *
+ * The pipeline scope consists of the direct children of the annotated loop (ignoring BlockRealize,
+ * Block, SeqStmt), and the number of children is denoted by `n` in the documentation.
+ *
+ * The following annotations are used to guide the loop transformation:
+ *
+ * 1) Loop annotation `software_pipeline_stage` defines the pipeline stage.
+ * An array of `n` integers, and each element should be in range [0, max_stage],
+ * where max_stage is the maximum (inclusive) stage.
+ * 2) Loop annotation `software_pipeline_order` defines the pipeline order.
+ * An array of `n` integers, a permutation of [0, 1, ..., num_components - 1];
+ * 3) Block annotation `double_buffer_scope` controls certain buffer sizes to allow decoupling of
+ * read/write dependency. It's an integer index of the write regions of the block.
+ *
+ * Every annotated loop is transformed into a loop with three blocks as its direct children:
+ *
+ * 1) Prologue block, where components whose stage is less than `max_stage` is executed;
+ *
+ * 2) Body block, where all the components are executed;
+ *
+ * 3) Epilogue block, where only components whose stage is greater than 0 will be executed.
+ * The execution order is controlled by the annotation `software_pipeline_order`,
+ * and thus could be different than the original order.
+ *
+ * Note: For nested software pipelines, the inner software pipeline will be generated first,
+ * which may affect the number of the direct children of the outer loop.
+ * In this case, the annotations for the outer software
+ * pipeline should include the result of the inner software pipeline,
+ * which is the three blocks as discussed above.
+ * Example:
+ *
+ * Before this pass, the TIR is:
+ *
+ * \code{.py}
+ * @T.prim_func
+ * def before_transform(A: T.Buffer[(16, 16), "float32"], C: T.Buffer[(16, 16), "float32"]) -> None:
+ *     for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+ *         for i in T.serial(0, 16,
+ *                           annotations={"software_pipeline_stage": [0, 1],
+ *                                        "software_pipeline_order": [0, 1]}
+ *                          ):
+ *             with T.block():
+ *                 T.reads(A[tx, i])
+ *                 T.writes(C[tx, i])
+ *                 B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+ *                 with T.block("B"):
+ *                     T.reads(A[tx, i])
+ *                     T.writes(B[tx, 0])
+ *                     B[tx, 0] = A[tx, i] * T.float32(2)
+ *                 with T.block("C"):
+ *                     T.reads(B[tx, 0])
+ *                     T.writes(C[tx, i])
+ *                     C[tx, i] = B[tx, 0] + T.float32(1)
+ * \endcode
+ *
+ * The TIR above annotates the loop as a two-stage pipeline with no reordering.
+ * After applying this pass, the TIR is transformed into:
+ *
+ * \code{.py}
+ * @T.prim_func
+ * def after_transform(A: T.Buffer[(16, 16), "float32"], C: T.Buffer[(16, 16), "float32"]) -> None:
+ *     for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+ *         with T.block():
+ *             T.reads([A[tx, 0:16]])
+ *             T.writes([C[tx, 0:16]])
+ *             B = T.alloc_buffer([2, 16, 1], dtype="float32", scope="shared")
+ *             with T.block("prologue"):
+ *                 T.reads([A[tx, 0]])
+ *                 T.writes([B[0, tx, 0]])
+ *                 B[0, tx, 0] = A[tx, 0] * T.float32(2)
+ *             with T.block("body"):
+ *                 T.reads([A[tx, 1:16], B[0:2, tx, 0]])
+ *                 T.writes([B[0:2, tx, 0], C[tx, 0:15]])
+ *                 for i in T.serial(0, 15):
+ *                     with T.block("B"):
+ *                         T.reads([A[tx, i + 1]])
+ *                         T.writes([B[(i + 1) % 2, tx, 0]])
+ *                         B[(i + 1) % 2, tx, 0] = A[tx, i + 1] * T.float32(2)
+ *                     with T.block("C"):
+ *                         T.reads([B[i % 2, tx, 0]])
+ *                         T.writes([C[tx, i]])
+ *                         C[tx, i] = B[i % 2, tx, 0] + T.float32(1)
+ *             with T.block("epilogue"):
+ *                 T.reads([B[1, tx, 0]])
+ *                 T.writes([C[tx, 15]])
+ *                 C[tx, 15] = B[1, tx, 0] + T.float32(1)
+ * \endcode
+ *
+ * The original loop has two blocks, B and C, as its direct children. The loop annotations indicate
+ * that block B has stage == 0, order == 0, block C has stage == 1, order == 1. Therefore, block B
+ * should be executed in advance of block C by one iteration. The order 0 and 1 specifies the order
+ * of block B and C inside the body block inside the result TIR.
+ *
+ * \return The IR transform pass.
+ */
+TVM_DLL Pass InjectSoftwarePipeline();
+
+TVM_DLL Pass BindParams(const Array<runtime::NDArray>& constants);
+
+/*!
+ * \brief Pass to collect tir non-scalar constants into module's 'Constants' attribute.
+ *
+ * \return The pass.
+ */
+TVM_DLL Pass ExtractPrimFuncConstants();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm
diff --git a/include/tvm/tir/usmp/utils.h b/include/tvm/tir/usmp/utils.h
index aa4f27de6d8b..d9f212f489f0 100644
--- a/include/tvm/tir/usmp/utils.h
+++ b/include/tvm/tir/usmp/utils.h
@@ -185,24 +185,24 @@ struct AllocatedPoolInfoNode : public Object {
   PoolInfo pool_info;
   /*! \brief The allocated size into this pool */
   Integer allocated_size;
-  /*! \brief An optional associated pool Var*/
-  Optional<Var> pool_var;
+  /*! \brief An optional associated pool Var index of PrimFunc params*/
+  Optional<Integer> pool_var_idx;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("pool_info", &pool_info);
     v->Visit("allocated_size", &allocated_size);
-    v->Visit("pool_var", &pool_var);
+    v->Visit("pool_var_idx", &pool_var_idx);
   }
 
   bool SEqualReduce(const AllocatedPoolInfoNode* other, SEqualReducer equal) const {
     return equal(pool_info, other->pool_info) && equal(allocated_size, other->allocated_size) &&
-           equal(pool_var, other->pool_var);
+           equal(pool_var_idx, other->pool_var_idx);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(pool_info);
     hash_reduce(allocated_size);
-    hash_reduce(pool_var);
+    hash_reduce(pool_var_idx);
   }
 
   static constexpr const char* _type_key = "tir.usmp.AllocatedPoolInfo";
@@ -211,7 +211,8 @@ struct AllocatedPoolInfoNode : public Object {
 
 class AllocatedPoolInfo : public ObjectRef {
  public:
-  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var = Var());
+  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size,
+                            Integer pool_var_idx = Integer());
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AllocatedPoolInfo, ObjectRef, AllocatedPoolInfoNode);
 };
 
diff --git a/include/tvm/tir/var.h b/include/tvm/tir/var.h
index 1ac58e18db3e..0a9000670a8e 100644
--- a/include/tvm/tir/var.h
+++ b/include/tvm/tir/var.h
@@ -38,7 +38,7 @@ namespace tir {
  *
  * A variable is uniquely identified by its address.
  *
- * Each variable is only binded once in the following nodes:
+ * Each variable is only bound once in the following nodes:
  * - Allocate
  * - For
  * - Let
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
index acff301f4f07..04027f8974fe 100644
--- a/include/tvm/topi/transform.h
+++ b/include/tvm/topi/transform.h
@@ -1608,7 +1608,11 @@ inline Tensor layout_transform(const Tensor& src, const std::string& src_layout,
       [&](const Array<Var>& dst_indices) {
         Array<PrimExpr> dst_indices_expr(dst_indices.begin(), dst_indices.end());
         Array<PrimExpr> src_indices = layout_converter.BackwardIndex(dst_indices_expr);
-        return src(src_indices);
+        PrimExpr in_range = PrimExpr(1) > PrimExpr(0);  // init with dtype=bool and value=true
+        for (size_t i = 0; i < src.ndim(); ++i) {
+          in_range = in_range && (src_indices[i] < src->shape[i]);
+        }
+        return if_then_else(in_range, src(src_indices), tvm::cast(src->dtype, PrimExpr(0)));
       },
       name, tag);
 }
diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py
index 755c1a32a3cf..6ebbbb653140 100644
--- a/python/tvm/autotvm/measure/measure_methods.py
+++ b/python/tvm/autotvm/measure/measure_methods.py
@@ -806,17 +806,22 @@ def check_remote(target, device_key, host=None, port=None, priority=100, timeout
     """
 
     def _check():
+        logger.debug("waiting for device...")
         remote = request_remote(device_key, host, port, priority)
         dev = remote.device(str(target))
         while not dev.exist:  # wait until we get an available device
             pass
+        logger.debug("device available")
 
     t = threading.Thread(
         target=_check,
     )
     t.start()
     t.join(timeout)
-    return not t.is_alive()
+
+    remote = request_remote(device_key, host, port, priority)
+    dev = remote.device(str(target))
+    return dev.exist
 
 
 def set_cuda_target_arch(arch):
diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py
index 06c33f2f7ae0..bd372572c403 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -21,7 +21,7 @@
 import multiprocessing
 import tvm
 from tvm import runtime, relay
-from tvm.contrib.nvcc import find_cuda_path, get_cuda_version
+from tvm.contrib.nvcc import get_cuda_version
 from .gen_gemm import CutlassGemmProfiler
 from .gen_conv2d import CutlassConv2DProfiler
 from .library import ConvKind
@@ -61,9 +61,8 @@ def _get_cutlass_compile_options(sm, threads, use_fast_math=False):
     ]
     if use_fast_math:
         kwargs["options"].append("-DCUTLASS_USE_TANH_FOR_SIGMOID")
-    cuda_path = find_cuda_path()
-    cuda_ver = get_cuda_version(cuda_path)
-    if cuda_ver >= 11.2:
+    cuda_ver = get_cuda_version()
+    if cuda_ver >= (11, 2):
         ncpu = multiprocessing.cpu_count() if threads < 0 else threads
         kwargs["options"].append("-t %d" % ncpu)
     return kwargs
diff --git a/python/tvm/contrib/ethosu/cascader/proposal.py b/python/tvm/contrib/ethosu/cascader/proposal.py
new file mode 100644
index 000000000000..13184108120e
--- /dev/null
+++ b/python/tvm/contrib/ethosu/cascader/proposal.py
@@ -0,0 +1,106 @@
+# 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.
+"""Proposal class to hold graph scheduling information."""
+from typing import Dict, FrozenSet, List
+import tvm._ffi
+from tvm.contrib.ethosu.cascader.plan import Plan
+
+from tvm.runtime import Object
+
+from . import _ffi_api
+from .graph import Tensor, Part, CascaderGraph
+from .tensor_config import TensorConfig, MemoryRegion
+
+
+@tvm._ffi.register_object("contrib.ethosu.cascader.Proposal")
+class Proposal(Object):
+    """A class which describes how to schedule a CascaderGraph as a series of disjoint Plans.
+
+    Attributes
+    ----------
+    graph : CascaderGraph
+        The CascaderGraph to which the Proposal applies.
+    part_group : FrozenSet[Part]
+        The Parts which are covered by the Proposal.
+    plans : List[Plan]
+        The Plans used in the Proposal.
+    input_tensor_configs : Dict[Tensor, TensorConfig]
+        The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan.
+    cascade_region : MemoryRegion
+        The MemoryRegion where cascading buffers should be homed.
+    memory_usage : int
+        The memory required to execute the Proposal in the cascading MemoryRegion.
+    cycles : int
+        The estimated cycles taken to execute the Proposal.
+
+    """
+
+    def __init__(
+        self,
+        graph: CascaderGraph,
+        part_group: FrozenSet[Part],
+        plans: List[Plan],
+        input_tensor_configs: Dict[Tensor, TensorConfig],
+        cascade_region: MemoryRegion,
+        memory_usage: Dict[MemoryRegion, int],
+        cycles: int,
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.Proposal,
+            graph,
+            list(part_group),
+            plans,
+            input_tensor_configs,
+            cascade_region,
+            memory_usage,
+            cycles,
+        )
+
+    @property
+    def graph(self) -> CascaderGraph:
+        """The CascaderGraph to which the Proposal applies."""
+        return self._graph
+
+    @property
+    def part_group(self) -> FrozenSet[Part]:
+        """The Parts which are covered by the Proposal."""
+        return frozenset(self._part_group)
+
+    @property
+    def plans(self) -> List[Plan]:
+        """The Plans used in the Proposal."""
+        return list(self._plans)
+
+    @property
+    def input_tensor_configs(self) -> Dict[Tensor, TensorConfig]:
+        """The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan."""
+        return dict(self._input_tensor_configs)
+
+    @property
+    def cascade_region(self) -> MemoryRegion:
+        """The MemoryRegion where cascading buffers should be homed."""
+        return self._cascade_region
+
+    @property
+    def memory_usage(self) -> int:
+        """The memory required to execute the Proposal in the cascading MemoryRegion."""
+        return int(self._memory_usage)
+
+    @property
+    def cycles(self) -> int:
+        """The estimated cycles taken to execute the Proposal."""
+        return int(self._cycles)
diff --git a/python/tvm/contrib/ethosu/cascader/proposal_generator.py b/python/tvm/contrib/ethosu/cascader/proposal_generator.py
new file mode 100644
index 000000000000..d79021a20539
--- /dev/null
+++ b/python/tvm/contrib/ethosu/cascader/proposal_generator.py
@@ -0,0 +1,58 @@
+# 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.
+"""Algorithms to generate Proposals for a Graph."""
+from typing import List, Dict, FrozenSet
+
+from . import _ffi_api
+from .cascader_options import CascaderOptions
+from .plan import Plan
+from .proposal import Proposal
+from .graph import CascaderGraph, Part
+
+
+def generate_proposals(
+    graph: CascaderGraph,
+    home_map: Dict[FrozenSet[Part], List[Plan]],
+    options: CascaderOptions,
+) -> List[Proposal]:
+    """Generate Pareto optimal Proposals for a CascaderGraph.
+
+    This algorithm takes a top-down dynamic programming approach to determining how
+    to optimally combine Plans into Proposals.
+
+    Parameters
+    ----------
+    graph : CascaderGraph
+        The CascaderGraph to generate Proposals for.
+    home_map : Dict[FrozenSet[Part], List[Plan]]
+        The Tensor homing map defining valid memory homes for Tensors.
+    options : CascaderOptions
+        The configuration options with which to run the generator.
+
+    Returns
+    ------
+    List[Proposal]
+        A list of Pareto optimal Proposals.
+
+    """
+    return list(
+        _ffi_api.GenerateProposals(
+            graph,
+            home_map,
+            options,
+        )
+    )
diff --git a/python/tvm/contrib/graph_executor.py b/python/tvm/contrib/graph_executor.py
index ac6803ca9842..e8e2798ef734 100644
--- a/python/tvm/contrib/graph_executor.py
+++ b/python/tvm/contrib/graph_executor.py
@@ -189,7 +189,7 @@ def set_input(self, key=None, value=None, **params):
             keys.sort(key=lambda x: -np.prod(params[x].shape))
             for k in keys:
                 # TODO(zhiics) Skip the weights for submodule in a better way.
-                # We should use MetadataModule for initialization and remove
+                # We should use ConstLoaderModule for initialization and remove
                 # params from set_input
                 val = self._get_input(k)
                 if val:
diff --git a/python/tvm/contrib/hexagon/__init__.py b/python/tvm/contrib/hexagon/__init__.py
index e728216853e5..2f568f921292 100644
--- a/python/tvm/contrib/hexagon/__init__.py
+++ b/python/tvm/contrib/hexagon/__init__.py
@@ -15,3 +15,4 @@
 # specific language governing permissions and limitations
 # under the License.
 """Hexagon APIs."""
+from .hexagon import ir_lower_vtcm_pass
diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py
index c413c60ce5c5..caaec6cf7b9a 100644
--- a/python/tvm/contrib/hexagon/session.py
+++ b/python/tvm/contrib/hexagon/session.py
@@ -17,6 +17,7 @@
 
 """Defines a Session class for Hexagon devices."""
 
+import os
 from tvm import rpc as _rpc
 
 
@@ -59,6 +60,7 @@ def __enter__(self):
                     "tvm.contrib.hexagon.create_hexagon_session",
                     self._session_name,
                     self._remote_stack_size_bytes,
+                    os.environ.get("HEXAGON_SIM_ARGS", ""),
                 ],
             )
             self.device = self._rpc.hexagon(0)
diff --git a/python/tvm/contrib/nvcc.py b/python/tvm/contrib/nvcc.py
index 19196dc3eefb..52b88d355602 100644
--- a/python/tvm/contrib/nvcc.py
+++ b/python/tvm/contrib/nvcc.py
@@ -140,27 +140,33 @@ def find_cuda_path():
     raise RuntimeError("Cannot find cuda path")
 
 
-def get_cuda_version(cuda_path):
+def get_cuda_version(cuda_path=None):
     """Utility function to get cuda version
 
     Parameters
     ----------
-    cuda_path : str
-        Path to cuda root.
+    cuda_path : Optional[str]
+
+        Path to cuda root.  If None is passed, will use
+        `find_cuda_path()` as default.
 
     Returns
     -------
     version : float
         The cuda version
+
     """
+    if cuda_path is None:
+        cuda_path = find_cuda_path()
+
     version_file_path = os.path.join(cuda_path, "version.txt")
     if not os.path.exists(version_file_path):
         # Debian/Ubuntu repackaged CUDA path
         version_file_path = os.path.join(cuda_path, "lib", "cuda", "version.txt")
     try:
         with open(version_file_path) as f:
-            version_str = f.readline().replace("\n", "").replace("\r", "")
-            return float(version_str.split(" ")[2][:2])
+            version_str = f.read().strip().split()[-1]
+            return tuple(int(field) for field in version_str.split("."))
     except FileNotFoundError:
         pass
 
@@ -171,9 +177,8 @@ def get_cuda_version(cuda_path):
     if proc.returncode == 0:
         release_line = [l for l in out.split("\n") if "release" in l][0]
         release_fields = [s.strip() for s in release_line.split(",")]
-        release_version = [f[1:] for f in release_fields if f.startswith("V")][0]
-        major_minor = ".".join(release_version.split(".")[:2])
-        return float(major_minor)
+        version_str = [f[1:] for f in release_fields if f.startswith("V")][0]
+        return tuple(int(field) for field in version_str.split("."))
     raise RuntimeError("Cannot read cuda version file")
 
 
@@ -206,7 +211,18 @@ def find_libdevice_path(arch):
     selected_ver = 0
     selected_path = None
     cuda_ver = get_cuda_version(cuda_path)
-    if cuda_ver in (9.0, 9.1, 10.0, 10.1, 10.2, 11.0, 11.1, 11.2, 11.3):
+    major_minor = (cuda_ver[0], cuda_ver[1])
+    if major_minor in (
+        (9, 0),
+        (9, 1),
+        (10, 0),
+        (10, 1),
+        (10, 2),
+        (11, 0),
+        (11, 1),
+        (11, 2),
+        (11, 3),
+    ):
         path = os.path.join(lib_path, "libdevice.10.bc")
     else:
         for fn in os.listdir(lib_path):
@@ -358,9 +374,8 @@ def have_tensorcore(compute_version=None, target=None):
 def have_cudagraph():
     """Either CUDA Graph support is provided"""
     try:
-        cuda_path = find_cuda_path()
-        cuda_ver = get_cuda_version(cuda_path)
-        if cuda_ver < 10.0:
+        cuda_ver = get_cuda_version()
+        if cuda_ver < (10, 0):
             return False
         return True
     except RuntimeError:
diff --git a/python/tvm/contrib/pipeline_executor.py b/python/tvm/contrib/pipeline_executor.py
index b858a209db83..bdf6019bc6c8 100644
--- a/python/tvm/contrib/pipeline_executor.py
+++ b/python/tvm/contrib/pipeline_executor.py
@@ -117,22 +117,18 @@ def __init__(self, module):
         # Get the packed functions from the pipeline executor.
         self._get_params_group_pipeline_map = self.module["get_params_group_pipeline_map"]
         self._run = self.module["run"]
-        self._stop = self.module["stop"]
         self._set_param = self.module["set_param"]
         self._set_input = self.module["set_input"]
         self._get_input = self.module["get_input"]
         self._get_output = self.module["get_output"]
         self._get_num_outputs = self.module["get_num_outputs"]
         self._get_input_pipeline_map = self.module["get_input_pipeline_map"]
+        self._get_pipe_execute_count = self.module["get_execute_count"]
 
     def run(self, sync=False):
         """Run the pipeline executor."""
         self._run(sync)
 
-    def stop(self):
-        """Stop the pipeline executor."""
-        self._stop()
-
     def get_input_pipeline_map(self, name):
         """Using the "name" to get the corresponding subgraph index and also get the "input name"
         of the corresponding subgraph interface.
@@ -213,6 +209,16 @@ def get_output(self):
         """
         return self._get_output()
 
+    @property
+    def num_executing_pipeline(self):
+        """Getting the count of running pipeline.
+        Returns
+        -------
+        count : int
+            The count of running pipeline.
+        """
+        return self._get_pipe_execute_count()
+
     @property
     def num_outputs(self):
         """Get the number of outputs.
diff --git a/python/tvm/contrib/relay_viz/__init__.py b/python/tvm/contrib/relay_viz/__init__.py
new file mode 100644
index 000000000000..32814b577d0d
--- /dev/null
+++ b/python/tvm/contrib/relay_viz/__init__.py
@@ -0,0 +1,105 @@
+# 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.
+"""Relay IR Visualizer"""
+from typing import Dict
+import tvm
+from tvm import relay
+from .interface import (
+    Plotter,
+    VizGraph,
+    VizParser,
+)
+from .terminal import (
+    TermPlotter,
+    TermVizParser,
+)
+
+
+class RelayVisualizer:
+    """Relay IR Visualizer
+
+    Parameters
+    ----------
+    relay_mod: tvm.IRModule
+        Relay IR module.
+    relay_param: None | Dict[str, tvm.runtime.NDArray]
+        Relay parameter dictionary. Default `None`.
+    plotter: Plotter
+        An instance of class inheriting from Plotter interface.
+        Default is an instance of `terminal.TermPlotter`.
+    parser: VizParser
+        An instance of class inheriting from VizParser interface.
+        Default is an instance of `terminal.TermVizParser`.
+    """
+
+    def __init__(
+        self,
+        relay_mod: tvm.IRModule,
+        relay_param: Dict[str, tvm.runtime.NDArray] = None,
+        plotter: Plotter = None,
+        parser: VizParser = None,
+    ):
+        self._plotter = plotter if plotter is not None else TermPlotter()
+        self._relay_param = relay_param if relay_param is not None else {}
+        self._parser = parser if parser is not None else TermVizParser()
+
+        global_vars = relay_mod.get_global_vars()
+        graph_names = []
+        # If we have main function, put it to the first.
+        # Then main function can be shown on the top.
+        for gv_node in global_vars:
+            if gv_node.name_hint == "main":
+                graph_names.insert(0, gv_node.name_hint)
+            else:
+                graph_names.append(gv_node.name_hint)
+
+        node_to_id = {}
+        # callback to generate an unique string-ID for nodes.
+        def traverse_expr(node):
+            if node in node_to_id:
+                return
+            node_to_id[node] = str(len(node_to_id))
+
+        for name in graph_names:
+            node_to_id.clear()
+            relay.analysis.post_order_visit(relay_mod[name], traverse_expr)
+            graph = self._plotter.create_graph(name)
+            self._add_nodes(graph, node_to_id)
+
+    def _add_nodes(self, graph: VizGraph, node_to_id: Dict[relay.Expr, str]):
+        """add nodes and to the graph.
+
+        Parameters
+        ----------
+        graph : VizGraph
+            a VizGraph for nodes to be added to.
+
+        node_to_id : Dict[relay.expr, str]
+            a mapping from nodes to an unique ID.
+
+        relay_param : Dict[str, tvm.runtime.NDarray]
+            relay parameter dictionary.
+        """
+        for node in node_to_id:
+            viz_node, viz_edges = self._parser.get_node_edges(node, self._relay_param, node_to_id)
+            if viz_node is not None:
+                graph.node(viz_node)
+            for edge in viz_edges:
+                graph.edge(edge)
+
+    def render(self, filename: str = None) -> None:
+        self._plotter.render(filename=filename)
diff --git a/python/tvm/contrib/relay_viz/interface.py b/python/tvm/contrib/relay_viz/interface.py
new file mode 100644
index 000000000000..6e52f024b1c5
--- /dev/null
+++ b/python/tvm/contrib/relay_viz/interface.py
@@ -0,0 +1,323 @@
+# 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.
+"""Abstract class used by :py:class:`tvm.contrib.relay_viz.RelayVisualizer`."""
+import abc
+from typing import (
+    Dict,
+    Union,
+    Tuple,
+    List,
+)
+
+import tvm
+from tvm import relay
+
+UNKNOWN_TYPE = "unknown"
+
+
+class VizNode:
+    """VizNode carry node information for `VizGraph` interface.
+
+    Parameters
+    ----------
+    node_id: str
+        Unique identifier for this node.
+    node_type: str
+        Type of this node.
+    node_detail: str
+        Any supplement for this node such as attributes.
+    """
+
+    def __init__(self, node_id: str, node_type: str, node_detail: str):
+        self._id = node_id
+        self._type = node_type
+        self._detail = node_detail
+
+    @property
+    def identity(self) -> Union[int, str]:
+        return self._id
+
+    @property
+    def type_name(self) -> str:
+        return self._type
+
+    @property
+    def detail(self) -> str:
+        return self._detail
+
+
+class VizEdge:
+    """VizEdge connect two `VizNode`.
+
+    Parameters
+    ----------
+    start_node: str
+        The identifier of the node starting the edge.
+    end_node: str
+        The identifier of the node ending the edge.
+    """
+
+    def __init__(self, start_node: str, end_node: str):
+        self._start_node = start_node
+        self._end_node = end_node
+
+    @property
+    def start(self) -> str:
+        return self._start_node
+
+    @property
+    def end(self) -> str:
+        return self._end_node
+
+
+class VizParser(abc.ABC):
+    """VizParser parses out a VizNode and VizEdges from a `relay.Expr`."""
+
+    @abc.abstractmethod
+    def get_node_edges(
+        self,
+        node: relay.Expr,
+        relay_param: Dict[str, tvm.runtime.NDArray],
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        """Get VizNode and VizEdges for a `relay.Expr`.
+
+        Parameters
+        ----------
+        node : relay.Expr
+            relay.Expr which will be parsed and generate a node and edges.
+
+        relay_param: Dict[str, tvm.runtime.NDArray]
+            relay parameters dictionary.
+
+        node_to_id : Dict[relay.Expr, str]
+            This is a mapping from relay.Expr to a unique id, generated by `RelayVisualizer`.
+
+        Returns
+        -------
+        rv1 : Union[VizNode, None]
+            VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node
+            to the graph, return None.
+
+        rv2 : List[VizEdge]
+            a list of VizEdges to describe the connectivity of the relay.Expr.
+            Can be empty list to indicate no connectivity.
+        """
+
+
+class VizGraph(abc.ABC):
+    """Abstract class for graph, which is composed of nodes and edges."""
+
+    @abc.abstractmethod
+    def node(self, viz_node: VizNode) -> None:
+        """Add a node to the underlying graph.
+        Nodes in a Relay IR Module are expected to be added in the post-order.
+
+        Parameters
+        ----------
+        viz_node : VizNode
+            A `VizNode` instance.
+        """
+
+    @abc.abstractmethod
+    def edge(self, viz_edge: VizEdge) -> None:
+        """Add an edge to the underlying graph.
+
+        Parameters
+        ----------
+        id_start : VizEdge
+            A `VizEdge` instance.
+        """
+
+
+class DefaultVizParser(VizParser):
+    """DefaultVizParser provde a set of logics to parse a various relay types.
+    These logics are inspired and heavily based on
+    `visualize` function in https://tvm.apache.org/2020/07/14/bert-pytorch-tvm
+    """
+
+    def get_node_edges(
+        self,
+        node: relay.Expr,
+        relay_param: Dict[str, tvm.runtime.NDArray],
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        if isinstance(node, relay.Function):
+            return self._function(node, node_to_id)
+        if isinstance(node, relay.expr.Call):
+            return self._call(node, node_to_id)
+        if isinstance(node, relay.expr.Var):
+            return self._var(node, relay_param, node_to_id)
+        if isinstance(node, relay.expr.Tuple):
+            return self._tuple(node, node_to_id)
+        if isinstance(node, relay.expr.TupleGetItem):
+            return self._tuple_get_item(node, node_to_id)
+        if isinstance(node, relay.expr.Constant):
+            return self._constant(node, node_to_id)
+        # GlobalVar possibly mean another global relay function,
+        # which is expected to in "Graph" level, not in "Node" level.
+        if isinstance(node, (relay.expr.GlobalVar, tvm.ir.Op)):
+            return None, []
+
+        viz_node = VizNode(node_to_id[node], UNKNOWN_TYPE, f"don't know how to parse {type(node)}")
+        viz_edges = []
+        return viz_node, viz_edges
+
+    def _var(
+        self,
+        node: relay.Expr,
+        relay_param: Dict[str, tvm.runtime.NDArray],
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        """Render rule for a relay var node"""
+
+        node_id = node_to_id[node]
+        name_hint = node.name_hint
+        node_detail = f"name_hint: {name_hint}"
+        node_type = "Var(Param)" if name_hint in relay_param else "Var(Input)"
+
+        if node.type_annotation is not None:
+            if hasattr(node.type_annotation, "shape"):
+                shape = tuple(map(int, node.type_annotation.shape))
+                dtype = node.type_annotation.dtype
+                node_detail = f"{node_detail}\nshape: {shape}\ndtype: {dtype}"
+            else:
+                node_detail = f"{node_detail}\ntype_annotation: {node.type_annotation}"
+
+        # only node
+        viz_node = VizNode(node_id, node_type, node_detail)
+        viz_edges = []
+        return viz_node, viz_edges
+
+    def _function(
+        self,
+        node: relay.Expr,
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        """Render rule for a relay function node"""
+        node_details = []
+        name = ""
+        func_attrs = node.attrs
+        if func_attrs:
+            node_details = [f"{k}: {func_attrs.get_str(k)}" for k in func_attrs.keys()]
+            # "Composite" might from relay.transform.MergeComposite
+            if "Composite" in func_attrs.keys():
+                name = func_attrs["Composite"]
+        node_id = node_to_id[node]
+
+        # Body -> FunctionNode
+        viz_node = VizNode(node_id, f"Func {name}", "\n".join(node_details))
+        viz_edges = [VizEdge(node_to_id[node.body], node_id)]
+        return viz_node, viz_edges
+
+    def _call(
+        self,
+        node: relay.Expr,
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        """Render rule for a relay call node"""
+        node_id = node_to_id[node]
+        op_name = UNKNOWN_TYPE
+        node_detail = []
+        if isinstance(node.op, tvm.ir.Op):
+            op_name = node.op.name
+            if node.attrs:
+                node_detail = [f"{k}: {node.attrs.get_str(k)}" for k in node.attrs.keys()]
+        elif isinstance(node.op, relay.Function):
+            func_attrs = node.op.attrs
+            op_name = "Anonymous Func"
+            if func_attrs:
+                node_detail = [f"{k}: {func_attrs.get_str(k)}" for k in func_attrs.keys()]
+                # "Composite" might from relay.transform.MergeComposite
+                if "Composite" in func_attrs.keys():
+                    op_name = func_attrs["Composite"]
+        elif isinstance(node.op, relay.GlobalVar):
+            op_name = "GlobalVar"
+            node_detail = [f"GlobalVar.name_hint: {node.op.name_hint}"]
+        else:
+            op_name = str(type(node.op)).split(".")[-1].split("'")[0]
+
+        # Arguments -> CallNode
+        viz_node = VizNode(node_id, f"Call {op_name}", "\n".join(node_detail))
+        args = [node_to_id[arg] for arg in node.args]
+        viz_edges = [VizEdge(arg, node_id) for arg in args]
+        return viz_node, viz_edges
+
+    def _tuple(
+        self,
+        node: relay.Expr,
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        node_id = node_to_id[node]
+
+        # Fields -> TupleNode
+        viz_node = VizNode(node_id, "Tuple", "")
+        viz_edges = [VizEdge(node_to_id[field], node_id) for field in node.fields]
+        return viz_node, viz_edges
+
+    def _tuple_get_item(
+        self,
+        node: relay.Expr,
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        node_id = node_to_id[node]
+
+        # Tuple -> TupleGetItemNode
+        viz_node = VizNode(node_id, f"TupleGetItem", "idx: {node.index}")
+        viz_edges = [VizEdge(node_to_id[node.tuple_value], node_id)]
+        return viz_node, viz_edges
+
+    def _constant(
+        self,
+        node: relay.Expr,
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        node_id = node_to_id[node]
+        node_detail = f"shape: {node.data.shape}, dtype: {node.data.dtype}"
+
+        # only node
+        viz_node = VizNode(node_id, "Const", node_detail)
+        viz_edges = []
+        return viz_node, viz_edges
+
+
+class Plotter(abc.ABC):
+    """Plotter can render a collection of Graph interfaces to a file."""
+
+    @abc.abstractmethod
+    def create_graph(self, name: str) -> VizGraph:
+        """Create a VizGraph
+
+        Parameters
+        ----------
+        name : str
+            the name of the graph
+
+        Return
+        ------
+        rv1: an instance of class inheriting from VizGraph interface.
+        """
+
+    @abc.abstractmethod
+    def render(self, filename: str) -> None:
+        """Render the graph as a file.
+
+        Parameters
+        ----------
+        filename : str
+            see the definition of implemented class.
+        """
diff --git a/python/tvm/contrib/relay_viz/terminal.py b/python/tvm/contrib/relay_viz/terminal.py
new file mode 100644
index 000000000000..7b72d9da4333
--- /dev/null
+++ b/python/tvm/contrib/relay_viz/terminal.py
@@ -0,0 +1,238 @@
+# 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.
+"""Visualize Relay IR in AST text-form."""
+
+from collections import deque
+from typing import (
+    Dict,
+    Union,
+    Tuple,
+    List,
+)
+import tvm
+from tvm import relay
+from .interface import (
+    DefaultVizParser,
+    Plotter,
+    VizEdge,
+    VizGraph,
+    VizNode,
+)
+
+
+class TermVizParser(DefaultVizParser):
+    """`TermVizParser` parse nodes and edges for `TermPlotter`."""
+
+    def __init__(self):
+        self._default_parser = DefaultVizParser()
+
+    def get_node_edges(
+        self,
+        node: relay.Expr,
+        relay_param: Dict[str, tvm.runtime.NDArray],
+        node_to_id: Dict[relay.Expr, str],
+    ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
+        """Parse a node and edges from a relay.Expr."""
+        if isinstance(node, relay.Call):
+            return self._call(node, node_to_id)
+        if isinstance(node, relay.Let):
+            return self._let(node, node_to_id)
+        if isinstance(node, relay.GlobalVar):
+            return self._global_var(node, node_to_id)
+        if isinstance(node, relay.If):
+            return self._if(node, node_to_id)
+        if isinstance(node, tvm.ir.Op):
+            return self._op(node, node_to_id)
+        if isinstance(node, relay.Function):
+            return self._function(node, node_to_id)
+
+        # Leverage logics from default parser.
+        return self._default_parser.get_node_edges(node, relay_param, node_to_id)
+
+    def _call(self, node, node_to_id):
+        node_id = node_to_id[node]
+        viz_node = VizNode(node_id, "Call", "")
+        viz_edges = [VizEdge(node_to_id[node.op], node_id)]
+        for arg in node.args:
+            arg_id = node_to_id[arg]
+            viz_edges.append(VizEdge(arg_id, node_id))
+        return viz_node, viz_edges
+
+    def _let(self, node, node_to_id):
+        node_id = node_to_id[node]
+        viz_node = VizNode(node_id, "Let", "(var, val, body)")
+        viz_edges = [
+            VizEdge(node_to_id[node.var], node_id),
+            VizEdge(node_to_id[node.value], node_id),
+            VizEdge(node_to_id[node.body], node_id),
+        ]
+        return viz_node, viz_edges
+
+    def _global_var(self, node, node_to_id):
+        node_id = node_to_id[node]
+        viz_node = VizNode(node_id, "GlobalVar", node.name_hint)
+        viz_edges = []
+        return viz_node, viz_edges
+
+    def _if(self, node, node_to_id):
+        node_id = node_to_id[node]
+        viz_node = VizNode(node_id, "If", "(cond, true, false)")
+        viz_edges = [
+            VizEdge(node_to_id[node.cond], node_id),
+            VizEdge(node_to_id[node.true_branch], node_id),
+            VizEdge(node_to_id[node.false_branch], node_id),
+        ]
+        return viz_node, viz_edges
+
+    def _op(self, node, node_to_id):
+        node_id = node_to_id[node]
+        op_name = node.name
+        viz_node = VizNode(node_id, op_name, "")
+        viz_edges = []
+        return viz_node, viz_edges
+
+    def _function(self, node, node_to_id):
+        node_id = node_to_id[node]
+        viz_node = VizNode(node_id, "Func", str(node.params))
+        viz_edges = [VizEdge(node_to_id[node.body], node_id)]
+        return viz_node, viz_edges
+
+
+class TermNode:
+    """TermNode is aimed to generate text more suitable for terminal visualization."""
+
+    def __init__(self, viz_node: VizNode):
+        self.type = viz_node.type_name
+        # We don't want too many lines in a terminal.
+        self.other_info = viz_node.detail.replace("\n", ", ")
+
+
+class TermGraph(VizGraph):
+    """Terminal graph for a relay IR Module
+
+    Parameters
+    ----------
+    name: str
+        name of this graph.
+    """
+
+    def __init__(self, name: str):
+        self._name = name
+        # A graph in adjacency list form.
+        # The key is source node, and the value is a list of destination nodes.
+        self._graph = {}
+        # a hash table for quick searching.
+        self._id_to_term_node = {}
+        # node_id in reversed post order
+        # That mean, root is the first node.
+        self._node_id_rpo = deque()
+
+    def node(self, viz_node: VizNode) -> None:
+        """Add a node to the underlying graph.
+        Nodes in a Relay IR Module are expected to be added in the post-order.
+
+        Parameters
+        ----------
+        viz_node : VizNode
+            A `VizNode` instance.
+        """
+
+        self._node_id_rpo.appendleft(viz_node.identity)
+
+        if viz_node.identity not in self._graph:
+            # Add the node into the graph.
+            self._graph[viz_node.identity] = []
+
+        # Create TermNode from VizNode
+        node = TermNode(viz_node)
+        self._id_to_term_node[viz_node.identity] = node
+
+    def edge(self, viz_edge: VizEdge) -> None:
+        """Add an edge to the terminal graph.
+
+        Parameters
+        ----------
+        id_start : VizEdge
+            A `VizEdge` instance.
+        """
+        # Take CallNode as an example, instead of "arguments point to CallNode",
+        # we want "CallNode points to arguments" in ast-dump form.
+        #
+        # The direction of edge is typically controlled by the implemented VizParser.
+        # Reverse start/end here simply because we leverage default parser implementation.
+        if viz_edge.end in self._graph:
+            self._graph[viz_edge.end].append(viz_edge.start)
+        else:
+            self._graph[viz_edge.end] = [viz_edge.start]
+
+    def render(self) -> str:
+        """Draw a terminal graph
+
+        Returns
+        -------
+        rv1: str
+            text representing a graph.
+        """
+        lines = []
+        seen_node = set()
+
+        def gen_line(indent, n_id):
+            if (indent, n_id) in seen_node:
+                return
+            seen_node.add((indent, n_id))
+
+            conn_symbol = ["|--", "`--"]
+            last = len(self._graph[n_id]) - 1
+            for i, next_n_id in enumerate(self._graph[n_id]):
+                node = self._id_to_term_node[next_n_id]
+                lines.append(
+                    f"{indent}{conn_symbol[1 if i==last else 0]}{node.type} {node.other_info}"
+                )
+                next_indent = indent
+                # increase indent for the next level.
+                next_indent += "   " if (i == last) else "|  "
+                gen_line(next_indent, next_n_id)
+
+        first_node_id = self._node_id_rpo[0]
+        first_node = self._id_to_term_node[first_node_id]
+        lines.append(f"@{self._name}({first_node.other_info})")
+        gen_line("", first_node_id)
+
+        return "\n".join(lines)
+
+
+class TermPlotter(Plotter):
+    """Terminal plotter"""
+
+    def __init__(self):
+        self._name_to_graph = {}
+
+    def create_graph(self, name):
+        self._name_to_graph[name] = TermGraph(name)
+        return self._name_to_graph[name]
+
+    def render(self, filename):
+        """If filename is None, print to stdio. Otherwise, write to the filename."""
+        lines = []
+        for name in self._name_to_graph:
+            text_graph = self._name_to_graph[name].render()
+            lines.append(text_graph)
+        if filename is None:
+            print("\n".join(lines))
+        else:
+            with open(filename, "w") as out_file:
+                out_file.write("\n".join(lines))
diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py
index d260c98b6721..df56e3b9825d 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -95,7 +95,8 @@ def add_compile_parser(subparsers, _):
         metavar=("name=value"),
         help="configurations to be used at compile time. This option can be provided multiple "
         "times, each one to set one configuration value, "
-        "e.g. '--pass-config relay.backend.use_auto_scheduler=0'.",
+        "e.g. '--pass-config relay.backend.use_auto_scheduler=0', "
+        "e.g. '--pass-config tir.add_lower_pass=opt_level1,pass1,opt_level2,pass2'.",
     )
 
     generate_target_args(parser)
diff --git a/python/tvm/driver/tvmc/pass_config.py b/python/tvm/driver/tvmc/pass_config.py
index 7cf0f0143e60..dde5b9c659d8 100644
--- a/python/tvm/driver/tvmc/pass_config.py
+++ b/python/tvm/driver/tvmc/pass_config.py
@@ -18,10 +18,41 @@
 TVMC PassContext Interface
 """
 
+import importlib
+
 import tvm
 from tvm.driver.tvmc import TVMCException
 
 
+def load_function(full_name):
+    """Dynamic loading a function by the full name.
+    Parameters
+    ----------
+    full_name: str
+        The name of a PackedFunc or a string of the form "path.to.module.func"
+        that indicates the module that can be imported.
+        You must be aware of the load order here, it first tries to find it via
+        TVM global function, if not find, try to import it by "importlib.import_module".
+    Returns
+    -------
+    func: function or PackedFunc
+        The loaded fucntion.
+    """
+    global_func = tvm.get_global_func(full_name, allow_missing=True)
+    if global_func is not None:
+        return global_func
+
+    # split full name "path.to.module.func" into two parts ["path.to.module", "func"]
+    module_name, func_name = full_name.rsplit(".", 1)
+
+    # import module and find the function
+    module = importlib.import_module(module_name)
+    if hasattr(module, func_name):
+        return getattr(module, func_name)
+
+    raise TVMCException(f"No function '{func_name}' found in module '{module_name}'.")
+
+
 def get_pass_config_value(name, value, config_type):
     """Get a PassContext configuration value, based on its config data type.
 
@@ -41,6 +72,8 @@ def get_pass_config_value(name, value, config_type):
         specified by config_type.
     """
 
+    parsed_value = None
+
     if config_type == "IntImm":
         # "Bool" configurations in the PassContext are recognized as
         # IntImm, so deal with this case here
@@ -56,11 +89,44 @@ def get_pass_config_value(name, value, config_type):
             parsed_value = mapping_values.get(value.lower(), None)
 
         if parsed_value is None:
-            raise TVMCException(f"Invalid value '{value}' for configuration '{name}'. ")
+            raise TVMCException(f"Invalid value '{value}' for configuration '{name}'.")
 
-    if config_type == "runtime.String":
+    elif config_type == "runtime.String":
         parsed_value = value
 
+    elif config_type == "Array":
+        if name == "tir.add_lower_pass":
+            pass_list = value.split(",")
+            if len(pass_list) % 2 != 0:
+                raise TVMCException(
+                    f"The configuration of '{name}' must be of the form "
+                    "'tir.add_lower_pass=opt_level1,pass1,opt_evel2,pass2'"
+                )
+
+            parsed_value = []
+            for i in range(0, len(pass_list), 2):
+                level, pass_func = pass_list[i].strip(), pass_list[i + 1].strip()
+                try:
+                    level = int(level)
+                except ValueError:
+                    raise TVMCException(f"Only integer is allow for configuration '{name}'.")
+
+                # TODO (@leeexyz) We should parse configurations of each tir Pass.
+                #     For now, we only use the defaults. Currently, There are four config nodes:
+                #     `tir.transform.LoopPartitionConfig`
+                #     `tir.transform.UnrollLoopConfig`
+                #     `tir.transform.HoistIfThenElseConfig`
+                #     `tir.transform.InjectDoubleBufferConfig`
+                # loading pass func and calling it to get the Pass
+                pass_func = load_function(pass_func)()
+                parsed_value.append((level, pass_func))
+        else:
+            raise TVMCException(f"Unsupported configuration '{name}' for '{config_type}' type.")
+
+    else:
+        # not raise here cause we alreay checked before calling this function
+        pass
+
     return parsed_value
 
 
@@ -81,7 +147,7 @@ def parse_configs(input_configs):
         return {}
 
     all_configs = tvm.ir.transform.PassContext.list_configs()
-    supported_config_types = ("IntImm", "runtime.String")
+    supported_config_types = ("IntImm", "runtime.String", "Array")
     supported_configs = [
         name for name in all_configs.keys() if all_configs[name]["type"] in supported_config_types
     ]
@@ -116,7 +182,13 @@ def parse_configs(input_configs):
                 f"The following configurations are supported: {', '.join(supported_configs)}"
             )
 
-        parsed_value = get_pass_config_value(name, value, all_configs[name]["type"])
-        pass_context_configs[name] = parsed_value
+        config_type = all_configs[name]["type"]
+        parsed_value = get_pass_config_value(name, value, config_type)
+
+        if config_type == "Array" and name in pass_context_configs:
+            # merge configs if the configuration exists
+            pass_context_configs[name].extend(parsed_value)
+        else:
+            pass_context_configs[name] = parsed_value
 
     return pass_context_configs
diff --git a/python/tvm/driver/tvmc/runner.py b/python/tvm/driver/tvmc/runner.py
index 4f1be94f6523..1ffb400fc9c9 100644
--- a/python/tvm/driver/tvmc/runner.py
+++ b/python/tvm/driver/tvmc/runner.py
@@ -103,6 +103,12 @@ def add_run_parser(subparsers, main_parser):
         "Profiling may also have an impact on inference time, "
         "making it take longer to be generated. (non-micro devices only)",
     )
+    parser.add_argument(
+        "--end-to-end",
+        action="store_true",
+        help="Measure data transfers as well as model execution. This can provide a "
+        "more realistic performance measurement in many cases.",
+    )
     parser.add_argument(
         "--repeat", metavar="N", type=int, default=1, help="run the model n times. Defaults to '1'"
     )
@@ -262,6 +268,7 @@ def drive_run(args):
         repeat=args.repeat,
         number=args.number,
         profile=args.profile,
+        end_to_end=args.end_to_end,
         options=options,
     )
 
@@ -400,6 +407,7 @@ def run_module(
     repeat: int = 10,
     number: int = 10,
     profile: bool = False,
+    end_to_end: bool = False,
     options: dict = None,
 ):
     """Run a compiled graph executor module locally or remotely with
@@ -435,6 +443,10 @@ def run_module(
         The number of runs to measure within each repeat.
     profile : bool
         Whether to profile the run with the debug runtime.
+    end_to_end : bool
+        Whether to measure the time of memory copies as well as model
+        execution. Turning this on can provide a more realistic estimate
+        of how long running the model in production would take.
 
     Returns
     -------
@@ -557,8 +569,13 @@ def run_module(
             module.run()
             times = []
         else:
-            # call the benchmarking function of the executor
-            times = module.benchmark(dev, number=number, repeat=repeat)
+            # Call the benchmarking function of the executor.
+            # Optionally measure e2e data transfers from the
+            # CPU to device memory overheads (e.g. PCIE
+            # overheads if the device is a discrete GPU).
+            if end_to_end:
+                dev = session.cpu()
+            times = module.benchmark(dev, number=number, repeat=repeat, end_to_end=end_to_end)
 
         logger.debug("Collecting the output tensors.")
         num_outputs = module.get_num_outputs()
diff --git a/python/tvm/meta_schedule/cost_model/xgb_model.py b/python/tvm/meta_schedule/cost_model/xgb_model.py
index ec6bb3c76d3d..60b139aab4ae 100644
--- a/python/tvm/meta_schedule/cost_model/xgb_model.py
+++ b/python/tvm/meta_schedule/cost_model/xgb_model.py
@@ -133,7 +133,7 @@ def obj_square_error(self, ys_pred: np.ndarray) -> Tuple[np.ndarray, np.ndarray]
         # Making prediction
         ys_pred = self.predict_with_score(ys_pred)
         # Propagate prediction to each block
-        ys_pred = ys_pred[self.ids]
+        ys_pred = ys_pred[self.ids]  # pylint: disable=invalid-sequence-index
         # The gradient and hessian
         ys = self.dmatrix.get_label()  # type: ignore # pylint: disable=invalid-name
         gradient = ys_pred - ys
@@ -158,7 +158,7 @@ def rmse(self, ys_pred: np.ndarray) -> Tuple[str, float]:
         # Making prediction
         ys_pred = self.predict_with_score(ys_pred)
         # Propagate prediction to each block
-        ys_pred = ys_pred[self.ids]
+        ys_pred = ys_pred[self.ids]  # pylint: disable=invalid-sequence-index
         # The RMSE
         ys = self.dmatrix.get_label()  # type: ignore # pylint: disable=invalid-name
         square_error = np.square(ys_pred - ys)
diff --git a/python/tvm/micro/model_library_format.py b/python/tvm/micro/model_library_format.py
index 3f64154b4edf..6b59b3443078 100644
--- a/python/tvm/micro/model_library_format.py
+++ b/python/tvm/micro/model_library_format.py
@@ -64,6 +64,10 @@ def generate_c_interface_header(
     return metadata_header
 
 
+# List of type_key for modules which are ephemeral and do not need to be exported.
+EPHEMERAL_MODULE_TYPE_KEYS = ("metadata_module",)
+
+
 def _populate_codegen_dir(mod, codegen_dir: str, module_name: str = None):
     """Populate the codegen sub-directory as part of a Model Library Format export.
 
@@ -79,6 +83,11 @@ def _populate_codegen_dir(mod, codegen_dir: str, module_name: str = None):
     """
     dso_modules = mod._collect_dso_modules()
     non_dso_modules = mod._collect_from_import_tree(lambda m: m not in dso_modules)
+
+    # Filter ephemeral modules which cannot be exported.
+    dso_modules = [m for m in dso_modules if m.type_key not in EPHEMERAL_MODULE_TYPE_KEYS]
+    non_dso_modules = [m for m in non_dso_modules if m.type_key not in EPHEMERAL_MODULE_TYPE_KEYS]
+
     if non_dso_modules:
         raise UnsupportedInModelLibraryFormatError(
             f"Don't know how to export non-c or non-llvm modules; found: {non_dso_modules!r}"
@@ -181,42 +190,26 @@ def _build_function_memory_map(function_metadata):
     """
     device_max_workspace = dict()
     main_func_metadata = function_metadata[MAIN_FUNC_NAME_STR]
-    num_targets = len(main_func_metadata.workspace_sizes.items())
-    from tvm.driver import tvmc  # pylint: disable=import-outside-toplevel
-
-    external_codegens = tvmc.composite_target.get_codegen_names()
     func_entries = []
     target_local_entries = dict()
-    for i in range(num_targets):
-        main_target = main_func_metadata.workspace_sizes.items()[i][0]
-        device_max_workspace[main_target] = 0
-        for func_name, finfo in function_metadata.items():
-            if func_name == MAIN_FUNC_NAME_STR:
-                continue
-            target_local_entries[func_name] = list()
 
-        for func_name, finfo in function_metadata.items():
-            # Skip a few unsupported cases:
-            # 1. The main function metadata is exported elsewhere.
-            # 2. BYOC operator implementations do not currently export useful FunctionInfo.
-            if func_name == MAIN_FUNC_NAME_STR or not finfo.tir_primfuncs:
-                continue
-            assert (
-                len(finfo.constant_sizes.items()) == num_targets
-            ), f"{func_name}: found {finfo.constant_sizes!r} vs {num_targets}"
-            assert len(finfo.io_sizes.items()) == num_targets
-            target = finfo.workspace_sizes.items()[i][0]
-            workspace_size = finfo.workspace_sizes.items()[i][1]
+    for func_name, finfo in function_metadata.items():
+        # Skip a few unsupported cases:
+        # 1. The main function metadata is exported elsewhere.
+        # 2. BYOC operator implementations do not currently export useful FunctionInfo.
+        if func_name == MAIN_FUNC_NAME_STR or not finfo.tir_primfuncs:
+            continue
+        if func_name not in target_local_entries.keys():
+            target_local_entries[func_name] = list()
+        for target in dict(finfo.workspace_sizes).keys():
+            workspace_size = finfo.workspace_sizes[target]
             target_entry = {
                 "device": int(target.kind.device_type),
                 "workspace_size_bytes": int(workspace_size),
             }
             target_local_entries[func_name].append(target_entry)
-            if workspace_size > device_max_workspace.get(target, 0):
-                device_max_workspace[target] = workspace_size
-            # TODO(Mousius) - Remove this massive hack when Targets are unified
-            if target.kind.name in external_codegens:
-                device_max_workspace[main_target] += int(workspace_size)
+            if workspace_size >= device_max_workspace.get(int(target.kind.device_type), 0):
+                device_max_workspace[int(target.kind.device_type)] = workspace_size
 
     for func_name, target_entries_ in target_local_entries.items():
         func_entry = {
@@ -225,25 +218,46 @@ def _build_function_memory_map(function_metadata):
         }
         func_entries.append(func_entry)
 
-    target_main_entries = list()
-    for i in range(num_targets):
-        target = main_func_metadata.workspace_sizes.items()[i][0]
-        main_func_local_workspace = main_func_metadata.workspace_sizes.items()[i][1]
-        main_func_constants = main_func_metadata.constant_sizes.items()[i][1]
-        main_func_io = main_func_metadata.io_sizes.items()[i][1]
-        target_main_entries.append(
-            {
-                "device": int(target.kind.device_type),
-                "workspace_size_bytes": int(device_max_workspace[target])
-                + int(main_func_local_workspace),
-                "constants_size_bytes": int(main_func_constants),
-                "io_size_bytes": int(main_func_io),
-            }
+    target_main_entries = dict()
+
+    def _create_empty_entry(target_device_type):
+        return {
+            "device": int(target_device_type),
+            "workspace_size_bytes": 0,
+            "constants_size_bytes": 0,
+            "io_size_bytes": 0,
+        }
+
+    for target in dict(main_func_metadata.workspace_sizes).keys():
+        main_func_local_workspace = main_func_metadata.workspace_sizes[target]
+        target_main_entries[int(target.kind.device_type)] = _create_empty_entry(
+            int(target.kind.device_type)
+        )
+        target_main_entries[int(target.kind.device_type)]["workspace_size_bytes"] = int(
+            device_max_workspace.get(int(target.kind.device_type), 0)
+        ) + int(main_func_local_workspace)
+
+    for target in dict(main_func_metadata.constant_sizes).keys():
+        if int(target.kind.device_type) not in target_main_entries.keys():
+            target_main_entries[int(target.kind.device_type)] = _create_empty_entry(
+                int(target.kind.device_type)
+            )
+        target_main_entries[int(target.kind.device_type)]["constants_size_bytes"] = int(
+            main_func_metadata.constant_sizes[target]
+        )
+
+    for target in dict(main_func_metadata.io_sizes).keys():
+        if int(target.kind.device_type) not in target_main_entries.keys():
+            target_main_entries[int(target.kind.device_type)] = _create_empty_entry(
+                int(target.kind.device_type)
+            )
+        target_main_entries[int(target.kind.device_type)]["io_size_bytes"] = int(
+            main_func_metadata.io_sizes[target]
         )
 
     ret = {
         "operator_functions": func_entries,
-        "main": target_main_entries,
+        "main": list(target_main_entries.values()),
     }
     return ret
 
@@ -270,9 +284,13 @@ def _get_inputs_and_outputs_from_module(mod):
     main_func = _get_main_relay_func(mod)
     inputs = [argument.name_hint for argument in main_func.params]
 
-    outputs = ["output"]
-    if isinstance(main_func.ret_type, TupleType):
-        outputs = _convert_tuple_to_outputs(main_func.ret_type)
+    if "output_tensor_names" in main_func.attrs:
+        outputs = main_func.attrs["output_tensor_names"]
+    else:
+        if isinstance(main_func.ret_type, TupleType):
+            outputs = _convert_tuple_to_outputs(main_func.ret_type)
+        else:
+            outputs = ["output"]
 
     return inputs, outputs
 
diff --git a/python/tvm/relay/backend/contrib/ethosu/legalize.py b/python/tvm/relay/backend/contrib/ethosu/legalize.py
index 64c6cefb8b58..704168eb34a5 100644
--- a/python/tvm/relay/backend/contrib/ethosu/legalize.py
+++ b/python/tvm/relay/backend/contrib/ethosu/legalize.py
@@ -16,7 +16,7 @@
 # under the License.
 # pylint: disable=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter
 """A set of passes to legalize some of operations for the NPU"""
-from typing import List, Type, Callable
+from typing import List, Type, Callable, Any, Dict
 import math
 
 import numpy as np  # type: ignore
@@ -145,16 +145,22 @@ def __call__(self, *args, **kwargs):
 
 
 def get_lut_from_func(
-    ifm_scale: float, ifm_zp: int, ofm_scale: float, ofm_zp: int, func: Callable[[float], float]
+    ifm_scale: float,
+    ifm_zp: int,
+    ofm_scale: float,
+    ofm_zp: int,
+    func: Callable[[float], float],
+    func_params: Dict[str, Any],
 ) -> List[int]:
-    """Method to calculate the values of the lookup table based on the calculation function"""
+    """Calculates the values of the lookup table based on the calculation function"""
+
     lut_values = list()
     # Only int8 is currently supported
     dtype = np.int8
     qmin, qmax = np.iinfo(dtype).min, np.iinfo(dtype).max
     for x in range(qmin, qmax + 1):
         x_real = ifm_scale * (x - ifm_zp)
-        out_real = func(x_real)
+        out_real = func(x_real, **func_params)
         lut_result = int(util.round_away_zero(ofm_zp + out_real / ofm_scale))
         lut_result = min(qmax, max(qmin, lut_result))
         lut_values.append(lut_result)
@@ -166,32 +172,58 @@ class LutActivationRewriter(DFPatternCallback):
     """A class to create an identity operator with the LUT"""
 
     def __init__(
-        self, params_class: Type, activation_type: str, calc_func: Callable[[float], float]
+        self,
+        params_class: Type,
+        activation_type: str,
+        calc_func: Callable[[float], float],
     ):
         super().__init__(require_type=True, rewrite_once=True)
+        self.params_class = params_class
         self.pattern = (wildcard().has_attr({"Composite": params_class.composite_name}))(wildcard())
         self.activation_type = activation_type
         self.calc_func = calc_func
 
+    def get_calc_func_params(self, expr: tvm.relay.Expr) -> Dict[str, Any]:
+        """
+        Overridable method that can be used to extract additional arguments
+        for passing to calc_func.
+
+        Parameters
+        ----------
+        expr : tvm.relay.Expr
+            The matched composite activation function.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Maps argument name to argument value.
+        """
+        return {}
+
     def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map):
-        id_input = post.args[0]
+        params = self.params_class(post.op.body)
+        params.ifm.tensor = post.args[0]
 
-        quantize_args = post.op.body.args
-        output_scale = float(quantize_args[1].data.asnumpy())
-        output_zp = int(quantize_args[2].data.asnumpy())
+        calc_func_params = self.get_calc_func_params(post.op)
 
-        dequantize_args = quantize_args[0].args[0].args
-        input_scale = float(dequantize_args[1].data.asnumpy())
-        input_zp = int(dequantize_args[2].data.asnumpy())
+        input_scale = float(params.ifm.q_params.scale_f32)
+        input_zp = int(params.ifm.q_params.zero_point)
+        output_scale = float(params.ofm.q_params.scale_f32)
+        output_zp = int(params.ofm.q_params.zero_point)
 
         lut_values = get_lut_from_func(
-            input_scale, input_zp, output_scale, output_zp, self.calc_func
+            input_scale,
+            input_zp,
+            output_scale,
+            output_zp,
+            self.calc_func,
+            calc_func_params,
         )
-        lut = relay.const(lut_values, dtype="uint8")
+        lut = relay.const(lut_values, dtype=params.ifm.dtype)
 
         # We baked the requantization into the LUT, so we don't requantize the identity operator
         identity = ethosu_ops.ethosu_identity(
-            ifm=id_input,
+            ifm=params.ifm.tensor,
             lut=lut,
             ifm_scale=input_scale,
             ifm_zero_point=input_zp,
@@ -270,6 +302,42 @@ def __call__(self, *args, **kwargs):
         pass
 
 
+def leaky_relu_calc_func(x: float, alpha: float) -> float:
+    """Function to calculate the values for leaky relu."""
+    return x if x >= 0 else x * alpha
+
+
+class LeakyReLURewriter(LutActivationRewriter):
+    """This pass adds leaky relu as a LUT for identity op."""
+
+    def __init__(self):
+        super().__init__(
+            params_class=ethosu_patterns.LeakyReLUParams,
+            activation_type="LUT",
+            calc_func=leaky_relu_calc_func,
+        )
+
+    def get_calc_func_params(self, expr: tvm.relay.Expr) -> Dict[str, Any]:
+        params = ethosu_patterns.LeakyReLUParams(expr.body)
+        return {"alpha": params.alpha}
+
+
+@ir.transform.module_pass(opt_level=1)
+class LegalizeLeakyReLU:
+    """This is the pass that wraps LeakyReLURewriter."""
+
+    def transform_module(
+        self, mod: tvm.ir.IRModule, ctx: tvm.ir.transform.PassContext
+    ) -> tvm.ir.IRModule:
+        for global_var, func in mod.functions.items():
+            func = rewrite(LeakyReLURewriter(), func)
+            mod.update_func(global_var, func)
+        return mod
+
+    def __call__(self, *args, **kwargs):
+        pass
+
+
 class Conv2DRewriter(DFPatternCallback):
     """Convert conv2d related composite functions into ethosu_conv2d operators"""
 
@@ -1536,6 +1604,7 @@ def transform_module(
         mod = LegalizeShl()(mod)
         mod = LegalizeAbs()(mod)
         mod = LegalizeTanh()(mod)
+        mod = LegalizeLeakyReLU()(mod)
         mod = LegalizeMean()(mod)
         mod = LegalizeConcat()(mod)
         mod = LegalizeSigmoid()(mod)
diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py b/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
index ee35da4cab61..bdc3b3186202 100644
--- a/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
+++ b/python/tvm/relay/backend/contrib/ethosu/tir/compiler.py
@@ -87,7 +87,9 @@ def lower_ethosu(sch, args, const_dict, name="main"):
         mod = ethosu_passes.ReplaceOperators()(mod)
         mod = tvm.tir.transform.RemoveNoOp()(mod)
         mod, const_dict = ethosu_passes.EncodeConstants(const_dict)(mod)
-        mod = tvm.tir.transform.StorageRewrite()(mod)
+        disable_storage_rewrite = curr_cfg.get("tir.disable_storage_rewrite", False)
+        if not disable_storage_rewrite:
+            mod = tvm.tir.transform.StorageRewrite()(mod)
         mod = tvm.tir.transform.RemoveNoOp()(mod)
         mod = ethosu_passes.AnnotateAllocates()(mod)
         mod, const_dict = ethosu_passes.CreatePrimFuncWithoutConstants(const_dict)(mod)
diff --git a/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py b/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py
index c55a6310ffa5..f642f5f7cfad 100644
--- a/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py
+++ b/python/tvm/relay/backend/contrib/ethosu/tir_to_cs_translator.py
@@ -42,15 +42,6 @@ class BufferType(Enum):
     shram = auto()
 
 
-_REGION_MAP = {
-    BufferType.constant: 0,
-    BufferType.scratch: 1,
-    BufferType.input: 3,
-    BufferType.output: 4,
-    BufferType.shram: int((1 << 8) | (3 << 0)),
-}
-
-
 class BufferInfo(NamedTuple):
     """A data structure to hold metadata of the buffer."""
 
@@ -81,6 +72,111 @@ def get_accelerator_arch_config(accel_type):
     return accel_config_str_map[accel_type]
 
 
+class RegionOffset(NamedTuple):
+    """A data structure to hold region and address offset corresponding to a tensor"""
+
+    region: int
+    offset: int
+
+
+def analyze_scratch_memory_acesses(mod: tvm.IRModule, candidate_regions_for_scratch: List[int]):
+    """
+    This function analyzes the IRModule for intermediary tensors that can be resulting
+    from a offset of pool variables (via Let nodes) and/or allocate nodes. The allocate
+    nodes will be folded into a single TVMBackendallocWorkspace call with offsets. Ultimately
+    this will produce a mapping from each such node to a RegionOffset named tuple that
+    has the region and the obtained offset, as mentioned above.
+
+    Parameters
+    ----------
+    mod: tvm.IRModule
+        The TIR module containing ethosu extern calls
+    candidate_regions_for_scratch: List[int]
+        A list of region integers that could be used for scratch regions
+
+    Returns
+    -------
+    scratch_region_map : Dict[tvm.tir.Var, RegionOffset]
+        A map between buffer vars to scratch regions they are assigned
+    tvm_backend_alloc_workspace_size : int
+        The size of tvm_backend_alloc_workspace call required to service
+        remaining allocate nodes if any
+    tvm_backend_alloc_workspace_region : int
+        The region associated with the tvm_backend_alloc_workspace
+    """
+    scratch_region_map = dict()
+    pool_var_region_map = dict()
+    # There should only be a single function
+    assert len(mod.functions.items()) == 1
+    primfunc = mod.functions.items()[0][1]
+    if "pool_args" in primfunc.attrs.keys():
+        pool_args = primfunc.attrs["pool_args"]
+        for pool_arg in pool_args:
+            pool_param = primfunc.params[int(pool_arg.pool_var_idx)]
+            pool_var_region_map[pool_param] = candidate_regions_for_scratch.pop()
+            scratch_region_map[pool_param] = RegionOffset(
+                region=pool_var_region_map[pool_param], offset=None
+            )
+
+    def analyze_pool_access(stmt):
+        if isinstance(stmt, tvm.tir.stmt.LetStmt):
+            call_address_of = stmt.value
+            load = call_address_of.args[0]
+            pool_var = load.buffer_var
+            scratch_region_map[stmt.var] = RegionOffset(
+                region=pool_var_region_map[pool_var], offset=int(load.index)
+            )
+
+    tvm.tir.stmt_functor.post_order_visit(primfunc.body, analyze_pool_access)
+
+    tvmbaw_region = None
+    if len(candidate_regions_for_scratch) > 0:
+        tvmbaw_region = candidate_regions_for_scratch.pop()
+        tvmbaw_size = 0
+
+        # If there are tir.Allocate remaining by now, they need to be serviced via
+        # TVMBAW calls.
+        def analyze_remaining_allocates(stmt):
+            nonlocal tvmbaw_size
+            if isinstance(stmt, tvm.tir.stmt.Allocate):
+                allocate = stmt
+                pointer_type = allocate.buffer_var.type_annotation
+                storage_scope = pointer_type.storage_scope
+                if storage_scope == "global":
+                    dtype_bytes = np.iinfo(np.dtype(allocate.dtype)).bits // 8
+                    size_in_bytes = int(dtype_bytes * np.prod(list(allocate.extents)))
+                    # Every memory address the NPU access have to be 16 byte aligned
+                    size_in_bytes = util.round_up(size_in_bytes, 16)
+                    address = tvmbaw_size
+                    tvmbaw_size += size_in_bytes
+                    scratch_region_map[allocate.buffer_var] = RegionOffset(
+                        region=tvmbaw_region, offset=address
+                    )
+
+        tvm.tir.stmt_functor.post_order_visit(primfunc.body, analyze_remaining_allocates)
+
+    return (
+        scratch_region_map,
+        tvmbaw_size,
+        tvmbaw_region,
+    )
+
+
+def _get_region(buffer_type, var=None, scratch_region_map=None):
+    """A helper to obtain regions for buffer_types and buffer vars"""
+    static_regions = {
+        BufferType.constant: 0,
+        BufferType.input: 3,
+        BufferType.output: 4,
+        BufferType.shram: int((1 << 8) | (3 << 0)),
+    }
+    if buffer_type in static_regions.keys():
+        return static_regions[buffer_type]
+    assert buffer_type == BufferType.scratch
+    assert var in scratch_region_map.keys(), f"{var} is not analyzed for scratch regions"
+    return scratch_region_map[var].region
+
+
 def translate(tir_module, params):
     """This will take an tir module for the NPU
     and compile to command stream
@@ -106,21 +202,31 @@ def translate(tir_module, params):
         base addresses to be used by the driver
     """
 
+    # The NPU has 6 usable regions ranging from 0-6
+    # The regions 0, 3, and 4 is already used for input,
+    # output and constant, respectively (See _get_regions()).
+    # Thus, for scratch we are left with 5, 2 and 1.
+    candidate_regions_for_scratch = [5, 2, 1]
+    (
+        scratch_region_map,
+        tvmbaw_workspace_size,
+        tvmbaw_region,
+    ) = analyze_scratch_memory_acesses(tir_module, candidate_regions_for_scratch)
     buffer_info = extract_buffer_info(tir_module, params)
     call_extern_list = extract_call_extern_list(tir_module)
     _npu_ops = list()
     for call_extern in call_extern_list:
         _npu_ops.append(translate_ethosu_tir_call_extern(call_extern))
-    _npu_ops, constant_data, scratch_size = assign_addresses(buffer_info, _npu_ops)
-    base_addresses = extract_param_base_addresses(tir_module, buffer_info)
-    if scratch_size > 0:
+    _npu_ops, constant_data = assign_addresses(buffer_info, _npu_ops, scratch_region_map)
+    base_addresses = extract_param_base_addresses(tir_module, buffer_info, scratch_region_map)
+    if tvmbaw_workspace_size:
         base_addresses.append(
             util.BaseAddress(
-                "scratch",
-                None,
-                _REGION_MAP[BufferType.scratch],
-                scratch_size,
-                True,
+                name="tvmbaw",
+                primfunc_param_idx=None,
+                region=tvmbaw_region,
+                size=tvmbaw_workspace_size,
+                is_runtime_allocation=True,
             )
         )
     target_accel_config = vela_api.get_accelerator_config()
@@ -129,7 +235,7 @@ def translate(tir_module, params):
     return payload.hex(), constant_data, base_addresses
 
 
-def extract_param_base_addresses(mod, buffer_info) -> List[util.BaseAddress]:
+def extract_param_base_addresses(mod, buffer_info, scratch_region_map) -> List[util.BaseAddress]:
     """This function extracts base addresses to be used by the driver
 
     Parameters
@@ -161,7 +267,12 @@ def extract_param_base_addresses(mod, buffer_info) -> List[util.BaseAddress]:
         element_size_bytes = np.iinfo(dtype).bits // 8
         size_bytes = element_size_bytes * np.prod(list(buffer.shape))
         base_addresses.append(
-            util.BaseAddress(param.name, idx, _REGION_MAP[buffer_info[param].btype], size_bytes)
+            util.BaseAddress(
+                param.name,
+                idx,
+                _get_region(buffer_info[param].btype, param, scratch_region_map),
+                size_bytes,
+            )
         )
         idx += 1
 
@@ -227,39 +338,42 @@ def extract_buffer_info(
             const_data, const_data.shape, const_data.dtype, BufferType.constant
         )
 
-    for param in primfunc.params:
+    pool_param_indices = list()
+    if "pool_args" in primfunc.attrs.keys():
+        pool_args = primfunc.attrs["pool_args"]
+        pool_param_indices = [allocated_pool_info.pool_var_idx for allocated_pool_info in pool_args]
+
+    for idx, param in enumerate(primfunc.params):
         if param not in buffer_info.keys():
+            if idx in pool_param_indices:
+                btype = BufferType.scratch
+            else:
+                btype = BufferType.input_or_output
             buffer_info[param] = BufferInfo(
                 None,
                 None,
                 None,
-                BufferType.input_or_output,
+                btype,
             )
 
     def populate_allocate_buffer_info(stmt):
         if isinstance(stmt, tvm.tir.stmt.Allocate):
             allocate = stmt
-            if "placeholder" in allocate.buffer_var.name:
-                storage_scope = allocate.buffer_var.name.split(".")[-1]
-            else:
-                storage_scope = "global"
-
+            pointer_type = allocate.buffer_var.type_annotation
+            storage_scope = pointer_type.storage_scope
             if storage_scope == "local":
-                buffer_type = BufferType.shram
-            else:
-                buffer_type = BufferType.scratch
-            buffer_info[allocate.buffer_var] = BufferInfo(
-                None,
-                allocate.extents,
-                allocate.dtype,
-                buffer_type,
-            )
+                buffer_info[allocate.buffer_var] = BufferInfo(
+                    None,
+                    allocate.extents,
+                    allocate.dtype,
+                    BufferType.shram,
+                )
 
     tvm.tir.stmt_functor.post_order_visit(primfunc.body, populate_allocate_buffer_info)
     return buffer_info
 
 
-def assign_addresses(buffer_info, npu_ops):
+def assign_addresses(buffer_info, npu_ops, scratch_region_map):
     """This function will assign addresses to tensors
     within two buffers : scratch and constants.
     The scratch is the buffer created to hold all intermediary data
@@ -272,14 +386,14 @@ def assign_addresses(buffer_info, npu_ops):
         The key is the buffer name to BufferInfo
     npu_ops : list
         A list of Vela NpuOps with tir.Loads for addresses
+    scratch_region_map : Dict[tvm.tir.Var, RegionOffset]
+        A buffer_var to region and offset map.
     Returns
     -------
     npu_ops : list
         A list of Vela NpuOps with addesses within scratch and constant buffers
     constant_tensor : NDArray
         A unified constant data array of uint8 as the constant buffer
-    scratch_size : int
-        The size of the scratch tensor.
     """
 
     def replace_npu_fm_with_address(npu_fm):
@@ -290,21 +404,37 @@ def replace_npu_fm_with_address(npu_fm):
         assert npu_fm.tiles.addresses[1:] == [0, 0, 0]
         npu_fm.tiles.addresses[1:] = [0, 0, 0]
         buffer = npu_fm.tiles.addresses[0].buffer_var
-        assert buffer in buffer_addresses.keys()
-        address, buffer_type = buffer_addresses[buffer]
+
+        if buffer in scratch_region_map.keys():
+            address = scratch_region_map[buffer].offset
+            region = scratch_region_map[buffer].region
+        else:
+            assert buffer in buffer_addresses.keys()
+            address, buffer_type = buffer_addresses[buffer]
+            region = _get_region(buffer_type)
+
         index = npu_fm.tiles.addresses[0].index * (
             np.iinfo(np.dtype(npu_fm.tiles.addresses[0])).bits // 8
         )
         npu_fm.tiles.addresses[0] = address + int(index)
-        npu_fm.region = _REGION_MAP[buffer_type]
+        npu_fm.region = region
         return npu_fm
 
     def replace_npu_address_range_with_address(npu_addr_range):
         assert isinstance(npu_addr_range.address, tvm.tir.Load)
         buffer = npu_addr_range.address.buffer_var
+        index = int(
+            npu_addr_range.address.index * (np.iinfo(np.dtype(npu_addr_range.address)).bits // 8)
+        )
+        if buffer in scratch_region_map.keys():
+            return vapi.NpuAddressRange(
+                scratch_region_map[buffer].region,
+                scratch_region_map[buffer].offset + index,
+                npu_addr_range.length,
+            )
         assert buffer in buffer_addresses.keys(), f"searching for buffer : {buffer}, but not found"
         address, buffer_type = buffer_addresses[buffer]
-        return vapi.NpuAddressRange(_REGION_MAP[buffer_type], address, npu_addr_range.length)
+        return vapi.NpuAddressRange(_get_region(buffer_type), address, npu_addr_range.length)
 
     def replace_tir_loads(npu_object):
         if isinstance(npu_object, vapi.NpuFeatureMap):
@@ -325,7 +455,6 @@ def classify_io(buffer):
 
         raise ValueError(f"Unused IO : {buffer} in tir module.")
 
-    scratch_size = 0
     constant_hex_data = []
     total_constant_len = 0
     buffer_addresses = dict()
@@ -345,8 +474,10 @@ def classify_io(buffer):
             constant_hex_data.append(constant_tensor)
             total_constant_len += len(constant_tensor) // 2
         else:
-            if info.btype == BufferType.input_or_output:
-                buffer_type = classify_io(_buffer)
+            if info.btype == BufferType.input_or_output or info.btype == BufferType.input:
+                buffer_type = info.btype
+                if info.btype == BufferType.input_or_output:
+                    buffer_type = classify_io(_buffer)
                 assert buffer_type in (BufferType.input, BufferType.output)
                 address = 0
                 buffer_addresses[_buffer] = (address, buffer_type)
@@ -359,14 +490,8 @@ def classify_io(buffer):
                 address = arch_config.lut_start_address
                 buffer_addresses[_buffer] = (address, info.btype)
             else:
-                dtype_bytes = np.iinfo(np.dtype(info.dtype)).bits // 8
-                size_in_bytes = int(dtype_bytes * np.prod(list(info.shape)))
-                # Every memory address the NPU access have to be 16 byte aligned
-                size_in_bytes = util.round_up(size_in_bytes, 16)
+                # These buffer_vars are already updated in scratch_region_map
                 assert info.btype == BufferType.scratch
-                address = scratch_size
-                scratch_size += size_in_bytes
-                buffer_addresses[_buffer] = (address, info.btype)
 
     for npu_op in npu_ops:
         for attr_name, attr in npu_op.__dict__.items():
@@ -379,11 +504,7 @@ def classify_io(buffer):
                 setattr(npu_op, attr_name, replace_tir_loads(attr))
 
     constant_data = "".join(constant_hex_data)
-    return (
-        npu_ops,
-        constant_data,
-        scratch_size,
-    )
+    return (npu_ops, constant_data)
 
 
 def translate_ethosu_tir_call_extern(tir_call_extern):
@@ -733,17 +854,18 @@ def _create_npu_rounding_mode(
 def _create_npu_dma_op(serial_copy):
     """This is a helper function to capture the list of arguments
     to create a NpuDmaOperation object"""
+    data_type_bytes = np.iinfo(np.dtype(serial_copy.read_address.dtype)).bits // 8
     src = vapi.NpuAddressRange(
         # region will be updated later
         region=0,
         address=serial_copy.read_address,
-        length=int(serial_copy.length.value),
+        length=int(serial_copy.length.value) * data_type_bytes,
     )
     dest = vapi.NpuAddressRange(
         # region will be updated later
         region=0,
         address=serial_copy.write_address,
-        length=int(serial_copy.length.value),
+        length=int(serial_copy.length.value) * data_type_bytes,
     )
     return vapi.NpuDmaOperation(src, dest)
 
diff --git a/python/tvm/relay/backend/executor_factory.py b/python/tvm/relay/backend/executor_factory.py
index 9ff7a7a8120b..eee3169400ff 100644
--- a/python/tvm/relay/backend/executor_factory.py
+++ b/python/tvm/relay/backend/executor_factory.py
@@ -109,6 +109,13 @@ def __init__(
         executor_codegen_metadata,
         devices,
     ):
+        fcreate = get_global_func("tvm.aot_executor_factory.create")
+        args = []
+        for k, v in params.items():
+            args.append(k)
+            args.append(ndarray.array(v))
+
+        self.module = fcreate(libmod, libmod_name, *args)
         self.ir_mod = ir_mod
         self.lowered_ir_mods = lowered_ir_mods
         self.target = target
@@ -134,6 +141,9 @@ def get_executor_config(self):
     def get_lib(self):
         return self.lib
 
+    def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
+        return self.module.export_library(file_name, fcompile, addons, **kwargs)
+
 
 class GraphExecutorFactoryModule(ExecutorFactoryModule):
     """Graph executor factory module.
diff --git a/python/tvm/relay/backend/te_compiler.py b/python/tvm/relay/backend/te_compiler.py
index def827c74131..654c8f66acf5 100644
--- a/python/tvm/relay/backend/te_compiler.py
+++ b/python/tvm/relay/backend/te_compiler.py
@@ -171,6 +171,8 @@ def select_implementation(op, attrs, inputs, out_type, target, use_autotvm=True)
         The best op implementation and the corresponding output tensors.
     """
     all_impls = get_valid_implementations(op, attrs, inputs, out_type, target)
+    if len(all_impls) == 0:
+        raise RuntimeError(f"No valid {op} implementations for {target}")
     best_plevel_impl = max(all_impls, key=lambda x: x.plevel)
 
     # Disable autotvm if auto_scheduler is enabled.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
index db3503dd9c9e..d529aec461eb 100644
--- a/python/tvm/relay/frontend/onnx.py
+++ b/python/tvm/relay/frontend/onnx.py
@@ -216,6 +216,15 @@ def get_scalar(x, params, dtype="float32"):
     return _op.cast(x, dtype)
 
 
+def get_scalar_or_1d_tensor(x, params, dtype="float32"):
+    """Helper to get a scalar value or 1D tensor for Quantized operators."""
+    if isinstance(x, _expr.Var) and x.name_hint in params:
+        return _op.const(params[x.name_hint].numpy(), dtype)
+    rank = len(infer_shape(x))
+    assert rank <= 1, "scale and zero_point input must be scalars or 1D tensors"
+    return _op.cast(x, dtype)
+
+
 def matmul_out_dtype(inputs, out_dtype):
     """Common function to handle MatMul and MatMulInteger16"""
     a_shape = shape_of(inputs[0])
@@ -264,23 +273,31 @@ def flatten_to_nd(x, x_shape, nd=3):
             b = _op.transpose(inputs[1])
             output = _op.nn.dense(a, b, out_dtype=out_dtype)
         else:
+            a = inputs[0]
+            b = inputs[1]
             # broadcast a and b
-            a_broadcasted_shape = _op.concatenate(
-                [
-                    out_batch,
-                    _op.strided_slice(a_shape, [a_rank - 2], [a_rank]),
-                ],
-                0,
+            a_broadcasted_shape = fold_constant(
+                _op.concatenate(
+                    [
+                        out_batch,
+                        _op.strided_slice(a_shape, [a_rank - 2], [a_rank]),
+                    ],
+                    0,
+                )
             )
-            b_broadcasted_shape = _op.concatenate(
-                [
-                    out_batch,
-                    _op.strided_slice(b_shape, [b_rank - 2], [b_rank]),
-                ],
-                0,
+            b_broadcasted_shape = fold_constant(
+                _op.concatenate(
+                    [
+                        out_batch,
+                        _op.strided_slice(b_shape, [b_rank - 2], [b_rank]),
+                    ],
+                    0,
+                )
             )
-            a = _op.transform.broadcast_to(inputs[0], fold_constant(a_broadcasted_shape))
-            b = _op.transform.broadcast_to(inputs[1], fold_constant(b_broadcasted_shape))
+            if not tvm.ir.structural_equal(a_shape, a_broadcasted_shape):
+                a = _op.transform.broadcast_to(a, a_broadcasted_shape)
+            if not tvm.ir.structural_equal(b_shape, b_broadcasted_shape):
+                b = _op.transform.broadcast_to(b, b_broadcasted_shape)
             # Convert a and b into 3 dimensional tensors.
             a = flatten_to_nd(a, shape_of(a), 3)
             b = flatten_to_nd(b, shape_of(b), 3)
@@ -3646,11 +3663,21 @@ def _impl_v10(cls, inputs, attr, params):
         x_scale = get_scalar(inputs[1], params)
         x_zero_point = get_scalar(inputs[2], params, "int32")
         weight = inputs[3]
-        w_scale = get_scalar(inputs[4], params)
-        w_zero_point = get_scalar(inputs[5], params, "int32")
+        w_scale = get_scalar_or_1d_tensor(inputs[4], params)
+        w_zero_point = get_scalar_or_1d_tensor(inputs[5], params, "int32")
         y_scale = fold_constant(get_scalar(inputs[6], params))
         y_zero_point = get_scalar(inputs[7], params, "int32")
 
+        # Check shapes for per channel quantization
+        w_scale_shape = infer_shape(w_scale)
+        w_zero_point_shape = infer_shape(w_zero_point)
+        if len(w_scale_shape) == 1 or len(w_zero_point_shape) == 1:
+            m = infer_shape(weight)[0]
+            if m != w_scale_shape[0] or m != w_zero_point_shape[0]:
+                raise tvm.error.OpAttributeInvalid(
+                    "The number of elements should be equal to the number of output channels"
+                )
+
         input_shape = infer_shape(data)
 
         ndim = len(input_shape)
@@ -3723,11 +3750,11 @@ def _impl_v10(cls, inputs, attr, params):
                 y_scale,
                 y_zero_point,
                 out_dtype=out_dtype,
-                axis=0,
+                axis=1,
             )
         else:
-            out = _qnn.op.dequantize(out, requantize_scale, _op.const(0, dtype="int32"), axis=0)
-            out = _qnn.op.quantize(out, y_scale, y_zero_point, axis=0, out_dtype=out_dtype)
+            out = _qnn.op.dequantize(out, requantize_scale, _op.const(0, dtype="int32"), axis=1)
+            out = _qnn.op.quantize(out, y_scale, y_zero_point, axis=1, out_dtype=out_dtype)
         return out
 
 
@@ -4830,7 +4857,32 @@ def from_onnx(self, graph, opset, get_output_expr=False):
             A dict of name: tvm.nd.array pairs, used as pretrained weights
         """
         self.opset = opset
-        # parse network inputs to relay, aka parameters
+        self._parse_graph_initializers(graph)
+        self._parse_graph_input(graph)
+        self._check_user_inputs_in_outermost_graph_scope()
+        self._check_for_unsupported_ops(graph)
+        self._construct_nodes(graph)
+
+        # now return the outputs
+        outputs = [self._nodes[self._parse_value_proto(i)] for i in graph.output]
+        outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
+        # If requested, directly return the converted expressions.
+        if get_output_expr:
+            return outputs
+        ## Maintain the order of inputs and parameters from the ONNX graph, but only include
+        ## those parameters that are needed to execute the relay graph
+        free_vars = analysis.free_vars(outputs)
+        nodes = {v: k for k, v in self._nodes.items()}
+        free_vars = [nodes[var] for var in free_vars]
+        for i_name in self._params:
+            if i_name in free_vars and i_name not in self._inputs:
+                self._inputs[i_name] = self._nodes[i_name]
+        # Create a function from our output expression and all input variables.
+        func = _function.Function([v for k, v in self._inputs.items()], outputs)
+        return IRModule.from_expr(func), self._params
+
+    def _parse_graph_initializers(self, graph):
+        """Parse network inputs to relay, aka parameters."""
         for init_tensor in graph.initializer:
             if not init_tensor.name.strip():
                 raise ValueError("Tensor's name is required.")
@@ -4844,6 +4896,8 @@ def from_onnx(self, graph, opset, get_output_expr=False):
                     shape=self._params[init_tensor.name].shape,
                     dtype=self._params[init_tensor.name].dtype,
                 )
+
+    def _parse_graph_input(self, graph):
         for i in graph.input:
             # from onnx v0.2, GraphProto.input has type ValueInfoProto,
             #  and the name is 'i.name'
@@ -4851,7 +4905,6 @@ def from_onnx(self, graph, opset, get_output_expr=False):
             if i_name in self._params:
                 # i is a param instead of input
                 self._num_param += 1
-                self._params[i_name] = self._params.pop(i_name)
                 self._nodes[i_name] = new_var(
                     i_name, shape=self._params[i_name].shape, dtype=self._params[i_name].dtype
                 )
@@ -4876,15 +4929,18 @@ def from_onnx(self, graph, opset, get_output_expr=False):
                     dtype = d_type
                 self._nodes[i_name] = new_var(i_name, shape=i_shape, dtype=dtype)
             self._inputs[i_name] = self._nodes[i_name]
-        # Only check user inputs in the outer-most graph scope.
+
+    def _check_user_inputs_in_outermost_graph_scope(self):
+        """Only check user inputs in the outer-most graph scope."""
         if self._old_manager is None:
             assert all(
                 [name in self._input_names for name in self._shape.keys()]
             ), "User specified the shape for inputs that weren't found in the graph: " + str(
                 self._shape
             )
-        # get list of unsupported ops
-        convert_map = _get_convert_map(opset)
+
+    def _check_for_unsupported_ops(self, graph):
+        convert_map = _get_convert_map(self.opset)
         unsupported_ops = set()
         for node in graph.node:
             op_name = node.op_type
@@ -4898,7 +4954,9 @@ def from_onnx(self, graph, opset, get_output_expr=False):
             msg = "The following operators are not supported for frontend ONNX: "
             msg += ", ".join(unsupported_ops)
             raise tvm.error.OpNotImplemented(msg)
-        # construct nodes, nodes are stored as directed acyclic graph
+
+    def _construct_nodes(self, graph):
+        """Nodes are stored as directed acyclic graph."""
         for node in graph.node:
             op_name = node.op_type
             attr = self._parse_attr(node.attribute)
@@ -4915,7 +4973,7 @@ def from_onnx(self, graph, opset, get_output_expr=False):
             attr["tvm_custom"]["name"] = i_name
             attr["tvm_custom"]["num_outputs"] = len(node_output)
 
-            op = self._convert_operator(op_name, inputs, attr, opset)
+            op = self._convert_operator(op_name, inputs, attr, self.opset)
             if not isinstance(op, _expr.TupleWrapper):
                 outputs_num = 1
             else:
@@ -4964,24 +5022,6 @@ def from_onnx(self, graph, opset, get_output_expr=False):
                 for k, i in zip(list(node_output), range(len(node_output))):
                     self._nodes[k] = op[i]
 
-        # now return the outputs
-        outputs = [self._nodes[self._parse_value_proto(i)] for i in graph.output]
-        outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
-        # If requested, directly return the converted expressions.
-        if get_output_expr:
-            return outputs
-        ## Maintain the order of inputs and parameters from the ONNX graph, but only include
-        ## those parameters that are needed to execute the relay graph
-        free_vars = analysis.free_vars(outputs)
-        nodes = {v: k for k, v in self._nodes.items()}
-        free_vars = [nodes[var] for var in free_vars]
-        for i_name in self._params:
-            if i_name in free_vars and i_name not in self._inputs:
-                self._inputs[i_name] = self._nodes[i_name]
-        # Create a function from our output expression and all input variables.
-        func = _function.Function([v for k, v in self._inputs.items()], outputs)
-        return IRModule.from_expr(func), self._params
-
     def _parse_value_proto(self, value_proto):
         """Parse ValueProto or raw str."""
         try:
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
index 3dcabe51cc59..ee75220ed392 100644
--- a/python/tvm/relay/frontend/pytorch.py
+++ b/python/tvm/relay/frontend/pytorch.py
@@ -972,19 +972,21 @@ def convolution(self, inputs, input_types):
             msg = "Data type %s could not be parsed in conv op" % (type(weight))
             raise AssertionError(msg)
 
-        # Transposed convolutions have IOHW layout.
-        if use_transpose:
-            weight_shape[0], weight_shape[1] = weight_shape[1], weight_shape[0]
-
-        channels = weight_shape[0]
         groups = int(inputs[8])
 
+        if use_transpose:
+            channels = weight_shape[1] * groups
+            in_channels = weight_shape[0]
+        else:
+            channels = weight_shape[0]
+            in_channels = weight_shape[1]
+
         # Check if this is depth wise convolution
         # We need to reshape weight so that Relay could recognize this is depth wise
         # weight_shape[1] is always in_channels // groups
         # For depthwise, in_channels == groups, so weight_shape[1] == 1
         # If groups > 1 but weight_shape[1] != 1, this is group convolution
-        if groups > 1 and weight_shape[1] == 1:
+        if groups > 1 and in_channels == 1:
             channel_multiplier = channels // groups
             new_weight_shape = (groups, channel_multiplier) + tuple(weight_shape[2:])
             weight = _op.transform.reshape(weight, new_weight_shape)
diff --git a/python/tvm/relay/frontend/qnn_torch.py b/python/tvm/relay/frontend/qnn_torch.py
index 83af057e8b5d..6a6dc467ab14 100644
--- a/python/tvm/relay/frontend/qnn_torch.py
+++ b/python/tvm/relay/frontend/qnn_torch.py
@@ -526,6 +526,7 @@ def inline_input_quant_params_for_fx(graph, params):
     %2403 : float = prim::Constant[value=1.]()
     %quantize_per_tensor_2 ...  = aten::quantize_per_tensor(..., %2403, %2402, ...)
     """
+    # pylint: disable=c-extension-no-member
     import torch
 
     def get_full_attr_name(current):
diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py
index b675dd56a7bb..adfa0a56f639 100644
--- a/python/tvm/relay/frontend/tflite.py
+++ b/python/tvm/relay/frontend/tflite.py
@@ -30,6 +30,7 @@
 from .. import function as _function
 from .. import op as _op
 from .. import qnn as _qnn
+from ..backend.name_transforms import sanitize_name
 from .common import ExprTable
 from .common import infer_shape as _infer_shape
 from .common import to_int_list
@@ -1666,7 +1667,7 @@ def _transform_mask(stride_dim, ellipsis_mask):
                             if begin[index] < 0
                             else begin[index]
                         )
-                        m_end[final_index] = begin[index] + 1
+                        m_end[final_index] = m_begin[final_index] + 1
                         m_stride[final_index] = 1
                         fshape_indices.append(-2)
                     else:
@@ -2704,9 +2705,9 @@ def convert_unpack(self, op):
         unpack_axis = unpack_options.Axis()
 
         # Relay doesn't support 'unpack' operator so we use 'split' & 'squeeze' instead.
-        # We have to do 'squeeze' along the split axis but Relay expects
-        # squeeze_axis to be either None or List.
-        squeeze_axis = None if unpack_axis == 0 else [unpack_axis]
+        # We have to do 'squeeze' along the split axis.
+        # Relay expects squeeze_axis to be List.
+        squeeze_axis = [unpack_axis]
 
         # Relay doesn't like TupleWrapper of 1 element so we isolate the case of unpacking
         # a tensor by an axis with len(axis) == 1. For reference see convert_split().
@@ -3769,6 +3770,15 @@ def from_tflite(model, shape_dict=None, dtype_dict=None, op_converter=OperatorCo
     params = {k: _nd.array(np.array(v)) for k, v in exp_tab.params.items()}
     outputs = [exp_tab.get_expr(get_tensor_name(subgraph, i)) for i in model_outputs]
     outputs = outputs[0] if len(outputs) == 1 else _expr.Tuple(outputs)
-    func = _function.Function(analysis.free_vars(outputs), outputs)
+    attrs = tvm.ir.make_node(
+        "DictAttrs",
+        **{
+            "output_tensor_names": [
+                sanitize_name(get_tensor_name(subgraph, model_output))
+                for model_output in model_outputs
+            ]
+        },
+    )
+    func = _function.Function(analysis.free_vars(outputs), outputs, attrs=attrs)
     mod = IRModule.from_expr(func)
     return mod, params
diff --git a/python/tvm/relay/op/contrib/dnnl.py b/python/tvm/relay/op/contrib/dnnl.py
index 05b588051a1c..6d4fe0d81260 100644
--- a/python/tvm/relay/op/contrib/dnnl.py
+++ b/python/tvm/relay/op/contrib/dnnl.py
@@ -67,22 +67,39 @@ def _func_wrapper(expr):
 
 
 _register_external_op_helper("nn.batch_norm")
+_register_external_op_helper("nn.conv1d")
 _register_external_op_helper("nn.conv2d")
+_register_external_op_helper("nn.conv3d")
+_register_external_op_helper("nn.conv2d_transpose")
+_register_external_op_helper("nn.conv3d_transpose")
 _register_external_op_helper("nn.dense")
+_register_external_op_helper("nn.max_pool2d")
+_register_external_op_helper("nn.avg_pool2d")
+_register_external_op_helper("nn.max_pool3d")
+_register_external_op_helper("nn.avg_pool3d")
+_register_external_op_helper("abs")
+_register_external_op_helper("clip")
+_register_external_op_helper("exp")
+_register_external_op_helper("log")
+_register_external_op_helper("sqrt")
+_register_external_op_helper("round")
+_register_external_op_helper("logsumexp")
 _register_external_op_helper("nn.relu")
+_register_external_op_helper("nn.leaky_relu")
 _register_external_op_helper("tanh")
 _register_external_op_helper("sigmoid")
+_register_external_op_helper("nn.softmax")
 _register_external_op_helper("add")
 _register_external_op_helper("multiply")
 
 
-def make_conv_pattern(with_bias=True, with_eltwise=None):
-    """Create patterns related to nn.conv2d.
+def make_conv_pattern(conv_name, with_bias=True, with_eltwise=None):
+    """Create patterns related to conv and deconv.
 
     Parameters
     ----------
     with_bias : bool
-        Whether attach `bias_add` to `nn.conv2d`.
+        Whether attach `bias_add` to `conv / deconv`.
     with_eltwise : str
         The attached elementwise post-op name.
     Returns
@@ -93,7 +110,7 @@ def make_conv_pattern(with_bias=True, with_eltwise=None):
     data = wildcard()
     weight = wildcard()
     bias = wildcard()
-    conv = is_op("nn.conv2d")(data, weight)
+    conv = is_op(conv_name)(data, weight)
     if with_bias:
         conv_out = is_op("add")(conv, bias)
     else:
@@ -146,15 +163,19 @@ def make_dnnl_pattern(op, with_bias, with_eltwise):
     pattern : Tuple(pattern_name, CallPattern)
         Created pattern name, along with its CallPattern.
     """
-    pat_name = "dnnl." + op
+    pat_name = op.replace("nn", "dnnl")
     pat_name += "_bias" if with_bias else ""
     pat_name += ("_" + with_eltwise.split(".")[-1]) if with_eltwise else ""
-    if op == "conv2d":
-        dnnl_pattern = (pat_name, make_conv_pattern(with_bias, with_eltwise))
-    elif op == "dense":
+    if "conv" in op:
+        dnnl_pattern = (pat_name, make_conv_pattern(op, with_bias, with_eltwise))
+    elif op == "nn.dense":
         dnnl_pattern = (pat_name, make_dense_pattern(with_bias, with_eltwise))
     else:
-        logger.warning("Currently, only conv2d and dense op are supported, but got %s.", op)
+        logger.warning(
+            "Currently, only conv1d, conv2d, conv2d_transpose, conv3d_transpose and "
+            "dense op are supported, but got %s.",
+            op,
+        )
         dnnl_pattern = ()
     return dnnl_pattern
 
@@ -174,8 +195,15 @@ def pattern_table():
         for elt in elt_list:
             if not with_bias and not elt:
                 return dnnl_patterns
-            dnnl_patterns.append(make_dnnl_pattern("conv2d", with_bias, elt))
-            dnnl_patterns.append(make_dnnl_pattern("dense", with_bias, elt))
+            for conv_name in [
+                "nn.conv1d",
+                "nn.conv2d",
+                "nn.conv3d",
+                "nn.conv2d_transpose",
+                "nn.conv3d_transpose",
+            ]:
+                dnnl_patterns.append(make_dnnl_pattern(conv_name, with_bias, elt))
+            dnnl_patterns.append(make_dnnl_pattern("nn.dense", with_bias, elt))
     return dnnl_patterns
 
 
diff --git a/python/tvm/relay/op/contrib/ethosu.py b/python/tvm/relay/op/contrib/ethosu.py
index fa11fca88390..d83c5d38f429 100644
--- a/python/tvm/relay/op/contrib/ethosu.py
+++ b/python/tvm/relay/op/contrib/ethosu.py
@@ -1074,8 +1074,28 @@ class LutActivationParams:
     """
 
     def __init__(self, func_body: Call):
-        self.ofm = TensorParams(func_body)
-        self.ifm = TensorParams(func_body.args[0].args[0].args[0])
+        from tvm.relay.backend.contrib.ethosu.util import QuantizeArgs
+        from tvm.relay.backend.contrib.ethosu.util import DequantizeArgs
+
+        layout = "NHWC"
+
+        quantize = func_body
+        activation = quantize.args[0]
+        dequantize = activation.args[0]
+        in_var = dequantize.args[0]
+
+        self.ifm = TensorParams(
+            in_var,
+            layout=layout,
+            scale=dequantize.args[DequantizeArgs.IFM_SCALE.value],
+            zero_point=dequantize.args[DequantizeArgs.IFM_ZERO_POINT.value],
+        )
+        self.ofm = TensorParams(
+            quantize,
+            layout=layout,
+            scale=quantize.args[QuantizeArgs.OFM_SCALE.value],
+            zero_point=quantize.args[QuantizeArgs.OFM_ZERO_POINT.value],
+        )
 
     def is_valid(self):
         """
@@ -1116,6 +1136,28 @@ def sigmoid_pattern():
     return quant
 
 
+class LeakyReLUParams(LutActivationParams):
+    """
+    This class will parse a call to ethos-u.leaky_relu composite function
+    and extract the parameter information.
+    """
+
+    composite_name = "ethos-u.leaky_relu"
+
+    def __init__(self, func_body: Call):
+        super().__init__(func_body)
+        self.alpha = func_body.args[0].attrs.alpha
+
+
+def leaky_relu_pattern() -> tvm.relay.dataflow_pattern.DFPattern:
+    """
+    This function creates the pattern for leaky relu.
+    """
+    dequantize = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
+    leaky_relu = is_op("nn.leaky_relu")(dequantize)
+    return is_op("qnn.quantize")(leaky_relu, is_constant(), is_constant())
+
+
 class MeanParams:
     """
     This class will parse a call to ethosu.mean composite function
@@ -1574,6 +1616,11 @@ def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Cal
             mean_pattern(),
             lambda pat: MeanParams(pat).is_valid(),
         ),
+        (
+            LeakyReLUParams.composite_name,
+            leaky_relu_pattern(),
+            lambda pat: LeakyReLUParams(pat).is_valid(),
+        ),
         (ConcatParams.composite_name, concat_pattern(), lambda pat: ConcatParams(pat).is_valid()),
         (
             SigmoidParams.composite_name,
diff --git a/python/tvm/relay/op/strategy/arm_cpu.py b/python/tvm/relay/op/strategy/arm_cpu.py
index 35db043e8c93..84b08dc4d9b7 100644
--- a/python/tvm/relay/op/strategy/arm_cpu.py
+++ b/python/tvm/relay/op/strategy/arm_cpu.py
@@ -94,12 +94,18 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target):
                     name="conv2d_nchw_spatial_pack.arm_cpu",
                 )
 
-                # Intel x86 conv2d schedule.
-                strategy.add_implementation(
-                    wrap_compute_conv2d(topi.x86.conv2d_nchw),
-                    wrap_topi_schedule(topi.x86.schedule_conv2d_nchw),
-                    name="conv2d_nchw.x86",
-                )
+                if topi.arm_cpu.is_int8_hw_support(data.dtype, kernel.dtype):
+                    strategy.add_implementation(
+                        wrap_compute_conv2d(topi.arm_cpu.conv2d_nchw_int8),
+                        wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_nchw_int8),
+                        name="conv2d_nchw_int8.arm_cpu",
+                    )
+                else:
+                    strategy.add_implementation(
+                        wrap_compute_conv2d(topi.x86.conv2d_nchw),
+                        wrap_topi_schedule(topi.x86.schedule_conv2d_nchw),
+                        name="conv2d_nchw.x86",
+                    )
 
                 # check if winograd algorithm is applicable
                 _, _, kh, kw = get_const_tuple(kernel.shape)
@@ -256,11 +262,19 @@ def conv2d_strategy_arm_cpu(attrs, inputs, out_type, target):
 def conv2d_NCHWc_strategy_arm_cpu(attrs, inputs, out_type, target):
     """conv2d_NCHWc adopted from x86"""
     strategy = _op.OpStrategy()
-    strategy.add_implementation(
-        wrap_compute_conv2d(topi.x86.conv2d_NCHWc, True, True),
-        wrap_topi_schedule(topi.x86.schedule_conv2d_NCHWc),
-        name="conv2d_NCHWc.x86",
-    )
+    data, kernel = inputs
+    if topi.arm_cpu.is_int8_hw_support(data.dtype, kernel.dtype):
+        strategy.add_implementation(
+            wrap_compute_conv2d(topi.arm_cpu.conv2d_NCHWc_int8, True, True),
+            wrap_topi_schedule(topi.arm_cpu.schedule_conv2d_NCHWc_int8),
+            name="conv2d_NCHWc_int8.arm_cpu",
+        )
+    else:
+        strategy.add_implementation(
+            wrap_compute_conv2d(topi.x86.conv2d_NCHWc, True, True),
+            wrap_topi_schedule(topi.x86.schedule_conv2d_NCHWc),
+            name="conv2d_NCHWc.x86",
+        )
     return strategy
 
 
diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py
index 730c3b4357ed..2d5b35e03ed2 100644
--- a/python/tvm/relay/op/strategy/cuda.py
+++ b/python/tvm/relay/op/strategy/cuda.py
@@ -666,6 +666,7 @@ def conv3d_strategy_cuda(attrs, inputs, out_type, target):
             and stride_w == 1
             and dilation_h == 1
             and dilation_w == 1
+            and attrs["groups"] == 1
         ):
             strategy.add_implementation(
                 wrap_compute_conv3d(topi.cuda.conv3d_ncdhw_winograd),
@@ -688,7 +689,7 @@ def conv3d_strategy_cuda(attrs, inputs, out_type, target):
                     (N % 16 == 0 and CI % 16 == 0 and CO % 16 == 0)
                     or (N % 8 == 0 and CI % 16 == 0 and CO % 32 == 0)
                     or (N % 32 == 0 and CI % 16 == 0 and CO % 8 == 0)
-                ):
+                ) and out_type == "float16":
                     strategy.add_implementation(
                         wrap_compute_conv3d(topi.cuda.conv3d_ndhwc_tensorcore),
                         wrap_topi_schedule(topi.cuda.schedule_conv3d_ndhwc_tensorcore),
@@ -921,7 +922,7 @@ def batch_matmul_strategy_cuda(attrs, inputs, out_type, target):
             wrap_compute_batch_matmul(topi.cuda.batch_matmul_cublas),
             wrap_topi_schedule(topi.generic.schedule_extern),
             name="batch_matmul_cublas.cuda",
-            plevel=15,
+            plevel=30,
         )
     if (
         target.kind.name == "cuda"
diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py
index 320401191530..0dbf0825edd4 100644
--- a/python/tvm/relay/op/strategy/generic.py
+++ b/python/tvm/relay/op/strategy/generic.py
@@ -545,10 +545,8 @@ def _compute_conv3d(attrs, inputs, out_type):
         (dilation_d, dilation_h, dilation_w) = dilation
         if dilation_d < 1 or dilation_h < 1 or dilation_w < 1:
             raise ValueError("Dilation should be positive value")
-        if groups != 1:
-            raise ValueError("Not support arbitrary group number for conv3d")
 
-        args = [inputs[0], inputs[1], strides, padding, dilation]
+        args = [inputs[0], inputs[1], strides, padding, dilation, groups]
         if need_layout:
             args.append(layout)
         args.append(out_dtype)
diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py
index 20b674966131..33e68627f656 100644
--- a/python/tvm/relay/op/strategy/x86.py
+++ b/python/tvm/relay/op/strategy/x86.py
@@ -530,11 +530,27 @@ def dense_strategy_cpu(attrs, inputs, out_type, target):
 def dense_pack_strategy_cpu(attrs, inputs, out_type, target):
     """dense_pack x86 strategy"""
     strategy = _op.OpStrategy()
-    strategy.add_implementation(
-        wrap_compute_dense(topi.x86.dense_pack),
-        wrap_topi_schedule(topi.x86.schedule_dense_pack),
-        name="dense_pack.x86",
-    )
+
+    if (
+        inputs[0].dtype == "uint8"
+        and inputs[1].dtype == "int8"
+        and out_type.dtype == "int32"
+        and attrs["weight_layout"] == "NC16n4c"
+    ):
+        strategy.add_implementation(
+            wrap_compute_dense(topi.x86.dense_vnni),
+            wrap_topi_schedule(topi.x86.schedule_dense_vnni),
+            name="dense_vnni.x86",
+            plevel=12,
+        )
+    else:
+        strategy.add_implementation(
+            wrap_compute_dense(topi.x86.dense_pack),
+            wrap_topi_schedule(topi.x86.schedule_dense_pack),
+            name="dense_pack.x86",
+            plevel=10,
+        )
+
     return strategy
 
 
diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py
index 68f684e3276a..27dfefbb7890 100644
--- a/python/tvm/relay/op/transform.py
+++ b/python/tvm/relay/op/transform.py
@@ -704,7 +704,6 @@ def tile(data, reps):
 
 def reverse(data, axis):
     """Reverses the order of elements along given axis while preserving array shape.
-    By default, repeat flattens the input array into 1-D and then repeats the elements.
 
     Parameters
     ----------
diff --git a/python/tvm/relay/qnn/op/legalizations.py b/python/tvm/relay/qnn/op/legalizations.py
index fd835d72fc09..11677315ac86 100644
--- a/python/tvm/relay/qnn/op/legalizations.py
+++ b/python/tvm/relay/qnn/op/legalizations.py
@@ -17,6 +17,7 @@
 # pylint: disable=invalid-name, unused-argument
 """Backend QNN related feature registration"""
 import numpy as np
+from scipy import special
 import tvm
 from tvm import relay
 from tvm._ffi.base import TVMError
@@ -47,19 +48,30 @@ def legalize_qnn_dense(attrs, inputs, types):
     return qnn_dense_legalize(attrs, inputs, types)
 
 
-# Registering QNN dense legalization function.
-@reg.register_qnn_legalize("qnn.rsqrt")
-def legalize_qnn_rsqrt(attrs, inputs, types):
-    return create_integer_lookup_op(
-        input_arg=inputs[0],
-        floating_point_func=lambda arr: 1 / np.sqrt(arr),
-        in_scale=inputs[1],
-        in_zero_point=inputs[2],
-        out_scale=inputs[3],
-        out_zero_point=inputs[4],
-        in_dtype=types[0].dtype,
-        out_dtype=types[0].dtype,
-    )
+def register_qnn_unary_op_legalize(op_name, floating_point_func):
+    """Register unary qnn op for legalization via table lookup op."""
+
+    def legalize_qnn_unary_op(attrs, inputs, types):
+        return create_integer_lookup_op(
+            input_arg=inputs[0],
+            floating_point_func=floating_point_func,
+            in_scale=inputs[1],
+            in_zero_point=inputs[2],
+            out_scale=inputs[3],
+            out_zero_point=inputs[4],
+            in_dtype=types[0].dtype,
+            out_dtype=types[0].dtype,
+        )
+
+    return reg.register_qnn_legalize(op_name, legalize_qnn_unary_op)
+
+
+register_qnn_unary_op_legalize("qnn.sqrt", np.sqrt)
+register_qnn_unary_op_legalize("qnn.rsqrt", lambda arr: 1 / np.sqrt(arr))
+register_qnn_unary_op_legalize("qnn.exp", np.exp)
+register_qnn_unary_op_legalize("qnn.erf", special.erf)
+register_qnn_unary_op_legalize("qnn.sigmoid", lambda arr: 1 / (1 + np.exp(-arr)))
+register_qnn_unary_op_legalize("qnn.tanh", np.tanh)
 
 
 # Default to None. If overridden by target, this will not be run.
@@ -183,7 +195,7 @@ def helper_no_fast_int8_hw_legalization(attrs, inputs, types, relay_op):
 
         shift_kernel = relay.nn.bias_add(
             relay.cast(kernel, dtype="int16"),
-            relay.cast(kernel_zero_point, dtype="int16"),
+            -relay.cast(kernel_zero_point, dtype="int16"),
             output_axis,
         )
     new_attrs = {k: attrs[k] for k in attrs.keys()}
diff --git a/python/tvm/relay/qnn/op/qnn.py b/python/tvm/relay/qnn/op/qnn.py
index 85629a9b5a5a..d8635a1c08d7 100644
--- a/python/tvm/relay/qnn/op/qnn.py
+++ b/python/tvm/relay/qnn/op/qnn.py
@@ -22,16 +22,16 @@
 import tvm
 import tvm.ir
 from tvm import relay
-from tvm.runtime import Object
 from tvm.relay.expr import Tuple, TupleWrapper
 from tvm.relay.op.nn.utils import get_pad_tuple2d
-from tvm.topi.nn.qnn import SQNN_DTYPE_TO_CODE
+from tvm.runtime import Object
 from tvm.target import Target
+from tvm.topi.nn.qnn import SQNN_DTYPE_TO_CODE
 from tvm.topi.x86.utils import target_has_sse41
+
 from ... import op as reg
 from ...op import OpPattern
-from . import _make
-from . import _requantize
+from . import _make, _requantize
 
 
 @tvm._ffi.register_object("relay.qnn.op.RequantizeConfig")
@@ -750,6 +750,111 @@ def mul(
     )
 
 
+def tanh(x, scale, zero_point, output_scale, output_zero_point):
+    """Quantized tanh.
+
+    Parameters
+    ----------
+    x : relay.Expr
+        The quantized input tensor.
+
+    scale: relay.Expr
+        The scale of the quantized expr.
+
+    zero_point: relay.Expr
+       The zero point of quantized expr.
+
+    output_scale: relay.Expr
+        The scale of the output quantized expr.
+
+    output_zero_point: relay.Expr
+       The zero point of output quantized expr.
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+
+    """
+    return _make.tanh(
+        x,
+        scale,
+        zero_point,
+        output_scale,
+        output_zero_point,
+    )
+
+
+def exp(x, scale, zero_point, output_scale, output_zero_point):
+    """Quantized exponential function.
+
+    Parameters
+    ----------
+    x : relay.Expr
+        The quantized input tensor.
+
+    scale: relay.Expr
+        The scale of the quantized expr.
+
+    zero_point: relay.Expr
+       The zero point of quantized expr.
+
+    output_scale: relay.Expr
+        The scale of the output quantized expr.
+
+    output_zero_point: relay.Expr
+       The zero point of output quantized expr.
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+
+    """
+    return _make.exp(
+        x,
+        scale,
+        zero_point,
+        output_scale,
+        output_zero_point,
+    )
+
+
+def sqrt(x, scale, zero_point, output_scale, output_zero_point):
+    """Quantized square root.
+
+    Parameters
+    ----------
+    x : relay.Expr
+        The quantized input tensor.
+
+    scale: relay.Expr
+        The scale of the quantized expr.
+
+    zero_point: relay.Expr
+       The zero point of quantized expr.
+
+    output_scale: relay.Expr
+        The scale of the output quantized expr.
+
+    output_zero_point: relay.Expr
+       The zero point of output quantized expr.
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+
+    """
+    return _make.sqrt(
+        x,
+        scale,
+        zero_point,
+        output_scale,
+        output_zero_point,
+    )
+
+
 def rsqrt(x, scale, zero_point, output_scale, output_zero_point):
     """Quantized reciprocal square root.
 
@@ -785,6 +890,76 @@ def rsqrt(x, scale, zero_point, output_scale, output_zero_point):
     )
 
 
+def erf(x, scale, zero_point, output_scale, output_zero_point):
+    """Quantized error function.
+
+    Parameters
+    ----------
+    x : relay.Expr
+        The quantized input tensor.
+
+    scale: relay.Expr
+        The scale of the quantized expr.
+
+    zero_point: relay.Expr
+       The zero point of quantized expr.
+
+    output_scale: relay.Expr
+        The scale of the output quantized expr.
+
+    output_zero_point: relay.Expr
+       The zero point of output quantized expr.
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+
+    """
+    return _make.erf(
+        x,
+        scale,
+        zero_point,
+        output_scale,
+        output_zero_point,
+    )
+
+
+def sigmoid(x, scale, zero_point, output_scale, output_zero_point):
+    """Quantized sigmoid.
+
+    Parameters
+    ----------
+    x : relay.Expr
+        The quantized input tensor.
+
+    scale: relay.Expr
+        The scale of the quantized expr.
+
+    zero_point: relay.Expr
+       The zero point of quantized expr.
+
+    output_scale: relay.Expr
+        The scale of the output quantized expr.
+
+    output_zero_point: relay.Expr
+       The zero point of output quantized expr.
+
+    Returns
+    -------
+    result : relay.Expr
+        The computed result.
+
+    """
+    return _make.sigmoid(
+        x,
+        scale,
+        zero_point,
+        output_scale,
+        output_zero_point,
+    )
+
+
 def subtract(
     lhs, rhs, lhs_scale, lhs_zero_point, rhs_scale, rhs_zero_point, output_scale, output_zero_point
 ):
diff --git a/python/tvm/relay/transform/fake_quantization_to_integer.py b/python/tvm/relay/transform/fake_quantization_to_integer.py
index e84ba5557a70..285ff3816582 100644
--- a/python/tvm/relay/transform/fake_quantization_to_integer.py
+++ b/python/tvm/relay/transform/fake_quantization_to_integer.py
@@ -130,20 +130,14 @@ def global_avgpool2d(expr, type_map):
     return [out, t]
 
 
-@register_fake_quantization_to_integer("rsqrt")
-def rsqrt(expr, type_map):
-    """Rewrite a rsqrt op"""
+@register_fake_quantization_to_integer("broadcast_to")
+def broadcast_to(expr, type_map):
+    """Rewrite a broadcast_to op"""
     arg = expr.args[0]
-    x_t = type_map[arg]
-    out_t = type_map[expr]
-    out = relay.qnn.op.rsqrt(
-        arg,
-        x_t.scale,
-        x_t.zero_point,
-        out_t.scale,
-        out_t.zero_point,
-    )
-    return [out, x_t]
+    t = type_map[arg]
+    shape = expr.attrs.shape
+    out = relay.op.broadcast_to(arg, shape)
+    return [out, t]
 
 
 @register_fake_quantization_to_integer("nn.bias_add")
@@ -471,3 +465,30 @@ def binary(expr, type_map):
 
 register_binary_identity("minimum", relay.op.minimum)
 register_binary_identity("maximum", relay.op.maximum)
+
+
+def register_unary_qnn(op_name, op):
+    """Rewrite a unary op"""
+
+    def unary(expr, type_map):
+        arg = expr.args[0]
+        x_t = type_map[arg]
+        out_t = type_map[expr]
+        out = op(
+            arg,
+            x_t.scale,
+            x_t.zero_point,
+            out_t.scale,
+            out_t.zero_point,
+        )
+        return [out, x_t]
+
+    return register_fake_quantization_to_integer(op_name, unary)
+
+
+register_unary_qnn("sqrt", relay.qnn.op.sqrt)
+register_unary_qnn("rsqrt", relay.qnn.op.rsqrt)
+register_unary_qnn("exp", relay.qnn.op.exp)
+register_unary_qnn("erf", relay.qnn.op.erf)
+register_unary_qnn("sigmoid", relay.qnn.op.sigmoid)
+register_unary_qnn("tanh", relay.qnn.op.tanh)
diff --git a/python/tvm/relay/transform/transform.py b/python/tvm/relay/transform/transform.py
index 14f1caae06eb..99c61c5bd96f 100644
--- a/python/tvm/relay/transform/transform.py
+++ b/python/tvm/relay/transform/transform.py
@@ -1238,7 +1238,7 @@ def AnnotateSpans():
     return _ffi_api.AnnotateSpans()
 
 
-def FakeQuantizationToInteger(hard_fail=False):
+def FakeQuantizationToInteger(hard_fail=False, use_qat=False):
     # pylint: disable=anomalous-backslash-in-string
     """
     Find regions of the graph of the form
@@ -1267,12 +1267,30 @@ def FakeQuantizationToInteger(hard_fail=False):
         If true, raise an error.
         If false, skip rewriting the subgraph.
 
+    use_qat : boolean
+        To perform an additional QAT pass - convert enabled operations with dequantized inputs.
+        Example: in the graph above op2 is not registered with the FakeQuantizationToInteger
+        attribute, op1 operation can still be converted. Converted pattern below:
+
+        .. code-block:: text
+
+            x    w
+            |    |
+            \\   /
+              op1
+              |
+              dq
+              |
+              op2
+              |
+              q
+
     Returns
     -------
     ret : tvm.transform.Pass
-        The registered SimplifyExpr pass.
+        The registered FakeQuantizationToInteger pass.
     """
-    return _ffi_api.FakeQuantizationToInteger(hard_fail)
+    return _ffi_api.FakeQuantizationToInteger(hard_fail, use_qat)
 
 
 def ToMixedPrecision(mixed_precision_type="float16", missing_op_mode=1):
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index b3504dbac506..ab0fc1709fa9 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -31,3 +31,5 @@
 from .module import load_module, enabled, system_lib
 from .container import String, ShapeTuple
 from .params import save_param_dict, load_param_dict
+
+from . import executor
diff --git a/tests/python/contrib/test_hexagon/rpc/conftest.py b/python/tvm/runtime/executor/__init__.py
similarity index 69%
rename from tests/python/contrib/test_hexagon/rpc/conftest.py
rename to python/tvm/runtime/executor/__init__.py
index 50c199a2f56f..ecc4097dbaa0 100644
--- a/tests/python/contrib/test_hexagon/rpc/conftest.py
+++ b/python/tvm/runtime/executor/__init__.py
@@ -15,17 +15,12 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import pytest
+"""This module contains Python wrappers for the TVM C++ Executor implementations.
 
+NOTE: at present, only AOT Executor is contained here. The others are:
+ - GraphExecutor, in python/tvm/contrib/graph_executor.py
+ - VM Executor, in python/tvm/runtime/vm.py
 
-def pytest_addoption(parser):
-    parser.addoption(
-        "--serial-number",
-        required=True,
-        help=("Android device serial number list from 'adb' command."),
-    )
-
-
-@pytest.fixture
-def android_serial_number(request):
-    return request.config.getoption("--serial-number")
+TODO(areusch): Consolidate these into this module.
+"""
+from .aot_executor import AotModule
diff --git a/python/tvm/runtime/executor/aot_executor.py b/python/tvm/runtime/executor/aot_executor.py
new file mode 100644
index 000000000000..91f056ff25fa
--- /dev/null
+++ b/python/tvm/runtime/executor/aot_executor.py
@@ -0,0 +1,182 @@
+# 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.
+"""A Python wrapper for the Module-based Model Runtime Interface for Ahead-of-Time compilation."""
+
+import numpy as np
+
+
+class AotModule(object):
+    """Wraps the AOT executor runtime.Module.
+
+    This is a thin wrapper of the underlying TVM module.
+    you can also directly call set_input, run, and get_output
+    of underlying module functions
+
+    Parameters
+    ----------
+    module : tvm.runtime.Module
+        The internal tvm module that holds the actual graph functions.
+
+    Attributes
+    ----------
+    module : tvm.runtime.Module
+        The internal tvm module that holds the actual graph functions.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import tvm
+        from tvm import relay
+        from tvm.contrib import graph_executor
+
+        # build the library using graph executor
+        lib = relay.build(...)
+        lib.export_library("compiled_lib.so")
+        # load it back as a runtime
+        lib: tvm.runtime.Module = tvm.runtime.load_module("compiled_lib.so")
+        # Call the library factory function for default and create
+        # a new runtime.Module, wrap with graph module.
+        gmod = graph_executor.GraphModule(lib["default"](dev))
+        # use the graph module.
+        gmod.set_input("x", data)
+        gmod.run()
+    """
+
+    def __init__(self, module):
+        self.module = module
+        self._set_input = module["set_input"]
+        self._run = module["run"]
+        self._get_output = module["get_output"]
+        self._get_input = module["get_input"]
+        self._get_num_outputs = module["get_num_outputs"]
+        self._get_input_index = module["get_input_index"]
+        self._get_num_inputs = module["get_num_inputs"]
+
+    def set_input(self, key=None, value=None, **params):
+        """Set inputs to the module via kwargs
+
+        Parameters
+        ----------
+        key : int or str
+           The input key
+
+        value : the input value.
+           The input key
+
+        params : dict of str to NDArray
+           Additional arguments
+        """
+        if key is not None:
+            v = self._get_input(key)
+            if v is None:
+                raise RuntimeError("Could not find '%s' in graph's inputs" % key)
+            v.copyfrom(value)
+
+        if params:
+            # upload big arrays first to avoid memory issue in rpc mode
+            keys = list(params.keys())
+            keys.sort(key=lambda x: -np.prod(params[x].shape))
+            for k in keys:
+                # TODO(zhiics) Skip the weights for submodule in a better way.
+                # We should use MetadataModule for initialization and remove
+                # params from set_input
+                val = self._get_input(k)
+                if val:
+                    self._get_input(k).copyfrom(params[k])
+
+    def run(self, **input_dict):
+        """Run forward execution of the graph
+
+        Parameters
+        ----------
+        input_dict: dict of str to NDArray
+            List of input values to be feed to
+        """
+        if input_dict:
+            self.set_input(**input_dict)
+        self._run()
+
+    def get_num_outputs(self):
+        """Get the number of outputs from the graph
+
+        Returns
+        -------
+        count : int
+            The number of outputs.
+        """
+        return self._get_num_outputs()
+
+    def get_num_inputs(self):
+        """Get the number of inputs to the graph
+
+        Returns
+        -------
+        count : int
+            The number of inputs.
+        """
+        return self._get_num_inputs()
+
+    def get_input(self, index, out=None):
+        """Get index-th input to out
+
+        Parameters
+        ----------
+        index : int
+            The input index
+
+        out : NDArray
+            The output array container
+        """
+        if out:
+            self._get_input(index).copyto(out)
+            return out
+
+        return self._get_input(index)
+
+    def get_input_index(self, name):
+        """Get inputs index via input name.
+
+        Parameters
+        ----------
+        name : str
+           The input key name
+
+        Returns
+        -------
+        index: int
+            The input index. -1 will be returned if the given input name is not found.
+        """
+        return self._get_input_index(name)
+
+    def get_output(self, index, out=None):
+        """Get index-th output to out
+
+        Parameters
+        ----------
+        index : int
+            The output index
+
+        out : NDArray
+            The output array container
+        """
+        if out:
+            self._get_output(index, out)
+            return out
+
+        return self._get_output(index)
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index 5395326805f0..27fd5af51a27 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -380,6 +380,7 @@ def __init__(self, exe, device, memory_cfg=None):
         self._get_num_outputs = self.module["get_num_outputs"]
         self._get_input_index = self.module["get_input_index"]
         self._set_input = self.module["set_input"]
+        self._set_one_input = self.module["set_one_input"]
         self._setup_device(device, memory_cfg)
 
     def _setup_device(self, dev, memory_cfg):
@@ -450,6 +451,30 @@ def set_input(self, func_name, *args, **kwargs):
         cargs = convert(args)
         self._set_input(func_name, *cargs)
 
+    def set_one_input(self, func_name, *args, **kwargs):
+        """Set the one input tensor with tag to a function.
+
+        Parameters
+        ----------
+        func_name : str
+            The name of the function.
+        args : [str or int, tvm.runtime.NDArray]
+            name or index of tensor and input tensor, optional
+        kwargs: dict of str or int to tvm.runtime.NDArray, optional
+            taged arguments to the function.
+        Only args or kwargs should exist
+        """
+        if kwargs:
+            assert len(kwargs) == 1
+            tag = next(iter(kwargs))
+            if isinstance(tag, str):
+                func_params = self._exec.get_function_params(func_name)
+                assert tag in func_params
+            self._set_one_input(func_name, tag, kwargs[tag])
+        else:
+            assert len(args) == 2
+            self._set_one_input(func_name, args[0], args[1])
+
     def invoke(self, func_name, *args, **kwargs):
         """Invoke a function.
 
diff --git a/python/tvm/script/tir/scope_handler.py b/python/tvm/script/tir/scope_handler.py
index fc953771bf21..418bae1a9d70 100644
--- a/python/tvm/script/tir/scope_handler.py
+++ b/python/tvm/script/tir/scope_handler.py
@@ -19,8 +19,9 @@
 from typing import Tuple, Any, Callable, Optional, List, Union, Mapping
 
 import synr
+import numpy as np
 import tvm.tir
-from tvm.runtime import Object
+from tvm.runtime import Object, String
 from tvm.ir import Span, Range
 from tvm.tir import Stmt, PrimExpr, IterVar, Var, Buffer, BufferRegion, ForKind
 
@@ -157,6 +158,56 @@ def setup_buffer_var(
         context.update_symbol(name, self.buffer_var, node)
 
 
+@register
+class AllocateConst(WithScopeHandler):
+    """With scope handler T.allocate_const(data, extents, dtype, condition)
+
+    TIR constant node to represent non-scalar constant
+    """
+
+    def __init__(self):
+        def allocate_const(raw_data, dtype, shape, span=None):
+            list_data = []
+            for i in raw_data:
+                list_data.append(i.value)
+            nd_data = tvm.nd.array(np.asarray(list_data, dtype=dtype))
+            n = tvm.tir.AllocateConst(self.buffer_var, dtype, shape, nd_data, self.body, span=span)
+            return n
+
+        super().__init__(allocate_const, concise_scope=True, def_symbol=True)
+        self.buffer_var = None
+
+    def enter_scope(
+        self,
+        node: synr.ast.Node,
+        context: ContextMaintainer,
+        arg_list: List[Any],
+        span: synr.ast.Span,
+    ):
+        # define buffer vars in symbol table
+        if isinstance(node, synr.ast.With):
+            vars = WithScopeHandler.get_optional_vars(node, context)
+            if len(vars) != 1:
+                context.report_error(f"Unexpected number of vars: 1 vs. {len(vars)}", node.span)
+            name = vars[0].id.name
+            var_span = vars[0].id.span
+        elif isinstance(node, synr.ast.Assign):
+            if len(node.lhs) != 1:
+                context.report_error(f"Unexpected number of vars: 1 vs. {len(node.lhs)}", node.span)
+            name = node.lhs[0].id.name
+            var_span = node.lhs[0].id.span
+        else:
+            raise Exception("Internal Bug")
+
+        def setup_buffer_var(data, dtype, shape, span: Span = None):
+            """Setup buffer var for a given type."""
+            buffer_ptr_type = tvm.ir.PointerType(tvm.ir.PrimType(dtype))
+            self.buffer_var = tvm.tir.Var(name, buffer_ptr_type, span)
+
+        setup_buffer_var(*arg_list, span=tvm_span_from_synr(var_span))
+        context.update_symbol(name, self.buffer_var, node)
+
+
 @register
 class LaunchThread(WithScopeHandler):
     """With scope handler T.launch_thread(env_var, extent)"""
@@ -483,8 +534,14 @@ def create_loop_info(
         """
         assert self.context and self.node, "call 'exit_scope' before 'enter_scope'"
         extent = end if begin == 0 else self.context.analyzer.simplify(end - begin)
-        self.annotations = annotations
-        self.loop_info.append(LoopInfo(begin, extent, kind, thread_binding, self.annotations))
+        self.annotations: Mapping[str, Object] = {}
+        if annotations is not None:
+            self.annotations = {
+                key: String(val) if isinstance(val, str) else val
+                for key, val in annotations.items()
+            }
+
+        self.loop_info.append(LoopInfo(begin, extent, kind, thread_binding, annotations))
 
 
 @register
diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py
index 308257085e51..aaad6e108e7b 100644
--- a/python/tvm/te/__init__.py
+++ b/python/tvm/te/__init__.py
@@ -31,7 +31,7 @@
 from .tensor import TensorSlice, Tensor
 from .tensor_intrin import decl_tensor_intrin
 from .tag import tag_scope
-from .operation import placeholder, compute, scan, extern, var, size_var
+from .operation import placeholder, compute, scan, extern, var, size_var, const
 from .operation import thread_axis, reduce_axis
 from .operation import create_prim_func, create_prim_func_from_outputs
 
diff --git a/python/tvm/te/hybrid/parser.py b/python/tvm/te/hybrid/parser.py
index 442aeb6f1027..655748783dd1 100644
--- a/python/tvm/te/hybrid/parser.py
+++ b/python/tvm/te/hybrid/parser.py
@@ -312,7 +312,7 @@ def visit_Assign(self, node):
                     "You should bind a pure name to the tensors",
                 )
                 self.add_symbol(node.targets[i].id, Symbol.GlobalBuffer, rhs.output(i))
-                rmap[rhs.outputs[i].op] = rhs.output(i)
+                rmap[rhs.symbolic_outputs[i].op] = rhs.output(i)
             return utils.replace_io(rhs.body, rmap)
 
         _internal_assert(len(node.targets) == 1, "So far only one-valued assignment is supported!")
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index dafbbfd15719..90d7cb5d75db 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -56,7 +56,7 @@ def placeholder(shape, dtype=None, name="placeholder"):
     return _ffi_api.Placeholder(shape, dtype, name)
 
 
-def compute(shape, fcompute, name="compute", tag="", attrs=None):
+def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=None):
     """Construct a new tensor by computing over the shape domain.
 
     The compute rule is result[axis] = fcompute(axis)
@@ -78,6 +78,10 @@ def compute(shape, fcompute, name="compute", tag="", attrs=None):
     attrs: dict, optional
         The additional auxiliary attributes about the compute.
 
+    varargs_names: list, optional
+        The names to use for each of the varargs. If not supplied, the varargs
+        will be called i1, i2, ...
+
     Returns
     -------
     tensor: Tensor
@@ -97,7 +101,16 @@ def compute(shape, fcompute, name="compute", tag="", attrs=None):
         arg_names = ["i%d" % i for i in range(out_ndim)]
     elif argspec.varargs is not None:
         # if there is a varargs, it takes the remaining dimensions of out_ndim
-        arg_names = argspec.args + [f"i{i}" for i in range(out_ndim - len(argspec.args))]
+        num_remaining_args = out_ndim - len(argspec.args)
+        if varargs_names is not None:
+            if len(varargs_names) != num_remaining_args:
+                raise RuntimeError(
+                    f"Number of varargs ({num_remaining_args}) does not match number"
+                    f"of varargs_names ({len(varargs_names)})"
+                )
+            arg_names = argspec.args + varargs_names
+        else:
+            arg_names = argspec.args + [f"i{i}" for i in range(out_ndim - len(argspec.args))]
     else:
         arg_names = argspec.args
         # if there are fewer args than out dimensions, the remaining dimensions
@@ -362,6 +375,28 @@ def var(name="tindex", dtype="int32", span=None):
     return tvm.tir.Var(name, dtype, span)
 
 
+def const(dtype="int32", span=None):
+    """Create a new constant with specified name and dtype
+
+    Parameters
+    ----------
+    name : str
+        The name
+
+    dtype : str
+        The data type
+
+    span : Optional[Span]
+        The location of this variable in the source.
+
+    Returns
+    -------
+    var : Var
+        The result symbolic variable.
+    """
+    return tvm.tir.const(dtype, span)
+
+
 def size_var(name="size", dtype="int32", span=None):
     """Create a new variable represents a tensor shape size, which is non-negative.
 
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index fc85d830c91a..b5609375b8aa 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -86,11 +86,12 @@ def __eq__(self, other):
             if isinstance(other, _expr.ExprOp):
                 return _expr.EqualOp(self, other)
             return False
+        if self.same_as(other):
+            return True
         if self.ndim == 0 and other.ndim == 0:
             raise ValueError(
                 "Equal == comparison among rank-0 tensor is ambiguous, "
-                "use Tensor.equal for content expression equvalence, "
-                "use Tensor.same_as for exact reference comparison"
+                "use Tensor.equal for content expression equvalence."
             )
         return _ffi_api.TensorEqual(self, other)
 
diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py
index dbaba46fdc9c..9ac1e9035235 100644
--- a/python/tvm/testing/utils.py
+++ b/python/tvm/testing/utils.py
@@ -635,6 +635,49 @@ def requires_nvptx(*args):
     return _compose(args, _requires_nvptx)
 
 
+def requires_nvcc_version(major_version, minor_version=0, release_version=0):
+    """Mark a test as requiring at least a specific version of nvcc.
+
+    Unit test marked with this decorator will run only if the
+    installed version of NVCC is at least `(major_version,
+    minor_version, release_version)`.
+
+    This also marks the test as requiring a cuda support.
+
+    Parameters
+    ----------
+    major_version: int
+
+        The major version of the (major,minor,release) version tuple.
+
+    minor_version: int
+
+        The minor version of the (major,minor,release) version tuple.
+
+    release_version: int
+
+        The release version of the (major,minor,release) version tuple.
+
+    """
+
+    try:
+        nvcc_version = nvcc.get_cuda_version()
+    except RuntimeError:
+        nvcc_version = (0, 0, 0)
+
+    min_version = (major_version, minor_version, release_version)
+    version_str = ".".join(str(v) for v in min_version)
+    requires = [
+        pytest.mark.skipif(nvcc_version < min_version, reason=f"Requires NVCC >= {version_str}"),
+        *requires_cuda(),
+    ]
+
+    def inner(func):
+        return _compose([func], requires)
+
+    return inner
+
+
 def requires_cudagraph(*args):
     """Mark a test as requiring the CUDA Graph Feature
 
diff --git a/python/tvm/tir/__init__.py b/python/tvm/tir/__init__.py
index 5854b9369c16..17f9aa3d9c60 100644
--- a/python/tvm/tir/__init__.py
+++ b/python/tvm/tir/__init__.py
@@ -28,7 +28,16 @@
 from .expr import Call, CallEffectKind, Let, IterVar, CommReducer, Any
 
 from .stmt import Stmt, LetStmt, AssertStmt, ForKind, For, While
-from .stmt import BufferStore, BufferRealize, Store, ProducerStore, Allocate, AttrStmt
+from .stmt import (
+    BufferStore,
+    BufferRealize,
+    Store,
+    ProducerStore,
+    Allocate,
+    AllocateConst,
+    AttrStmt,
+)
+
 from .stmt import ProducerRealize, SeqStmt
 from .stmt import IfThenElse, Evaluate, Prefetch, stmt_seq, stmt_list
 from .stmt import BufferRegion, MatchBufferRegion, Block, BlockRealize
diff --git a/python/tvm/tir/stmt.py b/python/tvm/tir/stmt.py
index de200d5eabdd..39831459f344 100644
--- a/python/tvm/tir/stmt.py
+++ b/python/tvm/tir/stmt.py
@@ -340,6 +340,40 @@ def __init__(self, buffer_var, dtype, extents, condition, body, annotations=None
         )
 
 
+@tvm._ffi.register_object("tir.AllocateConst")
+class AllocateConst(Stmt):
+    """Allocate constant node.
+
+    Parameters
+    ----------
+    buffer_var : Var
+        The buffer variable.
+
+    data : NDarray
+        The data associated with the constant
+
+    dtype : str
+        The data type of the buffer.
+
+    extents : list of Expr
+        The extents of the allocate
+
+    condition : PrimExpr
+        The condition.
+
+    body : Stmt
+        The body statement.
+
+    span : Optional[Span]
+        The location of this itervar in the source code.
+    """
+
+    def __init__(self, buffer_var, dtype, extents, condition, body, span=None):
+        self.__init_handle_by_constructor__(
+            _ffi_api.AllocateConst, buffer_var, dtype, extents, condition, body, span
+        )
+
+
 @tvm._ffi.register_object("tir.AttrStmt")
 class AttrStmt(Stmt):
     """AttrStmt node.
diff --git a/python/tvm/tir/transform/transform.py b/python/tvm/tir/transform/transform.py
index e2bcd6cf795b..d5f1a9ae979b 100644
--- a/python/tvm/tir/transform/transform.py
+++ b/python/tvm/tir/transform/transform.py
@@ -760,3 +760,25 @@ def ConvertForLoopsToSerial():
         The result pass
     """
     return _ffi_api.ConvertForLoopsToSerial()  # type: ignore
+
+
+def InjectSoftwarePipeline():
+    """Transform annotated loops into pipelined one that parallelize producers and consumers
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.InjectSoftwarePipeline()  # type: ignore
+
+
+def ExtractPrimFuncConstants():
+    """Collects and unificates tir non-scalar constants to module's attr 'Constants' array.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.ExtractPrimFuncConstants()  # type: ignore
diff --git a/python/tvm/topi/arm_cpu/arm_utils.py b/python/tvm/topi/arm_cpu/arm_utils.py
index 52ed7a18df81..4ab72178b30b 100644
--- a/python/tvm/topi/arm_cpu/arm_utils.py
+++ b/python/tvm/topi/arm_cpu/arm_utils.py
@@ -64,6 +64,12 @@ def is_aarch64_arm():
     return "aarch64" in target.attrs.get("mtriple", "")
 
 
+def is_neon_available():
+    """Check if neon instructions are available"""
+    target = tvm.target.Target.current(allow_none=False)
+    return "+neon" in target.mattr
+
+
 def get_tiling_B_interleaved_t(interleave_A):
     """Compute the tiling information for matrix B', where B'
     is the transposed and interleaved version of matrix B in C=A*B.
diff --git a/python/tvm/topi/arm_cpu/conv2d_alter_op.py b/python/tvm/topi/arm_cpu/conv2d_alter_op.py
index cbe8644c885f..409768fc8f75 100644
--- a/python/tvm/topi/arm_cpu/conv2d_alter_op.py
+++ b/python/tvm/topi/arm_cpu/conv2d_alter_op.py
@@ -24,10 +24,12 @@
 from tvm import relay
 from tvm import autotvm
 
-from ..nn import conv2d_alter_layout
+from ..nn import conv2d_alter_layout, conv2d_legalize
 from ..utils import get_const_tuple
 from ..x86.conv2d import _get_default_config as _get_x86_default_config
+from .conv2d_int8 import is_int8_hw_support
 from .arm_utils import get_tiling_B_interleaved_t
+from ..generic.conv2d import conv2d_alter_int8_common
 
 logger = logging.getLogger("topi")
 
@@ -257,7 +259,15 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         assert data_layout == "NCHW" and kernel_layout == "OIHW"
         if cfg.is_fallback:
             _get_x86_default_config(
-                cfg, data_tensor, kernel_tensor, strides, padding, out_dtype, False, data_layout
+                cfg,
+                data_tensor,
+                kernel_tensor,
+                strides,
+                padding,
+                dilation,
+                out_dtype,
+                False,
+                data_layout,
             )
         batch_size, in_channel, height, width = get_const_tuple(data_tensor.shape)
         out_channel, _, kh, kw = get_const_tuple(kernel_tensor.shape)
@@ -333,6 +343,57 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         )
         dispatch_ctx.update(target, new_workload, cfg)
         return relay.nn.contrib_depthwise_conv2d_nchwc(*inputs, **new_attrs)
+
+    if topi_tmpl == "conv2d_NCHWc_int8.arm_cpu":
+        assert data_layout == "NCHW" and kernel_layout == "OIHW"
+        if cfg.is_fallback:
+            _get_default_config_int8(
+                cfg,
+                data_tensor,
+                kernel_tensor,
+                strides,
+                padding,
+                dilation,
+                out_dtype,
+                False,
+                data_layout,
+            )
+
+        batch_size, in_channel, height, width = get_const_tuple(data_tensor.shape)
+        out_channel, channel_multiplier, kh, kw = get_const_tuple(kernel_tensor.shape)
+        ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
+        n_elems = 8
+
+        # update new attrs
+        new_attrs["channels"] = out_channel
+        new_attrs["data_layout"] = "NCHW%dc" % ic_bn
+        new_attrs["kernel_layout"] = "OIHW{:n}i{:n}o{:n}i".format(ic_bn // n_elems, oc_bn, n_elems)
+        new_attrs["out_layout"] = "NCHW%dc" % oc_bn
+
+        # Store altered operator's config.
+        new_data = te.placeholder(
+            (batch_size, in_channel // ic_bn, height, width, ic_bn), dtype=data_dtype
+        )
+        new_kernel = te.placeholder(
+            (out_channel // oc_bn, in_channel // ic_bn, kh, kw, ic_bn // n_elems, oc_bn, n_elems),
+            dtype=kernel_dtype,
+        )
+        new_workload = autotvm.task.args_to_workload(
+            [
+                new_data,
+                new_kernel,
+                strides,
+                padding,
+                dilation,
+                new_attrs["data_layout"],
+                new_attrs["out_layout"],
+                out_dtype,
+            ],
+            topi_tmpl,
+        )
+        dispatch_ctx.update(target, new_workload, cfg)
+        return relay.nn.contrib_conv2d_nchwc(*inputs, **new_attrs)
+
     if topi_tmpl == "conv2d_NHWC_quantized_interleaved.arm_cpu":
         assert data_layout == "NHWC" and kernel_layout == "HWIO"
         KH, KW, _, OC = get_const_tuple(kernel.shape)
@@ -365,3 +426,43 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
             inputs[0], new_kernel_expr, **new_attrs
         )
     return None
+
+
+@conv2d_legalize.register("arm_cpu")
+def _conv2d_legalize(attrs, inputs, arg_types):
+    """Legalizes Conv2D op.
+
+    Parameters
+    ----------
+    attrs : tvm.ir.Attrs
+        Attributes of current convolution
+    inputs : list of tvm.relay.Expr
+        The args of the Relay expr to be legalized
+    types : list of types
+        List of input and output types
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The legalized expr
+    """
+    # Collect the input tensors.
+    data_tensor, kernel_tensor = arg_types[0], arg_types[1]
+    data_dtype = data_tensor.dtype
+    kernel_dtype = kernel_tensor.dtype
+
+    # Collect the output tensor.
+    output_tensor = arg_types[2]
+
+    # Collect the input exprs.
+    data, kernel = inputs
+
+    # ARM vector instructions operate on the same dtype for data and kernel, we
+    # provide those here and conv2d_alter_int8_common will convert to the
+    # correct datatype.
+    if is_int8_hw_support(kernel_dtype, kernel_dtype):
+        # ARM intrinsics need the datatypes of data and kernel to be the same
+        return conv2d_alter_int8_common(
+            data, data_tensor, kernel, kernel_tensor, output_tensor, attrs, kernel_dtype, 8, 8
+        )
+    return None
diff --git a/python/tvm/topi/arm_cpu/conv2d_int8.py b/python/tvm/topi/arm_cpu/conv2d_int8.py
index bf4c03a6e5ed..8d9c47966113 100644
--- a/python/tvm/topi/arm_cpu/conv2d_int8.py
+++ b/python/tvm/topi/arm_cpu/conv2d_int8.py
@@ -16,20 +16,21 @@
 # under the License.
 # pylint: disable=invalid-name,unused-variable,unused-argument,no-member
 """Conv2D int8 schedule on ARM"""
-from tvm import te
-from tvm import autotvm
-from .. import tag
+from tvm import te, target, autotvm
 from ..utils import traverse_inline, get_const_tuple
 from ..generic import conv2d as conv2d_generic
 from .. import nn
-from ..nn.conv2d import _get_workload as _get_conv2d_workload
-from .tensor_intrin import dot_int8_int8_int32
+from ...target import codegen
+from ..nn.conv2d import _get_workload as _get_conv2d_workload, unpack_NCHWc_to_nchw
+from ..x86.conv2d_int8 import _pack_data
+from ..nn.utils import get_pad_tuple
+from .tensor_intrin import dot_int8_int8_int32_neon_82, dot_int8_int8_int32_neon
 from .conv2d_gemm import (
     compute_conv2d_gemm_without_weight_transform,
     schedule_conv2d_gemm_interleaved,
     schedule_conv2d_gemm_native,
 )
-from .arm_utils import get_tiling_B_interleaved_t
+from .arm_utils import get_tiling_B_interleaved_t, is_dotprod_available, is_neon_available
 
 
 def _get_default_config(cfg, data, kernel, strides, padding, dilation, out_dtype):
@@ -39,10 +40,10 @@ def _get_default_config(cfg, data, kernel, strides, padding, dilation, out_dtype
     wkl = _get_conv2d_workload(data, kernel, strides, padding, dilation, out_dtype)
     is_kernel_1x1 = wkl.kernel_h == 1 and wkl.kernel_w == 1
     if is_kernel_1x1:
-        conv2d_generic.fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes=2, num_int8_elements=4)
+        conv2d_generic.fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes=4, num_int8_elements=4)
     else:
         conv2d_generic.fallback_schedule_cpu_common_int8(
-            cfg, wkl, int32_lanes=2, num_int8_elements=4
+            cfg, wkl, int32_lanes=4, num_int8_elements=4
         )
 
 
@@ -51,11 +52,40 @@ def conv2d_NCHWc_int8(cfg, data, kernel, strides, padding, dilation, layout, out
     """Compute conv2d int8 with NCHWc layout"""
     # layout and out_layout are not used here,
     # we keep them for debug convenience when dumping autotvm workload
-    n, ic_chunk, ih, iw, ic_bn = get_const_tuple(data.shape)
-    in_channel = ic_chunk * ic_bn
 
-    oc_chunk, ic_chunk, kh, kw, ic_bn, oc_bn, n_elems = get_const_tuple(kernel.shape)
-    num_filter = oc_chunk * oc_bn
+    if len(data.shape) == 5:  # data is in nchwc
+        n, ic_chunk, ih, iw, ic_bn = get_const_tuple(data.shape)
+        in_channel = ic_chunk * ic_bn
+
+        oc_chunk, ic_chunk, kh, kw, ic_bn, oc_bn, n_elems = get_const_tuple(kernel.shape)
+        num_filter = oc_chunk * oc_bn
+    else:
+        # data is nchw, implicitly treat it as nchw1c
+        n, in_channel, ih, iw = get_const_tuple(data.shape)
+        num_filter, _, kh, kw = get_const_tuple(kernel.shape)
+
+    # Define autotvm tuning space
+    is_kernel_1x1 = kh == 1 and kw == 1
+    pt, pl, pb, pr = get_pad_tuple(padding, (kh, kw))
+    sh, sw = strides if isinstance(strides, (tuple, list)) else (strides, strides)
+    dh, dw = dilation if isinstance(dilation, (tuple, list)) else (dilation, dilation)
+    dilated_kernel_h = (kh - 1) * dh + 1
+    dilated_kernel_w = (kw - 1) * dw + 1
+    oh = (ih - dilated_kernel_h + pt + pb) // sh + 1
+    ow = (iw - dilated_kernel_w + pl + pr) // sw + 1
+
+    # input and output should be a multiple of 8 (intrinsics are 8 lanes)
+    cfg.define_split(
+        "tile_ic", in_channel, num_outputs=2, filter=lambda y: y.size[-1] % min(8, in_channel) == 0
+    )
+    cfg.define_split(
+        "tile_oc", num_filter, num_outputs=2, filter=lambda y: y.size[-1] % min(8, num_filter) == 0
+    )
+    cfg.define_split("tile_ow", ow, num_outputs=2, filter=lambda y: y.size[-1] <= 64)
+    if is_kernel_1x1:
+        cfg.define_knob("tile_oh", [1, 2] if oh > 1 else [1])
+    else:
+        cfg.define_knob("unroll_kw", [True, False])
 
     # If no config was set, we can fallback to NCHW config.
     if cfg.is_fallback:
@@ -68,27 +98,42 @@ def conv2d_NCHWc_int8(cfg, data, kernel, strides, padding, dilation, layout, out
             dilation,
             out_dtype,
         )
-    return nn.conv2d_NCHWc_int8_compute(
+    # Pack data if raw 4-D data is provided.
+    # This can only happen when autotuning.
+    if len(data.shape) == 4:
+        data, kernel = _pack_data(cfg, data, kernel)
+
+    return nn.conv2d_NCHWc_int8(
         data, kernel, strides, padding, dilation, layout, out_layout, out_dtype
     )
 
 
+def is_int8_hw_support(data_dtype, kernel_dtype):
+    """
+    Checks to ensure that we can use int8 on arm
+    1) The datatypes are correct.
+    2) LLVM version has support for the instructions.
+    """
+    # 1) Check datatypes
+    is_dtype_support = data_dtype == kernel_dtype and "int8" in data_dtype
+
+    # 2) Check LLVM support
+    llvm_version = codegen.llvm_version_major()
+    is_llvm_support = llvm_version >= 8
+
+    # 3) Check target
+    is_target_support = is_neon_available() or is_dotprod_available()
+
+    return is_dtype_support and is_llvm_support
+
+
 @autotvm.register_topi_schedule("conv2d_NCHWc_int8.arm_cpu")
 def schedule_conv2d_NCHWc_int8(cfg, outs):
     """Create schedule for tensors"""
     s = te.create_schedule([x.op for x in outs])
     scheduled_ops = []
 
-    def traverse(op):
-        """Traverse operators from computation graph"""
-        # inline all one-to-one-mapping operators except the last stage (output)
-        if tag.is_broadcast(op.tag):
-            if op not in s.outputs:
-                s[op].compute_inline()
-            for tensor in op.input_tensors:
-                if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
-                    traverse(tensor.op)
-
+    def _callback(op):
         if "conv2d_NCHWc_int8" in op.tag:
             conv_out = op.output(0)
             kernel_vec = conv_out.op.input_tensors[1]
@@ -106,21 +151,55 @@ def traverse(op):
             # int8 conv kernel is 7-dim
             _, _, kh, kw, _, _, _ = get_const_tuple(kernel_vec.shape)
             dtype = "uint" if data.dtype == "uint8" else "int"
+            if is_dotprod_available():
+                intrin = dot_int8_int8_int32_neon_82(int32_lanes=4)
+            elif is_neon_available():
+                intrin = dot_int8_int8_int32_neon()
+            else:
+                raise RuntimeError(
+                    "Cannot schedule schedule_NCHWc_int8 without neon or arm v8.2 neon support"
+                )
+            # On raspberry pi 4s, we see poor performance when the fused
+            # operations are inlined into the main computation body. These
+            # fused ops dominated the runtime on small convolutions repeatedly
+            # blow the cache. Using workloads from resnet50, inceptionv3, and
+            # mobilenetv3, we empirically determine the size at which inline is
+            # not worth it to be kernel heigh * kernel width < 500. These tests
+            # were only run on raspberry pi 4, other arm cpus may have larger
+            # caches where inlining has good performance.
+            if target.Target.current().mcpu == "cortex-a72" and kh * kw < 500:
+                inline_fused = False
+            else:
+                inline_fused = True
             if kh == 1 and kw == 1:
                 conv2d_generic.schedule_conv_NCHWc_cpu_1x1_int8(
-                    *args, int32_lanes=4, intrin=dot_int8_int8_int32(int32_lanes=4, dtype=dtype)
+                    *args, int32_lanes=4, int8_elems=4, intrin=intrin, inline_fused=inline_fused
                 )
             else:
                 conv2d_generic.schedule_conv_NCHWc_cpu_common_int8(
-                    *args, int32_lanes=4, intrin=dot_int8_int8_int32(int32_lanes=4, dtype=dtype)
+                    *args, int32_lanes=4, int8_elems=4, intrin=intrin, inline_fused=inline_fused
                 )
 
-        scheduled_ops.append(op)
-
-    traverse(outs[0].op)
+    traverse_inline(s, outs[0].op, _callback)
     return s
 
 
+def conv2d_nchw_int8(data, kernel, strides, padding, dilation, out_dtype):
+    """Compute conv2d with NCHW layout and int8 dtype"""
+    layout = "NCHW"
+    # pylint: disable=no-value-for-parameter
+    packed_out = conv2d_NCHWc_int8(
+        data, kernel, strides, padding, dilation, layout, layout, out_dtype
+    )
+    return unpack_NCHWc_to_nchw(packed_out, out_dtype)
+
+
+def schedule_conv2d_nchw_int8(outs):
+    """Create the schedule for conv2d_nchw_int8"""
+    # pylint: disable=no-value-for-parameter
+    return schedule_conv2d_NCHWc_int8(outs)
+
+
 def _compute_conv2d_NHWC_quantized(
     cfg, data, kernel, strides, padding, dilation, out_dtype, interleave_A
 ):
diff --git a/python/tvm/topi/arm_cpu/injective.py b/python/tvm/topi/arm_cpu/injective.py
index 330144b33fb6..7c3ea5261f5e 100644
--- a/python/tvm/topi/arm_cpu/injective.py
+++ b/python/tvm/topi/arm_cpu/injective.py
@@ -16,8 +16,8 @@
 # under the License.
 # pylint: disable=invalid-name, unused-variable
 """Schedule for pooling operators"""
-import tvm
 import numpy as np
+import tvm
 from tvm import te
 from ..utils import is_empty_shape
 
diff --git a/python/tvm/topi/arm_cpu/tensor_intrin.py b/python/tvm/topi/arm_cpu/tensor_intrin.py
index bc9b0491947b..6e0315233d87 100644
--- a/python/tvm/topi/arm_cpu/tensor_intrin.py
+++ b/python/tvm/topi/arm_cpu/tensor_intrin.py
@@ -428,7 +428,7 @@ def _instr():
     )
 
 
-def dot_int8_int8_int32(int32_lanes, dtype="uint"):
+def dot_int8_int8_int32_neon_82(int32_lanes, dtype="uint"):
     """
     Int8 dot product by every 4 elements using ARM v8.2 udot.
     This function takes two arrays of int8 datatype -- data[4] and
@@ -533,6 +533,138 @@ def _instr(index):
     )
 
 
+def dot_int8_int8_int32_neon():
+    """
+    Int8 dot product using vmlal instructions
+
+    .. code-block:: c
+
+        void dot_int8_int8_int32(int8 data[4], int8 kernel[4][4], int32 output[4]){
+            for (int i = 0; i < 4; i++){
+                out[i] = 0;
+                for (int k = 0; k < 4; k++){
+                    out[i] += data[k] * kernel[i][k]
+                }
+            }
+        }
+
+    We use the smull and saddlp instructions to compute the dot product.
+    smull : int8x16 -> int8x16 -> int16x8 elementwise multiplication
+    saddlp: int16x8 -> int32x4 pairwise addition of elements
+
+    Data is broadcast across the register
+    int8 elements
+    |         data      |         data      |
+    |  0 |  1 |  2 |  3 |  4 |  5 |  6 |  7 |
+
+                      smull
+
+    int8 elements
+    |     kernel[i]     |     kernel[i+1]   |
+    |  0 |  1 |  2 |  3 |  4 |  5 |  6 |  7 |
+
+                        =
+
+    int16 elements
+    |               data * kernel[i]        |         data * kernel[i+1]            |
+    |    0    |    1    |    2    |    3    |    4    |    5    |    6    |    7    |
+
+                                          saddlp =
+
+    int32 elements
+    |    partial sum(data * kernel[i])      |  partial sum(data * kernel[i+1])      |
+    |         0         |         1         |         2         |         3         |
+
+
+    We apply the above kernel twice and use addp to compute the second set of pairwise additions
+
+    int32 elements (narrowed for so they fit on a line)
+    |    psum d*k[i]    |   psum d*k[i+1]   |           |   psum d*k[i+2]   |   psum d*k[i+3]   |
+    |    0    |    1    |    2    |    3    |   addp    |    4    |    5    |    6    |    7    |
+                                                 =
+    |sum d*ki |sum d*ki1|sum d*ki2|sum d*ki3|
+    |    0    |    1    |    2    |    3    |
+
+
+    """
+    int32_lanes = 4  # 4 int32 lanes = 128
+    num_int8_elements = 4  # 4 int8 elements in int32
+    data = te.placeholder((num_int8_elements,), dtype="int8", name="data")
+    kernel = te.placeholder((int32_lanes, num_int8_elements), dtype="int8", name="kernel")
+    k = te.reduce_axis((0, num_int8_elements), name="k")
+    C = te.compute(
+        (int32_lanes,),
+        lambda i: te.sum(data[k].astype("int32") * kernel[i, k].astype("int32"), axis=k),
+        name="C",
+    )
+
+    a_buffer = tvm.tir.decl_buffer(
+        data.shape, dtype="int8", name="a_buffer", offset_factor=1, strides=[1]
+    )
+    b_buffer = tvm.tir.decl_buffer(
+        kernel.shape, dtype="int8", name="b_buffer", offset_factor=1, strides=[te.var("ldw"), 1]
+    )
+
+    def _intrin_func(ins, outs):
+        def _instr(index):
+            int_8xl = "int8x8"
+            int_32xl = "int32x4"
+            ib = tvm.tir.ir_builder.create()
+            if index == 1:
+                ib.emit(outs[0].vstore(0, tvm.tir.const(0, int_32xl)))
+                return ib.get()
+
+            def pairwise_add_mul(idx):
+                # this broadcasts data to the vector size
+                a_int8 = ins[0].vload([0], "int8x4")
+                re_int32 = tvm.tir.call_intrin("int32", "tir.reinterpret", a_int8)
+                vec_ai32 = re_int32.astype("int32x2")
+                vec_a = tvm.tir.call_intrin(int_8xl, "tir.reinterpret", vec_ai32)
+
+                vec_b = ins[1].vload([idx * 2, 0], int_8xl)  # we take two inputs at a time
+
+                multiply = tvm.tir.call_llvm_pure_intrin(
+                    "int16x8",
+                    "llvm.aarch64.neon.smull.v8i16",  # saturating pairwise multiplication
+                    tvm.tir.const(2, "uint32"),
+                    vec_a,
+                    vec_b,
+                )
+                pairwise_reduction = tvm.tir.call_llvm_pure_intrin(
+                    "int32x4",
+                    "llvm.aarch64.neon.saddlp.v4i32.v8i16",
+                    tvm.tir.const(1, "uint32"),
+                    multiply,
+                )
+                return pairwise_reduction
+
+            pair_1 = pairwise_add_mul(0)
+            pair_2 = pairwise_add_mul(1)
+            quad_reduction = tvm.tir.call_llvm_pure_intrin(
+                "int32x4",
+                "llvm.aarch64.neon.addp.v4i32",
+                tvm.tir.const(2, "uint32"),
+                pair_1,
+                pair_2,
+            )
+            if index == 0:
+                ib.emit(outs[0].vstore(0, quad_reduction))
+            else:
+                ib.emit(outs[0].vstore(0, quad_reduction + outs[0].vload([0], int_32xl)))
+            return ib.get()
+
+        # body, reset, update
+        return _instr(0), _instr(1), _instr(2)
+
+    buffer_params = {"offset_factor": 1}
+    return te.decl_tensor_intrin(
+        C.op,
+        _intrin_func,
+        binds={data: a_buffer, kernel: b_buffer},
+        default_buffer_params=buffer_params,
+    )
+
+
 def select_word(vec, lane, dtype_vec):
     """
     Utility function used to select a int8x4 word within a int8x16 vector
diff --git a/python/tvm/topi/cuda/conv3d.py b/python/tvm/topi/cuda/conv3d.py
index 51f1f7a27c7e..6b602384b4fe 100644
--- a/python/tvm/topi/cuda/conv3d.py
+++ b/python/tvm/topi/cuda/conv3d.py
@@ -26,7 +26,7 @@
 
 
 @autotvm.register_topi_compute("conv3d_ncdhw.cuda")
-def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, out_dtype="float32"):
+def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, groups, out_dtype="float32"):
     """Conv3D operator in NCDHW layout for cuda backend.
 
     Parameters
@@ -49,6 +49,9 @@ def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, out_dtype="float
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups
+
     out_dtype: str
         The output type. This is used for mixed precision.
 
@@ -57,7 +60,7 @@ def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, out_dtype="float
     output : tvm.te.Tensor
         5-D with shape [batch, out_channel, out_depth, out_height, out_width]
     """
-    return nn.conv3d_ncdhw(data, kernel, strides, padding, dilation, out_dtype)
+    return nn.conv3d_ncdhw(data, kernel, strides, padding, dilation, groups, out_dtype)
 
 
 @autotvm.register_topi_schedule("conv3d_ncdhw.cuda")
@@ -82,7 +85,7 @@ def schedule_conv3d_ncdhw(cfg, outs):
     s = te.create_schedule([x.op for x in outs])
 
     def _callback(op):
-        if op.tag == "conv3d_ncdhw":
+        if "conv3d_ncdhw" in op.tag:
             schedule_direct_conv3d_cuda(cfg, s, op.output(0), "NCDHW", "conv3d_ncdhw.cuda")
 
     traverse_inline(s, outs[0].op, _callback)
@@ -90,7 +93,7 @@ def _callback(op):
 
 
 @autotvm.register_topi_compute("conv3d_ndhwc.cuda")
-def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype="float32"):
+def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype="float32"):
     """Conv3d operator in NDHWC layout for cuda backend.
 
     Parameters
@@ -110,12 +113,15 @@ def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype="float
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups
+
     Returns
     -------
     Output : tvm.te.Tensor
         5-D with shape [batch, out_depth, out_height, out_width, out_channel]
     """
-    return nn.conv3d_ndhwc(data, kernel, strides, padding, dilation, out_dtype)
+    return nn.conv3d_ndhwc(data, kernel, strides, padding, dilation, groups, out_dtype)
 
 
 @autotvm.register_topi_schedule("conv3d_ndhwc.cuda")
@@ -140,7 +146,7 @@ def schedule_conv3d_ndhwc(cfg, outs):
     s = te.create_schedule([x.op for x in outs])
 
     def _callback(op):
-        if op.tag == "conv3d_ndhwc":
+        if "conv3d_ndhwc" in op.tag:
             schedule_direct_conv3d_cuda(cfg, s, op.output(0), "NDHWC", "conv3d_ndhwc.cuda")
 
     traverse_inline(s, outs[0].op, _callback)
@@ -149,7 +155,7 @@ def _callback(op):
 
 @autotvm.register_topi_compute("conv3d_cudnn.cuda")
 def conv3d_cudnn(
-    cfg, data, kernel, strides, padding, dilation, layout="NCDHW", out_dtype="float32"
+    cfg, data, kernel, strides, padding, dilation, groups, layout="NCDHW", out_dtype="float32"
 ):
     """Conv3D operator for cuda backend.
 
@@ -194,6 +200,8 @@ def conv3d_cudnn(
         raise ValueError("Unsupported layout %s in cudnn" % layout)
     CO, CI, KD, KH, KW = get_const_tuple(kernel.shape)
 
+    assert groups == 1, "conv3d_cudnn does not support groups"
+
     # handle dilation
     stride_d, stride_h, stride_w = (
         (strides, strides, strides) if isinstance(strides, int) else strides
diff --git a/python/tvm/topi/cuda/conv3d_ndhwc_tensorcore.py b/python/tvm/topi/cuda/conv3d_ndhwc_tensorcore.py
index efb25744b802..cf96794dd4f0 100644
--- a/python/tvm/topi/cuda/conv3d_ndhwc_tensorcore.py
+++ b/python/tvm/topi/cuda/conv3d_ndhwc_tensorcore.py
@@ -335,8 +335,9 @@ def get_strides(extents):
 
 
 @autotvm.register_topi_compute("conv3d_ndhwc_tensorcore.cuda")
-def conv3d_ndhwc_tensorcore(cfg, data, kernel, strides, padding, dilation, out_dtype):
+def conv3d_ndhwc_tensorcore(cfg, data, kernel, strides, padding, dilation, groups, out_dtype):
     """Compute conv3d with tensorcore for NDHWC layout"""
+    assert groups == 1, "tensorcore conv3d does not support groups"
     return ndhwc_tensorcore_cuda(cfg, data, kernel, strides, padding, dilation, out_dtype)
 
 
diff --git a/python/tvm/topi/cuda/conv3d_winograd.py b/python/tvm/topi/cuda/conv3d_winograd.py
index 2134ee9178b8..2f53d0458af1 100644
--- a/python/tvm/topi/cuda/conv3d_winograd.py
+++ b/python/tvm/topi/cuda/conv3d_winograd.py
@@ -620,7 +620,9 @@ def schedule_winograd_no_depth_cuda(cfg, s, output, pre_computed):
 
 
 @autotvm.register_topi_compute("conv3d_ncdhw_winograd.cuda")
-def conv3d_ncdhw_winograd(cfg, data, kernel, strides, padding, dilation, out_dtype):
+def conv3d_ncdhw_winograd(cfg, data, kernel, strides, padding, dilation, groups, out_dtype):
+    """Conv3d NCDHW using winograd optimization"""
+    assert groups == 1, "conv3d_ncdhw_winograd only supports a single group"
     CO, CI, KD, KH, KW = get_const_tuple(kernel.shape)
     # Check if we can transform depth.
     if 2 < KD < 8 and KD == KH:
@@ -650,9 +652,12 @@ def _callback(op):
 
 @autotvm.register_topi_compute("conv3d_ncdhw_winograd_without_weight_transform.cuda")
 def conv3d_ncdhw_winograd_without_weight_transform(
-    cfg, data, kernel, strides, padding, dilation, out_dtype
+    cfg, data, kernel, strides, padding, dilation, groups, out_dtype
 ):
     """Conv3d NCDHW winograd without weight transform."""
+    assert (
+        groups == 1
+    ), "conv3d_ncdhw_winograd_without_weight_transform does not support more than one group"
     A, B, C, _, _ = get_const_tuple(kernel.shape)
     # Check if we can transform depth.
     if A == B == C:
diff --git a/python/tvm/topi/generic/conv2d.py b/python/tvm/topi/generic/conv2d.py
index dc70e0ed89f9..4ecf4613cf8a 100644
--- a/python/tvm/topi/generic/conv2d.py
+++ b/python/tvm/topi/generic/conv2d.py
@@ -19,8 +19,10 @@
 """Generic convolution schedules"""
 from tvm import te
 from tvm import autotvm
+from tvm import relay
 from tvm.autotvm.task.space import SplitEntity, OtherOptionEntity
 from ..utils import get_const_tuple, traverse_inline
+from ..nn.utils import get_pad_tuple
 
 
 def fallback_schedule_cpu_common_int8(cfg, wkl, int32_lanes, num_int8_elements):
@@ -120,7 +122,16 @@ def fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes, num_int8_elements):
 
 
 def schedule_conv_NCHWc_cpu_common_int8(
-    s, cfg, data_vec, kernel_vec, conv_out, last, int32_lanes=16, int8_elems=4, intrin=None
+    s,
+    cfg,
+    data_vec,
+    kernel_vec,
+    conv_out,
+    last,
+    int32_lanes=16,
+    int8_elems=4,
+    intrin=None,
+    inline_fused=False,
 ):
     """
     Defines the schedule for INT8 for Intel and ARM machines
@@ -176,14 +187,16 @@ def schedule_conv_NCHWc_cpu_common_int8(
     if C == O:
         s[C].parallel(parallel_axis)
 
-    s[CC].compute_at(s[C], ow_chunk)
+    s[CC].compute_at(s[C], parallel_axis)
     _, oc_chunk, oh, ow, oc_block = s[CC].op.axis
     kh, kw, ic_outer, ic_f_inner, ic_s_inner = s[CC].op.reduce_axis
 
     ow_chunk, ow_block = s[CC].split(ow, factor=reg_n)
 
-    assert oc_bn % int32_lanes == 0
-    assert ic_bn % int8_elems == 0  # (u)int8 elements in (u)int32
+    assert oc_bn % int32_lanes == 0, f"oc_bn={oc_bn} % int32_lanes={int32_lanes} != 0"
+    assert (
+        ic_bn % int8_elems == 0
+    ), f"ic_bn={ic_bn} % int8_elems={int8_elems} != 0"  # (u)int8 elements in (u)int32
 
     oc_f_inner, oc_s_inner = s[CC].split(oc_block, factor=int32_lanes)
 
@@ -228,27 +241,35 @@ def schedule_conv_NCHWc_cpu_common_int8(
             batch, oc_chunk, oh, ow, oc_block = s[O].op.axis
             ow_chunk, ow_block = s[O].split(ow, factor=reg_n)
             s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block)
-            s[C].compute_at(s[O], ow_block)
-            parallel_axis = s[O].fuse(batch, oc_chunk, oh)
-            s[O].vectorize(oc_block)
-            s[O].parallel(parallel_axis)
         elif out_ndim == 4:
             batch, oc, oh, ow = s[O].op.axis
             ow_chunk, ow_block = s[O].split(ow, factor=reg_n)
             oc_chunk, oc_block = s[O].split(oc, factor=oc_bn)
             s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block)
-            s[C].compute_at(s[O], ow_block)
-            parallel_axis = s[O].fuse(batch, oc_chunk, oh)
-            s[O].vectorize(oc_block)
-            s[O].parallel(parallel_axis)
         else:
             raise ValueError("Unsupported output ndim: %s" % out_ndim)
+        parallel_axis = s[O].fuse(batch, oc_chunk, oh)
+        if inline_fused:
+            s[C].compute_at(s[O], ow_block)
+        else:
+            s[C].compute_at(s[O], parallel_axis)
+        s[O].vectorize(oc_block)
+        s[O].parallel(parallel_axis)
 
     return s
 
 
 def schedule_conv_NCHWc_cpu_1x1_int8(
-    s, cfg, data_vec, kernel_vec, conv_out, last, int32_lanes=16, int8_elems=4, intrin=None
+    s,
+    cfg,
+    data_vec,
+    kernel_vec,
+    conv_out,
+    last,
+    int32_lanes=16,
+    int8_elems=4,
+    intrin=None,
+    inline_fused=False,
 ):
     """
     Defines the 1x1 conv schedule for INT8 for Intel and ARM machines
@@ -301,10 +322,10 @@ def schedule_conv_NCHWc_cpu_1x1_int8(
     s[C].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block)
     s[C].vectorize(oc_block)
 
-    s[CC].compute_at(s[C], ow_inner)
     parallel_axis = s[C].fuse(batch, oc_chunk, oh_outer)
     if C == O:
         s[C].parallel(parallel_axis)
+    s[CC].compute_at(s[C], parallel_axis)  # good perf on mobilenet, but not on individuals?
 
     _, oc_chunk, oh, ow, oc_block = s[CC].op.axis
     kh, kw, ic_outer, ic_f_inner, ic_s_inner = s[CC].op.reduce_axis
@@ -344,25 +365,22 @@ def schedule_conv_NCHWc_cpu_1x1_int8(
             batch, oc_chunk, oh, ow, oc_block = s[O].op.axis
             oh_outer, oh_inner = s[O].split(oh, factor=oh_factor)
             ow_outer, ow_inner = s[O].split(ow, factor=ow_factor)
-            s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block)
-
-            s[C].compute_at(s[O], ow_inner)
-            parallel_axis = s[O].fuse(batch, oc_chunk, oh_outer)
-            s[O].vectorize(oc_block)
-            s[O].parallel(parallel_axis)
         elif out_ndim == 4:
             batch, oc, oh, ow = s[O].op.axis
             oc_chunk, oc_block = s[O].split(oc, factor=oc_bn)
             oh_outer, oh_inner = s[O].split(oh, factor=oh_factor)
             ow_outer, ow_inner = s[O].split(ow, factor=ow_factor)
-            s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block)
+        else:
+            raise ValueError("Unsupported output ndim: %s" % out_ndim)
 
+        s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block)
+        parallel_axis = s[O].fuse(batch, oc_chunk, oh_outer)
+        if inline_fused:
             s[C].compute_at(s[O], ow_inner)
-            parallel_axis = s[O].fuse(batch, oc_chunk, oh_outer)
-            s[O].vectorize(oc_block)
-            s[O].parallel(parallel_axis)
         else:
-            raise ValueError("Unsupported output ndim: %s" % out_ndim)
+            s[C].compute_at(s[O], parallel_axis)
+        s[O].vectorize(oc_block)
+        s[O].parallel(parallel_axis)
 
     return s
 
@@ -394,3 +412,171 @@ def _callback(op):
 
     traverse_inline(s, outs[0].op, _callback)
     return s
+
+
+def conv2d_alter_int8_common(
+    data,
+    data_tensor,
+    kernel,
+    kernel_tensor,
+    output_tensor,
+    attrs,
+    data_dtype: str,
+    in_channel_vector_length: int,
+    out_channel_vector_length: int,
+):
+    """
+    Convert TE inputs/outputs so that they are suitable for fast Int8 instructions.
+
+    Int8 instructions require input channels and output channels to be a
+    multiple of the vector length. For input channels, we pad both the inputs
+    and weights channels. For output channels, we pad the weight and
+    stride_slice the output.
+
+    Arguments
+    ---------
+    data: Expr
+        Data Expr
+    data_tensor: Tensor
+        Data tensor
+    kernel: Expr
+        Kernel Expr
+    kernel_tensor: Tensor
+        Kernel tensor
+    output_tensor: Tensor
+        Output tensor
+    attrs: Conv2dAttrs
+        Attributes of the computation
+    data_dtype: "int8" or "uint8"
+        Desired dtype of data. Data will be converted to this dtype before the main computation.
+    in_channel_vector_length: int
+        Length of vector units on target hardware. Input channels are padded to this length.
+    out_channel_vector_length: int
+        Output size of vector instruction. Output channels are padded to this length.
+
+    Returns
+    -------
+    out : Tensor
+        Conv2d computation with inputs in the correct order for tensorization.
+    """
+    # Dilation not supported yet. Return None if dilation is not (1, 1)
+    dilation = attrs.get_int_tuple("dilation")
+    if not (dilation[0] == 1 and dilation[1] == 1):
+        return None
+
+    # No legalization for depthwise convolutions yet.
+    groups = attrs.get_int("groups")
+    if groups != 1:
+        return None
+
+    # Get the conv attrs
+    new_attrs = {k: attrs[k] for k in attrs.keys()}
+
+    padding = attrs.get_int_tuple("padding")
+    kh, kw = attrs.get_int_tuple("kernel_size")
+    pt, pl, pb, pr = get_pad_tuple(padding, (kh, kw))
+
+    if data_tensor.dtype != data_dtype:
+        # How to convert data to int8
+        # Original --> C = A (conv) B
+        # A and B are int8
+        #   C = (A + 128 - 128) (conv) B
+        #   C = (A' conv B) - 128 (conv) B
+        # where A' = A + 128
+        # and 128 (conv) B is basically a reduce on CRS axis for weights.
+        #
+        # How to convert data to uint8
+        #   C = (A - 128 + 128) (conv) B
+        #   C = (A' conv B) + 128 (conv) B
+        # where A' = A - 128
+        if data_dtype == "int8":
+            # shift data to int8
+            before_shift = relay.add
+            after_shift = relay.subtract
+        else:
+            # shift data to uint8
+            before_shift = relay.subtract
+            after_shift = relay.add
+
+        if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
+            adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(0, 1, 2))
+            pad_width = ((0, 0), (pt, pb), (pl, pr), (0, 0))
+        elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
+            pad_width = ((0, 0), (0, 0), (pt, pb), (pl, pr))
+            adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(1, 2, 3))
+            adjust_shift = relay.expand_dims(adjust_shift, axis=1, num_newaxis=2)
+        else:
+            return None
+
+        data = relay.cast(data, "int32")
+        data = before_shift(data, relay.const(128, "int32"))
+        data = relay.cast(data, data_dtype)
+
+        # Do external padding as pad value has to be 128.
+        if any(padding):
+            data = relay.nn.pad(data, pad_width=pad_width, pad_value=128)
+        new_attrs["padding"] = (0, 0)
+
+        # Multiply 128 to adjust shift.
+        adjust_shift = relay.multiply(adjust_shift, relay.const(128, "int32"))
+
+    # Flags to remember if the expr is modified
+    ic_modified = False
+    oc_modified = False
+
+    # Find the value of input and output channel.
+    in_channel = -1
+    out_channel = -1
+    if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
+        in_channel = data_tensor.shape[3].value
+        out_channel = kernel_tensor.shape[3].value
+    elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
+        in_channel = data_tensor.shape[1].value
+        out_channel = kernel_tensor.shape[0].value
+    else:
+        return None
+
+    if in_channel % in_channel_vector_length != 0:
+        new_in_channel = (
+            (in_channel + in_channel_vector_length) // in_channel_vector_length
+        ) * in_channel_vector_length
+        diff = new_in_channel - in_channel
+        if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
+            data = relay.nn.pad(data, pad_width=((0, 0), (0, 0), (0, 0), (0, diff)))
+            kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, diff), (0, 0)))
+            ic_modified = True
+        elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
+            pad_width = ((0, 0), (0, diff), (0, 0), (0, 0))
+            data = relay.nn.pad(data, pad_width=pad_width)
+            kernel = relay.nn.pad(kernel, pad_width=pad_width)
+            ic_modified = True
+        else:
+            return None
+
+    new_out_channel = out_channel
+    if out_channel % out_channel_vector_length != 0:
+        new_out_channel = (
+            (out_channel + out_channel_vector_length) // out_channel_vector_length
+        ) * out_channel_vector_length
+        diff = new_out_channel - out_channel
+        if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
+            kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, 0), (0, diff)))
+            oc_modified = True
+        elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
+            kernel = relay.nn.pad(kernel, pad_width=((0, diff), (0, 0), (0, 0), (0, 0)))
+            oc_modified = True
+        else:
+            return None
+
+    if oc_modified:
+        new_attrs["channels"] = new_out_channel
+        out = relay.nn.conv2d(data, kernel, **new_attrs)
+        original_out_shape = [x.value for x in output_tensor.shape]
+        out = relay.strided_slice(out, begin=[0, 0, 0, 0], end=original_out_shape)
+    else:
+        out = relay.nn.conv2d(data, kernel, **new_attrs)
+
+    if data_tensor.dtype != data_dtype:
+        out = after_shift(out, adjust_shift)
+
+    return out
diff --git a/python/tvm/topi/nn/conv2d.py b/python/tvm/topi/nn/conv2d.py
index 97317fa0f9cd..34357508122f 100644
--- a/python/tvm/topi/nn/conv2d.py
+++ b/python/tvm/topi/nn/conv2d.py
@@ -21,7 +21,7 @@
 
 from collections import namedtuple
 import re
-from typing import Union, Sequence
+from typing import Union, Sequence, Optional
 import numpy as np
 
 import tvm
@@ -313,63 +313,18 @@ def conv2d_nhwc(
     output : tvm.te.Tensor
         4-D with shape [batch, out_height, out_width, out_channel]
     """
-    assert isinstance(stride, int) or len(stride) == 2
-    assert isinstance(dilation, int) or len(dilation) == 2
-
-    if isinstance(stride, int):
-        stride_h = stride_w = stride
-    else:
-        stride_h, stride_w = stride
-
-    if isinstance(dilation, int):
-        dilation_h = dilation_w = dilation
-    else:
-        dilation_h, dilation_w = dilation
-
-    if auto_scheduler_rewritten_layout:
-        # Infer shape for the rewritten layout
-        kernel_h, kernel_w, channel, num_filter = auto_scheduler.get_shape_from_rewritten_layout(
-            auto_scheduler_rewritten_layout, ["ry", "rx", "rc", "ff"]
-        )
-        auto_scheduler.remove_index_check(Filter)
-    else:
-        kernel_h, kernel_w, channel, num_filter = Filter.shape
-
-    batch, in_height, in_width, in_channel = Input.shape
-    # compute the output shape
-    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
-    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
-    pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
-        padding, (dilated_kernel_h, dilated_kernel_w)
-    )
-    out_channel = num_filter
-    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
-    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
-    pad_before = [0, pad_top, pad_left, 0]
-    pad_after = [0, pad_down, pad_right, 0]
-    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
-    rc = te.reduce_axis((0, in_channel), name="rc")
-    ry = te.reduce_axis((0, kernel_h), name="ry")
-    rx = te.reduce_axis((0, kernel_w), name="rx")
-    Output = te.compute(
-        (batch, out_height, out_width, out_channel),
-        lambda nn, yy, xx, ff: te.sum(
-            PaddedInput[
-                nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc
-            ].astype(out_dtype)
-            * Filter[ry, rx, rc, ff].astype(out_dtype),
-            axis=[ry, rx, rc],
-        ),
-        name="Conv2dOutput",
-        tag="conv2d_nhwc",
-        attrs={"layout_free_placeholders": [Filter]},
+    return conv(
+        Input,
+        Filter,
+        stride,
+        padding,
+        dilation,
+        1,
+        "NHWC",
+        out_dtype,
+        auto_scheduler_rewritten_layout,
     )
 
-    if auto_scheduler_rewritten_layout:
-        Output = auto_scheduler.rewrite_compute_body(Output, auto_scheduler_rewritten_layout)
-
-    return Output
-
 
 def conv2d_NCHWc(data, kernel, stride, padding, dilation, layout, out_layout, out_dtype="float32"):
     """Conv2D operator for nChw[x]c layout.
@@ -758,6 +713,7 @@ def conv(
     groups: int,
     order: str,
     out_dtype: Union[str, None] = None,
+    auto_scheduler_rewritten_layout: Optional[str] = None,
 ):
     """Convolution operator in NCHW or NHWC layout.
 
@@ -796,6 +752,9 @@ def conv(
         Elements are converted to this type before elementwise multiplication
         and summation.
 
+    auto_scheduler_rewritten_layout: str
+        Layout from autoscheduler's layout rewritting.
+
     Returns
     -------
     Output : tvm.te.Tensor
@@ -840,6 +799,15 @@ def conv(
         permutation_to_kernel
     ].tolist()
 
+    # Autoscheduler may have messed with the input layout, so we extract the
+    # dimensions that it gives us
+    if auto_scheduler_rewritten_layout:
+        num_filter, _, *kernel_dimensions = auto_scheduler.get_shape_from_rewritten_layout(
+            auto_scheduler_rewritten_layout,
+            ["ff", "rc"] + [f"r{i}" for i in ["y", "x", "z"][: len(kernel_dimensions)]],
+        )
+        auto_scheduler.remove_index_check(filt)
+
     assert in_channel % groups == 0, "input channels must divide group size"
     assert num_filter % groups == 0, "output channels must divide group size"
 
@@ -858,15 +826,21 @@ def conv(
     pad_after = list(np.array([0, 0] + pad_end)[permutation_from])
     temp = pad(inp, pad_before, pad_after, name="pad_temp")
     rc = te.reduce_axis((0, in_channel // groups), name="rc")
-    rs = [te.reduce_axis((0, k), name=f"r{i}") for i, k in enumerate(kernel_dimensions)]
+    rs = [te.reduce_axis((0, k), name=f"r{i}") for i, k in zip(["y", "x", "z"], kernel_dimensions)]
 
     def compute(*args):
         nn, ff, *dim_indices = list(np.array(args)[permutation_to])
+
+        if groups == 1:
+            simplified_channel_index = rc
+        else:
+            simplified_channel_index = ff // (num_filter // groups) * (in_channel // groups) + rc
+
         return te.sum(
             temp.__getitem__(
                 tuple(
                     np.array(
-                        [nn, ff // (num_filter // groups) * (in_channel // groups) + rc]
+                        [nn, simplified_channel_index]
                         + [
                             di * stride + r * dil
                             for di, stride, r, dil in zip(dim_indices, strides, rs, dilations)
@@ -882,13 +856,20 @@ def compute(*args):
             axis=np.array([rc, *rs])[permutation_from_reductions].tolist(),
         )
 
-    return te.compute(
+    out = te.compute(
         list(np.array([batch, out_channel] + out_dimensions)[permutation_from]),
         compute,
         # tag is expected to be lowercase
         tag=f"{'group_' if groups > 1 else ''}conv{dim}d_{order.lower()}",
         name=f"{'group_' if groups > 1 else ''}conv{dim}d_{order.lower()}",
+        attrs={"layout_free_placeholders": [filt]},
+        varargs_names=list(np.array(["nn", "ff", "yy", "xx", "zz"])[permutation_from]),
     )
+    # if we used autoscheduler's changed layout we need to rewrite the ordering
+    # of the output dimensions
+    if auto_scheduler_rewritten_layout:
+        out = auto_scheduler.rewrite_compute_body(out, auto_scheduler_rewritten_layout)
+    return out
 
 
 def group_conv2d_nhwc(Input, Filter, stride, padding, dilation, groups, out_dtype=None):
diff --git a/python/tvm/topi/nn/conv3d.py b/python/tvm/topi/nn/conv3d.py
index 26795882973f..2915b886a5a8 100644
--- a/python/tvm/topi/nn/conv3d.py
+++ b/python/tvm/topi/nn/conv3d.py
@@ -18,15 +18,14 @@
 # pylint: disable=unused-argument, redefined-builtin, no-else-return
 """Conv3D operators"""
 import tvm
-from tvm import te, auto_scheduler
+from tvm import te
 
-from .pad import pad
-from .utils import get_pad_tuple3d
-from ..utils import simplify, get_const_tuple
+from ..utils import get_const_tuple
 from .winograd_util import winograd_transform_matrices
+from .conv2d import conv
 
 
-def conv3d_ncdhw(Input, Filter, stride, padding, dilation, out_dtype=None):
+def conv3d_ncdhw(Input, Filter, stride, padding, dilation, groups, out_dtype=None):
     """Conv3D operator in NCDHW layout.
 
     Parameters
@@ -46,62 +45,15 @@ def conv3d_ncdhw(Input, Filter, stride, padding, dilation, out_dtype=None):
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups.
+
     Returns
     -------
     Output : tvm.te.Tensor
         5-D with shape [batch, out_channel, out_depth, out_height, out_width]
     """
-    if out_dtype is None:
-        out_dtype = Input.dtype
-    assert isinstance(stride, int) or len(stride) == 3
-    assert isinstance(dilation, int) or len(dilation) == 3
-    if isinstance(stride, int):
-        stride_d = stride_h = stride_w = stride
-    else:
-        stride_d, stride_h, stride_w = stride
-
-    if isinstance(dilation, int):
-        dilation_d = dilation_h = dilation_w = dilation
-    else:
-        dilation_d, dilation_h, dilation_w = dilation
-
-    batch, in_channel, in_depth, in_height, in_width = Input.shape
-    num_filter, channel, kernel_d, kernel_h, kernel_w = Filter.shape
-    # compute the output shape
-    dilated_kernel_d = (kernel_d - 1) * dilation_d + 1
-    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
-    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
-    pad_front, pad_top, pad_left, pad_back, pad_down, pad_right = get_pad_tuple3d(
-        padding, (dilated_kernel_d, dilated_kernel_h, dilated_kernel_w)
-    )
-    out_channel = num_filter
-    out_depth = simplify((in_depth - dilated_kernel_d + pad_front + pad_back) // stride_d + 1)
-    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
-    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
-    # compute graph
-    pad_before = [0, 0, pad_front, pad_top, pad_left]
-    pad_after = [0, 0, pad_back, pad_down, pad_right]
-    temp = pad(Input, pad_before, pad_after, name="pad_temp")
-    rc = te.reduce_axis((0, in_channel), name="rc")
-    rz = te.reduce_axis((0, kernel_d), name="rz")
-    ry = te.reduce_axis((0, kernel_h), name="ry")
-    rx = te.reduce_axis((0, kernel_w), name="rx")
-
-    return te.compute(
-        (batch, out_channel, out_depth, out_height, out_width),
-        lambda nn, ff, zz, yy, xx: te.sum(
-            temp[
-                nn,
-                rc,
-                zz * stride_d + rz * dilation_d,
-                yy * stride_h + ry * dilation_h,
-                xx * stride_w + rx * dilation_w,
-            ].astype(out_dtype)
-            * Filter[ff, rc, rz, ry, rx].astype(out_dtype),
-            axis=[rc, rz, ry, rx],
-        ),
-        tag="conv3d_ncdhw",
-    )
+    return conv(Input, Filter, stride, padding, dilation, groups, "NCDHW", out_dtype)
 
 
 def conv3d_ndhwc(
@@ -110,6 +62,7 @@ def conv3d_ndhwc(
     stride,
     padding,
     dilation,
+    groups,
     out_dtype="float32",
     auto_scheduler_rewritten_layout="",
 ):
@@ -132,6 +85,9 @@ def conv3d_ndhwc(
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups.
+
     out_dtype: str = "float32",
         The type of output tensor
 
@@ -143,78 +99,18 @@ def conv3d_ndhwc(
     Output : tvm.te.Tensor
         5-D with shape [batch, out_depth, out_height, out_width, out_channel]
     """
-    assert isinstance(stride, int) or len(stride) == 3
-    assert isinstance(dilation, int) or len(dilation) == 3
-
-    if isinstance(stride, int):
-        stride_d = stride_h = stride_w = stride
-    else:
-        stride_d, stride_h, stride_w = stride
-
-    if isinstance(dilation, int):
-        dilation_d = dilation_h = dilation_w = dilation
-    else:
-        dilation_d, dilation_h, dilation_w = dilation
-
-    batch, in_depth, in_height, in_width, in_channel = Input.shape
-
-    if auto_scheduler_rewritten_layout:
-        # Infer shape for the rewritten layout
-        (
-            kernel_d,
-            kernel_h,
-            kernel_w,
-            channel,
-            num_filter,
-        ) = auto_scheduler.get_shape_from_rewritten_layout(
-            auto_scheduler_rewritten_layout, ["rd", "rh", "rw", "rc", "cc"]
-        )
-        auto_scheduler.remove_index_check(Filter)
-    else:
-        kernel_d, kernel_h, kernel_w, channel, num_filter = Filter.shape
-
-    # compute the output shape
-    dilated_kernel_d = (kernel_d - 1) * dilation_d + 1
-    dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
-    dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
-
-    pad_front, pad_top, pad_left, pad_back, pad_down, pad_right = get_pad_tuple3d(
-        padding, (dilated_kernel_d, dilated_kernel_h, dilated_kernel_w)
-    )
-    out_channel = num_filter
-    out_depth = simplify((in_depth - dilated_kernel_d + pad_front + pad_back) // stride_d + 1)
-    out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
-    out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
-    pad_before = [0, pad_front, pad_top, pad_left, 0]
-    pad_after = [0, pad_back, pad_down, pad_right, 0]
-    PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
-    rd = te.reduce_axis((0, kernel_d), name="rd")
-    rh = te.reduce_axis((0, kernel_h), name="rh")
-    rw = te.reduce_axis((0, kernel_w), name="rw")
-    rc = te.reduce_axis((0, in_channel), name="rc")
-    Output = te.compute(
-        (batch, out_depth, out_height, out_width, out_channel),
-        lambda nn, dd, hh, ww, cc: te.sum(
-            PaddedInput[
-                nn,
-                dd * stride_d + rd * dilation_d,
-                hh * stride_h + rh * dilation_h,
-                ww * stride_w + rw * dilation_w,
-                rc,
-            ].astype(out_dtype)
-            * Filter[rd, rh, rw, rc, cc].astype(out_dtype),
-            axis=[rd, rh, rw, rc],
-        ),
-        name="Conv3dOutput",
-        tag="conv3d_ndhwc",
-        attrs={"layout_free_placeholders": [Filter]},
+    return conv(
+        Input,
+        Filter,
+        stride,
+        padding,
+        dilation,
+        groups,
+        "NDHWC",
+        out_dtype,
+        auto_scheduler_rewritten_layout,
     )
 
-    if auto_scheduler_rewritten_layout:
-        Output = auto_scheduler.rewrite_compute_body(Output, auto_scheduler_rewritten_layout)
-
-    return Output
-
 
 def conv3d_winograd_weight_transform(kernel, tile_size):
     """Weight transformation for 3D winograd
diff --git a/python/tvm/topi/nn/dense.py b/python/tvm/topi/nn/dense.py
index 58c458a7d676..69fac92c7cff 100644
--- a/python/tvm/topi/nn/dense.py
+++ b/python/tvm/topi/nn/dense.py
@@ -81,7 +81,9 @@ def matmul(
         out_dim, red_dim = tensor_b.shape
     else:
         red_dim, out_dim = tensor_b.shape
-    assert in_dim == red_dim
+
+    # cmp should be done by values
+    assert int(in_dim) == int(red_dim)
 
     k = te.reduce_axis((0, in_dim), name="k")
     if (transpose_a, transpose_b) == (True, True):
@@ -274,3 +276,22 @@ def dense_alter_layout(attrs, inputs, tinfos, out_type):
     """
     # not to change by default
     return None
+
+
+@tvm.target.generic_func
+def batch_matmul_legalize(attrs, inputs, types):
+    """Legalizes batch_matmul op.
+    Parameters
+    ----------
+    attrs : tvm.ir.Attrs
+        Attributes of current batch_matmul
+    inputs : list of tvm.relay.Expr
+        The args of the Relay expr to be legalized
+    types : list of types
+        List of input and output types
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The legalized expr
+    """
+    return None
diff --git a/python/tvm/topi/testing/conv3d_ndhwc_python.py b/python/tvm/topi/testing/conv3d_ndhwc_python.py
index 46f04f630863..8ca1f7da0bcc 100644
--- a/python/tvm/topi/testing/conv3d_ndhwc_python.py
+++ b/python/tvm/topi/testing/conv3d_ndhwc_python.py
@@ -21,7 +21,7 @@
 from tvm.topi.nn.utils import get_pad_tuple3d
 
 
-def conv3d_ndhwc_python(a_np, w_np, stride, padding):
+def _conv3d_ndhwc_python(a_np, w_np, stride, padding):
     """Convolution 3D operator in NDHWC layout.
 
     Parameters
@@ -37,8 +37,6 @@ def conv3d_ndhwc_python(a_np, w_np, stride, padding):
 
     padding : int or str or a list/tuple of three ints
         Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
-    groups : int
-        Number of groups
 
     Returns
     -------
@@ -66,13 +64,15 @@ def conv3d_ndhwc_python(a_np, w_np, stride, padding):
     # change the layout from NHWC to NCHW
     at = a_np.transpose((0, 4, 1, 2, 3))
     wt = w_np.transpose((4, 3, 0, 1, 2))
-    bt = np.zeros((batch, out_channel, out_depth, out_height, out_width))
+    bt = np.zeros((batch, out_channel, out_depth, out_height, out_width), dtype=a_np.dtype)
     # computation
     for n in range(batch):
         for f in range(out_channel):
             for c in range(in_channel):
                 if pad_d > 0 or pad_h > 0 or pad_w > 0:
-                    apad = np.zeros((in_depth + pad_d, in_height + pad_h, in_width + pad_w))
+                    apad = np.zeros(
+                        (in_depth + pad_d, in_height + pad_h, in_width + pad_w), dtype=a_np.dtype
+                    )
                     apad[
                         pad_front : pad_front + in_depth,
                         pad_top : pad_top + in_height,
@@ -83,3 +83,38 @@ def conv3d_ndhwc_python(a_np, w_np, stride, padding):
                 out = scipy.signal.convolve(apad, np.flip(wt[f, c]), mode="valid")
                 bt[n, f] += out[::stride_d, ::stride_h, ::stride_w]
     return bt.transpose((0, 2, 3, 4, 1))
+
+
+def conv3d_ndhwc_python(a_np, w_np, stride, padding, groups=1):
+    """Convolution 3D operator in NDHWC layout.
+
+    Parameters
+    ----------
+    a_np : numpy.ndarray
+        5-D with shape [batch, in_channel, in_depth, in_height, in_width]
+
+    w_np : numpy.ndarray
+        5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
+
+    stride : int or a list/tuple of three ints
+        Stride size, or [stride_depth, stride_height, stride_width]
+
+    padding : int or str or a list/tuple of three ints
+        Padding size, or ['VALID', 'SAME'], or [pad_depth, pad_height, pad_width]
+
+    groups : int
+        Number of groups
+
+    Returns
+    -------
+    b_np : np.ndarray
+        5-D with shape [batch, out_channel, out_depth, out_height, out_width]
+    """
+    a_slices = np.array_split(a_np, groups, axis=4)
+    w_slices = np.array_split(w_np, groups, axis=4)
+    b_slices = [
+        _conv3d_ndhwc_python(a_slice, w_slice, stride, padding)
+        for a_slice, w_slice in zip(a_slices, w_slices)
+    ]
+    b_np = np.concatenate(b_slices, axis=4)
+    return b_np
diff --git a/python/tvm/topi/x86/batch_matmul.py b/python/tvm/topi/x86/batch_matmul.py
index 917734b13e78..b446c1f0115c 100644
--- a/python/tvm/topi/x86/batch_matmul.py
+++ b/python/tvm/topi/x86/batch_matmul.py
@@ -16,12 +16,67 @@
 # under the License.
 # pylint: disable=invalid-name,too-many-locals,unused-variable
 """x86 batch_matmul operators"""
+import tvm
 from tvm import te
 from tvm import autotvm
 from tvm.autotvm.task.space import SplitEntity
 from tvm.contrib import cblas, mkl
 from .. import generic, nn
+from ..transform import layout_transform
 from ..utils import traverse_inline, get_const_tuple, get_max_power2_factor
+from .utils import target_has_vnni
+from .dense import dense_vnni_schedule
+from .injective import schedule_injective_from_existing
+
+
+def batch_matmul_vnni_compute(cfg, x, y):
+    """Compute for uint8 x int8 -> int32 batch_matmul"""
+    batch, m, k = x.shape
+    packed_y_layout = "BNK16n4k"
+    packed_y = layout_transform(y, "BNK", packed_y_layout)
+    _, n_o, _, n_i, _ = packed_y.shape
+    ak = te.reduce_axis((0, k), name="k")
+
+    z = te.compute(
+        (batch, m, n_o * n_i),
+        lambda b, i, j: te.sum(
+            x[b, i, ak].astype("int32")
+            * packed_y[b, tvm.tir.indexdiv(j, 16), tvm.tir.indexdiv(ak, 4), j % 16, ak % 4].astype(
+                "int32"
+            ),
+            axis=ak,
+        ),
+        tag="batch_matmul_vnni",
+    )
+
+    _, a_y, _ = z.op.axis
+    cfg.define_split("tile_y", a_y, num_outputs=2)
+    cfg.define_knob("layout_trans_compute_root", [0, 1])
+
+    return z
+
+
+def batch_matmul_vnni_schedule(cfg, s, C, O, layout_trans):
+    """Schedule batch_matmul compute using VNNI vpdpbusd instruction"""
+    # C: The output of batched GEMM
+    # O: The output of the fused op
+
+    # Schedule the GEMM part
+    s, fused_inner = dense_vnni_schedule(cfg, s, C, O, do_parallel=False)
+    # Parallelize over batch
+    fused = s[O].fuse(O.op.axis[0], fused_inner)
+    s[O].parallel(fused)
+
+    if cfg["layout_trans_compute_root"].val:
+        s[layout_trans].compute_root()
+        schedule_injective_from_existing(s, layout_trans)
+    else:
+        s[layout_trans].compute_at(s[O], fused)
+        _, _, _, ni, ki = s[layout_trans].op.axis
+        s[layout_trans].vectorize(ki)
+        s[layout_trans].unroll(ni)
+
+    return s
 
 
 @autotvm.register_topi_compute("batch_matmul.x86")
@@ -62,6 +117,19 @@ def batch_matmul(
     output : tvm.te.Tensor
         3-D with shape [batch, M, N]
     """
+    mcpu = tvm.target.Target.current().mcpu
+
+    if (
+        not transpose_a
+        and transpose_b
+        and target_has_vnni(mcpu)
+        and tensor_a.dtype == "uint8"
+        and tensor_b.dtype == "int8"
+        and tensor_b.shape[-2] % 16 == 0
+        and tensor_b.shape[-1] % 4 == 0
+    ):
+        return batch_matmul_vnni_compute(cfg, tensor_a, tensor_b)
+
     if cfg.is_fallback:
         if transpose_a:
             _, K, M = get_const_tuple(tensor_a.shape)
@@ -102,7 +170,10 @@ def schedule_batch_matmul(cfg, outs):
     s = te.create_schedule([x.op for x in outs])
 
     def _callback(op):
-        if "batch_matmul" in op.tag:
+        if "batch_matmul_vnni" in op.tag:
+            layout_trans = op.input_tensors[1]
+            batch_matmul_vnni_schedule(cfg, s, op.output(0), outs[0], layout_trans)
+        elif "batch_matmul" in op.tag:
             C = op.output(0)
             A, B = op.input_tensors
             if len(B.op.input_tensors) == 1 and B.op.input_tensors[0] == A:
diff --git a/python/tvm/topi/x86/conv2d_alter_op.py b/python/tvm/topi/x86/conv2d_alter_op.py
index 3f2df655a615..9234581f1d5b 100644
--- a/python/tvm/topi/x86/conv2d_alter_op.py
+++ b/python/tvm/topi/x86/conv2d_alter_op.py
@@ -28,7 +28,7 @@
 from .conv2d_int8 import is_int8_hw_support, _get_default_config_int8
 from ..utils import get_const_tuple
 from ..nn import conv2d_legalize, conv2d_alter_layout
-from ..nn.utils import get_pad_tuple
+from ..generic.conv2d import conv2d_alter_int8_common
 
 logger = logging.getLogger("topi")
 
@@ -164,25 +164,12 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         batch_size, in_channel, height, width = get_const_tuple(data_tensor.shape)
         out_channel, channel_multiplier, kh, kw = get_const_tuple(kernel_tensor.shape)
         ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
-        n_elems = 4
-
-        # convert kernel data layout from 4D to 7D
-        data_expr, kernel_expr = inputs
-        kernel_IHWO = relay.transpose(kernel_expr, axes=(1, 2, 3, 0))
-        kernel_IHWOo = relay.reshape(kernel_IHWO, (in_channel, kh, kw, out_channel // oc_bn, oc_bn))
-        kernel_OHWoI = relay.transpose(kernel_IHWOo, axes=(3, 1, 2, 4, 0))
-        kernel_OHWoIi = relay.reshape(
-            kernel_OHWoI, (out_channel // oc_bn, kh, kw, oc_bn, in_channel // ic_bn, ic_bn)
-        )
-        kernel_OHWoIie = relay.reshape(
-            kernel_OHWoIi,
-            (out_channel // oc_bn, kh, kw, oc_bn, in_channel // ic_bn, ic_bn // n_elems, n_elems),
-        )
-        kernel_OIHWioe = relay.transpose(kernel_OHWoIie, axes=(0, 4, 1, 2, 5, 3, 6))
 
         # update new attrs
+        n_elems = 4
         new_attrs["channels"] = out_channel
         new_attrs["data_layout"] = "NCHW%dc" % ic_bn
+        new_attrs["kernel_layout"] = "OIHW{:n}i{:n}o{:n}i".format(ic_bn // n_elems, oc_bn, n_elems)
         new_attrs["out_layout"] = "NCHW%dc" % oc_bn
 
         # Store altered operator's config.
@@ -208,7 +195,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         )
         dispatch_ctx.update(target, new_workload, cfg)
 
-        return relay.nn.contrib_conv2d_nchwc(data_expr, kernel_OIHWioe, **new_attrs)
+        return relay.nn.contrib_conv2d_nchwc(*inputs, **new_attrs)
 
     if topi_tmpl == "depthwise_conv2d_NCHWc.x86":
         if data_layout == "NCHW" and kernel_layout == "OIHW":
@@ -283,17 +270,6 @@ def _conv2d_legalize(attrs, inputs, arg_types):
     result : tvm.relay.Expr
         The legalized expr
     """
-
-    # Dilation not supported yet. Return None if dilation is not (1, 1)
-    dilation = attrs.get_int_tuple("dilation")
-    if not (dilation[0] == 1 and dilation[1] == 1):
-        return None
-
-    # No legalization for depthwise convolutions yet.
-    groups = attrs.get_int("groups")
-    if groups != 1:
-        return None
-
     # Collect the input tensors.
     data_tensor, kernel_tensor = arg_types[0], arg_types[1]
     data_dtype = data_tensor.dtype
@@ -305,108 +281,11 @@ def _conv2d_legalize(attrs, inputs, arg_types):
     # Collect the input exprs.
     data, kernel = inputs
 
-    # Get the conv attrs
-    new_attrs = {k: attrs[k] for k in attrs.keys()}
-
-    is_int8_inputs = False
-    # If both the inputs are int8, we can add 128 to make the input dtype uint8, and then adjust the
-    # output. This will help picking up Intel VNNI instructions.
-    # Original --> C = A (conv) B
-    # A and B are int8
-    #   C = (A + 128 - 128) (conv) B
-    #   C = (A' conv B) - 128 (conv) B
-    # where A' = A + 128
-    # and 128 (conv) B is basically a reduce on CRS axis for weights.
+    # Intel vector intructions require data and kernel to have different dtypes.
     if data_tensor.dtype == "int8" and kernel_tensor.dtype == "int8":
-        is_int8_inputs = True
-        padding = attrs.get_int_tuple("padding")
-        kh, kw = attrs.get_int_tuple("kernel_size")
-        pt, pl, pb, pr = get_pad_tuple(padding, (kh, kw))
-
-        if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
-            adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(0, 1, 2))
-            pad_width = ((0, 0), (pt, pb), (pl, pr), (0, 0))
-        elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
-            pad_width = ((0, 0), (0, 0), (pt, pb), (pl, pr))
-            adjust_shift = relay.sum(relay.cast(kernel, dtype="int32"), axis=(1, 2, 3))
-            adjust_shift = relay.expand_dims(adjust_shift, axis=1, num_newaxis=2)
-        else:
-            return None
-
-        data = relay.cast(data, "int32")
-        data = relay.add(data, relay.const(128, "int32"))
-        data = relay.cast(data, "uint8")
-
-        # Do external padding as pad value has to be 128.
-        if any(padding):
-            data = relay.nn.pad(data, pad_width=pad_width, pad_value=128)
-        new_attrs["padding"] = (0, 0)
-
-        # The data type is now shifted to uint8
         data_dtype = "uint8"
-
-        # Multiply 128 to adjust shift.
-        adjust_shift = relay.multiply(adjust_shift, relay.const(128, "int32"))
-
-    # Legalize if the datatypes are suitable for fast Int8 instructions.  Int8 instructions require
-    # input channel to be a multiple of 4 and output channels to be a multiple of 16. For input
-    # channels, we pad both the inputs and weights input channels. For output channels, we pad the
-    # weight and stride_slice the output.
     if is_int8_hw_support(data_dtype, kernel_dtype):
-        # Flags to remember if the expr is modified
-        ic_modified = False
-        oc_modified = False
-
-        # Find the value of input and output channel.
-        in_channel = -1
-        out_channel = -1
-        if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
-            in_channel = data_tensor.shape[3].value
-            out_channel = kernel_tensor.shape[3].value
-        elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
-            in_channel = data_tensor.shape[1].value
-            out_channel = kernel_tensor.shape[0].value
-        else:
-            return None
-
-        if in_channel % 4 != 0:
-            new_in_channel = ((in_channel + 4) // 4) * 4
-            diff = new_in_channel - in_channel
-            if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
-                data = relay.nn.pad(data, pad_width=((0, 0), (0, 0), (0, 0), (0, diff)))
-                kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, diff), (0, 0)))
-                ic_modified = True
-            elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
-                pad_width = ((0, 0), (0, diff), (0, 0), (0, 0))
-                data = relay.nn.pad(data, pad_width=pad_width)
-                kernel = relay.nn.pad(kernel, pad_width=pad_width)
-                ic_modified = True
-            else:
-                return None
-
-        new_out_channel = out_channel
-        if out_channel % 16 != 0:
-            new_out_channel = ((out_channel + 16) // 16) * 16
-            diff = new_out_channel - out_channel
-            if attrs["data_layout"] == "NHWC" and attrs["kernel_layout"] == "HWIO":
-                kernel = relay.nn.pad(kernel, pad_width=((0, 0), (0, 0), (0, 0), (0, diff)))
-                oc_modified = True
-            elif attrs["data_layout"] == "NCHW" and attrs["kernel_layout"] == "OIHW":
-                kernel = relay.nn.pad(kernel, pad_width=((0, diff), (0, 0), (0, 0), (0, 0)))
-                oc_modified = True
-            else:
-                return None
-
-        if oc_modified:
-            new_attrs["channels"] = new_out_channel
-            out = tvm.relay.nn.conv2d(data, kernel, **new_attrs)
-            original_out_shape = [x.value for x in output_tensor.shape]
-            out = relay.strided_slice(out, begin=[0, 0, 0, 0], end=original_out_shape)
-        else:
-            out = relay.nn.conv2d(data, kernel, **new_attrs)
-
-        if is_int8_inputs:
-            out = relay.subtract(out, adjust_shift)
-
-        return out
+        return conv2d_alter_int8_common(
+            data, data_tensor, kernel, kernel_tensor, output_tensor, attrs, data_dtype, 4, 16
+        )
     return None
diff --git a/python/tvm/topi/x86/conv3d.py b/python/tvm/topi/x86/conv3d.py
index c4194167ce47..557418655cfe 100644
--- a/python/tvm/topi/x86/conv3d.py
+++ b/python/tvm/topi/x86/conv3d.py
@@ -53,7 +53,7 @@
 
 
 @autotvm.register_topi_compute("conv3d_ndhwc.x86")
-def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
+def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype):
     """3D convolution forward operator.
 
     Parameters
@@ -74,6 +74,9 @@ def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups
+
     Returns
     -------
     output : tvm.te.Tensor
@@ -84,14 +87,14 @@ def conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     strides = strides if isinstance(strides, (tuple, list)) else (strides, strides, strides)
     dilation = dilation if isinstance(dilation, (tuple, list)) else (dilation, dilation, dilation)
 
-    _create_tuning_space(cfg, data, kernel, strides, padding, dilation, layout)
+    _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout)
     if cfg.is_fallback:
-        _get_default_config(cfg, data, kernel, strides, padding, out_dtype, layout)
-    return _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype)
+        _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout)
+    return _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype)
 
 
 @autotvm.register_topi_compute("conv3d_ncdhw.x86")
-def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, out_dtype):
+def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, groups, out_dtype):
     """3D convolution forward operator.
 
     Parameters
@@ -112,20 +115,24 @@ def conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, out_dtype):
     dilation: int or a list/tuple of three ints
         dilation size, or [dilation_depth, dilation_height, dilation_width]
 
+    groups: int
+        Number of groups
+
     Returns
     -------
     output : tvm.te.Tensor
         5-D with shape [batch, out_channel, out_depth, out_height, out_width] for NCDHW layout
     """
+    # assert groups == 1, "conv3d_ncdhw.x86 does not support groups"
     layout = "NCDHW"
     out_dtype = data.dtype if out_dtype is None else out_dtype
     strides = strides if isinstance(strides, (tuple, list)) else (strides, strides, strides)
     dilation = dilation if isinstance(dilation, (tuple, list)) else (dilation, dilation, dilation)
 
-    _create_tuning_space(cfg, data, kernel, strides, padding, dilation, layout)
+    _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout)
     if cfg.is_fallback:
-        _get_default_config(cfg, data, kernel, strides, padding, out_dtype, layout)
-    return _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dtype)
+        _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout)
+    return _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, groups, out_dtype)
 
 
 @autotvm.register_topi_schedule("conv3d_ndhwc.x86")
@@ -208,7 +215,7 @@ def _traverse(op):
     return s
 
 
-def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
+def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, groups, out_dtype):
     out_dtype = data.dtype if out_dtype is None else out_dtype
 
     assert isinstance(dilation, int) or len(dilation) == 3
@@ -221,6 +228,9 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     batch_size, in_depth, in_height, in_width, in_channel = get_const_tuple(data.shape)
     kernel_depth, kernel_height, kernel_width, _, num_filter = get_const_tuple(kernel.shape)
 
+    assert in_channel % groups == 0, "input channels must be a multiple of group size"
+    assert num_filter % groups == 0, "number of filters must be a multiple of group size"
+
     dilated_kernel_d = (kernel_depth - 1) * dilation_d + 1
     dilated_kernel_h = (kernel_height - 1) * dilation_h + 1
     dilated_kernel_w = (kernel_width - 1) * dilation_w + 1
@@ -255,6 +265,8 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
 
     # fetch schedule
     ic_bn, oc_bn = cfg["tile_ic"].size[-1], cfg["tile_oc"].size[-1]
+    assert groups == 1 or ic_bn <= groups
+    assert groups == 1 or oc_bn <= groups
     shape = (batch_size, in_channel // ic_bn, pad_depth, pad_height, ic_bn, pad_width)
     data_vec = te.compute(
         shape, lambda n, C, d, h, c, w: data_pad[n, d, h, w, C * ic_bn + c], name="data_vec"
@@ -263,7 +275,7 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     # pack kernel
     shape = (
         num_filter // oc_bn,
-        in_channel // ic_bn,
+        in_channel // groups // ic_bn,
         kernel_depth,
         kernel_height,
         kernel_width,
@@ -280,7 +292,7 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     oshape = (batch_size, num_filter // oc_bn, out_depth, out_height, out_width, oc_bn)
     unpack_shape = (batch_size, out_depth, out_height, out_width, num_filter)
 
-    ic = te.reduce_axis((0, in_channel), name="ic")
+    ic = te.reduce_axis((0, in_channel // groups), name="ic")
     kh = te.reduce_axis((0, kernel_height), name="kh")
     kw = te.reduce_axis((0, kernel_width), name="kw")
     kd = te.reduce_axis((0, kernel_depth), name="kd")
@@ -292,10 +304,18 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
         lambda n, oc_chunk, od, oh, ow, oc_block: te.sum(
             data_vec[
                 n,
-                idxdiv(ic, ic_bn),
+                idxdiv(
+                    (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups)
+                    + ic,
+                    ic_bn,
+                ),
                 od * DSTR + kd * dilation_d,
                 oh * HSTR + kh * dilation_h,
-                idxmod(ic, ic_bn),
+                idxmod(
+                    (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups)
+                    + ic,
+                    ic_bn,
+                ),
                 ow * WSTR + kw * dilation_w,
             ].astype(out_dtype)
             * kernel_vec[
@@ -316,7 +336,7 @@ def _conv3d_ndhwc(cfg, data, kernel, strides, padding, dilation, out_dtype):
     return conv_unpacked
 
 
-def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dtype):
+def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, groups, out_dtype):
     out_dtype = data.dtype if out_dtype is None else out_dtype
 
     assert isinstance(dilation, int) or len(dilation) == 3
@@ -372,7 +392,7 @@ def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dty
     # pack kernel
     shape = (
         num_filter // oc_bn,
-        in_channel // ic_bn,
+        in_channel // groups // ic_bn,
         kernel_depth,
         kernel_height,
         kernel_width,
@@ -389,7 +409,7 @@ def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dty
     oshape = (batch_size, num_filter // oc_bn, out_depth, out_height, out_width, oc_bn)
     unpack_shape = (batch_size, num_filter, out_depth, out_height, out_width)
 
-    ic = te.reduce_axis((0, in_channel), name="ic")
+    ic = te.reduce_axis((0, in_channel // groups), name="ic")
     kh = te.reduce_axis((0, kernel_height), name="kh")
     kw = te.reduce_axis((0, kernel_width), name="kw")
     kd = te.reduce_axis((0, kernel_depth), name="kd")
@@ -401,10 +421,18 @@ def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dty
         lambda n, oc_chunk, od, oh, ow, oc_block: te.sum(
             data_vec[
                 n,
-                idxdiv(ic, ic_bn),
+                idxdiv(
+                    (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups)
+                    + ic,
+                    ic_bn,
+                ),
                 od * DSTR + kd * dilation_d,
                 oh * HSTR + kh * dilation_h,
-                idxmod(ic, ic_bn),
+                idxmod(
+                    (oc_chunk * oc_bn + oc_block) // (num_filter // groups) * (in_channel // groups)
+                    + ic,
+                    ic_bn,
+                ),
                 ow * WSTR + kw * dilation_w,
             ].astype(out_dtype)
             * kernel_vec[
@@ -425,7 +453,7 @@ def _conv3d_ncdhw(cfg, data, kernel, strides, padding, dilation, layout, out_dty
     return conv_unpacked
 
 
-def _create_tuning_space(cfg, data, kernel, strides, padding, dilation, layout):
+def _create_tuning_space(cfg, data, kernel, strides, padding, dilation, groups, layout):
     """Create schedule configuration from input arguments"""
     dshape = get_const_tuple(data.shape)
     kshape = get_const_tuple(kernel.shape)
@@ -452,7 +480,7 @@ def _create_tuning_space(cfg, data, kernel, strides, padding, dilation, layout):
     cfg.define_knob("unroll_kw", [True, False])
 
 
-def _get_default_config(cfg, data, kernel, strides, padding, out_dtype, layout):
+def _get_default_config(cfg, data, kernel, strides, padding, groups, out_dtype, layout):
     """
     Get default schedule config for the workload
     """
@@ -466,11 +494,11 @@ def _get_default_config(cfg, data, kernel, strides, padding, out_dtype, layout):
         else:
             static_data_shape.append(dim)
     data = te.placeholder(static_data_shape, dtype=data.dtype)
-    wkl = _get_conv3d_workload(data, kernel, strides, padding, out_dtype, layout)
+    wkl = _get_conv3d_workload(data, kernel, strides, padding, groups, out_dtype, layout)
     _fallback_schedule(cfg, wkl)
 
 
-def _get_conv3d_workload(data, kernel, stride, padding, out_dtype, data_layout="NCHW"):
+def _get_conv3d_workload(data, kernel, stride, padding, groups, out_dtype, data_layout="NCHW"):
     """Get the workload structure."""
     if data_layout == "NCDHW":
         _, CI, ID, IH, IW = get_const_tuple(data.shape)
@@ -487,7 +515,6 @@ def _get_conv3d_workload(data, kernel, stride, padding, out_dtype, data_layout="
     DPAD = pad_front + pad_back
     HPAD = pad_top + pad_down
     WPAD = pad_left + pad_right
-    GRPS = CI // CIG
     if isinstance(stride, (tuple, list)):
         DSTR, HSTR, WSTR = stride
     else:
@@ -505,7 +532,7 @@ def _get_conv3d_workload(data, kernel, stride, padding, out_dtype, data_layout="
         IH,
         IW,
         CI,
-        GRPS,
+        groups,
         CO,
         KD,
         KH,
diff --git a/python/tvm/topi/x86/conv3d_transpose.py b/python/tvm/topi/x86/conv3d_transpose.py
index e743f02fb063..cb814a29d60d 100644
--- a/python/tvm/topi/x86/conv3d_transpose.py
+++ b/python/tvm/topi/x86/conv3d_transpose.py
@@ -30,7 +30,7 @@ def conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_pad
     )
 
     # reuse conv3d_ncdhw implementation
-    return conv3d_ncdhw(data_pad, kernel_transform, (1, 1, 1), (0, 0, 0), (1, 1, 1), out_dtype)
+    return conv3d_ncdhw(data_pad, kernel_transform, (1, 1, 1), (0, 0, 0), (1, 1, 1), 1, out_dtype)
 
 
 def schedule_conv3d_transpose_ncdhw(outs):
diff --git a/python/tvm/topi/x86/dense.py b/python/tvm/topi/x86/dense.py
index 9799ec02d644..c8574b971003 100644
--- a/python/tvm/topi/x86/dense.py
+++ b/python/tvm/topi/x86/dense.py
@@ -29,6 +29,7 @@
 from .utils import get_simd_32bit_lanes
 from .. import generic, tag
 from ..utils import traverse_inline, get_const_tuple
+from .tensor_intrin import dot_16x1x16_uint8_int8_int32_cascadelake
 
 
 def _schedule_dense_pack_template(cfg, s, C, O):
@@ -279,6 +280,100 @@ def _callback(op):
     return s
 
 
+def dense_vnni_compute(cfg, X, packed_w, bias=None):
+    """Compute for uint8 x int8 -> int32 dense"""
+    m, k = X.shape
+    n_o, _, n_i, _ = packed_w.shape
+    ak = te.reduce_axis((0, k), name="k")
+
+    C = te.compute(
+        (m, n_o * n_i),
+        lambda i, j: te.sum(
+            X[i, ak].astype("int32")
+            * packed_w[tvm.tir.indexdiv(j, 16), tvm.tir.indexdiv(ak, 4), j % 16, ak % 4].astype(
+                "int32"
+            ),
+            axis=ak,
+        ),
+        tag="dense_vnni",
+    )
+
+    if bias is not None:
+        C = te.compute(C.shape, lambda i, j: C[i, j] + bias[j], tag=tag.BROADCAST)
+
+    a_y, _ = C.op.axis
+    cfg.define_split("tile_y", a_y, num_outputs=2)
+
+    return C
+
+
+def dense_vnni_schedule(cfg, s, C, O, do_parallel=True):
+    """Schedule dense compute using VNNI vpdpbusd instruction"""
+    # C: The output of GEMM
+    # O: The output of the fused op
+    def split_y(out):
+        default_y_split_factor = 32
+        a_y = out.op.axis[-2]
+
+        if cfg.is_fallback:
+            return s[out].split(a_y, factor=default_y_split_factor)
+
+        return cfg["tile_y"].apply(s, out, a_y)
+
+    (a_k,) = C.op.reduce_axis
+
+    a_yo, a_yi = split_y(C)
+    a_xo, a_xi = s[C].split(C.op.axis[-1], factor=16)
+    a_ko, a_ki = s[C].split(a_k, factor=4)
+
+    s[C].reorder(a_yo, a_xo, a_yi, a_ko, a_xi, a_ki)
+
+    pc = dot_16x1x16_uint8_int8_int32_cascadelake()
+    s[C].tensorize(a_xi, pc)
+
+    if C == O:
+        fused = s[O].fuse(a_yo, a_xo)
+    else:
+        a_yo, a_yi = split_y(O)
+        a_xo, a_xi = s[O].split(O.op.axis[-1], factor=16)
+
+        s[O].reorder(a_yo, a_xo, a_yi, a_xi)
+        s[O].vectorize(a_xi)
+        s[C].compute_at(s[O], a_yi)
+
+        fused = s[O].fuse(a_yo, a_xo)
+
+    if do_parallel:
+        s[O].parallel(fused)
+
+    return s, fused
+
+
+@autotvm.register_topi_compute("dense_vnni.x86")
+def dense_vnni(cfg, data, weight, bias=None, out_dtype=None):
+    """Compute for uint8 x int8 -> int32 dense"""
+    if out_dtype is None:
+        out_dtype = data.dtype
+    assert len(weight.shape) == 4
+    assert data.dtype == "uint8" and weight.dtype == "int8"
+    _, _, n_inner, k_inner = get_const_tuple(weight.shape)  # out_dim
+    assert n_inner == 16 and k_inner == 4
+    return dense_vnni_compute(cfg, data, weight, bias)
+
+
+@autotvm.register_topi_schedule("dense_vnni.x86")
+def schedule_dense_vnni(cfg, outs):
+    """Create a schedule for dense_vnni"""
+    s = te.create_schedule([x.op for x in outs])
+
+    def _callback(op):
+        if "dense_vnni" in op.tag:
+            dense_vnni_schedule(cfg, s, op.output(0), outs[0])
+
+    traverse_inline(s, outs[0].op, _callback)
+    return s
+
+
 def matmul_blas_common(cfg, tensor_a, tensor_b, bias, out_dtype, transpose_a, transpose_b, lib):
     """Compute matmul/dense using a BLAS library"""
     M, K = get_const_tuple(tensor_a.shape)
diff --git a/python/tvm/topi/x86/dense_alter_op.py b/python/tvm/topi/x86/dense_alter_op.py
index 1d64261a50d7..0b195f487b7f 100644
--- a/python/tvm/topi/x86/dense_alter_op.py
+++ b/python/tvm/topi/x86/dense_alter_op.py
@@ -24,6 +24,19 @@
 from .dense import _default_dense_pack_config
 from ..utils import get_const_tuple
 from ..nn import dense_alter_layout
+from .utils import target_has_vnni
+from .. import nn
+
+
+def check_vnni_applicable(x, y):
+    mcpu = tvm.target.Target.current().mcpu
+    return (
+        target_has_vnni(mcpu)
+        and "int8" in x.dtype
+        and "int8" in y.dtype
+        and y.shape[-2] % 16 == 0
+        and y.shape[-1] % 4 == 0
+    )
 
 
 @dense_alter_layout.register(["cpu", "arm_cpu"])
@@ -35,7 +48,11 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type):
     M, K = get_const_tuple(data_tensor.shape)
     N, _ = get_const_tuple(weight_tensor.shape)
 
-    impl, outs = relay.backend.te_compiler.select_implementation(
+    if check_vnni_applicable(data_tensor, weight_tensor) and data_tensor.dtype == "uint8":
+        weight_layout = "NC16n4c"
+        return relay.nn.contrib_dense_pack(inputs[0], inputs[1], weight_layout, None, out_dtype)
+
+    _, outs = relay.backend.te_compiler.select_implementation(
         relay.op.get("nn.dense"), attrs, tinfos, out_type, target
     )
     workload = autotvm.task.get_workload(outs)
@@ -66,3 +83,37 @@ def _alter_dense_layout(attrs, inputs, tinfos, out_type):
             return relay.nn.contrib_dense_pack(inputs[0], inputs[1], weight_layout, None, out_dtype)
 
     return None
+
+
+def vnni_legalize(inputs, arg_types, op, attrs, need_expand=False):
+    """Legalizes s8, s8 -> s32 GEMM op for VNNI."""
+    if check_vnni_applicable(arg_types[0], arg_types[1]) and arg_types[0].dtype == "int8":
+        x, y = inputs
+        x = relay.cast(x, "int32")
+        x = relay.add(x, relay.const(128, "int32"))
+        x = relay.cast(x, "uint8")
+
+        adjust_shift = relay.const(128, "int32") * relay.sum(relay.cast(y, "int32"), axis=[-1])
+
+        if need_expand:
+            adjust_shift = relay.expand_dims(adjust_shift, axis=1)
+
+        out = op(x, y, **attrs)
+
+        return relay.subtract(out, adjust_shift)
+
+    return None
+
+
+@nn.dense_legalize.register("cpu")
+def _dense_legalize(attrs, inputs, arg_types):
+    """Legalizes s8, s8 -> s32 dense for VNNI."""
+    return vnni_legalize(inputs, arg_types, relay.nn.dense, attrs)
+
+
+@nn.batch_matmul_legalize.register("cpu")
+def _batch_matmul_legalize(attrs, inputs, arg_types):
+    """Legalizes s8, s8 -> s32 batch_matmul for VNNI."""
+    if attrs["transpose_a"] or not attrs["transpose_b"]:
+        return None
+    return vnni_legalize(inputs, arg_types, relay.nn.batch_matmul, attrs, need_expand=True)
diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc
index 29153037b9fa..67b9ffffe21f 100644
--- a/src/arith/canonical_simplify.cc
+++ b/src/arith/canonical_simplify.cc
@@ -982,9 +982,13 @@ SplitExpr CanonicalSimplifier::Impl::SplitModConst(SplitExpr lhs, int64_t cval,
     return lhs;
   }
   if (cval % lhs->scale == 0) {
-    // (x * c1) % (c2 * c1) => (x % c2) * c1
+    // The rationale:
+    //   (index % upper) / lower * scale % cval, given cval = scaled_cval * scale
+    //   by the rule (x * c1) % (c2 * c1) => (x % c2) * c1,
+    // = (index % upper) / lower % scaled_cval * scale
+    //   by the rule (x / c1) % c2  =>  (x % (c1 * c2)) / c1,
+    // = (index % upper) % (new_upper_factor) / lower * scale
     int64_t scaled_cval = cval / lhs->scale;
-    //  (x / c1) % c2  =>  (x % (c1 * c2)) / c2
     int64_t new_upper_factor = lhs->lower_factor * scaled_cval;
     // try to see if we can reduce the existing upper modular.
     if (lhs->upper_factor == SplitExprNode::kPosInf || lhs->upper_factor % new_upper_factor == 0) {
@@ -995,11 +999,13 @@ SplitExpr CanonicalSimplifier::Impl::SplitModConst(SplitExpr lhs, int64_t cval,
       if (new_upper_factor < lhs->upper_factor && lhs->upper_factor != SplitExprNode::kPosInf) {
         auto updated = ToSplitExpr(this->VisitExpr(
             ModImpl(lhs->index, make_const(lhs.dtype(), new_upper_factor), div_mode)));
-        updated.CopyOnWrite()->scale = lhs->scale;
         // re-apply the lower_factor
         if (lhs->lower_factor != 1) {
-          return SplitDivConst(updated, lhs->lower_factor, div_mode);
+          auto ret = SplitDivConst(updated, lhs->lower_factor, div_mode);
+          ret.CopyOnWrite()->MulToSelf(lhs->scale);
+          return ret;
         } else {
+          updated.CopyOnWrite()->MulToSelf(lhs->scale);
           return updated;
         }
       } else {
diff --git a/src/arith/int_set.cc b/src/arith/int_set.cc
index 94316c1e485f..9f5ef644d2bb 100644
--- a/src/arith/int_set.cc
+++ b/src/arith/int_set.cc
@@ -453,6 +453,24 @@ class IntervalSetEvaluator : public ExprFunctor<IntervalSet(const PrimExpr&)> {
     return IntervalSet(min_value, max_value);
   }
 
+  IntervalSet VisitExpr_(const BufferLoadNode* op) final {
+    if (!(op->dtype.is_int() || op->dtype.is_uint())) {
+      DLOG(WARNING) << "cannot evaluate set BufferLoad which loads from a " << op->dtype
+                    << " buffer";
+      return IntervalSet::Everything();
+    }
+    // If the indices do not contain any variables to be relaxed, return the BufferLoad itself.
+    // Otherwise return `IntervalSet::everything()` since we have no knowledge on the buffer data.
+    for (const PrimExpr& index : op->indices) {
+      if (UsesVar(index, [dom_map = &this->dom_map_](const VarNode* var) {
+            return dom_map->find(GetRef<Var>(var)) != dom_map->end();
+          })) {
+        return IntervalSet::Everything();
+      }
+    }
+    return IntervalSet::SinglePoint(GetRef<PrimExpr>(op));
+  }
+
   IntervalSet VisitExprDefault_(const Object* op) final {
     DLOG(WARNING) << "cannot evaluate set type " << op->GetTypeKey();
     return IntervalSet::Everything();
diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc
index 4a99e10211b7..9f4cdc49fc96 100644
--- a/src/arith/rewrite_simplify.cc
+++ b/src/arith/rewrite_simplify.cc
@@ -461,7 +461,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const DivNode* op) {
 
   // x / 2.0 = x * 0.5
   if (const FloatImmNode* ptr = op->b.as<FloatImmNode>()) {
-    ICHECK(op->dtype.is_float() ||
+    ICHECK(op->dtype.is_float() || op->dtype.is_bfloat16() ||
            datatype::Registry::Global()->GetTypeRegistered(op->dtype.code()));
     return op->a * make_const(op->b.dtype(), 1.0 / ptr->value);
   }
@@ -762,6 +762,20 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
         if (c2val % c1val == 0) return floordiv(x, floordiv(c2, c1)).Eval();
       }
     }
+    if (floordiv(x * c1 + c2, c3).Match(ret)) {
+      int64_t c1val = c1.Eval()->value;
+      int64_t c2val = c2.Eval()->value;
+      int64_t c3val = c3.Eval()->value;
+      if (c1val > 0 && c3val > 0 && c3val % c1val == 0 && floormod(c2val, c3val) < c1val) {
+        // assume c3 == a * c1, x == a * y + b, c2 = d * c3 + e then
+        // (x * c1 + c2) // c3
+        // ==> ((a * y + b) * c1 + d * a * c1 + e) // (a * c1)
+        // ==> y + d + (b * c1 + e) // c3
+        // ==> y + d since 0 <= b * c1 <= (a-1) * c1, 0 <= e < c1
+        // ==> x // (c3 // c1) + (c2 // c3)
+        return (floordiv(x, floordiv(c3, c1)) + floordiv(c2, c3)).Eval();
+      }
+    }
 
     TVM_TRY_REWRITE(floordiv(x, x), OneWithTypeLike(x));
     TVM_TRY_REWRITE(floordiv(x * c1, x), c1);
diff --git a/src/auto_scheduler/feature.cc b/src/auto_scheduler/feature.cc
index aaf7d48b10c5..5809888543c6 100755
--- a/src/auto_scheduler/feature.cc
+++ b/src/auto_scheduler/feature.cc
@@ -246,14 +246,14 @@ int64_t GetLoopExtent(const ForNode* node) {
 // Count math ops in an expr
 class MathOpCounter : public StmtExprVisitor {
  public:
-#define VisitBinary(Type, float_ct, int_ct) \
-  void VisitExpr_(const Type* op) final {   \
-    if (op->a.dtype().is_float()) {         \
-      float_ct++;                           \
-    } else {                                \
-      int_ct++;                             \
-    }                                       \
-    StmtExprVisitor::VisitExpr_(op);        \
+#define VisitBinary(Type, float_ct, int_ct)                        \
+  void VisitExpr_(const Type* op) final {                          \
+    if (op->a.dtype().is_float() || op->a.dtype().is_bfloat16()) { \
+      float_ct++;                                                  \
+    } else {                                                       \
+      int_ct++;                                                    \
+    }                                                              \
+    StmtExprVisitor::VisitExpr_(op);                               \
   }
 
   VisitBinary(AddNode, float_addsub, int_addsub);
@@ -299,13 +299,13 @@ class MathOpCounter : public StmtExprVisitor {
         effect_kind == CallEffectKind::kPure || effect_kind == CallEffectKind::kExprAnnotation;
 
     if (is_pure) {
-      if (op->dtype.is_float()) {
+      if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
         float_math_func++;
       } else {
         int_math_func++;
       }
     } else {
-      if (op->dtype.is_float()) {
+      if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
         float_other_func++;
       } else {
         int_other_func++;
diff --git a/src/autotvm/touch_extractor.h b/src/autotvm/touch_extractor.h
index 313e4d78d6e1..83260e1e0633 100644
--- a/src/autotvm/touch_extractor.h
+++ b/src/autotvm/touch_extractor.h
@@ -87,27 +87,37 @@ class TouchExtractor : public FeatureVisitor {
 
   // arithmetic stats
   void VisitExpr_(const AddNode* op) final {
-    if (op->dtype.is_float()) itervar_map[itervar_stack_.back()].add_ct++;
+    if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
+      itervar_map[itervar_stack_.back()].add_ct++;
+    }
     FeatureVisitor::VisitExpr_(op);
   }
 
   void VisitExpr_(const SubNode* op) final {
-    if (op->dtype.is_float()) itervar_map[itervar_stack_.back()].add_ct++;
+    if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
+      itervar_map[itervar_stack_.back()].add_ct++;
+    }
     FeatureVisitor::VisitExpr_(op);
   }
 
   void VisitExpr_(const MulNode* op) final {
-    if (op->dtype.is_float()) itervar_map[itervar_stack_.back()].mul_ct++;
+    if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
+      itervar_map[itervar_stack_.back()].mul_ct++;
+    }
     FeatureVisitor::VisitExpr_(op);
   }
 
   void VisitExpr_(const DivNode* op) final {
-    if (op->dtype.is_float()) itervar_map[itervar_stack_.back()].div_ct++;
+    if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
+      itervar_map[itervar_stack_.back()].div_ct++;
+    }
     FeatureVisitor::VisitExpr_(op);
   }
 
   void VisitExpr_(const ModNode* op) final {
-    if (op->dtype.is_float()) itervar_map[itervar_stack_.back()].div_ct++;
+    if (op->dtype.is_float() || op->dtype.is_bfloat16()) {
+      itervar_map[itervar_stack_.back()].div_ct++;
+    }
     FeatureVisitor::VisitExpr_(op);
   }
 
diff --git a/src/contrib/ethosu/cascader/pareto.cc b/src/contrib/ethosu/cascader/pareto.cc
index 255719088cb0..52ea729bffa2 100644
--- a/src/contrib/ethosu/cascader/pareto.cc
+++ b/src/contrib/ethosu/cascader/pareto.cc
@@ -29,6 +29,8 @@
 
 #include "common.h"
 #include "plan.h"
+#include "proposal.h"
+#include "tensor_config.h"
 
 namespace tvm {
 namespace contrib {
@@ -106,6 +108,31 @@ std::vector<Plan> ParetoCullPlans(std::vector<Plan> plans, size_t max_plans) {
   return ThinVector(optimal_plans, max_plans);
 }
 
+std::vector<Proposal> ParetoCullProposals(std::vector<Proposal> proposals, size_t max_proposals) {
+  std::sort(proposals.begin(), proposals.end(), [](const Proposal& a, const Proposal& b) -> bool {
+    return a->GetMemoryUsage() < b->GetMemoryUsage();
+  });
+  std::vector<std::array<float, 2>> costs;
+  for (const auto& proposal : proposals) {
+    std::array<float, 2> cost = {static_cast<float>(proposal->GetMemoryUsage()),
+                                 static_cast<float>(proposal->GetCycles())};
+    costs.emplace_back(cost);
+  }
+  std::vector<bool> is_optimal = GetParetoFrontier<2>(costs);
+  std::vector<Proposal> optimal_proposals;
+  size_t i = 0;
+  for (bool optimal : is_optimal) {
+    if (optimal) {
+      optimal_proposals.push_back(proposals[i]);
+    }
+    i++;
+  }
+  if (optimal_proposals.size() <= max_proposals) {
+    return optimal_proposals;
+  }
+  return ThinVector(optimal_proposals, max_proposals);
+}
+
 TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GetParetoFrontier")
     .set_body_typed([](Array<Array<FloatImm>> tcosts) {
       std::vector<std::array<float, 2>> costs;
diff --git a/src/contrib/ethosu/cascader/pareto.h b/src/contrib/ethosu/cascader/pareto.h
index 6b0cd4e294ac..511da6c2712f 100644
--- a/src/contrib/ethosu/cascader/pareto.h
+++ b/src/contrib/ethosu/cascader/pareto.h
@@ -37,6 +37,8 @@ namespace ethosu {
 namespace cascader {
 
 class Plan;
+class MemoryRegion;
+class Proposal;
 
 /*!
  * \brief Determine the Pareto optimal points.
@@ -65,6 +67,8 @@ std::vector<T> ThinVector(const std::vector<T>& vec, size_t max_size);
  */
 std::vector<Plan> ParetoCullPlans(std::vector<Plan> plans, size_t max_plans);
 
+std::vector<Proposal> ParetoCullProposals(std::vector<Proposal> proposals, size_t max_proposals);
+
 }  // namespace cascader
 }  // namespace ethosu
 }  // namespace contrib
diff --git a/src/contrib/ethosu/cascader/proposal.cc b/src/contrib/ethosu/cascader/proposal.cc
new file mode 100644
index 000000000000..e96be3466e10
--- /dev/null
+++ b/src/contrib/ethosu/cascader/proposal.cc
@@ -0,0 +1,82 @@
+/*
+ * 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.
+ */
+#include "proposal.h"
+
+#include <tvm/runtime/container/array.h>
+#include <tvm/runtime/container/map.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+
+#include <algorithm>
+#include <utility>
+#include <vector>
+
+#include "plan.h"
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+void ProposalNode::VisitAttrs(AttrVisitor* v) {
+  v->Visit("_graph", &graph_);
+  Array<Part> tmp_parts(part_group_.begin(), part_group_.end());
+  v->Visit("_part_group", &tmp_parts);
+  Array<Plan> tmp_plans(plans_.begin(), plans_.end());
+  v->Visit("_plans", &tmp_plans);
+  Map<Tensor, TensorConfig> tmp_tmap(input_tensor_configs_.begin(), input_tensor_configs_.end());
+  v->Visit("_input_tensor_configs", &tmp_tmap);
+  v->Visit("_cascade_region", &cascade_region_);
+  v->Visit("_memory_usage", &memory_usage_);
+  v->Visit("_cycles", &cycles_);
+}
+
+Proposal::Proposal(const CascaderGraph& graph, const std::vector<Part>& part_group,
+                   const std::vector<Plan>& plans, const TensorConfigMap& input_tensor_configs,
+                   const MemoryRegion& cascade_region, int memory_usage, int cycles) {
+  auto n = make_object<ProposalNode>();
+  n->graph_ = std::move(graph);
+  n->part_group_ = std::move(part_group);
+  std::sort(n->part_group_.begin(), n->part_group_.end());
+  n->plans_ = std::move(plans);
+  n->input_tensor_configs_ = std::move(input_tensor_configs);
+  n->cascade_region_ = std::move(cascade_region);
+  n->memory_usage_ = std::move(memory_usage);
+  n->cycles_ = cycles;
+  data_ = std::move(n);
+}
+
+TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.Proposal")
+    .set_body_typed([](CascaderGraph graph, Array<Part> part_group, Array<Plan> plans,
+                       Map<Tensor, TensorConfig> input_tensor_configs, MemoryRegion cascade_region,
+                       int memory_usage, int cycles) {
+      std::vector<Part> spart_group(part_group.begin(), part_group.end());
+      std::vector<Plan> vplans(plans.begin(), plans.end());
+      TensorConfigMap minput_tensor_configs(input_tensor_configs.begin(),
+                                            input_tensor_configs.end());
+      return Proposal(graph, spart_group, vplans, minput_tensor_configs, cascade_region,
+                      memory_usage, cycles);
+    });
+
+TVM_REGISTER_NODE_TYPE(ProposalNode);
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm
diff --git a/src/contrib/ethosu/cascader/proposal.h b/src/contrib/ethosu/cascader/proposal.h
new file mode 100644
index 000000000000..e5db0328b731
--- /dev/null
+++ b/src/contrib/ethosu/cascader/proposal.h
@@ -0,0 +1,105 @@
+/*
+ * 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 src/contrib/ethosu/cascader/proposal.h
+ * \brief Proposal object for the NPU cascader
+ */
+#ifndef TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_
+#define TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_
+
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/object.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "graph.h"
+#include "plan.h"
+#include "tensor_config.h"
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+using MemoryUsageMap = std::unordered_map<MemoryRegion, int, ObjectPtrHash, ObjectPtrEqual>;
+using TensorConfigMap = std::unordered_map<Tensor, TensorConfig, ObjectPtrHash, ObjectPtrEqual>;
+
+/*! \brief Node to represent a Proposal */
+class ProposalNode : public Object {
+ public:
+  void VisitAttrs(AttrVisitor* v);
+
+  /*! \return The CascaderGraph to which the Proposal applies */
+  const CascaderGraph GetGraph() const { return graph_; }
+  /*! \return The Parts which are covered by the Proposal */
+  const std::vector<Part> GetPartGroup() const { return part_group_; }
+  /*! \return The Plans used in the Proposal */
+  const std::vector<Plan> GetPlans() const { return plans_; }
+  /*! \return The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan */
+  const TensorConfigMap GetInputTensorConfigs() const { return input_tensor_configs_; }
+  /*! \return The MemoryRegion where cascading buffers should be homed */
+  const MemoryRegion GetCascadeRegion() const { return cascade_region_; }
+  /*! \return The memory required to execute the Proposal in the cascading MemoryRegion */
+  const int GetMemoryUsage() const { return memory_usage_; }
+  /*! \return The estimated cycles taken to execute the Proposal */
+  int GetCycles() const { return cycles_; }
+
+  static constexpr const char* _type_key = "contrib.ethosu.cascader.Proposal";
+  TVM_DECLARE_FINAL_OBJECT_INFO(ProposalNode, Object);
+
+ protected:
+  friend class Proposal;
+
+  /*! \brief The CascaderGraph to which the Proposal applies */
+  CascaderGraph graph_;
+  /*! \brief The Parts which are covered by the Proposal */
+  std::vector<Part> part_group_;
+  /*! \brief The Plans used in the Proposal */
+  std::vector<Plan> plans_;
+  /*! \brief The TensorConfigs indexed by Tensor in the Proposal which aren't produced by a Plan */
+  TensorConfigMap input_tensor_configs_;
+  /*! \brief The MemoryRegion where cascading buffers should be homed */
+  MemoryRegion cascade_region_;
+  /*! \brief The memory required to execute the Proposal in the cascading MemoryRegion */
+  int memory_usage_;
+  /*! \brief The estimated cycles taken to execute the Proposal */
+  int cycles_;
+};
+
+/*!
+ * \brief A class which describes how to schedule a CascaderGraph as a series of disjoint Plans.
+ */
+class Proposal : public ObjectRef {
+ public:
+  Proposal(const CascaderGraph& graph, const std::vector<Part>& part_group,
+           const std::vector<Plan>& plans, const TensorConfigMap& input_tensor_configs,
+           const MemoryRegion& cascade_region, int memory_usage, int cycles);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(Proposal, ObjectRef, ProposalNode);
+};
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm
+
+#endif  // TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_H_
diff --git a/src/contrib/ethosu/cascader/proposal_generator.cc b/src/contrib/ethosu/cascader/proposal_generator.cc
new file mode 100644
index 000000000000..ce709cbaa657
--- /dev/null
+++ b/src/contrib/ethosu/cascader/proposal_generator.cc
@@ -0,0 +1,219 @@
+/*
+ * 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.
+ */
+#include <tvm/runtime/container/array.h>
+#include <tvm/runtime/container/map.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/support/parallel_for.h>
+
+#include <algorithm>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "cascader_options.h"
+#include "graph.h"
+#include "pareto.h"
+#include "plan.h"
+#include "plan_generator.h"
+#include "proposal.h"
+#include "stripe_config.h"
+#include "tensor_config.h"
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+std::unordered_set<TensorConfig> GetPlanBoundaryConfigs(const Plan& plan) {
+  std::unordered_set<TensorConfig> boundary_configs;
+  for (const auto& config : plan->GetTensorConfigs()) {
+    if (config->GetState() == TensorConfigState::BOUNDARY) {
+      boundary_configs.insert(config);
+    }
+  }
+  return boundary_configs;
+}
+
+bool IsPlanCompatible(const Proposal& proposal, const std::vector<Part>& plan_part_group,
+                      const std::unordered_set<TensorConfig>& plan_boundary_configs) {
+  // Check the Plan Part group is disjoint with the Proposal Part group
+  for (const auto& plan_part : plan_part_group) {
+    for (const auto& proposal_part : proposal->GetPartGroup()) {
+      if (plan_part == proposal_part) {
+        return false;
+      }
+    }
+  }
+  // If the Plan and Proposal disagree on the memory home of a Tensor, they
+  // are incompatible and can't be used to create a new Proposal
+  auto tensor_configs = proposal->GetInputTensorConfigs();
+  for (const auto& plan_config : plan_boundary_configs) {
+    if (tensor_configs.find(plan_config->GetTensor()) != tensor_configs.end()) {
+      auto proposal_config = tensor_configs.at(plan_config->GetTensor());
+      if (proposal_config->GetHomeRegion() != plan_config->GetHomeRegion()) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual> CreatePlansByPart(
+    const std::unordered_map<std::vector<Part>, std::vector<Plan>>& plans_by_group,
+    const CascaderGraph& graph) {
+  std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual> plans_by_part;
+  for (const auto& it : plans_by_group) {
+    auto part_group = it.first;
+    auto plans = it.second;
+    int highest_index = 0;
+    Part& index_part = part_group.front();
+    // Determine the Part in the Part group with the highest ID - this will be used to index
+    // the Plans
+    for (const auto& part : part_group) {
+      int pid = graph->GetPartID(part);
+      if (pid >= highest_index) {
+        index_part = part;
+        highest_index = pid;
+      }
+    }
+    plans_by_part[index_part].insert(plans_by_part[index_part].begin(), plans.begin(), plans.end());
+  }
+  return plans_by_part;
+}
+
+Proposal AddPlanToProposal(const Proposal& proposal, const Plan& plan,
+                           const std::unordered_set<TensorConfig>& plan_boundary_configs) {
+  std::vector<Plan> new_plans = proposal->GetPlans();
+  new_plans.push_back(plan);
+  TensorConfigMap new_configs = proposal->GetInputTensorConfigs();
+  // Add input configs from the Plan if they're homed in the cascade region
+  for (const auto& config : plan_boundary_configs) {
+    if (config->GetHomeRegion() == proposal->GetCascadeRegion()) {
+      new_configs[config->GetTensor()] = config;
+    }
+  }
+  // Remove the Plan's output config from the new_configs if it's present because
+  // it won't be an input to the Proposal any more
+  if (new_configs.find(plan->GetOutputConfig()->GetTensor()) != new_configs.end()) {
+    new_configs.erase(plan->GetOutputConfig()->GetTensor());
+  }
+  // The updated memory usage is the memory required to run the Plan plus the
+  // non-local memory that's required in the Proposal at that point in time
+  int new_memory_usage = plan->GetMemoryUsage();
+  for (const auto& it : new_configs) {
+    if (plan_boundary_configs.find(it.second) == plan_boundary_configs.end()) {
+      new_memory_usage += it.first->GetSize();
+    }
+  }
+  new_memory_usage = std::max(new_memory_usage, proposal->GetMemoryUsage());
+  int new_cycles = proposal->GetCycles() + plan->GetCycles();
+  std::vector<Part> new_part_group = proposal->GetPartGroup();
+  new_part_group.insert(new_part_group.end(), plan->GetPartGroup().begin(),
+                        plan->GetPartGroup().end());
+  std::sort(new_part_group.begin(), new_part_group.end());
+  return Proposal(proposal->GetGraph(), new_part_group, new_plans, new_configs,
+                  proposal->GetCascadeRegion(), new_memory_usage, new_cycles);
+}
+
+std::vector<Proposal> GeneratePartialProposals(
+    const CascaderGraph& graph, const HomeMap& home_map, const CascaderOptions options,
+    const std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual>& plans_by_part,
+    const std::vector<Part>& partial_proposal_group,
+    std::unordered_map<std::vector<Part>, std::vector<Proposal>>* proposals_by_group) {
+  if (proposals_by_group->find(partial_proposal_group) != proposals_by_group->end()) {
+    return proposals_by_group->at(partial_proposal_group);
+  }
+  if (partial_proposal_group.size() == 0) {
+    (*proposals_by_group)[partial_proposal_group] =
+        std::vector<Proposal>{Proposal(graph, std::vector<Part>(), std::vector<Plan>(),
+                                       TensorConfigMap(), options->cascade_region, 0, 0)};
+  } else {
+    Part part = partial_proposal_group.back();
+    const auto& plans = plans_by_part.at(part);
+    for (const auto& plan : plans) {
+      if (plan->GetInteriorRegion() == options->cascade_region) {
+        // Doing this isn't very efficient, but it improves the performance of the Plan
+        // generator
+        std::unordered_set<TensorConfig> plan_boundary_configs = GetPlanBoundaryConfigs(plan);
+        // The residual_proposal_group is a Part group indicating the Parts which aren't
+        // covered by the current Plan. It's the group for which we must find 'residual
+        // Proposals', meaning Proposals which cover the rest of the CascaderGraph assuming we
+        // pick the current Plan.
+        std::vector<Part> residual_proposal_group;
+        std::copy_if(partial_proposal_group.begin(), partial_proposal_group.end(),
+                     std::back_inserter(residual_proposal_group), [&plan](Part value) {
+                       return std::find(plan->GetPartGroup().begin(), plan->GetPartGroup().end(),
+                                        value) == plan->GetPartGroup().end();
+                     });
+        // std::sort(residual_proposal_group.begin(), residual_proposal_group.end());
+        const auto& residual_proposals = GeneratePartialProposals(
+            graph, home_map, options, plans_by_part, residual_proposal_group, proposals_by_group);
+        auto plan_output_tensor = plan->GetOutputConfig()->GetTensor();
+        ICHECK_LE(plan_output_tensor->GetProducers().size(), 1)
+            << "All tensors must have at most one producer.";
+        for (const auto& residual_proposal : residual_proposals) {
+          if (IsPlanCompatible(residual_proposal, plan->GetPartGroup(), plan_boundary_configs)) {
+            (*proposals_by_group)[partial_proposal_group].push_back(
+                AddPlanToProposal(residual_proposal, plan, plan_boundary_configs));
+          }
+        }
+      }
+    }
+    (*proposals_by_group)[partial_proposal_group] =
+        ParetoCullProposals(proposals_by_group->at(partial_proposal_group), options->max_proposals);
+  }
+  return proposals_by_group->at(partial_proposal_group);
+}
+
+std::vector<Proposal> GenerateProposals(const CascaderGraph& graph, const HomeMap& home_map,
+                                        const CascaderOptions options) {
+  // First generate all the Pareto optimal Plans for the CascaderGraph
+  auto plans_by_group = GenerateGraphPlans(graph, home_map, options);
+  // First create a map between every Part in the CascaderGraph and all the Plans for which that
+  // Part is the lowest ID Part within the Plan's Part group
+  std::unordered_map<Part, std::vector<Plan>, ObjectPtrHash, ObjectPtrEqual> plans_by_part =
+      CreatePlansByPart(plans_by_group, graph);
+  // The Part group that partial Proposals are current being generated for
+  std::vector<Part> partial_proposal_group = graph->GetPartOrder();
+  // A map of Proposals indexed by the Part group they cover
+  std::unordered_map<std::vector<Part>, std::vector<Proposal>> proposals_by_group;
+  return GeneratePartialProposals(graph, home_map, options, plans_by_part, partial_proposal_group,
+                                  &proposals_by_group);
+}
+
+TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.GenerateProposals")
+    .set_body_typed([](CascaderGraph graph, Map<Tensor, Array<MemoryRegion>> home_map,
+                       CascaderOptions options) {
+      std::unordered_map<Tensor, std::vector<MemoryRegion>, ObjectPtrHash, ObjectPtrEqual>
+          mhome_map;
+      for (const auto& it : home_map) {
+        std::vector<MemoryRegion> home_regions;
+        for (const auto& i : it.second) {
+          home_regions.push_back(i);
+        }
+        mhome_map[it.first] = home_regions;
+      }
+      return Array<Proposal>(GenerateProposals(graph, mhome_map, options));
+    });
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm
diff --git a/src/contrib/ethosu/cascader/proposal_generator.h b/src/contrib/ethosu/cascader/proposal_generator.h
new file mode 100644
index 000000000000..624dc5b98a85
--- /dev/null
+++ b/src/contrib/ethosu/cascader/proposal_generator.h
@@ -0,0 +1,86 @@
+/*
+ * 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 src/contrib/ethosu/cascader/proposal_generator.h
+ * \brief Algorithm to generate possible Proposals in the NPU cascader
+ */
+#ifndef TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_GENERATOR_H_
+#define TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_GENERATOR_H_
+
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/object.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+class CascaderGraph;
+class MemoryRegion;
+class Tensor;
+class Proposal;
+class CascaderOptions;
+
+using HomeMap =
+    std::unordered_map<Tensor, std::vector<MemoryRegion>, ObjectPtrHash, ObjectPtrEqual>;
+
+/*!
+ * \brief Generate Pareto optimal Proposals for a CascaderGraph.
+ * \param graph The CascaderGraph to generate Proposals for.
+ * \param home_map The Tensor homing map defining valid memory homes for Tensors.
+ * \param options The configuration options with which to run the generator.
+ * \return A vector of Pareto optimal Proposals.
+ * \note This algorithm takes a top-down dynamic programming approach to determining how
+ * to optimally combine Plans into Proposals. It does the following:
+ *
+ * First, run GenerateGraphPlans to generate the Pareto optimal Plans that cover all the
+ * Part groups in the CascaderGraph.
+ *
+ * Solve the problem recursively, generating optimal Proposals for increasingly small
+ * portions of the overall graph.
+ *
+ * Take the first Part in the graph:
+ *   1. Find all the Plans for which the Part is both in the Plan's Part group and has the
+ *      highest Part ID of any Part in the Part group (i.e. it's the 'first' Part in the
+ *      group).
+ *   For each Plan:
+ *     2. Get the Part group covered by the Plan and subtract it from the 'total Part group'
+ *        covering all the Parts. This forms a 'residual Part group'.
+ *     3. Recursively, determine the optimal Proposals for the 'residual Part group' (the graph
+ *        minus the Parts included in the Plan). Memoize the results.
+ *     For each residual Proposal:
+ *       4. Create a new Proposal by adding the current Plan to the residual Proposal.
+ *   5. Pareto cull all the newly created Proposals (which all share the same Part group).
+ * 6. Return the Proposals which cover all the Parts in the CascaderGraph.
+ *
+ */
+std::vector<Proposal> GenerateProposals(const CascaderGraph& graph, const HomeMap& home_map,
+                                        const CascaderOptions& options);
+
+}  // namespace cascader
+}  // namespace ethosu
+}  // namespace contrib
+}  // namespace tvm
+
+#endif  // TVM_CONTRIB_ETHOSU_CASCADER_PROPOSAL_GENERATOR_H_
diff --git a/src/contrib/hybrid/codegen_hybrid.cc b/src/contrib/hybrid/codegen_hybrid.cc
index 54edbaee35cd..5872a49968cb 100644
--- a/src/contrib/hybrid/codegen_hybrid.cc
+++ b/src/contrib/hybrid/codegen_hybrid.cc
@@ -69,6 +69,9 @@ void CodeGenHybrid::PrintType(DataType t, std::ostream& os) {
   } else if (t.is_int()) {
     os << "int";
     ICHECK(t.bits() == 8 || t.bits() == 16 || t.bits() == 32 || t.bits() == 64);
+  } else if (t.is_bfloat16()) {
+    os << "bfloat";
+    ICHECK(t.bits() == 16);
   } else {
     ICHECK(t.is_uint()) << "Unsupported type " << t;
     os << "uint";
diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc
index 5431499d7c9f..2a0c2f73f2ba 100644
--- a/src/driver/driver_api.cc
+++ b/src/driver/driver_api.cc
@@ -249,6 +249,7 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   pass_list.push_back(tir::transform::UnifyThreadBinding());
   pass_list.push_back(tir::transform::CompactBufferAllocation());
   pass_list.push_back(tir::transform::LowerMatchBuffer());
+  pass_list.push_back(tir::transform::InjectSoftwarePipeline());
   pass_list.push_back(tir::transform::FlattenBuffer());
   pass_list.push_back(tir::transform::BF16Legalize());
   pass_list.push_back(tir::transform::NarrowDataType(32));
diff --git a/src/ir/module.cc b/src/ir/module.cc
index 6f2c9f9fe994..f5ec65e4fbca 100644
--- a/src/ir/module.cc
+++ b/src/ir/module.cc
@@ -444,6 +444,13 @@ IRModule IRModule::FromText(const String& text, const String& source_path) {
   return tvm::parser::ParseModule(source_path, text);
 }
 
+LinkedParam::LinkedParam(int64_t id, tvm::runtime::NDArray param) {
+  auto n = make_object<LinkedParamNode>();
+  n->id = id;
+  n->param = param;
+  data_ = std::move(n);
+}
+
 TVM_REGISTER_NODE_TYPE(IRModuleNode);
 
 TVM_REGISTER_GLOBAL("ir.IRModule")
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index 0e42b35349cd..1a5cd4737fb5 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -318,62 +318,6 @@ class ModulePass : public Pass {
   TVM_DEFINE_OBJECT_REF_METHODS(ModulePass, Pass, ModulePassNode);
 };
 
-/*!
- * \brief The SequentialNode contains a set of passes that transform Relay
- * programs from one AST to another semantically equivalent one.
- *
- * One example of this level of pass is that the pass manager needs to correctly
- * perform a host of optimizations with a given optimization level and disabled
- * passes.
- */
-class SequentialNode : public PassNode {
- public:
-  /* \brief The pass meta data.*/
-  PassInfo pass_info;
-
-  /*! \brief A list of passes that used to compose a sequential pass. */
-  tvm::Array<Pass> passes;
-
-  void VisitAttrs(tvm::AttrVisitor* v) {
-    v->Visit("pass_info", &pass_info);
-    v->Visit("passes", &passes);
-  }
-
-  /*!
-   * \brief Get the pass information/meta data.
-   */
-  PassInfo Info() const override { return pass_info; }
-
-  /*!
-   * \brief Resolve the pass dependency. It globs all required passes by
-   *        a given pass and executes them.
-   *
-   * \param mod The module that an optimization pass runs on.
-   *
-   * \return The updated module after resolving pass dependencies.
-   *
-   * TODO(zhiics) Build a dependency graph among the passes using provided
-   * metadata, i.e. required_passes. Likely, we can have a data structure, i.e.
-   * PassInfo, to store the relevant information including the parent passes.
-   */
-  void ResolveDependency(const IRModule& mod);
-
-  /*!
-   * \brief Perform optimizations on a series of passes. The aforementioned
-   *        typical pass manager jobs could be done by it. This function could
-   *        be overloaded to focus on different metrics, i.e. performance,
-   *        memory footprint, etc.
-   *
-   * \param mod The module that these passes are applied on.
-   * \param pass_ctx The context that these passes execute on.
-   *
-   * \return Return the updated module.
-   */
-  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
-
-  static constexpr const char* _type_key = "transform.Sequential";
-  TVM_DECLARE_FINAL_OBJECT_INFO(SequentialNode, PassNode);
-};
 
 PassInfo::PassInfo(int opt_level, String name, tvm::Array<runtime::String> required) {
   auto pass_info = make_object<PassInfoNode>();
diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc
index f5344ab9126e..05899e4465f9 100644
--- a/src/node/structural_hash.cc
+++ b/src/node/structural_hash.cc
@@ -326,44 +326,42 @@ struct ADTObjTrait {
 
 TVM_REGISTER_REFLECTION_VTABLE(runtime::ADTObj, ADTObjTrait);
 
-struct NDArrayContainerTrait {
-  static constexpr const std::nullptr_t VisitAttrs = nullptr;
-
-  static void SHashReduce(const runtime::NDArray::Container* key, SHashReducer hash_reduce) {
-    ICHECK_EQ(key->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
-    ICHECK(runtime::IsContiguous(key->dl_tensor)) << "Can only hash contiguous tensor";
-    hash_reduce(runtime::DataType(key->dl_tensor.dtype));
-    hash_reduce(key->dl_tensor.ndim);
-    for (int i = 0; i < key->dl_tensor.ndim; ++i) {
-      hash_reduce(key->dl_tensor.shape[i]);
-    }
-    hash_reduce->SHashReduceHashedValue(runtime::String::HashBytes(
-        static_cast<const char*>(key->dl_tensor.data), runtime::GetDataSize(key->dl_tensor)));
+void NDArrayContainerTrait::SHashReduce(const runtime::NDArray::Container* key,
+                                        SHashReducer hash_reduce) {
+  ICHECK_EQ(key->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
+  ICHECK(runtime::IsContiguous(key->dl_tensor)) << "Can only hash contiguous tensor";
+  hash_reduce(runtime::DataType(key->dl_tensor.dtype));
+  hash_reduce(key->dl_tensor.ndim);
+  for (int i = 0; i < key->dl_tensor.ndim; ++i) {
+    hash_reduce(key->dl_tensor.shape[i]);
   }
+  hash_reduce->SHashReduceHashedValue(runtime::String::HashBytes(
+      static_cast<const char*>(key->dl_tensor.data), runtime::GetDataSize(key->dl_tensor)));
+}
 
-  static bool SEqualReduce(const runtime::NDArray::Container* lhs,
-                           const runtime::NDArray::Container* rhs, SEqualReducer equal) {
-    if (lhs == rhs) return true;
-
-    auto ldt = lhs->dl_tensor.dtype;
-    auto rdt = rhs->dl_tensor.dtype;
-    ICHECK_EQ(lhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
-    ICHECK_EQ(rhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
-    ICHECK(runtime::IsContiguous(lhs->dl_tensor)) << "Can only compare contiguous tensor";
-    ICHECK(runtime::IsContiguous(rhs->dl_tensor)) << "Can only compare contiguous tensor";
-
-    if (lhs->dl_tensor.ndim != rhs->dl_tensor.ndim) return false;
-    for (int i = 0; i < lhs->dl_tensor.ndim; ++i) {
-      if (!equal(lhs->dl_tensor.shape[i], rhs->dl_tensor.shape[i])) return false;
-    }
-    if (ldt.code == rdt.code && ldt.lanes == rdt.lanes && ldt.bits == rdt.bits) {
-      size_t data_size = runtime::GetDataSize(lhs->dl_tensor);
-      return std::memcmp(lhs->dl_tensor.data, rhs->dl_tensor.data, data_size) == 0;
-    } else {
-      return false;
-    }
+bool NDArrayContainerTrait::SEqualReduce(const runtime::NDArray::Container* lhs,
+                                         const runtime::NDArray::Container* rhs,
+                                         SEqualReducer equal) {
+  if (lhs == rhs) return true;
+
+  auto ldt = lhs->dl_tensor.dtype;
+  auto rdt = rhs->dl_tensor.dtype;
+  ICHECK_EQ(lhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
+  ICHECK_EQ(rhs->dl_tensor.device.device_type, kDLCPU) << "can only compare CPU tensor";
+  ICHECK(runtime::IsContiguous(lhs->dl_tensor)) << "Can only compare contiguous tensor";
+  ICHECK(runtime::IsContiguous(rhs->dl_tensor)) << "Can only compare contiguous tensor";
+
+  if (lhs->dl_tensor.ndim != rhs->dl_tensor.ndim) return false;
+  for (int i = 0; i < lhs->dl_tensor.ndim; ++i) {
+    if (!equal(lhs->dl_tensor.shape[i], rhs->dl_tensor.shape[i])) return false;
   }
-};
+  if (ldt.code == rdt.code && ldt.lanes == rdt.lanes && ldt.bits == rdt.bits) {
+    size_t data_size = runtime::GetDataSize(lhs->dl_tensor);
+    return std::memcmp(lhs->dl_tensor.data, rhs->dl_tensor.data, data_size) == 0;
+  } else {
+    return false;
+  }
+}
 
 TVM_REGISTER_REFLECTION_VTABLE(runtime::NDArray::Container, NDArrayContainerTrait);
 
diff --git a/src/parser/parser.cc b/src/parser/parser.cc
index 44aeeb3bdee1..9b15893092f7 100644
--- a/src/parser/parser.cc
+++ b/src/parser/parser.cc
@@ -31,6 +31,7 @@
 #include <tvm/runtime/logging.h>
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/target/virtual_device.h>
 
 #include <fstream>
 
@@ -442,16 +443,26 @@ class Parser {
    */
   Expr LookupGraphBinding(const Token& token) {
     auto graph_no = token.ToNumber();
-    return this->graph_ctx.at(graph_no);
+    auto it = this->graph_ctx.find(graph_no);
+    if (it != this->graph_ctx.end()) {
+      return it->second;
+    } else {
+      LOG(FATAL) << "Local variable %" << graph_no << " has not yet been defined";
+      throw;
+    }
   }
 
   /*! \brief Bind a local variable in the expression scope.
    *
    * "x" -> Var("x"), these are needed to map from the raw string names
    * to unique variable nodes.
+   * If a virtual device is specified, sets the virtual device of the variable.
    */
-  Var BindVar(const std::string& name, const relay::Type& type_annotation) {
+  Var BindVar(const std::string& name, const relay::Type& type_annotation,
+              Optional<VirtualDevice> virtual_device = Optional<VirtualDevice>()) {
     auto var = Var(name, type_annotation);
+    var->virtual_device_ = virtual_device.value_or(VirtualDevice::FullyUnconstrained());
+    VLOG(1) << "Binding var named " << name << " to variable node " << PrettyPrint(var);
     this->expr_scopes.Add(name, var);
     return var;
   }
@@ -1106,11 +1117,26 @@ class Parser {
           [&]() {
             auto token = Match(TokenType::kLocal);
             auto string = token.ToString();
+
+            // The fake attributes where the virtual device is specified.
+            VirtualDevice virtual_device;
+            if (WhenMatch(TokenType::kLCurly)) {
+              Map<String, ObjectRef> fake_attrs = ParseAttrs();
+              VLOG(9) << "Fake attributes for function parameter: " << fake_attrs;
+              Match(TokenType::kRCurly);
+              if (fake_attrs.size() == 1 && fake_attrs.count(kVirtualDevice)) {
+                ICHECK(fake_attrs[kVirtualDevice].as<VirtualDeviceNode>())
+                    << "Expected the " << kVirtualDevice
+                    << " to have type VirtualDeviceNode, but got " << virtual_device->GetTypeKey();
+                virtual_device = Downcast<VirtualDevice>(fake_attrs[kVirtualDevice]);
+              }
+            }
+
             Type type;
             if (WhenMatch(TokenType::kColon)) {
               type = ParseType();
             }
-            return BindVar(string, type);
+            return BindVar(string, type, virtual_device);
           },
           [&] {
             auto is_ident = Lookahead(1)->token_type == TokenType::kIdentifier;
@@ -1137,7 +1163,26 @@ class Parser {
 
       // TODO(@jroesch): attributes should never be null, they should always be empty.
       if (raw_attrs.size()) {
-        return relay::Function(params, body, ret_type, generics, DictAttrs(raw_attrs));
+        // Promote kVirtualDevice to first-class
+        if (raw_attrs.count(kVirtualDevice)) {
+          ObjectRef vid = raw_attrs.at(kVirtualDevice);
+          ICHECK(vid.as<VirtualDeviceNode>())
+              << "Expected the " << kVirtualDevice << " to have type VirtualDeviceNode, but got "
+              << vid->GetTypeKey();
+
+          DictAttrs attrs;
+          // Don't fill the raw_attrs in if there's nothing other than kVirtualDevice in the
+          // attributes
+          if (raw_attrs.size() > 1) {
+            raw_attrs.erase(kVirtualDevice);
+            attrs = DictAttrs(raw_attrs);
+          }
+          Function func = relay::Function(params, body, ret_type, generics, attrs);
+          func->virtual_device_ = vid;
+          return func;
+        } else {
+          return relay::Function(params, body, ret_type, generics, DictAttrs(raw_attrs));
+        }
       } else {
         return relay::Function(params, body, ret_type, generics, tvm::DictAttrs());
       }
diff --git a/src/printer/relay_text_printer.cc b/src/printer/relay_text_printer.cc
index fdc6c37e527a..97231931ad88 100644
--- a/src/printer/relay_text_printer.cc
+++ b/src/printer/relay_text_printer.cc
@@ -220,9 +220,13 @@ Doc RelayTextPrinter::AllocVar(const Var& var) {
   }
   Doc val = GetUniqueName("%" + name);
   memo_[var] = val;
+  if (!var->virtual_device()->IsFullyUnconstrained()) {
+    val << " {" << kVirtualDevice << "=" << PrintAttributeValue(var->virtual_device()) << "}";
+  }
   if (var->type_annotation.defined()) {
     val << ": " << Print(var->type_annotation);
   }
+
   val << PrintOptionalInfo(var);
   return val;
 }
@@ -445,6 +449,11 @@ Doc RelayTextPrinter::PrintFunc(const Doc& prefix, const relay::Function& fn) {
   for (const Doc& d : PrintDictAttrs(fn->attrs)) {
     params.push_back(d);
   }
+  if (!fn->virtual_device()->IsFullyUnconstrained()) {
+    Doc vid_doc;
+    vid_doc << kVirtualDevice << "=" << PrintAttributeValue(fn->virtual_device());
+    params.push_back(vid_doc);
+  }
   doc << Doc::Concat(params) << ") ";
   if (fn->ret_type.defined()) {
     doc << "-> " << Print(fn->ret_type) << " ";
@@ -515,11 +524,7 @@ Doc RelayTextPrinter::VisitExpr_(const CallNode* op) {
   for (const Expr& arg : op->args) {
     args.push_back(Print(arg));
   }
-#if TVM_LOG_DEBUG
-  for (const Type& type_arg : op->type_args) {
-    args.push_back(Print(type_arg));
-  }
-#endif
+
   for (const Doc& d : PrintCallAttrs(op->attrs, op->op)) {
     args.push_back(d);
   }
diff --git a/src/printer/text_printer.h b/src/printer/text_printer.h
index a4d0ff30fa62..c34c4a5b6dbe 100644
--- a/src/printer/text_printer.h
+++ b/src/printer/text_printer.h
@@ -359,6 +359,7 @@ class TIRTextPrinter : public StmtFunctor<Doc(const Stmt&)>,
   Doc VisitStmt_(const BufferRealizeNode* op) override;
   Doc VisitStmt_(const ProducerRealizeNode* op) override;
   Doc VisitStmt_(const AllocateNode* op) override;
+  Doc VisitStmt_(const AllocateConstNode* op) override;
   Doc VisitStmt_(const IfThenElseNode* op) override;
   Doc VisitStmt_(const SeqStmtNode* op) override;
   Doc VisitStmt_(const EvaluateNode* op) override;
@@ -398,6 +399,7 @@ class TIRTextPrinter : public StmtFunctor<Doc(const Stmt&)>,
   static Doc PrintConstScalar(DataType dtype, const T& data);
   Doc GetUniqueName(std::string prefix);
   Doc AllocVar(const Var& var);
+  Doc AllocConst(const AllocateConst& var);
   Doc AllocBuf(const Buffer& buffer);
   Doc AllocProducer(const DataProducer& buffer);
   /*!
diff --git a/src/printer/tir_text_printer.cc b/src/printer/tir_text_printer.cc
index a9804229da91..e229da4c26d9 100644
--- a/src/printer/tir_text_printer.cc
+++ b/src/printer/tir_text_printer.cc
@@ -519,6 +519,19 @@ Doc TIRTextPrinter::VisitStmt_(const AllocateNode* op) {
   return doc;
 }
 
+Doc TIRTextPrinter::VisitStmt_(const AllocateConstNode* op) {
+  Doc doc;
+  doc << "constant(" << Print(op->buffer_var) << ", " << PrintDType(op->dtype) << ", "
+      << Print(op->extents) << ")";
+
+  if (op->body->IsInstance<SeqStmtNode>()) {
+    doc << PrintBody(op->body);
+  } else {
+    doc << ";" << Doc::NewLine() << Print(op->body);
+  }
+  return doc;
+}
+
 Doc TIRTextPrinter::VisitStmt_(const IfThenElseNode* op) {
   Doc doc;
   doc << "if " << Print(op->condition) << PrintBody(op->then_case);
diff --git a/src/printer/tvmscript_printer.cc b/src/printer/tvmscript_printer.cc
index 0d6c6e5deeba..fe85eb3cd593 100644
--- a/src/printer/tvmscript_printer.cc
+++ b/src/printer/tvmscript_printer.cc
@@ -241,6 +241,7 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const Stmt&)>,
   Doc VisitStmt_(const BufferStoreNode* op) override;
   Doc VisitStmt_(const BufferRealizeNode* op) override;
   Doc VisitStmt_(const AllocateNode* op) override;
+  Doc VisitStmt_(const AllocateConstNode* op) override;
   Doc VisitStmt_(const IfThenElseNode* op) override;
   Doc VisitStmt_(const SeqStmtNode* op) override;
   Doc VisitStmt_(const ForNode* op) override;
@@ -410,6 +411,26 @@ class TVMScriptPrinter : public StmtFunctor<Doc(const Stmt&)>,
   }
 };
 
+/*!
+ * \brief special method to print NDArray in TIR
+ * \param arr the NDArray to be printed
+ * \param os the output stream where the NDArray will be printed to
+ */
+template <typename T>
+void NDArrayToTIR(::tvm::runtime::NDArray arr, std::ostream& os) {
+  int ndim = arr->ndim;
+  int tot_dim = 1;
+  for (int i = 0; i < ndim; i++) {
+    tot_dim *= arr->shape[i];
+  }
+  T* data_ptr = reinterpret_cast<T*>(arr->data);
+  os << "[";
+  for (int i = 0; i < tot_dim; i++) {
+    os << (i != 0 ? ", " : "") << data_ptr[i];
+  }
+  os << "]";
+}
+
 Doc TVMScriptPrinter::GetUniqueName(std::string prefix) {
   std::replace(prefix.begin(), prefix.end(), '.', '_');
   std::string unique_prefix = prefix;
@@ -1015,6 +1036,50 @@ Doc TVMScriptPrinter::VisitStmt_(const AllocateNode* op) {
   return doc;
 }
 
+Doc TVMScriptPrinter::VisitStmt_(const AllocateConstNode* alloc) {
+  std::stringstream ss;
+  ICHECK(alloc->data) << "Should be presented";
+  const auto& data = alloc->data.value();
+
+  if (alloc->dtype.is_int()) {
+    if (alloc->dtype.bits() == 8) {
+      NDArrayToTIR<int8_t>(data, ss);
+    } else if (alloc->dtype.bits() == 16) {
+      NDArrayToTIR<int16_t>(data, ss);
+    } else if (alloc->dtype.bits() == 32) {
+      NDArrayToTIR<int32_t>(data, ss);
+    } else {
+      LOG(FATAL) << "DataType not supported";
+    }
+  } else if (alloc->dtype.is_float()) {
+    if (alloc->dtype.bits() == 16) {
+      NDArrayToTIR<int16_t>(data, ss);
+    } else if (alloc->dtype.bits() == 32) {
+      NDArrayToTIR<float>(data, ss);
+    } else if (alloc->dtype.bits() == 64) {
+      NDArrayToTIR<double>(data, ss);
+    } else {
+      LOG(FATAL) << "DataType not supported";
+    }
+  } else {
+    LOG(FATAL) << "DataType not supported";
+  }
+  auto ndarray_str = ss.str();
+
+  Doc doc;
+  var_not_in_headers_.insert(alloc->buffer_var.get());
+  if (current_num_ != num_child_ - 1) {
+    doc << "with tir.allocate_const(" << ndarray_str << ", " << PrintDType(alloc->dtype) << ", "
+        << Print(alloc->extents) << ")";
+    doc << Doc::Indent(4, Doc::NewLine() << PrintBody(alloc->body));
+  } else {
+    doc << Print(alloc->buffer_var) << " = tir.allocate_const(" << ndarray_str << ", "
+        << PrintDType(alloc->dtype) << ", " << Print(alloc->extents);
+    doc << ")" << Doc::NewLine() << PrintBody(alloc->body);
+  }
+  return doc;
+}
+
 Doc TVMScriptPrinter::VisitStmt_(const IfThenElseNode* op) {
   Doc doc;
   doc << "if " << Print(op->condition) << ":";
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
index a14cef98669a..e78f4d4aabda 100644
--- a/src/relay/backend/aot_executor_codegen.cc
+++ b/src/relay/backend/aot_executor_codegen.cc
@@ -27,6 +27,7 @@
 #include <tvm/relay/attrs/call.h>
 #include <tvm/relay/executor.h>
 #include <tvm/relay/expr_functor.h>
+#include <tvm/relay/runtime.h>
 #include <tvm/runtime/device_api.h>
 #include <tvm/runtime/object.h>
 #include <tvm/tir/analysis.h>
@@ -70,6 +71,7 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
   void Run(const Function& func) { VisitExpr(func); }
 
   std::vector<int> GetReturnIds() const { return return_ids_; }
+  std::vector<TensorType> GetReturnTtypes() const { return return_ttypes_; }
 
   StorageMap GetStorageMap() const { return storage_device_map_; }
 
@@ -177,6 +179,12 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
       for (auto sid : sinfo->storage_ids) {
         return_ids_.push_back(sid);
       }
+      return_ttypes_.clear();
+      auto ttypes = FlattenTupleType(e->checked_type());
+      return_ttypes_.reserve(ttypes.size());
+      for (auto ttype : ttypes) {
+        return_ttypes_.push_back(ttype);
+      }
     }
   }
   /*!
@@ -252,11 +260,176 @@ class AOTOnDemandAllocator : public transform::DeviceAwareExprVisitor {
   int next_available_sid_{0};
   /*! \brief the set of intermediate tensors that are return variables */
   std::vector<int> return_ids_;
+  /*! \brief the data types of the return values */
+  std::vector<TensorType> return_ttypes_;
 };
 
-/*! \brief Code generator for AOT executor */
-class AOTExecutorCodegen : public MixedModeVisitor {
- protected:
+namespace {
+
+/*!
+ * \brief Utility function to convert a concrete integer to a PrimExpr.
+ * \param num the number to convert
+ * \return PrimExpr representing num
+ */
+inline PrimExpr ConstInt32(size_t num) {
+  ICHECK_LE(num, std::numeric_limits<int>::max());
+  return tir::make_const(DataType::Int(32), static_cast<int>(num));
+}
+
+/*!
+ * \brief Emit a call to the C Device API.
+ * \param device_name Name of the device, used to prefix the function name.
+ * \param hook Name of the Device API function.
+ * \param context void* context arg passed to this API function.
+ */
+tir::Stmt MakeDeviceHookCall(const std::string& device_name, const std::string& hook, PrimExpr context) {
+  Array<String> sections = {"Device", device_name, hook};
+  String device_hook = ToCFunctionStyle(PrefixName(sections));
+
+  return tir::Evaluate(tir::Call(DataType::Int(32), tvm::tir::builtin::call_extern(),
+                                 {tvm::tir::StringImm(device_hook), context}));
+}
+}  // namespace
+
+class AOTCallGenerator {
+ public:
+  AOTCallGenerator(std::string func_name) :
+      func_name_{func_name}, args_{tvm::tir::StringImm(func_name)} {}
+
+  tir::Var PushArg(PrimExpr arg) {
+    if (!arg->IsInstance<tir::VarNode>()) {
+      arg = MakeLetBind(arg);
+    }
+    args_.push_back(arg);
+    return Downcast<tir::Var>(arg);
+  }
+
+  void PushStackDLTensor(const TensorType& ttype, PrimExpr data) {
+    auto dltensor_var = MakeLetBind(StackAlloca("array", 1));
+    auto shape_var = MakeLetBind(StackAlloca("shape", ttype->shape.size()));
+
+    // Populate DLTensor.data
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrData, data})));
+
+    // Populate DLTensor.device
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrDeviceType, kDLCPU})));
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrDeviceId, 0})));
+
+    // Populate DLTensor.ndim
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrNDim, static_cast<int32_t>(ttype->shape.size())})));
+
+    // Populate DLTensor.dtype
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrTypeCode, IntImm(DataType(kDLUInt, 8, 1), ttype->dtype.code())})));
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrTypeBits, IntImm(DataType(kDLUInt, 8, 1), ttype->dtype.bits())})));
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrTypeLanes, IntImm(DataType(kDLUInt, 16, 1), ttype->dtype.lanes())})));
+
+    // Populate DLTensor.shape
+    for (size_t i = 0; i < ttype->shape.size(); ++i) {
+      prep_stmts_.push_back(
+        tvm::tir::Store(shape_var, IntImm(DataType(kDLInt, 64, 1), Downcast<IntImm>(ttype->shape[i])->value), IntImm(DataType(kDLUInt, 64, 1), i), tir::const_true()));
+    }
+
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrShape, shape_var})));
+
+    // Populate DLTensor.strides. DNS -- TODO actually pull correct byte_offset
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrStrides, IntImm(DataType(kDLUInt, 64, 1), 0)})));
+
+    // Populate DLTensor.byte_offset. DNS -- TODO actually pull correct byte_offset
+    prep_stmts_.push_back(
+      tir::Evaluate(
+        tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                       {dltensor_var, 0, tir::builtin::kArrByteOffset, IntImm(DataType(kDLUInt, 64, 1), 0)})));
+
+    args_.push_back(dltensor_var);
+  }
+
+  void PushStackDLTensors(const Expr& expr, std::vector<tir::Var> sids) {
+    const TupleNode* t = expr.as<TupleNode>();
+    if (t != nullptr) {
+      CHECK_EQ(sids.size(), t->fields.size()) << "Relay tuple does not map 1:1 into TIR; AOT can't "
+                                                 "handle this type of Relay Expr in a CallNode.";
+      for (size_t i = 0; i < sids.size(); i++) {
+        PushStackDLTensor(Downcast<TensorType>(t->fields[i]->checked_type()), sids[i]);
+      }
+    } else {
+      PushStackDLTensor(Downcast<TensorType>(expr->checked_type()), sids[0]);
+    }
+  }
+
+  tir::Stmt GenerateUnpacked(std::string device_name, PrimExpr device_context) {
+    auto make_call = [this] {
+      return tir::Evaluate(tvm::tir::Call(DataType::Int(32), tvm::tir::builtin::call_extern(), args_));
+    };
+    if (device_context.defined()) {
+      tir::Var context_var = PushArg(device_context);
+      return Generate(tir::SeqStmt({
+            MakeDeviceHookCall(device_name, "Open", context_var),
+            make_call(),
+            MakeDeviceHookCall(device_name, "Close", context_var),
+          }));
+    } else {
+      return Generate(make_call());
+    }
+  }
+
+  tir::Stmt GeneratePacked() {
+    return Generate(
+      tir::Evaluate(tvm::tir::Call(DataType::Int(32), tir::builtin::tvm_call_packed(), args_)));
+  }
+
+  tir::Stmt GenerateCPacked() {
+    // call_cpacked calling convention does not use a context
+    PushArg(tir::make_zero(DataType::Handle()));
+    return Generate(
+      tir::Evaluate(tvm::tir::Call(DataType::Int(32), tir::builtin::tvm_call_cpacked(), args_)));
+  }
+
+ private:
+  tir::Stmt Generate(tir::Stmt call_stmts) {
+    tir::Stmt body = tir::SeqStmt::Flatten(prep_stmts_, call_stmts);
+
+    for (auto bind : let_binds_) {
+      body = tir::LetStmt(bind.first, bind.second, body);
+    }
+
+    return body;
+  }
+
+  tir::Var MakeLetBind(PrimExpr expr) {
+    std::stringstream ss;
+    ss << func_name_ << "_let" << let_binds_.size();
+    tir::Var v{ss.str(), DataType::Handle()};
+    let_binds_.emplace_back(std::make_pair(v, expr));
+    return v;
+  }
+
   /*!
    * \brief Utility function to allocate a DLTensor or TVMValue
    * \param  type the type of allocation
@@ -268,15 +441,82 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     return tir::Call(DataType::Handle(), tir::builtin::tvm_stack_alloca(), args);
   }
 
-  /*!
-   * \brief Utility function to convert a concrete integer to a PrimExpr.
-   * \param num the number to convert
-   * \return PrimExpr representing num
-   */
-  inline PrimExpr ConstInt32(size_t num) {
-    ICHECK_LE(num, std::numeric_limits<int>::max());
-    return tir::make_const(DataType::Int(32), static_cast<int>(num));
-  }
+  std::string func_name_;
+  tvm::Array<PrimExpr> args_;
+  std::vector<std::pair<tir::Var, PrimExpr>> let_binds_;
+  Array<tir::Stmt> prep_stmts_;
+};
+
+/*! \brief Code generator for AOT executor */
+class AOTExecutorCodegen : public MixedModeVisitor {
+ protected:
+
+  /*! \brief Describes the type of kernel call emitted. */
+  enum CallType {
+    /*!
+     * \brief Emit PackedFunc calls bound just-in-time using TVMBackend* functions.
+     *
+     * When this type is selected, assumes all operators must be called via TVMFuncCall. Given the
+     * implementation of TVMFuncCall in the C++ runtime, this in practice implies that those
+     * functions are of type TVMBackendPackedCFunc.
+     *
+     * The following code is emitted at call sites to call a function named `func`:
+     * void* func_ptr = TVMBackendGetFuncFromEnv("func");
+     * TVMFuncCall(func_ptr, values, tcodes, num_args, ret_values, ret_tcodes)
+     *
+     * The arguments given to the tir::Call node are encoded into `values`, `tcodes`, and `num_args`
+     * by LowerTVMBuiltin TIR transform.
+     *
+     * If `resource_handle` is passed to `func`, it is determined by TVMFuncCall (often,
+     * `resource_handle` is registered with the C++ runtime to provide a `this` equivalent when
+     * `func` is implemented in C).
+     *
+     * Compatible with both C++ and C runtimes, implemented with the C runtime only.
+     */
+    kPacked,  // Emit tir.call_packed and wrap all arguments in DLTensor.
+
+    /*!
+     * \brief Directly call a TVMBackendPackedCFunc named according to the tir::Call.
+     *
+     * When this type is selected, assumes all operators are implemented in functions of type
+     * `TVMBackendPackedCFunc` and should be called directly. That is, presumes at the time of
+     * downstream compilation that there is a symbol named after the 0th arg to tir::Call of
+     * type `TVMBackendPackedCFunc`. This situation should occur when target_host == target.
+     *
+     * The following code is emitted at call sites to call a function named `func`:
+     * func(values, tcodes, num_args, ret_values, ret_tcodes, resource_handle)
+     *
+     * The arguments given to the tir::Call node are encoded into `values`, `tcodes`, and `num_args`
+     * by LowerTVMBuiltin TIR transform.
+     *
+     * `resource_handle` is encoded as the final argument to the tir::Call node. In practice, it is
+     * always the device context parameter when not null. At present, the implementation does not
+     * support forwarding device context parameters to CPacked.
+     *
+     * Compatible with the C runtime and C++ runtime (so long as target_host == target). Implemented
+     * in the same scenarios.
+     */
+    kCPacked,  // Emit tir.call_cpacked and wrap all arguments in DLTensor.
+
+    /*! \brief Directly call a function accepting the `data` arrays as args.
+     *
+     * When this type is selected, assumes all operaotrs are implemented in C functions whose
+     * arguments are 1-to-1 with those in the tir::Call. DLTensor arguments are encoded as just the
+     * `data` parameters (i.e. no DLTensor object is passed along).
+     *
+     * The following code is emitted at call sites to a function named `func`:
+     * func(void* arg0, void* arg1, ..., void* argN) // no resource_handle
+     * -or-
+     * func(void* arg0, void* arg1, ..., void* argN, void* resource_handle) // with resource_handle
+     *
+     * `resource_handle` is encoded as the final argument to the tir::Call node. In practice, it is
+     * always the device context parameter when not null.
+     *
+     * Compatible with the C runtime and C++ runtime (so long as target_host == target). Implemented
+     * with the C runtime only.
+     */
+    kUnpacked,   // Emit tir.call_extern passing only the `data` part of DLTensors.
+  };
 
   /*!
    * \brief Return a vector of variables that represents the sids for the given Relay Expr
@@ -317,6 +557,21 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     }
   }
 
+  /*!
+   * \brief Reverse lookup the device name in devices_ map.
+   * \param device_context Value in devices_ to find.
+   * \return Key matching device_context in devices_.
+   */
+  std::string FindDeviceName(tir::Var device_context) {
+    for (std::pair<String, tir::Var> kv : devices_) {
+      if (kv.second->name_hint == device_context->name_hint) {
+        return kv.first;
+      }
+    }
+    ICHECK(false) << "Did not find a device name associated with " << device_context;
+    return "";
+  }
+
   /*!
    * brief Create a function call
    * \param call_lowered_props The lowered function and the arguments to call it with
@@ -324,74 +579,58 @@ class AOTExecutorCodegen : public MixedModeVisitor {
    * ids to hold the result).
    */
   void CreateFuncCall(CallLoweredProps call_lowered_props, const Expr& result_expr) {
-    std::string func_name = call_lowered_props.lowered_func->name_hint;
-    tvm::Array<PrimExpr> args{tvm::tir::StringImm(func_name)};
-    std::vector<tir::Stmt> create_func_call_stmts;
+    AOTCallGenerator call_gen{call_lowered_props.lowered_func->name_hint};
 
     // Pack the inputs
     for (const Expr& arg : call_lowered_props.arguments) {
       if (params_by_expr_.find(arg) != params_by_expr_.end()) {
         auto param_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::lookup_param(),
                                            {tir::StringImm(params_by_expr_[arg])});
-        args.push_back(tvm::tir::Cast(DataType::Handle(), param_handle));
+        // NOTE: this cast looks like a no-op, but is required for compilation downstream.
+        // Because DataType::Handle has default bits=64, but CodeGenC does not observe this field,
+        // adding this cast forces the codegen to insert the cast. In this case, a cast is required
+        // because param_handle is actually code-generated as `const void*`, and the `const` piece
+        // needs to be removed.
+        call_gen.PushStackDLTensor(Downcast<TensorType>(arg->checked_type()),
+                                   tvm::tir::Cast(DataType::Handle(32, 1), param_handle));
       } else {
-        auto var_arg = FindExpr(arg);
-        for (const auto& var : var_arg) {
-          args.push_back(var);
-        }
+        auto sids = FindExpr(arg);
+        call_gen.PushStackDLTensors(arg, sids);
       }
     }
 
     // Pack the return(s) value. A call node can produce multiple outputs
-    for (const auto& var : PackSid(result_expr)) {
-      args.push_back(var);
-    }
-
-    // Use tvm_call_packed to execute the function unless we're calling directly
-    auto calling_pattern = tvm::tir::builtin::tvm_call_cpacked();
-    if (use_unpacked_api_) {
-      calling_pattern = tvm::tir::builtin::call_extern();
-    }
+    auto result_expr_sid = PackSid(result_expr);
+    call_gen.PushStackDLTensors(result_expr, result_expr_sid);
 
     GlobalVar global_var = call_lowered_props.lowered_func;
     tir::Var empty_var("no_device_context", DataType::Handle());
     bool has_c_device_api_context = device_contexts_.count(global_var) != 0;
-    bool use_cpacked_api = !use_unpacked_api_;
-
-    // The device context is passed to the operator in one of the following calling patterns:
-    //  * Unpacked / direct function call with context:
-    //      operator(arg0, arg1, device_context);
-    //  * Unpacked / direct function call without context:
-    //      operator(arg0, arg1);
-    //  * Type-erased packed function call with context:
-    //      operator(args, type_codes, int num_args, out_ret_value, out_ret_tcode,
-    //      device_context_my_device)
-    //  * Type-erased packed function call without context (we create an empty var for codegen):
-    //      operator(args, type_codes, int num_args, out_ret_value, out_ret_tcode,
-    //      no_device_context)
-    if (has_c_device_api_context) {
-      // call_extern calling convention with context
-      tir::Var context = device_contexts_.Get(global_var).value();
-      args.push_back(context);
-
-      tir::Evaluate func_call(tvm::tir::Call(DataType::Int(32), calling_pattern, args));
-      create_func_call_stmts.push_back(tir::SeqStmt({
-          GenerateDeviceHook(context, "Open"),
-          func_call,
-          GenerateDeviceHook(context, "Close"),
-      }));
-    } else if (use_cpacked_api) {
-      // call_cpacked calling convention needs a blank context
-      args.push_back(tir::make_zero(DataType::Handle()));
-      tir::Evaluate func_call(tvm::tir::Call(DataType::Int(32), calling_pattern, args));
-      create_func_call_stmts.push_back(func_call);
-    } else {
-      // call_extern calling convention without context
-      tir::Evaluate func_call(tvm::tir::Call(DataType::Int(32), calling_pattern, args));
-      create_func_call_stmts.push_back(func_call);
+
+    tir::Stmt body;
+    switch (call_type_) {
+    case CallType::kUnpacked: {
+      ICHECK(has_c_device_api_context) << "Only support unpacked calls with device context";
+      // call_extern calling convention with optional context
+      auto device_context = device_contexts_.Get(global_var).value();
+      std::string device_name = FindDeviceName(device_context);
+      body = call_gen.GenerateUnpacked(device_name, device_context);
+      break;
+    }
+    case CallType::kCPacked:
+      ICHECK(!has_c_device_api_context) << "Only support CPacked calls with no device context";
+      body = call_gen.GenerateCPacked();
+      break;
+    case CallType::kPacked:
+      // call_packed does not accept a device context.
+      CHECK(!has_c_device_api_context) << "CallType::kPacked does not accept a device context";
+      body = call_gen.GeneratePacked();
+      break;
+    default:
+      ICHECK(false) << "Unknown CallType: " << call_type_;
     }
 
-    tir::Stmt body = tir::SeqStmt(create_func_call_stmts);
+    LOG(INFO) << "CreateFuncCall: " << call_lowered_props.lowered_func->name_hint << " -> " << body;
     stmts_.push_back(body);
   }
 
@@ -478,11 +717,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       return it.second->name_hint == context->name_hint;
     });
     const String& device_name = (*it).first;
-    Array<String> sections = {"Device", device_name, hook};
-    String device_hook = ToCFunctionStyle(PrefixName(sections));
-
-    return tir::Evaluate(tir::Call(DataType::Int(32), tvm::tir::builtin::call_extern(),
-                                   {tvm::tir::StringImm(device_hook), context}));
+    return MakeDeviceHookCall(device_name, hook, context);
   }
 
   /*!
@@ -545,22 +780,24 @@ class AOTExecutorCodegen : public MixedModeVisitor {
 
   void VisitExpr_(const ConstantNode* op) override {
     Expr expr = GetRef<Expr>(op);
-    size_t index = params_.size();
-    std::string name = "p" + std::to_string(index);
     StorageInfo& sinfo = storage_device_map_[expr];
-    param_storage_ids_[name] = sinfo->storage_ids[0];
-    params_[name] = op->data;
-    params_by_expr_.Set(expr, name);
+    std::stringstream ss;
+    ss << "constant_" << constant_map_.size();
+
+    tir::Var constant(ss.str(), PointerType(PrimType(DataType(op->data->dtype))));
+    constant_map_[constant] = op;
+    auto sid = sinfo->storage_ids[0];
+    sids_table_[sid] = constant;
 
     // If the Constant node is an output node we need to copy the content of the parameter to the
-    // output A Var node can only produce a single output
-    auto output_iter = std::find(return_sid_.begin(), return_sid_.end(), sinfo->storage_ids[0]);
+    // output. A node can only produce a single output
+    auto output_iter = std::find(return_sid_.begin(), return_sid_.end(), sid);
     if (output_iter != return_sid_.end()) {
       int output_index = std::distance(return_sid_.begin(), output_iter);
       auto param_handle = tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::lookup_param(),
-                                         {tir::StringImm(params_by_expr_[expr])});
-      CopyToOutput(GetBufferVarForIO(input_vars_.size() + output_index), param_handle, false,
-                   sinfo->storage_sizes_in_bytes[0]);
+                                         {tir::StringImm(ss.str())});
+      CopyToOutput(GetBufferVarForIO(input_vars_.size() + output_index), constant,
+                   /* pack_input */ false, sinfo->storage_sizes_in_bytes[0]);
     }
   }
 
@@ -608,7 +845,6 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   // runner function needs to be legalized by the LegalizePackedCalls pass.
   tir::PrimFunc CreateMainFunc(String mod_name, unsigned int relay_params) {
     tir::Stmt body = tir::SeqStmt(stmts_);
-
     // Allocate the sids
     std::unordered_map<int, bool> allocated;
 
@@ -629,6 +865,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
           continue;
         }
 
+        allocated[sid] = constant_map_.count(sids_table_[sid]);
+
         // TODO(giuseros): we should allocate this once outside the PrimFunc
         // so we don't pay the price of allocation for every inference
         if (!allocated[sid]) {
@@ -640,10 +878,23 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       }
     }
 
+    for (auto kv : constant_map_) {
+      auto buffer_var = kv.first;
+      auto dtype = DataType(kv.second->data->dtype);
+
+      int ndim = kv.second->data->ndim;
+      Array<PrimExpr> extents;
+
+      for (int i = 0; i < ndim; i++) {
+        int shape = kv.second->data->shape[i];
+        extents.push_back(tir::make_const(DataType::Int(32), shape));
+      }
+      body = tir::AllocateConst(buffer_var, dtype, extents, kv.second->data, body);
+    }
+
     // Define the PrimFunc attributes
     Map<String, ObjectRef> dict_attrs;
-    String run_func_name =
-        runtime::get_name_mangled(mod_name, runtime::symbol::tvm_run_func_suffix);
+    String run_func_name = runtime::get_name_mangled(mod_name, runtime::symbol::tvm_module_main);
     dict_attrs.Set("global_symbol", run_func_name);
     dict_attrs.Set("runner_function", Bool(true));
     dict_attrs.Set(tvm::attr::kTarget, target_host_);
@@ -683,9 +934,39 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       tir::Buffer buffer = tir::Buffer(buffer_var, elem_type, tensor_type->shape, {}, 0,
                                        name + "_buffer", 16, 1, tir::BufferType::kDefault);
       main_buffer_map_.Set(var, buffer);
+      io_tensor_types_.Set(var, Downcast<TensorType>(expr->checked_type()));
     }
   }
 
+  /*!
+   * brief Calculate workspace sizes for PrimFuncs in the IRModule
+   */
+  Map<String, FunctionInfo> CalculateWorkspaceSizes(
+      const IRModule& lowered_mod, const Map<String, FunctionInfo>& function_metadata) {
+    Executor executor_config = lowered_mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
+    Integer workspace_byte_alignment =
+        executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
+    Map<String, FunctionInfo> updated_function_metadata;
+    for (const auto& kv : lowered_mod->functions) {
+      GlobalVar global_var = kv.first;
+      BaseFunc base_func = kv.second;
+      if (base_func->IsInstance<tir::PrimFuncNode>()) {
+        tir::PrimFunc pfunc = Downcast<tir::PrimFunc>(base_func);
+        Target tgt = pfunc->GetAttr<Target>(tvm::attr::kTarget).value();
+        const auto& ws = CalculateWorkspaceBytes(pfunc, workspace_byte_alignment);
+        if (function_metadata.count(global_var->name_hint)) {
+          updated_function_metadata.Set(global_var->name_hint,
+                                        function_metadata[global_var->name_hint]);
+          updated_function_metadata[global_var->name_hint]->workspace_sizes.Set(tgt, ws);
+        } else {
+          FunctionInfo finfo{{{tgt, ws}}, {}, {}, {{tgt, pfunc}}, {}};
+          updated_function_metadata.Set(global_var->name_hint, finfo);
+        }
+      }
+    }
+    return updated_function_metadata;
+  }
+
   /*!
    * brief Run USMP to plan memory for lowered IRModule
    */
@@ -694,17 +975,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     Integer workspace_byte_alignment =
         executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
     IRModule lowered_mod = mod->ShallowCopy();
+    function_metadata_ = CalculateWorkspaceSizes(lowered_mod, function_metadata_);
     lowered_mod = tir::transform::UnifiedStaticMemoryPlanner()(lowered_mod);
-    // Update workspace size based on the pool allocations.
-    for (const auto& kv : function_metadata_) {
-      if (lowered_mod->ContainGlobalVar(kv.first) &&
-          lowered_mod->Lookup(kv.first)->IsInstance<tir::PrimFuncNode>()) {
-        tir::PrimFunc pfunc = Downcast<tir::PrimFunc>(lowered_mod->Lookup(kv.first));
-        Target tgt = pfunc->GetAttr<Target>(tvm::attr::kTarget).value();
-        const auto& ws = CalculateWorkspaceBytes(pfunc, workspace_byte_alignment);
-        kv.second->workspace_sizes.Set(tgt, ws);
-      }
-    }
     Optional<Array<tir::usmp::AllocatedPoolInfo>> allocated_pool_infos =
         lowered_mod->GetAttr<Array<tir::usmp::AllocatedPoolInfo>>(tvm::attr::kPoolArgs);
     backend::FunctionInfo main_func_info =
@@ -736,17 +1008,18 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     Integer workspace_byte_alignment =
         executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
     IRModule lowered_mod = mod->ShallowCopy();
+    function_metadata_ = CalculateWorkspaceSizes(lowered_mod, function_metadata_);
     // Running StorageRewrite just on the main function
     tir::PrimFunc tir_main_func =
-        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
     IRModule main_func_mod;
-    main_func_mod->Update(lowered_mod->GetGlobalVar(::tvm::runtime::symbol::tvm_run_func_suffix),
+    main_func_mod->Update(lowered_mod->GetGlobalVar(::tvm::runtime::symbol::tvm_module_main),
                           tir_main_func);
     main_func_mod = tir::transform::StorageRewrite()(main_func_mod);
-    lowered_mod->Update(lowered_mod->GetGlobalVar(::tvm::runtime::symbol::tvm_run_func_suffix),
-                        main_func_mod->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+    lowered_mod->Update(lowered_mod->GetGlobalVar(::tvm::runtime::symbol::tvm_module_main),
+                        main_func_mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
     tir_main_func =
-        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
     // Use the PrimFunc to calculate the workspace required to service the allocates
     Integer main_workspace_size_bytes =
         CalculateWorkspaceBytes(tir_main_func, workspace_byte_alignment);
@@ -770,19 +1043,17 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   Array<tir::Var> main_signature_;
   /*! \brief input and output variables belonging to the main function signature */
   Map<tir::Var, tir::Buffer> main_buffer_map_;
+  /*! \brief maps input and output variables to TensorType which describe them */
+  Map<tir::Var, TensorType> io_tensor_types_;
   /*! \brief target device */
   tec::TargetMap targets_;
   /*! \brief target host */
   Target target_host_;
   /*!
-   * \brief unpacked api toggle
-   * When set to true the code generated will use unpacked calls to functions:
-   * func(void* arg0, void* arg1)
-   * Rather than packed calls:
-   * func(void* args)
-   * Defaults to using the packed calling convention
+   * \brief The type of kernel call to be emitted.
+   * See CallType for more documentation.
    */
-  Bool use_unpacked_api_;
+  CallType call_type_;
 
   /*!
    * \brief parameters (i.e. ConstantNodes found in the graph).
@@ -795,11 +1066,13 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   Map<Expr, String> params_by_expr_;
   /*! \brief mapping between parameter names ("p0", "p1", etc..) and storage identifiers*/
   std::unordered_map<std::string, int64_t> param_storage_ids_;
+  std::unordered_map<const tir::Var, const ConstantNode*, ObjectPtrHash, ObjectPtrEqual>
+      constant_map_;
 
   /*! \brief plan memory of device result */
   StorageMap storage_device_map_;
   /*! \brief mapping sid -> tir::Var */
-  std::unordered_map<int, te::Var> sids_table_;
+  std::unordered_map<int, tir::Var> sids_table_;
   /*! \brief lowered funcs */
   Map<String, FunctionInfo> function_metadata_;
   /*! \brief the set of statements that make the program */
@@ -809,7 +1082,9 @@ class AOTExecutorCodegen : public MixedModeVisitor {
 
  public:
   AOTExecutorCodegen(runtime::Module* mod, const tec::TargetMap& targets, Target target_host)
-      : mod_(mod), targets_(targets), target_host_(target_host), use_unpacked_api_(Bool(false)) {}
+      : mod_(mod),
+        targets_(targets),
+        target_host_(target_host) {}
 
   LoweredOutput Codegen(IRModule mod, relay::Function func, String mod_name) {
     VLOG_CONTEXT << "AOT";
@@ -819,11 +1094,44 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     ICHECK(target_host_.defined()) << "require a target_host to be given for AOT codegen";
     VLOG(1) << "target host: " << target_host_->ToDebugString();
 
+    Runtime runtime_config = mod->GetAttr<Runtime>(tvm::attr::kRuntime).value();
     Executor executor_config = mod->GetAttr<Executor>(tvm::attr::kExecutor).value();
-    String interface_api = executor_config->GetAttr<String>("interface-api").value_or("packed");
+    std::string interface_api = executor_config->GetAttr<String>("interface-api").value_or("packed");
     Integer workspace_byte_alignment =
         executor_config->GetAttr<Integer>("workspace-byte-alignment").value_or(16);
-    use_unpacked_api_ = executor_config->GetAttr<Bool>("unpacked-api").value_or(Bool(false));
+    bool unpacked_api = executor_config->GetAttr<Bool>("unpacked-api").value_or(Bool(false));
+
+    // Validate choice of use_unpacked_api_ and use_call_cpacked_
+    if (runtime_config->name == kTvmRuntimeCrt) {
+      if (unpacked_api == true && interface_api == "c") {
+        call_type_ = CallType::kUnpacked;
+      } else if (unpacked_api == false && interface_api == "c") {
+        call_type_ = CallType::kCPacked;
+      } else if (unpacked_api == false && interface_api == "packed") {
+        call_type_ = CallType::kPacked;
+      } else {
+        CHECK(interface_api == "packed" || unpacked_api == true)
+          << "Either need interface-api == \"packed\" (got: " << interface_api
+          << ") or unpacked-api == true (got: " << unpacked_api
+          << ") when targeting c runtime";
+        ICHECK(false) << "Unhandled executor option config: interface-api=" << interface_api
+                      << ", unpacked-api=" << unpacked_api;
+      }
+    } else if (runtime_config->name == kTvmRuntimeCpp) {
+      if (unpacked_api == false && interface_api == "packed") {
+        call_type_ = CallType::kCPacked;
+      } else {
+        CHECK(static_cast<bool>(unpacked_api) == false && interface_api == "packed")
+          << "Need unpacked-api == false (got: " << unpacked_api
+          << ") and interface-api == \"packed\" (got: " << interface_api
+          << ") when targeting c++ runtime";
+        ICHECK(false) << "Unhandled executor option config: interface-api=" << interface_api
+                      << ", unpacked-api=" << unpacked_api;
+      }
+    } else {
+      ICHECK(false) << "runtime_config (" << runtime_config->name
+                    << ") is not one of the expected values";
+    }
 
     // TODO(mbs): Plumb from compiler config
     VirtualDevice host_virtual_device = VirtualDevice::ForTarget(target_host_);
@@ -891,7 +1199,6 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     // because the packed calls arguments are not wrapped in TVMValues. To make this happen we need
     // to run the LegalizePackedCalls pass.
     LoweredOutput ret;
-
     ret.params = std::unordered_map<std::string, std::pair<int, const tvm::runtime::NDArray>>();
     for (auto param : params_) {
       ret.params.emplace(std::make_pair(
@@ -903,7 +1210,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     // function and replacing it with its TIR version. We should try to make this a Pass.
     lowered_mod->Remove(lowered_mod->GetGlobalVar("main"));
     auto prim_func = CreateMainFunc(mod_name, lowered_main_func->params.size());
-    lowered_mod->Update(GlobalVar(::tvm::runtime::symbol::tvm_run_func_suffix), prim_func);
+    lowered_mod->Update(GlobalVar(::tvm::runtime::symbol::tvm_module_main), prim_func);
     // Parallel for loops are not supported in AoT codegen.
     lowered_mod = tir::transform::ConvertForLoopsToSerial()(lowered_mod);
 
@@ -918,10 +1225,10 @@ class AOTExecutorCodegen : public MixedModeVisitor {
 
     // Legalize AOT if needed. This means that all the packed calls
     // need to be wrapped in TVMValues (unless use_unpacked_api is set)
-    if (!use_unpacked_api_) {
-      auto pack_calls = tir::transform::LegalizePackedCalls();
-      lowered_mod = pack_calls(lowered_mod);
-    }
+//    if (!use_unpacked_api_) {
+//      auto pack_calls = tir::transform::LegalizePackedCalls();
+//      lowered_mod = pack_calls(lowered_mod);
+//    }
 
     Optional<Array<tvm::runtime::Module>> external_modules =
         lowered_mod->GetAttr<Array<tvm::runtime::Module>>("external_mods");
@@ -943,13 +1250,14 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     Map<tir::Var, tir::usmp::AllocatedPoolInfo> pool_var_info;
     std::vector<tir::Var> pool_vars;
     tir::PrimFunc tir_main_func =
-        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+        Downcast<tir::PrimFunc>(lowered_mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
     Optional<Array<tir::usmp::AllocatedPoolInfo>> allocated_pool_infos =
         tir_main_func->GetAttr<Array<tir::usmp::AllocatedPoolInfo>>(tvm::attr::kPoolArgs);
     if (allocated_pool_infos) {
       for (const tir::usmp::AllocatedPoolInfo& allocated_pool_info : allocated_pool_infos.value()) {
-        pool_vars.push_back(allocated_pool_info->pool_var.value());
-        pool_var_info.Set(allocated_pool_info->pool_var.value(), allocated_pool_info);
+        int pool_var_index = allocated_pool_info->pool_var_idx.value()->value;
+        pool_vars.push_back(tir_main_func->params[pool_var_index]);
+        pool_var_info.Set(tir_main_func->params[pool_var_index], allocated_pool_info);
       }
     }
     Array<String> devices = ListDevices();
@@ -957,9 +1265,36 @@ class AOTExecutorCodegen : public MixedModeVisitor {
         Array<tir::Var>(tir_main_func->params.begin(),
                         tir_main_func->params.begin() + tir_main_func->params.size() -
                             return_sid_.size() - pool_vars.size() - devices.size());
-    ret.metadata = ExecutorCodegenMetadata(inputs, pool_vars, devices, return_sid_.size(),
-                                           runtime::kTvmExecutorAot, mod_name, interface_api,
-                                           use_unpacked_api_, pool_var_info);
+
+    Array<TensorType> input_tensor_types;
+    for (auto i : inputs) {
+      input_tensor_types.push_back(io_tensor_types_[i]);
+    }
+
+    std::vector<String> output_var_names;
+    if (auto opt = func->GetAttr<Array<String>>("output_tensor_names")) {
+      Array<String> output_tensor_names = opt.value();
+      for (size_t i = 0; i < output_tensor_names.size(); ++i) {
+        output_var_names.push_back(output_tensor_names[i]);
+      }
+    }
+
+    // If output names have not been specified then generate default output names
+    if (output_var_names.size() == 0) {
+      if (return_sid_.size() == 1) {
+        output_var_names.push_back(String("output"));
+      } else {
+        for (size_t i = 0; i < return_sid_.size(); ++i) {
+          output_var_names.push_back(String("output" + std::to_string(i)));
+        }
+      }
+    }
+
+    Array<TensorType> output_tensor_types{final_aot_allocator.GetReturnTtypes()};
+
+    ret.metadata = ExecutorCodegenMetadata(
+        inputs, input_tensor_types, output_var_names, output_tensor_types, pool_vars, devices,
+        runtime::kTvmExecutorAot, mod_name, interface_api, unpacked_api, pool_var_info);
     return ret;
   }
 
diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index aa9c084de4f7..89ee61c83f7c 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -409,7 +409,9 @@ class RelayBuildModule : public runtime::ModuleNode {
    */
   void BuildRelay(IRModule relay_module, const String& mod_name) {
     // Relay IRModule -> IRModule optimizations.
-    relay_module = OptimizeImpl(std::move(relay_module));
+    IRModule module = WithAttrs(
+        relay_module, {{tvm::attr::kExecutor, executor_}, {tvm::attr::kRuntime, runtime_}});
+    relay_module = OptimizeImpl(std::move(module));
 
     // Get the updated function and new IRModule to build.
     // Instead of recreating the IRModule, we should look at the differences between this and the
@@ -437,31 +439,6 @@ class RelayBuildModule : public runtime::ModuleNode {
 
     const Target& host_target = config_->host_virtual_device->target;
     const runtime::PackedFunc* pf = runtime::Registry::Get("codegen.LLVMModuleCreate");
-
-    // Generate a placeholder function that attaches linked params as its arguments.
-    Bool should_link_params = func_module->ShouldLinkParameters();
-    if (should_link_params) {
-      CHECK(pf != nullptr) << "Unable to link-params with no target_host and no llvm codegen.";
-      auto param_ids = executor_codegen_->GetParamIds();
-      auto link_params = Map<String, tir::LinkedParam>();
-      for (auto param : ret_.params) {
-        link_params.Set(param.first, tir::LinkedParam(param_ids[param.first], param.second));
-      }
-
-      Map<String, ObjectRef> dict;
-      dict.Set(tvm::tir::attr::kLinkedParams, link_params);
-      dict.Set(tvm::attr::kGlobalSymbol, String(::tvm::runtime::symbol::tvm_lookup_linked_param));
-      DictAttrs attrs{dict};
-      auto prim = tir::PrimFunc(Array<tir::Var>(), tir::SeqStmt(Array<tir::Stmt>()), VoidType(),
-                                Map<tir::Var, tir::Buffer>(), attrs);
-      if (lowered_funcs.find(host_target) == lowered_funcs.end()) {
-        lowered_funcs.Set(host_target,
-                          IRModule(Map<GlobalVar, BaseFunc>({}), {}, {}, {}, func_module->attrs));
-      }
-      lowered_funcs[host_target]->Add(GlobalVar(::tvm::runtime::symbol::tvm_lookup_linked_param),
-                                      prim);
-    }
-
     // When there is no lowered_funcs due to reasons such as optimization.
     if (lowered_funcs.size() == 0) {
       if (host_target->kind->name == "llvm") {
diff --git a/src/relay/backend/contrib/codegen_c/codegen_c.h b/src/relay/backend/contrib/codegen_c/codegen_c.h
index 3c6f810534ea..49a5bca068d1 100644
--- a/src/relay/backend/contrib/codegen_c/codegen_c.h
+++ b/src/relay/backend/contrib/codegen_c/codegen_c.h
@@ -363,6 +363,8 @@ class CodegenCBase {
       dtype = "float";
     } else if (runtime::TypeMatch(ttype->dtype, kDLFloat, 16)) {
       dtype = "half";
+    } else if (runtime::TypeMatch(ttype->dtype, kDLBfloat, 16)) {
+      dtype = "bfloat";
     } else if (runtime::TypeMatch(ttype->dtype, kDLInt, 32)) {
       dtype = "int";
     } else if (runtime::TypeMatch(ttype->dtype, kDLInt, 64)) {
diff --git a/src/relay/backend/contrib/dnnl/codegen.cc b/src/relay/backend/contrib/dnnl/codegen.cc
index b1b2f580cf94..7971b9cf67d2 100644
--- a/src/relay/backend/contrib/dnnl/codegen.cc
+++ b/src/relay/backend/contrib/dnnl/codegen.cc
@@ -31,6 +31,7 @@
 
 #include <fstream>
 #include <numeric>
+#include <regex>
 #include <sstream>
 
 #include "../../utils.h"
@@ -439,6 +440,23 @@ class DNNLJSONSerializer : public backend::contrib::JSONSerializer {
   using JSONGraphNode = tvm::runtime::json::JSONGraphNode;
   using JSONGraphNodeEntry = tvm::runtime::json::JSONGraphNodeEntry;
 
+  std::map<std::string, std::string> op_map{
+      {"bias", "add"},
+      {"relu", "nn.relu"},
+      {"tanh", "tanh"},
+      {"sigmoid", "sigmoid"},
+  };
+
+  std::vector<std::string> ParsingOpList(std::string op, std::string pattern_name) {
+    std::vector<std::string> op_list = {"nn." + op};
+    for (auto& t : op_map) {
+      if (pattern_name.find(t.first) != std::string::npos) {
+        op_list.push_back(t.second);
+      }
+    }
+    return op_list;
+  }
+
  public:
   DNNLJSONSerializer(const std::string& symbol, const Expr& expr) : JSONSerializer(symbol, expr) {}
 
@@ -453,28 +471,29 @@ class DNNLJSONSerializer : public backend::contrib::JSONSerializer {
       ICHECK(comp.defined()) << "DNNL JSON runtime only supports composite functions.";
       name = comp.value();
 
-      if (name == "dnnl.conv2d_bias_relu") {
-        call = GetRootCall(fn->body.as<CallNode>(), 2, {"nn.conv2d", "add", "nn.relu"});
-      } else if (name == "dnnl.conv2d_bias_tanh") {
-        call = GetRootCall(fn->body.as<CallNode>(), 2, {"nn.conv2d", "add", "tanh"});
-        ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.conv2d_bias_sigmoid") {
-        call = GetRootCall(fn->body.as<CallNode>(), 2, {"nn.conv2d", "add", "sigmoid"});
+      if (name.find("dnnl.conv2d_transpose") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("conv2d_transpose", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.conv2d_bias") {
-        call = GetRootCall(fn->body.as<CallNode>(), 1, {"nn.conv2d", "add"});
+      } else if (name.find("dnnl.conv3d_transpose") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("conv3d_transpose", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.conv2d_relu") {
-        call = GetRootCall(fn->body.as<CallNode>(), 1, {"nn.conv2d", "nn.relu"});
+      } else if (name.find("dnnl.conv1d") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("conv1d", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.conv2d_tanh") {
-        call = GetRootCall(fn->body.as<CallNode>(), 1, {"nn.conv2d", "tanh"});
+      } else if (name.find("dnnl.conv2d") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("conv2d", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.conv2d_sigmoid") {
-        call = GetRootCall(fn->body.as<CallNode>(), 1, {"nn.conv2d", "sigmoid"});
+      } else if (name.find("dnnl.conv3d") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("conv3d", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
-      } else if (name == "dnnl.dense_bias") {
-        call = GetRootCall(fn->body.as<CallNode>(), 1, {"nn.dense", "add"});
+      } else if (name.find("dnnl.dense") != std::string::npos) {
+        std::vector<std::string> op_list = ParsingOpList("dense", name);
+        call = GetRootCall(fn->body.as<CallNode>(), op_list.size() - 1, op_list);
         ICHECK(call->op.as<OpNode>()) << "Not op node";
       } else {
         LOG(FATAL) << "Unrecognized DNNL pattern: " << name;
diff --git a/src/relay/backend/contrib/ethosn/codegen.cc b/src/relay/backend/contrib/ethosn/codegen.cc
index 2de36b543414..a9fe0aab02ec 100644
--- a/src/relay/backend/contrib/ethosn/codegen.cc
+++ b/src/relay/backend/contrib/ethosn/codegen.cc
@@ -718,6 +718,7 @@ TVM_REGISTER_GLOBAL("relay.ethos-n.support.reshape")
     .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
       Call call = args[0];
       ReshapeParams params;
+      EthosnAPI::DefaultInputTensor(call);
       auto err = EthosnAPI::Reshape(call, &params);
       err += EthosnCompiler::SupportedSetup();
       char reason[kReasonMaxLength];
@@ -784,6 +785,7 @@ TVM_REGISTER_GLOBAL("relay.ethos-n.support.split")
     .set_body([](tvm::TVMArgs args, tvm::TVMRetValue* rv) {
       Call call = args[0];
       SplitParams params;
+      EthosnAPI::DefaultInputTensor(call);
       auto err = EthosnAPI::Split(call, &params);
       err += EthosnCompiler::SupportedSetup();
       char reason[kReasonMaxLength];
diff --git a/src/relay/backend/contrib/ethosn/ethosn_api.cc b/src/relay/backend/contrib/ethosn/ethosn_api.cc
index 0c716d0a86db..88cdc790a60c 100644
--- a/src/relay/backend/contrib/ethosn/ethosn_api.cc
+++ b/src/relay/backend/contrib/ethosn/ethosn_api.cc
@@ -45,6 +45,14 @@ namespace relay {
 namespace contrib {
 namespace ethosn {
 
+sl::TensorInfo EthosnAPI::DefaultInputTensor(const Expr& expr) {
+  Call call = Downcast<Call>(expr);
+  const auto* dtype = call->args[0]->checked_type().as<TensorTypeNode>();
+  sl::DataType data_type;
+  Tvm2Npu(dtype->dtype, &data_type);
+  return sl::TensorInfo({}, data_type, sl::DataFormat::NHWC, {});
+}
+
 EthosnError EthosnAPI::QnnConv2d(const Expr& expr, ConvolutionParams* params) {
   Call requantize = Downcast<Call>(expr);
   Call bias_add = Downcast<Call>(requantize->args[0]);
@@ -108,7 +116,7 @@ EthosnError EthosnAPI::QnnConv2d(const Expr& expr, ConvolutionParams* params) {
   sl::Stride stride;
   err += Tvm2Npu(conv_attr->strides, &stride);
   // Dilation is not supported
-  std::array<uint32_t, 4> dilation = {1, 1, 1, 1};
+  std::array<uint32_t, 2> dilation = {1, 1};
   AsArray(conv_attr->dilation, &dilation);
   if (conv_attr->dilation.size() != 2 || dilation[0] != 1 || dilation[1] != 1) {
     err +=
@@ -485,9 +493,6 @@ EthosnError EthosnAPI::DepthToSpace(const Expr& expr, DepthToSpaceParams* params
   EthosnError err = Tvm2Npu(input_dtype->shape, &input_tensor_shape);
   err += Tvm2Npu(input_dtype->dtype, &input_data_type);
   err += Tvm2Npu(attrs->layout, &input_data_format);
-  if (input_data_format != sl::DataFormat::NHWC) {
-    err += EthosnError(ErrStrm() << "layout=" << attrs->layout << ", layout must = NHWC");
-  }
   params->input_info = sl::TensorInfo(input_tensor_shape, input_data_type, input_data_format,
                                       params->input_info.m_QuantizationInfo);
   return err;
@@ -517,11 +522,11 @@ EthosnError EthosnAPI::Tvm2Npu(const Array<IndexExpr>& padding, sl::Padding* npu
   }
   switch (padding.size()) {
     case 1:
-      *npu_padding = sl::Padding(dim[0], dim[0], dim[0], dim[0]);
+      *npu_padding = sl::Padding(dim[3], dim[3], dim[3], dim[3]);
       break;
     case 2:
       // Height, width -> top, bottom, left, right
-      *npu_padding = sl::Padding(dim[0], dim[0], dim[1], dim[1]);
+      *npu_padding = sl::Padding(dim[3], dim[3], dim[2], dim[2]);
       break;
     case 4:
       // Top, left, bottom, right -> top, bottom, left, right
@@ -538,7 +543,7 @@ EthosnError EthosnAPI::Tvm2Npu(const Array<IndexExpr>& strides, sl::Stride* npu_
   if (strides.size() != 2) {
     return EthosnError(ErrStrm() << "stride size=" << strides.size() << ", stride size must = 2");
   }
-  std::array<uint32_t, 4> dim;
+  std::array<uint32_t, 2> dim;
   if (EthosnError err = AsArray<IndexExpr, uint32_t>(strides, &dim)) {
     return err;
   }
@@ -550,7 +555,7 @@ EthosnError EthosnAPI::Tvm2Npu(const Array<IndexExpr>& size, uint32_t* x, uint32
   if (size.size() != 2) {
     return EthosnError(ErrStrm() << "dimensions=" << size.size() << ", dimensions must = 2");
   }
-  std::array<uint32_t, 4> dim;
+  std::array<uint32_t, 2> dim;
   if (EthosnError err = AsArray<IndexExpr, uint32_t>(size, &dim)) {
     return err;
   }
@@ -647,11 +652,12 @@ EthosnError EthosnAPI::Tvm2Npu(const Array<Array<Integer>>& padding, sl::Padding
 // Convert an array of IntImmNodes into ValueT
 // IndexT type of Array indexing variable
 // ValueT type of resulting value
-template <typename IndexT, typename ValueT>
-EthosnError EthosnAPI::AsArray(const Array<IndexT>& arr, std::array<ValueT, 4>* v) {
-  if (arr.size() > 4)
-    return EthosnError(ErrStrm() << "dimensions=" << arr.size() << ", dimensions must be <= 4");
-  for (size_t i = 0; i < std::min(arr.size(), 4ul); i++) {
+// N The size of the output array
+template <typename IndexT, typename ValueT, size_t N>
+EthosnError EthosnAPI::AsArray(const Array<IndexT>& arr, std::array<ValueT, N>* v) {
+  if (arr.size() > N)
+    return EthosnError(ErrStrm() << "dimensions=" << arr.size() << ", dimensions must be <= " << N);
+  for (size_t i = 0; i < arr.size(); i++) {
     const PrimExpr& a = arr[i];
     const auto* intImm = a.as<IntImmNode>();
     if (intImm->value > std::numeric_limits<ValueT>::max()) {
diff --git a/src/relay/backend/contrib/ethosn/ethosn_api.h b/src/relay/backend/contrib/ethosn/ethosn_api.h
index 8f0508626d4b..b790f2daf00e 100644
--- a/src/relay/backend/contrib/ethosn/ethosn_api.h
+++ b/src/relay/backend/contrib/ethosn/ethosn_api.h
@@ -179,6 +179,9 @@ class EthosnError {
  */
 class EthosnAPI {
  public:
+  /*! \brief Create a default input tensor */
+  static sl::TensorInfo DefaultInputTensor(const Expr& expr);
+
   /*! \brief Extract the Support Library convolution params from an ethos-n.qnn_conv2d func */
   static EthosnError QnnConv2d(const Expr& expr, ConvolutionParams* params);
   /*! \brief Extract the Support Library dense params from an ethos-n.qnn_fc func */
@@ -235,8 +238,8 @@ class EthosnAPI {
   // Convert an array of IntImmNodes into ValueT
   // IndexT type of Array indexing variable
   // ValueT type of resulting value
-  template <typename IndexT, typename ValueT>
-  static EthosnError AsArray(const Array<IndexT>& arr, std::array<ValueT, 4>* v);
+  template <typename IndexT, typename ValueT, size_t N>
+  static EthosnError AsArray(const Array<IndexT>& arr, std::array<ValueT, N>* v);
 
   // Get a T from a constant represented by a NDArray.
   template <typename T>
diff --git a/src/relay/backend/executor.cc b/src/relay/backend/executor.cc
index 7c0c690c07aa..581fbdf2cdf1 100644
--- a/src/relay/backend/executor.cc
+++ b/src/relay/backend/executor.cc
@@ -88,6 +88,7 @@ ExecutorRegEntry& ExecutorRegEntry::RegisterOrGet(const String& name) {
 /**********  Register Executors and options  **********/
 
 TVM_REGISTER_EXECUTOR("aot")
+    .add_attr_option<Bool>("link-params", Bool(true))
     .add_attr_option<Bool>("unpacked-api")
     .add_attr_option<String>("interface-api")
     .add_attr_option<Integer>("workspace-byte-alignment");
diff --git a/src/relay/backend/graph_executor_codegen.cc b/src/relay/backend/graph_executor_codegen.cc
index 227f7bbfdf31..30bc0beeebce 100644
--- a/src/relay/backend/graph_executor_codegen.cc
+++ b/src/relay/backend/graph_executor_codegen.cc
@@ -290,25 +290,15 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     // This is the point where we separate the functions in the module by target
     ret.lowered_funcs = tec::GetPerTargetModules(lowered_mod);
     ret.external_mods = external_modules.value();
+    ret.metadata =
+        ExecutorCodegenMetadata({} /* inputs */, {} /* input_tensor_types */, {} /* outputs */,
+                                {} /* output_tensor_types */, {} /* pools */, {} /* devices */,
+                                runtime::kTvmExecutorGraph /* executor */, mod_name_ /* mod_name */,
+                                "packed" /* interface_api */, Bool(false) /* unpacked_api */);
     return ret;
   }
 
  protected:
-  /*!
-   * \brief Extract shape from expr to vector<int64_t>
-   *
-   * \param shape
-   * \return std::vector<int64_t>
-   */
-  std::vector<int64_t> _ShapeToJSON(tvm::Array<IndexExpr> shape) {
-    std::vector<int64_t> ret;
-    for (IndexExpr dim : shape) {
-      const int64_t* pval = tir::as_const_int(dim);
-      ret.push_back(*pval);
-    }
-    return ret;
-  }
-
   /*!
    * \brief Add node to graph
    *
@@ -352,7 +342,7 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
       for (size_t i = 0; i < tuple_type->fields.size(); ++i) {
         if (const auto* typ = tuple_type->fields[i].as<TensorTypeNode>()) {
           ret.push_back(GraphNodeRef(node_id, i));
-          shape.emplace_back(_ShapeToJSON(typ->shape));
+          shape.emplace_back(ShapeToJSON(typ->shape));
           dtype.emplace_back(DType2String(typ->dtype));
         } else {
           LOG(FATAL) << "type " << checked_type->GetTypeKey() << " not supported";
@@ -369,7 +359,7 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator<std::vector<
     if (const auto* tensor_type = checked_type.as<TensorTypeNode>()) {
       ShapeVector shape;
       std::vector<std::string> dtype;
-      shape.emplace_back(_ShapeToJSON(tensor_type->shape));
+      shape.emplace_back(ShapeToJSON(tensor_type->shape));
       dtype.emplace_back(DType2String(tensor_type->dtype));
       node->attrs_["shape"] = shape;
       node->attrs_["dtype"] = dtype;
diff --git a/src/relay/backend/graph_plan_memory.cc b/src/relay/backend/graph_plan_memory.cc
index 2ad27a0d20b0..0019b22f1a8f 100644
--- a/src/relay/backend/graph_plan_memory.cc
+++ b/src/relay/backend/graph_plan_memory.cc
@@ -279,6 +279,8 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
       smap.Set(GetRef<Expr>(kv.first), storage_info);
     }
     // Either all or none of the nodes should be annotated.
+    VLOG(1) << "num annotated nodes / num_nodes: " << num_annotated_nodes << " / " << num_nodes
+            << std::endl;
     if (num_annotated_nodes != 0 && num_annotated_nodes != num_nodes) {
       LOG(FATAL) << num_annotated_nodes << " out of " << num_nodes
                  << "expressions are assigned with virtual device types. Either all "
diff --git a/src/relay/backend/interpreter.cc b/src/relay/backend/interpreter.cc
index 2bea8101d645..c4b1673e0731 100644
--- a/src/relay/backend/interpreter.cc
+++ b/src/relay/backend/interpreter.cc
@@ -1108,7 +1108,8 @@ TypedPackedFunc<ObjectRef(Array<Expr>)> EvalFunction(IRModule mod, Expr expr, De
 }
 
 ObjectRef Eval(Expr expr, Map<GlobalTypeVar, TypeData> type_definitions,
-               std::unordered_set<String> import_set, Device device, Target target) {
+               std::unordered_set<String> import_set, Device device, Target target,
+               Map<String, ObjectRef> attrs) {
   ICHECK_EQ(device.device_type, target->kind->device_type);
   TargetMap targets;
   targets.Set(device.device_type, target);
@@ -1118,7 +1119,7 @@ ObjectRef Eval(Expr expr, Map<GlobalTypeVar, TypeData> type_definitions,
   std::pair<IRModule, GlobalVar> mod_and_global =
       IRModule::FromExprInContext(expr, /*global_funcs=*/{}, type_definitions, import_set);
 
-  IRModule mod = Prepare(mod_and_global.first, config);
+  IRModule mod = Prepare(WithAttrs(mod_and_global.first, {attrs}), config);
 
   Interpreter intrp(mod, config, device);
   Expr expr_to_eval = mod->GetGlobalVar(mod_and_global.second->name_hint);
diff --git a/src/relay/backend/runtime.cc b/src/relay/backend/runtime.cc
index 786d6f937f14..923c9b2d5f65 100644
--- a/src/relay/backend/runtime.cc
+++ b/src/relay/backend/runtime.cc
@@ -88,9 +88,9 @@ RuntimeRegEntry& RuntimeRegEntry::RegisterOrGet(const String& name) {
 
 /**********  Register Runtimes and options  **********/
 
-TVM_REGISTER_RUNTIME("crt").add_attr_option<Bool>("system-lib");
+TVM_REGISTER_RUNTIME(kTvmRuntimeCrt).add_attr_option<Bool>("system-lib");
 
-TVM_REGISTER_RUNTIME("cpp").add_attr_option<Bool>("system-lib");
+TVM_REGISTER_RUNTIME(kTvmRuntimeCpp).add_attr_option<Bool>("system-lib");
 
 /**********  Registry  **********/
 
diff --git a/src/relay/backend/te_compiler.cc b/src/relay/backend/te_compiler.cc
index 46051b06c0d7..cfc0ad0087fc 100644
--- a/src/relay/backend/te_compiler.cc
+++ b/src/relay/backend/te_compiler.cc
@@ -33,6 +33,7 @@
 #include <tvm/runtime/registry.h>
 #include <tvm/te/schedule.h>
 #include <tvm/te/schedule_pass.h>
+#include <tvm/tir/transform.h>
 #include <tvm/topi/tags.h>
 
 #include <functional>
@@ -330,12 +331,18 @@ class TECompilerImpl : public TECompilerNode {
       for (te::Tensor arg : value->cached_func->outputs) {
         all_args.push_back(arg);
       }
+      Array<runtime::NDArray> all_consts;
+      for (auto kv : value->cached_func->constant_tensors) {
+        all_args.push_back(kv.second);
+        all_consts.push_back(kv.first->data);
+      }
       // lower the function
       std::unordered_map<te::Tensor, tir::Buffer> binds;
       auto func_name = value->cached_func->prim_fn_var->name_hint;
       VLOG(1) << "scheduling";
       IRModule scheduled_module =
           tvm::LowerSchedule(value->cached_func->schedule, all_args, func_name, binds);
+      scheduled_module->Update(tir::transform::BindParams(all_consts)(scheduled_module));
       // Unfortunately the above machinery creates its own GlobalVars instead of using *the*
       // GlobalVar we established above. Fix this before the confusion spreads any further.
       // TODO(mbs): LowerSchedule should be given prim_fn_gvar instead of func_name.
@@ -1179,7 +1186,8 @@ Pass LowerTEPass(const String& module_name, ProcessFn process_fn,
 
   return tvm::transform::Sequential(
       {tvm::relay::transform::RelayToTIRTargetHook(),
-       tvm::transform::CreateModulePass(pass_func, 0, "LowerTE", {"InferType"}), InferType()});
+       tvm::transform::CreateModulePass(pass_func, 0, "LowerTE", {"InferType"}), InferType(),
+       tvm::tir::transform::ExtractPrimFuncConstants()});
 }
 }  // namespace tec
 }  // namespace relay
diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc
index 15ef22942869..34ce3f24da27 100644
--- a/src/relay/backend/te_compiler_cache.cc
+++ b/src/relay/backend/te_compiler_cache.cc
@@ -71,15 +71,18 @@ CCacheKey::CCacheKey(Function source_func, Target target) {
 CachedFunc::CachedFunc(tvm::Target target, GlobalVar prim_fn_var, tvm::Array<te::Tensor> inputs,
                        tvm::Array<te::Tensor> outputs, te::Schedule schedule,
                        tir::PrimFunc prim_func, tvm::Array<Integer> shape_func_param_states,
-                       IRModule funcs) {
+                       IRModule funcs,
+                       std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors) {
   auto n = make_object<CachedFuncNode>();
   n->target = target;
   n->prim_fn_var = prim_fn_var;
   n->inputs = inputs;
   n->outputs = outputs;
   n->schedule = schedule;
+  n->prim_func = prim_func;
   n->shape_func_param_states = shape_func_param_states;
   n->funcs = funcs;
+  n->constant_tensors = constant_tensors;
   data_ = std::move(n);
 }
 
@@ -206,7 +209,8 @@ class ScheduleBuilder : public backend::MemoizedExprTranslator<Array<te::Tensor>
       }
     }
 
-    return CachedFunc(target_, prim_fn_var, fn_inputs, outputs, schedule, prim_func, {});
+    return CachedFunc(target_, prim_fn_var, fn_inputs, outputs, schedule, prim_func, {},
+                      IRModule(Map<GlobalVar, BaseFunc>({})), constant_tensors_);
   }
 
   Array<te::Tensor> VisitExpr_(const VarNode* op) final {
@@ -216,30 +220,40 @@ class ScheduleBuilder : public backend::MemoizedExprTranslator<Array<te::Tensor>
 
   Array<te::Tensor> VisitExpr_(const ConstantNode* op) final {
     using tir::make_const;
-    ICHECK(op->is_scalar());
     void* data = op->data->data;
     DataType dtype = DataType(op->data->dtype);
-    auto value = te::compute(
-        {},
-        [&](const Array<tvm::tir::Var>&) {
-          if (dtype == DataType::Int(32)) {
-            return make_const(dtype, static_cast<const int32_t*>(data)[0]);
-          } else if (dtype == DataType::Int(64)) {
-            return make_const(dtype, static_cast<const int64_t*>(data)[0]);
-          } else if (dtype == DataType::Float(32)) {
-            return make_const(dtype, static_cast<const float*>(data)[0]);
-          } else if (dtype == DataType::Float(64)) {
-            return make_const(dtype, static_cast<const double*>(data)[0]);
-          } else if (dtype == DataType::Bool()) {
-            return make_const(dtype, static_cast<const uint8_t*>(data)[0]);
-          } else {
-            LOG(FATAL) << "not handled";
-            return tvm::PrimExpr();
-          }
-        },
-        "compile_engine_const", topi::kBroadcast);
-    scalars_.push_back(value->op);
-    return {value};
+    if (op->is_scalar()) {
+      auto value = te::compute(
+          {},
+          [&](const Array<tvm::tir::Var>&) {
+            if (dtype == DataType::Int(16)) {
+              return make_const(dtype, static_cast<const int16_t*>(data)[0]);
+            } else if (dtype == DataType::Int(32)) {
+              return make_const(dtype, static_cast<const int32_t*>(data)[0]);
+            } else if (dtype == DataType::Int(64)) {
+              return make_const(dtype, static_cast<const int64_t*>(data)[0]);
+            } else if (dtype == DataType::Float(32)) {
+              return make_const(dtype, static_cast<const float*>(data)[0]);
+            } else if (dtype == DataType::Float(64)) {
+              return make_const(dtype, static_cast<const double*>(data)[0]);
+            } else if (dtype == DataType::Bool()) {
+              return make_const(dtype, static_cast<const uint8_t*>(data)[0]);
+            } else {
+              LOG(FATAL) << dtype << " not handled";
+              return tvm::PrimExpr();
+            }
+          },
+          "compile_engine_const", topi::kBroadcast);
+      scalars_.push_back(value->op);
+      return {value};
+    } else {
+      const auto* ttype = op->checked_type().as<TensorTypeNode>();
+      std::stringstream ss;
+      ss << "constant_" << const_index++;
+      tvm::te::Tensor tensor = tvm::te::placeholder(GetShape(ttype->shape), ttype->dtype, ss.str());
+      constant_tensors_[op] = tensor;
+      return {tensor};
+    }
   }
 
   Array<te::Tensor> VisitExpr_(const CallNode* call_node) final {
@@ -344,14 +358,19 @@ class ScheduleBuilder : public backend::MemoizedExprTranslator<Array<te::Tensor>
   OpImplementation anchor_implementation_;
   std::ostringstream readable_name_stream_;
   Array<te::Operation> scalars_;
+  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors_;
   bool use_auto_scheduler_;
   bool use_meta_schedule_;
   // Cache device copy op for equivalence checking to reduce registry lookup
   // overhead for each invocation of call node when retrieving schedules.
   const Op& device_copy_op_;
   bool create_schedule_;
+  // Index of the global constants
+  static int const_index;
 };
 
+int ScheduleBuilder::const_index = 0;
+
 /*!
  * \brief Create schedule for target.
  * \param source_func The primitive function to be lowered.
diff --git a/src/relay/backend/te_compiler_cache.h b/src/relay/backend/te_compiler_cache.h
index 2171880fd6a5..2ffca1aa6be7 100644
--- a/src/relay/backend/te_compiler_cache.h
+++ b/src/relay/backend/te_compiler_cache.h
@@ -145,6 +145,7 @@ struct CachedFuncNode : public Object {
   tvm::Array<Integer> shape_func_param_states;
   /*! \brief The lowered functions to support the function. */
   IRModule funcs = IRModule(Map<GlobalVar, BaseFunc>({}));
+  std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors;
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("target", &target);
@@ -166,7 +167,8 @@ class CachedFunc : public ObjectRef {
   CachedFunc(tvm::Target target, GlobalVar prim_fn_name, tvm::Array<te::Tensor> inputs,
              tvm::Array<te::Tensor> outputs, te::Schedule schedule, tir::PrimFunc prim_func,
              tvm::Array<Integer> shape_func_param_states,
-             IRModule funcs = IRModule(Map<GlobalVar, BaseFunc>({})));
+             IRModule funcs = IRModule(Map<GlobalVar, BaseFunc>({})),
+             std::unordered_map<const ConstantNode*, te::Tensor> constant_tensors = {});
 
  public:
   TVM_DEFINE_OBJECT_REF_METHODS(CachedFunc, ObjectRef, CachedFuncNode);
diff --git a/src/relay/backend/utils.cc b/src/relay/backend/utils.cc
index 3aa65259d57a..990d619e4231 100644
--- a/src/relay/backend/utils.cc
+++ b/src/relay/backend/utils.cc
@@ -179,14 +179,17 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
     });
 
 ExecutorCodegenMetadata::ExecutorCodegenMetadata(
-    Array<tir::Var> inputs, Array<tir::Var> pools, Array<String> devices, Integer num_outputs,
+    Array<tir::Var> inputs, Array<TensorType> input_tensor_types, Array<String> outputs,
+    Array<TensorType> output_tensor_types, Array<tir::Var> pools, Array<String> devices,
     String executor, String mod_name, String interface_api, bool unpacked_api,
     Map<tir::Var, tir::usmp::AllocatedPoolInfo> pool_inputs) {
   auto n = make_object<ExecutorCodegenMetadataNode>();
   n->inputs = inputs;
+  n->input_tensor_types = input_tensor_types;
+  n->outputs = outputs;
+  n->output_tensor_types = output_tensor_types;
   n->pools = pools;
   n->devices = devices;
-  n->num_outputs = num_outputs;
   n->executor = executor;
   n->interface_api = interface_api;
   n->unpacked_api = unpacked_api;
@@ -294,6 +297,15 @@ void UpdateAutoSchedulerOpWeights(const IRModule& module) {
   (*te_compiler_update_weights)(weight_map);
 }
 
+std::vector<int64_t> ShapeToJSON(tvm::Array<IndexExpr> shape) {
+  std::vector<int64_t> ret;
+  for (IndexExpr dim : shape) {
+    const int64_t* pval = tir::as_const_int(dim);
+    ret.push_back(*pval);
+  }
+  return ret;
+}
+
 }  // namespace backend
 }  // namespace relay
 }  // namespace tvm
diff --git a/src/relay/backend/utils.h b/src/relay/backend/utils.h
index 595527b2aabf..3b4d4c18de89 100644
--- a/src/relay/backend/utils.h
+++ b/src/relay/backend/utils.h
@@ -35,6 +35,8 @@
 #include <tvm/te/operation.h>
 #include <tvm/tir/usmp/utils.h>
 
+#include <iostream>
+#include <sstream>
 #include <string>
 #include <typeinfo>
 #include <unordered_map>
@@ -43,6 +45,7 @@
 #include <vector>
 
 #include "../../runtime/meta_data.h"
+#include "../../target/metadata.h"
 
 namespace tvm {
 namespace relay {
@@ -61,10 +64,14 @@ class ExecutorCodegenMetadataNode : public Object {
  public:
   /*! \brief input information for the main function */
   Array<tir::Var> inputs;
+  /*! \brief input tensor type information */
+  Array<TensorType> input_tensor_types;
+  /*! \brief output information for the main function */
+  Array<String> outputs;
+  /*! \brief output tensor type information */
+  Array<TensorType> output_tensor_types;
   /*! \brief pool information for the main function */
   Array<tir::Var> pools;
-  /*! \brief number of outputs of the main function */
-  Integer num_outputs = 1;
   /*! \brief device contexts information for the main function */
   Array<String> devices;
   /*! \brief the executor to be used to run the model */
@@ -80,8 +87,10 @@ class ExecutorCodegenMetadataNode : public Object {
 
   void VisitAttrs(tvm::AttrVisitor* v) {
     v->Visit("inputs", &inputs);
+    v->Visit("input_tensor_types", &input_tensor_types);
+    v->Visit("outputs", &outputs);
+    v->Visit("output_tensor_types", &output_tensor_types);
     v->Visit("pools", &pools);
-    v->Visit("num_outputs", &num_outputs);
     v->Visit("devices", &devices);
     v->Visit("executor", &executor);
     v->Visit("unpacked_api", &unpacked_api);
@@ -97,8 +106,9 @@ class ExecutorCodegenMetadataNode : public Object {
  */
 class ExecutorCodegenMetadata : public ObjectRef {
  public:
-  TVM_DLL ExecutorCodegenMetadata(Array<tir::Var> inputs, Array<tir::Var> pools,
-                                  Array<String> devices, Integer num_outputs, String executor,
+  TVM_DLL ExecutorCodegenMetadata(Array<tir::Var> inputs, Array<TensorType> input_tensor_types,
+                                  Array<String> outputs, Array<TensorType> output_tensor_types,
+                                  Array<tir::Var> pools, Array<String> devices, String executor,
                                   String mod_name, String interface_api = "packed",
                                   bool unpacked_api = false,
                                   Map<tir::Var, tir::usmp::AllocatedPoolInfo> pool_inputs =
@@ -357,6 +367,8 @@ inline std::string DType2String(const tvm::DataType dtype) {
     os << "int";
   } else if (dtype.is_uint()) {
     os << "uint";
+  } else if (dtype.is_bfloat16()) {
+    os << "bfloat";
   } else if ((*GetPackedFunc("runtime._datatype_get_type_registered"))(dtype.code())) {
     os << "custom["
        << (*GetPackedFunc("runtime._datatype_get_type_name"))(dtype.code()).operator std::string()
@@ -522,7 +534,7 @@ Array<Pass> GetPassPrefix(bool is_homogenous, bool is_vm);
 /*! \brief Target hash function */
 struct TargetStrHash {
   /*!
-   * \brief Calculate the hash code of a Target based on the string value of the Target.
+   * \brief Calculate the hash code of a Target based on the string value of the Target KIND.
    Note that this hash should NOT be used in new usecases, equality of targets based on their
    value is not well-defined.
    This will be removed when maps from Targets to IRModules are removed from the codebase.
@@ -530,7 +542,8 @@ struct TargetStrHash {
    * \return String hash of the target
    */
   size_t operator()(const Target& target) const {
-    return String::HashBytes(target->str().c_str(), target->str().size());
+    std::string s(target->kind->name);
+    return String::HashBytes(s.c_str(), s.size());
   }
 };
 
@@ -582,6 +595,14 @@ Map<Target, IRModule> TargetStrModuleMapToTargetModuleMap(
  */
 void UpdateAutoSchedulerOpWeights(const IRModule& module);
 
+/*!
+ * \brief Extract shape from expr to vector<int64_t>
+ *
+ * \param shape
+ * \return std::vector<int64_t>
+ */
+std::vector<int64_t> ShapeToJSON(tvm::Array<IndexExpr> shape);
+
 }  // namespace backend
 }  // namespace relay
 }  // namespace tvm
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index ef8bccf42d34..e94919de7f20 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -252,21 +252,16 @@ class VMFunctionCompiler : DeviceAwareExprFunctor<void(const Expr& n)> {
       // Do that flattening on-the-fly here.
       Function inner_func = Downcast<Function>(func->body);
       std::vector<Var> params;
-      std::vector<VirtualDevice> param_virtual_devices;
       params.reserve(func->params.size() + inner_func->params.size());
-      param_virtual_devices.reserve(func->params.size() + inner_func->params.size());
       param_device_indexes.reserve(func->params.size() + inner_func->params.size());
       for (size_t i = 0; i < func->params.size(); ++i) {
         params.emplace_back(func->params[i]);
-        VirtualDevice param_virtual_device = GetFunctionParamVirtualDevice(func.get(), i);
-        param_virtual_devices.push_back(param_virtual_device);
-        param_device_indexes.push_back(GetDeviceIndex(param_virtual_device));
+        param_device_indexes.push_back(GetDeviceIndex(func->params[i]->virtual_device()));
       }
       for (size_t i = 0; i < inner_func->params.size(); ++i) {
         params.emplace_back(inner_func->params[i]);
-        VirtualDevice param_virtual_device = GetFunctionParamVirtualDevice(inner_func.get(), i);
-        param_virtual_devices.push_back(param_virtual_device);
-        param_device_indexes.push_back(GetDeviceIndex(param_virtual_device));
+
+        param_device_indexes.push_back(GetDeviceIndex(inner_func->params[i]->virtual_device()));
       }
       std::vector<TypeVar> type_params;
       type_params.reserve(func->type_params.size() + inner_func->type_params.size());
@@ -278,13 +273,12 @@ class VMFunctionCompiler : DeviceAwareExprFunctor<void(const Expr& n)> {
       }
       Function flattened_func = Function(params, inner_func->body, inner_func->ret_type,
                                          type_params, func->attrs, func->span);
-      VisitExpr(MaybeFunctionOnDevice(flattened_func, param_virtual_devices,
-                                      GetFunctionResultVirtualDevice(inner_func.get())));
+      flattened_func->virtual_device_ = inner_func->virtual_device();
+      VisitExpr(flattened_func);
     } else {
       param_device_indexes.reserve(func->params.size());
       for (size_t i = 0; i < func->params.size(); ++i) {
-        param_device_indexes.push_back(
-            GetDeviceIndex(GetFunctionParamVirtualDevice(func.get(), i)));
+        param_device_indexes.push_back(GetDeviceIndex(func->params[i]->virtual_device()));
       }
       VisitExpr(func);
     }
diff --git a/src/relay/backend/vm/lambda_lift.cc b/src/relay/backend/vm/lambda_lift.cc
index f2bd9e6b9a8a..b2912f6263dc 100644
--- a/src/relay/backend/vm/lambda_lift.cc
+++ b/src/relay/backend/vm/lambda_lift.cc
@@ -111,7 +111,6 @@ class LambdaLifter : public transform::DeviceAwareExprMutator {
     auto free_type_vars = FreeTypeVars(func, module_);
 
     Array<Var> captured_vars;
-    std::vector<VirtualDevice> captured_var_virtual_devices;
     bool recursive = false;
     for (const auto& var : free_vars) {
       if (!letrec_.empty() && var == letrec_.back()) {
@@ -119,7 +118,6 @@ class LambdaLifter : public transform::DeviceAwareExprMutator {
         continue;
       }
       captured_vars.push_back(var);
-      captured_var_virtual_devices.push_back(GetVirtualDevice(var));
     }
 
     // Freshen all the captured vars.
@@ -127,6 +125,7 @@ class LambdaLifter : public transform::DeviceAwareExprMutator {
     Map<Var, Expr> rebinding_map;
     for (auto free_var : captured_vars) {
       auto var = Var(free_var->name_hint(), free_var->checked_type());
+      var->virtual_device_ = GetVirtualDevice(free_var);
       typed_captured_vars.push_back(var);
       rebinding_map.Set(free_var, var);
     }
@@ -173,6 +172,8 @@ class LambdaLifter : public transform::DeviceAwareExprMutator {
     if (captured_vars.empty() && free_type_vars.empty()) {
       lifted_func = Function(body->params, body->body, body->ret_type, body->type_params,
                              body->attrs, body->span);
+      // We also need to copy the virtual device
+      lifted_func->virtual_device_ = body->virtual_device();
     } else {
       // When a closure is locally bound in a program, we have its full type information
       // avalible to us.
@@ -187,14 +188,14 @@ class LambdaLifter : public transform::DeviceAwareExprMutator {
       // construct the "closure" function with fully annotated arguments, no longer relying
       // on type inference.
       size_t before_arity = body->params.size();
+      VLOG(9) << "Binding " << rebinding_map << " into\n" << PrettyPrint(body->body);
       auto rebound_body = WithFields(func, func->params, Bind(body->body, rebinding_map));
       size_t after_arity = rebound_body->params.size();
       CHECK_EQ(before_arity, after_arity);
       lifted_func =
           Function(typed_captured_vars, rebound_body, /*ret_type=*/func->func_type_annotation(),
                    free_type_vars, /*attrs=*/{}, func->span);
-      lifted_func =
-          MaybeFunctionOnDevice(lifted_func, captured_var_virtual_devices, result_virtual_device);
+      lifted_func->virtual_device_ = result_virtual_device;
       lifted_func = MarkClosure(lifted_func);
     }
 
diff --git a/src/relay/ir/expr_functor.cc b/src/relay/ir/expr_functor.cc
index b710c2791acf..98e2ac0433b0 100644
--- a/src/relay/ir/expr_functor.cc
+++ b/src/relay/ir/expr_functor.cc
@@ -472,24 +472,25 @@ class ExprBinder : public MixedModeMutator, PatternMutator {
   const tvm::Map<Var, Expr>& args_map_;
 };
 
+// This function should be called SubstAndBind, since it assumes any variables introduced
+// in the substitution right hand side should be implicitly bound in the function.
 Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
   if (const FunctionNode* func = expr.as<FunctionNode>()) {
     Expr new_body = ExprBinder(args_map).VisitExpr(func->body);
     Array<Var> new_params;
-    std::vector<VirtualDevice> new_param_virtual_devices;
     for (size_t i = 0; i < func->params.size(); ++i) {
       if (!args_map.count(func->params[i])) {
         new_params.push_back(func->params[i]);
-        new_param_virtual_devices.push_back(GetFunctionParamVirtualDevice(func, i));
       }
     }
     if (new_body.same_as(func->body) && new_params.size() == func->params.size()) {
       return expr;
     }
+
     auto ret =
         Function(new_params, new_body, func->ret_type, func->type_params, func->attrs, func->span);
-    ret =
-        MaybeFunctionOnDevice(ret, new_param_virtual_devices, GetFunctionResultVirtualDevice(func));
+    ret->virtual_device_ = func->virtual_device();
+
     std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual> set;
     for (const auto& v : FreeVars(expr)) {
       set.insert(v);
@@ -497,20 +498,16 @@ Expr Bind(const Expr& expr, const tvm::Map<Var, Expr>& args_map) {
     for (const auto& v : FreeVars(ret)) {
       if (set.count(v) == 0) {
         new_params.push_back(v);
-        if (!GetFunctionResultVirtualDevice(func)->IsFullyUnconstrained()) {
-          // TODO(mbs): The function has been annotated with a device, which means we are supposed
-          // to be preserving device annotations on every transformation. However there's no
-          // such context for the free vars in args_map.
-          LOG(WARNING) << "introduced free var '" << PrettyPrint(v)
-                       << "' into function body but no device is known for it";
-        }
-        new_param_virtual_devices.push_back(VirtualDevice::FullyUnconstrained());
       }
     }
+
     ret =
         Function(new_params, new_body, func->ret_type, func->type_params, func->attrs, func->span);
-    ret =
-        MaybeFunctionOnDevice(ret, new_param_virtual_devices, GetFunctionResultVirtualDevice(func));
+    ret->virtual_device_ = func->virtual_device();
+
+    VLOG(4) << "Expr:\n" << expr;
+    VLOG(4) << "Ret:\n" << ret;
+
     ICHECK_EQ(FreeVars(expr).size(), FreeVars(ret).size());
     return std::move(ret);
   } else {
@@ -528,6 +525,27 @@ TVM_REGISTER_GLOBAL("relay.ir.Bind").set_body([](TVMArgs args, TVMRetValue* ret)
   }
 });
 
+Function SubstituteBoundVars(const Function& func, const tvm::Map<Var, Expr>& args_map) {
+  Expr new_body = ExprBinder(args_map).VisitExpr(func->body);
+  Array<Var> new_params;
+  for (size_t i = 0; i < func->params.size(); i++) {
+    if (!args_map.count(func->params[i])) {
+      new_params.push_back(func->params[i]);
+    } else {
+      if (const VarNode* var = args_map[func->params[i]].as<VarNode>()) {
+        new_params.push_back(GetRef<Var>(var));
+      } else {
+        ICHECK(false) << "Expected all values in args_map to be vars, but found "
+                      << args_map[func->params[i]]->GetTypeKey();
+      }
+    }
+  }
+  auto ret =
+      Function(new_params, new_body, func->ret_type, func->type_params, func->attrs, func->span);
+  ret->virtual_device_ = func->virtual_device();
+  return ret;
+}
+
 void ExpandANormalForm(const LetNode* op, std::function<void(const LetNode*)> pre_visit,
                        std::function<void(const LetNode*)> post_visit) {
   std::stack<const LetNode*> stack;
diff --git a/src/relay/ir/transform.cc b/src/relay/ir/transform.cc
index f8bd7081521e..1a16cc9becf1 100644
--- a/src/relay/ir/transform.cc
+++ b/src/relay/ir/transform.cc
@@ -127,7 +127,7 @@ IRModule FunctionPassNode::operator()(IRModule mod, const PassContext& pass_ctx)
   IRModule updated_mod = mod->ShallowCopy();
 
   std::vector<std::pair<GlobalVar, Function> > updates;
-  for (const auto& kv : updated_mod->functions) {
+  for (const auto& kv : mod->functions) {
     // only process optimizable Relay Functions
     if (const auto* function_node = AsOptimizableFunctionNode(kv.second)) {
       Function updated_func = pass_func(GetRef<Function>(function_node), updated_mod, pass_ctx);
diff --git a/src/relay/op/memory/on_device.cc b/src/relay/op/memory/on_device.cc
index 48e93ccf654d..155b6daf0848 100644
--- a/src/relay/op/memory/on_device.cc
+++ b/src/relay/op/memory/on_device.cc
@@ -29,6 +29,7 @@
 #include <tvm/relay/expr.h>
 #include <tvm/relay/op.h>
 #include <tvm/relay/op_attr_types.h>
+#include <tvm/relay/transform.h>
 
 #include "../../transforms/infer_layout_utils.h"
 #include "../type_relations.h"
@@ -142,47 +143,5 @@ OnDeviceProps GetOnDeviceProps(const Expr& expr) {
   return {};
 }
 
-Function FunctionOnDevice(Function function, Array<VirtualDevice> param_virtual_devices,
-                          VirtualDevice result_virtual_device) {
-  return WithAttrs(std::move(function),
-                   {{tvm::attr::kParamVirtualDevice, std::move(param_virtual_devices)},
-                    {tvm::attr::kResultVirtualDevice, std::move(result_virtual_device)}});
-}
-
-TVM_REGISTER_GLOBAL("relay.op.annotation._make.FunctionOnDevice").set_body_typed(FunctionOnDevice);
-
-Function MaybeFunctionOnDevice(Function function, Array<VirtualDevice> param_virtual_devices,
-                               VirtualDevice result_virtual_device) {
-  if (std::all_of(param_virtual_devices.begin(), param_virtual_devices.end(),
-                  [](const VirtualDevice& virtual_device) {
-                    return virtual_device->IsFullyUnconstrained();
-                  }) &&
-      result_virtual_device->IsFullyUnconstrained()) {
-    // Nothing to annotate.
-    return function;
-  }
-  return FunctionOnDevice(function, std::move(param_virtual_devices),
-                          std::move(result_virtual_device));
-}
-
-VirtualDevice GetFunctionResultVirtualDevice(const FunctionNode* function_node) {
-  auto opt_virtual_device = function_node->GetAttr<VirtualDevice>(tvm::attr::kResultVirtualDevice);
-  return opt_virtual_device.value_or(VirtualDevice::FullyUnconstrained());
-}
-
-VirtualDevice GetFunctionParamVirtualDevice(const FunctionNode* function_node, size_t i) {
-  ICHECK_LT(i, function_node->params.size())
-      << "param index " << i << " out of range for function of arity "
-      << function_node->params.size();
-  auto opt_array = function_node->GetAttr<Array<VirtualDevice>>(tvm::attr::kParamVirtualDevice);
-  if (!opt_array) {
-    // No annotation.
-    return VirtualDevice::FullyUnconstrained();
-  }
-  ICHECK_EQ(opt_array.value().size(), function_node->params.size())
-      << "annotation parameters do not match function arity";
-  return opt_array.value()[i];
-}
-
 }  // namespace relay
 }  // namespace tvm
diff --git a/src/relay/op/memory/on_device.h b/src/relay/op/memory/on_device.h
index 7489e3b62b0c..b597af8fc7fa 100644
--- a/src/relay/op/memory/on_device.h
+++ b/src/relay/op/memory/on_device.h
@@ -154,33 +154,6 @@ const NodeType* AsIgnoringOnDevice(const Expr& expr) {
   return props.body.as<NodeType>();
 }
 
-/*!
- * \brief Returns \p function annotated with "param_virtual_devices" and "result_virtual_device"
- * attributes capturing parameter and result \p VirtualDevices respectively.
- */
-Function FunctionOnDevice(Function function, Array<VirtualDevice> param_virtual_devices,
-                          VirtualDevice body_virtual_device);
-
-/*!
- * \brief As for \p FunctionOnDevice, but returns \p function unchanged if all parameters and
- * result \p VirtualDevices are unconstrained.
- */
-Function MaybeFunctionOnDevice(Function function, Array<VirtualDevice> param_virtual_devices,
-                               VirtualDevice result_virtual_device);
-
-/*!
- * \brief Returns the \p VirtualDevice for the resut of \p function_node, or the unconstrained
- * \p VirtualDevice if function does not have the "result_virtual_device" annotation.
- */
-VirtualDevice GetFunctionResultVirtualDevice(const FunctionNode* function_node);
-
-/*!
- * \brief Returns the \p VirtualDevice for the \p i'th parameter of \p function_node, or
- * the unconstrained \p VirtualDevice if function does not have the "param_virtual_devices"
- * annotation.
- */
-VirtualDevice GetFunctionParamVirtualDevice(const FunctionNode* function_node, size_t i);
-
 }  // namespace relay
 }  // namespace tvm
 
diff --git a/src/relay/op/nn/convolution.cc b/src/relay/op/nn/convolution.cc
index 446f34340caf..447bbd4926b4 100644
--- a/src/relay/op/nn/convolution.cc
+++ b/src/relay/op/nn/convolution.cc
@@ -54,6 +54,96 @@ Expr MakeConvGemmWeightTransform(Expr weight, int tile_rows, int tile_cols, std:
 // relay.nn.conv1d
 TVM_REGISTER_NODE_TYPE(Conv1DAttrs);
 
+bool Conv1DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+               const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+  static const Layout kNCW("NCW");
+  static const Layout kOIW("OIW");
+
+  const auto* param = attrs.as<Conv1DAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv only support input layouts that are convertible from NCW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv only support kernel layouts that are convertible from OIW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv only support output layouts that are convertible from NCW."
+      << " But got " << out_layout;
+
+  Array<IndexExpr> dshape_ncw = trans_in_layout.ForwardShape(data->shape);
+
+  IndexExpr channels, dilated_ksize;
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    Array<IndexExpr> wshape;
+
+    wshape = {{param->channels, indexdiv(dshape_ncw[1], param->groups), param->kernel_size[0]}};
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    channels = param->channels;
+    dilated_ksize = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+    // assign result to reporter
+    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
+    if (param->kernel_size.defined()) {
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]))
+          << "Conv1D: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
+    }
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
+          << "Conv1D: shape of weight is inconsistent with channels, "
+          << " channels=" << param->channels << " wshape=" << wshape;
+    }
+    if (!dshape_ncw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(dshape_ncw[1], wshape[1]));
+    }
+    channels = wshape[0];
+    dilated_ksize = 1 + (wshape[2] - 1) * param->dilation[0];
+  }
+  // dilation
+  Array<IndexExpr> oshape({dshape_ncw[0], channels, 0});
+
+  if (!dshape_ncw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, indexdiv(dshape_ncw[2] + param->padding[0] + param->padding[1] - dilated_ksize,
+                           param->strides[0]) +
+                      1);
+  } else {
+    oshape.Set(2, dshape_ncw[2]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv1d")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -82,12 +172,190 @@ with the layer input to produce a tensor of outputs.
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(2)
-    .add_type_rel("Conv1D", Conv1DRel<Conv1DAttrs>)
+    .add_type_rel("Conv1D", Conv1DRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv1DAttrs>);
 
 // relay.nn.conv2d
 TVM_REGISTER_NODE_TYPE(Conv2DAttrs);
 
+bool Conv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+               const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+  static const Layout kNCHW("NCHW");
+  static const Layout kOIHW("OIHW");
+
+  const auto* param = attrs.as<Conv2DAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
+  if (!trans_in_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "conv2d only support input layouts that are convertible from NCHW."
+        << " The provided layout is: " << in_layout);
+    return false;
+  }
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIHW);
+  if (!trans_kernel_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "conv2d only support kernel layouts that are convertible from OIHW."
+        << " The provided layout is: " << kernel_layout);
+    return false;
+  }
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
+  if (!trans_out_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "conv2d only support output layouts that are convertible from NCHW."
+        << "The provided layout is: " << out_layout);
+    return false;
+  }
+
+  Array<IndexExpr> dshape_nchw = trans_in_layout.ForwardShape(data->shape);
+  bool is_depthwise = false;
+  if (param->groups > 1) {
+    if (!(weight && weight->shape.defined())) {
+      reporter->GetDiagCtx().Emit(
+          Diagnostic::Error(reporter->GetSpan())
+          << "Weight shape must be specified when groups is greater than 1.");
+      return false;
+    }
+
+    Array<IndexExpr> wshape_oihw = trans_kernel_layout.ForwardShape(weight->shape);
+    if (tvm::tir::ExprDeepEqual()(param->groups, dshape_nchw[1]) &&
+        tvm::tir::ExprDeepEqual()(param->groups, wshape_oihw[0])) {
+      is_depthwise = true;
+    }
+  }
+
+  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 2);
+    ICHECK_EQ(param->dilation.size(), 2);
+    Array<IndexExpr> wshape;
+
+    if (is_depthwise) {
+      // infer weight's shape for depthwise convolution
+      wshape = {{dshape_nchw[1], indexdiv(param->channels, dshape_nchw[1]), param->kernel_size[0],
+                 param->kernel_size[1]}};
+    } else {
+      wshape = {{param->channels, indexdiv(dshape_nchw[1], param->groups), param->kernel_size[0],
+                 param->kernel_size[1]}};
+    }
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    channels = param->channels;
+    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+
+    if (param->auto_scheduler_rewritten_layout.size() == 0) {
+      // Normal case: assign result to reporter
+      reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+    } else {
+      // If the layout is rewritten by auto-scheduler,
+      // we just forcly apply the layout provided by auto-scheduler and
+      // skip the normal inference logic.
+      {}  // do nothing
+    }
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+
+    Array<PrimExpr> wshape;
+    if (param->auto_scheduler_rewritten_layout.size() == 0) {
+      wshape = weight->shape;
+    } else {
+      // works for the default kernel layout "HWIO"
+      ICHECK_EQ(param->kernel_layout, "HWIO");
+      wshape = auto_scheduler::GetShapeFromRewrittenLayout(param->auto_scheduler_rewritten_layout,
+                                                           {"ry", "rx", "rc", "ff"});
+    }
+
+    wshape = trans_kernel_layout.ForwardShape(wshape);
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 2);
+
+      if (!reporter->AssertEQ(param->kernel_size[0], wshape[2])) {
+        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
+                                    << "Conv2D: shape of weight is inconsistent with kernel_size,"
+                                    << " kernel_size=" << param->kernel_size
+                                    << " wshape=" << wshape);
+      }
+
+      if (!reporter->AssertEQ(param->kernel_size[1], wshape[3])) {
+        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
+                                    << "Conv2D: shape of weight is inconsistent with kernel_size,"
+                                    << " kernel_size=" << param->kernel_size
+                                    << " wshape=" << wshape);
+        return false;
+      }
+    }
+
+    if (param->channels.defined() && !reporter->AssertEQ(param->channels, wshape[0])) {
+      reporter->GetDiagCtx().Emit(
+          Diagnostic::Error(reporter->GetSpan())
+          << "conv2D: the first dimensions of the weight tensor (" << wshape << ")"
+          << "does not match the number of channels (" << param->channels << ").");
+      return false;
+    }
+
+    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
+      if (!reporter->AssertEQ(indexdiv(dshape_nchw[1], param->groups), wshape[1])) {
+        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
+                                    << "conv2d: requires that `"
+                                    << indexdiv(dshape_nchw[1], param->groups) << "`,"
+                                    << " the input channels (" << dshape_nchw[1] << ")"
+                                    << " divided by groups (" << param->groups << ")"
+                                    << ",\n must match the input channels"
+                                    << " of the weight `" << wshape[1]
+                                    << "`, where the weight shape is (" << wshape << ").");
+        return false;
+      }
+    }
+    channels = wshape[0];
+    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
+  }
+  // dilation
+  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
+
+  IndexExpr pad_h, pad_w;
+  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
+  if (!dshape_nchw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, indexdiv(dshape_nchw[2] + pad_h - dilated_ksize_y, param->strides[0]) + 1);
+  } else {
+    oshape.Set(2, dshape_nchw[2]);
+  }
+
+  if (!dshape_nchw[3].as<tir::AnyNode>()) {
+    oshape.Set(3, indexdiv(dshape_nchw[3] + pad_w - dilated_ksize_x, param->strides[1]) + 1);
+  } else {
+    oshape.Set(3, dshape_nchw[3]);
+  }
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -116,12 +384,152 @@ with the layer input to produce a tensor of outputs.
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(2)
-    .add_type_rel("Conv2D", Conv2DRel<Conv2DAttrs>)
+    .add_type_rel("Conv2D", Conv2DRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv2DAttrs>);
 
 // relay.nn.conv3d
 TVM_REGISTER_NODE_TYPE(Conv3DAttrs);
 
+bool Conv3DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+               const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+  static const Layout kNCDHW("NCDHW");
+  static const Layout kOIDHW("OIDHW");
+
+  const auto* param = attrs.as<Conv3DAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv only support input layouts that are convertible from NCDHW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv only support kernel layouts that are convertible from OIDHW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv only support output layouts that are convertible from NCDHW."
+      << " But got " << out_layout;
+
+  Array<IndexExpr> dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
+
+  IndexExpr channels, dilated_ksize_z, dilated_ksize_y, dilated_ksize_x;
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 3);
+    ICHECK_EQ(param->dilation.size(), 3);
+    Array<IndexExpr> wshape;
+    tvm::tir::ExprDeepEqual expr_equal;
+
+    if (expr_equal(param->channels, param->groups) && !expr_equal(param->channels, 1)) {
+      // infer weight's shape for depthwise convolution
+      wshape = {{dshape_ncdhw[1], indexdiv(param->groups, dshape_ncdhw[1]), param->kernel_size[0],
+                 param->kernel_size[1], param->kernel_size[2]}};
+    } else {
+      wshape = {{param->channels, indexdiv(dshape_ncdhw[1], param->groups), param->kernel_size[0],
+                 param->kernel_size[1], param->kernel_size[2]}};
+    }
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    channels = param->channels;
+    dilated_ksize_z = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+    dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+
+    if (param->auto_scheduler_rewritten_layout.size() == 0) {
+      // Normal case: assign result to reporter
+      reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+    } else {
+      // If the layout is rewritten by auto-scheduler,
+      // we just forcly apply the layout provided by auto-scheduler and
+      // skip the normal inference logic.
+      {}  // do nothing
+    }
+
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+
+    Array<PrimExpr> wshape;
+    if (param->auto_scheduler_rewritten_layout.size() == 0) {
+      wshape = weight->shape;
+    } else {
+      // works for the default kernel layout "DHWIO"
+      ICHECK_EQ(param->kernel_layout, "DHWIO");
+      wshape = auto_scheduler::GetShapeFromRewrittenLayout(param->auto_scheduler_rewritten_layout,
+                                                           {"rd", "rh", "rw", "rc", "cc"});
+    }
+
+    wshape = trans_kernel_layout.ForwardShape(wshape);
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 3);
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
+             reporter->AssertEQ(param->kernel_size[1], wshape[3]) &&
+             reporter->AssertEQ(param->kernel_size[2], wshape[4]))
+          << "Conv3D: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
+    }
+
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
+          << "Conv3D: shape of weight is inconsistent with channels, "
+          << " channels=" << param->channels << " wshape=" << wshape;
+    }
+
+    if (!dshape_ncdhw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncdhw[1], param->groups), wshape[1]));
+    }
+    channels = wshape[0];
+    dilated_ksize_z = 1 + (wshape[2] - 1) * param->dilation[0];
+    dilated_ksize_y = 1 + (wshape[3] - 1) * param->dilation[1];
+    dilated_ksize_x = 1 + (wshape[4] - 1) * param->dilation[2];
+  }
+  // dilation
+  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
+
+  IndexExpr pad_d, pad_h, pad_w;
+  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
+  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, indexdiv(dshape_ncdhw[2] + pad_d - dilated_ksize_z, param->strides[0]) + 1);
+  } else {
+    oshape.Set(2, dshape_ncdhw[2]);
+  }
+
+  if (!dshape_ncdhw[3].as<tir::AnyNode>()) {
+    oshape.Set(3, indexdiv(dshape_ncdhw[3] + pad_h - dilated_ksize_y, param->strides[1]) + 1);
+  } else {
+    oshape.Set(3, dshape_ncdhw[3]);
+  }
+
+  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
+    oshape.Set(4, indexdiv(dshape_ncdhw[4] + pad_w - dilated_ksize_x, param->strides[2]) + 1);
+  } else {
+    oshape.Set(4, dshape_ncdhw[4]);
+  }
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv3d")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -151,12 +559,128 @@ with the layer input to produce a tensor of outputs.
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(2)
-    .add_type_rel("Conv3D", Conv3DRel<Conv3DAttrs>)
+    .add_type_rel("Conv3D", Conv3DRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv3DAttrs>);
 
 // relay.nn.conv3d_transpose
 TVM_REGISTER_NODE_TYPE(Conv3DTransposeAttrs);
 
+template <typename AttrType>
+bool Conv3DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                        const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  static const Layout kNCDHW("NCDHW");
+  static const Layout kOIDHW("OIDHW");
+
+  const Conv3DTransposeAttrs* param = attrs.as<Conv3DTransposeAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv3d_transpose only support input layouts that are convertible from NCDHW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv3d_transpose only support kernel layouts that are convertible from OIDHW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv3d_transpose only support output layouts that are convertible from NCDHW."
+      << " But got " << out_layout;
+
+  IndexExpr channels, dilated_ksize_d, dilated_ksize_y, dilated_ksize_x;
+
+  auto dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
+
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 3);
+    ICHECK_EQ(param->dilation.size(), 3);
+
+    Array<IndexExpr> wshape({dshape_ncdhw[1], indexdiv(param->channels, param->groups),
+                             param->kernel_size[0], param->kernel_size[1], param->kernel_size[2]});
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    dilated_ksize_d = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+    dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
+    channels = param->channels;
+
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+    // assign result to reporter
+    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 3);
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
+             reporter->AssertEQ(param->kernel_size[1], wshape[3]) &&
+             reporter->AssertEQ(param->kernel_size[2], wshape[4]))
+          << "Conv3D: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
+    }
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(param->channels, wshape[1]))
+          << "Conv3D: shape of weight is inconsistent with channels, "
+          << " channels=" << param->channels << " wshape=" << Array<IndexExpr>(wshape);
+    }
+    if (!dshape_ncdhw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncdhw[1], param->groups), wshape[0]));
+    }
+    channels = wshape[1];
+    dilated_ksize_d = 1 + (wshape[2] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
+    dilated_ksize_y = 1 + (wshape[4] - 1) * param->dilation[2];
+  }
+
+  // dilation
+  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
+  IndexExpr pad_d, pad_h, pad_w;
+  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
+
+  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, (param->strides[0] * (dshape_ncdhw[2] - 1) + dilated_ksize_d - pad_d +
+                   param->output_padding[0]));
+  } else {
+    oshape.Set(2, dshape_ncdhw[2]);
+  }
+  if (!dshape_ncdhw[3].as<tir::AnyNode>()) {
+    oshape.Set(3, (param->strides[1] * (dshape_ncdhw[3] - 1) + dilated_ksize_y - pad_h +
+                   param->output_padding[1]));
+  } else {
+    oshape.Set(3, dshape_ncdhw[3]);
+  }
+  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
+    oshape.Set(4, (param->strides[2] * (dshape_ncdhw[4] - 1) + dilated_ksize_x - pad_w +
+                   param->output_padding[2]));
+  } else {
+    oshape.Set(4, dshape_ncdhw[4]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv3d_transpose")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -202,6 +726,115 @@ said convolution.
 // relay.nn.conv2d_transpose
 TVM_REGISTER_NODE_TYPE(Conv2DTransposeAttrs);
 
+bool Conv2DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                        const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  static const Layout kNCHW("NCHW");
+  static const Layout kIOHW("IOHW");
+
+  const Conv2DTransposeAttrs* param = attrs.as<Conv2DTransposeAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv2DTransposed only support input layouts that are convertible from NCHW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kIOHW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv2DTransposed only support kernel layouts that are convertible from IOHW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv2DTransposed only support output layouts that are convertible from NCHW."
+      << " But got " << out_layout;
+
+  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
+
+  auto dshape_nchw = trans_in_layout.ForwardShape(data->shape);
+
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 2);
+    ICHECK_EQ(param->dilation.size(), 2);
+
+    Array<IndexExpr> wshape({dshape_nchw[1], indexdiv(param->channels, param->groups),
+                             param->kernel_size[0], param->kernel_size[1]});
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+    channels = param->channels;
+
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+    // assign result to reporter
+    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 2);
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
+             reporter->AssertEQ(param->kernel_size[1], wshape[3]))
+          << "Conv2DTransposed: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
+    }
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(indexdiv(param->channels, param->groups), wshape[1]))
+          << "Conv2DTransposed: shape of weight is inconsistent with out_channels, "
+          << " out_channels // groups != weight.shape[1] "
+          << " out_channels=" << param->channels << " groups=" << param->groups
+          << " weight.shape=" << Array<IndexExpr>(wshape);
+    }
+    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(dshape_nchw[1], wshape[0]))
+          << "Conv2DTransposed: shape of weight is inconsistent with in_channels."
+          << " data.shape= " << Array<IndexExpr>(dshape_nchw) << " groups= " << param->groups
+          << " weight.shape= " << Array<IndexExpr>(wshape);
+    }
+    channels = wshape[1];
+    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
+  }
+  // dilation
+  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
+  IndexExpr pad_h, pad_w;
+  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
+  if (!dshape_nchw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, (param->strides[0] * (dshape_nchw[2] - 1) + dilated_ksize_y - pad_h +
+                   param->output_padding[0]));
+  } else {
+    oshape.Set(2, dshape_nchw[2]);
+  }
+  if (!dshape_nchw[3].as<tir::AnyNode>()) {
+    oshape.Set(3, (param->strides[1] * (dshape_nchw[3] - 1) + dilated_ksize_x - pad_w +
+                   param->output_padding[1]));
+  } else {
+    oshape.Set(3, dshape_nchw[3]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d_transpose")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -241,11 +874,106 @@ v            (batch_size, channels, out_height, out_width) if `layout` is `NCHW`
     .set_support_level(2)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout",
                                    ConvInferCorrectLayout<Conv2DTransposeAttrs>)
-    .add_type_rel("Conv2DTranspose", Conv2DTransposeRel<Conv2DTransposeAttrs>);
+    .add_type_rel("Conv2DTranspose", Conv2DTransposeRel);
 
 // relay.nn.conv1d_transpose
 TVM_REGISTER_NODE_TYPE(Conv1DTransposeAttrs);
 
+bool Conv1DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                        const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[1].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  static const Layout kNCW("NCW");
+  static const Layout kOIW("OIW");
+
+  const Conv1DTransposeAttrs* param = attrs.as<Conv1DTransposeAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv only support input layouts that are convertible from NCW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv only support kernel layouts that are convertible from OIW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv only support output layouts that are convertible from NCW."
+      << " But got " << out_layout;
+
+  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
+
+  auto dshape_ncw = trans_in_layout.ForwardShape(data->shape);
+
+  // infer weight if the kernel_size and channels are defined
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 1);
+    ICHECK_EQ(param->dilation.size(), 1);
+
+    Array<IndexExpr> wshape(
+        {dshape_ncw[1], indexdiv(param->channels, param->groups), param->kernel_size[0]});
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    dilated_ksize_x = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    channels = param->channels;
+
+    DataType weight_dtype = data->dtype;
+    if (weight != nullptr) {
+      weight_dtype = weight->dtype;
+    }
+    // assign result to reporter
+    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 1);
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]))
+          << "Conv1D: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
+    }
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(param->channels, wshape[1]))
+          << "Conv1D: shape of weight is inconsistent with channels, "
+          << " channels=" << param->channels << " wshape=" << Array<IndexExpr>(wshape);
+    }
+    if (!dshape_ncw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncw[1], param->groups), wshape[0]));
+    }
+    channels = wshape[1];
+    dilated_ksize_x = 1 + (wshape[2] - 1) * param->dilation[0];
+  }
+  // dilation
+  IndexExpr pad_w;
+  GetPaddingWidth(param->padding, &pad_w);
+  Array<IndexExpr> oshape({dshape_ncw[0], channels, 0});
+  if (!dshape_ncw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, (param->strides[0] * (dshape_ncw[2] - 1) + dilated_ksize_x - pad_w +
+                   param->output_padding[0]));
+  } else {
+    oshape.Set(2, dshape_ncw[2]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.conv1d_transpose")
     .set_body_typed([](Expr data, Expr weight, Array<IndexExpr> strides, Array<IndexExpr> padding,
                        Array<IndexExpr> dilation, int groups, IndexExpr channels,
@@ -282,7 +1010,7 @@ said convolution.
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(2)
-    .add_type_rel("Conv1DTranspose", Conv1DTransposeRel<Conv1DTransposeAttrs>);
+    .add_type_rel("Conv1DTranspose", Conv1DTransposeRel);
 
 // relay.nn.contrib_conv2d_winograd_without_weight_transform
 TVM_REGISTER_NODE_TYPE(Conv2DWinogradAttrs);
@@ -322,6 +1050,28 @@ RELAY_REGISTER_OP("nn.contrib_conv2d_winograd_without_weight_transform")
 // relay.nn.contrib_conv2d_winograd_weight_transform
 TVM_REGISTER_NODE_TYPE(ConvWinogradWeightTransformAttrs);
 
+bool Conv2DWinogradWeightTransformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                                      const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  const ConvWinogradWeightTransformAttrs* param = attrs.as<ConvWinogradWeightTransformAttrs>();
+  ICHECK(param != nullptr);
+
+  ICHECK_EQ(data->shape.size(), 4) << "Only support NCHW normal kernel layout";
+
+  std::vector<IndexExpr> oshape{
+      param->tile_size + data->shape[2] - 1,
+      param->tile_size + data->shape[3] - 1,
+      data->shape[0],
+      data->shape[1],
+  };
+
+  reporter->Assign(types[1], TensorType(Array<IndexExpr>(oshape), data->dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_conv2d_winograd_weight_transform")
     .set_body_typed([](Expr weight, int tile_size) {
       return MakeConvWinogradWeightTransform(weight, tile_size,
@@ -345,6 +1095,88 @@ weight transformation in advance.
 // relay.nn.contrib_conv3d_winograd_without_weight_transform
 TVM_REGISTER_NODE_TYPE(Conv3DWinogradAttrs);
 
+bool Conv3DWinogradRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                       const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+  static const Layout kNCDHW("NCDHW");
+  static const Layout kOIDHW("OIDHW");
+
+  const auto* param = attrs.as<Conv3DWinogradAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
+  ICHECK(trans_in_layout.defined())
+      << "Conv only support input layouts that are convertible from NCDHW."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv only support kernel layouts that are convertible from OIDHW."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
+  ICHECK(trans_out_layout.defined())
+      << "Conv only support output layouts that are convertible from NCDHW."
+      << " But got " << out_layout;
+
+  Array<IndexExpr> dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
+
+  IndexExpr channels, dilated_ksize_d, dilated_ksize_y, dilated_ksize_x;
+
+  ICHECK(param->kernel_size.defined() && param->channels.defined())
+      << "The kernel size and channels of a Conv must be set or inferred by previous pass";
+
+  ICHECK_EQ(param->kernel_size.size(), 3);
+  ICHECK_EQ(param->dilation.size(), 3);
+
+  channels = param->channels;
+  dilated_ksize_d = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+  dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+  dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
+
+  // NOTE: Do not check weight shape here!
+  // Different backend requires different layout to compute
+  // the batch gemm stage in winograd efficiently, but we want to
+  // make this op work for all backends.
+  // So we accept all weight shapes, and assume the TOPI developers
+  // can handle this correctly in alter_op_layout.
+
+  // dilation
+  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
+
+  IndexExpr pad_d, pad_h, pad_w;
+  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
+  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
+    oshape.Set(2, (dshape_ncdhw[2] + pad_d - dilated_ksize_d) / param->strides[0] + 1);
+  } else {
+    oshape.Set(2, dshape_ncdhw[2]);
+  }
+  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
+    oshape.Set(3, (dshape_ncdhw[3] + pad_h - dilated_ksize_y) / param->strides[1] + 1);
+  } else {
+    oshape.Set(3, dshape_ncdhw[3]);
+  }
+  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
+    oshape.Set(4, (dshape_ncdhw[4] + pad_w - dilated_ksize_x) / param->strides[2] + 1);
+  } else {
+    oshape.Set(4, dshape_ncdhw[4]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_conv3d_winograd_without_weight_transform")
     .set_body_typed([](Expr data, Expr weight, int tile_size, Array<IndexExpr> strides,
                        Array<IndexExpr> padding, Array<IndexExpr> dilation, int groups,
@@ -373,7 +1205,7 @@ RELAY_REGISTER_OP("nn.contrib_conv3d_winograd_without_weight_transform")
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(10)
-    .add_type_rel("Conv3DWinograd", Conv3DWinogradRel<Conv3DWinogradAttrs>)
+    .add_type_rel("Conv3DWinograd", Conv3DWinogradRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout",
                                    ConvInferCorrectLayout<Conv3DWinogradAttrs>);
 
@@ -384,6 +1216,35 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_conv3d_winograd_weight_transform"
                                              "nn.contrib_conv3d_winograd_weight_transform");
     });
 
+bool Conv3DWinogradWeightTransformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                                      const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+
+  const ConvWinogradWeightTransformAttrs* param = attrs.as<ConvWinogradWeightTransformAttrs>();
+  ICHECK(param != nullptr);
+
+  ICHECK_EQ(data->shape.size(), 5) << "Only support NCDHW normal kernel layout";
+
+  // Shape of packed weights depends on whether depth is being transformed or not.
+  Array<IndexExpr> oshape({0, 0, 0, data->shape[0], data->shape[1]});
+  auto* depth_imm = data->shape[2].as<IntImmNode>();
+  bool transform_depth = (depth_imm->value > 2) && (depth_imm->value < 8);
+  if (transform_depth) {
+    oshape.Set(0, param->tile_size + data->shape[2] - 1);
+    oshape.Set(1, param->tile_size + data->shape[3] - 1);
+    oshape.Set(2, param->tile_size + data->shape[4] - 1);
+  } else {
+    oshape.Set(0, param->tile_size + data->shape[3] - 1);
+    oshape.Set(1, param->tile_size + data->shape[4] - 1);
+    oshape.Set(2, data->shape[2]);
+  }
+
+  reporter->Assign(types[1], TensorType(oshape, data->dtype));
+  return true;
+}
+
 RELAY_REGISTER_OP("nn.contrib_conv3d_winograd_weight_transform")
     .describe(R"code(Weight transformation of winograd fast 3d convolution algorithm.
 
@@ -401,6 +1262,35 @@ weight transformation in advance.
 // relay.nn.contrib_conv2d_winograd_nnpack_weight_transform
 TVM_REGISTER_NODE_TYPE(Conv2DWinogradNNPACKWeightTransformAttrs);
 
+bool Conv2DWinogradNNPACKWeightTransformRel(const Array<Type>& types, int num_inputs,
+                                            const Attrs& attrs, const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) {
+    return false;
+  }
+
+  const Conv2DWinogradNNPACKWeightTransformAttrs* param =
+      attrs.as<Conv2DWinogradNNPACKWeightTransformAttrs>();
+  ICHECK(param != nullptr);
+
+  ICHECK_EQ(data->shape.size(), 4) << "Only support NCHW normal kernel layout";
+
+  std::vector<IndexExpr> oshape{
+      data->shape[0],
+      data->shape[1],
+      8,
+      8,
+  };
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  reporter->Assign(types[1], TensorType(Array<IndexExpr>(oshape), out_dtype));
+  return true;
+}
+
 Expr MakeConv2DWinogradNNPACKWeightTransform(Expr weight, int convolution_algorithm,
                                              DataType out_dtype) {
   auto attrs = make_object<Conv2DWinogradNNPACKWeightTransformAttrs>();
@@ -438,6 +1328,77 @@ TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_conv2d_gemm_without_weight_transf
           kernel_layout, out_layout, out_dtype, "nn.contrib_conv2d_gemm_without_weight_transform");
     });
 
+bool Conv2DGemmRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                   const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 3);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) return false;
+  static const Layout kNHWC("NHWC");
+  static const Layout kHWIO("HWIO");
+
+  const auto* param = attrs.as<Conv2DAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNHWC);
+  ICHECK(trans_in_layout.defined())
+      << "Conv only support input layouts that are convertible from NHWC."
+      << " But got " << in_layout;
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kHWIO);
+  ICHECK(trans_kernel_layout.defined())
+      << "Conv only support kernel layouts that are convertible from HWIO."
+      << " But got " << kernel_layout;
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNHWC);
+  ICHECK(trans_out_layout.defined())
+      << "Conv only support output layouts that are convertible from NHWC."
+      << " But got " << out_layout;
+
+  Array<IndexExpr> dshape_nhwc = trans_in_layout.ForwardShape(data->shape);
+
+  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
+
+  ICHECK(param->kernel_size.defined() && param->channels.defined())
+      << "The kernel size and channels of a Conv must be set or inferred by previous pass";
+
+  ICHECK_EQ(param->kernel_size.size(), 2);
+  ICHECK_EQ(param->dilation.size(), 2);
+
+  channels = param->channels;
+  dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+  dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+
+  // NOTE: Do not check weight shape here!
+
+  // dilation
+  Array<IndexExpr> oshape({dshape_nhwc[0], 0, 0, channels});
+
+  IndexExpr pad_h, pad_w;
+  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
+  if (!dshape_nhwc[2].as<tir::AnyNode>()) {
+    oshape.Set(1, (dshape_nhwc[1] + pad_h - dilated_ksize_y) / param->strides[0] + 1);
+  } else {
+    oshape.Set(1, dshape_nhwc[1]);
+  }
+  if (!dshape_nhwc[3].as<tir::AnyNode>()) {
+    oshape.Set(2, (dshape_nhwc[2] + pad_w - dilated_ksize_x) / param->strides[1] + 1);
+  } else {
+    oshape.Set(2, dshape_nhwc[2]);
+  }
+
+  DataType out_dtype = param->out_dtype;
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+  oshape = trans_out_layout.BackwardShape(oshape);
+  // assign output type
+  reporter->Assign(types[2], TensorType(oshape, out_dtype));
+  return true;
+}
+
 RELAY_REGISTER_OP("nn.contrib_conv2d_gemm_without_weight_transform")
     .describe(R"code(Compute conv2d with gemm algorithm. Only supports NHWC layout.
                  This operator assumes the weight tensor is already pre-transformed by
@@ -455,13 +1416,72 @@ RELAY_REGISTER_OP("nn.contrib_conv2d_gemm_without_weight_transform")
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(10)
-    .add_type_rel("Conv2DGemm", Conv2DGemmRel<Conv2DAttrs>)
+    .add_type_rel("Conv2DGemm", Conv2DGemmRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv2DAttrs>);
 
 // relay.nn.contrib_conv2d_gemm_weight_transform
 
 TVM_REGISTER_NODE_TYPE(ConvGemmWeightTransformAttrs);
 
+// Gemm convolution shape relations
+// In order to run GEMM we need to block-transpose and interleave the K x N weights matrix W.
+// The high level idea is to subdivide W in tiles of tile_cols x tile_rows, and transpose and
+// interleave them. The final output is a [N//tile_rows, K//tile_cols, tile_rows, tile_cols]
+// matrix that we call W_interleaved_t.
+//
+// In the following picture, we show how the first [tile_cols,tile_rows] block of W is transformed
+// for tile_rows = 4 and tile_cols = 16
+//
+//              W[0,0,:,:]                        W_interleaved_t[0,0,:,:]
+//  +-------------------------------+     +----------------------------------- +
+//  |W[0,0]  W[0,1]  W[0,2]  W[0,3] |     |W[0,0]  W[1,0]  W[2,0]  ...  W[15,0]|
+//  |W[1,0]  W[1,1]  W[1,2]  W[1,3] | --\ |W[0,1]  W[1,1]  W[2,1]  ...  W[15,1]|
+//  |W[2,0]  W[2,1]  W[2,2]  W[2,3] | --/ |W[0,2]  W[1,2]  W[2,2]  ...  W[15,2]|
+//  |  ...     ...    ...      ...  |     |W[0,3]  W[1,3]  W[2,3]  ...  W[15,3]|
+//  |  ...     ...    ...      ...  |     +------------------------------------+
+//  |W[15,0] W[15,1] W[15,2] W[15,3]|
+//  +-------------------------------+
+//
+// Tile columns is usually the direction of the reduction. So, if our target can reduce k elements
+// at the time, we should set tile_cols = k.
+// Tile rows is connected with the number of registers available for the given target.
+//
+bool Conv2DGemmWeightTransformRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                                  const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 2);
+  const auto* weight = types[0].as<TensorTypeNode>();
+  if (weight == nullptr) return false;
+
+  const ConvGemmWeightTransformAttrs* param = attrs.as<ConvGemmWeightTransformAttrs>();
+  ICHECK(param != nullptr);
+  int n = param->tile_rows;
+  int k = param->tile_cols;
+
+  ICHECK_EQ(weight->shape.size(), 4) << "Only support HWIO kernel layout";
+
+  const auto K = weight->shape[0] * weight->shape[1] * weight->shape[2];
+  const auto N = weight->shape[3];
+
+  auto K_mod_k = indexmod(K, k);
+  auto N_mod_n = indexmod(N, n);
+
+  auto pad_K = tvm::if_then_else(K_mod_k != 0, k - K_mod_k, tir::make_zero(DataType::Int(32)));
+  auto pad_N = tvm::if_then_else(N_mod_n != 0, n - N_mod_n, tir::make_zero(DataType::Int(32)));
+
+  const auto N_padded = N + pad_N;
+  const auto K_padded = K + pad_K;
+
+  Array<IndexExpr> oshape{
+      indexdiv(N_padded, n),
+      indexdiv(K_padded, k),
+      n,
+      k,
+  };
+
+  reporter->Assign(types[1], TensorType(oshape, weight->dtype));
+  return true;
+}
+
 TVM_REGISTER_GLOBAL("relay.op.nn._make.contrib_conv2d_gemm_weight_transform")
     .set_body_typed([](Expr weights, int tile_rows, int tile_cols) {
       return MakeConvGemmWeightTransform(weights, tile_rows, tile_cols,
@@ -532,11 +1552,133 @@ RELAY_REGISTER_OP("nn.contrib_depthwise_conv2d_NCHWc")
     .add_argument("data", "Tensor", "The input tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(10)
-    .add_type_rel("Conv2D", Conv2DRel<Conv2DAttrs>)
+    .add_type_rel("Conv2D", Conv2DRel)
     .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ConvInferCorrectLayout<Conv2DAttrs>);
 
 TVM_REGISTER_NODE_TYPE(DeformableConv2DAttrs);
 
+// Deformable Convolution shape relations.
+bool DeformableConv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                         const TypeReporter& reporter) {
+  ICHECK_EQ(types.size(), 4);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto* weight = types[2].as<TensorTypeNode>();
+
+  ICHECK(data);
+  static const Layout kNCHW("NCHW");
+  static const Layout kOIHW("OIHW");
+
+  auto* param = attrs.as<DeformableConv2DAttrs>();
+  ICHECK(param != nullptr);
+  const Layout in_layout(param->data_layout);
+  const Layout kernel_layout(param->kernel_layout);
+
+  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
+  if (!trans_in_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "deformable_conv2d only support input layouts that are convertible from NCHW."
+        << " The provided layout is: " << in_layout);
+    return false;
+  }
+
+  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIHW);
+  if (!trans_kernel_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "deformable_conv2d only support kernel layouts that are convertible from OIHW."
+        << " The provided layout is: " << kernel_layout);
+    return false;
+  }
+
+  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
+  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
+  if (!trans_out_layout.defined()) {
+    reporter->GetDiagCtx().Emit(
+        Diagnostic::Error(reporter->GetSpan())
+        << "deformable_conv2d only support output layouts that are convertible from NCHW."
+        << "The provided layout is: " << out_layout);
+    return false;
+  }
+
+  Array<IndexExpr> dshape_nchw = trans_in_layout.ForwardShape(data->shape);
+
+  IndexExpr channels, dilated_ksize_y, dilated_ksize_x, ksize_y, ksize_x;
+
+  // infer weight shape if kernel_size and channels are defiend
+  if (param->kernel_size.defined() && param->channels.defined()) {
+    ICHECK_EQ(param->kernel_size.size(), 2);
+    ICHECK_EQ(param->dilation.size(), 2);
+    Array<IndexExpr> wshape({param->channels, indexdiv(dshape_nchw[1], param->groups),
+                             param->kernel_size[0], param->kernel_size[1]});
+
+    wshape = trans_kernel_layout.BackwardShape(wshape);
+    channels = param->channels;
+    ksize_y = param->kernel_size[0];
+    ksize_x = param->kernel_size[1];
+    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
+    // assign result to reporter
+    reporter->Assign(types[2], TensorType(wshape, data->dtype));
+  } else {
+    // use weight to infer the conv shape.
+    if (weight == nullptr) return false;
+    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
+
+    if (param->kernel_size.defined()) {
+      ICHECK_EQ(param->kernel_size.size(), 2);
+      // check the size
+      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
+             reporter->AssertEQ(param->kernel_size[1], wshape[3]))
+          << "DeformableConv2D: shape of weight is inconsistent with kernel_size, "
+          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
+    }
+    if (param->channels.defined()) {
+      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
+          << "DeformableConv2D: shape of weight is inconsistent with channels, "
+          << " channels=" << param->channels << " wshape=" << wshape;
+    }
+    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
+      ICHECK(reporter->AssertEQ(indexdiv(dshape_nchw[1], param->groups), wshape[1]));
+    }
+    channels = wshape[0];
+    ksize_y = wshape[2];
+    ksize_x = wshape[3];
+    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
+    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
+  }
+  // dilation
+  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
+
+  IndexExpr pad_h, pad_w;
+  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
+  oshape.Set(2, indexdiv(dshape_nchw[2] + pad_h - dilated_ksize_y, param->strides[0]) + 1);
+  oshape.Set(3, indexdiv(dshape_nchw[3] + pad_w - dilated_ksize_x, param->strides[1]) + 1);
+  DataType out_dtype = param->out_dtype;
+
+  // infer offset shape
+  Array<IndexExpr> offset_shape(
+      {dshape_nchw[0], 2 * ksize_y * ksize_x * param->deformable_groups, oshape[2], oshape[3]});
+  offset_shape = trans_in_layout.BackwardShape(offset_shape);
+  reporter->Assign(types[1], TensorType(offset_shape, data->dtype));
+  if (out_dtype.bits() == 0) {
+    out_dtype = data->dtype;
+  }
+
+  oshape = trans_out_layout.BackwardShape(oshape);
+  reporter->Assign(types[3], TensorType(oshape, out_dtype));
+  return true;
+}
+
+InferCorrectLayoutOutput DeformableConvInferCorrectLayout(
+    const Attrs& attrs, const Array<Layout>& new_in_layouts, const Array<Layout>& old_in_layouts,
+    const Array<tvm::relay::Type>& old_in_types) {
+  const auto* params = attrs.as<DeformableConv2DAttrs>();
+  return InferCorrectLayoutOutput(
+      {params->data_layout, params->data_layout, params->kernel_layout},
+      {params->out_layout == "" ? params->data_layout : params->out_layout}, attrs);
+}
+
 RELAY_REGISTER_OP("nn.deformable_conv2d")
     .describe(R"code(Compute 2-D deformable convolution on 4-D input.
 The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
@@ -563,9 +1705,8 @@ by concating all the *g* results.
     .add_argument("offset", "Tensor", "The offset tensor.")
     .add_argument("weight", "Tensor", "The weight tensor.")
     .set_support_level(5)
-    .add_type_rel("DeformableConv2D", DeformableConv2DRel<DeformableConv2DAttrs>)
-    .set_attr<FInferCorrectLayout>("FInferCorrectLayout",
-                                   DeformableConvInferCorrectLayout<DeformableConv2DAttrs>);
+    .add_type_rel("DeformableConv2D", DeformableConv2DRel)
+    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", DeformableConvInferCorrectLayout);
 
 // Positional relay function to create deformable_conv2d operator
 // used by frontend FFI.
diff --git a/src/relay/op/nn/convolution.h b/src/relay/op/nn/convolution.h
index 954daeaa86cf..62552ee4783e 100644
--- a/src/relay/op/nn/convolution.h
+++ b/src/relay/op/nn/convolution.h
@@ -37,558 +37,11 @@
 namespace tvm {
 namespace relay {
 
-// Standard convolution operator shape relations
-template <typename AttrType>
-bool Conv1DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-               const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-  static const Layout kNCW("NCW");
-  static const Layout kOIW("OIW");
-
-  const AttrType* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv only support input layouts that are convertible from NCW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv only support kernel layouts that are convertible from OIW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv only support output layouts that are convertible from NCW."
-      << " But got " << out_layout;
-
-  Array<IndexExpr> dshape_ncw = trans_in_layout.ForwardShape(data->shape);
-
-  IndexExpr channels, dilated_ksize;
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    Array<IndexExpr> wshape;
-
-    wshape = {{param->channels, indexdiv(dshape_ncw[1], param->groups), param->kernel_size[0]}};
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    channels = param->channels;
-    dilated_ksize = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
-    // assign result to reporter
-    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
-    if (param->kernel_size.defined()) {
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]))
-          << "Conv1D: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
-    }
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
-          << "Conv1D: shape of weight is inconsistent with channels, "
-          << " channels=" << param->channels << " wshape=" << wshape;
-    }
-    if (!dshape_ncw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(dshape_ncw[1], wshape[1]));
-    }
-    channels = wshape[0];
-    dilated_ksize = 1 + (wshape[2] - 1) * param->dilation[0];
-  }
-  // dilation
-  Array<IndexExpr> oshape({dshape_ncw[0], channels, 0});
-
-  if (!dshape_ncw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, indexdiv(dshape_ncw[2] + param->padding[0] + param->padding[1] - dilated_ksize,
-                           param->strides[0]) +
-                      1);
-  } else {
-    oshape.Set(2, dshape_ncw[2]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  // assign output type
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
 bool Conv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-               const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-  static const Layout kNCHW("NCHW");
-  static const Layout kOIHW("OIHW");
-
-  const AttrType* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
-  if (!trans_in_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "conv2d only support input layouts that are convertible from NCHW."
-        << " The provided layout is: " << in_layout);
-    return false;
-  }
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIHW);
-  if (!trans_kernel_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "conv2d only support kernel layouts that are convertible from OIHW."
-        << " The provided layout is: " << kernel_layout);
-    return false;
-  }
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
-  if (!trans_out_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "conv2d only support output layouts that are convertible from NCHW."
-        << "The provided layout is: " << out_layout);
-    return false;
-  }
-
-  Array<IndexExpr> dshape_nchw = trans_in_layout.ForwardShape(data->shape);
-  bool is_depthwise = false;
-  if (param->groups > 1) {
-    if (!(weight && weight->shape.defined())) {
-      reporter->GetDiagCtx().Emit(
-          Diagnostic::Error(reporter->GetSpan())
-          << "Weight shape must be specified when groups is greater than 1.");
-      return false;
-    }
-
-    Array<IndexExpr> wshape_oihw = trans_kernel_layout.ForwardShape(weight->shape);
-    if (tvm::tir::ExprDeepEqual()(param->groups, dshape_nchw[1]) &&
-        tvm::tir::ExprDeepEqual()(param->groups, wshape_oihw[0])) {
-      is_depthwise = true;
-    }
-  }
-
-  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 2);
-    ICHECK_EQ(param->dilation.size(), 2);
-    Array<IndexExpr> wshape;
-
-    if (is_depthwise) {
-      // infer weight's shape for depthwise convolution
-      wshape = {{dshape_nchw[1], indexdiv(param->channels, dshape_nchw[1]), param->kernel_size[0],
-                 param->kernel_size[1]}};
-    } else {
-      wshape = {{param->channels, indexdiv(dshape_nchw[1], param->groups), param->kernel_size[0],
-                 param->kernel_size[1]}};
-    }
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    channels = param->channels;
-    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
-
-    if (param->auto_scheduler_rewritten_layout.size() == 0) {
-      // Normal case: assign result to reporter
-      reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-    } else {
-      // If the layout is rewritten by auto-scheduler,
-      // we just forcly apply the layout provided by auto-scheduler and
-      // skip the normal inference logic.
-      {}  // do nothing
-    }
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-
-    Array<PrimExpr> wshape;
-    if (param->auto_scheduler_rewritten_layout.size() == 0) {
-      wshape = weight->shape;
-    } else {
-      // works for the default kernel layout "HWIO"
-      ICHECK_EQ(param->kernel_layout, "HWIO");
-      wshape = auto_scheduler::GetShapeFromRewrittenLayout(param->auto_scheduler_rewritten_layout,
-                                                           {"ry", "rx", "rc", "ff"});
-    }
-
-    wshape = trans_kernel_layout.ForwardShape(wshape);
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 2);
-
-      if (!reporter->AssertEQ(param->kernel_size[0], wshape[2])) {
-        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
-                                    << "Conv2D: shape of weight is inconsistent with kernel_size,"
-                                    << " kernel_size=" << param->kernel_size
-                                    << " wshape=" << wshape);
-      }
-
-      if (!reporter->AssertEQ(param->kernel_size[1], wshape[3])) {
-        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
-                                    << "Conv2D: shape of weight is inconsistent with kernel_size,"
-                                    << " kernel_size=" << param->kernel_size
-                                    << " wshape=" << wshape);
-        return false;
-      }
-    }
-
-    if (param->channels.defined() && !reporter->AssertEQ(param->channels, wshape[0])) {
-      reporter->GetDiagCtx().Emit(
-          Diagnostic::Error(reporter->GetSpan())
-          << "conv2D: the first dimensions of the weight tensor (" << wshape << ")"
-          << "does not match the number of channels (" << param->channels << ").");
-      return false;
-    }
-
-    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
-      if (!reporter->AssertEQ(indexdiv(dshape_nchw[1], param->groups), wshape[1])) {
-        reporter->GetDiagCtx().Emit(Diagnostic::Error(reporter->GetSpan())
-                                    << "conv2d: requires that `"
-                                    << indexdiv(dshape_nchw[1], param->groups) << "`,"
-                                    << " the input channels (" << dshape_nchw[1] << ")"
-                                    << " divided by groups (" << param->groups << ")"
-                                    << ",\n must match the input channels"
-                                    << " of the weight `" << wshape[1]
-                                    << "`, where the weight shape is (" << wshape << ").");
-        return false;
-      }
-    }
-    channels = wshape[0];
-    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
-  }
-  // dilation
-  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
-
-  IndexExpr pad_h, pad_w;
-  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
-  if (!dshape_nchw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, indexdiv(dshape_nchw[2] + pad_h - dilated_ksize_y, param->strides[0]) + 1);
-  } else {
-    oshape.Set(2, dshape_nchw[2]);
-  }
-
-  if (!dshape_nchw[3].as<tir::AnyNode>()) {
-    oshape.Set(3, indexdiv(dshape_nchw[3] + pad_w - dilated_ksize_x, param->strides[1]) + 1);
-  } else {
-    oshape.Set(3, dshape_nchw[3]);
-  }
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  // assign output type
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
-bool Conv3DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-               const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-  static const Layout kNCDHW("NCDHW");
-  static const Layout kOIDHW("OIDHW");
-
-  const AttrType* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv only support input layouts that are convertible from NCDHW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv only support kernel layouts that are convertible from OIDHW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv only support output layouts that are convertible from NCDHW."
-      << " But got " << out_layout;
-
-  Array<IndexExpr> dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
-
-  IndexExpr channels, dilated_ksize_z, dilated_ksize_y, dilated_ksize_x;
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 3);
-    ICHECK_EQ(param->dilation.size(), 3);
-    Array<IndexExpr> wshape;
-    tvm::tir::ExprDeepEqual expr_equal;
-
-    if (expr_equal(param->channels, param->groups) && !expr_equal(param->channels, 1)) {
-      // infer weight's shape for depthwise convolution
-      wshape = {{dshape_ncdhw[1], indexdiv(param->groups, dshape_ncdhw[1]), param->kernel_size[0],
-                 param->kernel_size[1], param->kernel_size[2]}};
-    } else {
-      wshape = {{param->channels, indexdiv(dshape_ncdhw[1], param->groups), param->kernel_size[0],
-                 param->kernel_size[1], param->kernel_size[2]}};
-    }
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    channels = param->channels;
-    dilated_ksize_z = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-    dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
+               const TypeReporter& reporter);
 
-    if (param->auto_scheduler_rewritten_layout.size() == 0) {
-      // Normal case: assign result to reporter
-      reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-    } else {
-      // If the layout is rewritten by auto-scheduler,
-      // we just forcly apply the layout provided by auto-scheduler and
-      // skip the normal inference logic.
-      {}  // do nothing
-    }
-
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-
-    Array<PrimExpr> wshape;
-    if (param->auto_scheduler_rewritten_layout.size() == 0) {
-      wshape = weight->shape;
-    } else {
-      // works for the default kernel layout "DHWIO"
-      ICHECK_EQ(param->kernel_layout, "DHWIO");
-      wshape = auto_scheduler::GetShapeFromRewrittenLayout(param->auto_scheduler_rewritten_layout,
-                                                           {"rd", "rh", "rw", "rc", "cc"});
-    }
-
-    wshape = trans_kernel_layout.ForwardShape(wshape);
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 3);
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
-             reporter->AssertEQ(param->kernel_size[1], wshape[3]) &&
-             reporter->AssertEQ(param->kernel_size[2], wshape[4]))
-          << "Conv3D: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
-    }
-
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
-          << "Conv3D: shape of weight is inconsistent with channels, "
-          << " channels=" << param->channels << " wshape=" << wshape;
-    }
-
-    if (!dshape_ncdhw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncdhw[1], param->groups), wshape[1]));
-    }
-    channels = wshape[0];
-    dilated_ksize_z = 1 + (wshape[2] - 1) * param->dilation[0];
-    dilated_ksize_y = 1 + (wshape[3] - 1) * param->dilation[1];
-    dilated_ksize_x = 1 + (wshape[4] - 1) * param->dilation[2];
-  }
-  // dilation
-  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
-
-  IndexExpr pad_d, pad_h, pad_w;
-  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
-  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, indexdiv(dshape_ncdhw[2] + pad_d - dilated_ksize_z, param->strides[0]) + 1);
-  } else {
-    oshape.Set(2, dshape_ncdhw[2]);
-  }
-
-  if (!dshape_ncdhw[3].as<tir::AnyNode>()) {
-    oshape.Set(3, indexdiv(dshape_ncdhw[3] + pad_h - dilated_ksize_y, param->strides[1]) + 1);
-  } else {
-    oshape.Set(3, dshape_ncdhw[3]);
-  }
-
-  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
-    oshape.Set(4, indexdiv(dshape_ncdhw[4] + pad_w - dilated_ksize_x, param->strides[2]) + 1);
-  } else {
-    oshape.Set(4, dshape_ncdhw[4]);
-  }
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  // assign output type
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-// Winograd convolution shape relations
-inline bool Conv2DWinogradWeightTransformRel(const Array<Type>& types, int num_inputs,
-                                             const Attrs& attrs, const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 2);
-  const auto* data = types[0].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-
-  const ConvWinogradWeightTransformAttrs* param = attrs.as<ConvWinogradWeightTransformAttrs>();
-  ICHECK(param != nullptr);
-
-  ICHECK_EQ(data->shape.size(), 4) << "Only support NCHW normal kernel layout";
-
-  std::vector<IndexExpr> oshape{
-      param->tile_size + data->shape[2] - 1,
-      param->tile_size + data->shape[3] - 1,
-      data->shape[0],
-      data->shape[1],
-  };
-
-  reporter->Assign(types[1], TensorType(Array<IndexExpr>(oshape), data->dtype));
-  return true;
-}
-
-// Gemm convolution shape relations
-// In order to run GEMM we need to block-transpose and interleave the K x N weights matrix W.
-// The high level idea is to subdivide W in tiles of tile_cols x tile_rows, and transpose and
-// interleave them. The final output is a [N//tile_rows, K//tile_cols, tile_rows, tile_cols]
-// matrix that we call W_interleaved_t.
-//
-// In the following picture, we show how the first [tile_cols,tile_rows] block of W is transformed
-// for tile_rows = 4 and tile_cols = 16
-//
-//              W[0,0,:,:]                        W_interleaved_t[0,0,:,:]
-//  +-------------------------------+     +----------------------------------- +
-//  |W[0,0]  W[0,1]  W[0,2]  W[0,3] |     |W[0,0]  W[1,0]  W[2,0]  ...  W[15,0]|
-//  |W[1,0]  W[1,1]  W[1,2]  W[1,3] | --\ |W[0,1]  W[1,1]  W[2,1]  ...  W[15,1]|
-//  |W[2,0]  W[2,1]  W[2,2]  W[2,3] | --/ |W[0,2]  W[1,2]  W[2,2]  ...  W[15,2]|
-//  |  ...     ...    ...      ...  |     |W[0,3]  W[1,3]  W[2,3]  ...  W[15,3]|
-//  |  ...     ...    ...      ...  |     +------------------------------------+
-//  |W[15,0] W[15,1] W[15,2] W[15,3]|
-//  +-------------------------------+
-//
-// Tile columns is usually the direction of the reduction. So, if our target can reduce k elements
-// at the time, we should set tile_cols = k.
-// Tile rows is connected with the number of registers available for the given target.
-//
-inline bool Conv2DGemmWeightTransformRel(const Array<Type>& types, int num_inputs,
-                                         const Attrs& attrs, const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 2);
-  const auto* weight = types[0].as<TensorTypeNode>();
-  if (weight == nullptr) return false;
-
-  const ConvGemmWeightTransformAttrs* param = attrs.as<ConvGemmWeightTransformAttrs>();
-  ICHECK(param != nullptr);
-  int n = param->tile_rows;
-  int k = param->tile_cols;
-
-  ICHECK_EQ(weight->shape.size(), 4) << "Only support HWIO kernel layout";
-
-  const auto K = weight->shape[0] * weight->shape[1] * weight->shape[2];
-  const auto N = weight->shape[3];
-
-  auto K_mod_k = indexmod(K, k);
-  auto N_mod_n = indexmod(N, n);
-
-  auto pad_K = tvm::if_then_else(K_mod_k != 0, k - K_mod_k, tir::make_zero(DataType::Int(32)));
-  auto pad_N = tvm::if_then_else(N_mod_n != 0, n - N_mod_n, tir::make_zero(DataType::Int(32)));
-
-  const auto N_padded = N + pad_N;
-  const auto K_padded = K + pad_K;
-
-  Array<IndexExpr> oshape{
-      indexdiv(N_padded, n),
-      indexdiv(K_padded, k),
-      n,
-      k,
-  };
-
-  reporter->Assign(types[1], TensorType(oshape, weight->dtype));
-  return true;
-}
-
-inline bool Conv3DWinogradWeightTransformRel(const Array<Type>& types, int num_inputs,
-                                             const Attrs& attrs, const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 2);
-  const auto* data = types[0].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-
-  const ConvWinogradWeightTransformAttrs* param = attrs.as<ConvWinogradWeightTransformAttrs>();
-  ICHECK(param != nullptr);
-
-  ICHECK_EQ(data->shape.size(), 5) << "Only support NCDHW normal kernel layout";
-
-  // Shape of packed weights depends on whether depth is being transformed or not.
-  Array<IndexExpr> oshape({0, 0, 0, data->shape[0], data->shape[1]});
-  auto* depth_imm = data->shape[2].as<IntImmNode>();
-  bool transform_depth = (depth_imm->value > 2) && (depth_imm->value < 8);
-  if (transform_depth) {
-    oshape.Set(0, param->tile_size + data->shape[2] - 1);
-    oshape.Set(1, param->tile_size + data->shape[3] - 1);
-    oshape.Set(2, param->tile_size + data->shape[4] - 1);
-  } else {
-    oshape.Set(0, param->tile_size + data->shape[3] - 1);
-    oshape.Set(1, param->tile_size + data->shape[4] - 1);
-    oshape.Set(2, data->shape[2]);
-  }
-
-  reporter->Assign(types[1], TensorType(oshape, data->dtype));
-  return true;
-}
-
-inline bool Conv2DWinogradNNPACKWeightTransformRel(const Array<Type>& types, int num_inputs,
-                                                   const Attrs& attrs,
-                                                   const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 2);
-  const auto* data = types[0].as<TensorTypeNode>();
-  if (data == nullptr) {
-    return false;
-  }
-
-  const Conv2DWinogradNNPACKWeightTransformAttrs* param =
-      attrs.as<Conv2DWinogradNNPACKWeightTransformAttrs>();
-  ICHECK(param != nullptr);
-
-  ICHECK_EQ(data->shape.size(), 4) << "Only support NCHW normal kernel layout";
-
-  std::vector<IndexExpr> oshape{
-      data->shape[0],
-      data->shape[1],
-      8,
-      8,
-  };
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  reporter->Assign(types[1], TensorType(Array<IndexExpr>(oshape), out_dtype));
-  return true;
-}
+bool Conv2DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+                        const TypeReporter& reporter);
 
 template <typename AttrType>
 bool Conv2DWinogradRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
@@ -667,608 +120,6 @@ bool Conv2DWinogradRel(const Array<Type>& types, int num_inputs, const Attrs& at
   return true;
 }
 
-template <typename AttrType>
-bool Conv2DGemmRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                   const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-  static const Layout kNHWC("NHWC");
-  static const Layout kHWIO("HWIO");
-
-  const AttrType* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNHWC);
-  ICHECK(trans_in_layout.defined())
-      << "Conv only support input layouts that are convertible from NHWC."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kHWIO);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv only support kernel layouts that are convertible from HWIO."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNHWC);
-  ICHECK(trans_out_layout.defined())
-      << "Conv only support output layouts that are convertible from NHWC."
-      << " But got " << out_layout;
-
-  Array<IndexExpr> dshape_nhwc = trans_in_layout.ForwardShape(data->shape);
-
-  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
-
-  ICHECK(param->kernel_size.defined() && param->channels.defined())
-      << "The kernel size and channels of a Conv must be set or inferred by previous pass";
-
-  ICHECK_EQ(param->kernel_size.size(), 2);
-  ICHECK_EQ(param->dilation.size(), 2);
-
-  channels = param->channels;
-  dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-  dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-
-  // NOTE: Do not check weight shape here!
-
-  // dilation
-  Array<IndexExpr> oshape({dshape_nhwc[0], 0, 0, channels});
-
-  IndexExpr pad_h, pad_w;
-  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
-  if (!dshape_nhwc[2].as<tir::AnyNode>()) {
-    oshape.Set(1, (dshape_nhwc[1] + pad_h - dilated_ksize_y) / param->strides[0] + 1);
-  } else {
-    oshape.Set(1, dshape_nhwc[1]);
-  }
-  if (!dshape_nhwc[3].as<tir::AnyNode>()) {
-    oshape.Set(2, (dshape_nhwc[2] + pad_w - dilated_ksize_x) / param->strides[1] + 1);
-  } else {
-    oshape.Set(2, dshape_nhwc[2]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  // assign output type
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
-bool Conv3DWinogradRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                       const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-  static const Layout kNCDHW("NCDHW");
-  static const Layout kOIDHW("OIDHW");
-
-  const AttrType* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv only support input layouts that are convertible from NCDHW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv only support kernel layouts that are convertible from OIDHW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv only support output layouts that are convertible from NCDHW."
-      << " But got " << out_layout;
-
-  Array<IndexExpr> dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
-
-  IndexExpr channels, dilated_ksize_d, dilated_ksize_y, dilated_ksize_x;
-
-  ICHECK(param->kernel_size.defined() && param->channels.defined())
-      << "The kernel size and channels of a Conv must be set or inferred by previous pass";
-
-  ICHECK_EQ(param->kernel_size.size(), 3);
-  ICHECK_EQ(param->dilation.size(), 3);
-
-  channels = param->channels;
-  dilated_ksize_d = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-  dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-  dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
-
-  // NOTE: Do not check weight shape here!
-  // Different backend requires different layout to compute
-  // the batch gemm stage in winograd efficiently, but we want to
-  // make this op work for all backends.
-  // So we accept all weight shapes, and assume the TOPI developers
-  // can handle this correctly in alter_op_layout.
-
-  // dilation
-  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
-
-  IndexExpr pad_d, pad_h, pad_w;
-  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
-  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, (dshape_ncdhw[2] + pad_d - dilated_ksize_d) / param->strides[0] + 1);
-  } else {
-    oshape.Set(2, dshape_ncdhw[2]);
-  }
-  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
-    oshape.Set(3, (dshape_ncdhw[3] + pad_h - dilated_ksize_y) / param->strides[1] + 1);
-  } else {
-    oshape.Set(3, dshape_ncdhw[3]);
-  }
-  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
-    oshape.Set(4, (dshape_ncdhw[4] + pad_w - dilated_ksize_x) / param->strides[2] + 1);
-  } else {
-    oshape.Set(4, dshape_ncdhw[4]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  // assign output type
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-// Transposed convolution shape relations
-template <typename AttrType>
-bool Conv1DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                        const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-
-  static const Layout kNCW("NCW");
-  static const Layout kOIW("OIW");
-
-  const Conv1DTransposeAttrs* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv only support input layouts that are convertible from NCW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv only support kernel layouts that are convertible from OIW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv only support output layouts that are convertible from NCW."
-      << " But got " << out_layout;
-
-  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
-
-  auto dshape_ncw = trans_in_layout.ForwardShape(data->shape);
-
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 1);
-    ICHECK_EQ(param->dilation.size(), 1);
-
-    Array<IndexExpr> wshape(
-        {dshape_ncw[1], indexdiv(param->channels, param->groups), param->kernel_size[0]});
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    dilated_ksize_x = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    channels = param->channels;
-
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
-    // assign result to reporter
-    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 1);
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]))
-          << "Conv1D: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
-    }
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[1]))
-          << "Conv1D: shape of weight is inconsistent with channels, "
-          << " channels=" << param->channels << " wshape=" << Array<IndexExpr>(wshape);
-    }
-    if (!dshape_ncw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncw[1], param->groups), wshape[0]));
-    }
-    channels = wshape[1];
-    dilated_ksize_x = 1 + (wshape[2] - 1) * param->dilation[0];
-  }
-  // dilation
-  IndexExpr pad_w;
-  GetPaddingWidth(param->padding, &pad_w);
-  Array<IndexExpr> oshape({dshape_ncw[0], channels, 0});
-  if (!dshape_ncw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, (param->strides[0] * (dshape_ncw[2] - 1) + dilated_ksize_x - pad_w +
-                   param->output_padding[0]));
-  } else {
-    oshape.Set(2, dshape_ncw[2]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
-bool Conv3DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                        const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-
-  static const Layout kNCDHW("NCDHW");
-  static const Layout kOIDHW("OIDHW");
-
-  const Conv3DTransposeAttrs* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCDHW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv3d_transpose only support input layouts that are convertible from NCDHW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIDHW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv3d_transpose only support kernel layouts that are convertible from OIDHW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCDHW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv3d_transpose only support output layouts that are convertible from NCDHW."
-      << " But got " << out_layout;
-
-  IndexExpr channels, dilated_ksize_d, dilated_ksize_y, dilated_ksize_x;
-
-  auto dshape_ncdhw = trans_in_layout.ForwardShape(data->shape);
-
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 3);
-    ICHECK_EQ(param->dilation.size(), 3);
-
-    Array<IndexExpr> wshape({dshape_ncdhw[1], indexdiv(param->channels, param->groups),
-                             param->kernel_size[0], param->kernel_size[1], param->kernel_size[2]});
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    dilated_ksize_d = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    dilated_ksize_y = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-    dilated_ksize_x = 1 + (param->kernel_size[2] - 1) * param->dilation[2];
-    channels = param->channels;
-
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
-    // assign result to reporter
-    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 3);
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
-             reporter->AssertEQ(param->kernel_size[1], wshape[3]) &&
-             reporter->AssertEQ(param->kernel_size[2], wshape[4]))
-          << "Conv3D: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
-    }
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[1]))
-          << "Conv3D: shape of weight is inconsistent with channels, "
-          << " channels=" << param->channels << " wshape=" << Array<IndexExpr>(wshape);
-    }
-    if (!dshape_ncdhw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(indexdiv(dshape_ncdhw[1], param->groups), wshape[0]));
-    }
-    channels = wshape[1];
-    dilated_ksize_d = 1 + (wshape[2] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
-    dilated_ksize_y = 1 + (wshape[4] - 1) * param->dilation[2];
-  }
-
-  // dilation
-  Array<IndexExpr> oshape({dshape_ncdhw[0], channels, 0, 0, 0});
-  IndexExpr pad_d, pad_h, pad_w;
-  GetPaddingDepthHeightWidth(param->padding, &pad_d, &pad_h, &pad_w);
-
-  if (!dshape_ncdhw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, (param->strides[0] * (dshape_ncdhw[2] - 1) + dilated_ksize_d - pad_d +
-                   param->output_padding[0]));
-  } else {
-    oshape.Set(2, dshape_ncdhw[2]);
-  }
-  if (!dshape_ncdhw[3].as<tir::AnyNode>()) {
-    oshape.Set(3, (param->strides[1] * (dshape_ncdhw[3] - 1) + dilated_ksize_y - pad_h +
-                   param->output_padding[1]));
-  } else {
-    oshape.Set(3, dshape_ncdhw[3]);
-  }
-  if (!dshape_ncdhw[4].as<tir::AnyNode>()) {
-    oshape.Set(4, (param->strides[2] * (dshape_ncdhw[4] - 1) + dilated_ksize_x - pad_w +
-                   param->output_padding[2]));
-  } else {
-    oshape.Set(4, dshape_ncdhw[4]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
-bool Conv2DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                        const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 3);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[1].as<TensorTypeNode>();
-  if (data == nullptr) return false;
-
-  static const Layout kNCHW("NCHW");
-  static const Layout kIOHW("IOHW");
-
-  const Conv2DTransposeAttrs* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
-  ICHECK(trans_in_layout.defined())
-      << "Conv2DTransposed only support input layouts that are convertible from NCHW."
-      << " But got " << in_layout;
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kIOHW);
-  ICHECK(trans_kernel_layout.defined())
-      << "Conv2DTransposed only support kernel layouts that are convertible from IOHW."
-      << " But got " << kernel_layout;
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
-  ICHECK(trans_out_layout.defined())
-      << "Conv2DTransposed only support output layouts that are convertible from NCHW."
-      << " But got " << out_layout;
-
-  IndexExpr channels, dilated_ksize_y, dilated_ksize_x;
-
-  auto dshape_nchw = trans_in_layout.ForwardShape(data->shape);
-
-  // infer weight if the kernel_size and channels are defined
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 2);
-    ICHECK_EQ(param->dilation.size(), 2);
-
-    Array<IndexExpr> wshape({dshape_nchw[1], indexdiv(param->channels, param->groups),
-                             param->kernel_size[0], param->kernel_size[1]});
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-    channels = param->channels;
-
-    DataType weight_dtype = data->dtype;
-    if (weight != nullptr) {
-      weight_dtype = weight->dtype;
-    }
-    // assign result to reporter
-    reporter->Assign(types[1], TensorType(wshape, weight_dtype));
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 2);
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
-             reporter->AssertEQ(param->kernel_size[1], wshape[3]))
-          << "Conv2DTransposed: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << Array<IndexExpr>(wshape);
-    }
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(indexdiv(param->channels, param->groups), wshape[1]))
-          << "Conv2DTransposed: shape of weight is inconsistent with out_channels, "
-          << " out_channels // groups != weight.shape[1] "
-          << " out_channels=" << param->channels << " groups=" << param->groups
-          << " weight.shape=" << Array<IndexExpr>(wshape);
-    }
-    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[0].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(dshape_nchw[1], wshape[0]))
-          << "Conv2DTransposed: shape of weight is inconsistent with in_channels."
-          << " data.shape= " << Array<IndexExpr>(dshape_nchw) << " groups= " << param->groups
-          << " weight.shape= " << Array<IndexExpr>(wshape);
-    }
-    channels = wshape[1];
-    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
-  }
-  // dilation
-  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
-  IndexExpr pad_h, pad_w;
-  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
-  if (!dshape_nchw[2].as<tir::AnyNode>()) {
-    oshape.Set(2, (param->strides[0] * (dshape_nchw[2] - 1) + dilated_ksize_y - pad_h +
-                   param->output_padding[0]));
-  } else {
-    oshape.Set(2, dshape_nchw[2]);
-  }
-  if (!dshape_nchw[3].as<tir::AnyNode>()) {
-    oshape.Set(3, (param->strides[1] * (dshape_nchw[3] - 1) + dilated_ksize_x - pad_w +
-                   param->output_padding[1]));
-  } else {
-    oshape.Set(3, dshape_nchw[3]);
-  }
-
-  DataType out_dtype = param->out_dtype;
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-  oshape = trans_out_layout.BackwardShape(oshape);
-  reporter->Assign(types[2], TensorType(oshape, out_dtype));
-  return true;
-}
-
-// Deformable Convolution shape relations.
-template <typename AttrType>
-bool DeformableConv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                         const TypeReporter& reporter) {
-  ICHECK_EQ(types.size(), 4);
-  const auto* data = types[0].as<TensorTypeNode>();
-  const auto* weight = types[2].as<TensorTypeNode>();
-
-  ICHECK(data);
-  static const Layout kNCHW("NCHW");
-  static const Layout kOIHW("OIHW");
-
-  auto* param = attrs.as<AttrType>();
-  ICHECK(param != nullptr);
-  const Layout in_layout(param->data_layout);
-  const Layout kernel_layout(param->kernel_layout);
-
-  const auto trans_in_layout = tir::BijectiveLayout(in_layout, kNCHW);
-  if (!trans_in_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "deformable_conv2d only support input layouts that are convertible from NCHW."
-        << " The provided layout is: " << in_layout);
-    return false;
-  }
-
-  const auto trans_kernel_layout = tir::BijectiveLayout(kernel_layout, kOIHW);
-  if (!trans_kernel_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "deformable_conv2d only support kernel layouts that are convertible from OIHW."
-        << " The provided layout is: " << kernel_layout);
-    return false;
-  }
-
-  Layout out_layout(param->out_layout == "" ? param->data_layout : param->out_layout);
-  const auto trans_out_layout = tir::BijectiveLayout(out_layout, kNCHW);
-  if (!trans_out_layout.defined()) {
-    reporter->GetDiagCtx().Emit(
-        Diagnostic::Error(reporter->GetSpan())
-        << "deformable_conv2d only support output layouts that are convertible from NCHW."
-        << "The provided layout is: " << out_layout);
-    return false;
-  }
-
-  Array<IndexExpr> dshape_nchw = trans_in_layout.ForwardShape(data->shape);
-
-  IndexExpr channels, dilated_ksize_y, dilated_ksize_x, ksize_y, ksize_x;
-
-  // infer weight shape if kernel_size and channels are defiend
-  if (param->kernel_size.defined() && param->channels.defined()) {
-    ICHECK_EQ(param->kernel_size.size(), 2);
-    ICHECK_EQ(param->dilation.size(), 2);
-    Array<IndexExpr> wshape({param->channels, indexdiv(dshape_nchw[1], param->groups),
-                             param->kernel_size[0], param->kernel_size[1]});
-
-    wshape = trans_kernel_layout.BackwardShape(wshape);
-    channels = param->channels;
-    ksize_y = param->kernel_size[0];
-    ksize_x = param->kernel_size[1];
-    dilated_ksize_y = 1 + (param->kernel_size[0] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (param->kernel_size[1] - 1) * param->dilation[1];
-    // assign result to reporter
-    reporter->Assign(types[2], TensorType(wshape, data->dtype));
-  } else {
-    // use weight to infer the conv shape.
-    if (weight == nullptr) return false;
-    auto wshape = trans_kernel_layout.ForwardShape(weight->shape);
-
-    if (param->kernel_size.defined()) {
-      ICHECK_EQ(param->kernel_size.size(), 2);
-      // check the size
-      ICHECK(reporter->AssertEQ(param->kernel_size[0], wshape[2]) &&
-             reporter->AssertEQ(param->kernel_size[1], wshape[3]))
-          << "DeformableConv2D: shape of weight is inconsistent with kernel_size, "
-          << " kernel_size=" << param->kernel_size << " wshape=" << wshape;
-    }
-    if (param->channels.defined()) {
-      ICHECK(reporter->AssertEQ(param->channels, wshape[0]))
-          << "DeformableConv2D: shape of weight is inconsistent with channels, "
-          << " channels=" << param->channels << " wshape=" << wshape;
-    }
-    if (!dshape_nchw[1].as<tir::AnyNode>() && !wshape[1].as<tir::AnyNode>()) {
-      ICHECK(reporter->AssertEQ(indexdiv(dshape_nchw[1], param->groups), wshape[1]));
-    }
-    channels = wshape[0];
-    ksize_y = wshape[2];
-    ksize_x = wshape[3];
-    dilated_ksize_y = 1 + (wshape[2] - 1) * param->dilation[0];
-    dilated_ksize_x = 1 + (wshape[3] - 1) * param->dilation[1];
-  }
-  // dilation
-  Array<IndexExpr> oshape({dshape_nchw[0], channels, 0, 0});
-
-  IndexExpr pad_h, pad_w;
-  GetPaddingHeightWidth(param->padding, &pad_h, &pad_w);
-  oshape.Set(2, indexdiv(dshape_nchw[2] + pad_h - dilated_ksize_y, param->strides[0]) + 1);
-  oshape.Set(3, indexdiv(dshape_nchw[3] + pad_w - dilated_ksize_x, param->strides[1]) + 1);
-  DataType out_dtype = param->out_dtype;
-
-  // infer offset shape
-  Array<IndexExpr> offset_shape(
-      {dshape_nchw[0], 2 * ksize_y * ksize_x * param->deformable_groups, oshape[2], oshape[3]});
-  offset_shape = trans_in_layout.BackwardShape(offset_shape);
-  reporter->Assign(types[1], TensorType(offset_shape, data->dtype));
-  if (out_dtype.bits() == 0) {
-    out_dtype = data->dtype;
-  }
-
-  oshape = trans_out_layout.BackwardShape(oshape);
-  reporter->Assign(types[3], TensorType(oshape, out_dtype));
-  return true;
-}
-
-template <typename AttrType>
-InferCorrectLayoutOutput DeformableConvInferCorrectLayout(
-    const Attrs& attrs, const Array<Layout>& new_in_layouts, const Array<Layout>& old_in_layouts,
-    const Array<tvm::relay::Type>& old_in_types) {
-  const AttrType* params = attrs.as<AttrType>();
-  return InferCorrectLayoutOutput(
-      {params->data_layout, params->data_layout, params->kernel_layout},
-      {params->out_layout == "" ? params->data_layout : params->out_layout}, attrs);
-}
-
 template <typename T>
 InferCorrectLayoutOutput ConvInferCorrectLayout(const Attrs& attrs,
                                                 const Array<Layout>& new_in_layouts,
diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc
index e0e05436a8e6..27f295b8b39d 100644
--- a/src/relay/op/nn/nn.cc
+++ b/src/relay/op/nn/nn.cc
@@ -259,7 +259,7 @@ bool DensePackRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   ICHECK(param != nullptr);
 
   ICHECK_EQ(data->shape.size(), 2) << "Only 2D data is supported";
-  ICHECK_EQ(weight->shape.size(), 3) << "Weight is not packed";
+  ICHECK(weight->shape.size() == 3 || weight->shape.size() == 4) << "Expect weight to be 3D or 4D";
 
   Array<tvm::PrimExpr> oshape = data->shape;
   oshape.Set(1, weight->shape[0] * weight->shape[2]);
@@ -1165,7 +1165,8 @@ bool NLLLossRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
                                      << ", weights shape = " << weights->shape);
     return false;
   }
-  if (!(predictions->dtype == weights->dtype && predictions->dtype.is_float())) {
+  if (!(predictions->dtype == weights->dtype &&
+        (predictions->dtype.is_float() || predictions->dtype.is_bfloat16()))) {
     reporter->GetDiagCtx().EmitFatal(Diagnostic::Error(reporter->GetSpan())
                                      << "NLLLossRel: predictions and weights should"
                                      << " be of the same floating type.");
diff --git a/src/relay/op/type_relations.cc b/src/relay/op/type_relations.cc
index 6e30ad9624c4..e8b58f414a43 100644
--- a/src/relay/op/type_relations.cc
+++ b/src/relay/op/type_relations.cc
@@ -107,7 +107,7 @@ bool BroadcastRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
       if (t0->dtype != t1->dtype) {
         reporter->GetDiagCtx().Emit(Diagnostic::Error(t0->span)
                                     << "data types " << t0->dtype << " and " << t1->dtype
-                                    << "do not match in BroadcastRel");
+                                    << " do not match in BroadcastRel");
       }
       reporter->Assign(
           types[2], ConcreteBroadcast(GetRef<TensorType>(t0), GetRef<TensorType>(t1), t0->dtype));
@@ -127,7 +127,7 @@ bool BroadcastCompRel(const Array<Type>& types, int num_inputs, const Attrs& att
       if (t0->dtype != t1->dtype) {
         reporter->GetDiagCtx().Emit(Diagnostic::Error(t0->span)
                                     << "data types " << t0->dtype << " and " << t1->dtype
-                                    << "do not match in BroadcastCompRel");
+                                    << " do not match in BroadcastCompRel");
       }
       reporter->Assign(types[2], ConcreteBroadcast(GetRef<TensorType>(t0), GetRef<TensorType>(t1),
                                                    DataType::Bool()));
diff --git a/src/relay/qnn/op/convolution.cc b/src/relay/qnn/op/convolution.cc
index 03fa770e404f..8a7521e8ee50 100644
--- a/src/relay/qnn/op/convolution.cc
+++ b/src/relay/qnn/op/convolution.cc
@@ -84,7 +84,7 @@ bool QnnConv2DRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
   // Collect the input tensor and output tensor devoid of scale and zero points to reuse Relay
   // Conv2D infer type function.
   Array<Type> tensor_types = {types[0], types[1], types[6]};
-  return Conv2DRel<Conv2DAttrs>(tensor_types, 3, attrs, reporter);
+  return Conv2DRel(tensor_types, 3, attrs, reporter);
 }
 
 InferCorrectLayoutOutput QnnConvInferCorrectLayout(const Attrs& attrs,
diff --git a/src/relay/qnn/op/convolution_transpose.cc b/src/relay/qnn/op/convolution_transpose.cc
index b9227ff96c2a..cdc6c8d98f3a 100644
--- a/src/relay/qnn/op/convolution_transpose.cc
+++ b/src/relay/qnn/op/convolution_transpose.cc
@@ -128,7 +128,7 @@ bool QnnConv2DTransposeRel(const Array<Type>& types, int num_inputs, const Attrs
   // Collect the input tensor and output tensor devoid of scale and zero points to reuse Relay
   // Conv2D infer type function.
   Array<Type> tensor_types = {types[0], types[1], types[6]};
-  return Conv2DTransposeRel<Conv2DTransposeAttrs>(tensor_types, 3, attrs, reporter);
+  return Conv2DTransposeRel(tensor_types, 3, attrs, reporter);
 }
 
 RELAY_REGISTER_OP("qnn.conv2d_transpose")
diff --git a/src/relay/qnn/op/op_common.h b/src/relay/qnn/op/op_common.h
index 4df958603ff2..de97fb860b8a 100644
--- a/src/relay/qnn/op/op_common.h
+++ b/src/relay/qnn/op/op_common.h
@@ -28,6 +28,7 @@
 #include <tvm/relay/op.h>
 #include <tvm/relay/op_attr_types.h>
 #include <tvm/relay/qnn/attrs.h>
+#include <tvm/relay/qnn/transform.h>
 
 #include <vector>
 
@@ -289,6 +290,98 @@ static inline bool QnnBroadcastRel(const Array<Type>& types, int num_inputs, con
       .set_attr<TNonComputational>("TNonComputational", true)                                      \
       .set_attr<FInferCorrectLayout>("FInferCorrectLayout", QnnBinaryBroadcastLayout)
 
+static inline bool QnnElementwiseUnaryFuncRel(const Array<Type>& types, int num_inputs,
+                                              const Attrs& attrs, const TypeReporter& reporter) {
+  // Expected Types: data, scale, zero_point, output_scale, output_zero_point
+  ICHECK_EQ(types.size(), 6);
+  const auto* x = types[0].as<TensorTypeNode>();
+  if (x == nullptr) return false;
+  ICHECK(x->dtype == DataType::Int(8) || x->dtype == DataType::UInt(8))
+      << "Expected quantized type(int8, uint8) for input but was " << x->dtype;
+
+  // Check the types of scale and zero points.
+  for (size_t i = 1; i < 5; ++i) {
+    if (types[i].as<IncompleteTypeNode>()) {
+      return false;
+    }
+  }
+  ICHECK(IsScalarType(types[1], DataType::Float(32)));  // scale
+  ICHECK(IsScalarType(types[2], DataType::Int(32)));    // zero_point
+  ICHECK(IsScalarType(types[3], DataType::Float(32)));  // output_scale
+  ICHECK(IsScalarType(types[4], DataType::Int(32)));    // output_zero_point
+
+  // Assign types for scale and zero points.
+  reporter->Assign(types[1], TensorType({}, DataType::Float(32)));  // scale
+  reporter->Assign(types[2], TensorType({}, DataType::Int(32)));    // zero_point
+  reporter->Assign(types[3], TensorType({}, DataType::Float(32)));  // output_scale
+  reporter->Assign(types[4], TensorType({}, DataType::Int(32)));    // output_zero_point
+
+  // Collect the input tensor and output tensor devoid of scale and zero points to reuse Relay
+  // IdentityRel infer type function.
+  Array<Type> tensor_types = {types[0], types[5]};
+  return IdentityRel(tensor_types, 2, attrs, reporter);
+}
+
+static inline Expr LegalizeExpr(const Expr& expr) {
+  // Canonicalizations should not contain qnn ops, so use this
+  // to lower expressions automatically after using things like qnn.dequantize
+  // in the lowering process.
+  auto mod = IRModule::FromExpr(expr);
+  mod = transform::Legalize()(mod);
+  if (expr.as<FunctionNode>()) {
+    return mod->Lookup("main");
+  } else {
+    return mod->Lookup("main").as<FunctionNode>()->body;
+  }
+}
+
+/*! Quick helper macro
+ * - Expose a positional make function to construct the node.
+ * - Register op to the registry.
+ *
+ * For Unary Operators which also take in QParams.
+ *
+ * \param OpName the name of registry.
+ */
+#define QNN_CREATE_UNARY_ELEMENTWISE_OP(OpName)                                                 \
+  TVM_REGISTER_GLOBAL("relay.qnn.op._make." OpName)                                             \
+      .set_body_typed(                                                                          \
+          [](Expr x, Expr scale, Expr zero_point, Expr output_scale, Expr output_zero_point) {  \
+            return Call(Op::Get("qnn." OpName),                                                 \
+                        {x, scale, zero_point, output_scale, output_zero_point}, Attrs(), {});  \
+          });                                                                                   \
+                                                                                                \
+  RELAY_REGISTER_OP("qnn." OpName)                                                              \
+      .describe("Elementwise " OpName " for quantized tensors.")                                \
+      .set_num_inputs(5)                                                                        \
+      .add_argument("data", "Quantized Tensor", "The input data.")                              \
+      .add_argument("scale", "Tensor", "The quantization scale of the input tensor.")           \
+      .add_argument("zero_point", "Tensor", "The quantization zero_point of the input tensor.") \
+      .add_argument("output_scale", "Tensor", "The quantization scale of the output tensor.")   \
+      .add_argument("output_zero_point", "Tensor",                                              \
+                    "The quantization zero_point of the output tensor.")                        \
+      .set_support_level(11)                                                                    \
+      .add_type_rel("qnn." OpName, QnnElementwiseUnaryFuncRel)                                  \
+      .set_attr<TNonComputational>("TNonComputational", true)
+
+/*! Quick helper macro
+ * Create a default canonicalization for a QNN operator, which dequantizes the operator
+ * runs the calculation using the provided Call func, and then requantizes.
+ *
+ * FloatingPointFunc is usually a handle from "src/relay/transforms/pattern_utils.h"
+ *
+ * \param FloatingPointFunc the floating point function with function signature `Expr Erf(Expr e)`
+ */
+#define QNN_UNARY_OP_DEFAULT_CANONICALIZATION(FloatingPointFunc)                                  \
+  [](const Attrs& attrs, const Array<Expr>& new_args, const Array<tvm::relay::Type>& arg_types) { \
+    QnnUnaryOpArguments args(new_args);                                                           \
+    QnnUnaryOpTensorType input_type(arg_types, 0);                                                \
+    Expr dequantized_arg = MakeDequantize(args.x, args.scale, args.zero_point, -1);               \
+    Expr output = FloatingPointFunc(dequantized_arg);                                             \
+    Expr result =                                                                                 \
+        MakeQuantize(output, args.output_scale, args.output_zero_point, -1, input_type.dtype);    \
+    return LegalizeExpr(result);                                                                  \
+  }
 }  // namespace qnn
 }  // namespace relay
 }  // namespace tvm
diff --git a/src/relay/qnn/op/rsqrt.cc b/src/relay/qnn/op/rsqrt.cc
deleted file mode 100644
index 93baa308a796..000000000000
--- a/src/relay/qnn/op/rsqrt.cc
+++ /dev/null
@@ -1,92 +0,0 @@
-/*
- * 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 src/relay/qnn/op/rsqrt.cc
- * \brief QNN rsqrt operator.
- */
-#include <tvm/relay/analysis.h>
-#include <tvm/relay/op_attr_types.h>
-
-#include "op_common.h"
-
-namespace tvm {
-namespace relay {
-namespace qnn {
-
-bool QnnRsqrtRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
-                 const TypeReporter& reporter) {
-  // Expected Types: data, scale, zero_point, output_scale, output_zero_point
-  ICHECK_EQ(types.size(), 6);
-  const auto* x = types[0].as<TensorTypeNode>();
-  if (x == nullptr) return false;
-  ICHECK(x->dtype == DataType::Int(8) || x->dtype == DataType::UInt(8))
-      << "Expected quantized rsqrt type(int8, uint8) for input but was " << x->dtype;
-
-  // Check the types of scale and zero points.
-  for (size_t i = 1; i < 5; ++i) {
-    if (types[i].as<IncompleteTypeNode>()) {
-      return false;
-    }
-  }
-  ICHECK(IsScalarType(types[1], DataType::Float(32)));  // scale
-  ICHECK(IsScalarType(types[2], DataType::Int(32)));    // zero_point
-  ICHECK(IsScalarType(types[3], DataType::Float(32)));  // output_scale
-  ICHECK(IsScalarType(types[4], DataType::Int(32)));    // output_zero_point
-
-  // Assign types for scale and zero points.
-  reporter->Assign(types[1], TensorType({}, DataType::Float(32)));  // scale
-  reporter->Assign(types[2], TensorType({}, DataType::Int(32)));    // zero_point
-  reporter->Assign(types[3], TensorType({}, DataType::Float(32)));  // output_scale
-  reporter->Assign(types[4], TensorType({}, DataType::Int(32)));    // output_zero_point
-
-  // Collect the input tensor and output tensor devoid of scale and zero points to reuse Relay
-  // IdentityRel infer type function.
-  Array<Type> tensor_types = {types[0], types[5]};
-  return IdentityRel(tensor_types, 2, attrs, reporter);
-}
-
-// Positional relay function to create quantized rsqrt operator used by frontend FFI.
-Expr MakeQuantizedRsqrt(Expr x, Expr scale, Expr zero_point, Expr output_scale,
-                        Expr output_zero_point) {
-  static const Op& op = Op::Get("qnn.rsqrt");
-  return Call(op, {x, scale, zero_point, output_scale, output_zero_point}, Attrs(), {});
-}
-
-// Translation to relay is done via canonicalization/legalization functions in python
-// e.g. python/tvm/relay/qnn/op/canonicalizations.py or
-//      python/tvm/relay/qnn/op/legalizations.py
-RELAY_REGISTER_OP("qnn.rsqrt")
-    .describe("Elementwise rsqrt for quantized tensors.")
-    .set_num_inputs(5)
-    .add_argument("data", "Quantized Tensor", "The input data.")
-    .add_argument("scale", "Tensor", "The quantization scale of the input tensor.")
-    .add_argument("zero_point", "Tensor", "The quantization zero_point of the input tensor.")
-    .add_argument("output_scale", "Tensor", "The quantization scale of the output tensor.")
-    .add_argument("output_zero_point", "Tensor",
-                  "The quantization zero_point of the output tensor.")
-    .set_support_level(11)
-    .add_type_rel("QRsqrt", QnnRsqrtRel)
-    .set_attr<TNonComputational>("TNonComputational", true);
-
-TVM_REGISTER_GLOBAL("relay.qnn.op._make.rsqrt").set_body_typed(MakeQuantizedRsqrt);
-
-}  // namespace qnn
-}  // namespace relay
-}  // namespace tvm
diff --git a/src/relay/qnn/op/unary_elementwise_op.cc b/src/relay/qnn/op/unary_elementwise_op.cc
new file mode 100644
index 000000000000..ff259d975230
--- /dev/null
+++ b/src/relay/qnn/op/unary_elementwise_op.cc
@@ -0,0 +1,49 @@
+/*
+ * 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 src/relay/qnn/op/unary_elementwise_op.cc
+ * \brief QNN unary elementwise operators.
+ */
+
+#include "op_common.h"
+
+namespace tvm {
+namespace relay {
+namespace qnn {
+
+QNN_CREATE_UNARY_ELEMENTWISE_OP("tanh").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Tanh));
+
+QNN_CREATE_UNARY_ELEMENTWISE_OP("exp").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Exp));
+
+QNN_CREATE_UNARY_ELEMENTWISE_OP("sqrt").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Sqrt));
+QNN_CREATE_UNARY_ELEMENTWISE_OP("rsqrt").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Rsqrt));
+
+QNN_CREATE_UNARY_ELEMENTWISE_OP("erf").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Erf));
+QNN_CREATE_UNARY_ELEMENTWISE_OP("sigmoid").set_attr<FTVMLegalize>(
+    "FTVMQnnCanonicalize", QNN_UNARY_OP_DEFAULT_CANONICALIZATION(Sigmoid));
+
+}  // namespace qnn
+}  // namespace relay
+}  // namespace tvm
diff --git a/src/relay/qnn/utils.h b/src/relay/qnn/utils.h
index 0f3645a9882a..d7769707f01e 100644
--- a/src/relay/qnn/utils.h
+++ b/src/relay/qnn/utils.h
@@ -131,6 +131,7 @@ static inline Expr Dequantize(const Expr& data, const Expr& input_scale,
 
   return DequantizeLower(data, input_scale, input_zero_point, types, attrs.operator->());
 }
+Expr MakeDequantize(Expr data, Expr input_scale, Expr input_zero_point, int axis);
 
 Expr QuantizeLower(const Expr& input_tensor, const Expr& output_scale,
                    const Expr& output_zero_point, const Array<tvm::relay::Type>& types,
@@ -145,6 +146,8 @@ static inline Expr Quantize(const Expr& data, const Expr& output_scale,
 
   return QuantizeLower(data, output_scale, output_zero_point, types, attrs.operator->());
 }
+Expr MakeQuantize(Expr data, Expr output_scale, Expr output_zero_point, int axis,
+                  DataType out_dtype);
 
 static inline int64_t get_const_int(const tvm::PrimExpr& x) {
   auto* value_ptr = tir::as_const_int(x);
diff --git a/src/relay/transforms/device_aware_visitors.cc b/src/relay/transforms/device_aware_visitors.cc
index 10584da51976..b5ad64add89a 100644
--- a/src/relay/transforms/device_aware_visitors.cc
+++ b/src/relay/transforms/device_aware_visitors.cc
@@ -38,7 +38,7 @@ LexicalOnDeviceMixin::LexicalOnDeviceMixin(const Optional<IRModule>& maybe_mod)
   if (maybe_mod) {
     for (const auto& kv : maybe_mod.value()->functions) {
       if (const auto* function_node = kv.second.as<FunctionNode>()) {
-        VirtualDevice virtual_device = GetFunctionResultVirtualDevice(function_node);
+        VirtualDevice virtual_device = function_node->virtual_device();
         if (!virtual_device->IsFullyUnconstrained()) {
           VLOG(2) << "global '" << kv.first->name_hint << "' has virtual device " << virtual_device;
           global_var_virtual_devices_.emplace(kv.first, virtual_device);
@@ -74,7 +74,7 @@ VirtualDevice LexicalOnDeviceMixin::GetVirtualDevice(const Expr& expr) const {
       }
       // else: fallthrough to unconstrained
     } else {
-      return GetFunctionResultVirtualDevice(function_node);
+      return function_node->virtual_device();
     }
   } else {
     if (!expr_virtual_devices_.empty()) {
@@ -132,11 +132,11 @@ void DeviceAwareExprVisitor::VisitExpr_(const FunctionNode* function_node) {
     DeviceAwareVisitExpr_(function_node);
   } else {
     // Function parameters come into scope.
-    for (size_t i = 0; i < function_node->params.size(); ++i) {
-      PushBoundVar(function_node->params[i], GetFunctionParamVirtualDevice(function_node, i));
+    for (auto param : function_node->params) {
+      PushBoundVar(param, param->virtual_device());
     }
     // Entering scope of function body.
-    PushVirtualDevice(GetFunctionResultVirtualDevice(function_node));
+    PushVirtualDevice(function_node->virtual_device());
     EnterFunctionBody();
 
     DeviceAwareVisitExpr_(function_node);
@@ -218,11 +218,11 @@ Expr DeviceAwareExprMutator::VisitExpr_(const FunctionNode* function_node) {
     return DeviceAwareVisitExpr_(function_node);
   } else {
     // Function parameters come into scope.
-    for (size_t i = 0; i < function_node->params.size(); ++i) {
-      PushBoundVar(function_node->params[i], GetFunctionParamVirtualDevice(function_node, i));
+    for (auto param : function_node->params) {
+      PushBoundVar(param, param->virtual_device());
     }
     // Entering scope of function body.
-    PushVirtualDevice(GetFunctionResultVirtualDevice(function_node));
+    PushVirtualDevice(function_node->virtual_device());
     EnterFunctionBody();
 
     Expr result = DeviceAwareVisitExpr_(function_node);
diff --git a/src/relay/transforms/device_aware_visitors.h b/src/relay/transforms/device_aware_visitors.h
index 9340c03fc2d5..8e669725b91d 100644
--- a/src/relay/transforms/device_aware_visitors.h
+++ b/src/relay/transforms/device_aware_visitors.h
@@ -145,11 +145,11 @@ class DeviceAwareExprFunctor<void(const Expr& n)> : public ExprFunctor<void(cons
       return DeviceAwareVisitExpr_(function_node);
     } else {
       // Function parameters come into scope.
-      for (size_t i = 0; i < function_node->params.size(); ++i) {
-        PushBoundVar(function_node->params[i], GetFunctionParamVirtualDevice(function_node, i));
+      for (auto param : function_node->params) {
+        PushBoundVar(param, param->virtual_device());
       }
       // Entering scope of function body.
-      VirtualDevice virtual_device = GetFunctionResultVirtualDevice(function_node);
+      VirtualDevice virtual_device = function_node->virtual_device();
       VLOG(2) << "entering " << virtual_device << " for function:" << std::endl
               << PrettyPrint(GetRef<Function>(function_node));
       PushVirtualDevice(virtual_device);
diff --git a/src/relay/transforms/device_planner.cc b/src/relay/transforms/device_planner.cc
index c7ca2227fd90..3562da3b0d6f 100644
--- a/src/relay/transforms/device_planner.cc
+++ b/src/relay/transforms/device_planner.cc
@@ -22,9 +22,6 @@
  * \brief Determines a unique \p VirtualDevice to hold the result of every Relay sub-expression.
  * This pass can be run multiple times, and can be run both before and after lowering.
  *
- * TODO(mbs): Rename VirtualDevice |-> VirtualDevice, and use 'virtual device' (or just 'device')
- * throughout.
- *
  * We say a Relay expression E is 'on device D' if the result of executing E is stored on D.
  * We represent D by an \p VirtualDevice, which means we can track anywhere from an arbitrary device
  * of some \p DLDeviceType to a specific memory scope on a specific (virtual) \p Device who's
@@ -423,7 +420,7 @@ class RewriteOnDevices : public ExprMutator {
  * \endcode
  * we discover \p %b must be on device \p d.
  */
-class DeviceAnalyzer : public ExprVisitor {
+class DeviceAnalyzer : public MixedModeVisitor {
  public:
   DeviceAnalyzer(IRModule mod, CompilationConfig config)
       : mod_(std::move(mod)), domains_(std::make_unique<DeviceDomains>(std::move(config))) {}
@@ -504,7 +501,6 @@ class DeviceAnalyzer : public ExprVisitor {
     args_and_result_domains.reserve(vanilla_call->args.size() + 1);
     for (const auto& arg : vanilla_call->args) {
       args_and_result_domains.emplace_back(domains_->DomainFor(arg));
-      VisitExpr(arg);
     }
     args_and_result_domains.emplace_back(domains_->DomainFor(call));
     auto implied_domain =
@@ -587,15 +583,14 @@ class DeviceAnalyzer : public ExprVisitor {
 
     // If the function already has VirtualDevice attributes then we can further constrain the
     // function's domain to match them.
-    if (!GetFunctionResultVirtualDevice(function_node)->IsFullyUnconstrained()) {
+    if (!function_node->virtual_device()->IsFullyUnconstrained()) {
       std::vector<DeviceDomainPtr> args_and_result;
-      for (size_t i = 0; i < function_node->params.size(); ++i) {
+      for (auto param : function_node->params) {
         args_and_result.emplace_back(
-            domains_->ForVirtualDevice(function_node->params[i]->checked_type(),
-                                       GetFunctionParamVirtualDevice(function_node, i)));
+            domains_->ForVirtualDevice(param->checked_type(), param->virtual_device()));
       }
-      args_and_result.emplace_back(domains_->ForVirtualDevice(
-          function_node->body->checked_type(), GetFunctionResultVirtualDevice(function_node)));
+      args_and_result.emplace_back(domains_->ForVirtualDevice(function_node->body->checked_type(),
+                                                              function_node->virtual_device()));
       auto annotation_domain = domains_->MakeHigherOrderDomain(std::move(args_and_result));
       if (domains_->UnifyOrNull(func_domain, annotation_domain) == nullptr) {  // higher-order
         // TODO(mbs): Proper diagnostics.
@@ -625,7 +620,6 @@ class DeviceAnalyzer : public ExprVisitor {
     for (size_t i = 0; i < tuple->fields.size(); i++) {
       auto domain = domains_->DomainFor(tuple->fields[i]);  // may be higher-order
       domains_->UnifyExprCollapsed(tuple, domain);          // collapse to first-order if needed
-      VisitExpr(tuple->fields[i]);
     }
   }
 
@@ -634,7 +628,6 @@ class DeviceAnalyzer : public ExprVisitor {
     auto domain = domains_->DomainFor(tuple_get_item);  // may be higher-order
     domains_->UnifyExprCollapsed(tuple_get_item_node->tuple,
                                  domain);  // collapse to first-order if needed
-    VisitExpr(tuple_get_item_node->tuple);
   }
 
   class DevicePatternAnalyzer : public PatternVisitor {
@@ -883,7 +876,6 @@ class DeviceDefaulter : public ExprVisitor {
 };
 
 /* =============== Phase 3 =============== */
-
 /*!
  * \brief Inserts missing "device_copy" CallNodes, and ensures the device type of every
  * sub-expression in a module can be easily recovered by a later transformation using simple
@@ -892,8 +884,9 @@ class DeviceDefaulter : public ExprVisitor {
  * - Discard any existing "on_device" CallNodes since their job is done. Similarly, discard
  *   any existing "device_copy" CallNodes which are no-ops.
  *
- * - Functions are given "param_virtual_devices" and "result_virtual_device" attributes to capture
- *   the device type for its parameters and result.
+ * - The result virtual device for a function is stored in the function's virtual_device_ field
+ *   and the virtual devices of the function's parameters are stored in the parameter's
+ *   virtual_device_ field.
  *
  * - Additional "device_copy" CallNodes are inserted wherever there's a transition between
  *   storage device types. Since the DeviceAnalyzer phase succeeded this can only happen
@@ -1009,14 +1002,25 @@ class DeviceCapturer : public ExprMutator {
     ICHECK_EQ(func_domain->function_arity(), function_node->params.size());
     VirtualDevice result_virtual_device = domains_->ResultVirtualDevice(func_domain);
     ICHECK(!result_virtual_device->IsFullyUnconstrained());
-    Array<VirtualDevice> param_virtual_devices;
-    param_virtual_devices.reserve(function_node->params.size());
+
+    // Map the function parameters to a new variable annotated with a virtual device so
+    // we can substitute them later.
+    Map<Var, Expr> annotated_bind_map;
+    Array<Var> annotated_params;
+    annotated_params.reserve(function_node->params.size());
     for (size_t i = 0; i < function_node->params.size(); ++i) {
       VirtualDevice param_virtual_device =
           domains_->ResultVirtualDevice(func_domain->function_param(i));
+      VLOG(4) << "Param: " << function_node->params[i];
+      Var annotated_var = WithFields(function_node->params[i], {}, {}, param_virtual_device);
+      VLOG(4) << "Annotated param: " << annotated_var;
+      VLOG(4) << "VirtualDevice: " << annotated_var->virtual_device();
       ICHECK(!param_virtual_device->IsFullyUnconstrained());
-      param_virtual_devices.push_back(param_virtual_device);
+      annotated_bind_map.Set(function_node->params[i], annotated_var);
+      annotated_params.push_back(annotated_var);
     }
+    // Eventually we probably want to bind before visiting, but for now this is causing an issue
+    // with the GetVirtualDevice utility, so leaving as is for now.
 
     // Rewrite the body. Note that the body may have begun with an "on_device" so
     // be prepared to insert a "device_copy".
@@ -1024,10 +1028,14 @@ class DeviceCapturer : public ExprMutator {
         /*lexical_virtual_device=*/result_virtual_device,
         /*expected_virtual_device=*/result_virtual_device,
         /*child_virtual_device=*/GetVirtualDevice(function_node->body), function_node->body);
-
+    VLOG(4) << "Visited body: " << body;
     Function func = WithFields(GetRef<Function>(function_node), function_node->params, body);
-    return FunctionOnDevice(func, std::move(param_virtual_devices),
-                            std::move(result_virtual_device));
+    VLOG(4) << "New function: " << func;
+    func = SubstituteBoundVars(func, annotated_bind_map);
+    VLOG(4) << "Func with bound params: " << func;
+    func->virtual_device_ = result_virtual_device;
+    VLOG(4) << "Func with bound params & result vid set: " << func;
+    return func;
   }
 
   Expr VisitExpr_(const CallNode* call_node) final {
@@ -1273,7 +1281,9 @@ class DeviceCapturer : public ExprMutator {
 /*! \brief Rewrite the "on_device" calls (and implicitly re-type-check). */
 tvm::transform::Pass Rewrite() {
   auto pass_func = [](Function f, IRModule m, transform::PassContext ctxt) {
-    return Downcast<Function>(RewriteOnDevices(std::move(m)).Mutate(f));
+    auto attrs = m->attrs;
+    auto r = Downcast<Function>(RewriteOnDevices(std::move(m)).Mutate(f));
+    return attrs.defined() ? WithAttrs(r, {attrs->dict}) : r;
   };
   return tvm::relay::transform::CreateFunctionPass(pass_func, 0, "PlanDevicesRewrite", {});
 }
diff --git a/src/relay/transforms/fake_quantization_to_integer.cc b/src/relay/transforms/fake_quantization_to_integer.cc
index 4273fc29cec8..4fd034edc199 100644
--- a/src/relay/transforms/fake_quantization_to_integer.cc
+++ b/src/relay/transforms/fake_quantization_to_integer.cc
@@ -25,6 +25,9 @@
 
 #include "fake_quantization_to_integer.h"
 
+#include <tvm/ir/affine_type.h>
+#include <tvm/relay/attrs/nn.h>
+#include <tvm/relay/dataflow_matcher.h>
 #include <tvm/relay/expr.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/qnn/attrs.h>
@@ -32,6 +35,8 @@
 
 #include <unordered_map>
 
+#include "../qnn/utils.h"
+
 namespace tvm {
 namespace relay {
 
@@ -69,6 +74,28 @@ namespace relay {
  *
  * After the second and third passes run, the first pass replaces the quantize with the
  * rewritten subgraph and the processing continues
+ *
+ *
+ * After that an additional QAT pass can be enabled by use_qat flag. The goal of the pass is to find
+ * operations in those regions(which were not successfully converted by the main pass) that can
+ * still be converted into quantized form. The idea is to find and transform operations with
+ * dequantized inputs one by one individually. Only operations for which all parameters can be
+ * explicitly calculated are allowed. For example, if on the above general  pattern op2 is not
+ * registered with the FTVMFakeQuantizationToInteger attribute, op1 operation can still be
+ * converted. Converted pattern below:
+ *
+ *   x    w
+ *   |    |
+ *    \   /
+ *     op1
+ *      |
+ *     dq
+ *      |
+ *     op2
+ *      |
+ *      q
+ *
+ * This pass works in the same multi-pass approach.
  */
 
 using ExprSet = std::unordered_set<Expr, ObjectPtrHash, ObjectPtrEqual>;
@@ -270,16 +297,250 @@ class FakeQuantizationRewriter : public MixedModeMutator {
   const bool hard_fail_;
 };
 
-Expr FakeQuantizationToInteger(const Expr& expr, const IRModule& mod, bool hard_fail) {
-  return FakeQuantizationRewriter(hard_fail).Mutate(expr);
+/* Checks if the operation to convert QAT pass is enabled.
+ * The following conditions must be satisfied:
+ * 1. operations registered for FTVMFakeQuantizationToInteger;
+ * 2. Unary operators or operators with with the TensorAffineType calculated during
+ * FTVMFakeQuantizationToInteger conversion;
+ * 3. Not one of the "key" operations: requantize,quantize and dequantize(they are at the boundaries
+ * of regions defined to be quantized).
+ */
+bool is_op_enabled_for_optional_fq2i(const CallNode* call_node) {
+  const Op op = Downcast<Op>(call_node->op);
+  static auto fqfq = Op::GetAttrMap<FTVMFakeQuantizationToInteger>("FTVMFakeQuantizationToInteger");
+  static std::unordered_set<relay::Expr, tvm::ObjectHash, tvm::ObjectEqual> ops = {
+      Op::Get("broadcast_to"),
+      Op::Get("clip"),
+      Op::Get("expand_dims"),
+      Op::Get("max"),
+      Op::Get("maximum"),
+      Op::Get("min"),
+      Op::Get("minimum"),
+      Op::Get("nn.avg_pool2d"),
+      Op::Get("nn.batch_flatten"),
+      Op::Get("nn.batch_matmul"),
+      Op::Get("nn.bias_add"),
+      Op::Get("nn.conv2d"),
+      Op::Get("nn.conv2d_transpose"),
+      Op::Get("nn.dense"),
+      Op::Get("nn.depth_to_space"),
+      Op::Get("nn.global_avg_pool2d"),
+      Op::Get("nn.max_pool2d"),
+      Op::Get("nn.pad"),
+      Op::Get("nn.relu"),
+      Op::Get("reshape"),
+      Op::Get("split"),
+      Op::Get("squeeze"),
+      Op::Get("strided_slice"),
+      Op::Get("transpose")};
+
+  return ops.find(call_node->op) != ops.end() && fqfq.count(Downcast<Op>(op));
+}
+
+class QATSubgraphExtractor : public ExprVisitor {
+ public:
+  const ExprSet GetSubgraph(const Expr& expr) {
+    expr_call_node_ = expr.as<CallNode>();
+    ICHECK(expr_call_node_ != nullptr);
+    ICHECK(is_op_enabled_for_optional_fq2i(expr_call_node_));
+
+    VisitExpr(expr);
+
+    ExprSet subgraph;
+    if (is_fake_quantized_) {
+      for (auto kv : this->visit_counter_) {
+        if (auto call_node = GetRef<ObjectRef>(kv.first).as<CallNode>()) {
+          if (call_node != expr_call_node_) {
+            subgraph.insert(Downcast<Expr>(GetRef<ObjectRef>(kv.first)));
+          }
+        }
+      }
+    }
+    return subgraph;
+  }
+  const AffineTypeMap GetAffineTypes() { return affine_types_; }
+  void VisitExpr(const Expr& expr) override {
+    // When looking for fake quantized subgraphs, we only support data-flow regions of the graph,
+    // i.e. call nodes/tuples/constants/etc. If we see anything else (like control flow) we
+    // abort the rewrite.
+    if (expr.as<CallNode>() == nullptr && expr.as<OpNode>() == nullptr &&
+        expr.as<TupleNode>() == nullptr && expr.as<TupleGetItemNode>() == nullptr &&
+        expr.as<ConstantNode>() == nullptr) {
+      DLOG(INFO) << "FakeQuantizationToInteger found a non - dataflow op inside a fake quantize "
+                    "region, aborting this rewrite";
+      is_fake_quantized_ = false;
+    } else {
+      ExprVisitor::VisitExpr(expr);
+    }
+  }
+
+ protected:
+  void VisitExpr_(const CallNode* call_node) override {
+    if (call_node->op == dequantize_op_) {
+      const auto* attrs = call_node->attrs.as<qnn::DequantizeAttrs>();
+      ICHECK(attrs != nullptr);
+
+      affine_types_.Set(
+          GetRef<Expr>(call_node),
+          TensorAffineType(
+              call_node->args[1], call_node->args[2],
+              tvm::relay::transform::InferTypeLocal(call_node->args[0]).as<TensorTypeNode>()->dtype,
+              attrs->axis));
+    } else if (call_node == expr_call_node_) {
+      for (auto arg : call_node->args) {
+        VisitExpr(arg);
+      }
+    } else {
+      // run normally on everything else.
+      ExprVisitor::VisitExpr_(call_node);
+    }
+  }
+
+  const Op dequantize_op_ = Op::Get("qnn.dequantize");
+  bool is_fake_quantized_ = true;
+  AffineTypeMap affine_types_;
+  const CallNode* expr_call_node_ = nullptr;
+};
+
+class QATSubgraphMutator : public ExprMutator {
+ public:
+  QATSubgraphMutator(ExprSet subgraph, AffineTypeMap affine_types, bool hard_fail)
+      : subgraph_(subgraph), affine_types_(affine_types), hard_fail_(hard_fail) {}
+
+  Expr MutateSubgraph(const Expr& expr) {
+    if (subgraph_.size() == 0) {
+      return expr;
+    }
+
+    quantize_node_ = expr.as<CallNode>();
+    ICHECK(quantize_node_);
+    ICHECK(is_op_enabled_for_optional_fq2i(quantize_node_));
+
+    for (auto node : subgraph_) {
+      const Op op = Downcast<Op>(node.as<CallNode>()->op);
+
+      if (node.as<CallNode>()->op != dequantize_op_) {
+        if (hard_fail_) {
+          LOG(FATAL) << "Not dequantization was found in the input arguments for"
+                     << AsText(op, false) << std::endl;
+        } else {
+          DLOG(INFO) << "Not dequantization was found in the input arguments for "
+                     << AsText(op, false) << std::endl;
+          return expr;
+        }
+      }
+    }
+    try {
+      return Mutate(expr);
+    } catch (std::exception& e) {
+      if (hard_fail_) {
+        throw e;
+      } else {
+        DLOG(INFO) << "Ran into an error rewriting a subgraph, skipping" << expr << std::endl;
+        return expr;
+      }
+    }
+  }
+
+ protected:
+  Expr VisitExpr_(const CallNode* call_node) {
+    Expr out;
+    static auto fqfq =
+        Op::GetAttrMap<FTVMFakeQuantizationToInteger>("FTVMFakeQuantizationToInteger");
+
+    Op op = Downcast<Op>(call_node->op);
+    if (fqfq.count(op)) {
+      Expr expr;
+      if (op == dequantize_op_) {
+        expr = GetRef<Expr>(call_node);
+      } else {
+        expr = ExprMutator::VisitExpr_(call_node);
+      }
+      // Call the rewrite
+      Array<ObjectRef> vals = fqfq[op](expr, affine_types_);
+      // Save the outputs of the rewrite
+      ICHECK(vals.size() == 2)
+          << "got the wrong number of returned arguments from FTVMFakeQuantizationToInteger for "
+          << AsText(op, false);
+      out = Downcast<Expr>(vals[0]);
+
+      affine_types_.Set(out, Downcast<AffineType>(vals[1]));
+
+      if (call_node == quantize_node_) {
+        out = qnn::MakeDequantize(out, vals[1].as<TensorAffineTypeNode>()->scale,
+                                  vals[1].as<TensorAffineTypeNode>()->zero_point,
+                                  vals[1].as<TensorAffineTypeNode>()->axis);
+      }
+    } else {
+      ICHECK(false) << "When rewriting a fake quantized graph, found an invalid node "
+                    << AsText(GetRef<Expr>(call_node), false);
+    }
+    return out;
+  }
+
+  Expr VisitExpr_(const TupleNode* node) {
+    Expr expr = ExprMutator::VisitExpr_(node);
+    auto new_node = expr.as<TupleNode>();
+    Array<TensorAffineType> types;
+    for (Expr field : new_node->fields) {
+      ICHECK(affine_types_[field].as<TensorAffineTypeNode>());
+      types.push_back(Downcast<TensorAffineType>(affine_types_[field]));
+    }
+    affine_types_.Set(expr, TupleAffineType(types));
+    return expr;
+  }
+
+  Expr VisitExpr_(const TupleGetItemNode* node) {
+    Expr expr = ExprMutator::VisitExpr_(node);
+    auto tuple_type = affine_types_[expr.as<TupleGetItemNode>()->tuple].as<TupleAffineTypeNode>();
+    affine_types_.Set(expr, tuple_type->types[node->index]);
+    return expr;
+  }
+
+  ExprSet subgraph_;
+  AffineTypeMap affine_types_;
+  const bool hard_fail_;
+  const Op dequantize_op_ = Op::Get("qnn.dequantize");
+  const CallNode* quantize_node_ = nullptr;
+};
+
+class QATRewriter : public MixedModeMutator {
+ public:
+  explicit QATRewriter(bool hard_fail) : hard_fail_(hard_fail) {}
+
+ protected:
+  Expr Rewrite_(const CallNode* pre, const Expr& post) override {
+    if (const CallNode* call_node = post.as<CallNode>()) {
+      const Op op = Downcast<Op>(call_node->op);
+      if (is_op_enabled_for_optional_fq2i(call_node)) {
+        QATSubgraphExtractor extractor;
+        ExprSet subgraph = extractor.GetSubgraph(post);
+        AffineTypeMap affine_types = extractor.GetAffineTypes();
+        Expr out = QATSubgraphMutator(subgraph, affine_types, hard_fail_).MutateSubgraph(post);
+        return out;
+      }
+    }
+    return post;
+  }
+  const bool hard_fail_;
+};
+
+Expr FakeQuantizationToInteger(const Expr& expr, const IRModule& mod, bool hard_fail,
+                               bool use_qat) {
+  auto fq_expr = FakeQuantizationRewriter(hard_fail).Mutate(expr);
+  if (use_qat) {
+    fq_expr = tvm::relay::InferType(fq_expr);
+    fq_expr = QATRewriter(hard_fail).Mutate(fq_expr);
+  }
+  return fq_expr;
 }
 
 namespace transform {
 
-Pass FakeQuantizationToInteger(bool hard_fail) {
+Pass FakeQuantizationToInteger(bool hard_fail, bool use_qat) {
   runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
       [=](Function f, IRModule m, PassContext pc) {
-        return Downcast<Function>(FakeQuantizationToInteger(f, m, hard_fail));
+        return Downcast<Function>(FakeQuantizationToInteger(f, m, hard_fail, use_qat));
       };
   return CreateFunctionPass(pass_func, 0, "FakeQuantizationToInteger", {"InferType"});
 }
diff --git a/src/relay/transforms/fold_constant.cc b/src/relay/transforms/fold_constant.cc
index dd8195797e8d..3a8391e05856 100644
--- a/src/relay/transforms/fold_constant.cc
+++ b/src/relay/transforms/fold_constant.cc
@@ -253,9 +253,10 @@ class ConstantFolder : public MixedModeMutator {
     // Use a fresh build context in case we are already in a build context.
     // needed for both execution and creation(due to JIT)
     With<transform::PassContext> fresh_build_ctx(transform::PassContext::Create());
-
-    Expr result = ObjectToExpr(
-        Eval(expr, module_->type_definitions, module_->Imports(), eval_cpu_dev_, eval_cpu_target_));
+    Map<String, ObjectRef> dict =
+        (module_->attrs.defined()) ? module_->attrs->dict : Map<String, ObjectRef>();
+    Expr result = ObjectToExpr(Eval(expr, module_->type_definitions, module_->Imports(),
+                                    eval_cpu_dev_, eval_cpu_target_, dict));
     VLOG(1) << "Evaluated to constant:" << std::endl << PrettyPrint(result);
     return result;
   }
diff --git a/src/relay/transforms/fuse_ops.cc b/src/relay/transforms/fuse_ops.cc
index 5037b32ce615..e25b8db152c4 100644
--- a/src/relay/transforms/fuse_ops.cc
+++ b/src/relay/transforms/fuse_ops.cc
@@ -25,6 +25,7 @@
  *   Fuse necessary ops into a single one.
  */
 #include <tvm/relay/analysis.h>
+#include <tvm/relay/executor.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/op_attr_types.h>
 #include <tvm/relay/transform.h>
@@ -85,6 +86,7 @@ constexpr uint32_t kMaxFusedOps = 256;
 static const Op& stop_fusion_op = Op::Get("annotation.stop_fusion");
 
 TVM_REGISTER_PASS_CONFIG_OPTION("relay.FuseOps.max_depth", Integer);
+TVM_REGISTER_PASS_CONFIG_OPTION("relay.FuseOps.link_params", Bool);
 
 /*!
  * \brief Indexed data flow graph in forward direction.
@@ -809,8 +811,19 @@ std::vector<GraphPartitioner::Group*> GraphPartitioner::Partition(
 
 class FuseMutator : private MixedModeMutator {
  public:
+  FuseMutator(int fuse_opt_level, size_t max_fuse_depth, bool link_params)
+      : fuse_opt_level_(fuse_opt_level),
+        max_fuse_depth_(max_fuse_depth),
+        link_params_(link_params) {}
+
+  // Run the transform
+  Expr Transform(const Expr& body) {
+    return Transform(body, fuse_opt_level_, max_fuse_depth_, link_params_);
+  }
+
+ protected:
   // Run the transform
-  Expr Transform(const Expr& body, int fuse_opt_level, size_t max_fuse_depth) {
+  Expr Transform(const Expr& body, int fuse_opt_level, size_t max_fuse_depth, bool link_params) {
     // setup the group map.
     auto graph = IndexedForwardGraph::Create(&arena_, body);
     auto groups = GraphPartitioner(&arena_, fuse_opt_level, max_fuse_depth).Partition(graph);
@@ -824,6 +837,10 @@ class FuseMutator : private MixedModeMutator {
   }
 
  private:
+  int fuse_opt_level_;
+  size_t max_fuse_depth_;
+  bool link_params_;
+
   using MixedModeMutator::VisitExpr_;
 
   /*! \brief Temporary information from each group. */
@@ -994,8 +1011,12 @@ class FuseMutator : private MixedModeMutator {
       auto type = arg->checked_type();
       Expr new_arg = this->Mutate(arg);
       if (current_group != arg_group) {
-        Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
-        new_args.push_back(param);
+        if (!link_params_ || new_arg.as<ConstantNode>() == nullptr) {
+          Var param = ginfo_[current_group].GetOrAllocParam(new_arg, type);
+          new_args.push_back(param);
+        } else {
+          new_args.push_back(new_arg);
+        }
       } else {
         new_args.push_back(new_arg);
       }
@@ -1017,8 +1038,9 @@ class FuseMutator : private MixedModeMutator {
   }
 };
 
-Expr FuseOps(const Expr& expr, int fuse_opt_level, size_t max_fuse_depth, const IRModule& module) {
-  return FuseMutator().Transform(expr, fuse_opt_level, max_fuse_depth);
+Expr FuseOps(const Expr& expr, int fuse_opt_level, size_t max_fuse_depth, bool link_params,
+             const IRModule& module) {
+  return FuseMutator(fuse_opt_level, max_fuse_depth, link_params).Transform(expr);
 }
 
 namespace transform {
@@ -1026,9 +1048,16 @@ namespace transform {
 Pass FuseOps(int fuse_opt_level) {
   runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
       [=](Function f, IRModule m, PassContext pc) {
+        bool link_params = false;
+        Executor executor =
+            m->GetAttr<Executor>(tvm::attr::kExecutor).value_or(NullValue<Executor>());
+        link_params = executor.defined()
+                          ? executor->attrs.GetAttr<Bool>("link-params").value_or(Bool(link_params))
+                          : link_params;
+        link_params = pc->GetConfig("relay.FuseOps.link_params", Bool(link_params)).value();
         int opt_level = fuse_opt_level == -1 ? pc->opt_level : fuse_opt_level;
         auto max_fuse_depth = pc->GetConfig("relay.FuseOps.max_depth", Integer(kMaxFusedOps));
-        return Downcast<Function>(FuseOps(f, opt_level, max_fuse_depth.value(), m));
+        return Downcast<Function>(FuseOps(f, opt_level, max_fuse_depth.value(), link_params, m));
       };
   return CreateFunctionPass(pass_func, 0, "FuseOps", {"InferType"});
 }
diff --git a/src/relay/transforms/memory_alloc.cc b/src/relay/transforms/memory_alloc.cc
index 900442e9b9a8..5b584e199dc7 100644
--- a/src/relay/transforms/memory_alloc.cc
+++ b/src/relay/transforms/memory_alloc.cc
@@ -140,6 +140,9 @@ class DialectRewriter : public transform::DeviceAwareExprMutator {
 
     VirtualDevice virtual_device = GetVirtualDevice(call);
     ICHECK(!virtual_device->IsFullyUnconstrained());
+    ICHECK(!scopes_.empty())
+        << "Calls out of a let block are not supported, do you forget to transform "
+        << "with ToANormalForm or set opt_level >= 1 in the pass context?";
     LetList& scope = scopes_.back();
 
     std::vector<Expr> new_args;
diff --git a/src/relay/transforms/pattern_utils.h b/src/relay/transforms/pattern_utils.h
index 7d2657eb04f2..cf97d2c25d2d 100644
--- a/src/relay/transforms/pattern_utils.h
+++ b/src/relay/transforms/pattern_utils.h
@@ -63,6 +63,9 @@ namespace relay {
   } else if (type == DataType::Float(16)) {                                           \
     typedef uint16_t DType;                                                           \
     { __VA_ARGS__ }                                                                   \
+  } else if (type == DataType::BFloat(16)) {                                          \
+    typedef uint16_t DType;                                                           \
+    { __VA_ARGS__ }                                                                   \
   } else if (type == DataType::Int(64)) {                                             \
     typedef int64_t DType;                                                            \
     { __VA_ARGS__ }                                                                   \
@@ -259,6 +262,11 @@ inline Constant MakeConstantScalar(DataType dtype, T value) {
       // storage is uint16_t
       *static_cast<DType*>(arr->data) =
           __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 10>(static_cast<float>(value));
+    } else if (dtype == DataType::BFloat(16)) {
+      // convert to bfloat16
+      // storage is uint16_t
+      *static_cast<DType*>(arr->data) =
+          __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 7>(static_cast<float>(value));
     } else {
       *static_cast<DType*>(arr->data) = value;
     }
@@ -286,6 +294,12 @@ static inline Constant MakeConstantTensor(DataType dtype, std::vector<int64_t> s
         *(static_cast<DType*>(arr->data) + i) =
             __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 10>(
                 static_cast<float>(value[i]));
+      } else if (dtype == DataType::BFloat(16)) {
+        // convert to bfloat16
+        // storage is uint16_t
+        *(static_cast<DType*>(arr->data) + i) =
+            __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 7>(
+                static_cast<float>(value[i]));
       } else {
         *(static_cast<DType*>(arr->data) + i) = value[i];
       }
@@ -314,6 +328,12 @@ static inline Constant MakeConstantTensor(DataType dtype, std::vector<int64_t> s
         *(static_cast<DType*>(arr->data) + i) =
             __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 10>(
                 static_cast<float>(value[i]));
+      } else if (dtype == DataType::BFloat(16)) {
+        // convert to bfloat16
+        // storage is uint16_t
+        *(static_cast<DType*>(arr->data) + i) =
+            __truncXfYf2__<float, uint32_t, 23, uint16_t, uint16_t, 7>(
+                static_cast<float>(value[i]));
       } else {
         *(static_cast<DType*>(arr->data) + i) = value[i];
       }
@@ -417,6 +437,12 @@ static inline dmlc::optional<long double> TryToScalar(const runtime::NDArray& ar
     } else if (array->dtype.bits == 64) {
       return dmlc::optional<long double>(reinterpret_cast<double*>(array->data)[i]);
     }
+  } else if (array->dtype.code == kDLBfloat) {
+    if (array->dtype.bits == 16) {
+      return dmlc::optional<long double>(
+          __extendXfYf2__<uint16_t, uint16_t, 7, float, uint32_t, 23>(
+              reinterpret_cast<uint16_t*>(array->data)[i]));
+    }
   }
   return dmlc::optional<long double>();
 }
@@ -500,6 +526,11 @@ inline Expr Exp(Expr e) {
   return Call(op, {e});
 }
 
+inline Expr Erf(Expr e) {
+  static const Op& op = Op::Get("erf");
+  return Call(op, {e});
+}
+
 inline Expr FastExp(Expr e) {
   static const Op& op = Op::Get("fast_exp");
   return Call(op, {e});
@@ -524,6 +555,11 @@ inline Expr Log(Expr e) {
   static const Op& op = Op::Get("log");
   return Call(op, {e});
 }
+
+inline Expr Tanh(Expr e) {
+  static const Op& op = Op::Get("tanh");
+  return Call(op, {e});
+}
 /*!
  * \brief Get an immediate scalar from a Constant expr.
  *
@@ -550,6 +586,11 @@ inline Expr Sqrt(Expr x) {
   return Call(op, {x}, Attrs(), {});
 }
 
+inline Expr Sigmoid(Expr x) {
+  static const Op& op = Op::Get("sigmoid");
+  return Call(op, {x}, Attrs(), {});
+}
+
 inline Expr Rsqrt(Expr x) {
   static const Op& op = Op::Get("rsqrt");
   return Call(op, {x}, Attrs(), {});
diff --git a/src/relay/transforms/split_args.cc b/src/relay/transforms/split_args.cc
index a5266df8b057..00b9a3be3b2e 100644
--- a/src/relay/transforms/split_args.cc
+++ b/src/relay/transforms/split_args.cc
@@ -85,7 +85,8 @@ namespace transform {
 Pass SplitArgs(int max_function_args) {
   runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
       [=](Function f, IRModule m, PassContext pc) {
-        return Downcast<Function>(SplitArgs(f, max_function_args));
+        auto r = Downcast<Function>(SplitArgs(f, max_function_args));
+        return m->attrs.defined() ? WithAttrs(r, {m->attrs->dict}) : r;
       };
   return CreateFunctionPass(pass_func, 1, "SplitArgs", {"InferType"});
 }
diff --git a/src/relay/transforms/to_a_normal_form.cc b/src/relay/transforms/to_a_normal_form.cc
index 2f6efb9cef9a..8319726b79c5 100644
--- a/src/relay/transforms/to_a_normal_form.cc
+++ b/src/relay/transforms/to_a_normal_form.cc
@@ -304,9 +304,9 @@ class Fill : ExprFunctor<Expr(const Expr&, const Var&)>, private transform::Lexi
     } else {
       // Keep track of expression and bound variable device types for lexically enclosing
       // sub-expressions.
-      PushVirtualDevice(GetFunctionResultVirtualDevice(f));
-      for (size_t i = 0; i < f->params.size(); ++i) {
-        PushBoundVar(f->params[i], GetFunctionParamVirtualDevice(f, i));
+      PushVirtualDevice(f->virtual_device());
+      for (auto param : f->params) {
+        PushBoundVar(param, param->virtual_device());
       }
       EnterFunctionBody();
       ret = WithFields(GetRef<Function>(f), f->params,
diff --git a/src/relay/transforms/type_infer.cc b/src/relay/transforms/type_infer.cc
index 456e210f7343..7de43eb36882 100644
--- a/src/relay/transforms/type_infer.cc
+++ b/src/relay/transforms/type_infer.cc
@@ -41,6 +41,7 @@
 #include <tvm/ir/transform.h>
 #include <tvm/ir/type_functor.h>
 #include <tvm/relay/analysis.h>
+#include <tvm/relay/dataflow_matcher.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/pattern_functor.h>
 #include <tvm/relay/transform.h>
@@ -914,15 +915,9 @@ Type InferTypeLocal(const Expr& expr) {
   */
   SameTypedSubgraphExtractor subgraph_extractor;
   Expr sub_graph = subgraph_extractor(expr);
-  auto mod = IRModule::FromExpr(sub_graph);
-  mod = transform::InferType()(mod);
 
   Type result_type;
-  if (expr.as<FunctionNode>()) {
-    result_type = mod->Lookup("main")->checked_type();
-  } else {
-    result_type = mod->Lookup("main").as<FunctionNode>()->body->checked_type();
-  }
+  result_type = relay::InferType(sub_graph)->checked_type();
 
   expr->checked_type_ = result_type;
   return result_type;
diff --git a/src/runtime/aot_executor/aot_executor.cc b/src/runtime/aot_executor/aot_executor.cc
new file mode 100644
index 000000000000..8dbf0b235b31
--- /dev/null
+++ b/src/runtime/aot_executor/aot_executor.cc
@@ -0,0 +1,204 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \brief Defines an implementation of Module-based Model Runtime Interface that works with
+ *        Ahead-of-Time compilation.
+ * \file aot_executor.cc
+ */
+
+#include "aot_executor.h"
+
+#include <tvm/runtime/c_runtime_api.h>
+
+#include <memory>
+
+#include "../meta_data.h"
+
+namespace tvm {
+namespace runtime {
+
+AotExecutor::AotExecutor(tvm::runtime::Module module, const std::vector<Device>& devs)
+    : module_{module}, devices_{devs} {
+  auto fmetadata = module->GetFunction("get_metadata");
+  CHECK(fmetadata != nullptr) << "Expected a module with PackedFunc get_metadata";
+  auto ret_value = fmetadata();
+  metadata_ = ret_value.AsObjectRef<tvm::runtime::metadata::Metadata>();
+
+
+  LOG(INFO) << "num inputs: " << metadata_->num_inputs();
+  LOG(INFO) << "num outputs: " << metadata_->num_outputs();
+  LOG(INFO) << metadata_->mod_name();
+
+  for (auto input : metadata_->inputs()) {
+    // TODO(areusch): Encode device information in Metadata.
+    args_.emplace_back(NDArray::Empty(ShapeTuple(input->shape().begin(), input->shape().end()),
+                                      input->dtype(), devices_[0]));
+    for (auto s : input->shape()) {
+      LOG(INFO) << s;
+    }
+  }
+
+  for (auto output : metadata_->outputs()) {
+    args_.emplace_back(NDArray::Empty(ShapeTuple(output->shape().begin(), output->shape().end()),
+                                      output->dtype(), devices_[0]));
+    for (auto s: output->shape()) {
+      LOG(INFO) << s;
+    }
+  }
+
+  for (auto pool : metadata_->pools()) {
+    args_.emplace_back(NDArray::Empty(ShapeTuple(pool->shape().begin(), pool->shape().end()),
+                                      pool->dtype(), devices_[0]));
+    for (auto s: pool->shape()) {
+      LOG(INFO) << s;
+    }
+  }
+}
+
+PackedFunc AotExecutor::GetFunction(const std::string& name,
+                                    const ObjectPtr<Object>& sptr_to_self) {
+  // Return member functions during query.
+  if (name == "set_input") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      if (String::CanConvertFrom(args[0])) {
+        int in_idx = this->GetInputIndex(args[0].operator String());
+        if (in_idx >= 0) this->SetInput(in_idx, args[1]);
+      } else {
+        this->SetInput(args[0], args[1]);
+      }
+    });
+  } else if (name == "set_input_zero_copy") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      if (String::CanConvertFrom(args[0])) {
+        int in_idx = this->GetInputIndex(args[0].operator String());
+        if (in_idx >= 0) this->SetInputZeroCopy(in_idx, args[1]);
+      } else {
+        this->SetInputZeroCopy(args[0], args[1]);
+      }
+    });
+  } else if (name == "set_output_zero_copy") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      if (String::CanConvertFrom(args[0])) {
+        int out_idx = this->GetOutputIndex(args[0].operator String());
+        if (out_idx >= 0) this->SetOutputZeroCopy(out_idx, args[1]);
+      } else {
+        this->SetOutputZeroCopy(args[0], args[1]);
+      }
+    });
+  } else if (name == "get_output") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      if (args.num_args == 2) {
+        this->CopyOutputTo(args[0], args[1]);
+      } else {
+        *rv = this->GetOutput(args[0]);
+      }
+    });
+  } else if (name == "get_input") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      int in_idx = 0;
+      if (String::CanConvertFrom(args[0])) {
+        in_idx = this->GetInputIndex(args[0].operator String());
+      } else {
+        in_idx = args[0];
+      }
+      if (in_idx >= 0) {
+        *rv = this->GetInput(in_idx);
+      }
+    });
+  } else if (name == "get_num_outputs") {
+    return PackedFunc(
+        [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->NumOutputs(); });
+  } else if (name == "get_num_inputs") {
+    return PackedFunc(
+        [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->NumInputs(); });
+  } else if (name == "run") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { this->Run(); });
+  } else if (name == "get_input_index") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      CHECK(String::CanConvertFrom(args[0])) << "Input key is not a string";
+      *rv = this->GetInputIndex(args[0].operator String());
+    });
+  } else {
+    return PackedFunc();
+  }
+}
+
+void AotExecutor::Run() {
+  auto pf = module_.GetFunction(
+      get_name_mangled(metadata_->mod_name(), ::tvm::runtime::symbol::tvm_module_main),
+      true /* query_imports */);
+  ICHECK(pf != nullptr) << "Module entrypoint is not defined";
+
+  const int num_args = args_.size();
+  ::std::unique_ptr<TVMValue> call_values{new TVMValue[num_args]};
+  ::std::unique_ptr<int> call_type_codes{new int[num_args]};
+  for (int i = 0; i < num_args; ++i) {
+    auto managed = args_[i].ToDLPack();
+    call_values.get()[i].v_handle = &managed->dl_tensor;
+    call_type_codes.get()[i] = kTVMDLTensorHandle;
+  }
+
+  TVMArgs args{call_values.get(), call_type_codes.get(), num_args};
+  TVMRetValue rv;
+  pf.CallPacked(args, &rv);
+}
+
+int AotExecutor::GetInputIndex(const std::string& name) {
+  auto inputs = metadata_->inputs();
+  for (unsigned int i = 0; i < inputs.size(); i++) {
+    if (inputs[i]->name() == name) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+int AotExecutor::GetOutputIndex(const std::string& name) {
+  auto outputs = metadata_->outputs();
+  for (unsigned int i = 0; i < outputs.size(); i++) {
+    if (outputs[i]->name() == name) {
+      return i;
+    }
+  }
+  return -1;
+}
+
+void AotExecutor::SetInput(int index, DLTensor* data_ref) { args_[index].CopyFrom(data_ref); }
+
+void AotExecutor::SetInputZeroCopy(int index, DLTensor* data_ref) {
+  ICHECK(false) << "not implemented";
+}
+
+void AotExecutor::SetOutputZeroCopy(int index, DLTensor* data_ref) {
+  ICHECK(false) << "not implemented";
+}
+
+int AotExecutor::NumOutputs() const { return metadata_->num_outputs(); }
+
+int AotExecutor::NumInputs() const { return metadata_->num_inputs(); }
+
+NDArray AotExecutor::GetInput(int index) const { return args_[index]; }
+
+NDArray AotExecutor::GetOutput(int index) const { return args_[metadata_->num_inputs() + index]; }
+
+void AotExecutor::CopyOutputTo(int index, DLTensor* data_out) { GetOutput(index).CopyTo(data_out); }
+
+}  // namespace runtime
+}  // namespace tvm
diff --git a/src/runtime/aot_executor/aot_executor.h b/src/runtime/aot_executor/aot_executor.h
new file mode 100644
index 000000000000..a213ef83bdc2
--- /dev/null
+++ b/src/runtime/aot_executor/aot_executor.h
@@ -0,0 +1,150 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \brief Defines an implementation of Module-based Model Runtime Interface that works with
+ *        Ahead-of-Time compilation.
+ * \file aot_executor.h
+ */
+#ifndef TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_H_
+#define TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_H_
+
+#include <tvm/runtime/metadata.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/object.h>
+#include <tvm/runtime/packed_func.h>
+
+#include <string>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+
+class TVM_DLL AotExecutor : public ModuleNode {
+ public:
+  /*!
+   * \brief Implements member function lookup for this Module for the frontend.
+   * \param name The name of the function.
+   * \param sptr_to_self The pointer to the module node.
+   * \return The corresponding member function.
+   */
+  PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) override;
+
+  /*!
+   * \return The type key of the executor.
+   */
+  const char* type_key() const final { return "AotExecutor"; }
+
+  void Run();
+
+  /*!
+   * \brief Initialize the AOT executor with metadata, runtime::Module, and device.
+   * \param module The module containing the compiled functions for the host
+   *  processor.
+   * \param devs The device of the host and devices where graph nodes will be
+   *  executed on.
+   * \param lookup_linked_param_func If given, a PackedFunc invoked to lookup linked parameters
+   *  by storage_id. If not given, linked parameters are looked-up using an internal implementation,
+   *  which is not compatible with RPCModules. Default is nullptr.
+   */
+  AotExecutor(tvm::runtime::Module module, const std::vector<Device>& devs);
+
+  /*!
+   * \brief Get the input index given the name of input.
+   * \param name The name of the input.
+   * \return The index of input.
+   */
+  int GetInputIndex(const std::string& name);
+
+  /*!
+   * \brief Get the output index given the name of output.
+   * \param name The name of the output.
+   * \return The index of output.
+   */
+  int GetOutputIndex(const std::string& name);
+
+  /*!
+   * \brief set index-th input to the graph.
+   * \param index The input index.
+   * \param data_in The input data.
+   */
+  void SetInput(int index, DLTensor* data_in);
+  /*!
+   * \brief set index-th input to the graph without copying the data
+   * \param index The input index.
+   * \param data_ref The input data that is referred.
+   */
+  void SetInputZeroCopy(int index, DLTensor* data_ref);
+  /*!
+   * \brief set index-th output to the graph without copying the data.
+   * \param index The output index.
+   * \param data_ref The output data that is referred.
+   */
+  void SetOutputZeroCopy(int index, DLTensor* data_ref);
+  /*!
+   * \brief Get the number of outputs
+   *
+   * \return The number of outputs from graph.
+   */
+  int NumOutputs() const;
+  /*!
+   * \brief Get the number of inputs
+   *
+   * \return The number of inputs to the graph.
+   */
+  int NumInputs() const;
+  /*!
+   * \brief Return NDArray for given input index.
+   * \param index The input index.
+   *
+   * \return NDArray corresponding to given input node index.
+   */
+  NDArray GetInput(int index) const;
+  /*!
+   * \brief Return NDArray for given output index.
+   * \param index The output index.
+   *
+   * \return NDArray corresponding to given output node index.
+   */
+  NDArray GetOutput(int index) const;
+  /*!
+   * \brief Copy index-th output to data_out.
+   * \param index The output index.
+   * \param data_out the output data.
+   */
+  void CopyOutputTo(int index, DLTensor* data_out);
+
+ private:
+  /*! \brief Metadata provided to the runtime from the compiler. */
+  metadata::Metadata metadata_;
+
+  /*! \brief Runtime module which contains the AOT top-level function. */
+  Module module_;
+
+  /*! \brief The devices which should be used to execute the computations. */
+  std::vector<Device> devices_;
+
+  /*! \brief Holds one NDArray per function argument in the same order. */
+  std::vector<NDArray> args_;
+};
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_H_
diff --git a/src/runtime/aot_executor/aot_executor_factory.cc b/src/runtime/aot_executor/aot_executor_factory.cc
new file mode 100644
index 000000000000..4cb3026991fe
--- /dev/null
+++ b/src/runtime/aot_executor/aot_executor_factory.cc
@@ -0,0 +1,132 @@
+/*
+ * 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 aot_executor_factory.cc
+ * \brief Graph executor factory implementations
+ */
+
+#include "./aot_executor_factory.h"
+
+#include <tvm/runtime/container/string.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/registry.h>
+
+#include <iterator>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+
+AotExecutorFactory::AotExecutorFactory(
+    const std::unordered_map<std::string, tvm::runtime::NDArray>& params,
+    const std::string& module_name) {
+  params_ = params;
+  module_name_ = module_name;
+}
+
+PackedFunc AotExecutorFactory::GetFunction(
+    const std::string& name, const tvm::runtime::ObjectPtr<tvm::runtime::Object>& sptr_to_self) {
+  if (name == module_name_) {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      ICHECK_GT(args.num_args, 0) << "Must supply at least one device argument";
+      std::vector<Device> devices;
+      for (int i = 0; i < args.num_args; ++i) {
+        devices.emplace_back(args[i].operator Device());
+      }
+      *rv = this->ExecutorCreate(devices);
+    });
+  } else if (name == "remove_params") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      std::unordered_map<std::string, tvm::runtime::NDArray> empty_params{};
+      auto exec = make_object<AotExecutorFactory>(empty_params, this->module_name_);
+      exec->Import(this->imports_[0]);
+      *rv = Module(exec);
+    });
+  } else {
+    return PackedFunc();
+  }
+}
+
+void AotExecutorFactory::SaveToBinary(dmlc::Stream* stream) {
+  std::vector<std::string> names;
+  std::vector<DLTensor*> arrays;
+  for (const auto& v : params_) {
+    names.emplace_back(v.first);
+    arrays.emplace_back(const_cast<DLTensor*>(v.second.operator->()));
+  }
+  uint64_t sz = arrays.size();
+  ICHECK(sz == names.size());
+  stream->Write(sz);
+  stream->Write(names);
+  for (size_t i = 0; i < sz; ++i) {
+    tvm::runtime::SaveDLTensor(stream, arrays[i]);
+  }
+  stream->Write(module_name_);
+}
+
+Module AotExecutorFactory::ExecutorCreate(const std::vector<Device>& devs) {
+  auto exec = make_object<AotExecutor>(this->imports_[0], devs);
+  // set params
+  SetParams(exec.get(), this->params_);
+  return Module(exec);
+}
+
+Module AotExecutorFactoryModuleLoadBinary(void* strm) {
+  dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
+  std::unordered_map<std::string, tvm::runtime::NDArray> params;
+  std::string module_name;
+  uint64_t sz;
+  ICHECK(stream->Read(&sz));
+  std::vector<std::string> names;
+  ICHECK(stream->Read(&names));
+  ICHECK(sz == names.size());
+  for (size_t i = 0; i < sz; ++i) {
+    tvm::runtime::NDArray temp;
+    temp.Load(stream);
+    params[names[i]] = temp;
+  }
+  ICHECK(stream->Read(&module_name));
+  auto exec = make_object<AotExecutorFactory>(params, module_name);
+  return Module(exec);
+}
+
+TVM_REGISTER_GLOBAL("tvm.aot_executor_factory.create").set_body([](TVMArgs args, TVMRetValue* rv) {
+  ICHECK_GE(args.num_args, 2) << "The expected number of arguments for "
+                                 "aot_executor_factory.create needs at least 2, "
+                                 "but it has "
+                              << args.num_args;
+  // The argument order is module, module_name, param0_name, param0_tensor,
+  // [param1_name, param1_tensor], ...
+  ICHECK_EQ((args.size() - 2) % 2, 0);
+  std::unordered_map<std::string, tvm::runtime::NDArray> params;
+  for (size_t i = 2; i < static_cast<size_t>(args.size()); i += 2) {
+    std::string name = args[i].operator String();
+    params[name] = args[i + 1].operator tvm::runtime::NDArray();
+  }
+  auto exec = make_object<AotExecutorFactory>(params, args[1]);
+  exec->Import(args[0]);
+  *rv = Module(exec);
+});
+
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_AotExecutorFactory")
+    .set_body_typed(AotExecutorFactoryModuleLoadBinary);
+
+}  // namespace runtime
+}  // namespace tvm
diff --git a/src/runtime/aot_executor/aot_executor_factory.h b/src/runtime/aot_executor/aot_executor_factory.h
new file mode 100644
index 000000000000..1d6a0a62776e
--- /dev/null
+++ b/src/runtime/aot_executor/aot_executor_factory.h
@@ -0,0 +1,119 @@
+/*
+ * 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 tvm/runtime/aot_executor/aot_executor_factory.h
+ * \brief Aot executor factory creating aot executor.
+ */
+
+#ifndef TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_FACTORY_H_
+#define TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_FACTORY_H_
+
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/packed_func.h>
+
+#include <algorithm>
+#include <functional>
+#include <numeric>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "./aot_executor.h"
+
+namespace tvm {
+namespace runtime {
+
+class TVM_DLL AotExecutorFactory : public runtime::ModuleNode {
+ public:
+  /*!
+   * \brief Construct the AotExecutorFactory.
+   * \param params The params of aot.
+   * \param module_name The module name of aot.
+   */
+  AotExecutorFactory(const std::unordered_map<std::string, tvm::runtime::NDArray>& params,
+                     const std::string& module_name);
+
+  /*!
+   * \brief Get member function to front-end
+   * \param name The name of the function.
+   * \param sptr_to_self The pointer to the module node.
+   * \return The corresponding member function.
+   */
+  PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
+
+  /*!
+   * \return The type key of the executor.
+   */
+  const char* type_key() const override { return "AotExecutorFactory"; }
+
+  /*!
+   * \brief Save the module to binary stream.
+   * \param stream The binary stream to save to.
+   */
+  void SaveToBinary(dmlc::Stream* stream) override;
+
+  /*!
+   * \brief Create a specific executor module
+   * \param devs The device of the host and devices where the model will be
+   *  executed.
+   * \return created executor module
+   */
+  Module ExecutorCreate(const std::vector<Device>& devs);
+
+  /*!
+   * \brief Set params.
+   * \param aot_executor The aot executor we want to set the params into.
+   * \param params The aot params value we want to set.
+   */
+  void SetParams(AotExecutor* aot_executor,
+                 const std::unordered_map<std::string, tvm::runtime::NDArray>& params) const {
+    std::unordered_map<std::string, tvm::runtime::NDArray> value = params;
+    // upload big arrays first to avoid memory issue in rpc mode
+    std::vector<std::string> keys;
+    for (const auto& p : value) {
+      keys.emplace_back(p.first);
+    }
+    std::sort(std::begin(keys), std::end(keys),
+              [&](const std::string& lhs, const std::string& rhs) -> bool {
+                auto lhs_size = GetDataSize(*value[lhs].operator->());
+                auto rhs_size = GetDataSize(*value[rhs].operator->());
+                return lhs_size > rhs_size;
+              });
+    for (const auto& key : keys) {
+      int in_idx = aot_executor->GetInputIndex(key);
+      if (in_idx >= 0) {
+        aot_executor->SetInput(in_idx, const_cast<DLTensor*>(value[key].operator->()));
+      }
+    }
+  }
+
+ protected:
+  /*! \brief The params. */
+  std::unordered_map<std::string, tvm::runtime::NDArray> params_;
+  /*! \brief module name */
+  std::string module_name_;
+};
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_AOT_EXECUTOR_AOT_EXECUTOR_FACTORY_H_
diff --git a/src/runtime/metadata_module.cc b/src/runtime/const_loader_module.cc
similarity index 67%
rename from src/runtime/metadata_module.cc
rename to src/runtime/const_loader_module.cc
index 7cb986bba62c..5496e161e57f 100644
--- a/src/runtime/metadata_module.cc
+++ b/src/runtime/const_loader_module.cc
@@ -18,13 +18,13 @@
  */
 
 /*!
- * \file src/runtime/metadata_module.cc
- * \brief A wrapper for initializing imported modules using metadata. This
+ * \file src/runtime/const_loader_module.cc
+ * \brief A wrapper for initializing imported modules using constant NDArray. This
  * module is intended to be used by various runtime in the TVM stack, i.e.
  * graph executor, relay VM, AOT runtime, and various user defined runtimes. It
  * paves the way to separate the code and metedata, which makes compilation
  * and/or interpretation more convenient. In addition, the clear separation of
- * code and metadata significantly reduces the efforts for handling external
+ * code and constants significantly reduces the efforts for handling external
  * codegen and runtimes.
  */
 #include <tvm/runtime/container/array.h>
@@ -42,18 +42,19 @@ namespace tvm {
 namespace runtime {
 
 /*!
- * \brief The metadata module is designed to manage initialization of the
- * imported submodules.
+ * \brief The const-loader module is designed to manage initialization of the
+ * imported submodules for the C++ runtime.
  */
-class MetadataModuleNode : public ModuleNode {
+class ConstLoaderModuleNode : public ModuleNode {
  public:
-  MetadataModuleNode(const std::unordered_map<std::string, NDArray>& metadata,
-                     const std::unordered_map<std::string, std::vector<std::string>>& sym_vars)
-      : metadata_(metadata), sym_vars_(sym_vars) {
+  ConstLoaderModuleNode(
+      const std::unordered_map<std::string, NDArray>& const_var_ndarray,
+      const std::unordered_map<std::string, std::vector<std::string>>& const_vars_by_symbol)
+      : const_var_ndarray_(const_var_ndarray), const_vars_by_symbol_(const_vars_by_symbol) {
     // Only the related submodules are cached to reduce the number of runtime
     // symbol lookup for initialization. Otherwise, symbols/primitives in the
     // DSO module will also be cached but they never need to be initialized.
-    for (const auto& it : sym_vars_) {
+    for (const auto& it : const_vars_by_symbol_) {
       initialized_[it.first] = false;
     }
   }
@@ -78,20 +79,20 @@ class MetadataModuleNode : public ModuleNode {
     return PackedFunc(nullptr);
   }
 
-  const char* type_key() const { return "metadata"; }
+  const char* type_key() const { return "const_loader"; }
 
   /*!
-   * \brief Get the list of metadata that is required by the given module.
+   * \brief Get the list of constants that is required by the given module.
    * \param symbol The symbol that is being queried.
    * \return The list of needed NDArray.
    */
-  Array<NDArray> GetRequiredMetadata(const std::string& symbol) {
+  Array<NDArray> GetRequiredConstants(const std::string& symbol) {
     Array<NDArray> ret;
-    ICHECK_GT(sym_vars_.count(symbol), 0U) << "No symbol is recorded for " << symbol;
-    std::vector<std::string> vars = sym_vars_[symbol];
+    ICHECK_GT(const_vars_by_symbol_.count(symbol), 0U) << "No symbol is recorded for " << symbol;
+    std::vector<std::string> vars = const_vars_by_symbol_[symbol];
     for (const auto& it : vars) {
-      ICHECK_GT(metadata_.count(it), 0U) << "Found not recorded constant variable: " << it;
-      ret.push_back(metadata_[it]);
+      ICHECK_GT(const_var_ndarray_.count(it), 0U) << "Found not recorded constant variable: " << it;
+      ret.push_back(const_var_ndarray_[it]);
     }
     return ret;
   }
@@ -102,12 +103,12 @@ class MetadataModuleNode : public ModuleNode {
    * for runtime lookup.
    *
    * \note  A module could be like the following:
-   *  MetadataModuleNode (contains all the metadata)
+   *  ConstLoaderModuleNode (contains all the constants)
    *    - CSourceModule
    *    - JSON runtime module
    *
    *  The initializer iterates through the imported modules and intilizes the
-   *  found module accordingly by passing the needed metadata into it.
+   *  found module accordingly by passing the needed constants into it.
    */
   void InitSubModule(const std::string& symbol) {
     PackedFunc init(nullptr);
@@ -116,8 +117,8 @@ class MetadataModuleNode : public ModuleNode {
       std::string init_name = "__init_" + symbol;
       init = it.GetFunction(init_name, false);
       if (init != nullptr) {
-        auto md = GetRequiredMetadata(symbol);
-        // Initialize the module with metadata.
+        auto md = GetRequiredConstants(symbol);
+        // Initialize the module with constants.
         int ret = init(md);
         // Report the error if initialization is failed.
         ICHECK_EQ(ret, 0) << TVMGetLastError();
@@ -128,32 +129,32 @@ class MetadataModuleNode : public ModuleNode {
 
   void SaveToBinary(dmlc::Stream* stream) final {
     std::vector<std::string> variables;
-    std::vector<NDArray> metadata;
-    for (const auto& it : metadata_) {
+    std::vector<NDArray> const_var_ndarray;
+    for (const auto& it : const_var_ndarray_) {
       String var_name = it.first;
       variables.push_back(var_name);
-      metadata.push_back(it.second);
+      const_var_ndarray.push_back(it.second);
     }
 
     // Save all variables in the function.
     stream->Write(variables);
     // Save all constant data.
-    uint64_t sz = static_cast<uint64_t>(metadata.size());
+    uint64_t sz = static_cast<uint64_t>(const_var_ndarray.size());
     stream->Write(sz);
     for (uint64_t i = 0; i < sz; i++) {
-      metadata[i].Save(stream);
+      const_var_ndarray[i].Save(stream);
     }
 
     // Save the symbol to list of required constant variables mapping
     std::vector<std::string> symbols;
     std::vector<std::vector<std::string>> const_vars;
-    for (const auto& it : sym_vars_) {
+    for (const auto& it : const_vars_by_symbol_) {
       symbols.push_back(it.first);
       const_vars.push_back(it.second);
     }
 
     stream->Write(symbols);
-    sz = static_cast<uint64_t>(sym_vars_.size());
+    sz = static_cast<uint64_t>(const_vars_by_symbol_.size());
     stream->Write(sz);
     for (uint64_t i = 0; i < sz; i++) {
       stream->Write(const_vars[i]);
@@ -165,9 +166,9 @@ class MetadataModuleNode : public ModuleNode {
 
     // Load the variables.
     std::vector<std::string> variables;
-    ICHECK(stream->Read(&variables)) << "Loading variables failed";
+    ICHECK(stream->Read(&variables)) << "Loading variable names failed";
     uint64_t sz;
-    ICHECK(stream->Read(&sz, sizeof(sz))) << "Loading metadata size failed";
+    ICHECK(stream->Read(&sz, sizeof(sz))) << "Loading number of vars failed";
     ICHECK_EQ(static_cast<size_t>(sz), variables.size())
         << "The number of variables and ndarray counts must match";
     // Load the list of ndarray.
@@ -178,10 +179,10 @@ class MetadataModuleNode : public ModuleNode {
       arrays.push_back(temp);
     }
 
-    std::unordered_map<std::string, NDArray> metadata;
+    std::unordered_map<std::string, NDArray> const_var_ndarray;
     for (uint64_t i = 0; i < sz; i++) {
-      ICHECK_EQ(metadata.count(variables[i]), 0U);
-      metadata[variables[i]] = arrays[i];
+      ICHECK_EQ(const_var_ndarray.count(variables[i]), 0U);
+      const_var_ndarray[variables[i]] = arrays[i];
     }
 
     // Load the symbol to list of required constant variables mapping
@@ -196,12 +197,12 @@ class MetadataModuleNode : public ModuleNode {
       const_vars.push_back(vars);
     }
 
-    std::unordered_map<std::string, std::vector<std::string>> sym_vars;
+    std::unordered_map<std::string, std::vector<std::string>> const_vars_by_symbol;
     for (uint64_t i = 0; i < sz; i++) {
-      sym_vars[symbols[i]] = const_vars[i];
+      const_vars_by_symbol[symbols[i]] = const_vars[i];
     }
 
-    auto n = make_object<MetadataModuleNode>(metadata, sym_vars);
+    auto n = make_object<ConstLoaderModuleNode>(const_var_ndarray, const_vars_by_symbol);
     return Module(n);
   }
 
@@ -212,19 +213,21 @@ class MetadataModuleNode : public ModuleNode {
    */
   std::unordered_map<std::string, bool> initialized_;
   /*! \brief Variable name to NDArray mapping. */
-  std::unordered_map<std::string, NDArray> metadata_;
+  std::unordered_map<std::string, NDArray> const_var_ndarray_;
   /*! \brief Symbol name to required constant variables mapping. */
-  std::unordered_map<std::string, std::vector<std::string>> sym_vars_;
+  std::unordered_map<std::string, std::vector<std::string>> const_vars_by_symbol_;
 };
 
-Module MetadataModuleCreate(
-    const std::unordered_map<std::string, NDArray>& metadata,
-    const std::unordered_map<std::string, std::vector<std::string>>& sym_vars) {
-  auto n = make_object<MetadataModuleNode>(metadata, sym_vars);
+Module ConstLoaderModuleCreate(
+    const std::unordered_map<std::string, NDArray>& const_var_ndarray,
+    const std::unordered_map<std::string, std::vector<std::string>>& const_vars_by_symbol) {
+  auto n = make_object<ConstLoaderModuleNode>(const_var_ndarray, const_vars_by_symbol);
   return Module(n);
 }
 
 TVM_REGISTER_GLOBAL("runtime.module.loadbinary_metadata")
-    .set_body_typed(MetadataModuleNode::LoadFromBinary);
+    .set_body_typed(ConstLoaderModuleNode::LoadFromBinary);
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_const_loader")
+    .set_body_typed(ConstLoaderModuleNode::LoadFromBinary);
 }  // namespace runtime
 }  // namespace tvm
diff --git a/src/runtime/const_loader_module.h b/src/runtime/const_loader_module.h
new file mode 100644
index 000000000000..eb548dfcf370
--- /dev/null
+++ b/src/runtime/const_loader_module.h
@@ -0,0 +1,53 @@
+/*
+ * 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 const_loader_module.h
+ * \brief Defines an interface to use the ConstLoaderModule.
+ */
+
+#ifndef TVM_RUNTIME_CONST_LOADER_MODULE_H_
+#define TVM_RUNTIME_CONST_LOADER_MODULE_H_
+
+#include <tvm/runtime/ndarray.h>
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace tvm {
+namespace runtime {
+
+/*!
+ * \brief Create a ConstLoader module object.
+ *
+ * \param const_var_ndarray Maps consts var name to NDArray containing data for the var.
+ * \param const_vars_by_symbol Maps the name of a module init function to a list of names of
+ * const vars whose data will be passed to that init function.
+ *
+ * \return The created ConstLoaderModule.
+ */
+Module ConstLoaderModuleCreate(
+    const std::unordered_map<std::string, NDArray>& const_var_ndarray,
+    const std::unordered_map<std::string, std::vector<std::string>>& const_vars_by_symbol);
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_CONST_LOADER_MODULE_H_
diff --git a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
index f9f1961e2697..6d5e5543cd40 100644
--- a/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
+++ b/src/runtime/contrib/dnnl/dnnl_json_runtime.cc
@@ -26,6 +26,7 @@
 #include <tvm/runtime/registry.h>
 
 #include <cstddef>
+#include <regex>
 #include <string>
 #include <vector>
 
@@ -90,44 +91,95 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
 
  private:
   // Build up the engine based on the input graph.
+  std::map<std::string, tag> layout_dict{
+      {"NCW", tag::ncw},       {"OIW", tag::oiw},     {"GOIW", tag::goiw},   {"NCHW", tag::nchw},
+      {"OIHW", tag::oihw},     {"GOIHW", tag::goihw}, {"NCDHW", tag::ncdhw}, {"OIDHW", tag::oidhw},
+      {"GOIDHW", tag::goidhw}, {"IOHW", tag::iohw},   {"GIOHW", tag::giohw}, {"IODHW", tag::iodhw},
+      {"GIODHW", tag::giodhw},
+  };
+
+  std::map<std::string, dnnl::algorithm> elt_name2algo{
+      {"abs", dnnl::algorithm::eltwise_abs},
+      {"exp", dnnl::algorithm::eltwise_exp},
+      {"log", dnnl::algorithm::eltwise_log},
+      {"sqrt", dnnl::algorithm::eltwise_sqrt},
+      {"round", dnnl::algorithm::eltwise_round},
+      {"logsumexp", dnnl::algorithm::eltwise_logsigmoid},
+      {"nn.relu", dnnl::algorithm::eltwise_relu},
+      {"nn.leaky_relu", dnnl::algorithm::eltwise_relu},
+      {"tanh", dnnl::algorithm::eltwise_tanh},
+      {"sigmoid", dnnl::algorithm::eltwise_logistic},
+      {"clip", dnnl::algorithm::eltwise_clip},
+  };
+
+  bool ParsingOpName(const std::string op_name, dnnl::primitive_attr attr) {
+    // Define RegExp.
+    std::regex bias_add_pat(".*_bias.*");
+    std::regex relu_pat(".*_relu.*");
+    std::regex tanh_pat(".*_tanh.*");
+    std::regex sigmoid_pat(".*_sigmoid.*");
+
+    // Parsing post-ops.
+    dnnl::post_ops ops;
+    if (std::regex_match(op_name, relu_pat)) {
+      ops.append_eltwise(1.f, dnnl::algorithm::eltwise_relu, 0.f, 0.f);
+    }
+    if (std::regex_match(op_name, tanh_pat)) {
+      ops.append_eltwise(1.f, dnnl::algorithm::eltwise_tanh, 0.f, 0.f);
+    }
+    if (std::regex_match(op_name, sigmoid_pat)) {
+      ops.append_eltwise(1.f, dnnl::algorithm::eltwise_logistic, 0.f, 0.f);
+    }
+    attr.set_post_ops(ops);
+
+    // Parsing bias_add.
+    return std::regex_match(op_name, bias_add_pat) ? true : false;
+  }
+
+  dnnl::memory::dims TransformStr2Dims(std::vector<std::string> strs, std::string str_name) {
+    dnnl::memory::dims out_dims;
+    if (str_name == "dilates") {
+      std::transform(strs.begin(), strs.end(), std::back_inserter(out_dims),
+                     [](const std::string& str) { return std::stoi(str) - 1; });
+    } else {
+      std::transform(strs.begin(), strs.end(), std::back_inserter(out_dims),
+                     [](const std::string& str) { return std::stoi(str); });
+    }
+    return out_dims;
+  }
+
   void BuildEngine() {
     engine_ = dnnl::engine(dnnl::engine::kind::cpu, 0);
     stream_ = dnnl::stream(engine_);
 
+    std::regex conv_pat(".*conv[1-3]d.*");
+    std::regex conv_tranpose_pat(".*conv[1-3]d_transpose.*");
+    std::regex dense_pat(".*dense.*");
+    std::regex max_pool_pat(".*max_pool[1-3]d");
+    std::regex avg_pool_pat(".*avg_pool[1-3]d");
+
     // Build subgraph engine.
     for (size_t nid = 0; nid < nodes_.size(); ++nid) {
       const auto& node = nodes_[nid];
       if (node.GetOpType() == "kernel") {
         ICHECK_EQ(node.GetOpType(), "kernel");
         auto op_name = node.GetOpName();
-        if ("nn.conv2d" == op_name) {
-          Conv2d(nid);
-        } else if ("dnnl.conv2d_relu" == op_name) {
-          Conv2d(nid, true, false, dnnl::algorithm::eltwise_relu);
-        } else if ("dnnl.conv2d_tanh" == op_name) {
-          Conv2d(nid, true, false, dnnl::algorithm::eltwise_tanh);
-        } else if ("dnnl.conv2d_sigmoid" == op_name) {
-          Conv2d(nid, true, false, dnnl::algorithm::eltwise_logistic);
-        } else if ("dnnl.conv2d_bias" == op_name) {
-          Conv2d(nid, false, true);
-        } else if ("dnnl.conv2d_bias_relu" == op_name) {
-          Conv2d(nid, true, true, dnnl::algorithm::eltwise_relu);
-        } else if ("dnnl.conv2d_bias_tanh" == op_name) {
-          Conv2d(nid, true, true, dnnl::algorithm::eltwise_tanh);
-        } else if ("dnnl.conv2d_bias_sigmoid" == op_name) {
-          Conv2d(nid, true, true, dnnl::algorithm::eltwise_logistic);
-        } else if ("nn.dense" == op_name) {
+        if (std::regex_match(op_name, conv_tranpose_pat)) {
+          Deconvolution(nid);
+        } else if (std::regex_match(op_name, conv_pat)) {
+          Convolution(nid);
+        } else if (std::regex_match(op_name, dense_pat)) {
           Dense(nid);
-        } else if ("dnnl.dense_bias" == op_name) {
-          Dense(nid, true);
         } else if ("nn.batch_norm" == op_name) {
           BatchNorm(nid);
-        } else if ("nn.relu" == op_name) {
-          Eltwise(nid, dnnl::algorithm::eltwise_relu);
-        } else if ("tanh" == op_name) {
-          Eltwise(nid, dnnl::algorithm::eltwise_tanh);
-        } else if ("sigmoid" == op_name) {
-          Eltwise(nid, dnnl::algorithm::eltwise_logistic);
+        } else if (std::regex_match(op_name, max_pool_pat)) {
+          Pooling(nid, dnnl::algorithm::pooling_max);
+        } else if (std::regex_match(op_name, avg_pool_pat)) {
+          Pooling(nid, dnnl::algorithm::pooling_avg);
+        } else if (elt_name2algo.count(op_name)) {
+          Eltwise(nid);
+        } else if ("nn.softmax" == op_name) {
+          Softmax(nid);
         } else if ("add" == op_name) {
           Binary(nid, dnnl::algorithm::binary_add);
         } else if ("multiply" == op_name) {
@@ -166,73 +218,73 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
     return entry_out_mem_[eid].first;
   }
 
-  void Conv2d(const size_t& nid, const bool has_elt = false, const bool has_bias = false,
-              dnnl::algorithm algo = dnnl::algorithm::eltwise_relu) {
+  void Convolution(const size_t& nid) {
     auto node = nodes_[nid];
+    auto op_name = node.GetOpName();
+    dnnl::primitive_attr attr;
+    bool has_bias = ParsingOpName(op_name, attr);
 
     // Setup attributes.
     auto data_entry = node.GetInputs()[0];
     auto weight_entry = node.GetInputs()[1];
+    JSONGraphNodeEntry out_entry(nid, 0);
     dnnl::memory::dims input_shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
     dnnl::memory::dims weight_shape = nodes_[weight_entry.id_].GetOpShape()[weight_entry.index_];
+    dnnl::memory::dims out_shape = nodes_[out_entry.id_].GetOpShape()[out_entry.index_];
+    dnnl::memory::dim channels =
+        node.GetAttr<std::vector<std::string>>("channels")[0] != ""
+            ? std::stoi(node.GetAttr<std::vector<std::string>>("channels")[0])
+            : out_shape[1];
     std::vector<std::string> str_strides = node.GetAttr<std::vector<std::string>>("strides");
     std::vector<std::string> str_dilates = node.GetAttr<std::vector<std::string>>("dilation");
     std::vector<std::string> str_padding = node.GetAttr<std::vector<std::string>>("padding");
+    std::vector<std::string> str_padding_l(str_padding.begin(),
+                                           str_padding.begin() + str_padding.size() / 2);
+    std::vector<std::string> str_padding_r(str_padding.end() - str_padding.size() / 2,
+                                           str_padding.end());
     dnnl::memory::dim groups = std::stoi(node.GetAttr<std::vector<std::string>>("groups")[0]);
+    std::string data_layout = node.GetAttr<std::vector<std::string>>("data_layout")[0];
+    std::string kernel_layout = node.GetAttr<std::vector<std::string>>("kernel_layout")[0];
 
-    dnnl::memory::dim N = input_shape[0],        // batch size
-        IC = input_shape[1],                     // input channels
-        IH = input_shape[2],                     // input height
-        IW = input_shape[3],                     // input width
-        OC = weight_shape[0],                    // output channels
-        KH = weight_shape[2],                    // weight height
-        KW = weight_shape[3],                    // weight width
-        PW_L = std::stoi(str_padding[1]),        // width padding: left
-        PW_R = std::stoi(str_padding[3]),        // width padding: right
-        PH_L = std::stoi(str_padding[0]),        // height padding: top
-        PH_R = std::stoi(str_padding[2]),        // height padding: bottom
-        SH = std::stoi(str_strides[0]),          // height-wise stride
-        SW = std::stoi(str_strides[1]),          // weight-wise stride
-        DH = std::stoi(str_dilates[0]) - 1,      // height-wise dilate
-        DW = std::stoi(str_dilates[1]) - 1,      // weight-wise dilate
-        DKH = 1 + (KH - 1) * (DH + 1),           // dilated weight height
-        DKW = 1 + (KW - 1) * (DW + 1),           // dilated weight width
-        OH = (IH - DKH + PH_L + PH_R) / SH + 1,  // output height
-        OW = (IW - DKW + PW_L + PW_R) / SW + 1;  // output width
+    // Check layout.
+    if (layout_dict.find(data_layout) == layout_dict.end() ||
+        layout_dict.find(kernel_layout) == layout_dict.end()) {
+      LOG(FATAL) << "Unsupported layout for conv: " << data_layout << " " << kernel_layout;
+    }
 
     // Memory shapes.
-    dnnl::memory::dims src_dims = {N, IC, IH, IW};
-    dnnl::memory::dims weights_dims = {OC, IC, KH, KW};
+    dnnl::memory::dims src_dims = input_shape;       // {N, IC, ID, IH, IW}
+    dnnl::memory::dims weights_dims = weight_shape;  // {OC, IC, KD, KH, KW}
     if (groups > 1) {
-      weights_dims = {groups, 1, IC / groups, KH, KW};
+      weights_dims = {groups, channels / groups, input_shape[1] / groups};
+      weights_dims.insert(weights_dims.end(), weight_shape.begin() + 2, weight_shape.end());
+      kernel_layout.insert(0, "G");
     }
-    dnnl::memory::dims bias_dims = {OC};
-    dnnl::memory::dims dst_dims = {N, OC, OH, OW};
-    dnnl::memory::dims strides_dims = {SH, SW};
-    dnnl::memory::dims dilates_dims = {DH, DW};
-    dnnl::memory::dims padding_dims_l = {PH_L, PW_L};
-    dnnl::memory::dims padding_dims_r = {PH_R, PW_R};
+    dnnl::memory::dims bias_dims = {channels};
+    dnnl::memory::dims dst_dims = out_shape;  // {N, OC, OD, OH, OW}
+    dnnl::memory::dims strides_dims = TransformStr2Dims(str_strides, "strides");
+    dnnl::memory::dims dilates_dims = TransformStr2Dims(str_dilates, "dilates");
+    dnnl::memory::dims padding_dims_l = TransformStr2Dims(str_padding_l, "padding");
+    dnnl::memory::dims padding_dims_r = TransformStr2Dims(str_padding_r, "padding");
 
     // Memory descriptions.
-    auto conv_src_md = dnnl::memory::desc(src_dims, dt::f32, tag::any);
-    auto conv_weights_md = dnnl::memory::desc(weights_dims, dt::f32, tag::any);
+    auto conv_src_md = dnnl::memory::desc(src_dims, dt::f32, layout_dict[data_layout]);
+    auto conv_weights_md = dnnl::memory::desc(weights_dims, dt::f32, layout_dict[kernel_layout]);
     auto conv_bias_md = dnnl::memory::desc(bias_dims, dt::f32, tag::any);
-    auto conv_dst_md = dnnl::memory::desc(dst_dims, dt::f32, tag::nchw);
+    auto conv_dst_md = dnnl::memory::desc(dst_dims, dt::f32, layout_dict[data_layout]);
 
     // Covn2d description.
-    auto conv_desc = dnnl::convolution_forward::desc(
-        dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_direct, conv_src_md,
-        conv_weights_md, conv_bias_md, conv_dst_md, strides_dims, dilates_dims, padding_dims_l,
-        padding_dims_r);
-
-    // Enable elementwise post-ops
-    dnnl::primitive_attr attr;
-    if (has_elt) {
-      dnnl::post_ops ops;
-      ops.append_eltwise(1.f, algo, 0.f, 0.f);
-      attr.set_post_ops(ops);
-    }
-
+    auto conv_desc =
+        has_bias ? dnnl::convolution_forward::desc(
+                       dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_direct,
+                       conv_src_md, conv_weights_md, conv_bias_md, conv_dst_md, strides_dims,
+                       dilates_dims, padding_dims_l, padding_dims_r)
+                 : dnnl::convolution_forward::desc(dnnl::prop_kind::forward_inference,
+                                                   dnnl::algorithm::convolution_direct, conv_src_md,
+                                                   conv_weights_md, conv_dst_md, strides_dims,
+                                                   dilates_dims, padding_dims_l, padding_dims_r);
+
+    // Enable elementwise post-ops.
     auto conv2d_prim_desc = dnnl::convolution_forward::primitive_desc(conv_desc, attr, engine_);
 
     // Push to the network.
@@ -240,53 +292,163 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
     net_.push_back(conv);
 
     // Data memory.
-    ICHECK_EQ(node.GetAttr<std::vector<std::string>>("data_layout")[0], "NCHW");
-    auto conv2d_src_memory = BindDNNLMemory(data_entry, {src_dims, dt::f32, tag::nchw});
+    auto conv2d_src_memory = BindDNNLMemory(data_entry, conv_src_md);
 
     // Weight memory.
-    ICHECK_EQ(node.GetAttr<std::vector<std::string>>("kernel_layout")[0], "OIHW");
-    auto conv2d_weights_memory = BindDNNLMemory(
-        weight_entry, {weights_dims, dt::f32, (groups > 1) ? tag::goihw : tag::oihw});
+    auto conv2d_weights_memory = BindDNNLMemory(weight_entry, conv_weights_md);
+
+    // Output memory.
+    auto conv2d_dst_memory = BindDNNLMemory(out_entry, conv2d_prim_desc.dst_desc());
 
     // Bias memory.
     auto conv2d_bias_memory = dnnl::memory({bias_dims, dt::f32, tag::x}, engine_);
     if (has_bias) {
       auto bias_entry = node.GetInputs()[2];
       BindDNNLMemory(bias_entry, conv2d_bias_memory);
+
+      // Bind memory buffers.
+      net_args_.push_back({{DNNL_ARG_SRC, conv2d_src_memory},
+                           {DNNL_ARG_WEIGHTS, conv2d_weights_memory},
+                           {DNNL_ARG_BIAS, conv2d_bias_memory},
+                           {DNNL_ARG_DST, conv2d_dst_memory}});
     } else {
-      float bias[OC] = {0};
-      write_to_dnnl_memory(bias, conv2d_bias_memory, OC * sizeof(float));
+      // Bind memory buffers.
+      net_args_.push_back({{DNNL_ARG_SRC, conv2d_src_memory},
+                           {DNNL_ARG_WEIGHTS, conv2d_weights_memory},
+                           {DNNL_ARG_DST, conv2d_dst_memory}});
     }
+  }
 
-    // Output memory.
+  void Deconvolution(const size_t& nid) {
+    auto node = nodes_[nid];
+    auto op_name = node.GetOpName();
+    dnnl::primitive_attr attr;
+    bool has_bias = ParsingOpName(op_name, attr);
+
+    // Setup attributes.
+    auto data_entry = node.GetInputs()[0];
+    auto weight_entry = node.GetInputs()[1];
     JSONGraphNodeEntry out_entry(nid, 0);
-    auto conv2d_dst_memory = BindDNNLMemory(out_entry, conv2d_prim_desc.dst_desc());
+    dnnl::memory::dims input_shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
+    dnnl::memory::dims weight_shape = nodes_[weight_entry.id_].GetOpShape()[weight_entry.index_];
+    dnnl::memory::dims out_shape = nodes_[out_entry.id_].GetOpShape()[out_entry.index_];
+    dnnl::memory::dim channels =
+        node.GetAttr<std::vector<std::string>>("channels")[0] != ""
+            ? std::stoi(node.GetAttr<std::vector<std::string>>("channels")[0])
+            : out_shape[1];
+    std::vector<std::string> str_strides = node.GetAttr<std::vector<std::string>>("strides");
+    std::vector<std::string> str_dilates = node.GetAttr<std::vector<std::string>>("dilation");
+    std::vector<std::string> str_padding = node.GetAttr<std::vector<std::string>>("padding");
+    std::vector<std::string> str_padding_l(str_padding.begin(),
+                                           str_padding.begin() + str_padding.size() / 2);
+    std::vector<std::string> str_padding_r(str_padding.end() - str_padding.size() / 2,
+                                           str_padding.end());
+    dnnl::memory::dim groups = std::stoi(node.GetAttr<std::vector<std::string>>("groups")[0]);
+    std::string data_layout = node.GetAttr<std::vector<std::string>>("data_layout")[0];
+    std::string kernel_layout = node.GetAttr<std::vector<std::string>>("kernel_layout")[0];
 
-    // Bind memory buffers.
-    net_args_.push_back({{DNNL_ARG_SRC, conv2d_src_memory},
-                         {DNNL_ARG_WEIGHTS, conv2d_weights_memory},
-                         {DNNL_ARG_BIAS, conv2d_bias_memory},
-                         {DNNL_ARG_DST, conv2d_dst_memory}});
+    // Check layout.
+    if (layout_dict.find(data_layout) == layout_dict.end() ||
+        layout_dict.find(kernel_layout) == layout_dict.end()) {
+      LOG(FATAL) << "Unsupported layout: " << data_layout << " " << kernel_layout;
+    }
+
+    // Memory shapes.
+    dnnl::memory::dims src_dims = input_shape;       // {N, IC, ID, IH, IW}
+    dnnl::memory::dims weights_dims = weight_shape;  // {OC, IC, KD, KH, KW}
+
+    // Check weight shape, transform to `OIHW`
+    if (weights_dims[0] == src_dims[1] && weights_dims[1] == channels) {
+      std::swap(weights_dims[0], weights_dims[1]);
+    }
+    if (kernel_layout == "OIDHW") {
+      kernel_layout = "IODHW";
+    }
+    if (groups > 1) {
+      weights_dims = {groups, channels / groups, input_shape[1] / groups};
+      weights_dims.insert(weights_dims.end(), weight_shape.begin() + 2, weight_shape.end());
+      kernel_layout.insert(0, "G");
+    }
+    dnnl::memory::dims bias_dims = {channels};
+    dnnl::memory::dims dst_dims = out_shape;  // {N, OC, OD, OH, OW}
+    dnnl::memory::dims strides_dims = TransformStr2Dims(str_strides, "strides");
+    dnnl::memory::dims dilates_dims = TransformStr2Dims(str_dilates, "dilates");
+    dnnl::memory::dims padding_dims_l = TransformStr2Dims(str_padding_l, "padding");
+    dnnl::memory::dims padding_dims_r = TransformStr2Dims(str_padding_r, "padding");
+
+    // Memory descriptions.
+    auto deconv_src_md = dnnl::memory::desc(src_dims, dt::f32, layout_dict[data_layout]);
+    auto deconv_weights_md = dnnl::memory::desc(weights_dims, dt::f32, layout_dict[kernel_layout]);
+    auto deconv_bias_md = dnnl::memory::desc(bias_dims, dt::f32, tag::any);
+    auto deconv_dst_md = dnnl::memory::desc(dst_dims, dt::f32, layout_dict[data_layout]);
+
+    // Transposed covn2d description.
+    auto deconv_desc =
+        has_bias ? dnnl::deconvolution_forward::desc(
+                       dnnl::prop_kind::forward_inference, dnnl::algorithm::deconvolution_direct,
+                       deconv_src_md, deconv_weights_md, deconv_bias_md, deconv_dst_md,
+                       strides_dims, dilates_dims, padding_dims_l, padding_dims_r)
+                 : dnnl::deconvolution_forward::desc(
+                       dnnl::prop_kind::forward_inference, dnnl::algorithm::deconvolution_direct,
+                       deconv_src_md, deconv_weights_md, deconv_dst_md, strides_dims, dilates_dims,
+                       padding_dims_l, padding_dims_r);
+
+    // Enable elementwise post-ops.
+    auto deconv2d_prim_desc =
+        dnnl::deconvolution_forward::primitive_desc(deconv_desc, attr, engine_);
+
+    // Push to the network.
+    auto deconv = dnnl::deconvolution_forward(deconv2d_prim_desc);
+    net_.push_back(deconv);
+
+    // Data memory.
+    auto deconv2d_src_memory = BindDNNLMemory(data_entry, deconv_src_md);
+
+    // Weight memory.
+    auto deconv2d_weights_memory = BindDNNLMemory(weight_entry, deconv_weights_md);
+
+    // Output memory.
+    auto deconv2d_dst_memory = BindDNNLMemory(out_entry, deconv2d_prim_desc.dst_desc());
+
+    // Bias memory.
+    auto deconv2d_bias_memory = dnnl::memory({bias_dims, dt::f32, tag::x}, engine_);
+    if (has_bias) {
+      auto bias_entry = node.GetInputs()[2];
+      BindDNNLMemory(bias_entry, deconv2d_bias_memory);
+
+      // Bind memory buffers.
+      net_args_.push_back({{DNNL_ARG_SRC, deconv2d_src_memory},
+                           {DNNL_ARG_WEIGHTS, deconv2d_weights_memory},
+                           {DNNL_ARG_BIAS, deconv2d_bias_memory},
+                           {DNNL_ARG_DST, deconv2d_dst_memory}});
+    } else {
+      // Bind memory buffers.
+      net_args_.push_back({{DNNL_ARG_SRC, deconv2d_src_memory},
+                           {DNNL_ARG_WEIGHTS, deconv2d_weights_memory},
+                           {DNNL_ARG_DST, deconv2d_dst_memory}});
+    }
   }
 
-  void Dense(const size_t& nid, const bool has_bias = false) {
+  void Dense(const size_t& nid) {
     auto node = nodes_[nid];
+    auto op_name = node.GetOpName();
+    dnnl::primitive_attr attr;
+    bool has_bias = ParsingOpName(op_name, attr);
 
     // Setup attributes.
     auto data_entry = node.GetInputs()[0];
     auto weight_entry = node.GetInputs()[1];
+    JSONGraphNodeEntry out_entry(nid, 0);
     dnnl::memory::dims input_shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
     dnnl::memory::dims weight_shape = nodes_[weight_entry.id_].GetOpShape()[weight_entry.index_];
-
-    dnnl::memory::dim B = input_shape[0],  // batch size
-        IC = input_shape[1],               // input channels
-        OC = weight_shape[0];              // output channels
+    dnnl::memory::dims out_shape = nodes_[out_entry.id_].GetOpShape()[out_entry.index_];
+    dnnl::memory::dim OC = out_shape[1];
 
     // Memory shapes.
-    dnnl::memory::dims data_dims = {B, IC};
-    dnnl::memory::dims weight_dims = {OC, IC};
+    dnnl::memory::dims data_dims = input_shape;
+    dnnl::memory::dims weight_dims = weight_shape;
     dnnl::memory::dims bias_dims = {OC};
-    dnnl::memory::dims out_dims = {B, OC};
+    dnnl::memory::dims out_dims = out_shape;
 
     // Memory descriptions.
     auto data_md = dnnl::memory::desc({data_dims, dt::f32, tag::nc});
@@ -297,7 +459,9 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
     // Dense description.
     auto dense_desc = dnnl::inner_product_forward::desc(dnnl::prop_kind::forward_inference, data_md,
                                                         weight_md, bias_md, dst_md);
-    auto dense_prim_desc = dnnl::inner_product_forward::primitive_desc(dense_desc, engine_);
+
+    // Enable elementwise post-ops.
+    auto dense_prim_desc = dnnl::inner_product_forward::primitive_desc(dense_desc, attr, engine_);
 
     auto dense = dnnl::inner_product_forward(dense_prim_desc);
     net_.push_back(dense);
@@ -317,7 +481,6 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
     }
 
     // Output memory.
-    JSONGraphNodeEntry out_entry(nid, 0);
     auto dst_memory = BindDNNLMemory(out_entry, dense_prim_desc.dst_desc());
 
     net_args_.push_back({{DNNL_ARG_SRC, data_memory},
@@ -368,15 +531,85 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
                          {DNNL_ARG_VARIANCE, variance_memory}});
   }
 
-  void Eltwise(const size_t& nid, dnnl::algorithm algo) {
+  void Pooling(const size_t& nid, dnnl::algorithm algo) {
+    auto node = nodes_[nid];
+
+    // Setup attributes.
+    auto data_entry = node.GetInputs()[0];
+    JSONGraphNodeEntry out_entry(nid, 0);
+    dnnl::memory::dims input_shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
+    dnnl::memory::dims out_shape = nodes_[out_entry.id_].GetOpShape()[out_entry.index_];
+    std::vector<std::string> str_kernel = node.GetAttr<std::vector<std::string>>("pool_size");
+    std::vector<std::string> str_strides = node.GetAttr<std::vector<std::string>>("strides");
+    std::vector<std::string> str_padding = node.GetAttr<std::vector<std::string>>("padding");
+    std::vector<std::string> str_padding_l(str_padding.begin(),
+                                           str_padding.begin() + str_padding.size() / 2);
+    std::vector<std::string> str_padding_r(str_padding.end() - str_padding.size() / 2,
+                                           str_padding.end());
+    std::vector<std::string> str_dilates = node.GetAttr<std::vector<std::string>>("dilation");
+    std::string layout = node.GetAttr<std::vector<std::string>>("layout")[0];
+
+    // Check layout.
+    if (layout_dict.find(layout) == layout_dict.end()) {
+      LOG(FATAL) << "Unsupported layout for pooling: " << layout;
+    }
+
+    // Attributes related to AvgPool
+    if (algo == dnnl::algorithm::pooling_avg) {
+      int int_countpad = std::stoi(node.GetAttr<std::vector<std::string>>("count_include_pad")[0]);
+      bool count_include_pad = int_countpad != 0 ? true : false;
+      algo = count_include_pad ? dnnl::algorithm::pooling_avg_include_padding
+                               : dnnl::algorithm::pooling_avg_exclude_padding;
+    }
+
+    dnnl::memory::dims src_dims = input_shape;
+    dnnl::memory::dims dst_dims = out_shape;
+    dnnl::memory::dims kernel_dims = TransformStr2Dims(str_kernel, "kernel");
+    dnnl::memory::dims strides_dims = TransformStr2Dims(str_strides, "strides");
+    dnnl::memory::dims dilates_dims = TransformStr2Dims(str_dilates, "dilates");
+    dnnl::memory::dims padding_dims_l = TransformStr2Dims(str_padding_l, "padding");
+    dnnl::memory::dims padding_dims_r = TransformStr2Dims(str_padding_r, "padding");
+
+    // Memory descriptions.
+    auto pool_src_md = dnnl::memory::desc(src_dims, dt::f32, layout_dict[layout]);
+    auto pool_dst_md = dnnl::memory::desc(dst_dims, dt::f32, tag::any);
+
+    // Pooling description.
+    auto pool_desc = dnnl::pooling_forward::desc(dnnl::prop_kind::forward_inference, algo,
+                                                 pool_src_md, pool_dst_md, strides_dims,
+                                                 kernel_dims, padding_dims_l, padding_dims_r);
+
+    auto pool_prim_desc = dnnl::pooling_forward::primitive_desc(pool_desc, engine_, true);
+    auto pool = dnnl::pooling_forward(pool_prim_desc);
+    net_.push_back(pool);
+
+    // Memories.
+    auto pool2d_src_memory = BindDNNLMemory(data_entry, pool_src_md);
+
+    auto pool2d_dst_memory = BindDNNLMemory(out_entry, pool_prim_desc.dst_desc());
+
+    // Bind memory buffers.
+    net_args_.push_back({{DNNL_ARG_SRC, pool2d_src_memory}, {DNNL_ARG_DST, pool2d_dst_memory}});
+  }
+
+  void Eltwise(const size_t& nid) {
     auto node = nodes_[nid];
+    auto op_name = node.GetOpName();
+    auto algo = elt_name2algo[op_name];
 
     auto data_entry = node.GetInputs()[0];
     dnnl::memory::dims shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
     dnnl::memory::desc data_md = GenDNNLMemDescByShape(shape, dt::f32);
+    float alpha = 0., beta = 0.;
+    if (op_name == "clip") {
+      alpha = std::stof(node.GetAttr<std::vector<std::string>>("a_min")[0]);
+      beta = std::stof(node.GetAttr<std::vector<std::string>>("a_max")[0]);
+    } else if (op_name == "nn.leaky_relu") {
+      alpha = std::stof(node.GetAttr<std::vector<std::string>>("alpha")[0]);
+    }
 
     auto elt_desc =
-        dnnl::eltwise_forward::desc(dnnl::prop_kind::forward_inference, algo, data_md, 0);
+        dnnl::eltwise_forward::desc(dnnl::prop_kind::forward_inference, algo, data_md, alpha, beta);
     auto elt_prim_desc = dnnl::eltwise_forward::primitive_desc(elt_desc, engine_);
     ICHECK(data_md == elt_prim_desc.dst_desc());
 
@@ -390,6 +623,32 @@ class DNNLJSONRuntime : public JSONRuntimeBase {
     net_args_.push_back({{DNNL_ARG_SRC, data_memory}, {DNNL_ARG_DST, out_memory}});
   }
 
+  void Softmax(const size_t& nid) {
+    auto node = nodes_[nid];
+
+    auto data_entry = node.GetInputs()[0];
+    dnnl::memory::dims shape = nodes_[data_entry.id_].GetOpShape()[data_entry.index_];
+    int axis = std::stoi(node.GetAttr<std::vector<std::string>>("axis")[0]);
+    if (axis < 0) {
+      axis = shape.size() + axis;
+    }
+    dnnl::memory::desc data_md = GenDNNLMemDescByShape(shape, dt::f32);
+
+    auto softmax_desc =
+        dnnl::softmax_forward::desc(dnnl::prop_kind::forward_inference, data_md, axis);
+    auto softmax_prim_desc = dnnl::softmax_forward::primitive_desc(softmax_desc, engine_);
+    ICHECK(data_md == softmax_prim_desc.dst_desc());
+
+    auto softmax = dnnl::softmax_forward(softmax_prim_desc);
+    net_.push_back(softmax);
+
+    auto data_memory = BindDNNLMemory(data_entry, data_md);
+    JSONGraphNodeEntry out_entry(nid, 0);
+    auto out_memory = BindDNNLMemory(out_entry, data_md);
+
+    net_args_.push_back({{DNNL_ARG_SRC, data_memory}, {DNNL_ARG_DST, out_memory}});
+  }
+
   void Binary(const size_t& nid, dnnl::algorithm algo) {
     auto node = nodes_[nid];
 
diff --git a/src/runtime/crt/common/packed_func.c b/src/runtime/crt/common/packed_func.c
index e946cda9d9ae..645b22f3b255 100644
--- a/src/runtime/crt/common/packed_func.c
+++ b/src/runtime/crt/common/packed_func.c
@@ -49,6 +49,9 @@ DLDataType String2DLDataType(const char* s) {
   } else if (!strncmp(s, "float", 5)) {
     t.code = kDLFloat;
     scan = s + 5;
+  } else if (!strncmp(s, "bfloat", 6)) {
+    t.code = kDLBfloat;
+    scan = s + 6;
   } else if (!strncmp(s, "handle", 6)) {
     t.code = kTVMOpaqueHandle;
     t.bits = 64;  // handle uses 64 bit by default.
diff --git a/src/runtime/hexagon/android/hexagon_module.cc b/src/runtime/hexagon/android/hexagon_module.cc
index 3dc8fe9087c1..b8af3698ab9b 100644
--- a/src/runtime/hexagon/android/hexagon_module.cc
+++ b/src/runtime/hexagon/android/hexagon_module.cc
@@ -510,9 +510,12 @@ std::shared_ptr<Device> Device::Global() {
 
 }  // namespace hexagon
 
-TVM_REGISTER_GLOBAL("runtime.module.loadfile_hexagon").set_body([](TVMArgs args, TVMRetValue* rv) {
-  *rv = HexagonModuleLoadFile(args[0], args[1]);
-});
+// Disable this: it conflicts with loadfile_hexagon from hexagon_common.cc
+// This was only used with offload on Android, which is being deprecated.
+// TVM_REGISTER_GLOBAL("runtime.module.loadfile_hexagon").set_body([](TVMArgs args, TVMRetValue* rv)
+// {
+//   *rv = HexagonModuleLoadFile(args[0], args[1]);
+// });
 
 }  // namespace runtime
 }  // namespace tvm
diff --git a/src/runtime/hexagon/hexagon/hexagon_buffer.cc b/src/runtime/hexagon/hexagon/hexagon_buffer.cc
index 752c5325af21..e4654a349dca 100644
--- a/src/runtime/hexagon/hexagon/hexagon_buffer.cc
+++ b/src/runtime/hexagon/hexagon/hexagon_buffer.cc
@@ -37,7 +37,7 @@ namespace tvm {
 namespace runtime {
 namespace hexagon {
 
-int hexagon_user_dma_1d_sync(void* src, void* dst, uint32_t length);
+int hexagon_user_dma_1d_sync(void* dst, void* src, uint32_t length);
 
 struct Allocation {
   Allocation(size_t allocation_nbytes, size_t alignment)
diff --git a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc
index 03033bc33840..a4f6803f682c 100644
--- a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc
+++ b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc
@@ -40,6 +40,8 @@ namespace tvm {
 namespace runtime {
 namespace hexagon {
 
+int hexagon_user_dma_1d_sync(void* dst, void* src, uint32_t length);
+
 HexagonDeviceAPIv2* HexagonDeviceAPIv2::Global() {
   static auto* inst = new HexagonDeviceAPIv2();
   return inst;
@@ -90,7 +92,8 @@ void HexagonDeviceAPIv2::FreeDataSpace(Device dev, void* ptr) {
 
 // WorkSpace: runtime allocations for Hexagon
 struct HexagonWorkspacePool : public WorkspacePool {
-  HexagonWorkspacePool() : WorkspacePool(kDLCPU, HexagonDeviceAPIv2::Global()) {}
+  HexagonWorkspacePool()
+      : WorkspacePool(static_cast<DLDeviceType>(kDLHexagon), HexagonDeviceAPIv2::Global()) {}
 };
 
 void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) {
@@ -148,6 +151,16 @@ void HexagonDeviceAPIv2::CopyDataFromTo(const void* from, size_t from_offset, vo
   memcpy(static_cast<char*>(to) + to_offset, static_cast<const char*>(from) + from_offset, size);
 }
 
+TVM_REGISTER_GLOBAL("device_api.hexagon.mem_copy").set_body([](TVMArgs args, TVMRetValue* rv) {
+  void* dst = args[0];
+  void* src = args[1];
+  int size = args[2];
+
+  hexagon_user_dma_1d_sync(dst, src, size);
+
+  *rv = static_cast<int32_t>(0);
+});
+
 TVM_REGISTER_GLOBAL("device_api.hexagon.v2").set_body([](TVMArgs args, TVMRetValue* rv) {
   DeviceAPI* ptr = HexagonDeviceAPIv2::Global();
   *rv = static_cast<void*>(ptr);
diff --git a/src/runtime/hexagon/hexagon/hexagon_user_dma.cc b/src/runtime/hexagon/hexagon/hexagon_user_dma.cc
index 18da4af61090..b286c8b3165b 100644
--- a/src/runtime/hexagon/hexagon/hexagon_user_dma.cc
+++ b/src/runtime/hexagon/hexagon/hexagon_user_dma.cc
@@ -81,6 +81,7 @@ int hexagon_user_dma_1d_sync(void* dst, void* src, uint32_t length) {
     return DMA_FAILURE;
   }
 
+  dma_desc_set_state(dma_desc, DESC_STATE_READY);
   dma_desc_set_next(dma_desc, DMA_NULL_PTR);
   dma_desc_set_length(dma_desc, length);
   dma_desc_set_desctype(dma_desc, DESC_DESCTYPE_1D);
@@ -89,13 +90,13 @@ int hexagon_user_dma_1d_sync(void* dst, void* src, uint32_t length) {
   dma_desc_set_bypassdst(dma_desc, DESC_BYPASS_OFF);
   dma_desc_set_bypasssrc(dma_desc, DESC_BYPASS_OFF);
   dma_desc_set_order(dma_desc, DESC_ORDER_ORDER);
-  dma_desc_set_dstate(dma_desc, DESC_DSTATE_INCOMPLETE);
+  dma_desc_set_done(dma_desc, DESC_DONE_INCOMPLETE);
   dma_desc_set_src(dma_desc, src32);
   dma_desc_set_dst(dma_desc, dst32);
 
   dmstart(dma_desc);
   unsigned int status = dmwait() & DM0_STATUS_MASK;
-  unsigned int done = dma_desc_get_dstate(dma_desc);
+  unsigned int done = dma_desc_get_done(dma_desc);
 
 #ifdef _WIN32
   _aligned_free(dma_desc);
@@ -103,7 +104,7 @@ int hexagon_user_dma_1d_sync(void* dst, void* src, uint32_t length) {
   free(dma_desc);
 #endif
 
-  if (status == DM0_STATUS_IDLE && done == DESC_DSTATE_COMPLETE) {
+  if (status == DM0_STATUS_IDLE && done == DESC_DONE_COMPLETE) {
     return DMA_SUCCESS;
   }
   return DMA_FAILURE;
diff --git a/src/runtime/hexagon/hexagon/hexagon_user_dma_descriptors.h b/src/runtime/hexagon/hexagon/hexagon_user_dma_descriptors.h
index a90c25e71199..cea91310dd94 100644
--- a/src/runtime/hexagon/hexagon/hexagon_user_dma_descriptors.h
+++ b/src/runtime/hexagon/hexagon/hexagon_user_dma_descriptors.h
@@ -27,7 +27,11 @@ namespace hexagon {
 // NOTE: Using 2D descriptor size even for 1D descriptors
 #define DMA_DESC_2D_SIZE 32
 
-// TODO(Straw): DMA Status [0][3:0]?
+// DMA State
+// desc[0][3:0]
+#define DESC_STATE_MASK 0x0000000F
+#define DESC_STATE_SHIFT 0
+#define DESC_STATE_READY 0
 
 // desc[0][31:4]
 // Descriptors addresses must be (minimum) 16 byte aligned
@@ -72,10 +76,10 @@ namespace hexagon {
 #define DESC_ORDER_ORDER 1
 
 // desc[1][31]
-#define DESC_DSTATE_MASK 0x80000000
-#define DESC_DSTATE_SHIFT 31
-#define DESC_DSTATE_INCOMPLETE 0
-#define DESC_DSTATE_COMPLETE 1
+#define DESC_DONE_MASK 0x80000000
+#define DESC_DONE_SHIFT 31
+#define DESC_DONE_INCOMPLETE 0
+#define DESC_DONE_COMPLETE 1
 
 // desc[2]
 #define DESC_SRC_MASK 0xFFFFFFFF
@@ -130,8 +134,8 @@ namespace hexagon {
 /* 1D (linear) descriptor */
 /**************************/
 struct dma_desc_1d_t {
-  unsigned int next;
-  unsigned int dstate_order_bypass_comp_desctype_length;
+  unsigned int next_state;
+  unsigned int done_order_bypass_comp_desctype_length;
   unsigned int src;
   unsigned int dst;
 };
@@ -140,8 +144,8 @@ struct dma_desc_1d_t {
 /* 2D (box) descriptor */
 /***********************/
 struct dma_desc_2d_t {
-  unsigned int next;
-  unsigned int dstate_order_bypass_comp_desctype_length;
+  unsigned int next_state;
+  unsigned int done_order_bypass_comp_desctype_length;
   unsigned int src;
   unsigned int dst;
   unsigned int allocation_padding;
@@ -150,26 +154,33 @@ struct dma_desc_2d_t {
   unsigned int dstwidthoffset_srcwidthoffset;
 };
 
+// desc[0][3:0]
+inline void dma_desc_set_state(void* dma_desc_ptr, unsigned int v) {
+  dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
+  (dma_desc_1d_ptr->next_state) &= ~DESC_STATE_MASK;
+  (dma_desc_1d_ptr->next_state) |= ((v << DESC_STATE_SHIFT) & DESC_STATE_MASK);
+}
+
 // desc[0][31:4]
 inline void dma_desc_set_next(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->next) &= ~DESC_NEXT_MASK;
-  (dma_desc_1d_ptr->next) |= ((v << DESC_NEXT_SHIFT) & DESC_NEXT_MASK);
+  (dma_desc_1d_ptr->next_state) &= ~DESC_NEXT_MASK;
+  (dma_desc_1d_ptr->next_state) |= ((v << DESC_NEXT_SHIFT) & DESC_NEXT_MASK);
 }
 
 // desc[1][23:0]
 inline void dma_desc_set_length(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_LENGTH_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_LENGTH_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_LENGTH_SHIFT) & DESC_LENGTH_MASK);
 }
 
 // desc[1][25:24]
 inline void dma_desc_set_desctype(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_DESCTYPE_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_DESCTYPE_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_DESCTYPE_SHIFT) & DESC_DESCTYPE_MASK);
 }
 
@@ -177,8 +188,8 @@ inline void dma_desc_set_desctype(void* dma_desc_ptr, unsigned int v) {
 // desc[1][26]
 inline void dma_desc_set_dstcomp(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_DSTCOMP_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_DSTCOMP_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_DSTCOMP_SHIFT) & DESC_DSTCOMP_MASK);
 }
 
@@ -186,48 +197,48 @@ inline void dma_desc_set_dstcomp(void* dma_desc_ptr, unsigned int v) {
 // desc[1][27]
 inline void dma_desc_set_srccomp(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_SRCCOMP_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_SRCCOMP_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_SRCCOMP_SHIFT) & DESC_SRCCOMP_MASK);
 }
 
 // desc[1][28]
 inline void dma_desc_set_bypassdst(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_BYPASSDST_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_BYPASSDST_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_BYPASSDST_SHIFT) & DESC_BYPASSDST_MASK);
 }
 
 // desc[1][29]
 inline void dma_desc_set_bypasssrc(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_BYPASSSRC_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_BYPASSSRC_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_BYPASSSRC_SHIFT) & DESC_BYPASSSRC_MASK);
 }
 
 // desc[1][30]
 inline void dma_desc_set_order(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_ORDER_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_ORDER_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
       ((v << DESC_ORDER_SHIFT) & DESC_ORDER_MASK);
 }
 
 // desc[1][31]
-inline void dma_desc_set_dstate(void* dma_desc_ptr, unsigned int v) {
+inline void dma_desc_set_done(void* dma_desc_ptr, unsigned int v) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) &= ~DESC_DSTATE_MASK;
-  (dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) |=
-      ((v << DESC_DSTATE_SHIFT) & DESC_DSTATE_MASK);
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) &= ~DESC_DONE_MASK;
+  (dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) |=
+      ((v << DESC_DONE_SHIFT) & DESC_DONE_MASK);
 }
 
 // desc[1][31]
-inline unsigned int dma_desc_get_dstate(void* dma_desc_ptr) {
+inline unsigned int dma_desc_get_done(void* dma_desc_ptr) {
   dma_desc_1d_t* dma_desc_1d_ptr = reinterpret_cast<dma_desc_1d_t*>(dma_desc_ptr);
-  return (((dma_desc_1d_ptr->dstate_order_bypass_comp_desctype_length) & DESC_DSTATE_MASK) >>
-          DESC_DSTATE_SHIFT);
+  return (((dma_desc_1d_ptr->done_order_bypass_comp_desctype_length) & DESC_DONE_MASK) >>
+          DESC_DONE_SHIFT);
 }
 
 // desc[2]
diff --git a/src/runtime/hexagon/rpc/simulator/hexagon_sim_proto.h b/src/runtime/hexagon/rpc/simulator/hexagon_sim_proto.h
new file mode 100644
index 000000000000..fec9d7329dbe
--- /dev/null
+++ b/src/runtime/hexagon/rpc/simulator/hexagon_sim_proto.h
@@ -0,0 +1,75 @@
+/*
+ * 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.
+ */
+
+#ifndef TVM_RUNTIME_HEXAGON_RPC_SIMULATOR_HEXAGON_SIM_PROTO_H_
+#define TVM_RUNTIME_HEXAGON_RPC_SIMULATOR_HEXAGON_SIM_PROTO_H_
+
+struct Message {
+  enum : uint32_t {
+    kNone = 0,
+    kAck,
+    kTerminate,
+    kReceiveStart,
+    kReceiveEnd,
+    kSendStart,
+    kSendEnd,
+  };
+  enum : uint32_t {
+    null_va = 0,
+  };
+
+  uint32_t code;
+  uint32_t len;
+  uint32_t va;
+} __attribute__((packed));
+
+// Protocol:
+//
+// Copying data from host to remote:
+//
+// Host >-- [ kReceiveStart, len,      null_va ] --> Remote
+// * Remote client prepares a buffer with at least `len` bytes.
+// Host <-- [ kAck,          buf_size, buf_ptr ] <-- Remote
+// * Host writes `nbytes` into buffer, `nbytes` <= `len`.
+// Host >-- [ kReceiveEnd,   nbytes,   buf_ptr ] --> Remote
+// * Remote client processes the data.
+// Host <-- [ kAck,          ___,      ___     ] <-- Remote
+//
+// Copying data from remote to host:
+//
+// Host >-- [ kSendStart,    len,      null_va ] --> Remote
+// * Remote client returns pointer to the buffer with the data to be read.
+// * There should be at least `len` bytes ready in the buffer.
+// Host <-- [ kAck,          buf_size, buf_ptr ] <-- Remote
+// * Host reads `nbytes` from buffer, `nbytes` <= `buf_size`.
+// Host >-- [ kSendEnd   ,   nbytes,   buf_ptr ] --> Remote
+// * Remote client processes the data.
+// Host <-- [ kAck,          ___,      ___     ] <-- Remote
+//
+// Teminating server:
+//
+// Host >-- [ kTerminate,    ___,      ___     ] --> Remote
+// Host <-- [ kAck,          ___,      ___     ] <-- Remote
+// * Host continues execution of the client.
+// * Client terminates.
+
+#define DISPATCH_FUNCTION_NAME dispatch_875b2e3a28186123
+#define MESSAGE_BUFFER_NAME message_buffer_71d6a7b93c318d7e
+
+#endif  // TVM_RUNTIME_HEXAGON_RPC_SIMULATOR_HEXAGON_SIM_PROTO_H_
diff --git a/src/runtime/hexagon/rpc/simulator/rpc_server.cc b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
new file mode 100644
index 000000000000..715d30f81e37
--- /dev/null
+++ b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
@@ -0,0 +1,319 @@
+/*
+ * 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.
+ */
+
+#include <HAP_farf.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstdlib>
+#include <sstream>
+#include <string>
+
+#include "../../../library_module.h"
+#include "../../../minrpc/minrpc_server.h"
+#include "../../hexagon/hexagon_common.h"
+#include "hexagon_sim_proto.h"
+#include "tvm/runtime/packed_func.h"
+#include "tvm/runtime/registry.h"
+
+namespace tvm {
+namespace runtime {
+namespace hexagon {
+
+class stringbuf_with_remote_access : public std::stringbuf {
+  constexpr static size_t zero_count = 1024;
+  constexpr static char zeros[zero_count] = {};
+
+ public:
+  char* reserve_for_remote_write(size_t size) {
+    // Reserve memory in the put area by adding zeros to it. The put
+    // area will automatically resize itself as needed. This is needed
+    // for the simulator to be able to write to program's memory.
+    size_t remaining = size;
+    while (remaining >= zero_count) {
+      sputn(zeros, zero_count);
+      remaining -= zero_count;
+    }
+    sputn(zeros, remaining);
+    return pptr() - size;
+  }
+
+  // Get area is the storage area with the data that will be read from
+  // the buffer. From the buffer's point of view, this is the area with
+  // the outgoing data.
+  char* get_area() { return gptr(); }
+
+  // Adjust the buffer state after the data has been read by the remote
+  // end. This means discarding given amount of bytes from the get area.
+  size_t acknowledge_remote_read(size_t size) {
+    size_t remaining = size;
+    do {
+      auto bump = std::min<size_t>(in_avail(), remaining);
+      setg(eback(), gptr() + bump, egptr());
+      remaining -= bump;
+      // The setg will cause in_avail to become 0. At this point calling
+      // bumpc will either recaulate the get area pointers (if there are
+      // more characters left), or it will return eof().
+      if (remaining == 0 || sbumpc() == traits_type::eof()) {
+        break;
+      }
+      remaining--;
+    } while (remaining > 0);
+
+    // This will return 0 on success, non-zero if underflow occured.
+    return size - remaining;
+  }
+};
+
+const char stringbuf_with_remote_access::zeros[zero_count];
+
+class SimulatorIOHandler {
+ public:
+  SimulatorIOHandler() = default;
+
+  void MessageStart(size_t message_size_bytes) {}
+  void MessageDone() {}
+
+  // Store data from the input buffer into 'buf'.
+  ssize_t PosixRead(uint8_t* buf, size_t read_len_bytes) {
+    auto char_buf = reinterpret_cast<decltype(inp_buffer_)::char_type*>(buf);
+    return static_cast<ssize_t>(inp_buffer_.sgetn(char_buf, read_len_bytes));
+  }
+
+  // Append 'write_len_bytes' starting at 'buf' to the output buffer.
+  ssize_t PosixWrite(const uint8_t* buf, size_t write_len_bytes) {
+    auto char_buf = reinterpret_cast<const decltype(out_buffer_)::char_type*>(buf);
+    return static_cast<ssize_t>(out_buffer_.sputn(char_buf, write_len_bytes));
+  }
+
+  void Close() {}
+  void Exit(int code) { exit(code); }
+
+  char* PrepareToReceiveFromRemote(size_t nbytes) {
+    // Reserve space in the incoming buffer.
+    return inp_buffer_.reserve_for_remote_write(nbytes);
+  }
+  bool CompleteReceiveFromRemote(size_t nbytes) {
+    // Nothing to do.
+    return true;
+  }
+  char* PrepareToSendToRemote(size_t nbytes) {
+    // Return the pointer to the data coming out of the buffer.
+    return out_buffer_.get_area();
+  }
+  bool CompleteSendToRemote(size_t nbytes) {
+    // Flush the read data from the outgoing buffer.
+    return out_buffer_.acknowledge_remote_read(nbytes) == 0;
+  }
+
+ private:
+  stringbuf_with_remote_access inp_buffer_;  // Data received from remote.
+  stringbuf_with_remote_access out_buffer_;  // Data to be sent to remote.
+};
+
+// Internal allocator that redirects alloc to TVM's C API.
+template <typename TIOHandler>
+class SimulatorPageAllocator {
+ public:
+  using ArenaPageHeader = tvm::support::ArenaPageHeader;
+
+  explicit SimulatorPageAllocator(TIOHandler* io) : io_(io) {}
+
+  ArenaPageHeader* allocate(size_t min_size) {
+    size_t npages = ((min_size + kPageSize - 1) / kPageSize);
+    void* data;
+
+    if (posix_memalign(&data, kPageAlign, npages * kPageSize) != 0) {
+      io_->Exit(static_cast<int>(RPCServerStatus::kAllocError));
+    }
+
+    ArenaPageHeader* header = static_cast<ArenaPageHeader*>(data);
+    header->size = npages * kPageSize;
+    header->offset = sizeof(ArenaPageHeader);
+    return header;
+  }
+
+  void deallocate(ArenaPageHeader* page) { free(page); }
+
+  static const constexpr int kPageSize = 2 << 10;
+  static const constexpr int kPageAlign = kPageSize;
+
+ private:
+  TIOHandler* io_;
+};
+
+class SimulatorRPCServer : public MinRPCServer<SimulatorIOHandler, SimulatorPageAllocator> {
+  using Base = MinRPCServer<SimulatorIOHandler, SimulatorPageAllocator>;
+
+ public:
+  SimulatorRPCServer() : Base(&io_) {}
+
+  char* PrepareToReceive(size_t nbytes) {
+    // Reserve receiving area.
+    return io_.PrepareToReceiveFromRemote(nbytes);
+  }
+  bool CompleteReceive(size_t nbytes) {
+    // Interpret the received message.
+    return io_.CompleteReceiveFromRemote(nbytes) && ProcessOnePacket();
+  }
+  char* PrepareToSend(size_t nbytes) {
+    // Identify the beginning of the outgoing data.
+    return io_.PrepareToSendToRemote(nbytes);
+  }
+  bool CompleteSend(size_t nbytes) {
+    // Flush the read data.
+    return io_.CompleteSendToRemote(nbytes);
+  }
+
+ private:
+  SimulatorIOHandler io_;
+};
+
+}  // namespace hexagon
+}  // namespace runtime
+}  // namespace tvm
+
+// Handling communication with the simulator.
+//
+// Simulator can read and write the memory of the process, but the process has
+// no way to find out that it's running in a simulation, nor can it actively
+// communicate with the simulator. The RPC server must be entirely passive,
+// allowing the simulator to perform data transfers.
+
+extern "C" {
+// The names of these symbols will "pollute" the global namespace in the final
+// binary, so make them unique (i.e. "more likely to be unique"). These names
+// will be referenced in the host's code (the code controlling the simulator)
+// as well, so to avoid repetition use human-friendly macros.
+int DISPATCH_FUNCTION_NAME(void*) __attribute__((noinline));
+alignas(8) volatile Message MESSAGE_BUFFER_NAME;
+}
+
+inline uint32_t va(const volatile void* p) {
+  static_assert(sizeof(p) == sizeof(uint32_t), "Pointers must be 32-bit long");
+  return static_cast<uint32_t>(reinterpret_cast<uintptr_t>(p));
+}
+
+// NOLINTNEXTLINE(runtime/references)
+__attribute__((__unused__)) static std::string to_string(const volatile Message& m) {
+  std::stringstream out;
+  out << "{code=";
+  switch (m.code) {
+    case Message::kNone:
+      out << "kNone";
+      break;
+    case Message::kAck:
+      out << "kAck";
+      break;
+    case Message::kTerminate:
+      out << "kTerminate";
+      break;
+    case Message::kReceiveStart:
+      out << "kReceiveStart";
+      break;
+    case Message::kReceiveEnd:
+      out << "kReceiveEnd";
+      break;
+    case Message::kSendStart:
+      out << "kSendStart";
+      break;
+    case Message::kSendEnd:
+      out << "kSendEnd";
+      break;
+    default:
+      out << "<unknown>(" << m.code << ")";
+      break;
+  }
+  out << ", len:" << m.len << ", va:" << std::hex << m.va << std::dec << '}';
+  return out.str();
+}
+
+static inline void setmsg(volatile Message* m, uint32_t code, uint32_t len, uint32_t va) {
+  m->code = code;
+  m->len = len;
+  m->va = va;
+}
+
+int DISPATCH_FUNCTION_NAME(void* serverp) {
+  static bool terminate = false;
+  if (terminate) {
+    return 1;
+  }
+
+  Message msg;
+  setmsg(&msg, MESSAGE_BUFFER_NAME.code, MESSAGE_BUFFER_NAME.len, MESSAGE_BUFFER_NAME.va);
+
+  auto& server = *reinterpret_cast<tvm::runtime::hexagon::SimulatorRPCServer*>(serverp);
+
+  switch (msg.code) {
+    case Message::kReceiveStart:
+      assert(msg.va == Message::null_va);
+      setmsg(&MESSAGE_BUFFER_NAME, Message::kAck, msg.len, va(server.PrepareToReceive(msg.len)));
+      break;
+    case Message::kReceiveEnd:
+      server.CompleteReceive(msg.len);
+      setmsg(&MESSAGE_BUFFER_NAME, Message::kAck, 0u, Message::null_va);
+      break;
+    case Message::kSendStart:
+      assert(msg.va == Message::null_va);
+      setmsg(&MESSAGE_BUFFER_NAME, Message::kAck, msg.len, va(server.PrepareToSend(msg.len)));
+      break;
+    case Message::kSendEnd:
+      server.CompleteSend(msg.len);
+      setmsg(&MESSAGE_BUFFER_NAME, Message::kAck, 0u, Message::null_va);
+      break;
+    case Message::kTerminate:
+      // Don't exit immediately, send response to simulator first.
+      terminate = true;
+      setmsg(&MESSAGE_BUFFER_NAME, Message::kAck, 0u, Message::null_va);
+      break;
+  }
+
+  return 0;
+}
+
+int main() {
+  const auto* api_v2 = tvm::runtime::Registry::Get("device_api.hexagon.v2");
+  ICHECK(api_v2 != nullptr);
+  tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(*api_v2);
+
+  tvm::runtime::hexagon::SimulatorRPCServer server;
+
+  // Hand-encode user-instruction:
+  // r17:16 = userinsn(r17:16, r17:16, #0)
+  // 1100 1111 00010000 11010000 00010000 : cf10d010
+
+  asm volatile("r17:16 = combine(%0,%1); .long 0xcf10d010"
+               :  // No outputs
+               : "r"(&MESSAGE_BUFFER_NAME), "r"(&DISPATCH_FUNCTION_NAME)
+               : "r16", "r17");
+
+  while (!DISPATCH_FUNCTION_NAME(&server)) {
+    // nothing
+  }
+
+  return 0;
+}
+
+TVM_REGISTER_GLOBAL("tvm.hexagon.load_module")
+    .set_body([](tvm::runtime::TVMArgs args, tvm::runtime::TVMRetValue* rv) {
+      std::string soname = args[0];
+      tvm::ObjectPtr<tvm::runtime::Library> n = tvm::runtime::CreateDSOLibraryObject(soname);
+      *rv = CreateModuleFromLibrary(n, tvm::runtime::hexagon::WrapPackedFunc);
+    });
diff --git a/src/runtime/hexagon/rpc/simulator/session.cc b/src/runtime/hexagon/rpc/simulator/session.cc
new file mode 100644
index 000000000000..164738a1d162
--- /dev/null
+++ b/src/runtime/hexagon/rpc/simulator/session.cc
@@ -0,0 +1,1322 @@
+/*
+ * 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.
+ */
+
+#include <HexagonWrapper.h>
+#include <dmlc/optional.h>
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/registry.h>
+// POSIX includes
+#include <dirent.h>
+#include <unistd.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <deque>
+#include <iterator>
+#include <map>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <utility>
+
+#include "../../../rpc/rpc_channel.h"
+#include "../../../rpc/rpc_endpoint.h"
+#include "../../../rpc/rpc_session.h"
+#include "hexagon_sim_proto.h"
+
+#define CHECKED_CALL(func, ...)                                               \
+  do {                                                                        \
+    HEXAPI_Status s = sim_->func(__VA_ARGS__);                                \
+    ICHECK_EQ(s, HEX_STAT_SUCCESS)                                            \
+        << self_name_ << ": " #func " failed with code " << Status_{s}.str(); \
+  } while (false)
+
+namespace tvm {
+namespace runtime {
+namespace hexagon {
+
+using string_list = std::deque<std::string>;
+
+namespace detail {
+
+// An "Optional" class, originally a replacement for llvm::Optional, then an
+// extension of dmlc::optional to make it compatible with C++17's std::optional.
+template <typename T>
+struct Optional : public dmlc::optional<T> {
+  using dmlc::optional<T>::optional;
+  using dmlc::optional<T>::operator=;
+  Optional(const T& val) : dmlc::optional<T>(val) {}  // NOLINT(*)
+
+  T* operator->() { return &this->operator*(); }
+  const T* operator->() const { return &this->operator*(); }
+};
+
+// Replacement for llvm::StringSwitch.
+template <typename T>
+class StringSwitch {
+ public:
+  explicit StringSwitch(const std::string& key) : key(key) {}
+  operator T() const {
+    auto f = map.find(key);
+    if (f != map.end()) {
+      return f->second;
+    }
+    ICHECK(static_cast<bool>(def_val)) << "default value not set";
+    return *def_val;
+  }
+  StringSwitch& Case(const std::string& key, T val) {
+    map.insert(std::make_pair(key, val));
+    return *this;
+  }
+  StringSwitch& Default(T val) {
+    ICHECK(!static_cast<bool>(def_val)) << "default value already set";
+    def_val = val;
+    return *this;
+  }
+
+ private:
+  const std::string key;
+  std::map<std::string, T> map;
+  Optional<T> def_val;
+};
+
+using MaybeString = Optional<std::string>;
+
+MaybeString front(const string_list& deq) {
+  return !deq.empty() ? MaybeString(deq.front()) : MaybeString();
+}
+
+MaybeString pop_front(string_list& deq) {  // NOLINT(*)
+  if (deq.empty()) return MaybeString();
+  std::string f = deq.front();
+  deq.pop_front();
+  return MaybeString(f);
+}
+
+// Functions used when parsing the argument string.
+
+Optional<int64_t> to_int(const MaybeString& str) {
+  auto none = Optional<int64_t>();
+  if (str.has_value()) {
+    try {
+      size_t pos;
+      int64_t val = std::stoll(*str, &pos, 0);
+      return pos == str->size() ? Optional<int64_t>(val) : none;
+    } catch (std::invalid_argument&) {
+    }
+  }
+  return none;
+}
+
+Optional<uint64_t> to_uint(const MaybeString& str) {
+  auto none = Optional<uint64_t>();
+  if (str.has_value()) {
+    try {
+      size_t pos;
+      uint64_t val = std::stoull(*str, &pos, 0);
+      return pos == str->size() ? Optional<uint64_t>(val) : none;
+    } catch (std::invalid_argument&) {
+    }
+  }
+  return none;
+}
+
+Optional<float> to_float(const MaybeString& str) {
+  auto none = Optional<float>();
+  if (str.has_value()) {
+    try {
+      size_t pos;
+      float val = std::stof(*str, &pos);
+      return pos == str->size() ? Optional<float>(val) : none;
+    } catch (std::invalid_argument&) {
+    }
+  }
+  return none;
+}
+
+Optional<bool> to_bool(const MaybeString& str) {
+  auto none = Optional<bool>();
+  if (auto num = to_int(str)) {
+    if (*num == 0) return false;
+    if (*num == 1) return true;
+    return none;
+  }
+  if (str) {
+    if (*str == "true" || *str == "TRUE") return true;
+    if (*str == "false" || *str == "FALSE") return false;
+  }
+  return none;
+}
+
+template <typename T>
+using MaybeRange = Optional<std::pair<T, T>>;
+
+template <typename T, Optional<T> Parse(const MaybeString&)>
+MaybeRange<T> to_range(const MaybeString& str) {
+  auto none = MaybeRange<T>();
+  if (str && !str->empty()) {
+    auto n = str->find('-', 1);
+    if (n != std::string::npos) {
+      auto begin = Parse(str->substr(0, n));
+      auto end = Parse(str->substr(n + 1, str->size() - n - 1));
+      if (begin && end) {
+        return std::make_pair(*begin, *end);
+      }
+    }
+  }
+  return none;
+}
+
+}  // namespace detail
+
+class SimulatorRPCChannel final : public RPCChannel {
+ public:
+  SimulatorRPCChannel(std::string args);
+  ~SimulatorRPCChannel() final;
+  size_t Send(const void* data, size_t size) final;
+  size_t Recv(void* data, size_t size) final;
+
+ private:
+  struct Status_ {
+    HEXAPI_Status s;
+    std::string str() const;
+  };
+  struct Core_ {
+    HEXAPI_CoreState s;
+    std::string str() const;
+  };
+  struct Cpu_ {
+    HEXAPI_Cpu c;
+    std::string str() const;
+  };
+  struct SDKInfo_ {
+    SDKInfo_(const std::string& sdk_root, const std::string& cpu);
+    std::string root;
+    std::string qurt_root;  // sdk_root/rtos/qurt/computevNN.
+    std::string runelf;     // Path to runelf.pbn.
+    std::string runmain;    // Path to run_main_on_hexagon.
+  };
+
+  Message SendMsg(Message msg);
+  Message SendMsg(uint32_t code, uint32_t len, uint32_t va);
+  void ReadFromProcess(void* host_dst, HEX_VA_t src, size_t len);
+  void WriteToProcess(HEX_VA_t dst, const void* host_src, size_t len);
+
+  static HEX_8u_t PassVirtAddrCallback(void* handle, int threadno, HEX_8u_t RssV, HEX_8u_t RttV,
+                                       HEX_8u_t RxxV, HEX_1u_t imm);
+
+  detail::Optional<HEXAPI_Cpu> GetCPU(const detail::MaybeString& cpu_str);
+
+  // File name templates for mkstemps.
+#define SUFFIX ".cfg"
+  static constexpr int template_length_ = sizeof("temp-xxxx-XXXXXX" SUFFIX);  // == strlen() + 1
+  char osam_file_[template_length_] = "temp-osam-XXXXXX" SUFFIX;
+  char cosim_file_[template_length_] = "temp-q6ss-XXXXXX" SUFFIX;
+  const int suffix_len_ = strlen(SUFFIX);
+#undef SUFFIX
+
+  static const constexpr char* self_name_ = "SimulatorRPCChannel";
+  static const constexpr char* default_cpu_ = "v68";
+  std::string cpu_;
+
+  HEX_VA_t dispatch_v_, message_buffer_v_;
+  std::unique_ptr<HexagonWrapper> sim_;
+
+  // Sim configuration routines.
+  bool Configure(string_list& opts);  // NOLINT(*)
+
+  bool HandleAHBBusPenalty(string_list& rest);      // NOLINT(*)
+  bool HandleAHBBusRatio(string_list& rest);        // NOLINT(*)
+  bool HandleAHBHighAddr(string_list& rest);        // NOLINT(*)
+  bool HandleAHBLowAddr(string_list& rest);         // NOLINT(*)
+  bool HandleAXI2BusPenalty(string_list& rest);     // NOLINT(*)
+  bool HandleAXI2BusRatio(string_list& rest);       // NOLINT(*)
+  bool HandleAXI2HighAddr(string_list& rest);       // NOLINT(*)
+  bool HandleAXI2LowAddr(string_list& rest);        // NOLINT(*)
+  bool HandleBuildTag(string_list& rest);           // NOLINT(*)
+  bool HandleBusPenalty(string_list& rest);         // NOLINT(*)
+  bool HandleBusRatio(string_list& rest);           // NOLINT(*)
+  bool HandleBusTrace(string_list& rest);           // NOLINT(*)
+  bool HandleBypassIdle(string_list& rest);         // NOLINT(*)
+  bool HandleConnectionTimeout(string_list& rest);  // NOLINT(*)
+  bool HandleCoprocTrace(string_list& rest);        // NOLINT(*)
+  bool HandleCoreDump(string_list& rest);           // NOLINT(*)
+  bool HandleCosimFile(string_list& rest);          // NOLINT(*)
+  bool HandleDCacheTrace(string_list& rest);        // NOLINT(*)
+  bool HandleDSPClock(string_list& rest);           // NOLINT(*)
+  bool HandleETMCFGBase(string_list& rest);         // NOLINT(*)
+  bool HandleGDBServ(string_list& rest);            // NOLINT(*)
+  bool HandleHVXLength(string_list& rest);          // NOLINT(*)
+  bool HandleICacheTrace(string_list& rest);        // NOLINT(*)
+  bool HandleL2CacheTrace(string_list& rest);       // NOLINT(*)
+  bool HandleL2CFGBase(string_list& rest);          // NOLINT(*)
+  bool HandleL2TCMBase(string_list& rest);          // NOLINT(*)
+  bool HandleMemFillRand(string_list& rest);        // NOLINT(*)
+  bool HandleMemFill(string_list& rest);            // NOLINT(*)
+  bool HandleMemTrace(string_list& rest);           // NOLINT(*)
+  bool HandleNullPtr(string_list& rest);            // NOLINT(*)
+  bool HandlePacketAnalyze(string_list& rest);      // NOLINT(*)
+  bool HandlePCFilter(string_list& rest);           // NOLINT(*)
+  bool HandlePCTraceMin(string_list& rest);         // NOLINT(*)
+  bool HandlePCTraceNano(string_list& rest);        // NOLINT(*)
+  bool HandlePCTrace(string_list& rest);            // NOLINT(*)
+  bool HandlePMUStatsFile(string_list& rest);       // NOLINT(*)
+  bool HandleProfile(string_list& rest);            // NOLINT(*)
+  bool HandleProfileTimeZero(string_list& rest);    // NOLINT(*)
+  bool HandleQuiet(string_list& rest);              // NOLINT(*)
+  bool HandleReconnect(string_list& rest);          // NOLINT(*)
+  bool HandleRTOS(string_list& rest);               // NOLINT(*)
+  bool HandleSimErr(string_list& rest);             // NOLINT(*)
+  bool HandleSimIn(string_list& rest);              // NOLINT(*)
+  bool HandleSimOut(string_list& rest);             // NOLINT(*)
+  bool HandleStackStart(string_list& rest);         // NOLINT(*)
+  bool HandleStallTrace(string_list& rest);         // NOLINT(*)
+  bool HandleStatsFile(string_list& rest);          // NOLINT(*)
+  bool HandleSubsystemBase(string_list& rest);      // NOLINT(*)
+  bool HandleSymFile(string_list& rest);            // NOLINT(*)
+  bool HandleTCM(string_list& rest);                // NOLINT(*)
+  bool HandleTCMHighAddr(string_list& rest);        // NOLINT(*)
+  bool HandleTCMLowAddr(string_list& rest);         // NOLINT(*)
+  bool HandleTimeFilterNS(string_list& rest);       // NOLINT(*)
+  bool HandleTiming(string_list& rest);             // NOLINT(*)
+  bool HandleUArchTrace(string_list& rest);         // NOLINT(*)
+  bool HandleUseFS(string_list& rest);              // NOLINT(*)
+  bool HandleV2PTranslation(string_list& rest);     // NOLINT(*)
+  bool HandleVerbose(string_list& rest);            // NOLINT(*)
+
+  using MaybeUInt64 = detail::Optional<uint64_t>;
+  using MaybeUIntRange = std::pair<MaybeUInt64, MaybeUInt64>;
+
+  bool should_parse_next(const string_list& rest);
+  detail::Optional<HEXAPI_Interval> to_interval(const detail::MaybeString& str);
+  detail::Optional<HEXAPI_TimingMode> to_timingmode(const detail::MaybeString& str);
+  detail::Optional<HEXAPI_VerboseMode> to_verbosemode(const detail::MaybeString& str);
+  detail::Optional<HEXAPI_Nullptr> to_nullptr(const detail::MaybeString& str);
+
+  MaybeUIntRange ahb_, axi2_;
+  detail::Optional<uint32_t> debug_port_;
+
+  using OptionHandler = bool (SimulatorRPCChannel::*)(string_list&);
+  static std::map<std::string, OptionHandler> opt_map_;
+};
+
+const constexpr char* SimulatorRPCChannel::self_name_;
+const constexpr char* SimulatorRPCChannel::default_cpu_;
+
+decltype(SimulatorRPCChannel::opt_map_) SimulatorRPCChannel::opt_map_ = {
+    {"--ahbbuspenalty", &SimulatorRPCChannel::HandleAHBBusPenalty},
+    {"--ahbbusratio", &SimulatorRPCChannel::HandleAHBBusRatio},
+    {"--ahb:highaddr", &SimulatorRPCChannel::HandleAHBHighAddr},
+    {"--ahb:lowaddr", &SimulatorRPCChannel::HandleAHBLowAddr},
+    {"--axi2buspenalty", &SimulatorRPCChannel::HandleAXI2BusPenalty},
+    {"--axi2busratio", &SimulatorRPCChannel::HandleAXI2BusRatio},
+    {"--axi2:highaddr", &SimulatorRPCChannel::HandleAXI2HighAddr},
+    {"--axi2:lowaddr", &SimulatorRPCChannel::HandleAXI2LowAddr},
+    {"-b", &SimulatorRPCChannel::HandleBusTrace},
+    {"--build_tag", &SimulatorRPCChannel::HandleBuildTag},
+    {"--buspenalty", &SimulatorRPCChannel::HandleBusPenalty},
+    {"--busratio", &SimulatorRPCChannel::HandleBusRatio},
+    {"--bustrace", &SimulatorRPCChannel::HandleBusTrace},
+    {"--bypass_idle", &SimulatorRPCChannel::HandleBypassIdle},
+    {"--connection_timeout", &SimulatorRPCChannel::HandleConnectionTimeout},
+    {"--coproctrace", &SimulatorRPCChannel::HandleCoprocTrace},
+    {"--coredump", &SimulatorRPCChannel::HandleCoreDump},
+    {"--cosim_file", &SimulatorRPCChannel::HandleCosimFile},
+    {"--dcachetrace", &SimulatorRPCChannel::HandleDCacheTrace},
+    {"--dsp_clock", &SimulatorRPCChannel::HandleDSPClock},
+    {"-E", &SimulatorRPCChannel::HandleSimErr},
+    {"--etm_base", &SimulatorRPCChannel::HandleETMCFGBase},
+    {"--etmcfg_base", &SimulatorRPCChannel::HandleETMCFGBase},
+    {"--gdbserv", &SimulatorRPCChannel::HandleGDBServ},
+    {"-G", &SimulatorRPCChannel::HandleGDBServ},
+    {"--hvx_length", &SimulatorRPCChannel::HandleHVXLength},
+    {"--icachetrace", &SimulatorRPCChannel::HandleICacheTrace},
+    {"-I", &SimulatorRPCChannel::HandleSimIn},
+    {"--l2cachetrace", &SimulatorRPCChannel::HandleL2CacheTrace},
+    {"--l2cfg_base", &SimulatorRPCChannel::HandleL2CFGBase},
+    {"--l2tcm_base", &SimulatorRPCChannel::HandleL2TCMBase},
+    {"--memfill", &SimulatorRPCChannel::HandleMemFill},
+    {"--memfill_rand", &SimulatorRPCChannel::HandleMemFillRand},
+    {"--memtrace", &SimulatorRPCChannel::HandleMemTrace},
+    {"-m", &SimulatorRPCChannel::HandleMemTrace},
+    {"--nullptr", &SimulatorRPCChannel::HandleNullPtr},
+    {"-O", &SimulatorRPCChannel::HandleSimOut},
+    {"--packet_analyze", &SimulatorRPCChannel::HandlePacketAnalyze},
+    {"--pcfilter", &SimulatorRPCChannel::HandlePCFilter},
+    {"--pctrace", &SimulatorRPCChannel::HandlePCTrace},
+    {"--pctrace_min", &SimulatorRPCChannel::HandlePCTraceMin},
+    {"--pctrace_nano", &SimulatorRPCChannel::HandlePCTraceNano},
+    {"-p", &SimulatorRPCChannel::HandleProfile},
+    {"--pmu_statsfile", &SimulatorRPCChannel::HandlePMUStatsFile},
+    {"--profile", &SimulatorRPCChannel::HandleProfile},
+    {"--profile_timezero", &SimulatorRPCChannel::HandleProfileTimeZero},
+    {"-q", &SimulatorRPCChannel::HandleQuiet},
+    {"--quiet", &SimulatorRPCChannel::HandleQuiet},
+    {"--reconnect", &SimulatorRPCChannel::HandleReconnect},
+    {"--rtos", &SimulatorRPCChannel::HandleRTOS},
+    {"-S", &SimulatorRPCChannel::HandleStatsFile},
+    {"--sim_err", &SimulatorRPCChannel::HandleSimErr},
+    {"--sim_in", &SimulatorRPCChannel::HandleSimIn},
+    {"--sim_out", &SimulatorRPCChannel::HandleSimOut},
+    {"--stackstart", &SimulatorRPCChannel::HandleStackStart},
+    {"--stalltrace", &SimulatorRPCChannel::HandleStallTrace},
+    {"--statsfile", &SimulatorRPCChannel::HandleStatsFile},
+    {"--subsystem_base", &SimulatorRPCChannel::HandleSubsystemBase},
+    {"--symfile", &SimulatorRPCChannel::HandleSymFile},
+    {"--tcm", &SimulatorRPCChannel::HandleTCM},
+    {"--tcm:highaddr", &SimulatorRPCChannel::HandleTCMHighAddr},
+    {"--tcm:lowaddr", &SimulatorRPCChannel::HandleTCMLowAddr},
+    {"-t", &SimulatorRPCChannel::HandlePCTrace},
+    {"--timefilter_ns", &SimulatorRPCChannel::HandleTimeFilterNS},
+    {"--timing", &SimulatorRPCChannel::HandleTiming},
+    {"--uarchtrace", &SimulatorRPCChannel::HandleUArchTrace},
+    {"-u", &SimulatorRPCChannel::HandlePCTraceMin},
+    {"--usefs", &SimulatorRPCChannel::HandleUseFS},
+    {"--v2p_translation", &SimulatorRPCChannel::HandleV2PTranslation},
+    {"--verbose", &SimulatorRPCChannel::HandleVerbose},
+};
+
+std::string SimulatorRPCChannel::Status_::str() const {
+  switch (s) {
+    case HEX_STAT_ERROR:
+      return "HEX_STAT_ERROR";
+    case HEX_STAT_SUCCESS:
+      return "HEX_STAT_SUCCESS";
+    case HEX_STAT_CANNOT_CONFIG:
+      return "HEX_STAT_CANNOT_CONFIG";
+    case HEX_STAT_INVALID_ARGS:
+      return "HEX_STAT_INVALID_ARGS";
+    case HEX_STAT_RANGE_ERROR:
+      return "HEX_STAT_RANGE_ERROR";
+    case HEX_STAT_FILE_ACCESS_ERROR:
+      return "HEX_STAT_FILE_ACCESS_ERROR";
+    case HEX_STAT_DEVICE_NOT_FOUND:
+      return "HEX_STAT_DEVICE_NOT_FOUND";
+    case HEX_STAT_MEM_ACCESS_ERROR:
+      return "HEX_STAT_MEM_ACCESS_ERROR";
+    case HEX_STAT_CANNOT_TRANSLATE:
+      return "HEX_STAT_CANNOT_TRANSLATE";
+    case HEX_STAT_NO_ACTIVE_THREADS:
+      return "HEX_STAT_NO_ACTIVE_THREADS";
+    case HEX_STAT_LOAD_ELF_ERROR:
+      return "HEX_STAT_LOAD_ELF_ERROR";
+    case HEX_STAT_CORE_RESET:
+      return "HEX_STAT_CORE_RESET";
+    default:
+      break;
+  }
+  return std::to_string(static_cast<int>(s));
+}
+
+std::string SimulatorRPCChannel::Core_::str() const {
+  switch (s) {
+    case HEX_CORE_SUCCESS:
+      return "HEX_CORE_SUCCESS";
+    case HEX_CORE_FINISHED:
+      return "HEX_CORE_FINISHED";
+    case HEX_CORE_RESET:
+      return "HEX_CORE_RESET";
+    case HEX_CORE_BREAKPOINT:
+      return "HEX_CORE_BREAKPOINT";
+    case HEX_CORE_ASYNCHRONOUS_BREAK:
+      return "HEX_CORE_ASYNCHRONOUS_BREAK";
+    case HEX_CORE_ERROR:
+      return "HEX_CORE_ERROR";
+    default:
+      break;
+  }
+  return std::to_string(static_cast<int>(s));
+}
+
+std::string SimulatorRPCChannel::Cpu_::str() const {
+  switch (c) {
+    case HEX_CPU_V65:
+      return "v65";
+    case HEX_CPU_V66:
+      return "v66";
+    case HEX_CPU_V68:
+      return "v68";
+    case HEX_CPU_V69:
+      return "v69";
+    default:
+      break;
+  }
+  return default_cpu_;
+}
+
+SimulatorRPCChannel::SDKInfo_::SDKInfo_(const std::string& sdk_root, const std::string& cpu)
+    : root(sdk_root) {
+  qurt_root = root + "/rtos/qurt/compute" + cpu;
+  runelf = qurt_root + "/sdksim_bin/runelf.pbn";
+
+  // The "run_main_on_hexagon_sim" binary lives in a subdirectory that looks
+  // like "[...]on_hexagon/ship/hexagon_toolv84_v68/run_main_on_hexagon_sim".
+  // We need to get the right "hexagon_toolv..." component, based mostly on
+  // the cpu version.
+  std::vector<std::string> dir_names;
+
+  DIR* dir = opendir((root + "/libs/run_main_on_hexagon/ship").c_str());
+  ICHECK(dir != nullptr) << "Cannot read directory " << root + "/libs/run_main_on_hexagon/ship";
+  while (dirent* d = readdir(dir)) {
+    if (d->d_type != DT_DIR) continue;
+
+    std::string name = d->d_name;
+    // Note: The first substr is always safe, and the second only executes
+    // when "name" is at least 13 characters long.
+    if (name.substr(0, 13) == "hexagon_toolv" && name.substr(name.size() - 3, 3) == cpu) {
+      dir_names.push_back(name);
+    }
+  }
+  closedir(dir);
+  ICHECK(!dir_names.empty());
+
+  auto max_it = std::max_element(dir_names.begin(), dir_names.end());
+  runmain = root + "/libs/run_main_on_hexagon/ship/" + *max_it + "/run_main_on_hexagon_sim";
+}
+
+HEX_8u_t SimulatorRPCChannel::PassVirtAddrCallback(void* handle, int threadno, HEX_8u_t RssV,
+                                                   HEX_8u_t RttV, HEX_8u_t RxxV, HEX_1u_t imm) {
+  // Rssv = combine(&message_buffer, &dispatch)
+  auto* rpc = reinterpret_cast<SimulatorRPCChannel*>(handle);
+  rpc->dispatch_v_ = RssV & ~0u;  // ~0u is uint32_t
+  rpc->message_buffer_v_ = RssV >> 32;
+
+  LOG(INFO) << "dispatch:" << reinterpret_cast<void*>(rpc->dispatch_v_)
+            << ", message buffer:" << reinterpret_cast<void*>(rpc->message_buffer_v_);
+  HEXAPI_Status s = rpc->sim_->SetBreakpoint(rpc->dispatch_v_);
+  ICHECK_EQ(s, HEX_STAT_SUCCESS) << self_name_ << ": SetBreakpoint failed with code "
+                                 << Status_{s}.str();
+  return RssV;
+}
+
+detail::Optional<HEXAPI_Cpu> SimulatorRPCChannel::GetCPU(const detail::MaybeString& cpu_str) {
+  auto none = detail::Optional<HEXAPI_Cpu>();
+  if (!cpu_str) return none;
+  return detail::StringSwitch<detail::Optional<HEXAPI_Cpu>>(*cpu_str)
+      .Case("v65", HEX_CPU_V65)
+      .Case("v66", HEX_CPU_V66)
+      .Case("v68", HEX_CPU_V68)
+      .Case("v69", HEX_CPU_V69)
+      .Default(none);
+}
+
+SimulatorRPCChannel::SimulatorRPCChannel(std::string args) {
+  const auto* api_v2 = tvm::runtime::Registry::Get("device_api.hexagon.v2");
+  ICHECK(api_v2 != nullptr);
+  tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(*api_v2);
+
+  const char* sdk_root_env = std::getenv("HEXAGON_SDK_ROOT");
+  ICHECK(sdk_root_env != nullptr) << "Please set HEXAGON_SDK_ROOT";
+  const char* toolchain_env = std::getenv("HEXAGON_TOOLCHAIN");
+  ICHECK(toolchain_env != nullptr) << "Please set HEXAGON_TOOLCHAIN";
+
+  std::string sdk_root(sdk_root_env);
+  std::string toolchain(toolchain_env);
+
+  auto sim_args_iss = std::istringstream(args);
+  using iterator = std::istream_iterator<std::string>;
+  auto sim_args = string_list(iterator(sim_args_iss), iterator());
+
+  detail::MaybeString target_str = detail::pop_front(sim_args);
+  auto maybe_cpu = GetCPU(target_str);
+  if (!maybe_cpu) {
+    if (!target_str || target_str->empty()) {
+      LOG(INFO) << "CPU not given, defaulting to " << default_cpu_;
+      maybe_cpu = GetCPU(std::string(default_cpu_));
+    } else {
+      LOG(FATAL) << "Invalid CPU name " << *target_str;
+    }
+  }
+  cpu_ = Cpu_{*maybe_cpu}.str();
+  sim_ = std::make_unique<HexagonWrapper>(*maybe_cpu, (toolchain + "/lib/iss").c_str());
+  SDKInfo_ sdk(sdk_root, cpu_);
+
+  // Prepare the osam.cfg file.
+  int fd_osam = mkstemps(osam_file_, suffix_len_);
+  ICHECK_GE(fd_osam, 0);
+  std::string osam_str = sdk.qurt_root + "/debugger/lnx64/qurt_model.so";
+  ICHECK_EQ(write(fd_osam, osam_str.c_str(), osam_str.size()), osam_str.size());
+  close(fd_osam);
+  // Prepare the q6ss.cfg file.
+  int fd_cosim = mkstemps(cosim_file_, suffix_len_);
+  ICHECK_GE(fd_cosim, 0);
+  std::string cosim_str =
+      toolchain +
+      "/lib/iss/qtimer.so --csr_base=0xFC900000 --irq_p=1 --freq=19200000 --cnttid=1\n" +
+      toolchain + "/lib/iss/l2vic.so 32 0xFC910000";
+  ICHECK_EQ(write(fd_cosim, cosim_str.c_str(), cosim_str.size()), cosim_str.size());
+  close(fd_cosim);
+
+  CHECKED_CALL(ConfigureL2tcmBase, 0xD800);
+  CHECKED_CALL(ConfigureARFilesystem, &std::string(".")[0]);
+  CHECKED_CALL(ConfigureOSAwareness, osam_file_);
+  CHECKED_CALL(ConfigureCosim, cosim_file_);
+  CHECKED_CALL(ConfigureExecutableBinary, sdk.runelf.c_str());
+
+  std::string cmdline = sdk.runelf + " " + sdk.runmain + " -- libhexagon_rpc_sim.so";
+  char* parg = &cmdline[0];
+  CHECKED_CALL(ConfigureAppCommandLine, 1, &parg);
+
+  // Configure the simulator.
+  Configure(sim_args);
+
+  CHECKED_CALL(EndOfConfiguration);
+  CHECKED_CALL(AddUserDefinedInstCallback, this, &PassVirtAddrCallback);
+
+  // Start the initial run, until the callback is executed.
+  HEX_4u_t result;
+  HEXAPI_CoreState core = sim_->Run(&result);
+  if (core != HEX_CORE_BREAKPOINT) {
+    LOG(FATAL) << self_name_ << ": Run not stopped on breakpoint, code=" << Core_{core}.str();
+  }
+
+  // At this point the simulator has executed the executable's initialization
+  // code that could have written to the SSR register.
+  // Enable UPCYCLE register.
+  HEX_4u_t thread_num;
+  CHECKED_CALL(GetCurrentHWThreadNum, &thread_num);
+  HEX_4u_t thread_ssr;
+  CHECKED_CALL(ReadThreadRegister, thread_num, TH_REG_SSR, &thread_ssr);
+  thread_ssr |= (1 << 23);
+  CHECKED_CALL(WriteThreadRegister, thread_num, TH_REG_SSR, thread_ssr);
+}
+
+SimulatorRPCChannel::~SimulatorRPCChannel() {
+  SendMsg(Message::kTerminate, 0, Message::null_va);
+
+  HEX_4u_t result;
+  HEXAPI_CoreState core = sim_->Run(&result);
+  ICHECK_EQ(core, HEX_CORE_FINISHED);
+
+  unlink(osam_file_);
+  unlink(cosim_file_);
+}
+
+size_t SimulatorRPCChannel::Send(const void* data, size_t size) {
+  ICHECK(size <= std::numeric_limits<uint32_t>::max());
+
+  Message reply_start =
+      SendMsg(Message::kReceiveStart, static_cast<uint32_t>(size), Message::null_va);
+  ICHECK_EQ(reply_start.code, Message::kAck);
+  ICHECK_GE(reply_start.len, size);
+  ICHECK_NE(reply_start.va, Message::null_va);
+
+  WriteToProcess(reply_start.va, data, size);
+
+  Message reply_end = SendMsg(Message::kReceiveEnd, static_cast<uint32_t>(size), reply_start.va);
+  ICHECK_EQ(reply_end.code, Message::kAck);
+  return size;
+}
+
+size_t SimulatorRPCChannel::Recv(void* data, size_t size) {
+  ICHECK(size <= std::numeric_limits<uint32_t>::max());
+
+  Message reply_start = SendMsg(Message::kSendStart, static_cast<uint32_t>(size), Message::null_va);
+  ICHECK_EQ(reply_start.code, Message::kAck);
+  ICHECK_GE(reply_start.len, size);
+  ICHECK_NE(reply_start.va, Message::null_va);
+
+  ReadFromProcess(data, reply_start.va, size);
+
+  Message reply_end = SendMsg(Message::kSendEnd, static_cast<uint32_t>(size), reply_start.va);
+  ICHECK_EQ(reply_end.code, Message::kAck);
+  return size;
+}
+
+Message SimulatorRPCChannel::SendMsg(Message msg) {
+  auto run = [this]() {
+    HEXAPI_CoreState core = HEX_CORE_RESET;
+    HEX_4u_t result;
+
+    core = sim_->Run(&result);
+    ICHECK_EQ(core, HEX_CORE_BREAKPOINT);
+  };
+
+  WriteToProcess(message_buffer_v_, &msg, sizeof msg);
+  run();
+
+  Message ret = {0};
+  ReadFromProcess(&ret, message_buffer_v_, sizeof ret);
+  return ret;
+}
+
+Message SimulatorRPCChannel::SendMsg(uint32_t code, uint32_t len, uint32_t va) {
+  Message m;
+  m.code = code;
+  m.len = len;
+  m.va = va;
+  return SendMsg(m);
+}
+
+void SimulatorRPCChannel::ReadFromProcess(void* host_dst, HEX_VA_t src, size_t len) {
+  if (len == 0) return;
+
+  auto* dst = reinterpret_cast<char*>(host_dst);
+
+  while (len > 0) {
+    uint32_t src_align = 1u << __builtin_ctz(src);  // Max pow-of-2 dividing src.
+    uint32_t len_align = 1u << __builtin_ctz(len);  // Max pow-of-2 dividing len.
+    // The transfer behaves as a bus transaction, so it has to be of a size
+    // that is a power of 2 not exceeding 8, and the remote address has to
+    // be aligned to at least the transfer size.
+    auto read_size = std::min<size_t>({src_align, len_align, 8});
+    CHECKED_CALL(ReadVirtual, src, /*asid*/ -1u, read_size, dst);
+
+    src += read_size;
+    dst += read_size;
+    len -= read_size;
+  }
+}
+
+void SimulatorRPCChannel::WriteToProcess(HEX_VA_t dst, const void* host_src, size_t len) {
+  if (len == 0) return;
+
+  auto* src = reinterpret_cast<const char*>(host_src);
+
+  while (len > 0) {
+    uint32_t dst_align = 1u << __builtin_ctz(dst);  // Max pow-of-2 dividing dst.
+    uint32_t len_align = 1u << __builtin_ctz(len);  // Max pow-of-2 dividing len.
+    // The transfer behaves as a bus transaction, so it has to be of a size
+    // that is a power of 2 not exceeding 8, and the remote address has to
+    // be aligned to at least the transfer size.
+    auto write_size = std::min<size_t>({dst_align, len_align, 8});
+    HEX_8u_t val = 0;
+    memcpy(&val, src, write_size);
+    CHECKED_CALL(WriteVirtual, dst, /*asid*/ -1u, write_size, val);
+
+    src += write_size;
+    dst += write_size;
+    len -= write_size;
+  }
+}
+
+// Configuration functions
+
+bool SimulatorRPCChannel::Configure(string_list& opts) {
+  while (!opts.empty()) {
+    std::string key = *detail::pop_front(opts);
+    auto f = opt_map_.find(key);
+    if (f == opt_map_.end()) {
+      LOG(FATAL) << "Unrecognized simulator option: " << key;
+      // unreachable
+    }
+    ICHECK((this->*f->second)(opts)) << "error handling option: " << key;
+  }
+
+  // Check AHB.
+  if (ahb_.first.has_value() && ahb_.second.has_value()) {
+    CHECKED_CALL(ConfigureAHB, *ahb_.first, *ahb_.second);
+  } else {
+    ICHECK(!ahb_.first.has_value() && !ahb_.second.has_value())
+        << self_name_ << ": please specify both low and high addresses for AHB";
+  }
+
+  // Check AXI2.
+  if (axi2_.first.has_value() && axi2_.second.has_value()) {
+    CHECKED_CALL(ConfigureAXI2, *axi2_.first, *axi2_.second);
+  } else {
+    ICHECK(!axi2_.first.has_value() && !axi2_.second.has_value())
+        << self_name_ << ": please specify both low and high addresses for AXI2";
+  }
+
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleAHBBusPenalty(string_list& rest) {
+  auto penalty = detail::to_uint(detail::pop_front(rest));
+  auto interval = to_interval(detail::pop_front(rest));
+  if (penalty && interval) {
+    CHECKED_CALL(ConfigureAHBBusPenalty, *penalty, *interval);
+  }
+  return static_cast<bool>(penalty) && static_cast<bool>(interval);
+}
+
+bool SimulatorRPCChannel::HandleAHBBusRatio(string_list& rest) {
+  auto ratio = detail::to_float(detail::pop_front(rest));
+  if (ratio) {
+    CHECKED_CALL(ConfigureAHBBusRatio, *ratio);
+  }
+  return static_cast<bool>(ratio);
+}
+
+bool SimulatorRPCChannel::HandleAHBHighAddr(string_list& rest) {
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  ICHECK(addr) << self_name_ << ": invalid value for AHB high adddress";
+  if (addr) {
+    ahb_.second = *addr;
+  }
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleAHBLowAddr(string_list& rest) {
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  ICHECK(addr) << self_name_ << ": invalid value for AHB low adddress";
+  if (addr) {
+    ahb_.first = *addr;
+  }
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleAXI2BusPenalty(string_list& rest) {
+  auto penalty = detail::to_uint(detail::pop_front(rest));
+  auto interval = to_interval(detail::pop_front(rest));
+  if (penalty && interval) {
+    CHECKED_CALL(ConfigureAXI2BusPenalty, *penalty, *interval);
+  }
+  return static_cast<bool>(penalty) && static_cast<bool>(interval);
+}
+
+bool SimulatorRPCChannel::HandleAXI2BusRatio(string_list& rest) {
+  auto ratio = detail::to_float(detail::pop_front(rest));
+  if (ratio) {
+    CHECKED_CALL(ConfigureAXI2BusRatio, *ratio);
+  }
+  return static_cast<bool>(ratio);
+}
+
+bool SimulatorRPCChannel::HandleAXI2HighAddr(string_list& rest) {
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  ICHECK(addr) << self_name_ << ": invalid value for AXI2 high adddress";
+  if (addr) {
+    axi2_.second = *addr;
+  }
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleAXI2LowAddr(string_list& rest) {
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  ICHECK(addr) << self_name_ << ": invalid value for AXI2 low adddress";
+  if (addr) {
+    axi2_.first = *addr;
+  }
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleBuildTag(string_list& rest) {
+  sim_->PrintBuildTag();
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleBusPenalty(string_list& rest) {
+  auto penalty = detail::to_uint(detail::pop_front(rest));
+  auto interval = to_interval(detail::pop_front(rest));
+  if (penalty && interval) {
+    CHECKED_CALL(ConfigureBusPenalty, *penalty, *interval);
+  }
+  return static_cast<bool>(penalty) && static_cast<bool>(interval);
+}
+
+bool SimulatorRPCChannel::HandleBusRatio(string_list& rest) {
+  auto ratio = detail::to_float(detail::pop_front(rest));
+  if (ratio) {
+    CHECKED_CALL(ConfigureBusRatio, *ratio);
+  }
+  return static_cast<bool>(ratio);
+}
+
+bool SimulatorRPCChannel::HandleBusTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_BUS, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleBypassIdle(string_list& rest) {
+  CHECKED_CALL(ConfigureBypassIdle, true);
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleConnectionTimeout(string_list& rest) {
+  auto time = detail::to_int(detail::pop_front(rest));
+  if (time) {
+    CHECKED_CALL(ConfigureConnectionTimeout, *time);
+  }
+  return static_cast<bool>(time);
+}
+
+bool SimulatorRPCChannel::HandleCoprocTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_COPROC, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleCoreDump(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureCoreDump, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleCosimFile(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureCosim, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleDCacheTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_DCACHE, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleDSPClock(string_list& rest) {
+  auto freq = detail::to_uint(detail::pop_front(rest));
+  if (freq) {
+    CHECKED_CALL(ConfigureCoreFrequency, *freq);
+  }
+  return static_cast<bool>(freq);
+}
+
+bool SimulatorRPCChannel::HandleETMCFGBase(string_list& rest) {
+  auto base = detail::to_uint(detail::pop_front(rest));
+  if (base) {
+    CHECKED_CALL(ConfigureEtmcfgBase, *base);
+  }
+  return static_cast<bool>(base);
+}
+
+bool SimulatorRPCChannel::HandleGDBServ(string_list& rest) {
+  auto port = detail::to_uint(detail::pop_front(rest));
+  if (port) {
+    CHECKED_CALL(ConfigureRemoteDebug, *port);
+    debug_port_ = *port;
+  }
+  return static_cast<bool>(port);
+}
+
+bool SimulatorRPCChannel::HandleHVXLength(string_list& rest) {
+  auto len = detail::to_int(detail::pop_front(rest));
+  if (len) {
+    CHECKED_CALL(ConfigureHVXLength, *len);
+  }
+  return static_cast<bool>(len);
+}
+
+bool SimulatorRPCChannel::HandleICacheTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_ICACHE, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleL2CacheTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_L2CACHE, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleL2CFGBase(string_list& rest) {
+  auto base = detail::to_uint(detail::pop_front(rest));
+  if (base) {
+    CHECKED_CALL(ConfigureL2cfgBase, *base);
+  }
+  return static_cast<bool>(base);
+}
+
+bool SimulatorRPCChannel::HandleL2TCMBase(string_list& rest) {
+  auto base = detail::to_uint(detail::pop_front(rest));
+  if (base) {
+    CHECKED_CALL(ConfigureL2tcmBase, *base);
+  }
+  return static_cast<bool>(base);
+}
+
+bool SimulatorRPCChannel::HandleMemFillRand(string_list& rest) {
+  auto seed = detail::to_uint(detail::pop_front(rest));
+  if (seed) {
+    CHECKED_CALL(ConfigureMemFillRandom, *seed);
+  }
+  return static_cast<bool>(seed);
+}
+
+bool SimulatorRPCChannel::HandleMemFill(string_list& rest) {
+  auto val = detail::to_uint(detail::pop_front(rest));
+  if (val) {
+    CHECKED_CALL(ConfigureMemFill, *val);
+  }
+  return static_cast<bool>(val);
+}
+
+bool SimulatorRPCChannel::HandleMemTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_MEM, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleNullPtr(string_list& rest) {
+  auto behavior = to_nullptr(detail::pop_front(rest));
+  if (behavior) {
+    CHECKED_CALL(ConfigureNULLPointerBehavior, *behavior);
+  }
+  return static_cast<bool>(behavior);
+}
+
+bool SimulatorRPCChannel::HandlePacketAnalyze(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigurePacketAnalysis, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandlePCFilter(string_list& rest) {
+  auto range = detail::to_range<uint64_t, detail::to_uint>(detail::pop_front(rest));
+  if (range) {
+    CHECKED_CALL(ConfigurePCRangeFilter, range->first, range->second);
+  }
+  return static_cast<bool>(range);
+}
+
+bool SimulatorRPCChannel::HandlePCTraceMin(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_PC_MIN, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandlePCTraceNano(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_PC_NANO, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandlePCTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_PC, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandlePMUStatsFile(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigurePmuStatisticsFile, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleProfile(string_list& rest) {
+  auto path = detail::pop_front(rest);
+  if (path) {
+    CHECKED_CALL(ConfigureGProf, path->c_str());
+  }
+  return static_cast<bool>(path);
+}
+
+bool SimulatorRPCChannel::HandleProfileTimeZero(string_list& rest) {
+  auto timezero = detail::to_bool(detail::pop_front(rest));
+  if (timezero) {
+    CHECKED_CALL(ConfigureProfileMode, *timezero);
+  }
+  return static_cast<bool>(timezero);
+}
+
+bool SimulatorRPCChannel::HandleQuiet(string_list& rest) {
+  sim_->VerboseMode(HEX_QUIET);
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleReconnect(string_list& rest) {
+  if (!debug_port_) {
+    LOG(FATAL) << "Reconnect error: --reconnect must be specified "
+                  "AFTER --gdbserv <port_num>";
+  }
+  CHECKED_CALL(ConfigureRemoteDebug, *debug_port_, true);
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleRTOS(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureOSAwareness, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleSimErr(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureSimStderr, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleSimIn(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureSimStdin, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleSimOut(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureSimStdout, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleStackStart(string_list& rest) {
+  auto base = detail::to_uint(detail::pop_front(rest));
+  auto size = detail::to_uint(detail::pop_front(rest));
+  if (base && size) {
+    CHECKED_CALL(ConfigureStackInfo, *base, *size);
+  }
+  return static_cast<bool>(base) && static_cast<bool>(size);
+}
+
+bool SimulatorRPCChannel::HandleStallTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_STALL, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleStatsFile(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureStatisticsFile, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleSubsystemBase(string_list& rest) {
+  auto base = detail::to_uint(detail::pop_front(rest));
+  if (base) {
+    CHECKED_CALL(ConfigureSubsystemBase, *base);
+  }
+  return static_cast<bool>(base);
+}
+
+bool SimulatorRPCChannel::HandleSymFile(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(AddSymbolFile, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleTCM(string_list& rest) {
+  CHECKED_CALL(ConfigureTimingMode, HEX_TIMING);
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleTCMHighAddr(string_list& rest) {
+  // This option takes an argument, but (the option) is ignored.
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleTCMLowAddr(string_list& rest) {
+  auto addr = detail::to_uint(detail::pop_front(rest));
+  if (addr) {
+    CHECKED_CALL(ConfigureTCM, *addr);
+  }
+  return static_cast<bool>(addr);
+}
+
+bool SimulatorRPCChannel::HandleTimeFilterNS(string_list& rest) {
+  auto range = detail::to_range<uint64_t, detail::to_uint>(detail::pop_front(rest));
+  if (range) {
+    CHECKED_CALL(ConfigureTimeRangeFilter, range->first, HEX_NANOSEC, range->second, HEX_NANOSEC);
+  }
+  return static_cast<bool>(range);
+}
+
+bool SimulatorRPCChannel::HandleTiming(string_list& rest) {
+  HEXAPI_TimingMode timing_mode = HEX_TIMING;
+  // The argument to --timing is optional.
+  if (should_parse_next(rest)) {
+    if (auto mode = to_timingmode(detail::pop_front(rest))) {
+      timing_mode = *mode;
+    } else {
+      return false;
+    }
+  }
+  CHECKED_CALL(ConfigureTimingMode, timing_mode);
+  return true;
+}
+
+bool SimulatorRPCChannel::HandleUArchTrace(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(SetTracing, HEX_TRACE_UARCH, file->c_str());
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleUseFS(string_list& rest) {
+  auto file = detail::pop_front(rest);
+  if (file) {
+    CHECKED_CALL(ConfigureARFilesystem, &(*file)[0]);
+  }
+  return static_cast<bool>(file);
+}
+
+bool SimulatorRPCChannel::HandleV2PTranslation(string_list& rest) {
+  auto enable = detail::to_bool(detail::pop_front(rest));
+  if (enable) {
+    CHECKED_CALL(EnableVirtualToPhysicalTranslation, *enable);
+  }
+  return static_cast<bool>(enable);
+}
+
+bool SimulatorRPCChannel::HandleVerbose(string_list& rest) {
+  auto mode = to_verbosemode(detail::pop_front(rest));
+  if (mode) {
+    sim_->VerboseMode(*mode);
+  }
+  return static_cast<bool>(mode);
+}
+
+bool SimulatorRPCChannel::should_parse_next(const string_list& rest) {
+  if (auto str = detail::front(rest)) {
+    return str->empty() || str->front() != '-';
+  }
+  return false;
+}
+
+detail::Optional<HEXAPI_Interval> SimulatorRPCChannel::to_interval(const detail::MaybeString& str) {
+  auto none = detail::Optional<HEXAPI_Interval>();
+  if (!str) return none;
+
+  if (auto val = detail::to_int(*str)) {
+    switch (*val) {
+      case HEX_MILLISEC:
+      case HEX_MICROSEC:
+      case HEX_NANOSEC:
+      case HEX_PICOSEC:
+      case HEX_PCYCLE:
+        return static_cast<HEXAPI_Interval>(*val);
+    }
+  }
+
+  return detail::StringSwitch<detail::Optional<HEXAPI_Interval>>(*str)
+      .Case("MILLISEC", HEX_MILLISEC)
+      .Case("MICROSEC", HEX_MICROSEC)
+      .Case("NANOSEC", HEX_NANOSEC)
+      .Case("PICOSEC", HEX_PICOSEC)
+      .Case("PCYCLE", HEX_PCYCLE)
+      .Default(none);
+}
+
+detail::Optional<HEXAPI_TimingMode> SimulatorRPCChannel::to_timingmode(
+    const detail::MaybeString& str) {
+  auto none = detail::Optional<HEXAPI_TimingMode>();
+  if (!str) return none;
+
+  if (auto val = detail::to_int(*str)) {
+    switch (*val) {
+      case HEX_NOTIMING:
+      case HEX_TIMING_NODBC:
+      case HEX_TIMING:
+      case HEX_TIMING_COHERENCY:
+        return static_cast<HEXAPI_TimingMode>(*val);
+    }
+  }
+
+  return detail::StringSwitch<detail::Optional<HEXAPI_TimingMode>>(*str)
+      .Case("NOTIMING", HEX_NOTIMING)
+      .Case("TIMING_NODBC", HEX_TIMING_NODBC)
+      .Case("TIMING", HEX_TIMING)
+      .Case("TIMING_COHERENCY", HEX_TIMING_COHERENCY)
+      .Default(none);
+}
+
+detail::Optional<HEXAPI_VerboseMode> SimulatorRPCChannel::to_verbosemode(
+    const detail::MaybeString& str) {
+  auto none = detail::Optional<HEXAPI_VerboseMode>();
+  if (!str) return none;
+
+  if (auto val = detail::to_int(*str)) {
+    switch (*val) {
+      case HEX_SILENT:
+      case HEX_QUIET:
+      case HEX_NORMAL:
+      case HEX_VERBOSE:
+      case HEX_REALLY_VERBOSE:
+        return static_cast<HEXAPI_VerboseMode>(*val);
+    }
+  }
+
+  return detail::StringSwitch<detail::Optional<HEXAPI_VerboseMode>>(*str)
+      .Case("SILENT", HEX_SILENT)
+      .Case("QUIET", HEX_QUIET)
+      .Case("NORMAL", HEX_NORMAL)
+      .Case("VERBOSE", HEX_VERBOSE)
+      .Case("REALLY_VERBOSE", HEX_REALLY_VERBOSE)
+      .Default(none);
+}
+
+detail::Optional<HEXAPI_Nullptr> SimulatorRPCChannel::to_nullptr(const detail::MaybeString& str) {
+  auto none = detail::Optional<HEXAPI_Nullptr>();
+  if (!str) return none;
+
+  if (auto val = detail::to_int(*str)) {
+    switch (*val) {
+      case HEX_NULLPTR_IGNORE:
+      case HEX_NULLPTR_WARN:
+      case HEX_NULLPTR_FATAL:
+      case HEX_NULLPTR_PCZERO:
+        return static_cast<HEXAPI_Nullptr>(*val);
+    }
+  }
+
+  return detail::StringSwitch<detail::Optional<HEXAPI_Nullptr>>(*str)
+      .Case("IGNORE", HEX_NULLPTR_IGNORE)
+      .Case("WARN", HEX_NULLPTR_WARN)
+      .Case("FATAL", HEX_NULLPTR_FATAL)
+      .Case("PCZERO", HEX_NULLPTR_PCZERO)
+      .Default(none);
+}
+
+TVM_REGISTER_GLOBAL("tvm.contrib.hexagon.create_hexagon_session")
+    .set_body([](TVMArgs args, TVMRetValue* rv) {
+      std::string session_name = args[0];
+      // For target, the second parameter is remote_stack_size_bytes, ignore it.
+      std::string sim_args = args[2];
+      auto channel = std::make_unique<SimulatorRPCChannel>(sim_args);
+      std::shared_ptr<RPCEndpoint> endpoint =
+          RPCEndpoint::Create(std::move(channel), session_name, "", nullptr);
+      std::shared_ptr<RPCSession> session = CreateClientSession(endpoint);
+      *rv = CreateRPCSessionModule(session);
+    });
+
+}  // namespace hexagon
+}  // namespace runtime
+}  // namespace tvm
diff --git a/src/runtime/logging.cc b/src/runtime/logging.cc
index ec71e4453957..8b41edbd9d44 100644
--- a/src/runtime/logging.cc
+++ b/src/runtime/logging.cc
@@ -191,6 +191,19 @@ constexpr const size_t kSrcPrefixLength = 5;
 constexpr const char* kDefaultKeyword = "DEFAULT";
 }  // namespace
 
+namespace {
+/*! \brief Convert __FILE__ to a vlog_level_map_ key, which strips any prefix ending iwth src/ */
+std::string FileToVLogMapKey(const std::string& filename) {
+  // Canonicalize the filename.
+  // TODO(mbs): Not Windows friendly.
+  size_t last_src = filename.rfind(kSrcPrefix, std::string::npos, kSrcPrefixLength);
+  // Strip anything before the /src/ prefix, on the assumption that will yield the
+  // TVM project relative filename. If no such prefix fallback to filename without
+  // canonicalization.
+  return (last_src == std::string::npos) ? filename : filename.substr(last_src + kSrcPrefixLength);
+}
+}  // namespace
+
 /* static */
 TvmLogDebugSettings TvmLogDebugSettings::ParseSpec(const char* opt_spec) {
   TvmLogDebugSettings settings;
@@ -220,6 +233,8 @@ TvmLogDebugSettings TvmLogDebugSettings::ParseSpec(const char* opt_spec) {
       return settings;
     }
 
+    name = FileToVLogMapKey(name);
+
     std::string level;
     if (!std::getline(spec_stream, level, ';')) {
       LOG(FATAL) << "TVM_LOG_DEBUG ill-formed, expecting level";
@@ -243,16 +258,8 @@ TvmLogDebugSettings TvmLogDebugSettings::ParseSpec(const char* opt_spec) {
 }
 
 bool TvmLogDebugSettings::VerboseEnabledImpl(const std::string& filename, int level) const {
-  // Canonicalize the filename.
-  // TODO(mbs): Not Windows friendly.
-  size_t last_src = filename.rfind(kSrcPrefix, std::string::npos, kSrcPrefixLength);
-  // Strip anything before the /src/ prefix, on the assumption that will yield the
-  // TVM project relative filename. If no such prefix fallback to filename without
-  // canonicalization.
-  std::string key =
-      last_src == std::string::npos ? filename : filename.substr(last_src + kSrcPrefixLength);
   // Check for exact match.
-  auto itr = vlog_level_map_.find(key);
+  auto itr = vlog_level_map_.find(FileToVLogMapKey(filename));
   if (itr != vlog_level_map_.end()) {
     return level <= itr->second;
   }
diff --git a/src/runtime/meta_data.h b/src/runtime/meta_data.h
index e83e1a3a7629..766b93261ac0 100644
--- a/src/runtime/meta_data.h
+++ b/src/runtime/meta_data.h
@@ -27,6 +27,7 @@
 #include <dmlc/io.h>
 #include <dmlc/json.h>
 #include <tvm/runtime/executor_info.h>
+#include <tvm/runtime/metadata.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/packed_func.h>
@@ -52,19 +53,19 @@ inline String get_name_mangled(const String& module_name, const String& name) {
 /*!
  * \brief Create a metadata module object.
  *
- * \param metadata The variable name to ndarray mapping.
- * \param sym_vars The symbol to the list of required constant variables
- * mapping.
+ * \param metadata Exported metadata structure.
  *
  * \return The created metadata module.
  */
-Module MetadataModuleCreate(
-    const std::unordered_map<std::string, NDArray>& metadata,
-    const std::unordered_map<std::string, std::vector<std::string>>& sym_vars);
+Module MetadataModuleCreate(metadata::Metadata metadata);
+
+namespace launch_param {
 
 /*! \brief A tag to specify whether or not dynamic shared memory is used */
 constexpr const char* kUseDynamicSharedMemoryTag = "tir.use_dyn_shared_memory";
 
+}  // namespace launch_param
+
 /*! \brief function information needed by device */
 struct FunctionInfo {
   std::string name;
diff --git a/src/runtime/metadata.cc b/src/runtime/metadata.cc
new file mode 100644
index 000000000000..d76d64fe47d2
--- /dev/null
+++ b/src/runtime/metadata.cc
@@ -0,0 +1,134 @@
+/*
+ * 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 src/runtime/metadata.cc
+ * \brief Defines implementations of TVM metadata which can exist in the runtime.
+ */
+
+#include <tvm/runtime/c_backend_api.h>
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/runtime/metadata.h>
+#include <tvm/runtime/registry.h>
+
+#include <string>
+
+namespace tvm {
+namespace runtime {
+namespace metadata {
+
+TVM_REGISTER_OBJECT_TYPE(MetadataBaseNode);
+
+ArrayAccessor<struct TVMTensorInfo, TensorInfo> MetadataNode::inputs() {
+  return ArrayAccessor<struct TVMTensorInfo, TensorInfo>(data_->inputs, data_->num_inputs);
+}
+ArrayAccessor<struct TVMTensorInfo, TensorInfo> MetadataNode::outputs() {
+  return ArrayAccessor<struct TVMTensorInfo, TensorInfo>(data_->outputs, data_->num_outputs);
+}
+ArrayAccessor<struct TVMTensorInfo, TensorInfo> MetadataNode::pools() {
+  return ArrayAccessor<struct TVMTensorInfo, TensorInfo>(data_->pools, data_->num_pools);
+}
+
+TVM_REGISTER_OBJECT_TYPE(MetadataBaseNode);
+
+MetadataArray::MetadataArray(Array<ObjectRef> array, MetadataKind kind,
+                             const char* struct_name)
+    : MetadataBase{make_object<MetadataArrayNode>(array, kind, struct_name)} {}
+
+const char* MetadataArrayNode::get_c_struct_name() const {
+  ICHECK(false) << "MetadataArrayNode get_c_struct_name is unimplemented";
+  return nullptr;
+}
+TVM_REGISTER_OBJECT_TYPE(MetadataArrayNode);
+
+Metadata::Metadata(const struct ::TVMMetadata* data)
+    : MetadataBase{make_object<MetadataNode>(data)} {}
+TVM_REGISTER_OBJECT_TYPE(MetadataNode);
+
+const char* MetadataNode::get_c_struct_name()  const { return "TVMMetadata"; }
+
+TensorInfo::TensorInfo(const struct ::TVMTensorInfo* data)
+    : MetadataBase{make_object<TensorInfoNode>(data)} {}
+TVM_REGISTER_OBJECT_TYPE(TensorInfoNode);
+
+const char* TensorInfoNode::get_c_struct_name() const { return "TVMTensorInfo"; }
+
+}  // namespace metadata
+
+class MetadataModuleNode : public ::tvm::runtime::ModuleNode {
+ public:
+  explicit MetadataModuleNode(runtime::metadata::Metadata metadata) {
+    // CHECK((metadata.defined() && code.size() > 0) || (!metadata.defined() && code.size() == 0))
+    //   << "metadata and code must both be either defined (when passed from compiler) or undefined
+    //   "
+    //   << "(when passed from runtime)";
+    metadata_ = metadata;
+    //    code_ = code;
+  }
+
+  const char* type_key() const { return "metadata_module"; }
+
+  static Module LoadFromBinary() {
+    return Module(make_object<MetadataModuleNode>(runtime::metadata::Metadata()));
+  }
+
+  void SaveToBinary(dmlc::Stream* stream) final {}
+
+  PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) {
+    if (name == "get_metadata") {
+      return PackedFunc([this, sptr_to_self](TVMArgs args, TVMRetValue* rv) {
+        if (!metadata_.defined()) {
+          TVMFunctionHandle f_handle;
+          int32_t ret_code = TVMBackendGetFuncFromEnv(this, "get_c_metadata", &f_handle);
+          CHECK_EQ(ret_code, 0) << "Unable to locate get_c_metadata PackedFunc";
+
+          TVMValue ret_value;
+          int ret_type_code;
+          ret_code = TVMFuncCall(f_handle, nullptr, nullptr, 0, &ret_value, &ret_type_code);
+          CHECK_EQ(ret_code, 0) << "Invoking get_c_metadata: TVMFuncCall returned " << ret_code;
+
+          CHECK_EQ(ret_type_code, kTVMOpaqueHandle)
+              << "Expected kOpaqueHandle returned; got " << ret_type_code;
+          CHECK(ret_value.v_handle != nullptr) << "get_c_metadata returned nullptr";
+
+          metadata_ = runtime::metadata::Metadata(
+              static_cast<const struct ::TVMMetadata*>(ret_value.v_handle));
+        }
+
+        *rv = metadata_;
+        return;
+      });
+    }
+
+    return PackedFunc();
+  }
+
+ private:
+  runtime::metadata::Metadata metadata_;
+};
+
+Module MetadataModuleCreate(metadata::Metadata metadata) {
+  return Module(make_object<MetadataModuleNode>(metadata));
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_metadata_module")
+    .set_body([](TVMArgs args, TVMRetValue* rv) { *rv = MetadataModuleNode::LoadFromBinary(); });
+
+}  // namespace runtime
+}  // namespace tvm
diff --git a/src/runtime/pipeline/pipeline_executor.cc b/src/runtime/pipeline/pipeline_executor.cc
index 2cad8cf3b060..ccf5e09ebcf1 100644
--- a/src/runtime/pipeline/pipeline_executor.cc
+++ b/src/runtime/pipeline/pipeline_executor.cc
@@ -79,8 +79,9 @@ PackedFunc PipelineExecutor::GetFunction(const std::string& name,
         [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetOutput(); });
   } else if (name == "run") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { this->Run(args[0]); });
-  } else if (name == "stop") {
-    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { this->Stop(); });
+  } else if (name == "get_execute_count") {
+    return PackedFunc(
+        [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetExecutionCount(); });
   } else {
     LOG(FATAL) << "Unknown packed function: " << name;
     return PackedFunc();
@@ -96,7 +97,6 @@ PackedFunc PipelineExecutor::GetFunction(const std::string& name,
 void PipelineExecutor::SetInput(std::string input_name, DLTensor* data_in) {
   std::pair<int, int> indexs = this->GetInputIndex(input_name);
   if (indexs.first < 0 || indexs.first >= static_cast<int>(runtimes_.size())) {
-    this->Stop();
     LOG(FATAL) << "input name " << input_name << " not found.";
   }
   runtimes_[indexs.first]->SetInput(indexs.second, data_in);
@@ -109,7 +109,6 @@ void PipelineExecutor::SetInput(std::string input_name, DLTensor* data_in) {
 NDArray PipelineExecutor::GetInput(std::string input_name) {
   std::pair<int, int> indexs = this->GetInputIndex(input_name);
   if (indexs.first < 0 || indexs.first >= static_cast<int>(runtimes_.size())) {
-    this->Stop();
     LOG(FATAL) << "input name " << input_name << " not found.";
   }
   return runtimes_[indexs.first]->GetInput(indexs.second);
@@ -153,11 +152,6 @@ void PipelineExecutor::Run(bool serialized_mode) {
  * \brief return A list of global output data.
  */
 Array<NDArray> PipelineExecutor::GetOutput(void) { return pipeline_scheduler_.PipelineGetOutput(); }
-/*!
- * \brief Stop the pipeline executor.
- */
-void PipelineExecutor::Stop() { pipeline_scheduler_.PipelineStop(); }
-
 /*!
  * \brief Use the mod_config information to create a graph runtime list.
  * \param mod_config The config information that generates by the export library function call.
@@ -242,6 +236,10 @@ std::pair<int, int> PipelineExecutor::GetInputIndex(const std::string& name) {
   auto gruntime = runtimes_[index.first];
   return std::make_pair(index.first, gruntime->GetInputIndex(index.second));
 }
+/*!
+ * \brief Getting the count of running pipeline.
+ */
+int PipelineExecutor::GetExecutionCount() { return runtimes_.back()->GetExecutionCount(); }
 /*!
  * \brief Initialize the pipeline executor with a list of modules to be pipelined
  *  and config in JSON format.
diff --git a/src/runtime/pipeline/pipeline_executor.h b/src/runtime/pipeline/pipeline_executor.h
index 7b9f5eadf92b..1d547f206a95 100644
--- a/src/runtime/pipeline/pipeline_executor.h
+++ b/src/runtime/pipeline/pipeline_executor.h
@@ -95,6 +95,10 @@ class TVM_DLL PipelineExecutor : public ModuleNode {
    * \return Return input data.
    */
   NDArray GetInput(std::string input_name);
+  /*!
+   * \brief Getting the count of running pipeline.
+   */
+  int GetExecutionCount();
   /*!
    * \brief Use the parameters group name to get the specific backend runtime then use
    *  the param_key_name to set param data for the said backend runtime.
@@ -114,10 +118,6 @@ class TVM_DLL PipelineExecutor : public ModuleNode {
    * \param serialized_mode Whether run the pipeline executor in serialized mode.
    */
   void Run(bool serialized_mode);
-  /*!
-   * \brief Stop the pipeline executor.
-   */
-  void Stop();
   /*!
    * \brief Get a list output data.
    * \return A list of output data.
diff --git a/src/runtime/pipeline/pipeline_scheduler.cc b/src/runtime/pipeline/pipeline_scheduler.cc
index 4a3368e32391..760bcd2c07a8 100644
--- a/src/runtime/pipeline/pipeline_scheduler.cc
+++ b/src/runtime/pipeline/pipeline_scheduler.cc
@@ -32,9 +32,10 @@ std::vector<std::shared_ptr<BackendRuntime>> PipelineScheduler::PipelineInit(
     const std::vector<Module>& modules, const ConfigPipelineExecution& pipeline_config) {
   std::vector<std::shared_ptr<BackendRuntime>> runtimes;
   graph_modules_ = modules;
+  // Creating a list of runtimes.
   for (size_t i = 0; i < graph_modules_.size(); i++) {
-    auto runItem = std::make_shared<BackendRuntime>(graph_modules_[i], i);
-    runtimes.push_back(runItem);
+    auto run_item = std::make_shared<BackendRuntime>(graph_modules_[i], i);
+    runtimes.push_back(run_item);
   }
   // Creating a list of NDArray in order to storage the outputs data.
   auto global_output_map = pipeline_config.GetGlobalConfigOutputBindings();
@@ -46,6 +47,10 @@ std::vector<std::shared_ptr<BackendRuntime>> PipelineScheduler::PipelineInit(
     NDArray output = runtimes[output_pair.first]->CreateFromOutput(output_pair.second);
     output_arrays_.push_back(output);
   }
+  // Initializing and then running the worker thread.
+  for (auto runtime : runtimes) {
+    runtime->InitializePipeline(pipeline_config, &runtimes);
+  }
   return runtimes;
 }
 /*!
@@ -101,18 +106,11 @@ void PipelineScheduler::PipelineRunSequential(
 void PipelineScheduler::PipelineRun(const std::vector<std::shared_ptr<BackendRuntime>>& runtimes,
                                     ConfigPipelineExecution pipeline_config, bool sequential_mode) {
   if (!sequential_mode) {
-    // TODO(huajsj) remove this check after all of pipeline features in.
-    LOG(FATAL) << "Currently only supports sequential mode.";
+    runtimes.front()->RunPipeline();
   } else {
     PipelineRunSequential(runtimes, pipeline_config);
   }
 }
-/*!
- * \brief Stop the pipeline exection.
- */
-void PipelineScheduler::PipelineStop() {
-  // TODO(huajsj) Add stop logic.
-}
 /*!
  * \brief Get a list of output.
  */
diff --git a/src/runtime/pipeline/pipeline_scheduler.h b/src/runtime/pipeline/pipeline_scheduler.h
index 6075747a6c7f..3339d4376083 100644
--- a/src/runtime/pipeline/pipeline_scheduler.h
+++ b/src/runtime/pipeline/pipeline_scheduler.h
@@ -58,10 +58,6 @@ class PipelineScheduler {
    */
   void PipelineRunSequential(const std::vector<std::shared_ptr<BackendRuntime>>& runtimes,
                              ConfigPipelineExecution pipeline_config);
-  /*!
-   * \brief Stop the pipeline exection.
-   */
-  void PipelineStop();
   /*!
    * \brief Get a list of outputs.
    */
diff --git a/src/runtime/pipeline/pipeline_struct.h b/src/runtime/pipeline/pipeline_struct.h
index 4002885f6702..05eeb221aa37 100644
--- a/src/runtime/pipeline/pipeline_struct.h
+++ b/src/runtime/pipeline/pipeline_struct.h
@@ -24,8 +24,13 @@
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/packed_func.h>
 
+#include <atomic>
+#include <condition_variable>
+#include <functional>
 #include <limits>
+#include <memory>
 #include <string>
+#include <thread>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -33,7 +38,13 @@ namespace tvm {
 namespace runtime {
 #define GLOBAL_MODULE_INDEX -1
 /*!
- *\brief The pair includes the module output index and the global output index.
+ *\brief The function is used to build the binding configuration for a runtime. The first
+ * 'int' is the output index of the current runtime, the second 'int' is the index of child
+ * runtime, and the 'string' is the input name of child runtime.
+ */
+using BindingConfigParseFunc = std::function<void(int, int, std::string)>;
+/*!
+ *\brief The 'pair' includes the module output index and the global output index.
  * The first 'int' is the module output index, and the second 'int' is the global output index.
  */
 using GlobalOutputPair = std::pair<int, int>;
@@ -42,6 +53,99 @@ using GlobalOutputPair = std::pair<int, int>;
  * The first 'int' is the module index, and the second 'int' is the module output index.
  */
 using ModuleOutputPair = std::pair<int, int>;
+/*!
+ *\brief The pair includes the runtime module index and the module input index.
+ * The first 'int' is the module index, and the second 'int' is the module input index.
+ */
+using ModuleInputPair = std::pair<int, int>;
+/*!\brief The runtime module interface type.*/
+enum InterfaceType {
+  INPUT = 0,
+  OUTPUT,
+};
+/*!
+ *\brief The structure includes the module index and the module output index.
+ */
+struct ModuleInterfaceID {
+  ModuleInterfaceID() : runtime_idx(0), runtime_interface_idx(0), interface_type(OUTPUT) { ; }
+  ModuleInterfaceID(int runtime_index, int runtime_interface_index, InterfaceType type = OUTPUT) {
+    runtime_idx = runtime_index;
+    runtime_interface_idx = runtime_interface_index;
+    interface_type = type;
+  }
+  int runtime_idx;
+  union {
+    /*!\brief The output interface index.*/
+    int runtime_output_idx;
+    /*!\brief The input interface index.*/
+    int runtime_input_idx;
+    /*!\brief The interface index.*/
+    int runtime_interface_idx;
+  };
+  /*!\brief The interface type*/
+  InterfaceType interface_type;
+};
+/*!\brief The data notification structure.*/
+class DataNotify {
+ private:
+  /*!\brief The 'contitional variable' is used to wait for notification.*/
+  std::condition_variable notify_cv_;
+  /*!\brief The mutex is used to protect the 'conditional variable'.*/
+  std::mutex mutex_;
+  /*!\brief Whether a data is ready or not.*/
+  bool data_ready_ = false;
+  /*!\brief Whether the thread should exit or not.*/
+  std::atomic<bool> exit_state_{false};
+  /*!
+   * \brief The 'ModuleInterfaceID' in which the data was ready and triggered this
+   *  notification.
+   */
+  ModuleInterfaceID notification_source_;
+
+ public:
+  /*!
+   * \brief Constructing the DataNotify class.
+   * \param parent_output_id The id of a runtime interface which is sending out the data
+   *  notification.
+   */
+  explicit DataNotify(ModuleInterfaceID parent_output_id) {
+    notification_source_ = parent_output_id;
+  }
+  /*!
+   * \brief Getting the notification source.
+   * \return The first 'int' is the runtime index, and the second 'int' is the output index.
+   */
+  ModuleInterfaceID GetNotifySource(void) { return notification_source_; }
+  /*!
+   *\brief Waiting for the notification.
+   *\return Returning the value 'false' when the notification is in a 'exit' state, else
+   * return true.
+   */
+  bool Wait(void) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    notify_cv_.wait(lock, [&] { return this->data_ready_; });
+    data_ready_ = false;
+    return !GetExitState();
+  }
+  /*!brief Sending the notification in which the related data is ready.*/
+  void Notify(void) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      data_ready_ = true;
+    }
+    notify_cv_.notify_one();
+  }
+  /*!brief Sending the notification when the notification state changes into 'exit'.*/
+  void ExitNotify(void) {
+    exit_state_.store(true, std::memory_order_release);
+    Notify();
+  }
+  /*!
+   *\brief Getting the 'exit state'.
+   *\return Returning the value of 'exit_state_'
+   */
+  bool GetExitState(void) { return exit_state_.load(std::memory_order_acquire); }
+};
 /*!
  * \brief All binding information of a output interface.
  */
@@ -53,6 +157,16 @@ class ConfigBindings {
   int GetGlobalOutputIndex() const { return global_output_index_; }
   /*!\brief Returning the binding configuration.*/
   std::unordered_map<int, std::string>& Get() { return bindings_; }
+  /*!
+   * \brief Enumerating the binding configuration.
+   * \param parse_function The function is used to parse the binding configuration.
+   * \param output_idx The index of output interface is used for parsing.
+   */
+  void VisitOutput(BindingConfigParseFunc parse_function, int output_idx) {
+    for (auto output : bindings_) {
+      parse_function(output_idx, output.first, output.second);
+    }
+  }
   /*!
    * \brief Create a module interface map from JSONReader.
    * \param reader JSON reader.
@@ -118,6 +232,15 @@ class ConfigOutputBindings {
     ICHECK(output_binding_map_.find(key) != output_binding_map_.end());
     return output_binding_map_[key];
   }
+  /*!
+   * \brief Enumerating the output configuration.
+   * \param parse_function The callback function is used to parse the binding configeration.
+   */
+  void VisitOutputConfig(BindingConfigParseFunc parse_function) {
+    for (auto output : output_binding_map_) {
+      output.second.VisitOutput(parse_function, output.first);
+    }
+  }
   /*!brief Return the variable "output_binding_map_".*/
   std::unordered_map<int, ConfigBindings> GetOutBindings() const { return output_binding_map_; }
   /*!
@@ -189,6 +312,18 @@ class ConfigPipelineExecution {
     ICHECK(config_.find(key) != config_.end());
     return config_[key];
   }
+  /*!
+   * \brief Enumerating the binding configuration for a specified runtime.
+   * \param parse_function The callback function is used to parse the binding configuration.
+   * \param runtime_index The index of a runtime is used to parse the binding configuration.
+   */
+  void VisitRuntimeOutputConfig(BindingConfigParseFunc parse_function, int runtime_index) {
+    auto config = config_.find(runtime_index);
+    if (config == config_.end()) {
+      LOG(FATAL) << "Do not finding the runtime " << runtime_index;
+    }
+    config->second.VisitOutputConfig(parse_function);
+  }
   /*
    *!\brief This function is used to verify whether config is loaded successfully.
    * \return Return "true" to indicate that this class has not been successfully loaded.
@@ -367,11 +502,21 @@ struct ParamConnectionConfig {
  *\brief Backend Runtime.
  */
 class BackendRuntime {
+  using ModuleInputPairList = std::vector<std::pair<std::shared_ptr<BackendRuntime>, int>>;
+
  private:
   /*\brief The index of runtime indicates the runtime position in the pipeline.*/
   int runtime_idx_;
   /*\brief The Runtime module of a backend graph executor.*/
   Module module_;
+  /*\brief The thread is associated with the current runtime*/
+  std::thread thread_;
+  /*\brief A list of runtime which depends on the current runtime.*/
+  std::unordered_map<int, ModuleInputPairList> children_;
+  /*\brief A map including the runtime input index and the notification data structure.*/
+  std::unordered_map<int, std::shared_ptr<DataNotify>> parents_notify_;
+  /*\brief The times of using pipeline function. */
+  uint32_t statistic_pipeline_execute_times_ = 0;
   /*!
    *\brief In order to transfer data from one backend runtime to another, we need a local
    * tensor variable as a medium. "input_tensor_local_copy_" is a map including
@@ -386,6 +531,101 @@ class BackendRuntime {
   tvm::runtime::PackedFunc get_num_inputs_;
   tvm::runtime::PackedFunc get_input_index_;
   tvm::runtime::PackedFunc run_;
+  /*!\brief The worker thread is used to execute the runtimes in pipeline.*/
+  void StartWorkThread() {
+    if (runtime_idx_ == 0) {
+      this->CreateParentsNotify(0, GLOBAL_MODULE_INDEX, 0);
+    } else {
+      // Only launching the worker thread for the runtimes after the first runtime.
+      thread_ = std::thread([&]() {
+        while (!this->WaitAndLoadPipelineData()) {
+          this->RunPipeline();
+        }
+        VLOG(1) << "Runtime " << this->runtime_idx_ << " exit.";
+      });
+    }
+    return;
+  }
+  /*!\brief Stopping the threads in pipeline.*/
+  void StopPipeline() {
+    for (auto notify : parents_notify_) {
+      notify.second->ExitNotify();
+    }
+    if (thread_.joinable()) {
+      thread_.join();
+    }
+  }
+  /*!
+   * \brief Waiting for the internal forwarding data.
+   * \return Returning 'true' when getting a 'exit' notification otherwise returning 'false'.
+   */
+  bool WaitAndLoadPipelineData() {
+    std::unordered_map<int, std::shared_ptr<DataNotify>> notifys = parents_notify_;
+    bool exit_notify = false;
+    while (!notifys.empty() && !exit_notify) {
+      auto notify = notifys.begin();
+      // Breaking the loop when the notification is in the exit state.
+      if ((exit_notify = notify->second->GetExitState())) break;
+      // Getting the source which sends this notification.
+      auto notify_source = notify->second->GetNotifySource();
+      // Loading the binding data.
+      while (!this->LoadBindingData(notify->first, notify_source.runtime_idx,
+                                    notify_source.runtime_output_idx)) {
+        // Waiting for the notification.
+        if (!notify->second->Wait()) {
+          VLOG(1) << "runtime index:" << runtime_idx_ << " receive exit notify.";
+          exit_notify = true;
+          break;
+        }
+        // TODO(huajsj): removing this 'break' after finishing the 'LoadBindingData'.
+        break;
+      }
+      VLOG(1) << "runtime_index.input_index:" << runtime_idx_ << "." << notify->first
+              << "from runtime_index.output_index:" << notify_source.runtime_idx << "."
+              << notify_source.runtime_output_idx;
+      notifys.erase(notify);
+    }
+    return exit_notify;
+  }
+  /*!
+   * \brief Loading the binding data.
+   * \param parent_idx The index of runtime which forwards data to current runtime.
+   * \param parent_output_idx The index of output where the forwarding data is coming from.
+   * \param input_idx The index of input where the data will be forwarding to.
+   * \return Returning 'true' when data is loaded successfully, otherwise returning 'false'.
+   */
+  bool LoadBindingData(int parent_idx, int parent_output_idx, int input_idx) {
+    // TODO(huajsj): Loading data.
+    return false;
+  }
+  /*!
+   * \brief Forwarding the output data into the child runtimes.
+   */
+  void ForwardingOutputDataToChildren(void) {
+    for (auto child : children_) {
+      // TODO(huajsj): Getting the output data from the current runtime in order to forward
+      // data to the child.
+
+      // Notifying the 'children runtime' that the forwarding data are ready.
+      for (auto module_pair : child.second) {
+        module_pair.first->ParentNotify(module_pair.second);
+      }
+    }
+  }
+  /*!
+   *\brief Creating a parent notification.
+   *\param input_index The input index of the 'current runtime'.
+   *\param parent_idx The index of 'parent runtime' which will send the notification.
+   *\param parent_output_idx The output index of the 'parent runtime' which will send
+   * the nofication.
+   */
+  void CreateParentsNotify(int input_index, int parent_idx, int parent_output_idx) {
+    if (parents_notify_.find(input_index) != parents_notify_.end()) {
+      LOG(FATAL) << "Not finding the input index " << input_index << " in runtime " << runtime_idx_;
+    }
+    parents_notify_[input_index] =
+        std::make_shared<DataNotify>(ModuleInterfaceID(parent_idx, parent_output_idx));
+  }
   /*!
    * \brief Copying from a given tensor and using 'CPU' as the device.
    */
@@ -439,11 +679,60 @@ class BackendRuntime {
     get_output_ = module_.GetFunction("get_output");
     run_ = module_.GetFunction("run");
   }
-  BackendRuntime(void) {}
   ~BackendRuntime() {
     for (auto data : input_tensor_local_copy_) {
       TVMArrayFree(data.second);
     }
+    StopPipeline();
+  }
+  /*!brief Getting the runtime index*/
+  int GetIndex() const { return runtime_idx_; }
+  /*!
+   * \brief Getting the times of using pipeline function.
+   * \return The times of using pipeline function.
+   */
+  int GetExecutionCount() const { return statistic_pipeline_execute_times_; }
+  /*!
+   * \brief Initializing data structures for the pipeline execution.
+   * \param config The pipeline configueration.
+   * \param runtimes A list of BackendRuntime.
+   */
+  void InitializePipeline(ConfigPipelineExecution config,
+                          std::vector<std::shared_ptr<BackendRuntime>>* runtimes) {
+    // Getting the 'binding configuration' for each runtime.
+    config.VisitRuntimeOutputConfig(
+        [&](int output_idx, int child_idx, std::string child_input_name) {
+          int runtime_idx_max = runtimes->size();
+          if (child_idx < 0 || child_idx >= runtime_idx_max) {
+            LOG(FATAL) << "The runtime index " << child_idx << " is out of the range.";
+          }
+          auto child_runtime = runtimes->at(child_idx);
+          int input_index = child_runtime->GetInputIndex(child_input_name);
+          if (input_index < 0) {
+            LOG(FATAL) << "Can not find the input " << input_index << "in runtime " << child_idx;
+          }
+          children_[output_idx].push_back(std::make_pair(child_runtime, input_index));
+          child_runtime->CreateParentsNotify(input_index, runtime_idx_, output_idx);
+          VLOG(1) << " parent_idx.output:" << runtime_idx_ << "." << output_idx << " child.input"
+                  << child_idx << "." << input_index;
+        },
+        runtime_idx_);
+
+    StartWorkThread();
+  }
+  /*!
+   * \brief Notifying a input is ready.
+   * \param input_index The index of 'input interface' which is ready for data.
+   */
+  void ParentNotify(int input_index) {
+    auto notify = parents_notify_.find(input_index);
+    if (notify == parents_notify_.end()) {
+      LOG(FATAL) << "Can not find the input for index " << input_index << " in runtime"
+                 << runtime_idx_;
+      return;
+    }
+    notify->second->Notify();
+    VLOG(1) << "Notification at runtime_index.input_index:" << runtime_idx_ << "." << input_index;
   }
   /*!\brief Creating a NDArray containing same shape and data type with a module output. */
   NDArray CreateFromOutput(int idx) {
@@ -475,6 +764,12 @@ class BackendRuntime {
   NDArray GetOutput(int index) { return get_output_(index); }
   /*!\brief Running the runtime.*/
   void Run() { run_(); }
+  /*!\brief Running the runtime in the pipeline mode.*/
+  void RunPipeline() {
+    Run();
+    ForwardingOutputDataToChildren();
+    statistic_pipeline_execute_times_++;
+  }
 };
 /*!
  * \brief The information used to initialize the graph executor module, the information
diff --git a/src/runtime/thread_storage_scope.h b/src/runtime/thread_storage_scope.h
index d577770db1a9..4122f9d0798e 100644
--- a/src/runtime/thread_storage_scope.h
+++ b/src/runtime/thread_storage_scope.h
@@ -24,6 +24,7 @@
 #ifndef TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
 #define TVM_RUNTIME_THREAD_STORAGE_SCOPE_H_
 
+#include <tvm/runtime/metadata.h>
 #include <tvm/runtime/packed_func.h>
 
 #include <string>
@@ -205,7 +206,7 @@ class LaunchParamConfig {
     std::vector<bool> filled(6, false);
     for (size_t i = 0; i < launch_param_tags.size(); ++i) {
       const std::string& tag = launch_param_tags[i];
-      if (tag == kUseDynamicSharedMemoryTag) {
+      if (tag == launch_param::kUseDynamicSharedMemoryTag) {
         ICHECK_EQ(i, launch_param_tags.size() - 1)
             << "kUseDynamicSharedMemoryTag should be the last tag in launch_param_tags.";
         use_dyn_shared_memory_ = true;
diff --git a/src/runtime/vm/bytecode.cc b/src/runtime/vm/bytecode.cc
index 8f7a3776aa16..424dfe87c79b 100644
--- a/src/runtime/vm/bytecode.cc
+++ b/src/runtime/vm/bytecode.cc
@@ -509,6 +509,9 @@ void DLDatatypePrint(std::ostream& os, const DLDataType& dtype) {
     case kDLFloat:
       os << "float";
       break;
+    case kDLBfloat:
+      os << "bfloat";
+      break;
   }
 
   os << int(dtype.bits);
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index dc8f572a13dc..a003367c3724 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -109,27 +109,20 @@ PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Obje
   }
 }
 
-int Executable::GetFunctionArity(std::string func_name) const {
+const VMFunction& Executable::GetVMFunctionWithName(const std::string& func_name) const {
   auto it = global_map.find(func_name);
-  if (it == global_map.end()) {
-    LOG(ERROR) << "Cannot find function " << func_name << " in executable";
-    return -1;
-  }
-  const auto& func = functions[it->second];
+  ICHECK(it != global_map.end()) << "Cannot find function " << func_name << " in executable";
+  return functions[it->second];
+}
+
+int Executable::GetFunctionArity(std::string func_name) const {
+  const auto& func = GetVMFunctionWithName(func_name);
   return func.params.size();
 }
 
 std::string Executable::GetFunctionParameterName(std::string func_name, uint32_t index) const {
-  auto it = global_map.find(func_name);
-  if (it == global_map.end()) {
-    LOG(ERROR) << "Cannot find function " << func_name << " in executable";
-    return "";
-  }
-  const auto& func = functions[it->second];
-  if (index > func.params.size()) {
-    LOG(ERROR) << "Invalid parameter index";
-    return "";
-  }
+  const auto& func = GetVMFunctionWithName(func_name);
+  ICHECK_LT(index, func.params.size()) << "Invalid parameter index";
   return func.params[index];
 }
 
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index f0d0a44908ae..51114681dd57 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -190,17 +190,7 @@ PackedFunc VirtualMachine::GetFunction(const std::string& name,
   } else if (name == "get_input_index") {
     return TypedPackedFunc<int64_t(std::string, std::string)>(
         [this](std::string input_name, std::string func_name) {
-          auto gvit = exec_->global_map.find(func_name);
-          ICHECK(gvit != exec_->global_map.end()) << "Cannot find function " << func_name;
-          auto func_index = gvit->second;
-          const auto& vm_func = exec_->functions[func_index];
-          const auto& param_names = vm_func.params;
-          for (uint64_t i = 0; i < param_names.size(); i++) {
-            if (input_name == param_names[i]) {
-              return static_cast<int64_t>(i);
-            }
-          }
-          return static_cast<int64_t>(-1);
+          return GetInputIndexFromVMFunction(func_name, input_name);
         });
   } else if (name == "init") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
@@ -221,6 +211,12 @@ PackedFunc VirtualMachine::GetFunction(const std::string& name,
   } else if (name == "set_input") {
     return PackedFunc(
         [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { SetInput(args[0], args, 1); });
+  } else if (name == "set_one_input") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      ICHECK_EQ(args.size(), 3) << "The expected number of arguments is 3 "
+                                << "(func_name, index or name, tensor)";
+      SetOneInput(args[0], args[1], args[2]);
+    });
   } else if (name == "load_late_bound_consts") {
     return PackedFunc([this](TVMArgs args, TVMRetValue* rv) {
       CHECK_EQ(args.size(), 1);
@@ -234,39 +230,91 @@ PackedFunc VirtualMachine::GetFunction(const std::string& name,
 }
 
 void VirtualMachine::SetInput(std::string func_name, TVMArgs args, int offset) {
-  ICHECK(exec_) << "The executable is not created yet.";
-  auto gvit = exec_->global_map.find(func_name);
-  ICHECK(gvit != exec_->global_map.end()) << "Cannot find function " << func_name;
-  auto func_index = gvit->second;
-  const auto& vm_func = exec_->functions[func_index];
-  const auto& param_names = vm_func.params;
-  ICHECK_EQ(args.size() - offset, param_names.size())
+  const auto& vm_func = CheckAndGetVMFunction(func_name);
+  size_t params_num = vm_func.params.size();
+  ICHECK_EQ(args.size() - offset, params_num)
       << "The number of provided parameters doesn't match the number of arguments";
-  ICHECK_EQ(param_names.size(), vm_func.param_device_indexes.size())
-      << "The number of provided parameters doesn't match the number of assigned devices";
-  std::vector<ObjectRef> func_args(param_names.size());
+  std::vector<ObjectRef> func_args(params_num);
   for (int i = offset; i < args.size(); ++i) {
-    Device dev = GetDevice(vm_func.param_device_indexes[i - offset]);
-
-    if (args[i].type_code() == kTVMDLTensorHandle) {
-      // Automatically convert input DLTensors to NDArray
-      DLTensor* tensor = args[i];
-      std::vector<int64_t> shape;
-      for (int64_t i = 0; i < tensor->ndim; i++) {
-        shape.push_back(tensor->shape[i]);
-      }
-      NDArray ary = NDArray::Empty(shape, tensor->dtype, dev);
-      ary.CopyFrom(tensor);
-      func_args[i - offset] = ary;
-    } else {
-      ObjectRef obj = CopyTo(args[i], dev);
-      func_args[i - offset] = obj;
-    }
+    int index = i - offset;
+    Device dev = GetDevice(vm_func.param_device_indexes[index]);
+    SetInputTensorWithIndex(func_args, args[i], index, dev);
   }
   inputs_.erase(func_name);
   inputs_.emplace(func_name, func_args);
 }
 
+void VirtualMachine::SetOneInput(std::string func_name, const TVMArgValue& tag,
+                                 const TVMArgValue& tensor) {
+  const auto& vm_func = CheckAndGetVMFunction(func_name);
+  size_t params_num = vm_func.params.size();
+
+  int inp_index;
+  if (tag.type_code() == kTVMArgInt) {
+    inp_index = tag;
+  } else if (tag.type_code() == kTVMStr) {
+    inp_index = static_cast<int>(GetInputIndexFromName(vm_func.params, tag));
+  } else {
+    LOG(FATAL) << "The type of input tensor tag (" << tag.type_code()
+               << ") doesn't match integer or string";
+  }
+  ICHECK_LT(inp_index, params_num);
+
+  CreateInputsOrCheckSize(func_name, params_num);
+  Device dev = GetDevice(vm_func.param_device_indexes[inp_index]);
+  SetInputTensorWithIndex(inputs_[func_name], tensor, inp_index, dev);
+}
+
+int64_t VirtualMachine::GetInputIndexFromVMFunction(const std::string& func_name,
+                                                    const std::string& input_name) const {
+  const auto& vm_func = CheckAndGetVMFunction(func_name);
+  return GetInputIndexFromName(vm_func.params, input_name);
+}
+
+int64_t VirtualMachine::GetInputIndexFromName(const std::vector<std::string>& params,
+                                              const std::string& input_name) const {
+  // TODO(vvchernov): excess integer type?
+  for (uint64_t i = 0; i < params.size(); i++) {
+    if (input_name == params[i]) {
+      return static_cast<int64_t>(i);
+    }
+  }
+  return static_cast<int64_t>(-1);
+}
+
+const VMFunction& VirtualMachine::CheckAndGetVMFunction(const std::string& func_name) const {
+  ICHECK(exec_) << "The executable is not created yet.";
+  return exec_->GetVMFunctionWithName(func_name);
+}
+
+void VirtualMachine::CreateInputsOrCheckSize(const std::string& func_name, size_t size) {
+  if (inputs_.count(func_name)) {
+    ICHECK_EQ(inputs_[func_name].size(), size)
+        << "The size of function" << func_name
+        << " doesn't match the number of provided parameters";
+  } else {
+    std::vector<ObjectRef> func_args(size);
+    inputs_.emplace(func_name, func_args);
+  }
+}
+
+void VirtualMachine::SetInputTensorWithIndex(std::vector<ObjectRef>& tensors,
+                                             const TVMArgValue& inp_tensor, int index, Device dev) {
+  if (inp_tensor.type_code() == kTVMDLTensorHandle) {
+    // Automatically convert input DLTensors to NDArray
+    DLTensor* tensor = inp_tensor;
+    std::vector<int64_t> shape;
+    for (int64_t i = 0; i < tensor->ndim; i++) {
+      shape.push_back(tensor->shape[i]);
+    }
+    NDArray ary = NDArray::Empty(shape, tensor->dtype, dev);
+    ary.CopyFrom(tensor);
+    tensors[index] = ary;
+  } else {
+    tensors[index] = CopyTo(inp_tensor, dev);
+  }
+}
+
 inline Device VirtualMachine::GetDevice(Index device_index) const {
   ICHECK_GE(devices_.size(), device_index) << "invalid device index: " << device_index;
   return devices_[device_index];
diff --git a/src/target/build_common.h b/src/target/build_common.h
index c66c2b52822e..6c94ec8703b7 100644
--- a/src/target/build_common.h
+++ b/src/target/build_common.h
@@ -58,7 +58,7 @@ inline std::unordered_map<std::string, runtime::FunctionInfo> ExtractFuncInfo(co
     }
     if (auto opt = f->GetAttr<Integer>(tir::attr::kDeviceUseDynSharedMemory)) {
       if (opt.value()) {
-        info.launch_param_tags.push_back(runtime::kUseDynamicSharedMemoryTag);
+        info.launch_param_tags.push_back(runtime::launch_param::kUseDynamicSharedMemoryTag);
       }
     }
     auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
diff --git a/src/target/llvm/codegen_amdgpu.cc b/src/target/llvm/codegen_amdgpu.cc
index 33a09b1ded66..9e5103ef7f39 100644
--- a/src/target/llvm/codegen_amdgpu.cc
+++ b/src/target/llvm/codegen_amdgpu.cc
@@ -80,7 +80,7 @@ class CodeGenAMDGPU : public CodeGenLLVM {
       buf = AllocateSharedMemory(op->dtype, 0, 3, std::min(info.alignment, 16),
                                  llvm::GlobalValue::ExternalLinkage);
     } else {
-      int32_t constant_size = op->constant_allocation_size();
+      size_t constant_size = op->ConstantAllocationSize();
       ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation in GPU";
 
       if (constant_size % 4 == 0 && info.alignment == 0) {
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 6d9d98072ee6..32314cf68176 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -22,8 +22,6 @@
  */
 #ifdef TVM_LLVM_VERSION
 
-#include "codegen_cpu.h"
-
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/tir/analysis.h>
 
@@ -32,6 +30,8 @@
 #include <unordered_map>
 
 #include "../func_registry_generator.h"
+#include "../metadata_utils.h"
+#include "codegen_cpu.h"
 
 namespace tvm {
 namespace codegen {
@@ -72,11 +72,11 @@ void CodeGenCPU::Init(const std::string& module_name, llvm::TargetMachine* tm,
   // typedef int (*TVMBackendPackedCFunc)(TVMValue* args, int* type_codes, int num_args,
   //                                      TVMValue* out_ret_value, int* out_ret_tcode,
   //                                      void* resource_handle);
-  ftype_tvm_backend_packed_c_func_ = llvm::FunctionType::get(
-      t_int_,
-      {t_tvm_func_handle_, t_tvm_value_->getPointerTo(), t_int_->getPointerTo(), t_int_,
-       t_tvm_value_->getPointerTo(), t_int_->getPointerTo(), t_void_p_},
-      false);
+  ftype_tvm_backend_packed_c_func_ =
+      llvm::FunctionType::get(t_int_,
+                              {t_void_p_, t_void_p_, t_int_,
+                               t_void_p_, t_void_p_, t_void_p_},
+                              false);
   t_tvm_crt_func_registry_ = llvm::StructType::create(
       {t_char_->getPointerTo(), ftype_tvm_backend_packed_c_func_->getPointerTo()});
   t_tvm_crt_module_ = llvm::StructType::create({t_tvm_crt_func_registry_->getPointerTo()});
@@ -798,10 +798,14 @@ llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) {
 
 CodeGenCPU::PackedCall CodeGenCPU::MakeCallPackedLowered(const Array<PrimExpr>& args,
                                                          const DataType& r_type,
-                                                         const int64_t begin, const int64_t end) {
+                                                         const int64_t begin, const int64_t end,
+                                                         bool use_string_lookup) {
   PackedCall pc;
   std::string func_name = args[0].as<StringImmNode>()->value;
-  llvm::Value* handle = GetPackedFuncHandle(func_name);
+  llvm::Value* handle = nullptr;
+  if (use_string_lookup) {
+    handle = GetPackedFuncHandle(func_name);
+  }
   // call the function
   int64_t nargs = end - begin;
   ICHECK_GE(nargs, 0);
@@ -816,14 +820,45 @@ CodeGenCPU::PackedCall CodeGenCPU::MakeCallPackedLowered(const Array<PrimExpr>&
       ConstInt32(end));
   TypedPointer ret_tcode = CreateBufferPtr(DataType::Int(32), stack_tcode, ConstInt32(end));
 
+  llvm::FunctionType* callee_ftype = nullptr;
+  llvm::Value* callee_value = nullptr;
+  std::vector<llvm::Value*> call_args;
+
+  if (use_string_lookup) {
+    callee_ftype = ftype_tvm_func_call_;
+    callee_value = RuntimeTVMFuncCall();
+    call_args.push_back(handle);
+  } else {
+    callee_ftype = ftype_tvm_backend_packed_c_func_;
+    callee_value = module_->getFunction(func_name);
+    if (callee_value == nullptr) {
+      callee_value =
+          llvm::Function::Create(ftype_tvm_backend_packed_c_func_, llvm::Function::ExternalLinkage,
+                                 func_name, module_.get());
+    }
+  }
+
+  if (use_string_lookup) {
+    call_args.insert(call_args.end(),
+                     {arg_value, arg_tcode.addr, ConstInt32(nargs), ret_value, ret_tcode.addr});
+  } else {
+    nargs -= 1;
+    call_args.insert(call_args.end(), {
+                                          builder_->CreateBitCast(arg_value, t_void_p_),
+                                          builder_->CreateBitCast(arg_tcode.addr, t_void_p_),
+                                          ConstInt32(nargs),
+                                          builder_->CreateBitCast(ret_value, t_void_p_),
+                                          builder_->CreateBitCast(ret_tcode.addr, t_void_p_),
+                                      });
+    call_args.push_back(llvm::ConstantPointerNull::get(t_void_p_));
+  }
 #if TVM_LLVM_VERSION >= 90
-  auto call_callee = llvm::FunctionCallee(ftype_tvm_func_call_, RuntimeTVMFuncCall());
+  auto call_callee = llvm::FunctionCallee(callee_ftype, callee_value);
 #else
-  auto call_callee = RuntimeTVMFuncCall();
+  auto call_callee = callee_value;
 #endif
-  llvm::Value* call = builder_->CreateCall(
-      call_callee,
-      {handle, arg_value, arg_tcode.addr, ConstInt32(nargs), ret_value, ret_tcode.addr});
+  llvm::Value* call = builder_->CreateCall(call_callee, call_args);
+
   llvm::BasicBlock* end_block = CheckCallSuccess(call);
 
   // Load the return value and cast it to the designated type (r_type).
@@ -852,17 +887,19 @@ CodeGenCPU::PackedCall CodeGenCPU::MakeCallPackedLowered(const Array<PrimExpr>&
   return pc;
 }
 
-llvm::Value* CodeGenCPU::CreateCallPacked(const CallNode* op) {
-  ICHECK_EQ(op->args.size(), 5U);
+llvm::Value* CodeGenCPU::CreateCallPacked(const CallNode* op, bool use_string_lookup) {
+  LOG(INFO) << "CreateCallPacked: " << GetRef<Call>(op);
+  auto expected_num_args = use_string_lookup ? 5U : 6U;
+  ICHECK_EQ(op->args.size(), expected_num_args);
   PackedCall pc = MakeCallPackedLowered(op->args, op->dtype, op->args[3].as<IntImmNode>()->value,
-                                        op->args[4].as<IntImmNode>()->value);
+                                        op->args[4].as<IntImmNode>()->value, use_string_lookup);
   return pc.ret_value;
 }
 
 llvm::Value* CodeGenCPU::CreateCallTracePacked(const CallNode* op) {
   ICHECK_EQ(op->args.size(), 6U);
   PackedCall pc = MakeCallPackedLowered(op->args, op->dtype, op->args[3].as<IntImmNode>()->value,
-                                        op->args[4].as<IntImmNode>()->value);
+                                        op->args[4].as<IntImmNode>()->value, true);
   // Get traced value.
   llvm::Value* traced_value = MakeValue(op->args[5]);
   // The update_block handles case when we need to update the return value.
@@ -908,6 +945,328 @@ llvm::Value* CodeGenCPU::RuntimeTVMParallelBarrier() {
   return GetContextPtr(gv_tvm_parallel_barrier_);
 }
 
+/*! \brief Defines LLVM Types for each Metadata member type. */
+struct MetadataLlvmTypes {
+  llvm::Type* t_float64;
+  llvm::Type* t_uint8;
+  llvm::Type* t_int64;
+  llvm::Type* t_bool;
+  llvm::Type* t_cstring;
+  llvm::Type* t_void_p;
+  llvm::StructType* t_data_type;
+
+  /*! \brief Maps a MetadataBase subclass' type_key to its corresponding LLVM StructType. */
+  ::std::unordered_map<std::string, llvm::StructType*> structs_by_type_key;
+};
+
+template <typename T>
+void PrintLLVM(const T* llvm_obj) {
+  std::string modpe_str;
+  llvm::raw_string_ostream rso(modpe_str);
+  llvm_obj->print(rso);
+  LOG(INFO) << rso.str();
+}
+
+class MetadataTypeDefiner : public AttrVisitor {
+public:
+  MetadataTypeDefiner(llvm::LLVMContext* ctx, struct MetadataLlvmTypes* llvm_types) : ctx_{ctx}, llvm_types_{llvm_types} {}
+
+  void Visit(const char* key, double* value) final {
+    elements_.emplace_back(llvm_types_->t_float64);
+  }
+  void Visit(const char* key, int64_t* value) final {
+    elements_.emplace_back(llvm_types_->t_int64);
+  }
+  void Visit(const char* key, uint64_t* value) final {
+    elements_.emplace_back(llvm_types_->t_int64);
+  }
+  void Visit(const char* key, int* value) final {
+    elements_.emplace_back(llvm_types_->t_int64);
+  }
+  void Visit(const char* key, bool* value) final {
+    elements_.emplace_back(llvm_types_->t_bool);
+  }
+  void Visit(const char* key, std::string* value) final {
+    elements_.emplace_back(llvm_types_->t_cstring);
+  }
+  void Visit(const char* key, void** value) final {
+    elements_.emplace_back(llvm_types_->t_void_p);
+  }
+  void Visit(const char* key, DataType* value) final {
+    elements_.emplace_back(llvm_types_->t_data_type);
+  }
+  void Visit(const char* key, runtime::NDArray* value) final {
+    CHECK(false) << "Do not support serializing NDArray";
+  }
+
+private:
+  void VisitMetadataBase(runtime::metadata::MetadataBase metadata) {
+    elements_.emplace_back(llvm::PointerType::getUnqual(llvm::StructType::create(*ctx_, metadata->get_c_struct_name())));
+    if (visited_.find(metadata->get_c_struct_name()) != visited_.end()) {
+      return;
+    }
+
+    if (to_visit_.find(metadata->get_c_struct_name()) != to_visit_.end()) {
+      return;
+    }
+    to_visit_[metadata->get_c_struct_name()] = metadata;
+  }
+
+public:
+  using MetadataKind = runtime::metadata::MetadataKind;
+
+  void VisitArray(const runtime::metadata::MetadataArrayNode* arr) {
+    switch (arr->kind) {
+    case MetadataKind::kUint64:  // LLVM encodes signed and unsigned with same types.
+    case MetadataKind::kInt64:
+      elements_.emplace_back(llvm::PointerType::getUnqual(llvm_types_->t_int64));
+      break;
+    case MetadataKind::kBool:
+      elements_.emplace_back(llvm::PointerType::getUnqual(llvm_types_->t_bool));
+      break;
+    case MetadataKind::kString:
+      elements_.emplace_back(llvm::PointerType::getUnqual(llvm_types_->t_cstring));
+      break;
+    case MetadataKind::kHandle:
+      CHECK(false) << "Do not support handle";
+      break;
+    case MetadataKind::kMetadata:
+      elements_.emplace_back(llvm::PointerType::getUnqual(llvm_types_->structs_by_type_key[arr->type_key]));
+      break;
+    default:
+      CHECK(false) << "Unsupported metadata kind " << arr->kind;
+      break;
+    }
+  }
+
+  void Visit(const char* key, ObjectRef* value) final {
+    const runtime::metadata::MetadataArrayNode* arr =
+        value->as<runtime::metadata::MetadataArrayNode>();
+    if (arr != nullptr) {
+      VisitArray(arr);
+    } else {
+      elements_.emplace_back(llvm::PointerType::getUnqual(llvm_types_->structs_by_type_key[(*value)->GetTypeKey()]));
+    }
+  }
+
+  void DefineType(runtime::metadata::MetadataBase metadata) {
+    ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+    LOG(INFO) << "Created type for " << metadata->GetTypeKey() << ":";
+    for (auto e : elements_) {
+      std::string value;
+      llvm::raw_string_ostream os(value);
+      e->print(os, true);
+//      LOG(INFO) << " - " << e << ", tyid=" << e->getTypeID() << " == " << value;
+//      e->dump();
+    }
+    llvm_types_->structs_by_type_key[metadata->GetTypeKey()] = llvm::StructType::create(*ctx_, elements_, metadata->get_c_struct_name());
+    elements_.clear();
+  }
+
+  llvm::LLVMContext* ctx_;
+  struct MetadataLlvmTypes* llvm_types_;
+  ::std::unordered_set<::std::string> visited_;
+  ::std::unordered_map<::std::string, runtime::metadata::MetadataBase> to_visit_;
+  ::std::vector<llvm::Type*> elements_;
+};
+
+
+class MetadataSerializerLLVM : public AttrVisitor {
+  using MetadataKind = runtime::metadata::MetadataKind;
+
+ public:
+  MetadataSerializerLLVM(CodeGenLLVM* codegen, struct MetadataLlvmTypes* llvm_types)
+      : codegen_{codegen}, llvm_types_{llvm_types} {}
+
+  void Visit(const char* key, double* value) final {
+    elements_.back().emplace_back(llvm::ConstantFP::get(llvm_types_->t_float64, *value));
+  }
+  void Visit(const char* key, int64_t* value) final {
+    elements_.back().emplace_back(llvm::ConstantInt::get(
+        llvm_types_->t_int64, static_cast<uint64_t>(*value), true /* isSigned */));
+  }
+  void Visit(const char* key, uint64_t* value) final {
+    elements_.back().emplace_back(
+        llvm::ConstantInt::get(llvm_types_->t_int64, *value, false /* isSigned */));
+  }
+  void Visit(const char* key, int* value) final {
+    elements_.back().emplace_back(
+        llvm::ConstantInt::get(llvm_types_->t_int64, *value, true /* isSigned */));
+  }
+  void Visit(const char* key, bool* value) final {
+    elements_.back().emplace_back(llvm::ConstantInt::get(
+        llvm_types_->t_uint8, static_cast<uint64_t>(*value), false /* isSigned */));
+  }
+  void Visit(const char* key, std::string* value) final {
+    elements_.back().emplace_back(codegen_->GetConstString(*value));
+  }
+  void Visit(const char* key, void** value) final {
+    CHECK(false) << "Do not support serializing void*";
+  }
+  void Visit(const char* key, DataType* value) final {
+    elements_.back().emplace_back(llvm::ConstantStruct::get(
+        llvm_types_->t_data_type,
+        {llvm::ConstantInt::get(llvm_types_->t_uint8, value->code(), false /* isSigned */),
+         llvm::ConstantInt::get(llvm_types_->t_uint8, value->bits(), false /* isSigned */),
+         llvm::ConstantInt::get(llvm_types_->t_uint8, value->lanes(), false /* isSigned */)}));
+  }
+
+  void Visit(const char* key, runtime::NDArray* value) final {
+    CHECK(false) << "Do not support serializing NDArray";
+  }
+
+  void VisitMetadata(runtime::metadata::MetadataBase metadata) {
+    elements_.emplace_back(std::vector<llvm::Constant*>());
+    ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+    auto struct_elements = elements_.back();
+    elements_.pop_back();
+    auto struct_ty = llvm_types_->structs_by_type_key[metadata->GetTypeKey()];
+    ICHECK(struct_ty != nullptr) << "Did not find LLVM StructType* for type_key=" << metadata->GetTypeKey();
+    std::string ty_value;
+    llvm::raw_string_ostream ty_os(ty_value);
+    struct_ty->print(ty_os, true);
+    LOG(INFO) << "Get LLVM ConstantStruct (" << struct_elements.size() << " elements)";
+    LOG(INFO) << "  Type (" << metadata->GetTypeKey() << "==" << struct_ty->getName().data() << "): " << ty_value;
+    for (auto e : struct_elements) {
+      std::string value;
+      llvm::raw_string_ostream os(value);
+      e->print(os);
+      LOG(INFO) << " - " << value;
+    }
+    CHECK_EQ(struct_elements.size(), struct_ty->getNumElements());
+    auto out = llvm::ConstantStruct::get(struct_ty, struct_elements);
+    if (elements_.size() > 0) {
+      elements_.back().push_back(out);
+    } else {
+      last_production_ = out;
+    }
+  }
+
+  void VisitArray(const runtime::metadata::MetadataArrayNode* arr) {
+    llvm::Type* element_type;
+    switch (arr->kind) {
+    case MetadataKind::kInt64:
+      element_type = llvm_types_->t_int64;
+      break;
+    case MetadataKind::kUint64:
+      element_type = llvm_types_->t_int64;
+      break;
+    case MetadataKind::kBool:
+      element_type = llvm_types_->t_uint8;
+      break;
+    case MetadataKind::kString:
+      element_type = llvm_types_->t_cstring;
+      break;
+    case MetadataKind::kMetadata: {
+      element_type = llvm_types_->structs_by_type_key[arr->type_key];
+      ICHECK(element_type != nullptr) << "Did not find LLVM StructType* for type_key=" << arr->type_key;
+      break;
+    }
+    default:
+      LOG(FATAL) << "unknown metadata kind " << arr->kind;
+      break;
+    }
+
+    elements_.emplace_back(std::vector<llvm::Constant*>());
+    for (auto o : arr->array) {
+      if (o->IsInstance<FloatImmNode>()) {
+        double value = Downcast<FloatImm>(o)->value;
+        Visit(nullptr, &value);
+      }
+      if (o->IsInstance<IntImmNode>()) {
+        auto value = Downcast<IntImm>(o)->value;
+        Visit(nullptr, &value);
+      } else if (o->IsInstance<StringObj>()) {
+        ::std::string value = Downcast<String>(o);
+        Visit(nullptr, &value);
+      } else {
+        // nested array not possible.
+        VisitMetadata(Downcast<runtime::metadata::MetadataBase>(o));
+      }
+    }
+
+    auto array = elements_.back();
+    elements_.pop_back();
+    CHECK(element_type != nullptr);
+    auto arr_ty = llvm::ArrayType::get(element_type, array.size());
+    auto llvm_arr = llvm::ConstantArray::get(arr_ty, array);
+
+    if (elements_.size() > 0) {
+      elements_.back().emplace_back(codegen_->GetGlobalConstant(llvm_arr, "", llvm::GlobalValue::PrivateLinkage));
+    } else {
+      last_production_ = llvm_arr;
+    }
+  }
+
+  void Visit(const char* key, ObjectRef* value) final {
+    const runtime::metadata::MetadataArrayNode* arr =
+        value->as<runtime::metadata::MetadataArrayNode>();
+    if (arr != nullptr) {
+      VisitArray(arr);
+      return;
+    }
+
+    runtime::metadata::MetadataBase metadata = Downcast<runtime::metadata::MetadataBase>(*value);
+    VisitMetadata(metadata);
+  }
+
+  llvm::Constant* Serialize(runtime::metadata::MetadataBase metadata) {
+    Visit(nullptr, &metadata);
+    ICHECK(last_production_);
+    return codegen_->GetGlobalConstant(last_production_);
+  }
+
+  CodeGenLLVM* codegen_;
+  MetadataLlvmTypes* llvm_types_;
+  llvm::LLVMContext* ctx_;
+  llvm::Module* module_;
+  std::vector<std::vector<llvm::Constant*>> elements_;
+  llvm::Constant* last_production_;
+};
+
+void CodeGenCPU::DefineMetadata(runtime::metadata::Metadata metadata) {
+  MetadataLlvmTypes llvm_types{
+      t_float64_ /* t_float64 */,
+      llvm::Type::getInt8Ty(*ctx_) /* t_uint8 */,
+      t_int64_ /* t_int64 */,
+      llvm::Type::getInt8Ty(*ctx_) /* t_bool */,
+      t_char_->getPointerTo() /* t_cstring */,
+      t_void_p_ /* t_void_p */,
+      llvm::StructType::create(*ctx_, {t_int8_, t_int8_, t_int8_}, "DLDataType") /* t_data_type */,
+  };
+
+  std::vector<runtime::metadata::MetadataBase> queue;
+  metadata::DiscoverComplexTypesVisitor discover_complex{&queue};
+  discover_complex.Discover(metadata);
+
+  MetadataTypeDefiner definer{ctx_, &llvm_types};
+  for (auto md : queue) {
+    if (md.defined()) {
+      definer.DefineType(md);
+    }
+  }
+
+  MetadataSerializerLLVM serializer{this, &llvm_types};
+  auto metadata_constant_gv = serializer.Serialize(metadata);
+
+  function_ =
+      llvm::Function::Create(ftype_tvm_backend_packed_c_func_, llvm::Function::ExternalLinkage,
+                             "get_c_metadata", module_.get());
+  function_->setCallingConv(llvm::CallingConv::C);
+  function_->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);
+
+  llvm::BasicBlock* entry_point_entry = llvm::BasicBlock::Create(*ctx_, "entry", function_);
+  builder_->SetInsertPoint(entry_point_entry);
+
+  auto ret_values_p = builder_->CreateBitCast(GetArg(function_, 3), t_void_p_->getPointerTo());
+  builder_->CreateStore(builder_->CreateBitCast(metadata_constant_gv, t_void_p_), ret_values_p);
+
+  auto ret_tcode = builder_->CreateBitCast(GetArg(function_, 4), t_int_->getPointerTo());
+  builder_->CreateStore(llvm::ConstantInt::get(t_int_, kTVMOpaqueHandle), ret_tcode);
+
+  builder_->CreateRet(ConstInt32(0));
+}
+
 void CodeGenCPU::DefineFunctionRegistry(Array<String> func_names) {
   ICHECK(is_system_lib_) << "Loading of --system-lib modules is yet to be defined for C runtime";
   Array<String> symbols;
@@ -974,9 +1333,11 @@ void CodeGenCPU::AddStartupFunction() {
 
 llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
   if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
-    return CreateCallPacked(op);
+    return CreateCallPacked(op, true /* use_string_lookup */);
   } else if (op->op.same_as(builtin::tvm_call_trace_packed_lowered())) {
     return CreateCallTracePacked(op);
+  } else if (op->op.same_as(builtin::tvm_call_cpacked_lowered())) {
+    return CreateCallPacked(op, false /* use_string_lookup */);
   } else if (op->op.same_as(builtin::tvm_static_handle())) {
     return CreateStaticHandle();
   } else if (op->op.same_as(builtin::tvm_throw_last_error())) {
@@ -1046,6 +1407,12 @@ void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
   builder_->CreateCondBr(cond, end_block, fail_block, md_very_likely_branch_);
   // fail condition.
   builder_->SetInsertPoint(fail_block);
+
+  std::vector<llvm::Type *> AsmArgTypes;
+  llvm::FunctionType *AsmFTy = llvm::FunctionType::get(t_void_, false);
+  llvm::InlineAsm *IA =
+    llvm::InlineAsm::get(AsmFTy, "int 3", "", true, /* IsAlignStack */ false, llvm::InlineAsm::AD_Intel);
+  builder_->CreateCall(AsmFTy, IA);
 #if TVM_LLVM_VERSION >= 90
   auto err_callee =
       llvm::FunctionCallee(ftype_tvm_api_set_last_error_, RuntimeTVMAPISetLastError());
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index 26f251f1a9c8..b3c9516cd0ad 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -56,6 +56,11 @@ class CodeGenCPU : public CodeGenLLVM {
    */
   void DefineFunctionRegistry(Array<String> func_names);
 
+  /*!
+   * \brief Serialize the metadata object as data, and implement get_c_metadata function.
+   * \param metadata The metadata which should be serialized.
+   */
+  void DefineMetadata(runtime::metadata::Metadata metadata);
  protected:
   void AddStartupFunction() final;
   // meta data
@@ -117,9 +122,9 @@ class CodeGenCPU : public CodeGenLLVM {
     llvm::BasicBlock* end_block;
   };
   PackedCall MakeCallPackedLowered(const Array<PrimExpr>& args, const DataType& r_type,
-                                   const int64_t begin, const int64_t end);
+                                   const int64_t begin, const int64_t end, bool use_string_lookup);
   // create call into tvm packed function.
-  llvm::Value* CreateCallPacked(const CallNode* op);
+  llvm::Value* CreateCallPacked(const CallNode* op, bool use_string_lookup);
   // Create trace call into tvm packed function.
   llvm::Value* CreateCallTracePacked(const CallNode* op);
   // Create static initialization
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index cc6fdc31c563..3268790898d5 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -22,6 +22,7 @@
  */
 #ifdef TVM_LLVM_VERSION
 // Part of the code are adapted from Halide's CodeGen_LLVM
+#include "llvm_common.h"
 #include "codegen_llvm.h"
 
 #include <tvm/runtime/c_runtime_api.h>
@@ -190,6 +191,20 @@ void CodeGenLLVM::AddFunctionInternal(const PrimFunc& f, bool ret_void) {
   }
 }
 
+llvm::GlobalVariable* CodeGenLLVM::GetLinkedParamSymbol(const std::string& param_name,
+                                                        llvm::ConstantArray* array) {
+  std::string symbol_name = std::string(::tvm::runtime::symbol::tvm_param_prefix) + param_name;
+  llvm::GlobalVariable* var = module_->getGlobalVariable(symbol_name, true /* AllowInternal */);
+  if (var == nullptr) {
+    CHECK(array != nullptr) << "Expect param symbol " << symbol_name
+                            << " to either be defined or for the array to be supplied";
+    var = new llvm::GlobalVariable(*module_, static_cast<llvm::Type*>(array->getType()), true,
+                                   llvm::GlobalValue::InternalLinkage, array, symbol_name);
+  }
+  //(array != nullptr ?  : static_cast<llvm::Type*>(t_void_p_)),
+  return var;
+}
+
 void CodeGenLLVM::LinkParameters(const Map<String, LinkedParam> params) {
   // It would be nice to de-dupe these declarations frm src/tir/transforms/make_packed_api.cc,
   // but they are at a different layer in the compiler...
@@ -208,22 +223,13 @@ void CodeGenLLVM::LinkParameters(const Map<String, LinkedParam> params) {
   llvm::BasicBlock* entry = llvm::BasicBlock::Create(*ctx_, "entry", function);
   builder_->SetInsertPoint(entry);
 
-  auto getArg = [function](int i) -> llvm::Argument* {
-#if TVM_LLVM_VERSION >= 100
-    return function->getArg(i);
-#elif TVM_LLVM_VERSION >= 50
-    return &function->arg_begin()[i];
-#else
-    return &*std::next(function->arg_begin(), i);
-#endif
-  };
-
   llvm::Type* t_int64_p = t_int64_->getPointerTo(GetGlobalAddressSpace());
-  llvm::Value* sid = builder_->CreateLoad(t_int64_, builder_->CreateBitCast(getArg(0), t_int64_p));
+  llvm::Value* sid =
+      builder_->CreateLoad(t_int64_, builder_->CreateBitCast(GetArg(function, 0), t_int64_p));
 
-  auto ret_tcode = builder_->CreateBitCast(getArg(4), t_int_p);
-  auto ret_value =
-      builder_->CreateBitCast(getArg(3), t_void_p_->getPointerTo(GetGlobalAddressSpace()));
+  auto ret_tcode = builder_->CreateBitCast(GetArg(function, 4), t_int_p);
+  auto ret_value = builder_->CreateBitCast(GetArg(function, 3),
+                                           t_void_p_->getPointerTo(GetGlobalAddressSpace()));
 
   llvm::BasicBlock* default_block = llvm::BasicBlock::Create(*ctx_, "default_block", function);
   llvm::SwitchInst* switch_inst = builder_->CreateSwitch(sid, default_block, params.size() + 1);
@@ -235,9 +241,7 @@ void CodeGenLLVM::LinkParameters(const Map<String, LinkedParam> params) {
   // Add data to the global section.
   for (auto kv : params) {
     auto array = NDArrayToLLVMArray(ctx_, kv.second->param);
-    std::string symbol_name = std::string(::tvm::runtime::symbol::tvm_param_prefix) + kv.first;
-    llvm::GlobalVariable* param_symbol = new llvm::GlobalVariable(
-        *module_, array->getType(), true, llvm::GlobalValue::InternalLinkage, array, symbol_name);
+    llvm::GlobalVariable* param_symbol = GetLinkedParamSymbol(kv.first, array);
     auto dtype = tvm::runtime::DataType(kv.second->param->dtype);
     size_t align = std::max(tvm::runtime::GetVectorBytes(dtype), tvm::runtime::kAllocAlignment);
 #if TVM_LLVM_VERSION >= 100
@@ -245,8 +249,10 @@ void CodeGenLLVM::LinkParameters(const Map<String, LinkedParam> params) {
 #else
     param_symbol->setAlignment(align);
 #endif
+    param_symbol->setInitializer(array);
 
-    llvm::BasicBlock* case_block = llvm::BasicBlock::Create(*ctx_, "case_" + symbol_name, function);
+    llvm::BasicBlock* case_block =
+        llvm::BasicBlock::Create(*ctx_, "case_" + param_symbol->getName(), function);
     switch_inst->addCase(
         llvm::cast<llvm::ConstantInt>(llvm::ConstantInt::get(t_int64_, kv.second->id)), case_block);
     builder_->SetInsertPoint(case_block);
@@ -387,6 +393,7 @@ void CodeGenLLVM::Optimize() {
     fpass.run(*it);
   }
   fpass.doFinalization();
+  // PrintModule(module_.get());
   mpass.run(*module_);
 }
 
@@ -762,21 +769,28 @@ llvm::Value* CodeGenLLVM::CreateCast(DataType from, DataType to, llvm::Value* va
   }
 }
 
-llvm::Constant* CodeGenLLVM::GetConstString(const std::string& str) {
-  auto it = str_map_.find(str);
-  if (it != str_map_.end()) return it->second;
-  llvm::Type* type = llvm::ArrayType::get(t_char_, str.length() + 1);
-  llvm::GlobalVariable* global = new llvm::GlobalVariable(
-      *module_, type, true, llvm::GlobalValue::PrivateLinkage, nullptr, ".str");
+// TODO: refactor GetConststring
+llvm::Constant* CodeGenLLVM::GetGlobalConstant(llvm::Constant* const_data, const std::string& name,
+                                               llvm::GlobalValue::LinkageTypes linkage_type) {
+  llvm::Type* ty = const_data->getType();
+  llvm::GlobalVariable* global =
+      new llvm::GlobalVariable(*module_, ty, true, linkage_type, const_data, name);
 #if TVM_LLVM_VERSION >= 100
   global->setAlignment(llvm::Align(1));
 #else
   global->setAlignment(1);
 #endif
-  global->setInitializer(llvm::ConstantDataArray::getString(*ctx_, str));
   llvm::Constant* zero = ConstInt32(0);
   llvm::Constant* indices[] = {zero, zero};
-  llvm::Constant* ptr = llvm::ConstantExpr::getGetElementPtr(type, global, indices);
+  llvm::Constant* ptr = llvm::ConstantExpr::getGetElementPtr(ty, global, indices);
+  return ptr;
+}
+
+llvm::Constant* CodeGenLLVM::GetConstString(const std::string& str) {
+  auto it = str_map_.find(str);
+  if (it != str_map_.end()) return it->second;
+  auto llvm_str = llvm::ConstantDataArray::getString(*ctx_, str);
+  auto ptr = GetGlobalConstant(llvm_str, ".str", llvm::GlobalValue::PrivateLinkage);
   str_map_[str] = ptr;
   return ptr;
 }
@@ -1306,9 +1320,13 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const LoadNode* op) {
 }
 
 llvm::Value* CodeGenLLVM::VisitExpr_(const CallNode* op) {
+  //  LOG(INFO) << "Visit Call:" << GetRef<Call>(op);
   if (auto* ptr_op = op->op.as<OpNode>()) {
     auto call_op = GetRef<Op>(ptr_op);
-    if (op->op.same_as(builtin_call_extern_) || op->op.same_as(builtin_call_pure_extern_)) {
+    if (op->op.same_as(builtin_lookup_param_)) {
+      //      return llvm::ConstantInt::get(t_void_p_, 0);
+      return GetLinkedParamSymbol(Downcast<StringImm>(op->args[0])->value, nullptr);
+    } else if (op->op.same_as(builtin_call_extern_) || op->op.same_as(builtin_call_pure_extern_)) {
       // call extern intrinsic
       ICHECK_GE(op->args.size(), 1U);
       auto global_symbol = Downcast<StringImm>(op->args[0]);
@@ -1319,7 +1337,10 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const CallNode* op) {
       return this->CreateCallExtern(GetType(GetRef<PrimExpr>(op)), op_attr_global_symbol_[call_op],
                                     op->args, false);
     } else {
-      return CreateIntrinsic(op);
+      VLOG(2) << "CreateIntrinsic: " << GetRef<Call>(op);
+      auto x = CreateIntrinsic(op);
+      VLOG(2) << "CreateIntrinsic done";
+      return x;
     }
   } else {
     ICHECK(op->op.as<GlobalVarNode>());
@@ -1480,11 +1501,22 @@ void CodeGenLLVM::VisitStmt_(const IfThenElseNode* op) {
   builder_->SetInsertPoint(end_block);
 }
 
+void CodeGenLLVM::VisitStmt_(const AllocateConstNode* op) {
+  auto data = op->data.value();
+  auto array = NDArrayToLLVMArray(ctx_, data);
+  std::string symbol_name = op->buffer_var->name_hint;
+  llvm::GlobalVariable* param_symbol = new llvm::GlobalVariable(
+      *module_, array->getType(), true, llvm::GlobalValue::InternalLinkage, array, symbol_name);
+
+  var_map_[op->buffer_var.operator->()] = param_symbol;
+  this->VisitStmt(op->body);
+}
+
 void CodeGenLLVM::VisitStmt_(const AllocateNode* op) {
   ICHECK(!is_zero(op->condition));
   llvm::Value* buf = nullptr;
 
-  int32_t constant_size = op->constant_allocation_size();
+  int32_t constant_size = op->ConstantAllocationSize();
   ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation";
   StorageInfo& info = alloc_storage_info_[op->buffer_var.get()];
   if (constant_size % 4 == 0 && info.alignment == 0) {
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index a40677c955f8..f12d729c3682 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -23,6 +23,7 @@
  */
 #ifndef TVM_TARGET_LLVM_CODEGEN_LLVM_H_
 #define TVM_TARGET_LLVM_CODEGEN_LLVM_H_
+#include <llvm/IR/GlobalValue.h>
 #ifdef TVM_LLVM_VERSION
 
 #include <tvm/arith/analyzer.h>
@@ -141,7 +142,7 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
    * \param e The expression to be created value for.
    * \return created value.
    */
-  llvm::Value* MakeValue(const PrimExpr& e) { return VisitExpr(e); }
+  llvm::Value* MakeValue(const PrimExpr& e) { auto a = VisitExpr(e); /* LOG(INFO) << "MakeValue (" << e << "): " << a; */ return a;  }
   // Short hande code to get a constant int 32
   llvm::Constant* ConstInt32(int64_t value) const {
     return llvm::ConstantInt::getSigned(t_int32_, value);
@@ -181,12 +182,20 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
   void VisitStmt_(const WhileNode* op) override;
   void VisitStmt_(const IfThenElseNode* op) override;
   void VisitStmt_(const AllocateNode* op) override;
+  void VisitStmt_(const AllocateConstNode* op) override;
   void VisitStmt_(const AttrStmtNode* op) override;
   void VisitStmt_(const AssertStmtNode* op) override;
   void VisitStmt_(const LetStmtNode* op) override;
   void VisitStmt_(const SeqStmtNode* op) override;
   void VisitStmt_(const EvaluateNode* op) override;
 
+  // Get constant string
+  llvm::Constant* GetConstString(const std::string& str);
+
+  llvm::Constant* GetGlobalConstant(
+      llvm::Constant* const_data, const std::string& name = "",
+      llvm::GlobalValue::LinkageTypes linkage_type = llvm::GlobalValue::InternalLinkage);
+
  protected:
   /*!
    * \brief Address and type pair to assist in handling opaque pointers.
@@ -308,6 +317,13 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
    */
   llvm::Function* GetIntrinsicDecl(llvm::Intrinsic::ID id, llvm::Type* ret_type,
                                    llvm::ArrayRef<llvm::Type*> arg_types);
+  /*!
+   * \brief Lookup or create a GlobalVariable whose content is the data field of a DLTensor for a
+   * given linked_param() CallNode.
+   * \param param_name Parameter name (e.g. unmangled, from lookup_param node).
+   * \return the GlobalVariable indicated in the brief.
+   */
+  llvm::GlobalVariable* GetLinkedParamSymbol(const ::std::string& param_name, llvm::ConstantArray* array);
   /*!
    * \brief Get the number of elements in the given vector value.
    * \param vec The value, must be of a vector type.
@@ -318,8 +334,6 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
   // Get alignment given index.
   void GetAlignment(DataType t, const VarNode* buf_var, const PrimExpr& index, int* p_alignment,
                     int* p_native_bits);
-  // Get constant string
-  llvm::Constant* GetConstString(const std::string& str);
   // do a scalarize call with f
   llvm::Value* CreateScalarizedCall(const CallNode* op, llvm::Function* f,
                                     const std::vector<llvm::Value*>& args);
@@ -353,6 +367,16 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
                                              unsigned int shared_address_space, int alignment,
                                              llvm::GlobalValue::LinkageTypes linkage);
 
+  llvm::Argument* GetArg(const llvm::Function* function, int i) const {
+#if TVM_LLVM_VERSION >= 100
+    return function->getArg(i);
+#elif TVM_LLVM_VERSION >= 50
+    return const_cast<llvm::Argument*>(&function->arg_begin()[i]);
+#else
+    return &*std::next(function->arg_begin(), i);
+#endif
+  }
+
   // The IRBuilder.
   using IRBuilder = llvm::IRBuilder<llvm::ConstantFolder, llvm::IRBuilderDefaultInserter>;
   // The current function
@@ -411,6 +435,8 @@ class CodeGenLLVM : public ExprFunctor<llvm::Value*(const PrimExpr&)>,
   const Op& builtin_call_pure_extern_ = builtin::call_pure_extern();
   const Op& builtin_call_llvm_intrin_ = builtin::call_llvm_intrin();
   const Op& builtin_call_llvm_pure_intrin_ = builtin::call_llvm_pure_intrin();
+  const Op& builtin_lookup_param_ = builtin::lookup_param();
+  const Op& builtin_tvm_call_cpacked_lowered_ = builtin::tvm_call_cpacked_lowered();
 
   /*! \brief Helper struct for debug infos. */
   struct DebugInfo {
@@ -442,9 +468,12 @@ void CodeGenLLVM::AddFunctionsOrdered(IterType begin, IterType end, ConvType pfu
   std::sort(funcs.begin(), funcs.end(), [](PrimFunc func_a, PrimFunc func_b) {
     std::string name_a = func_a->GetAttr<String>(tvm::attr::kGlobalSymbol).value();
     std::string name_b = func_b->GetAttr<String>(tvm::attr::kGlobalSymbol).value();
-    return name_a < name_b;
+    // TODO(masahi,areusch): AoT codegen should not depend on the sort order
+    return name_a > name_b;
   });
   for (auto& f : funcs) {
+    auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
+    LOG(INFO) << "Adding " << static_cast<std::string>(global_symbol.value());
     AddFunction(f);
   }
 }
diff --git a/src/target/llvm/codegen_nvptx.cc b/src/target/llvm/codegen_nvptx.cc
index 24b7bd2b6acc..01a3191cc220 100644
--- a/src/target/llvm/codegen_nvptx.cc
+++ b/src/target/llvm/codegen_nvptx.cc
@@ -60,7 +60,7 @@ class CodeGenNVPTX : public CodeGenLLVM {
       buf =
           AllocateSharedMemory(op->dtype, 0, 3, info.alignment, llvm::GlobalValue::ExternalLinkage);
     } else {
-      int32_t constant_size = op->constant_allocation_size();
+      size_t constant_size = op->ConstantAllocationSize();
       ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation in GPU";
 
       if (constant_size % 4 == 0 && info.alignment == 0) {
diff --git a/src/target/llvm/llvm_common.cc b/src/target/llvm/llvm_common.cc
index 06b2be2d9fb6..f13e8563e053 100644
--- a/src/target/llvm/llvm_common.cc
+++ b/src/target/llvm/llvm_common.cc
@@ -189,6 +189,13 @@ std::string LLVMTargetToString(const Target& target) {
   return os.str();
 }
 
+void PrintModule(const llvm::Module* mod) {
+  std::string modpe_str;
+  llvm::raw_string_ostream rso(modpe_str);
+  mod->print(rso, nullptr);
+  LOG(INFO) << rso.str();
+}
+
 }  // namespace codegen
 }  // namespace tvm
 #endif  // TVM_LLVM_VERSION
diff --git a/src/target/llvm/llvm_common.h b/src/target/llvm/llvm_common.h
index 556f05d2e33a..732e0f4ab1df 100644
--- a/src/target/llvm/llvm_common.h
+++ b/src/target/llvm/llvm_common.h
@@ -126,6 +126,9 @@ std::unique_ptr<llvm::TargetMachine> GetLLVMTargetMachine(const Target& target,
  */
 std::string LLVMTargetToString(const Target& target);
 
+void PrintModule(const llvm::Module* mod);
+
+
 }  // namespace codegen
 }  // namespace tvm
 
diff --git a/src/target/llvm/llvm_module.cc b/src/target/llvm/llvm_module.cc
index cf8b59357b47..b1da1fccfa6e 100644
--- a/src/target/llvm/llvm_module.cc
+++ b/src/target/llvm/llvm_module.cc
@@ -308,14 +308,14 @@ class LLVMModuleNode final : public runtime::ModuleNode {
 
     cg->SetFastMathFlag(fmf);
 
+    if (found_linked_params) {
+      cg->LinkParameters(linked_params);
+    }
     cg->AddFunctionsOrdered(funcs.begin(), funcs.end());
     if (entry_func.length() != 0) {
       cg->AddMainFunction(entry_func);
     }
 
-    if (found_linked_params) {
-      cg->LinkParameters(linked_params);
-    }
     module_ = cg->Finish();
     module_->addModuleFlag(llvm::Module::Warning, "tvm_target",
                            llvm::MDString::get(*ctx_, LLVMTargetToString(target)));
@@ -527,6 +527,45 @@ TVM_REGISTER_GLOBAL("codegen.codegen_blob")
       return runtime::Module(n);
     });
 
+runtime::Module CreateLLVMCppMetadataModule(runtime::metadata::Metadata metadata, Target target,
+                                            tvm::relay::Runtime runtime) {
+  InitializeLLVM();
+  auto tm = GetLLVMTargetMachine(target);
+  bool system_lib = runtime->GetAttr<Bool>("system-lib").value_or(Bool(false));
+  auto ctx = std::make_shared<llvm::LLVMContext>();
+  std::unique_ptr<CodeGenCPU> cg{new CodeGenCPU()};
+
+  cg->Init("TVMMetadataMod", tm.get(), ctx.get(), system_lib, system_lib, false /* target_c_runtime */);
+
+  cg->DefineMetadata(metadata);
+  auto mod = cg->Finish();
+  mod->addModuleFlag(llvm::Module::Warning, "tvm_target",
+                     llvm::MDString::get(*ctx, LLVMTargetToString(target)));
+  mod->addModuleFlag(llvm::Module::Override, "Debug Info Version", llvm::DEBUG_METADATA_VERSION);
+
+  if (tm->getTargetTriple().isOSDarwin()) {
+    mod->addModuleFlag(llvm::Module::Override, "Dwarf Version", 2);
+  }
+
+  std::string verify_errors_storage;
+  llvm::raw_string_ostream verify_errors(verify_errors_storage);
+  LOG_IF(FATAL, llvm::verifyModule(*mod, &verify_errors))
+      << "LLVM module verification failed with the following errors: \n"
+      << verify_errors.str();
+
+  // std::string tmp;
+  // llvm::raw_string_ostream stream(tmp);
+  // mod->print(stream, nullptr);
+  // LOG(INFO) << "LLVM metadata IR: " << stream.str();
+
+  auto n = make_object<LLVMModuleNode>();
+  n->Init(std::move(mod), ctx);
+
+  auto meta_mod = MetadataModuleCreate(metadata);
+  meta_mod->Import(runtime::Module(n));
+  return meta_mod;
+}
+
 runtime::Module CreateLLVMCrtMetadataModule(const Array<runtime::Module>& modules, Target target,
                                             tvm::relay::Runtime runtime) {
   Array<String> func_names;
diff --git a/src/target/llvm/llvm_module.h b/src/target/llvm/llvm_module.h
index 933030e213d2..660d81400b0d 100644
--- a/src/target/llvm/llvm_module.h
+++ b/src/target/llvm/llvm_module.h
@@ -33,6 +33,9 @@
 namespace tvm {
 namespace codegen {
 
+runtime::Module CreateLLVMCppMetadataModule(runtime::metadata::Metadata metadata, Target target,
+                                            tvm::relay::Runtime runtime);
+
 runtime::Module CreateLLVMCrtMetadataModule(const Array<runtime::Module>& modules, Target target,
                                             tvm::relay::Runtime runtime);
 
diff --git a/src/target/metadata.cc b/src/target/metadata.cc
new file mode 100644
index 000000000000..adf4cba3e610
--- /dev/null
+++ b/src/target/metadata.cc
@@ -0,0 +1,47 @@
+/*
+ * 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 metadata.cc
+ * \brief Implementations of the compiler extensions for Metadata.
+ */
+
+#include "metadata.h"
+
+#include <tvm/node/reflection.h>
+
+namespace tvm {
+namespace target {
+namespace metadata {
+
+TVM_REGISTER_REFLECTION_VTABLE(VisitableMetadataNode,
+                               ::tvm::detail::ReflectionTrait<VisitableMetadataNode>)
+    .set_creator([](const std::string&) -> ObjectPtr<Object> {
+      return ::tvm::runtime::make_object<VisitableMetadataNode>();
+    });
+
+TVM_REGISTER_REFLECTION_VTABLE(VisitableTensorInfoNode,
+                               ::tvm::detail::ReflectionTrait<VisitableTensorInfoNode>)
+    .set_creator([](const std::string&) -> ObjectPtr<Object> {
+      return ::tvm::runtime::make_object<VisitableTensorInfoNode>();
+    });
+
+}  // namespace metadata
+}  // namespace target
+}  // namespace tvm
diff --git a/src/target/metadata.h b/src/target/metadata.h
new file mode 100644
index 000000000000..c803408298e5
--- /dev/null
+++ b/src/target/metadata.h
@@ -0,0 +1,194 @@
+/*
+ * 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 tvm/target/metadata.h
+ * \brief Extends Metadata for use in the compiler.
+ */
+#ifndef TVM_TARGET_METADATA_H_
+#define TVM_TARGET_METADATA_H_
+
+#include <tvm/runtime/metadata.h>
+
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace tvm {
+namespace target {
+namespace metadata {
+
+/*!
+ * \brief Subclass of MetadataNode that implements the VisitAttrs reflection method.
+ *
+ * This implementation (and other such Visitable subclasses) is compiled into libtvm.so, but not
+ * libtvm_runtime.so, because reflection is not supported in libtvm_runtime.so over code size
+ * concerns. It is used during compilation by the generic metadata code-generators.
+ */
+class VisitableMetadataNode : public ::tvm::runtime::metadata::MetadataNode {
+ public:
+  explicit VisitableMetadataNode(const struct ::TVMMetadata* data) : MetadataNode{data} {}
+  VisitableMetadataNode() : MetadataNode{nullptr} {}
+
+  void VisitAttrs(AttrVisitor* v) {
+    int64_t version_cpp{version()};
+    v->Visit("version", &version_cpp);
+    auto inputs_array = Array<ObjectRef>();
+    auto inputs_accessor = inputs();
+    inputs_array.reserve(num_inputs());
+    for (int64_t i = 0; i < num_inputs(); ++i) {
+      inputs_array.push_back(::tvm::runtime::metadata::TensorInfo{inputs_accessor[i]});
+    }
+    ::tvm::runtime::metadata::MetadataArray inputs_metadata_array{
+        inputs_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, ::tvm::runtime::metadata::TensorInfoNode::_type_key};
+    v->Visit("inputs", &inputs_metadata_array);
+    int64_t num_inputs_cpp = num_inputs();
+    v->Visit("num_inputs", &num_inputs_cpp);
+    auto outputs_array = Array<ObjectRef>();
+    auto outputs_accessor = outputs();
+    outputs_array.reserve(num_outputs());
+    for (int64_t i = 0; i < num_outputs(); ++i) {
+      outputs_array.push_back(::tvm::runtime::metadata::TensorInfo{outputs_accessor[i]});
+    }
+    ::tvm::runtime::metadata::MetadataArray outputs_metadata_array{
+      outputs_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, ::tvm::runtime::metadata::TensorInfoNode::_type_key};
+    v->Visit("outputs", &outputs_metadata_array);
+    int64_t num_outputs_cpp = num_outputs();
+    v->Visit("num_outputs", &num_outputs_cpp);
+    auto pools_array = Array<ObjectRef>();
+    auto pools_accessor = pools();
+    pools_array.reserve(num_pools());
+    for (int64_t i = 0; i < num_pools(); ++i) {
+      pools_array.push_back(::tvm::runtime::metadata::TensorInfo{pools_accessor[i]});
+    }
+    ::tvm::runtime::metadata::MetadataArray pools_metadata_array{
+        pools_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, ::tvm::runtime::metadata::TensorInfoNode::_type_key};
+    v->Visit("pools", &pools_metadata_array);
+    int64_t num_pools_cpp = num_pools();
+    v->Visit("num_pools", &num_pools_cpp);
+    ::std::string mod_name_cpp{data()->mod_name};
+    v->Visit("mod_name", &mod_name_cpp);
+  }
+};
+
+/*!
+ * \brief Subclass of MetadataNode which also owns the backing C structures.
+ *
+ * This class (and other InMemory subclasses) are used during compilation to instantiate Metadata
+ * instances whose storage lives outside of .rodata. This class exists because the Module returned
+ * from tvm.relay.build must also be ready to run inference.
+ */
+class InMemoryMetadataNode : public ::tvm::target::metadata::VisitableMetadataNode {
+ public:
+  InMemoryMetadataNode()
+      : InMemoryMetadataNode(0 /* version */, {} /* inputs */, {} /* outputs */, {} /* pools */,
+                             "" /* mod_name */) {}
+  InMemoryMetadataNode(int64_t version,
+                       const ::std::vector<::tvm::runtime::metadata::TensorInfo>& inputs,
+                       const ::std::vector<::tvm::runtime::metadata::TensorInfo>& outputs,
+                       const ::std::vector<::tvm::runtime::metadata::TensorInfo>& pools,
+                       const ::tvm::runtime::String mod_name)
+      : VisitableMetadataNode{&storage_},
+        inputs_{new struct TVMTensorInfo[inputs.size()]()},
+        inputs_objs_{inputs},
+        outputs_{new struct TVMTensorInfo[outputs.size()]()},
+        outputs_objs_{outputs},
+        pools_{new struct TVMTensorInfo[pools.size()]()},
+        pools_objs_{pools},
+        mod_name_{mod_name},
+        storage_{version, nullptr, 0, nullptr, 0, nullptr, 0, mod_name_.c_str()} {
+    storage_.inputs = inputs_.get();
+    storage_.num_inputs = inputs.size();
+    for (unsigned int i = 0; i < inputs.size(); ++i) {
+      inputs_.get()[i] = *inputs[i]->data();
+    }
+    storage_.outputs = outputs_.get();
+    storage_.num_outputs = outputs.size();
+    for (unsigned int i = 0; i < outputs.size(); ++i) {
+      outputs_.get()[i] = *outputs[i]->data();
+    }
+    storage_.pools = pools_.get();
+    storage_.num_pools = pools.size();
+    for (unsigned int i = 0; i < pools.size(); ++i) {
+      pools_.get()[i] = *pools[i]->data();
+    }
+  }
+
+ private:
+  ::std::unique_ptr<struct TVMTensorInfo> inputs_;
+  std::vector<::tvm::runtime::metadata::TensorInfo> inputs_objs_;
+  ::std::unique_ptr<struct TVMTensorInfo> outputs_;
+  std::vector<::tvm::runtime::metadata::TensorInfo> outputs_objs_;
+  ::std::unique_ptr<struct TVMTensorInfo> pools_;
+  std::vector<::tvm::runtime::metadata::TensorInfo> pools_objs_;
+  ::std::string mod_name_;
+  struct ::TVMMetadata storage_;
+};
+
+class VisitableTensorInfoNode : public ::tvm::runtime::metadata::TensorInfoNode {
+ public:
+  explicit VisitableTensorInfoNode(const struct ::TVMTensorInfo* data) : TensorInfoNode{data} {}
+  VisitableTensorInfoNode() : TensorInfoNode{nullptr} {}
+
+  void VisitAttrs(AttrVisitor* v) {
+    ::std::string name_cpp{data()->name};
+    v->Visit("name", &name_cpp);
+    auto shape_array = Array<ObjectRef>();
+    auto shape_accessor = shape();
+    shape_array.reserve(num_shape());
+    for (int64_t i = 0; i < num_shape(); ++i) {
+      shape_array.push_back(::tvm::Integer{static_cast<int>(shape_accessor[i])});
+    }
+    ::tvm::runtime::metadata::MetadataArray shape_metadata_array{
+        shape_array, ::tvm::runtime::metadata::MetadataKind::kInt64, nullptr};
+    v->Visit("shape", &shape_metadata_array);
+    int64_t num_shape_cpp = num_shape();
+    v->Visit("num_shape", &num_shape_cpp);
+    ::tvm::runtime::DataType dtype_cpp{dtype()};
+    v->Visit("dtype", &dtype_cpp);
+  }
+};
+
+class InMemoryTensorInfoNode : public ::tvm::target::metadata::VisitableTensorInfoNode {
+ public:
+  InMemoryTensorInfoNode() : InMemoryTensorInfoNode("", {}, ::tvm::runtime::DataType(0, 0, 0)) {}
+  InMemoryTensorInfoNode(const ::tvm::runtime::String& name, const ::std::vector<int64_t>& shape,
+                         ::tvm::runtime::DataType dtype)
+      : VisitableTensorInfoNode{&storage_},
+        name_{name},
+        shape_{new int64_t[shape.size()]()},
+        storage_{name_.c_str(), nullptr, 0, dtype} {
+    storage_.shape = shape_.get();
+    storage_.num_shape = shape.size();
+    for (unsigned int i = 0; i < shape.size(); ++i) {
+      shape_.get()[i] = shape[i];
+    }
+  }
+
+ private:
+  ::std::string name_;
+  ::std::unique_ptr<int64_t> shape_;
+  struct ::TVMTensorInfo storage_;
+};
+
+}  // namespace metadata
+}  // namespace target
+}  // namespace tvm
+
+#endif  // TVM_TARGET_METADATA_H_
diff --git a/src/target/metadata_module.cc b/src/target/metadata_module.cc
index 2facf1de64d5..0b5af02b2f4d 100644
--- a/src/target/metadata_module.cc
+++ b/src/target/metadata_module.cc
@@ -25,8 +25,10 @@
 
 #include <tvm/relay/runtime.h>
 
+#include <utility>
 #include <vector>
 
+#include "../runtime/const_loader_module.h"
 #include "../runtime/meta_data.h"
 #include "llvm/llvm_module.h"
 #include "source/source_module.h"
@@ -34,8 +36,142 @@
 namespace tvm {
 namespace codegen {
 
+static runtime::Module CreateCrtMetadataModule(
+    runtime::Module target_module, Target target, relay::Runtime runtime,
+    relay::backend::ExecutorCodegenMetadata metadata,
+    Array<runtime::Module> non_crt_exportable_modules,
+    Array<runtime::Module> crt_exportable_modules,
+    const std::unordered_map<std::string, runtime::NDArray>& const_var_ndarray) {
+  if (!non_crt_exportable_modules.empty()) {
+    std::string non_exportable_modules;
+    for (unsigned int i = 0; i < non_crt_exportable_modules.size(); i++) {
+      if (i > 0) {
+        non_exportable_modules += ", ";
+      }
+      auto mod = non_crt_exportable_modules[i];
+      auto pf_sym = mod.GetFunction("get_symbol");
+      if (pf_sym != nullptr) {
+        non_exportable_modules += pf_sym().operator std::string();
+      } else {
+        non_exportable_modules +=
+            std::string{"(module type_key="} + mod->type_key() + std::string{")"};
+      }
+    }
+    CHECK(false) << "These " << non_crt_exportable_modules.size()
+                 << " modules are not exportable to C-runtime: " << non_exportable_modules;
+  }
+
+  if (target->kind->name == "c") {
+    crt_exportable_modules.push_back(target_module);
+    target_module =
+        CreateCSourceCrtMetadataModule(crt_exportable_modules, target, runtime, metadata);
+  } else if (target->kind->name == "llvm") {
+#ifdef TVM_LLVM_VERSION
+    crt_exportable_modules.push_back(target_module);
+    target_module = CreateLLVMCrtMetadataModule(crt_exportable_modules, target, runtime);
+#else   // TVM_LLVM_VERSION
+    LOG(FATAL) << "TVM was not built with LLVM enabled.";
+#endif  // TVM_LLVM_VERSION
+  }
+
+  return target_module;
+}
+
+// TODO(areusch,masahi): Unify metadata representation and remove the need for this function
+static runtime::metadata::Metadata ConvertMetaData(
+    relay::backend::ExecutorCodegenMetadata metadata) {
+  std::vector<runtime::metadata::TensorInfo> inputs;
+  ICHECK(metadata.defined());
+  for (size_t i = 0; i < metadata->inputs.size(); ++i) {
+    auto v = metadata->inputs[i];
+    auto ttype = metadata->input_tensor_types[i];
+    inputs.push_back(
+        runtime::metadata::TensorInfo(make_object<target::metadata::InMemoryTensorInfoNode>(
+            v->name_hint, relay::backend::ShapeToJSON(ttype->shape), ttype->dtype)));
+  }
+
+  std::vector<runtime::metadata::TensorInfo> outputs;
+  auto output_ttypes = metadata->output_tensor_types;
+  for (size_t i = 0; i < output_ttypes.size(); ++i) {
+    auto ttype = output_ttypes[i];
+    std::stringstream name;
+    name << "output" << i;
+    outputs.push_back(
+        runtime::metadata::TensorInfo(make_object<target::metadata::InMemoryTensorInfoNode>(
+            name.str(), relay::backend::ShapeToJSON(ttype->shape), ttype->dtype)));
+  }
+
+  std::vector<runtime::metadata::TensorInfo> pools;
+  for (size_t i = 0; i < metadata->pools.size(); ++i) {
+    auto var = metadata->pools[i];
+    pools.push_back(
+        runtime::metadata::TensorInfo(make_object<target::metadata::InMemoryTensorInfoNode>(
+            var->name_hint,
+            std::vector<int64_t>{metadata->pool_inputs.value()[var]->allocated_size},
+            tvm::runtime::DataType{kDLUInt, 8, 1})));
+  }
+
+  auto n = make_object<target::metadata::InMemoryMetadataNode>(
+      runtime::metadata::kMetadataVersion, inputs, outputs, pools, metadata->mod_name);
+
+  return runtime::metadata::Metadata(std::move(n));
+}
+
+static runtime::Module CreateCppMetadataModule(
+    runtime::Module target_module, Target target, relay::Runtime runtime,
+    relay::backend::ExecutorCodegenMetadata metadata,
+    const std::unordered_map<std::string, std::vector<std::string>>& const_vars_by_symbol,
+    Array<runtime::Module> non_crt_exportable_modules,
+    Array<runtime::Module> crt_exportable_modules,
+    const std::unordered_map<std::string, runtime::NDArray>& const_var_ndarray) {
+  if (!non_crt_exportable_modules.empty()) {
+    runtime::Module const_loader_mod =
+        runtime::ConstLoaderModuleCreate(const_var_ndarray, const_vars_by_symbol);
+    const_loader_mod.Import(target_module);
+    for (const auto& it : non_crt_exportable_modules) {
+      const_loader_mod.Import(it);
+    }
+    target_module = const_loader_mod;
+  }
+
+  if (metadata.defined()) {
+    runtime::metadata::Metadata runtime_metadata = ConvertMetaData(metadata);
+
+    if (metadata->executor == runtime::kTvmExecutorAot && runtime->name == relay::kTvmRuntimeCpp) {
+      if (target->kind->name == "c") {
+        auto metadata_module = CreateCSourceCppMetadataModule(runtime_metadata);
+        metadata_module->Import(target_module);
+        target_module = metadata_module;
+      }
+#ifdef TVM_LLVM_VERSION
+      else if (target->kind->name == "llvm") {
+        auto metadata_module = CreateLLVMCppMetadataModule(runtime_metadata, target, runtime);
+        metadata_module->Import(target_module);
+        target_module = metadata_module;
+      }
+#endif  // TVM_LLVM_VERSION
+      else {
+        CHECK(false) << "Don't know how to create MetadataModule for target type " << target->str();
+      }
+    }
+  }
+
+  return target_module;
+}
+
+/*!
+ * \brief Create a metadata module wrapper. The helper is used by different
+ *        codegens, such as graph executor codegen and the vm compiler.
+ *
+ * \param params The metadata for initialization of all modules.
+ * \param target_module the internal module that is compiled by tvm.
+ * \param ext_modules The external modules that needs to be imported inside the metadata
+ * module(s).
+ * \param target The target that all the modules are compiled for
+ * \return The created metadata module that manages initialization of metadata.
+ */
 runtime::Module CreateMetadataModule(
-    const std::unordered_map<std::string, runtime::NDArray>& params,
+    const std::unordered_map<std::string, runtime::NDArray>& const_var_ndarray,
     tvm::runtime::Module target_module, const Array<runtime::Module>& ext_modules, Target target,
     tvm::relay::Runtime runtime, relay::backend::ExecutorCodegenMetadata metadata) {
   // Here we split modules into two groups:
@@ -52,19 +188,19 @@ runtime::Module CreateMetadataModule(
   bool is_targeting_crt = runtime->name == "crt";
 
   // Wrap all submodules in the initialization wrapper.
-  std::unordered_map<std::string, std::vector<std::string>> sym_metadata;
+  std::unordered_map<std::string, std::vector<std::string>> const_vars_by_symbol;
   for (tvm::runtime::Module mod : ext_modules) {
     auto pf_sym = mod.GetFunction("get_symbol");
     auto pf_var = mod.GetFunction("get_const_vars");
-    std::vector<std::string> arrays;
+    std::vector<std::string> symbol_const_vars;
     if (pf_sym != nullptr && pf_var != nullptr) {
       String symbol = pf_sym();
       Array<String> variables = pf_var();
       for (size_t i = 0; i < variables.size(); i++) {
-        arrays.push_back(variables[i].operator std::string());
+        symbol_const_vars.push_back(variables[i].operator std::string());
       }
-      ICHECK_EQ(sym_metadata.count(symbol), 0U) << "Found duplicated symbol: " << symbol;
-      sym_metadata[symbol] = arrays;
+      ICHECK_EQ(const_vars_by_symbol.count(symbol), 0U) << "Found duplicated symbol: " << symbol;
+      const_vars_by_symbol[symbol] = symbol_const_vars;
     }
     // We only need loading of serialized constant data
     // if there are constants present and required by the
@@ -74,7 +210,7 @@ runtime::Module CreateMetadataModule(
 
     // TODO(@manupa-arm) : we should be able to use csource_metadata
     // if the variables are empty when all the runtime modules implement get_func_names
-    if (arrays.empty() && is_targeting_crt && DSOExportable(mod) &&
+    if (symbol_const_vars.empty() && is_targeting_crt && DSOExportable(mod) &&
         (target->kind->name == "c" || target->kind->name == "llvm")) {
       crt_exportable_modules.push_back(mod);
     } else {
@@ -83,49 +219,16 @@ runtime::Module CreateMetadataModule(
   }
 
   if (is_targeting_crt) {
-    if (!non_crt_exportable_modules.empty()) {
-      std::string non_exportable_modules;
-      for (unsigned int i = 0; i < non_crt_exportable_modules.size(); i++) {
-        if (i > 0) {
-          non_exportable_modules += ", ";
-        }
-        auto mod = non_crt_exportable_modules[i];
-        auto pf_sym = mod.GetFunction("get_symbol");
-        if (pf_sym != nullptr) {
-          non_exportable_modules += pf_sym().operator std::string();
-        } else {
-          non_exportable_modules +=
-              std::string{"(module type_key="} + mod->type_key() + std::string{")"};
-        }
-      }
-      CHECK(false) << "These " << non_crt_exportable_modules.size()
-                   << " modules are not exportable to C-runtime: " << non_exportable_modules;
-    }
-
-    if (target->kind->name == "c") {
-      crt_exportable_modules.push_back(target_module);
-      target_module =
-          CreateCSourceCrtMetadataModule(crt_exportable_modules, target, runtime, metadata);
-    } else if (target->kind->name == "llvm") {
-#ifdef TVM_LLVM_VERSION
-      crt_exportable_modules.push_back(target_module);
-      target_module = CreateLLVMCrtMetadataModule(crt_exportable_modules, target, runtime);
-#else   // TVM_LLVM_VERSION
-      LOG(FATAL) << "TVM was not built with LLVM enabled.";
-#endif  // TVM_LLVM_VERSION
-    }
+    return CreateCrtMetadataModule(target_module, target, runtime, metadata,
+                                   non_crt_exportable_modules, crt_exportable_modules,
+                                   const_var_ndarray);
   } else {
-    if (!non_crt_exportable_modules.empty()) {
-      runtime::Module binary_meta_mod = runtime::MetadataModuleCreate(params, sym_metadata);
-      binary_meta_mod.Import(target_module);
-      for (const auto& it : non_crt_exportable_modules) {
-        binary_meta_mod.Import(it);
-      }
-      return binary_meta_mod;
-    }
+    return CreateCppMetadataModule(target_module, target, runtime, metadata, const_vars_by_symbol,
+                                   non_crt_exportable_modules, crt_exportable_modules,
+                                   const_var_ndarray);
   }
-  return target_module;
 }
 
 }  // namespace codegen
+
 }  // namespace tvm
diff --git a/src/target/metadata_utils.cc b/src/target/metadata_utils.cc
new file mode 100644
index 000000000000..6a059c9fe561
--- /dev/null
+++ b/src/target/metadata_utils.cc
@@ -0,0 +1,157 @@
+/*
+ * 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 tvm/target/metadata_utils.cc
+ * \brief Defines utility functions and classes for emitting metadata.
+ */
+#include "metadata_utils.h"
+
+namespace tvm {
+namespace codegen {
+namespace metadata {
+
+DiscoverArraysVisitor::DiscoverArraysVisitor(std::vector<DiscoveredArray>* queue) : queue_{queue} {}
+
+std::string address_from_parts(const std::vector<std::string>& parts) {
+  std::stringstream ss;
+  for (unsigned int i = 0; i < parts.size(); ++i) {
+    if (i > 0) {
+      ss << "_";
+    }
+    ss << parts[i];
+  }
+  return ss.str();
+}
+
+void DiscoverArraysVisitor::Visit(const char* key, double* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, int64_t* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, uint64_t* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, int* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, bool* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, std::string* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, DataType* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, runtime::NDArray* value) {}
+void DiscoverArraysVisitor::Visit(const char* key, void** value) {}
+
+void DiscoverArraysVisitor::Visit(const char* key, ObjectRef* value) {
+  address_parts_.push_back(key);
+  if (value->as<runtime::metadata::MetadataBaseNode>() != nullptr) {
+    auto metadata = Downcast<runtime::metadata::MetadataBase>(*value);
+    const runtime::metadata::MetadataArrayNode* arr =
+      value->as<runtime::metadata::MetadataArrayNode>();
+
+    if (arr != nullptr) {
+      for (unsigned int i = 0; i < arr->array.size(); i++) {
+        ObjectRef o = arr->array[i];
+        if (o.as<runtime::metadata::MetadataBaseNode>() != nullptr) {
+          std::stringstream ss;
+          ss << i;
+          address_parts_.push_back(ss.str());
+          runtime::metadata::MetadataBase metadata = Downcast<runtime::metadata::MetadataBase>(o);
+          ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+          address_parts_.pop_back();
+        }
+      }
+
+      queue_->push_back(std::make_tuple(address_from_parts(address_parts_),
+                                        Downcast<runtime::metadata::MetadataArray>(metadata)));
+    } else {
+      ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+    }
+  }
+  address_parts_.pop_back();
+}
+
+
+
+void DiscoverComplexTypesVisitor::Visit(const char* key, double* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, int64_t* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, uint64_t* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, int* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, bool* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, std::string* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, DataType* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, runtime::NDArray* value) {}
+void DiscoverComplexTypesVisitor::Visit(const char* key, void** value) {}
+
+bool DiscoverComplexTypesVisitor::DiscoverType(std::string type_key) {
+  VLOG(2) << "DiscoverType " << type_key;
+  auto position_it = type_key_to_position_.find(type_key);
+  if (position_it != type_key_to_position_.end()) {
+    return false;
+  }
+
+  queue_->emplace_back(tvm::runtime::metadata::MetadataBase());
+  type_key_to_position_[type_key] = queue_->size() - 1;
+  return true;
+}
+
+void DiscoverComplexTypesVisitor::DiscoverInstance(runtime::metadata::MetadataBase md) {
+  auto position_it = type_key_to_position_.find(md->GetTypeKey());
+  ICHECK(position_it != type_key_to_position_.end())
+    << "DiscoverInstance requires that DiscoverType has already been called: type_key="
+    << md->GetTypeKey();
+
+  int queue_position = (*position_it).second;
+  if (!(*queue_)[queue_position].defined() && md.defined()) {
+    VLOG(2) << "DiscoverInstance  " << md->GetTypeKey() << ":" << md;
+    (*queue_)[queue_position] = md;
+  }
+}
+
+void DiscoverComplexTypesVisitor::Visit(const char* key, ObjectRef* value) {
+  ICHECK_NOTNULL(value->as<runtime::metadata::MetadataBaseNode>());
+
+  auto metadata = Downcast<runtime::metadata::MetadataBase>(*value);
+  const runtime::metadata::MetadataArrayNode* arr =
+    value->as<runtime::metadata::MetadataArrayNode>();
+
+  if (arr == nullptr) {
+    VLOG(2) << "No array, object-traversing " << metadata->GetTypeKey();
+    ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+    DiscoverType(metadata->GetTypeKey());
+    DiscoverInstance(metadata);
+    return;
+  }
+
+  if (arr->kind != tvm::runtime::metadata::MetadataKind::kMetadata) {
+    return;
+  }
+
+  bool needs_instance = DiscoverType(arr->type_key);
+  for (unsigned int i = 0; i < arr->array.size(); i++) {
+    tvm::runtime::metadata::MetadataBase o = Downcast<tvm::runtime::metadata::MetadataBase>(arr->array[i]);
+    if (needs_instance) {
+      DiscoverInstance(o);
+      needs_instance = false;
+    }
+    ReflectionVTable::Global()->VisitAttrs(o.operator->(), this);
+  }
+}
+
+void DiscoverComplexTypesVisitor::Discover(runtime::metadata::MetadataBase metadata) {
+  ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+  DiscoverType(metadata->GetTypeKey());
+  DiscoverInstance(metadata);
+}
+
+}  // namespace metadata
+}  // namespace codegen
+}  // namespace tvm
diff --git a/src/target/metadata_utils.h b/src/target/metadata_utils.h
new file mode 100644
index 000000000000..6848e25264d1
--- /dev/null
+++ b/src/target/metadata_utils.h
@@ -0,0 +1,122 @@
+/*
+ * 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 tvm/target/metadata_utils.h
+ * \brief Declares utilty functions and classes for emitting metadata.
+ */
+#ifndef TVM_TARGET_METADATA_UTILS_H_
+#define TVM_TARGET_METADATA_UTILS_H_
+
+#include <string>
+#include <tuple>
+#include <vector>
+#include <tvm/runtime/data_type.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/object.h>
+#include "metadata.h"
+
+namespace tvm {
+namespace codegen {
+
+namespace metadata {
+
+std::string address_from_parts(const std::vector<std::string>& parts);
+static constexpr const char* kMetadataGlobalSymbol = "kTvmgenMetadata";
+
+/*!
+ * \brief Post-order traverse metadata to discover arrays which need to be forward-defined.
+ */
+class DiscoverArraysVisitor : public AttrVisitor {
+ public:
+  /*! \brief Models a single array discovered in this visitor.
+   * Conatains two fields:
+   *  0. An address which uniquely identifies the array in this Metadata instance.
+   *  1. The discovered MetadataArray.
+   */
+  using DiscoveredArray = std::tuple<std::string, runtime::metadata::MetadataArray>;
+  explicit DiscoverArraysVisitor(std::vector<DiscoveredArray>* queue);
+
+  void Visit(const char* key, double* value) final;
+  void Visit(const char* key, int64_t* value) final;
+  void Visit(const char* key, uint64_t* value) final;
+  void Visit(const char* key, int* value) final;
+  void Visit(const char* key, bool* value) final;
+  void Visit(const char* key, std::string* value) final;
+  void Visit(const char* key, DataType* value) final;
+  void Visit(const char* key, runtime::NDArray* value) final;
+  void Visit(const char* key, void** value) final;
+
+  void Visit(const char* key, ObjectRef* value) final;
+ private:
+  /*! \brief The queue to be filled with discovered arrays. */
+  std::vector<DiscoveredArray>* queue_;
+
+  /*! \brief Tracks the preceding address pieces. */
+  std::vector<std::string> address_parts_;
+};
+
+/*!
+ * \brief Post-order traverse Metadata to discover all complex types which need to be forward-defined.
+ * This visitor finds one defined() MetadataBase instance for each unique subclass present inside
+ * Metadata in the order in which the subclass was first discovered.
+ */
+class DiscoverComplexTypesVisitor : public AttrVisitor {
+ public:
+  /*! \brief Models a single complex type discovered in this visitor.
+   * Contains two fields:
+   *  0. The struct_name for this Metadata instance.
+   *  1. The discovered MetadataArray.
+   */
+  using DiscoveredComplexType = std::tuple<std::string, runtime::metadata::MetadataBase>;
+
+  /*! \brief Construct a new instance.
+   * \param queue An ordered map which holds the
+   */
+  explicit DiscoverComplexTypesVisitor(std::vector<runtime::metadata::MetadataBase>* queue) : queue_{queue} {}
+
+  void Visit(const char* key, double* value) final;
+  void Visit(const char* key, int64_t* value) final;
+  void Visit(const char* key, uint64_t* value) final;
+  void Visit(const char* key, int* value) final;
+  void Visit(const char* key, bool* value) final;
+  void Visit(const char* key, std::string* value) final;
+  void Visit(const char* key, DataType* value) final;
+  void Visit(const char* key, runtime::NDArray* value) final;
+  void Visit(const char* key, void** value) final;
+
+  void Visit(const char* key, ObjectRef* value) final;
+
+  void Discover(runtime::metadata::MetadataBase metadata);
+ private:
+  bool DiscoverType(std::string type_key);
+
+  void DiscoverInstance(runtime::metadata::MetadataBase md);
+
+  std::vector<runtime::metadata::MetadataBase>* queue_;
+
+  /*! \brief map type_index to index in queue_. */
+  std::unordered_map<std::string, int> type_key_to_position_;
+};
+
+}  // namespace metadata
+}  // namespace codegen
+}  // namespace tvm
+
+#endif  // TVM_TARGET_METADATA_UTILS_H_
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index e6f81646242d..01c1c911b7de 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -28,6 +28,7 @@
 #include <iomanip>
 
 #include "../../arith/pattern_match.h"
+#include "codegen_params.h"
 
 namespace tvm {
 namespace codegen {
@@ -648,6 +649,37 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op, DataType t, PrimExpr lhs,
   }
 }
 
+void CodeGenC::VisitStmt_(const AllocateConstNode* op) {
+  std::string symbol_name = op->buffer_var->name_hint;
+  int64_t num_elements = 1;
+  const auto& data = op->data.value();
+
+  for (int64_t dim : data.Shape()) {
+    num_elements *= dim;
+  }
+
+  decl_stream << "\n"
+              << "#ifdef __cplusplus\n"
+              << "extern \"C\" {\n"
+              << "#endif\n"
+              << "static const ";
+
+  PrintType(data.DataType(), decl_stream);
+
+  // Allocate the global static variable
+  decl_stream << " __attribute__((section(\".rodata.tvm\"), "
+              << "aligned(" << constants_byte_alignment_->value << "))) " << symbol_name << "["
+              << num_elements << "] = {\n";
+  NDArrayDataToC(data, 4, decl_stream);
+
+  decl_stream << "};\n"
+              << "#ifdef __cplusplus\n"
+              << "}  // extern \"C\"\n"
+              << "#endif\n";
+  var_idmap_[op->buffer_var.operator->()] = symbol_name;
+  this->PrintStmt(op->body);
+}
+
 void CodeGenC::VisitExpr_(const LoadNode* op, std::ostream& os) {  // NOLINT(*)
   int lanes = op->dtype.lanes();
   // delcare type.
@@ -820,7 +852,7 @@ void CodeGenC::VisitStmt_(const AllocateNode* op) {
   std::string vid = AllocVarID(op->buffer_var.get());
 
   this->PrintIndent();
-  int32_t constant_size = op->constant_allocation_size();
+  size_t constant_size = op->ConstantAllocationSize();
   ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";
 
   auto scope = GetPtrStorageScope(op->buffer_var);
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index 3b042b9fbd2c..2af77bb28b53 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -163,6 +163,7 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
   void VisitStmt_(const AssertStmtNode* op) override;
   void VisitStmt_(const EvaluateNode* op) override;
   void VisitStmt_(const SeqStmtNode* op) override;
+  void VisitStmt_(const AllocateConstNode* op) override;
 
   /*!
    * \brief Print expr representing the thread tag
@@ -192,6 +193,10 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
   // Print restrict keyword for a given Var if applicable
   virtual void PrintRestrict(const Var& v, std::ostream& os);
 
+  virtual void SetConstantsByteAlignment(Integer constants_byte_alignment) {
+    constants_byte_alignment_ = constants_byte_alignment;
+  }
+
  protected:
   // Print reference to struct location
   std::string GetStructRef(DataType t, const PrimExpr& buffer, const PrimExpr& index, int kind);
@@ -262,6 +267,7 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
   // cache commonly used ops
   const Op& builtin_call_extern_ = builtin::call_extern();
   const Op& builtin_call_pure_extern_ = builtin::call_pure_extern();
+  Integer constants_byte_alignment_ = 16;
 
  private:
   /*! \brief whether to print in SSA form */
diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc
index 515cdccb88fb..71563df636f0 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -23,6 +23,7 @@
 #include "codegen_c_host.h"
 
 #include <tvm/relay/executor.h>
+#include <tvm/relay/runtime.h>
 #include <tvm/runtime/crt/error_codes.h>
 #include <tvm/runtime/module.h>
 #include <tvm/target/codegen.h>
@@ -51,6 +52,10 @@ void CodeGenCHost::Init(bool output_ssa, bool emit_asserts, std::string target_s
   CodeGenC::Init(output_ssa);
 }
 
+void CodeGenCHost::InitGlobalContext() {
+  decl_stream << "void* " << tvm::runtime::symbol::tvm_module_ctx << " = NULL;\n";
+}
+
 void CodeGenCHost::DefineModuleName() { decl_stream << "void* " << module_name_ << " = NULL;\n"; }
 
 void CodeGenCHost::AddFunction(const PrimFunc& f) {
@@ -73,56 +78,6 @@ void CodeGenCHost::AddFunction(const PrimFunc& f) {
   }
 }
 
-void CodeGenCHost::DeclareParameters(Map<String, LinkedParam> params,
-                                     const Integer& constants_byte_alignment) {
-  for (auto kv : params) {
-    decl_stream << "\n"
-                << "#ifdef __cplusplus\n"
-                << "extern \"C\" {\n"
-                << "#endif\n"
-                << "static const ";
-    int64_t num_elements = 1;
-    for (int64_t dim : kv.second->param.Shape()) {
-      num_elements *= dim;
-    }
-    PrintType(kv.second->param.DataType(), decl_stream);
-    decl_stream << " __attribute__((section(\".rodata.tvm\"), "
-                << "aligned(" << constants_byte_alignment->value << "))) "
-                << ::tvm::runtime::symbol::tvm_param_prefix << kv.first << "[" << num_elements
-                << "] = {\n";
-    NDArrayDataToC(kv.second->param, 4, decl_stream);
-    decl_stream << "};\n"
-                << "#ifdef __cplusplus\n"
-                << "}  // extern \"C\"\n"
-                << "#endif\n";
-  }
-}
-
-void CodeGenCHost::LinkParameters(Map<String, LinkedParam> params) {
-  PrintFuncPrefix();
-  stream << " " << tvm::runtime::symbol::tvm_lookup_linked_param
-         << "(void* args, int* arg_type_ids, int num_args, void* out_ret_value, "
-         << "int* out_ret_tcode, void* resource_handle) {\n";
-  ICHECK_EQ(GetUniqueName(tvm::runtime::symbol::tvm_lookup_linked_param),
-            tvm::runtime::symbol::tvm_lookup_linked_param)
-      << "builtin PackedFunc name already taken: " << tvm::runtime::symbol::tvm_lookup_linked_param;
-  stream << "    switch (((int64_t*) args)[0]) {\n"
-         << "    default:\n"
-         << "        out_ret_tcode[0] = " << kTVMNullptr << ";\n"
-         << "        return 0;\n";
-
-  function_names_.push_back(tvm::runtime::symbol::tvm_lookup_linked_param);
-  for (auto kv : params) {
-    stream << "    case " << kv.second->id << ":\n"
-           << "        ((uint64_t*)out_ret_value)[0] = (uint64_t) (uintptr_t) "
-           << ::tvm::runtime::symbol::tvm_param_prefix << kv.first << ";\n"
-           << "        out_ret_tcode[0] = " << kTVMOpaqueHandle << ";\n"
-           << "        return 0;\n";
-  }
-  stream << "    }\n"
-         << "}\n";
-}
-
 void CodeGenCHost::PrintFuncPrefix() {  // NOLINT(*)
   stream << "#ifdef __cplusplus\n"
          << "extern \"C\"\n"
@@ -389,26 +344,14 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare,
 runtime::Module BuildCHost(IRModule mod, Target target) {
   using tvm::runtime::Registry;
   bool output_ssa = false;
-  bool emit_asserts = false;
+  bool emit_asserts = true;
   CodeGenCHost cg;
   cg.Init(output_ssa, emit_asserts, target->str());
-
+  cg.SetConstantsByteAlignment(target->GetAttr<Integer>("constants-byte-alignment").value_or(16));
   Map<String, LinkedParam> linked_params;
-  bool found_linked_params = false;
-  bool could_have_linked_params = mod->ShouldLinkParameters();
   PrimFunc aot_executor_fn;
 
   for (auto kv : mod->functions) {
-    if (could_have_linked_params &&
-        kv.first->name_hint == ::tvm::runtime::symbol::tvm_lookup_linked_param) {
-      Map<String, ObjectRef> attrs_dict = Downcast<Map<String, ObjectRef>>(kv.second->attrs->dict);
-      CHECK(attrs_dict.find(::tvm::tir::attr::kLinkedParams) != attrs_dict.end())
-          << "no " << ::tvm::tir::attr::kLinkedParams << " attribute found!";
-      linked_params =
-          Downcast<Map<String, LinkedParam>>(attrs_dict[::tvm::tir::attr::kLinkedParams]);
-      found_linked_params = true;
-      continue;
-    }
     // Make sure that the executor function is the last one to be code generated so that all the
     // symbols are available to tvm_run_func
     auto fun_name = std::string(kv.first->name_hint);
@@ -424,17 +367,21 @@ runtime::Module BuildCHost(IRModule mod, Target target) {
     cg.AddFunction(f);
   }
 
-  auto constants_byte_alignment = target->GetAttr<Integer>("constants-byte-alignment").value_or(16);
-
-  if (could_have_linked_params && !aot_executor_fn.defined()) {
-    ICHECK(found_linked_params) << "-link-params given but none found";
-    cg.DeclareParameters(linked_params, constants_byte_alignment);
-    cg.LinkParameters(linked_params);
+  if (aot_executor_fn.defined()) {
+    cg.AddFunction(aot_executor_fn);
   }
 
-  if (could_have_linked_params && aot_executor_fn.defined()) {
-    cg.DeclareParameters(linked_params, constants_byte_alignment);
-    cg.AddFunction(aot_executor_fn);
+  // NOTE: it's possible that kRuntime attr is not attached when the mod was built with tvm.build().
+  // See issue #10373.
+  auto opt_runtime = mod->GetAttr<relay::Runtime>(tvm::attr::kRuntime);
+  relay::Runtime runtime;
+  if (opt_runtime.get() != nullptr) {
+    runtime = opt_runtime.value();
+  } else {
+    runtime = relay::Runtime::Create("cpp", {});
+  }
+  if (aot_executor_fn.defined() && runtime->name == relay::kTvmRuntimeCpp) {
+    cg.InitGlobalContext();
   }
 
   if (target->GetAttr<Bool>("system-lib").value_or(Bool(false))) {
diff --git a/src/target/source/codegen_c_host.h b/src/target/source/codegen_c_host.h
index c94612cfeac3..c0e4ee9a263c 100644
--- a/src/target/source/codegen_c_host.h
+++ b/src/target/source/codegen_c_host.h
@@ -40,14 +40,11 @@ class CodeGenCHost : public CodeGenC {
   CodeGenCHost();
   void Init(bool output_ssa, bool emit_asserts, std::string target_str);
 
+  void InitGlobalContext();
   void AddFunction(const PrimFunc& f);
 
   void DefineModuleName();
 
-  /*! \brief Add linked parameters, if they are present. */
-  void DeclareParameters(Map<String, LinkedParam> params, const Integer& constants_byte_alignment);
-  void LinkParameters(Map<String, LinkedParam> params);
-
   void PrintType(DataType t, std::ostream& os) final;  // NOLINT(*)
   void PrintFuncPrefix() final;                        // NOLINT(*)
   void PrintFinalReturn() final;                       // NOLINT(*)
diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc
index a1f257391db4..984f8a13351e 100644
--- a/src/target/source/codegen_cuda.cc
+++ b/src/target/source/codegen_cuda.cc
@@ -803,7 +803,7 @@ void CodeGenCUDA::VisitStmt_(const AllocateNode* op) {
   if (scope == "shared.dyn") {
     stream << ' ' << vid << "[];\n";
   } else {
-    int32_t constant_size = op->constant_allocation_size();
+    size_t constant_size = op->ConstantAllocationSize();
     ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";
 
     if (scope.find("wmma.") == 0) {
@@ -1048,7 +1048,7 @@ void CodeGenCUDA::PrintWmmaScope(const std::string& scope, DataType t, const Var
                                  std::ostream& os) {
   std::stringstream type;
   PrintType(t, type);
-  std::string shape_str = fragment_shapes[variable];
+  std::string shape_str = fragment_shapes.at(variable);
   if ((t.is_int() || t.is_uint()) && t.bits() < 8 && t.lanes() == 1) {
     type.str(std::string());
     if (t.is_int()) {
@@ -1084,18 +1084,27 @@ void CodeGenCUDA::PrintWmmaScope(const std::string& scope, DataType t, const Var
   }
 }
 
+int stoi(const std::string& str) {
+  try {
+    return std::stoi(str);
+  } catch (std::invalid_argument& e) {
+    LOG(FATAL) << "Cannot convert \"" << str << "\" to int";
+    throw;
+  }
+}
+
 int32_t CodeGenCUDA::GetWmmaFragmentSize(const std::string& scope, const VarNode* variable,
                                          int32_t size) {
-  std::string shape_str = fragment_shapes[variable];
+  std::string shape_str = fragment_shapes.at(variable);
   size_t m, n, k;
   size_t last_pos = 0, pos = 0;
   pos = shape_str.find(", ", last_pos);
-  m = std::stoi(shape_str.substr(last_pos, pos - last_pos));
+  m = tvm::codegen::stoi(shape_str.substr(last_pos, pos - last_pos));
   last_pos = pos + 2;
   pos = shape_str.find(", ", last_pos);
-  n = std::stoi(shape_str.substr(last_pos, pos - last_pos));
+  n = tvm::codegen::stoi(shape_str.substr(last_pos, pos - last_pos));
   last_pos = pos + 2;
-  k = std::stoi(shape_str.substr(last_pos, shape_str.length() - last_pos));
+  k = tvm::codegen::stoi(shape_str.substr(last_pos, shape_str.length() - last_pos));
   if (scope == "wmma.matrix_a") {
     return size / m / k;
   } else if (scope == "wmma.matrix_b") {
diff --git a/src/target/source/codegen_opencl.cc b/src/target/source/codegen_opencl.cc
index 507a6243cb0c..a9cd9d8ae930 100644
--- a/src/target/source/codegen_opencl.cc
+++ b/src/target/source/codegen_opencl.cc
@@ -364,8 +364,7 @@ void CodeGenOpenCL::VisitExpr_(const CastNode* op, std::ostream& os) {
 }
 
 void CodeGenOpenCL::VisitStmt_(const AllocateNode* op) {
-  allocation_size_.insert(
-      {op->buffer_var.get(), op->constant_allocation_size() * op->dtype.lanes()});
+  allocation_size_.insert({op->buffer_var.get(), op->ConstantAllocationSize() * op->dtype.lanes()});
   CodeGenC::VisitStmt_(op);
 }
 
diff --git a/src/target/source/codegen_opencl.h b/src/target/source/codegen_opencl.h
index 8c36a817753c..2670c601c43c 100644
--- a/src/target/source/codegen_opencl.h
+++ b/src/target/source/codegen_opencl.h
@@ -83,7 +83,7 @@ class CodeGenOpenCL final : public CodeGenC {
   bool need_texture_ssa_{true};
   // Mapping from buffer to allocation size.
   // Useful to track when a scalar store of a vectorized texture load is required.
-  std::unordered_map<const Object*, int32_t> allocation_size_;
+  std::unordered_map<const Object*, size_t> allocation_size_;
 };
 
 }  // namespace codegen
diff --git a/src/target/source/codegen_params.cc b/src/target/source/codegen_params.cc
index cc7695abfd25..798ef73f0fa8 100644
--- a/src/target/source/codegen_params.cc
+++ b/src/target/source/codegen_params.cc
@@ -66,14 +66,11 @@ void PrintIntegralArray(void* data, size_t num_elements, int indent_chars, std::
     }
   }
 
-  int elements_per_row = ComputeNumElementsPerRow(one_element_size_bytes, indent_chars);
+  size_t elements_per_row = ComputeNumElementsPerRow(one_element_size_bytes, indent_chars);
   std::string indent_str(indent_chars, ' ');
 
   for (size_t i = 0; i < num_elements; i++) {
     if ((i % elements_per_row) == 0) {
-      if (i != 0) {
-        os << std::endl;
-      }
       os << indent_str;
     }
     int64_t elem = static_cast<T*>(data)[i];
@@ -99,6 +96,9 @@ void PrintIntegralArray(void* data, size_t num_elements, int indent_chars, std::
     if (i < num_elements - 1) {
       os << ", ";
     }
+    if ((i % elements_per_row) == elements_per_row - 1) {
+      os << "\n";
+    }
   }
 
   if ((num_elements % elements_per_row) != 0) {
@@ -117,7 +117,7 @@ void PrintFloatingPointArray(void* data, size_t num_elements, int indent_chars,
     one_element_size_bytes += 1; /* extra decimal digit in exponent, relative to bits / 4 */
   }
 
-  int elements_per_row = ComputeNumElementsPerRow(one_element_size_bytes, indent_chars);
+  size_t elements_per_row = ComputeNumElementsPerRow(one_element_size_bytes, indent_chars);
   std::string indent_str(indent_chars, ' ');
 
   std::stringstream ss;
@@ -130,9 +130,6 @@ void PrintFloatingPointArray(void* data, size_t num_elements, int indent_chars,
   }
   for (size_t i = 0; i < num_elements; i++) {
     if ((i % elements_per_row) == 0) {
-      if (i != 0) {
-        os << std::endl;
-      }
       os << indent_str;
     }
 
@@ -151,6 +148,9 @@ void PrintFloatingPointArray(void* data, size_t num_elements, int indent_chars,
     if (i < num_elements - 1) {
       os << ", ";
     }
+    if ((i % elements_per_row) == elements_per_row - 1) {
+      os << "\n";
+    }
   }
 
   if ((num_elements % elements_per_row) != 0) {
diff --git a/src/target/source/codegen_source_base.h b/src/target/source/codegen_source_base.h
index 8f8f9e1b8bf2..133886758430 100644
--- a/src/target/source/codegen_source_base.h
+++ b/src/target/source/codegen_source_base.h
@@ -25,6 +25,7 @@
 #ifndef TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_
 #define TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_
 
+#include <tvm/runtime/metadata.h>
 #include <tvm/target/codegen.h>
 #include <tvm/tir/expr.h>
 #include <tvm/tir/op.h>
@@ -157,6 +158,19 @@ runtime::Module CSourceModuleCreate(const String& code, const String& fmt,
                                     const Array<String>& func_names,
                                     const Array<String>& const_vars = {});
 
+/*!
+ * \brief Wrap the submodules in a metadata module.
+ * \param params The variable to constant mapping that is collected by the host
+ *        module.
+ * \param target_module The main TIR-lowered internal runtime module
+ * \param modules All the external modules that needs to be imported inside the metadata module(s).
+ * \param target The target that all the modules are compiled for
+ * \return The wrapped module.
+ */
+runtime::Module CreateMetadataModule(
+    const std::unordered_map<std::string, runtime::NDArray>& params, runtime::Module target_module,
+    const Array<runtime::Module>& ext_modules, Target target, runtime::metadata::Metadata metadata);
+
 /*!
  * \brief Create a source module for viewing and limited saving for device.
  * \param data The code data to be viewed.
@@ -169,6 +183,16 @@ runtime::Module DeviceSourceModuleCreate(
     std::string data, std::string fmt, std::unordered_map<std::string, runtime::FunctionInfo> fmap,
     std::string type_key, std::function<std::string(const std::string&)> fget_source = nullptr);
 
+/*!
+ * \brief Wrap the submodules that are to be wrapped in a c-source metadata module for C runtime.
+ * \param modules The modules to be wrapped.
+ * \param target the target the modules are compiled for.
+ * \param metadata the metadata needed for code generation.
+ * \return The wrapped module.
+ */
+runtime::Module CreateCSourceCrtMetadataModule(const Array<runtime::Module>& modules, Target target,
+                                               runtime::metadata::Metadata metadata);
+
 }  // namespace codegen
 }  // namespace tvm
 #endif  // TVM_TARGET_SOURCE_CODEGEN_SOURCE_BASE_H_
diff --git a/src/target/source/literal/cuda_half_t.h b/src/target/source/literal/cuda_half_t.h
index 0ecefc3178f6..c3a1c66e5874 100644
--- a/src/target/source/literal/cuda_half_t.h
+++ b/src/target/source/literal/cuda_half_t.h
@@ -338,7 +338,7 @@ __pack_nv_bfloat162(const nv_bfloat16 x, const nv_bfloat16 y) {
 // so we define them here to make sure the generated CUDA code
 // is valid.
 #define CUDA_UNSUPPORTED_HALF_MATH_BINARY(HALF_MATH_NAME, FP32_MATH_NAME) \
-static inline __device__ __host__ half HALF_MATH_NAME(half x, half y) {   \
+static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x, nv_bfloat16 y) {   \
   float tmp_x = __bfloat162float(x);                                      \
   float tmp_y = __bfloat162float(y);                                      \
   float result = FP32_MATH_NAME(tmp_x, tmp_y);                            \
@@ -346,7 +346,7 @@ static inline __device__ __host__ half HALF_MATH_NAME(half x, half y) {   \
 }
 
 #define CUDA_UNSUPPORTED_HALF_MATH_UNARY(HALF_MATH_NAME, FP32_MATH_NAME) \
-static inline __device__ __host__ half HALF_MATH_NAME(half x) {          \
+static inline __device__ __host__ nv_bfloat16 HALF_MATH_NAME(nv_bfloat16 x) {          \
   float tmp_x = __bfloat162float(x);                                     \
   float result = FP32_MATH_NAME(tmp_x);                                  \
   return __float2bfloat16(result);                                       \
diff --git a/src/target/source/source_module.cc b/src/target/source/source_module.cc
index 29d867ad719d..61aa912ae0b8 100644
--- a/src/target/source/source_module.cc
+++ b/src/target/source/source_module.cc
@@ -23,18 +23,23 @@
  */
 #include "source_module.h"
 
+#include <tvm/runtime/metadata.h>
 #include <tvm/runtime/ndarray.h>
 #include <tvm/runtime/packed_func.h>
 #include <tvm/runtime/registry.h>
 
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
+#include <vector>
 
 #include "../../relay/backend/name_transforms.h"
 #include "../../runtime/file_utils.h"
 #include "../../support/str_escape.h"
 #include "../func_registry_generator.h"
+#include "../metadata.h"
+#include "../metadata_utils.h"
 #include "codegen_source_base.h"
 
 namespace tvm {
@@ -273,7 +278,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
         }
         call_args_ss << " " << input_var->name_hint << ",";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         call_args_ss << "void* output" << i << ",";
       }
       for (const tir::Var& pool_var : metadata_->pools) {
@@ -300,7 +305,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       for (unsigned int i = 0; i < metadata_->inputs.size(); ++i) {
         call_args_ss << "((DLTensor*)(((TVMValue*)args)[" << i << "].v_handle))[0].data,";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         int j = metadata_->inputs.size() + i;
         call_args_ss << "((DLTensor*)(((TVMValue*)args)[" << j << "].v_handle))[0].data,";
       }
@@ -328,7 +333,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       entrypoint_arg_count++;
       run_func_arg_count++;
     }
-    for (int i = 0; i < metadata_->num_outputs->value; i++) {
+    for (unsigned int i = 0; i < metadata_->outputs.size(); i++) {
       run_func_to_entry_point_args[run_func_arg_count] = Integer(entrypoint_arg_count);
       entrypoint_arg_count++;
       run_func_arg_count++;
@@ -356,7 +361,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
 
     // We are creating a copy of the set of pointers
     size_t number_of_io_tensors =
-        metadata_->inputs.size() + metadata_->num_outputs->value + metadata_->pools.size();
+        metadata_->inputs.size() + metadata_->outputs.size() + metadata_->pools.size();
     code_ << "TVMValue tensors[" << number_of_io_tensors << "];\n";
 
     std::unordered_map<int, ObjectRef> run_func_to_entry_point_args =
@@ -395,7 +400,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
         }
         call_args_ss << " " << relay::backend::SanitizeName(input_var->name_hint) << ",";
       }
-      for (int i = 0; i < metadata_->num_outputs->value; ++i) {
+      for (unsigned int i = 0; i < metadata_->outputs.size(); ++i) {
         call_args_ss << "void* output" << i << ",";
       }
       for (const tir::Var& pool_var : metadata_->pools) {
@@ -449,13 +454,11 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
       for (const auto& input : metadata_->inputs) {
         call_args_ss << "inputs->" << relay::backend::SanitizeName(input->name_hint) << ",";
       }
-      if (metadata_->num_outputs->value == 1) {
-        call_args_ss << "outputs->output,";
-      } else {
-        for (int i = 0; i < metadata_->num_outputs->value; ++i) {
-          call_args_ss << "outputs->output" << i << ",";
-        }
+      for (const auto& output : metadata_->outputs) {
+        call_args_ss << "outputs->" << relay::backend::SanitizeName(output);
+        call_args_ss << ",";
       }
+
       for (const tir::Var& pool_var : metadata_->pools) {
         String pool_name = metadata_->pool_inputs.value()[pool_var]->pool_info->pool_name;
         if (IsInternalWorkspaceBuffer(pool_var)) {
@@ -476,7 +479,7 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
   }
 
   void GenerateAOTDescriptor() {
-    const std::string run_func_suffix = ::tvm::runtime::symbol::tvm_run_func_suffix;
+    const std::string run_func_suffix = ::tvm::runtime::symbol::tvm_module_main;
     const std::string tvm_entrypoint_suffix = ::tvm::runtime::symbol::tvm_entrypoint_suffix;
     const std::string run_func_mangled =
         runtime::get_name_mangled(metadata_->mod_name, run_func_suffix);
@@ -520,6 +523,230 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
   }
 };
 
+class MetadataSerializer : public AttrVisitor {
+ public:
+  using MetadataKind = ::tvm::runtime::metadata::MetadataKind;
+
+  MetadataSerializer() : is_first_item_{true} {}
+
+  void WriteComma() {
+    if (is_first_item_) {
+      is_first_item_ = false;
+    } else {
+      code_ << ", " << std::endl;
+    }
+  }
+
+  void WriteKey(const char* key) {
+    if (key != nullptr) {
+      code_ << " /* " << key << "*/";
+    }
+  }
+
+  void Visit(const char* key, double* value) final {
+    WriteComma();
+    code_.setf(std::ios::hex | std::ios::showbase | std::ios::fixed | std::ios::scientific,
+               std::ios::basefield | std::ios::showbase | std::ios::floatfield);
+    code_ << *value;
+    WriteKey(key);
+  }
+
+  void Visit(const char* key, int64_t* value) final {
+    WriteComma();
+    code_ << *value << "L";
+    WriteKey(key);
+  }
+
+  void Visit(const char* key, uint64_t* value) final {
+    WriteComma();
+    code_ << *value << "UL";
+    WriteKey(key);
+  }
+  void Visit(const char* key, int* value) final {
+    WriteComma();
+    code_ << *value;
+    WriteKey(key);
+  }
+  void Visit(const char* key, bool* value) final {
+    WriteComma();
+    code_ << *value;
+    WriteKey(key);
+  }
+  void Visit(const char* key, std::string* value) final {
+    WriteComma();
+    code_ << "\"" << *value << "\"";
+    WriteKey(key);
+  }
+  void Visit(const char* key, void** value) final {
+    WriteComma();
+    code_ << *value;
+    WriteKey(key);
+  }
+  void Visit(const char* key, DataType* value) final {
+    WriteComma();
+    code_ << "{" << value->code() << ", " << value->bits() << ", " << value->lanes() << "}";
+    WriteKey(key);
+  }
+
+  void Visit(const char* key, runtime::NDArray* value) final {
+    // TODO(areusch): probably we could consolidate --link-params here, tho...
+    ICHECK(false) << "do not support serializing NDArray as metadata";
+  }
+
+  void VisitArray(runtime::metadata::MetadataArray array) {
+    auto old_is_first_item = is_first_item_;
+    is_first_item_ = true;
+    for (unsigned int i = 0; i < array->array.size(); ++i) {
+      ObjectRef o = array->array[i];
+
+      switch (array->kind) {
+      case MetadataKind::kUint64: {
+        int64_t i = Downcast<Integer>(o);
+        CHECK_GT(i, 0) << "Metadata is of type uint64_t, but array type contains a negative number";
+        uint64_t ui = static_cast<uint64_t>(i);
+        Visit(nullptr, &ui);
+        continue;
+      }
+      case MetadataKind::kInt64: {
+        int64_t i = Downcast<Integer>(o);
+        Visit(nullptr, &i);
+        continue;
+      }
+      case MetadataKind::kBool: {
+        bool b = Downcast<Bool>(o);
+        Visit(nullptr, &b);
+        break;
+      }
+      case MetadataKind::kString: {
+        std::string s = Downcast<String>(o);
+        Visit(nullptr, &s);
+        break;
+      }
+      case MetadataKind::kHandle:
+        CHECK(false) << "Don't know how to serialize handle";
+        break;
+
+      case MetadataKind::kMetadata: {
+        runtime::metadata::MetadataBase metadata = Downcast<runtime::metadata::MetadataBase>(o);
+        std::stringstream i_str;
+        i_str << i;
+        address_.push_back(i_str.str());
+        Visit(nullptr, &metadata);
+        address_.pop_back();
+        break;
+      }
+      default:
+        CHECK(false) << "Unknown MetadataKind for array: " << array->kind;
+        break;
+      }
+      is_first_item_ = false;
+    }
+    is_first_item_ = old_is_first_item;
+  }
+
+  void Visit(const char* key, ObjectRef* value) final {
+    const runtime::metadata::MetadataArrayNode* arr =
+        value->as<runtime::metadata::MetadataArrayNode>();
+    std::cout << "Is array? " << arr << std::endl;
+    if (arr != nullptr) {
+      WriteComma();
+      if (key != nullptr) {
+        address_.push_back(key);
+      }
+      code_ << metadata::address_from_parts(address_);
+      if (key != nullptr) {
+        address_.pop_back();
+      }
+      return;
+    }
+
+    runtime::metadata::MetadataBase metadata = Downcast<runtime::metadata::MetadataBase>(*value);
+    if (key != nullptr) {  // NOTE: outermost call passes nullptr key
+      address_.push_back(key);
+    }
+    ReflectionVTable::Global()->VisitAttrs(metadata.operator->(), this);
+    if (key != nullptr) {  // NOTE: outermost call passes nullptr key
+      address_.pop_back();
+    }
+  }
+
+private:
+  void EmitCType(const runtime::metadata::MetadataArrayNode* arr, std::ostream& os) {
+    switch (arr->kind) {
+      case MetadataKind::kUint64:
+        os << "uint64_t";
+        break;
+      case MetadataKind::kInt64:
+        os << "int64_t";
+        break;
+      case MetadataKind::kBool:
+        os << "bool";
+        break;
+      case MetadataKind::kString:
+        os << "const char*";
+        break;
+      case MetadataKind::kHandle:
+        os << "void*";
+        break;
+      case MetadataKind::kMetadata:
+        os << "struct " << arr->get_element_c_struct_name();
+        break;
+      default:
+        CHECK(false) << "Unknown kind in MetadataArray: " << arr->kind
+                     << " (struct_name=" << arr->get_c_struct_name() << ")";
+        break;
+    }
+  }
+public:
+
+  void CodegenMetadata(::tvm::runtime::metadata::Metadata metadata) {
+    decl_ << "#include <inttypes.h>" << std::endl
+          << "#include <tvm/runtime/metadata.h>" << std::endl
+          << "#include <tvm/runtime/c_runtime_api.h>" << std::endl;
+    std::vector<metadata::DiscoverArraysVisitor::DiscoveredArray> queue;
+    metadata::DiscoverArraysVisitor array_discover{&queue};
+    array_discover.Visit(metadata::kMetadataGlobalSymbol, &metadata);
+
+    for (auto item : queue) {
+      auto struct_address = std::get<0>(item);
+      address_.push_back(struct_address);
+
+      auto arr = std::get<1>(item);
+
+      // Prepend const with everything except C-string, which needs appending.
+      if (arr->kind != MetadataKind::kString) {
+        code_ << "const ";
+      }
+      EmitCType(arr.operator->(), code_);
+      if (arr->kind == MetadataKind::kString) {
+        code_ << " const";
+      }
+      code_ << " " << struct_address << "[" << arr->array.size() << "] = {" << std::endl;
+      is_first_item_ = true;
+
+      VisitArray(arr);
+      address_.pop_back();
+      code_ << "};" << std::endl;
+    }
+
+    // Finally, emit overall struct.
+    code_ << "const struct TVMMetadata " << metadata::kMetadataGlobalSymbol << " = {" << std::endl;
+    Visit(nullptr, &metadata);
+    code_ << "};" << std::endl;
+  }
+
+  std::string GetOutput() { return decl_.str() + code_.str(); }
+
+ private:
+  std::vector<std::string> address_;
+  std::stringstream decl_;
+  std::stringstream code_;
+  bool is_first_item_;
+  std::unordered_set<std::string> generated_struct_decls_;
+  std::vector<bool> is_defining_struct_;
+
+};
+
 runtime::Module CreateCSourceCrtMetadataModule(const Array<runtime::Module>& modules, Target target,
                                                relay::Runtime runtime,
                                                relay::backend::ExecutorCodegenMetadata metadata) {
@@ -541,6 +768,32 @@ runtime::Module CreateCSourceCrtMetadataModule(const Array<runtime::Module>& mod
   return std::move(csrc_metadata_module);
 }
 
+runtime::Module CreateCSourceCppMetadataModule(runtime::metadata::Metadata metadata) {
+  MetadataSerializer serializer;
+  serializer.CodegenMetadata(metadata);
+  std::stringstream lookup_func;
+  lookup_func << "#ifdef __cplusplus\n"
+              << "extern \"C\"\n"
+              << "#endif\n";
+
+  lookup_func << "TVM_DLL int32_t get_c_metadata(TVMValue* arg_values, int* arg_tcodes, int "
+                 "num_args, TVMValue* ret_values, int* ret_tcodes, void* resource_handle) {"
+              << std::endl;
+  lookup_func << "    ret_values[0].v_handle = (void*) &" << metadata::kMetadataGlobalSymbol
+              << ";" << std::endl;
+  lookup_func << "    ret_tcodes[0] = kTVMOpaqueHandle;" << std::endl;
+  lookup_func << "    return 0;" << std::endl;
+  lookup_func << "};" << std::endl;
+
+  auto mod = MetadataModuleCreate(metadata);
+  std::vector<String> func_names{"get_c_metadata"};
+  // definer.GetOutput() +
+  auto c = CSourceModuleCreate(serializer.GetOutput() + lookup_func.str(), "c", func_names,
+                               Array<String>());
+  mod->Import(c);
+  return mod;
+}
+
 // supports limited save without cross compile
 class DeviceSourceModuleNode final : public runtime::ModuleNode {
  public:
diff --git a/src/target/source/source_module.h b/src/target/source/source_module.h
index 3b482a107600..2a63a8eeb814 100644
--- a/src/target/source/source_module.h
+++ b/src/target/source/source_module.h
@@ -26,6 +26,7 @@
 #define TVM_TARGET_SOURCE_SOURCE_MODULE_H_
 
 #include <tvm/relay/runtime.h>
+#include <tvm/runtime/metadata.h>
 #include <tvm/runtime/module.h>
 #include <tvm/target/target.h>
 
@@ -36,17 +37,24 @@ namespace tvm {
 namespace codegen {
 
 /*!
+
  * \brief Wrap the submodules that are to be wrapped in a c-source metadata module for C runtime.
  * \param modules The modules to be wrapped.
  * \param target the target the modules are compiled for.
  * \param runtime the runtime to code generate against
- * \param metadata the metadata needed for code generation.
+ * \param metadata Compiler-generated metadata exported to runtime.
  * \return The wrapped module.
  */
 runtime::Module CreateCSourceCrtMetadataModule(const Array<runtime::Module>& modules, Target target,
                                                relay::Runtime runtime,
                                                relay::backend::ExecutorCodegenMetadata metadata);
 
+/*!
+ * \brief Create C++-runtime targeted metadata module for "c" backend.
+ * \param metadata Compiler-generated metadata.
+ */
+runtime::Module CreateCSourceCppMetadataModule(runtime::metadata::Metadata metadata);
+
 }  // namespace codegen
 }  // namespace tvm
 
diff --git a/src/target/spirv/codegen_spirv.cc b/src/target/spirv/codegen_spirv.cc
index 66952dae269e..1d30b9bfd63a 100644
--- a/src/target/spirv/codegen_spirv.cc
+++ b/src/target/spirv/codegen_spirv.cc
@@ -638,7 +638,7 @@ void CodeGenSPIRV::VisitStmt_(const IfThenElseNode* op) {
 void CodeGenSPIRV::VisitStmt_(const AllocateNode* op) {
   ICHECK(!is_zero(op->condition));
   ICHECK(!op->dtype.is_handle());
-  int32_t constant_size = op->constant_allocation_size();
+  size_t constant_size = op->ConstantAllocationSize();
   ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation in GPU";
 
   spirv::Value buf;
diff --git a/src/te/operation/extern_op.cc b/src/te/operation/extern_op.cc
index b602efcfc28b..84869ccd7775 100644
--- a/src/te/operation/extern_op.cc
+++ b/src/te/operation/extern_op.cc
@@ -90,6 +90,7 @@ Operation ExternOpNode::ReplaceInputs(const Operation& self,
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<ExternOpNode>(*this);
   n->body = ReplaceTensor(this->body, rmap);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->inputs.size(); ++i) {
     Tensor t = n->inputs[i];
     if (rmap.count(t)) {
diff --git a/src/te/operation/hybrid_op.cc b/src/te/operation/hybrid_op.cc
index 5d2412abb3d2..5bb645822996 100644
--- a/src/te/operation/hybrid_op.cc
+++ b/src/te/operation/hybrid_op.cc
@@ -49,13 +49,13 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 
 TVM_REGISTER_NODE_TYPE(HybridOpNode);
 
-int HybridOpNode::num_outputs() const { return static_cast<int>(outputs.size()); }
+int HybridOpNode::num_outputs() const { return static_cast<int>(symbolic_outputs.size()); }
 
 Array<IterVar> HybridOpNode::root_iter_vars() const { return this->axis; }
 
-DataType HybridOpNode::output_dtype(size_t i) const { return outputs[i]->dtype; }
+DataType HybridOpNode::output_dtype(size_t i) const { return symbolic_outputs[i]->dtype; }
 
-Array<PrimExpr> HybridOpNode::output_shape(size_t i) const { return outputs[i]->shape; }
+Array<PrimExpr> HybridOpNode::output_shape(size_t i) const { return symbolic_outputs[i]->shape; }
 
 HybridOp::HybridOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
                    Array<Tensor> inputs, Array<Tensor> outputs, Stmt body) {
@@ -67,7 +67,7 @@ HybridOp::HybridOp(std::string name, std::string tag, Map<String, ObjectRef> att
   n->tag = std::move(tag);
   n->attrs = std::move(attrs);
   n->inputs = std::move(inputs);
-  n->outputs = std::move(outputs);
+  n->symbolic_outputs = std::move(outputs);
   n->axis = te::GatherLoopVars(body);
   n->body = std::move(body);
   data_ = std::move(n);
@@ -104,6 +104,7 @@ Operation HybridOpNode::ReplaceInputs(const Operation& self,
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<HybridOpNode>(*this);
   n->body = te::ReplaceTensor(this->body, rmap);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->inputs.size(); ++i) {
     Tensor t = n->inputs[i];
     if (rmap.count(t)) {
@@ -166,7 +167,7 @@ Stmt HybridOpNode::BuildProvide(const Stage& stage,
   Stmt ret = AttrStmt(make_zero(DataType::Int(32)), tir::attr::extern_scope, 0, this->body);
   std::unordered_map<Tensor, Tensor> rmap;
   for (int i = 0; i < this->num_outputs(); ++i) {
-    rmap[outputs[i]] = stage->op.output(i);
+    rmap[symbolic_outputs[i]] = stage->op.output(i);
   }
   auto n = make_object<HybridOpNode>(*this);
   /* This is a story little bit complicated.
diff --git a/src/te/operation/scan_op.cc b/src/te/operation/scan_op.cc
index 39689bd9654a..a29cc6601795 100644
--- a/src/te/operation/scan_op.cc
+++ b/src/te/operation/scan_op.cc
@@ -40,6 +40,7 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
 TVM_REGISTER_NODE_TYPE(ScanOpNode);
 
 int ScanOpNode::num_outputs() const { return static_cast<int>(update.size()); }
+
 Array<IterVar> ScanOpNode::root_iter_vars() const {
   Array<IterVar> ret{scan_axis};
   for (IterVar iv : spatial_axis_) {
@@ -143,6 +144,7 @@ Operation ScanOpNode::ReplaceInputs(const Operation& self,
                                     const std::unordered_map<Tensor, Tensor>& rmap) const {
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<ScanOpNode>(*this);
+  n->outputs = Array<Tensor>();
   for (size_t i = 0; i < n->init.size(); ++i) {
     if (rmap.count(n->init[i])) {
       n->init.Set(i, rmap.at(n->init[i]));
diff --git a/src/te/operation/tensor_compute_op.cc b/src/te/operation/tensor_compute_op.cc
index 262e5a2b97f4..432df75f03b9 100644
--- a/src/te/operation/tensor_compute_op.cc
+++ b/src/te/operation/tensor_compute_op.cc
@@ -85,6 +85,7 @@ Operation TensorComputeOpNode::ReplaceInputs(const Operation& self,
                                              const std::unordered_map<Tensor, Tensor>& rmap) const {
   ICHECK_EQ(self.operator->(), this);
   auto n = make_object<TensorComputeOpNode>(*this);
+  n->outputs = Array<Tensor>();
   auto intrin = make_object<TensorIntrinNode>(*(this->intrin.operator->()));
   intrin->body = ReplaceTensor(this->intrin->body, rmap);
   if (intrin->reduce_init.defined()) {
diff --git a/src/te/schedule/schedule_dataflow_rewrite.cc b/src/te/schedule/schedule_dataflow_rewrite.cc
index fae826b926e3..e09cdfe146e1 100644
--- a/src/te/schedule/schedule_dataflow_rewrite.cc
+++ b/src/te/schedule/schedule_dataflow_rewrite.cc
@@ -616,7 +616,7 @@ void InjectInline(ScheduleNode* sch, bool feature_extraction_mode) {
       const HybridOpNode* hybrid = sch->stages[i]->op.as<HybridOpNode>();
       ICHECK(hybrid);
       Operation op = HybridOp(hybrid->name, hybrid->tag, hybrid->attrs, hybrid->inputs,
-                              hybrid->outputs, new_hybrid_body[i]);
+                              hybrid->symbolic_outputs, new_hybrid_body[i]);
       op = op->ReplaceInputs(op, repl);
       for (int idx = 0; idx < s->op->num_outputs(); ++idx) {
         repl[s->op.output(idx)] = op.output(idx);
diff --git a/src/te/schedule/schedule_ops.cc b/src/te/schedule/schedule_ops.cc
index 1568df4670af..99e02ccaf943 100644
--- a/src/te/schedule/schedule_ops.cc
+++ b/src/te/schedule/schedule_ops.cc
@@ -298,9 +298,15 @@ class SchedulePostProc : public StmtExprMutator {
  private:
   void AddReplace(Tensor src, Tensor dst, Tensor repl_realize = Tensor(),
                   Operation repl_op = Operation()) {
-    replace_buffer_[src] = dst;
-    replace_realize_[src] = repl_realize;
-    replace_op_[src->op.get()] = repl_op;
+    if (!src.same_as(dst)) {
+      replace_buffer_[src] = dst;
+    }
+    if (!src.same_as(repl_realize)) {
+      replace_realize_[src] = repl_realize;
+    }
+    if (!src->op.same_as(repl_op)) {
+      replace_op_[src->op.get()] = repl_op;
+    }
   }
   // The thread extent scope.
   std::unordered_map<const Object*, PrimExpr> thread_extent_scope_;
diff --git a/src/te/tensor.cc b/src/te/tensor.cc
index 1d75761216f1..1f43714ea107 100644
--- a/src/te/tensor.cc
+++ b/src/te/tensor.cc
@@ -78,13 +78,22 @@ String TensorNode::GetNameHint() const {
   return op->num_outputs() == 1 ? op->name : (op->name + ".v" + std::to_string(value_index));
 }
 
-Tensor Operation::output(size_t i) const {
-  auto node = make_object<TensorNode>();
-  node->op = *this;
-  node->value_index = i;
-  node->dtype = (*this)->output_dtype(i);
-  node->shape = (*this)->output_shape(i);
-  return Tensor(node);
+Tensor Operation::output(size_t n) const {
+  // cache the output tensors if empty
+  if ((*this)->outputs.empty()) {
+    auto* ptr = static_cast<OperationNode*>(get_mutable());
+    size_t num = static_cast<size_t>((*this)->num_outputs());
+    for (size_t i = 0; i < num; ++i) {
+      auto node = make_object<TensorNode>();
+      node->op = *this;
+      node->value_index = i;
+      node->dtype = (*this)->output_dtype(i);
+      node->shape = (*this)->output_shape(i);
+      ptr->outputs.push_back(Tensor(node));
+    }
+  }
+  ICHECK_LT(n, (*this)->outputs.size());
+  return (*this)->outputs[n];
 }
 
 Tensor::Tensor(Array<PrimExpr> shape, DataType dtype, Operation op, int value_index) {
diff --git a/src/tir/analysis/verify_gpu_code.cc b/src/tir/analysis/verify_gpu_code.cc
index dc1ed1c193e8..c1579c21f249 100644
--- a/src/tir/analysis/verify_gpu_code.cc
+++ b/src/tir/analysis/verify_gpu_code.cc
@@ -63,10 +63,10 @@ class GPUCodeVerifier : public StmtExprVisitor {
     auto scope = GetPtrStorageScope(op->buffer_var);
     // visit an allocation of a buffer in shared memory, record its size
     if (scope == "local") {
-      size_t size = static_cast<size_t>(op->constant_allocation_size());
+      size_t size = static_cast<size_t>(op->ConstantAllocationSize());
       local_memory_per_block_ += size * op->dtype.bytes() * op->dtype.lanes();
     } else if (scope == "shared") {
-      size_t size = static_cast<size_t>(op->constant_allocation_size());
+      size_t size = static_cast<size_t>(op->ConstantAllocationSize());
       shared_memory_per_block_ += size * op->dtype.bytes() * op->dtype.lanes();
     }
     if (op->dtype.lanes() > 1) {
diff --git a/src/tir/ir/data_layout.cc b/src/tir/ir/data_layout.cc
index 65be0900f9e3..070cd7077d18 100644
--- a/src/tir/ir/data_layout.cc
+++ b/src/tir/ir/data_layout.cc
@@ -131,7 +131,6 @@ Layout::Layout(const std::string& name) {  // NOLINT(*)
     ICHECK_EQ(axis_str.size(), 1);
     char axis = axis_str[0];
     ICHECK((axis >= 'a' && axis <= 'z') || (axis >= 'A' && axis <= 'Z'));
-    ICHECK(!exist_axis[axis]) << "Invalid layout " << name << ": duplicate axis " << axis;
     exist_axis[axis] = true;
   }
   for (const IterVar& v : node->axes) {
@@ -182,15 +181,20 @@ Layout Layout::Split(const LayoutAxis& axis, size_t target_pos, int32_t factor)
 int32_t Layout::FactorOf(const LayoutAxis& axis) const {
   if (!defined()) return -1;
   const LayoutAxis& sub = axis.ToSubordinate();
-  if (!this->defined()) return -1;
+
+  int32_t factor = 1;
+  bool has_sub = false;
   for (const IterVar& itvar : operator->()->axes) {
     if (sub == LayoutAxis::Get(itvar)) {
-      const auto* factor = itvar->dom->extent.as<IntImmNode>();
-      ICHECK(factor);
-      return factor->value;
+      has_sub = true;
+      int32_t val = itvar->dom->extent.as<IntImmNode>()->value;
+      ICHECK(val);
+      factor *= val;
     }
   }
-  return -1;
+  factor = has_sub ? factor : -1;
+
+  return factor;
 }
 
 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
@@ -199,8 +203,8 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << "Layout(" << l->name << ")";
     });
 
-inline bool GetStoreRule(Array<PrimExpr>* rule, const Layout& src_layout,
-                         const Layout& dst_layout) {
+inline bool GetStoreRule(Array<PrimExpr>* index_rule, Array<PrimExpr>* shape_rule,
+                         const Layout& src_layout, const Layout& dst_layout) {
   if (!src_layout.defined() || src_layout.name().empty() || !dst_layout.defined() ||
       dst_layout.name().empty()) {
     return false;
@@ -208,7 +212,7 @@ inline bool GetStoreRule(Array<PrimExpr>* rule, const Layout& src_layout,
   for (size_t i = 0; i < dst_layout.ndim(); ++i) {
     const auto& store_axis = dst_layout[i];
     const IterVar& store_axis_impl = dst_layout->axes[i];
-    PrimExpr store(0);
+    PrimExpr index_store(0);
 
     for (size_t j = 0; j < src_layout.ndim(); ++j) {
       const auto& orig_axis = src_layout[j];
@@ -220,28 +224,63 @@ inline bool GetStoreRule(Array<PrimExpr>* rule, const Layout& src_layout,
           if (factor > 0) {
             orig_var = orig_var * factor;
           }
-          store = store + orig_var;
+          index_store = index_store + orig_var;
         } else {
-          store = store + orig_axis_impl->var;
+          PrimExpr factor(1);
+          for (size_t k = j + 1; k < src_layout.ndim(); ++k) {
+            if (LayoutAxis::Get(orig_axis_impl) == LayoutAxis::Get(src_layout->axes[k])) {
+              factor = factor * src_layout->axes[k]->dom->extent;
+            }
+          }
+          index_store = index_store + orig_axis_impl->var * factor;
         }
       }
     }
-    if (tir::is_zero(store)) {
+    if (tir::is_zero(index_store)) {
       // Not convertible
       return false;
     }
 
+    PrimExpr shape_store = index_store;
     if (store_axis.IsPrimal()) {
       const int32_t factor = dst_layout.FactorOf(store_axis);
       if (factor > 0) {
-        store = indexdiv(store, PrimExpr(factor));
+        shape_store = shapediv(index_store, PrimExpr(factor));
+        index_store = indexdiv(index_store, PrimExpr(factor));
       }
     } else {
-      store = indexmod(store, store_axis_impl->dom->extent);
+      PrimExpr stride(1);
+      PrimExpr factor(1);
+      for (size_t j = i; j < dst_layout.ndim(); ++j) {
+        if (LayoutAxis::Get(store_axis_impl) == LayoutAxis::Get(dst_layout->axes[j])) {
+          stride = stride * dst_layout->axes[j]->dom->extent;
+          if (j > i) {
+            factor = factor * dst_layout->axes[j]->dom->extent;
+          }
+        }
+      }
+      shape_store = indexdiv(indexmod(index_store, stride), factor);
+      index_store = indexdiv(indexmod(index_store, stride), factor);
     }
 
-    rule->push_back(store);
+    index_rule->push_back(index_store);
+    shape_rule->push_back(shape_store);
+  }
+
+  std::stringstream ss;
+  ss << "index rule for " << src_layout.name() << "-->" << dst_layout.name() << ": [ ";
+  for (const auto& r : *index_rule) {
+    ss << r << ", ";
   }
+  ss << "]" << std::endl;
+
+  ss << "shape rule for " << src_layout.name() << "-->" << dst_layout.name() << ": [ ";
+  for (const auto& r : *shape_rule) {
+    ss << r << ", ";
+  }
+  ss << "]" << std::endl;
+  VLOG(1) << std::endl << ss.str();
+
   return true;
 }
 
@@ -265,7 +304,7 @@ Array<PrimExpr> BijectiveLayout::ForwardIndex(const Array<PrimExpr>& src_index)
   const BijectiveLayoutNode* self = operator->();
   ICHECK_EQ(src_index.size(), self->src_layout->axes.size())
       << "Input mismatch with layout " << self->src_layout;
-  return TransformIndex(src_index, self->src_layout->axes, self->forward_rule);
+  return TransformIndex(src_index, self->src_layout->axes, self->index_forward_rule);
 }
 
 Array<PrimExpr> BijectiveLayout::BackwardIndex(const Array<PrimExpr>& dst_index) const {
@@ -273,7 +312,7 @@ Array<PrimExpr> BijectiveLayout::BackwardIndex(const Array<PrimExpr>& dst_index)
   const BijectiveLayoutNode* self = operator->();
   ICHECK_EQ(dst_index.size(), self->dst_layout->axes.size())
       << "Output mismatch with layout " << self->dst_layout;
-  return TransformIndex(dst_index, self->dst_layout->axes, self->backward_rule);
+  return TransformIndex(dst_index, self->dst_layout->axes, self->index_backward_rule);
 }
 
 inline Array<PrimExpr> TransformShape(const Array<PrimExpr>& src_shape,
@@ -331,19 +370,41 @@ inline Array<PrimExpr> TransformShape(const Array<PrimExpr>& src_shape,
       }
     }
   }
+
+  std::stringstream ss;
+  ss << "shape rule for " << Layout(src_axis).name() << "-->" << Layout(target_axis).name()
+     << ": [ ";
+  for (const auto& r : transform_rule) {
+    ss << r << ", ";
+  }
+  ss << "]" << std::endl;
+
+  ss << "shape transform: [ ";
+  for (const auto& s : src_shape) {
+    ss << s << ", ";
+  }
+  ss << "] --> [ ";
+  for (const auto& r : result) {
+    ss << r << ", ";
+  }
+  ss << "]" << std::endl;
+  VLOG(1) << std::endl << ss.str();
+
   return result;
 }
 
 Array<PrimExpr> BijectiveLayout::ForwardShape(const Array<PrimExpr>& shape) const {
   ICHECK(defined()) << "Cannot operate on an undefined bijective layout.";
   const BijectiveLayoutNode* self = operator->();
-  return TransformShape(shape, self->src_layout->axes, self->dst_layout->axes, self->forward_rule);
+  return TransformShape(shape, self->src_layout->axes, self->dst_layout->axes,
+                        self->shape_forward_rule);
 }
 
 Array<PrimExpr> BijectiveLayout::BackwardShape(const Array<PrimExpr>& shape) const {
   ICHECK(defined()) << "Cannot operate on an undefined bijective layout.";
   const BijectiveLayoutNode* self = operator->();
-  return TransformShape(shape, self->dst_layout->axes, self->src_layout->axes, self->backward_rule);
+  return TransformShape(shape, self->dst_layout->axes, self->src_layout->axes,
+                        self->shape_backward_rule);
 }
 
 BijectiveLayout::BijectiveLayout(Layout src_layout, Layout dst_layout) {
@@ -354,8 +415,9 @@ BijectiveLayout::BijectiveLayout(Layout src_layout, Layout dst_layout) {
 
   // To be consistent with previous behavior, a nullptr layout is created
   // when argument is invalid.
-  if (GetStoreRule(&n->forward_rule, n->src_layout, n->dst_layout)) {
-    ICHECK(GetStoreRule(&n->backward_rule, n->dst_layout, n->src_layout));
+  if (GetStoreRule(&n->index_forward_rule, &n->shape_forward_rule, n->src_layout, n->dst_layout)) {
+    ICHECK(GetStoreRule(&n->index_backward_rule, &n->shape_backward_rule, n->dst_layout,
+                        n->src_layout));
     data_ = std::move(n);
   }
 }
diff --git a/src/tir/ir/function.cc b/src/tir/ir/function.cc
index 1c34e34468b5..f58dd8aa820c 100644
--- a/src/tir/ir/function.cc
+++ b/src/tir/ir/function.cc
@@ -27,14 +27,6 @@
 
 namespace tvm {
 namespace tir {
-
-LinkedParam::LinkedParam(int64_t id, ::tvm::runtime::NDArray param) {
-  auto n = make_object<LinkedParamNode>();
-  n->id = id;
-  n->param = param;
-  data_ = std::move(n);
-}
-
 // Get the function type of a PrimFunc
 PrimFunc::PrimFunc(Array<tir::Var> params, Stmt body, Type ret_type,
                    Map<tir::Var, Buffer> buffer_map, DictAttrs attrs, Span span) {
diff --git a/src/tir/ir/stmt.cc b/src/tir/ir/stmt.cc
index 078561c447ad..1269607fd334 100644
--- a/src/tir/ir/stmt.cc
+++ b/src/tir/ir/stmt.cc
@@ -366,19 +366,19 @@ Allocate::Allocate(Var buffer_var, DataType dtype, Array<PrimExpr> extents, Prim
   data_ = std::move(node);
 }
 
-int32_t AllocateNode::constant_allocation_size(const Array<PrimExpr>& extents) {
+int64_t AllocateNode::ConstantAllocationSize(const Array<PrimExpr>& extents) {
   int64_t result = 1;
   for (size_t i = 0; i < extents.size(); ++i) {
     if (const IntImmNode* int_size = extents[i].as<IntImmNode>()) {
       result *= int_size->value;
-      if (result > std::numeric_limits<int32_t>::max()) {
+      if (result > std::numeric_limits<int64_t>::max()) {
         return 0;
       }
     } else {
       return 0;
     }
   }
-  return static_cast<int32_t>(result);
+  return static_cast<int64_t>(result);
 }
 
 TVM_REGISTER_GLOBAL("tir.Allocate")
@@ -409,6 +409,79 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->Print(op->body);
     });
 
+// Const
+// The constructor to create a IRNode with constant data
+// depending on the type of ObjectRef, it will either
+// create AllocateConstNode with irmod_storage_idx or data
+AllocateConst::AllocateConst(Var buffer_var, DataType dtype, Array<PrimExpr> extents,
+                             ObjectRef data_or_idx, Stmt body, Span span) {
+  ICHECK(IsPointerType(buffer_var->type_annotation, dtype))
+      << "The allocated data type (" << dtype
+      << ") does not match the type annotation of the buffer " << buffer_var << " ("
+      << buffer_var->type_annotation
+      << "). The data type should be an element of the pointer type.";
+
+  for (size_t i = 0; i < extents.size(); ++i) {
+    ICHECK(extents[i].defined());
+    ICHECK(extents[i].dtype().is_scalar());
+  }
+  ICHECK(body.defined());
+  ICHECK(data_or_idx.defined());
+
+  ObjectPtr<AllocateConstNode> node = make_object<AllocateConstNode>();
+  node->buffer_var = std::move(buffer_var);
+  node->dtype = dtype;
+  node->extents = std::move(extents);
+  node->body = std::move(body);
+  node->span = std::move(span);
+  if (data_or_idx->IsInstance<runtime::NDArray::ContainerType>()) {
+    node->data = Optional<tvm::runtime::NDArray>(Downcast<runtime::NDArray>(data_or_idx));
+    node->irmod_storage_idx = Optional<Integer>();
+  } else if (data_or_idx->IsInstance<IntImmNode>()) {
+    node->data = Optional<tvm::runtime::NDArray>();
+    node->irmod_storage_idx = Optional<Integer>(Downcast<Integer>(data_or_idx));
+  } else {
+    LOG(FATAL) << "Data type not supported: " << data_or_idx->GetTypeKey();
+  }
+  data_ = std::move(node);
+}
+
+int64_t AllocateConstNode::ConstantAllocationSize(const Array<PrimExpr>& extents) {
+  int64_t result = 1;
+  for (size_t i = 0; i < extents.size(); ++i) {
+    if (const IntImmNode* int_size = extents[i].as<IntImmNode>()) {
+      result *= int_size->value;
+      if (result > std::numeric_limits<int64_t>::max()) {
+        return 0;
+      }
+    } else {
+      return 0;
+    }
+  }
+  return static_cast<int64_t>(result);
+}
+TVM_REGISTER_GLOBAL("tir.AllocateConst")
+    .set_body_typed([](Var buffer_var, DataType dtype, Array<PrimExpr> extents,
+                       ObjectRef data_or_idx, Stmt body, Span span) {
+      return AllocateConst(buffer_var, dtype, extents, data_or_idx, body, span);
+    });
+
+TVM_REGISTER_NODE_TYPE(AllocateConstNode);
+
+TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
+    .set_dispatch<AllocateConstNode>([](const ObjectRef& node, ReprPrinter* p) {
+      auto* op = static_cast<const AllocateConstNode*>(node.get());
+      p->PrintIndent();
+      p->stream << "constant " << op->buffer_var << "[" << op->dtype;
+      for (size_t i = 0; i < op->extents.size(); ++i) {
+        p->stream << " * ";
+        p->Print(op->extents[i]);
+      }
+      p->stream << "]";
+      p->stream << "\n";
+      p->Print(op->body);
+    });
+
 // ProducerRealize
 ProducerRealize::ProducerRealize(DataProducer producer, Region bounds, PrimExpr condition,
                                  Stmt body, String storage_scope, Span span) {
diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc
index d60ec72a7589..949e8a1312aa 100644
--- a/src/tir/ir/stmt_functor.cc
+++ b/src/tir/ir/stmt_functor.cc
@@ -58,6 +58,11 @@ void StmtVisitor::VisitStmt_(const AllocateNode* op) {
   this->VisitExpr(op->condition);
 }
 
+void StmtVisitor::VisitStmt_(const AllocateConstNode* op) {
+  VisitArray(op->extents, [this](const PrimExpr& e) { this->VisitExpr(e); });
+  this->VisitStmt(op->body);
+}
+
 void StmtVisitor::VisitStmt_(const StoreNode* op) {
   this->VisitExpr(op->value);
   this->VisitExpr(op->index);
@@ -319,6 +324,20 @@ Stmt StmtMutator::VisitStmt_(const AllocateNode* op) {
   }
 }
 
+Stmt StmtMutator::VisitStmt_(const AllocateConstNode* op) {
+  Array<PrimExpr> extents = Internal::Mutate(this, op->extents);
+  Stmt body = this->VisitStmt(op->body);
+
+  if (extents.same_as(op->extents) && body.same_as(op->body)) {
+    return GetRef<Stmt>(op);
+  } else {
+    auto n = CopyOnWrite(op);
+    n->extents = std::move(extents);
+    n->body = std::move(body);
+    return Stmt(n);
+  }
+}
+
 Stmt StmtMutator::VisitStmt_(const IfThenElseNode* op) {
   PrimExpr condition = this->VisitExpr(op->condition);
   Stmt then_case = this->VisitStmt(op->then_case);
diff --git a/src/tir/op/builtin.cc b/src/tir/op/builtin.cc
index 5fccae040dc9..0e767ead4e6b 100644
--- a/src/tir/op/builtin.cc
+++ b/src/tir/op/builtin.cc
@@ -260,6 +260,9 @@ TIR_DEFINE_BUILTIN_FUNC(texture2d_load)
     .set_attr<TVectorizable>("TVectorizable", true)
     .set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kOpaque));
 
+TIR_DEFINE_BUILTIN_FUNC(mem_copy).set_attr<TCallEffectKind>("TCallEffectKind",
+                                                            Integer(CallEffectKind::kOpaque));
+
 }  // namespace builtin
 }  // namespace tir
 }  // namespace tvm
diff --git a/src/tir/op/op.cc b/src/tir/op/op.cc
index d08bef2ab91a..696d82be721f 100644
--- a/src/tir/op/op.cc
+++ b/src/tir/op/op.cc
@@ -128,6 +128,16 @@ void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Span span) {  // NOLINT(*)
              !rtype.is_float()) {
     // Cast int->float when the other operand is a float
     rhs = cast(ltype, rhs);
+  } else if (!ltype.is_bfloat16() &&
+             (rtype.is_bfloat16() ||
+              datatype::Registry::Global()->GetTypeRegistered(rtype.code()))) {
+    // Cast int->bfloat16 when the other operand is a bfloat16
+    lhs = cast(rtype, lhs);
+  } else if ((ltype.is_bfloat16() ||
+              datatype::Registry::Global()->GetTypeRegistered(ltype.code())) &&
+             !rtype.is_bfloat16()) {
+    // Cast int->bfloat16 when the other operand is a bfloat16
+    rhs = cast(ltype, rhs);
   } else if ((ltype.is_int() && rtype.is_int()) || (ltype.is_uint() && rtype.is_uint())) {
     // Promote int to higher bits e.g. int8 + int16 --> int16 + int16
     if (ltype.bits() < rtype.bits()) {
@@ -186,6 +196,8 @@ PrimExpr max_value(const DataType& dtype, Span span) {
     } else if (dtype.bits() == 16) {
       return FloatImm(dtype, 65504.0, span);
     }
+  } else if (dtype.is_bfloat16()) {
+    return FloatImm(dtype, std::numeric_limits<float>::max(), span);
   }
   LOG(FATAL) << "Cannot decide max_value for type" << dtype;
   return PrimExpr();
@@ -219,6 +231,8 @@ PrimExpr min_value(const DataType& dtype, Span span) {
     } else if (dtype.bits() == 16) {
       return FloatImm(dtype, -65504.0, span);
     }
+  } else if (dtype.is_bfloat16()) {
+    return FloatImm(dtype, std::numeric_limits<float>::lowest(), span);
   }
   LOG(FATAL) << "Cannot decide min_value for type" << dtype;
   return PrimExpr();
@@ -369,6 +383,8 @@ PrimExpr operator%(PrimExpr a, PrimExpr b) { return truncmod(a, b); }
 // TODO(tqchen): switch to floordiv
 PrimExpr indexdiv(PrimExpr a, PrimExpr b, Span span) { return floordiv(a, b, span); }
 
+PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span) { return ceildiv(a, b, span); }
+
 PrimExpr indexmod(PrimExpr a, PrimExpr b, Span span) { return floormod(a, b, span); }
 
 PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span) {
@@ -380,6 +396,15 @@ PrimExpr floordiv(PrimExpr a, PrimExpr b, Span span) {
   return tir::FloorDiv(a, b, span);
 }
 
+PrimExpr ceildiv(PrimExpr a, PrimExpr b, Span span) {
+  ICHECK(a.dtype().is_int() || a.dtype().is_uint()) << a;
+  ICHECK(b.dtype().is_int() || b.dtype().is_uint()) << b;
+  BinaryOpMatchTypes(a, b, span);
+  PrimExpr ret = arith::TryConstFold<tir::FloorDiv>(a + b - 1, b);
+  if (ret.defined()) return ret;
+  return tir::FloorDiv(a + b - 1, b, span);
+}
+
 PrimExpr floormod(PrimExpr a, PrimExpr b, Span span) {
   ICHECK(a.dtype().is_int() || a.dtype().is_uint()) << a;
   ICHECK(b.dtype().is_int() || b.dtype().is_uint()) << b;
diff --git a/src/tir/schedule/instruction_traits.h b/src/tir/schedule/instruction_traits.h
index 91a79456e579..14d05a4a340c 100644
--- a/src/tir/schedule/instruction_traits.h
+++ b/src/tir/schedule/instruction_traits.h
@@ -197,6 +197,8 @@ class PythonAPICall {
   inline void Input(String arg_name, int arg);
   /*! \brief Add an integer input */
   inline void Input(String arg_name, int64_t arg);
+  /*! \brief Add a bool input */
+  inline void Input(String arg_name, bool arg);
   /*! \brief Add a double input */
   inline void Input(String arg_name, double arg);
   /*! \brief Add an input random variable */
@@ -462,6 +464,17 @@ void PythonAPICall::Input(String arg_name, int64_t arg) {
   args_.push_back(std::to_string(arg));
 }
 
+void PythonAPICall::Input(String arg_name, bool arg) {
+  static const char* true_str = "True";
+  static const char* false_str = "False";
+  arg_names_.emplace_back(std::move(arg_name));
+  if (arg) {
+    args_.push_back(true_str);
+  } else {
+    args_.push_back(false_str);
+  }
+}
+
 void PythonAPICall::Input(String arg_name, double arg) {
   arg_names_.emplace_back(std::move(arg_name));
   std::ostringstream os;
diff --git a/src/tir/schedule/primitive/compute_at.cc b/src/tir/schedule/primitive/compute_at.cc
index 88bf7663e85d..b811afb23614 100644
--- a/src/tir/schedule/primitive/compute_at.cc
+++ b/src/tir/schedule/primitive/compute_at.cc
@@ -202,6 +202,10 @@ struct BlockVarDomainInfo {
     if (analyzer->CanProveEqual(dom.min(), intersect.min()) &&
         analyzer->CanProveEqual(dom.max(), intersect.max())) {
       bound = arith::IntSet::Nothing();
+    } else if (analyzer->CanProveEqual(bound.min(), intersect.min()) &&
+               analyzer->CanProveEqual(bound.max(), intersect.max())) {
+      dom = bound;
+      bound = arith::IntSet::Nothing();
     }
   }
 };
diff --git a/src/tir/transforms/arg_binder.cc b/src/tir/transforms/arg_binder.cc
index d3ab32cbd7f9..1e566a980463 100644
--- a/src/tir/transforms/arg_binder.cc
+++ b/src/tir/transforms/arg_binder.cc
@@ -169,7 +169,8 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
                        IntImm(DataType::UInt(8), dtype.bits()) &&
                    TVMArrayGet(DataType::UInt(16), handle, builtin::kArrTypeLanes) ==
                        IntImm(DataType::UInt(16), dtype.lanes()));
-  if (!(dtype == DataType::Int(4) || dtype == DataType::UInt(4) || dtype == DataType::Int(1))) {
+  if (!(dtype == DataType::Int(1) || dtype == DataType::Int(4) || dtype == DataType::UInt(4) ||
+        dtype == DataType::UInt(16))) {
     auto type_msg = tvm::tir::StringImm(type_err_msg.str());
     asserts_.emplace_back(AssertStmt(a_ndim == v_ndim, msg, nop));
     asserts_.emplace_back(AssertStmt(cond, type_msg, nop));
diff --git a/src/tir/transforms/bf16_legalize.cc b/src/tir/transforms/bf16_legalize.cc
index 76845cbebd2a..79c406818185 100644
--- a/src/tir/transforms/bf16_legalize.cc
+++ b/src/tir/transforms/bf16_legalize.cc
@@ -45,26 +45,6 @@ class BF16PromoteRewriter : public StmtExprMutator {
 
   Stmt operator()(Stmt s) { return VisitStmt(s); }
 
-  std::tuple<PrimExpr, PrimExpr> DoCast(PrimExpr orig_a, PrimExpr orig_b, bool* is_bfloat16) {
-    auto a = this->VisitExpr(orig_a);
-    auto b = this->VisitExpr(orig_b);
-    *is_bfloat16 = false;
-    if (a->dtype.is_bfloat16()) {
-      ICHECK(b->dtype.is_bfloat16());
-      *is_bfloat16 = true;
-    } else if (b->dtype.is_bfloat16()) {
-      ICHECK(a->dtype.is_bfloat16());
-      *is_bfloat16 = true;
-    }
-
-    if (*is_bfloat16) {
-      DataType fp32ty(kDLFloat, 32, 1);
-      a = Cast(fp32ty, a);
-      b = Cast(fp32ty, b);
-    }
-    return std::make_tuple(a, b);
-  }
-
   PrimExpr VisitExpr_(const AddNode* op) final;
   PrimExpr VisitExpr_(const SubNode* op) final;
   PrimExpr VisitExpr_(const MulNode* op) final;
@@ -77,45 +57,36 @@ class BF16PromoteRewriter : public StmtExprMutator {
   PrimExpr VisitExpr_(const GENode* op) final;
 };
 
-#define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC)  \
-  PrimExpr BF16PromoteRewriter::VisitExpr_(const OP* op) { \
-    PrimExpr a, b;                                         \
-    bool is_bfloat16;                                      \
-    std::tie(a, b) = DoCast(op->a, op->b, &is_bfloat16);   \
-    if (a.same_as(op->a) && b.same_as(op->b)) {            \
-      return GetRef<PrimExpr>(op);                         \
-    } else {                                               \
-      auto ret = FUNC(a, b);                               \
-      if (!is_bfloat16)                                    \
-        return ret;                                        \
-      else                                                 \
-        return Cast(DataType(kDLBfloat, 16, 1), ret);      \
-    }                                                      \
-  }
-
-#define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH_NO_CAST(OP, FUNC) \
-  PrimExpr BF16PromoteRewriter::VisitExpr_(const OP* op) {        \
-    PrimExpr a, b;                                                \
-    bool is_bfloat16;                                             \
-    std::tie(a, b) = DoCast(op->a, op->b, &is_bfloat16);          \
-    if (a.same_as(op->a) && b.same_as(op->b)) {                   \
-      return GetRef<PrimExpr>(op);                                \
-    } else {                                                      \
-      auto ret = FUNC(a, b);                                      \
-      return ret;                                                 \
-    }                                                             \
+#define DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(OP, FUNC, NEEDCAST)                \
+  PrimExpr BF16PromoteRewriter::VisitExpr_(const OP* op) {                         \
+    PrimExpr origin_a = this->VisitExpr(op->a);                                    \
+    PrimExpr origin_b = this->VisitExpr(op->b);                                    \
+    bool a_is_bfloat16 = origin_a->dtype.is_bfloat16();                            \
+    bool b_is_bfloat16 = origin_b->dtype.is_bfloat16();                            \
+    bool both_bfloat16 = a_is_bfloat16 && b_is_bfloat16;                           \
+    bool none_bfloat16 = !(a_is_bfloat16 || b_is_bfloat16);                        \
+    if (none_bfloat16) {                                                           \
+      return GetRef<PrimExpr>(op);                                                 \
+    }                                                                              \
+    DataType float32_dtype(kDLFloat, 32, 1);                                       \
+    PrimExpr float32_a = a_is_bfloat16 ? Cast(float32_dtype, origin_a) : origin_a; \
+    PrimExpr float32_b = b_is_bfloat16 ? Cast(float32_dtype, origin_b) : origin_b; \
+    PrimExpr result = FUNC(float32_a, float32_b);                                  \
+    DataType bfloat16_dtype(kDLBfloat, 16, 1);                                     \
+    bool do_cast = both_bfloat16 && NEEDCAST;                                      \
+    return do_cast ? Cast(bfloat16_dtype, result) : result;                        \
   }
 
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(AddNode, operator+)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(SubNode, operator-)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MulNode, operator*)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(DivNode, div)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MinNode, min)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MaxNode, max)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH_NO_CAST(LTNode, operator<)   // NOLINT(*)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH_NO_CAST(LENode, operator<=)  // NOLINT(*)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH_NO_CAST(GTNode, operator>)   // NOLINT(*)
-DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH_NO_CAST(GENode, operator>=)  // NOLINT(*)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(AddNode, operator+, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(SubNode, operator-, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MulNode, operator*, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(DivNode, div, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MinNode, min, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(MaxNode, max, true)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LTNode, operator<, false)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(LENode, operator<=, false)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GTNode, operator>, false)
+DEFINE_BIOP_EXPR_MUTATE_WITH_TYPE_MATCH(GENode, operator>=, false)
 
 /*
  * Eliminate verbose casting between fp32 and bf16
@@ -179,25 +150,23 @@ class BF16LowerRewriter : public StmtExprMutator {
   using StmtExprMutator::operator();
 
   PrimExpr VisitExpr_(const CastNode* op) final {
-    auto op_val = StmtExprMutator::VisitExpr(op->value);
-    if (op->value->dtype.is_bfloat16()) {
-      // if is cast_from_bf16, check if is to fp32
-      ICHECK(op->dtype.is_float() && op->dtype.bits() == 32);
-      auto uint32_dtype = DataType(kDLUInt, 32, op_val->dtype.lanes());
-      auto uint32_v = Cast(uint32_dtype, op_val);
-      // to be endian invariant.
-      return Call(op->dtype, builtin::reinterpret(), {uint32_v << 16});
-    } else if (op->dtype.is_bfloat16()) {
-      // if is cast_to_bf16, check if op->value is fp32
-      ICHECK(op->value->dtype.is_float() && op->value->dtype.bits() == 32);
-      auto uint32_dtype = DataType(kDLUInt, 32, op_val->dtype.lanes());
-      auto uint32_v = Call(uint32_dtype, builtin::reinterpret(), {op_val});
-      auto uint16_dtype = DataType(kDLUInt, 16, op_val->dtype.lanes());
+    PrimExpr op_val = StmtExprMutator::VisitExpr(op->value);
+    DataType uint32_dtype(kDLUInt, 32, op_val->dtype.lanes());
+    DataType float32_dtype(kDLFloat, 32, op_val->dtype.lanes());
+    if (op->value->dtype.is_bfloat16()) {  // cast from bf16
+      PrimExpr uint32_v = Cast(uint32_dtype, op_val);
+      PrimExpr float32_v = Call(float32_dtype, builtin::reinterpret(), {uint32_v << 16});
+      bool is_to_float32 = op->dtype.is_float() && op->dtype.bits() == 32;
+      return is_to_float32 ? float32_v : Cast(op->dtype, float32_v);
+    } else if (op->dtype.is_bfloat16()) {  // cast to bf16
+      bool is_from_float32 = op->value->dtype.is_float() && op->value->dtype.bits() == 32;
+      PrimExpr float32_v = is_from_float32 ? op_val : Cast(float32_dtype, op_val);
+      PrimExpr uint32_v = Call(uint32_dtype, builtin::reinterpret(), {float32_v});
+      DataType uint16_dtype(kDLUInt, 16, op_val->dtype.lanes());
       /* the following TIR is equivalent to the C++ code below:
       uint32_t rounding_bias = ((U32 >> 16) & 1) + UINT32_C(0x7FFF);
       return static_cast<uint16_t>((U32 + rounding_bias) >> 16);*/
-      auto rounding_bias = ((uint32_v >> 16) & 1) + make_const(uint16_dtype, 0x7FFF);
-      // to be endian invariant.
+      PrimExpr rounding_bias = ((uint32_v >> 16) & 1) + make_const(uint16_dtype, 0x7FFF);
       return Cast(uint16_dtype, {(uint32_v + rounding_bias) >> 16});
     }
     if (op->value.same_as(op_val)) return GetRef<PrimExpr>(op);
diff --git a/src/tir/transforms/bind_params.cc b/src/tir/transforms/bind_params.cc
new file mode 100644
index 000000000000..944a67a879fd
--- /dev/null
+++ b/src/tir/transforms/bind_params.cc
@@ -0,0 +1,133 @@
+/*
+ * 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 storage_rewrite.cc
+ * \brief Memory access pattern analysis and optimization.
+ *  Re-write data access to enable memory sharing when possible.
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/ir/type.h>
+#include <tvm/relay/expr.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/target/target_info.h>
+#include <tvm/tir/analysis.h>
+#include <tvm/tir/builtin.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/function.h>
+#include <tvm/tir/stmt_functor.h>
+#include <tvm/tir/transform.h>
+
+#include <map>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "../../runtime/thread_storage_scope.h"
+#include "ir_utils.h"
+
+namespace tvm {
+namespace tir {
+
+class ParamsCollector : public StmtExprVisitor {
+ public:
+  explicit ParamsCollector(const Map<tir::Var, runtime::NDArray>& constant_map)
+      : constant_map_(constant_map) {}
+  std::vector<const tir::VarNode*> CollectParams(tir::Stmt body) {
+    this->VisitStmt(body);
+    return constant_list_;
+  }
+
+  void VisitExpr_(const LoadNode* ln) {
+    if (constant_map_.find(ln->buffer_var) != constant_map_.end()) {
+      auto it =
+          std::find(constant_list_.begin(), constant_list_.end(), ln->buffer_var.operator->());
+      if (it == constant_list_.end()) {
+        constant_list_.push_back(ln->buffer_var.operator->());
+      }
+    }
+    StmtExprVisitor::VisitExpr_(ln);
+  }
+
+  void VisitExpr_(const CallNode* cn) {
+    if (cn->op.same_as(builtin::tvm_access_ptr())) {
+      ICHECK_EQ(cn->args.size(), 5U);
+      const Var& var = Downcast<Var>(cn->args[1]);
+      const VarNode* buffer = cn->args[1].as<VarNode>();
+      auto it = constant_map_.find(var);
+      if (it != constant_map_.end()) {
+        auto it = std::find(constant_list_.begin(), constant_list_.end(), buffer);
+        if (it == constant_list_.end()) {
+          constant_list_.push_back(buffer);
+        }
+      }
+    }
+    StmtExprVisitor::VisitExpr_(cn);
+  }
+
+ private:
+  std::vector<const tir::VarNode*> constant_list_;
+  Map<tir::Var, runtime::NDArray> constant_map_;
+  bool first_for_ = true;
+};
+
+namespace transform {
+
+Pass BindParams(const Array<runtime::NDArray>& constants) {
+  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
+    Map<tir::Var, runtime::NDArray> constant_map;
+
+    // Remove constants from the primfunc signature
+    size_t num_constants = constants.size();
+    size_t start = f->params.size() - num_constants;
+    Array<tir::Var> params;
+    for (unsigned i = 0; i < start; i++) {
+      params.push_back(f->params[i]);
+    }
+
+    auto* n = f.CopyOnWrite();
+    for (unsigned i = start; i < f->params.size(); i++) {
+      tir::Var p = n->params[i];
+      tir::Var b = n->buffer_map[p]->data;
+      n->buffer_map.erase(p);
+      constant_map.Set(b, constants[i - start]);
+    }
+    n->params = params;
+    auto constant_list = ParamsCollector(constant_map).CollectParams(n->body);
+
+    // Allocate constants within the primfunc
+    for (auto i : constant_list) {
+      auto var = GetRef<Var>(i);
+      int ndim = constant_map[var]->ndim;
+      Array<PrimExpr> extents;
+
+      for (int i = 0; i < ndim; i++) {
+        int shape = constant_map[var]->shape[i];
+        extents.push_back(make_const(DataType::Int(32), shape));
+      }
+      DataType dtype = DataType(constant_map[var]->dtype);
+      n->body = tir::AllocateConst(var, dtype, extents, constant_map[var], n->body);
+    }
+    return f;
+  };
+  return CreatePrimFuncPass(pass_func, 0, "tir.BindParams", {});
+}
+}  // namespace transform
+
+}  // namespace tir
+}  // namespace tvm
diff --git a/src/tir/transforms/extract_constants.cc b/src/tir/transforms/extract_constants.cc
new file mode 100644
index 000000000000..237f923516da
--- /dev/null
+++ b/src/tir/transforms/extract_constants.cc
@@ -0,0 +1,115 @@
+/*
+ * 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 extract_constants.cc
+ * \brief Collects PrimFunc's constant data into mod's 'tvm::attr::kConstantsArray' attrs array,
+ * sets irmod_storage_idx as index in this array.
+ * For more information, see the RFC:
+ * https://github.com/apache/tvm-rfcs/blob/main/rfcs/0022-tir-non-scalar-constants.md
+ */
+#include <tvm/arith/analyzer.h>
+#include <tvm/ir/transform.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/tir/stmt_functor.h>
+
+#include "ir_utils.h"
+
+namespace tvm {
+namespace tir {
+
+using ConstArrayType = Array<runtime::NDArray>;
+class Applicator : public tir::StmtMutator {
+ protected:
+  // returns index of the a in constant_array_, if not found - appends
+  size_t DeDup(const runtime::NDArray& a) {
+    tvm::SEqualReducer eql;
+    auto it = std::find_if(
+        constant_array_.begin(), constant_array_.end(), [&eql, a](const runtime::NDArray& v) {
+          return NDArrayContainerTrait::SEqualReduce(a.as<runtime::NDArray::Container>(),
+                                                     v.as<runtime::NDArray::Container>(), eql);
+        });
+    if (it != constant_array_.end()) {
+      return it - constant_array_.begin();
+    }
+    constant_array_.push_back(std::move(a));
+    return constant_array_.size() - 1;
+  }
+
+ public:
+  Stmt Apply(tir::Stmt body, const ConstArrayType& constant_array) {
+    constant_array_ = constant_array;
+    return this->VisitStmt(body);
+  }
+
+  Stmt VisitStmt_(const tir::AllocateConstNode* acn) override {
+    // Check whether the data already defined within the module's attrs
+    // and add array index.
+    ICHECK(acn->data) << "data field should be defined";
+    auto node = CopyOnWrite(acn);
+    node->irmod_storage_idx = Optional<Integer>(Integer(DeDup(node->data.value())));
+    return Stmt(node);
+  }
+
+  ConstArrayType constant_array_;
+};
+
+namespace transform {
+
+tvm::transform::Pass ExtractPrimFuncConstants() {
+  auto prim_func_pass = [=](PrimFunc foo, IRModule m, tvm::transform::PassContext ctx) {
+    auto* func = foo.CopyOnWrite();
+    if (!m->attrs.defined()) {
+      m->attrs = DictAttrs(Map<String, ObjectRef>());
+    }
+    auto* attrs = m->attrs.CopyOnWrite();
+    ConstArrayType constant_array_ =
+        (attrs->dict.count(tvm::attr::kConstantsArray))
+            ? Downcast<ConstArrayType>(attrs->dict[tvm::attr::kConstantsArray])
+            : ConstArrayType();
+    Applicator a = Applicator();
+    func->body = a.Apply(func->body, constant_array_);
+    const ConstArrayType constant_list = a.constant_array_;
+    if (constant_list.size()) {
+      attrs->dict.Set(tvm::attr::kConstantsArray, constant_list);
+    }
+    return GetRef<PrimFunc>(func);
+  };
+
+  auto pass_func = [=](IRModule module, tvm::transform::PassContext pc) {
+    auto m = GetRef<IRModule>(module.CopyOnWrite());
+    for (const auto& kv : m->functions) {
+      BaseFunc f = kv.second;
+      if (f->IsInstance<PrimFuncNode>()) {
+        m->Update(kv.first, prim_func_pass(GetRef<PrimFunc>(f.as<PrimFuncNode>()), m, pc));
+      }
+    }
+    return m;
+  };
+
+  return tvm::transform::CreateModulePass(pass_func, 0, "tir.ExtractPrimFuncConstants", {});
+}
+
+TVM_REGISTER_GLOBAL("tir.transform.ExtractPrimFuncConstants")
+    .set_body_typed(ExtractPrimFuncConstants);
+
+}  // namespace transform
+
+}  // namespace tir
+}  // namespace tvm
diff --git a/src/tir/transforms/inject_copy_intrin.cc b/src/tir/transforms/inject_copy_intrin.cc
index f99cbd5b5a05..9e74b8cd1fdb 100644
--- a/src/tir/transforms/inject_copy_intrin.cc
+++ b/src/tir/transforms/inject_copy_intrin.cc
@@ -45,26 +45,35 @@ class CopyIntrinInjector : public StmtMutator {
   Stmt VisitStmt_(const AttrStmtNode* op) final {
     if (op->attr_key == pragma_key_) {
       Stmt ret;
-      ICHECK(MatchCopyPattern(op->body, &ret)) << "Cannot match copy pattern of " << op->body;
+      std::string error_info;
+      ICHECK(MatchCopyPattern(op->body, &ret, &error_info))
+          << "Cannot match copy pattern. The error is " << error_info << " The body is "
+          << op->body;
       return ret;
     }
     return StmtMutator::VisitStmt_(op);
   }
 
  private:
-  bool MatchCopyPattern(Stmt stmt, Stmt* out) {
+  bool MatchCopyPattern(Stmt stmt, Stmt* out, std::string* error_info) {
     using namespace arith;
     Stmt body = stmt;
 
     // strip the loops
     std::vector<const ForNode*> loops;
     while (const ForNode* op = body.as<ForNode>()) {
-      if (!is_zero(op->min)) return false;
+      if (!is_zero(op->min)) {
+        *error_info = "the 'min' value of body 'Fonode' is 0.";
+        return false;
+      }
       loops.push_back(op);
       body = op->body;
     }
     const StoreNode* store = body.as<StoreNode>();
-    if (store == nullptr) return false;
+    if (store == nullptr) {
+      *error_info = "the 'StoreNode' of body is a nullptr.";
+      return false;
+    }
     // Expr sel_cond, sel_true_value, sel_false_value;
     // match select or if
     PVar<PrimExpr> sel_cond, sel_true_value, sel_false_value;
@@ -84,7 +93,10 @@ class CopyIntrinInjector : public StmtMutator {
     if (cast != nullptr) {
       load = cast->value.as<LoadNode>();
     }
-    if (load == nullptr) return false;
+    if (load == nullptr) {
+      *error_info = "the 'LoadNode' of body is a nullptr.";
+      return false;
+    }
     if (load->dtype.lanes() != 1) return false;
     Array<Var> loop_vars;
     for (const ForNode* op : loops) {
@@ -109,7 +121,10 @@ class CopyIntrinInjector : public StmtMutator {
     if (has_cond) {
       Array<PrimExpr> clip_bound = arith::DetectClipBound(sel_cond.Eval(), loop_vars);
       pad_value = sel_false_value.Eval();
-      if (clip_bound.size() == 0) return false;
+      if (clip_bound.size() == 0) {
+        *error_info = "the size of clip bound is 0.";
+        return false;
+      }
       ICHECK_EQ(src_shape.size(), loop_vars.size());
       ICHECK_EQ(clip_bound.size(), loop_vars.size() * 2);
       for (size_t i = 0; i < src_shape.size(); ++i) {
@@ -150,7 +165,10 @@ class CopyIntrinInjector : public StmtMutator {
     Buffer src = Buffer(load->buffer_var, load->dtype, src_shape, src_strides, src_elem_offset,
                         load->buffer_var->name_hint, 0, 0, kDefault);
     *out = flower_copy_fromto_(src, dst, pad_before, pad_after, pad_value);
-    ICHECK(out->defined()) << "flower function did not return correct stmt";
+    if (!out->defined()) {
+      *error_info = "flower function did not return correct stmt";
+      return false;
+    }
     return true;
   }
 
diff --git a/src/tir/transforms/inject_rolling_buffer.cc b/src/tir/transforms/inject_rolling_buffer.cc
index bc4012cf0556..0b70cf6c0818 100644
--- a/src/tir/transforms/inject_rolling_buffer.cc
+++ b/src/tir/transforms/inject_rolling_buffer.cc
@@ -172,7 +172,7 @@ class RollingBufferInjector : public StmtExprMutator {
 
           auto it{std::find_if(
               bound_iter_vars.begin(), bound_iter_vars.end(),
-              [&](Optional<Var> var) { return var && (var.value().get() == loop_var.get()); })};
+              [&](Optional<Var> var) { return var && (var.get() == loop_var.get()); })};
 
           if (it != bound_iter_vars.end()) {
             auto i{std::distance(bound_iter_vars.begin(), it)};
diff --git a/src/tir/transforms/inject_software_pipeline.cc b/src/tir/transforms/inject_software_pipeline.cc
new file mode 100644
index 000000000000..b607ba485a6a
--- /dev/null
+++ b/src/tir/transforms/inject_software_pipeline.cc
@@ -0,0 +1,777 @@
+/*
+ * 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 inject_software_pipeline.cc
+ * \brief Transform annotated loops into pipelined one that parallelize producers and consumers
+ */
+#include <tvm/target/target.h>
+#include <tvm/tir/builtin.h>
+#include <tvm/tir/transform.h>
+
+#include "../../support/utils.h"
+#include "../schedule/utils.h"
+#include "./ir_utils.h"
+
+namespace tvm {
+namespace tir {
+
+namespace software_pipeline {
+
+/*!
+ * \brief Create a block and infer the access region with the given body.
+ *
+ * The result is a opaque block that doesn't contain any block iter vars. In case the body is a
+ * block realize without predicate, it is unnecessary to create a new block, the block of the block
+ * realize will be returned.
+ *
+ * \param body The body of the block.
+ * \param buffer_data_to_buffer The map from buffer data to buffer.
+ * \return The result block.
+ */
+Block MakeBlock(const Stmt& body, const Map<Var, Buffer>& buffer_data_to_buffer) {
+  if (const BlockRealizeNode* block_realize = body.as<BlockRealizeNode>()) {
+    if (is_one(block_realize->predicate)) {
+      // no need to create a new block
+      return block_realize->block;
+    }
+  }
+  Block block(/*iter_vars=*/{}, /*reads=*/{}, /*writes=*/{}, /*name_hint=*/"", /*body*/ body);
+  Array<Array<BufferRegion>> access = GetBlockReadWriteRegion(block, buffer_data_to_buffer);
+  BlockNode* n = block.CopyOnWrite();
+  n->reads = access[0];
+  n->writes = access[1];
+  return block;
+}
+
+/*! Structure that represents the stage and order of the software pipeline component. */
+struct PipelineStageOrder {
+  int stage;
+  int order;
+};
+
+using PipelineInfo = std::unordered_map<Block, PipelineStageOrder, ObjectPtrHash, ObjectPtrEqual>;
+
+struct BufferAccessInfo {
+  int def = -1;  // the defining stage of the buffer
+  int use = -1;  // the last using stage of the buffer
+};
+
+class PipelineOpaqueAccessRewriter {
+ public:
+  /*!
+   * \brief Constructor
+   * \param buffer_data_to_buffer The map from buffer data to buffer.
+   * \param buffer_remap The map from original buffer to the buffer with updated shape for
+   *        multi-versioning in the software pipeline.
+   * \param pipeline_loop The original loop to be software pipelined.
+   * \param fragment_info Information about tensor core fragment
+   */
+  PipelineOpaqueAccessRewriter(
+      const Map<Var, Buffer>& buffer_data_to_buffer, const Map<Buffer, Buffer>& buffer_remap,
+      const For& pipeline_loop,
+      const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info)
+      : buffer_data_to_buffer_(buffer_data_to_buffer),
+        buffer_remap_(buffer_remap),
+        pipeline_loop_(pipeline_loop),
+        fragment_info_(fragment_info) {}
+
+  PrimExpr Rewrite(const Call& call) {
+    // Intrinsic calls should be handled explicitly here as they are opaque accesses to
+    // buffer.
+    static const auto& load_matrix_sync = builtin::tvm_load_matrix_sync();
+    static const auto& store_matrix_sync = builtin::tvm_store_matrix_sync();
+    static const auto& mma_sync = builtin::tvm_mma_sync();
+    static const auto& access_ptr = builtin::tvm_access_ptr();
+    if (call->op.same_as(load_matrix_sync) || call->op.same_as(store_matrix_sync)) {
+      const Buffer& buffer = buffer_data_to_buffer_.at(Downcast<Var>(call->args[0]));
+      auto it = buffer_remap_.find(buffer);
+      if (it != buffer_remap_.end()) {
+        Array<PrimExpr> new_args = call->args;
+        const Buffer& new_buffer = (*it).second;
+        new_args.Set(4, RewriteWmmaFragmentIndex(buffer, new_buffer, call->args[4]));
+        return Call(call->dtype, call->op, new_args, call->span);
+      }
+    } else if (call->op.same_as(mma_sync)) {
+      Array<PrimExpr> new_args = call->args;
+      for (int i = 0; i < 4; i++) {
+        const Var& buffer_var = Downcast<Var>(call->args[i * 2]);
+        const PrimExpr& index = call->args[i * 2 + 1];
+        const Buffer& buffer = buffer_data_to_buffer_.at(buffer_var);
+        auto it = buffer_remap_.find(buffer);
+        if (it != buffer_remap_.end()) {
+          PrimExpr new_index = RewriteWmmaFragmentIndex(buffer, (*it).second, index);
+          new_args.Set(i * 2 + 1, new_index);
+        }
+      }
+      return Call(call->dtype, call->op, new_args, call->span);
+    } else if (call->op.same_as(access_ptr)) {
+      const Buffer& buffer = buffer_data_to_buffer_.at(Downcast<Var>(call->args[1]));
+      auto it = buffer_remap_.find(buffer);
+      if (it != buffer_remap_.end()) {
+        Array<PrimExpr> new_args = call->args;
+        const Buffer& new_buffer = (*it).second;
+        const PrimExpr& old_index = call->args[2];
+        PrimExpr offset;
+        if (new_buffer->strides.empty()) {
+          offset = foldl([](PrimExpr a, PrimExpr b, Span span) { return mul(a, b, span); },
+                         make_const(DataType::Int(32), 1), buffer->shape);
+        } else {
+          offset = new_buffer->strides[0];
+        }
+        PrimExpr new_index = old_index + floormod(pipeline_loop_->loop_var, 2) * offset;
+        new_args.Set(2, new_index);
+        return Call(call->dtype, call->op, new_args, call->span);
+      }
+    }
+    return call;
+  }
+
+ private:
+  int GetWmmaFragmentSize(const Buffer& buffer) {
+    auto it = fragment_info_.find(buffer->data.get());
+    ICHECK(it != fragment_info_.end());
+    const FragmentInfo& info = (*it).second;
+    return info.GetSize();
+  }
+
+  PrimExpr RewriteWmmaFragmentIndex(const Buffer& old_buffer, const Buffer& new_buffer,
+                                    const PrimExpr& old_index) {
+    PrimExpr new_buffer_offset = old_index;
+
+    int fragment_size = GetWmmaFragmentSize(old_buffer);
+    PrimExpr offset =
+        floordiv(foldl([](PrimExpr a, PrimExpr b, Span span) { return mul(a, b, span); },
+                       make_const(DataType::Int(32), 1), old_buffer->shape),
+                 fragment_size);
+    new_buffer_offset +=
+        floormod(pipeline_loop_->loop_var - pipeline_loop_->min, new_buffer->shape[0]) * offset;
+    return new_buffer_offset;
+  }
+
+  const Map<Var, Buffer>& buffer_data_to_buffer_;
+  const Map<Buffer, Buffer>& buffer_remap_;
+  const For& pipeline_loop_;
+  const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info_;
+};
+
+/*!
+ * \brief Rewriter for the body of the software pipeline. This pass inserts `floormod` to indices
+ * of the remapped buffer to select the version corresponding to the pipeline stage.
+ */
+class PipelineBodyRewriter : public StmtExprMutator {
+ public:
+  /*!
+   * \brief Constructor of PipelineBodyRewriter.
+   * \param buffer_data_to_buffer The map from buffer data to buffer.
+   * \param buffer_remap The map from original buffer to the buffer with updated shape for
+   *        multi-versioning in the software pipeline.
+   * \param pipeline_loop The original loop to be software pipelined.
+   * \param access_all_versions Whether all versions the the buffers in the software pipeline are
+   *        accessed. This will be used to update block access region. In the prologue and epilogue
+   *        of a two-stage software pipeline, only one version of these buffers are accessed.
+   * \param fragment_info Information about tensor core fragment
+   */
+  PipelineBodyRewriter(const Map<Var, Buffer>& buffer_data_to_buffer,
+                       const Map<Buffer, Buffer>& buffer_remap, For pipeline_loop,
+                       bool access_all_versions,
+                       const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info)
+      : buffer_data_to_buffer_(buffer_data_to_buffer),
+        buffer_remap_(buffer_remap),
+        pipeline_loop_(pipeline_loop),
+        access_all_versions_(access_all_versions),
+        opaque_access_rewriter_(buffer_data_to_buffer_, buffer_remap_, pipeline_loop_,
+                                fragment_info) {}
+
+ private:
+  BufferRegion RewritePipelineBufferRegion(const BufferRegion& buffer_region) const {
+    auto it = buffer_remap_.find(buffer_region->buffer);
+    if (it != buffer_remap_.end()) {
+      Region new_region = buffer_region->region;
+      const Buffer& new_buffer = (*it).second;
+      // For pipeline buffers, relax the access region of the first dimension to full extent
+      // if access_all_versions == true
+      Range accessed_version =
+          access_all_versions_
+              ? Range::FromMinExtent(0, new_buffer->shape[0])
+              : Range::FromMinExtent(floormod((pipeline_loop_->loop_var - pipeline_loop_->min),
+                                              new_buffer->shape[0]),
+                                     Integer(1));
+      new_region.insert(new_region.begin(), accessed_version);
+      return BufferRegion(new_buffer, new_region);
+    }
+    return buffer_region;
+  }
+
+  Stmt VisitStmt_(const BlockNode* op) final {
+    for (const Buffer& alloc_buffer : op->alloc_buffers) {
+      buffer_data_to_buffer_.Set(alloc_buffer->data, alloc_buffer);
+    }
+    Block block = Downcast<Block>(StmtExprMutator::VisitStmt_(op));
+    BlockNode* n = block.CopyOnWrite();
+    n->reads.MutateByApply([this](const BufferRegion& buffer_region) {
+      return RewritePipelineBufferRegion(buffer_region);
+    });
+    n->writes.MutateByApply([this](const BufferRegion& buffer_region) {
+      return RewritePipelineBufferRegion(buffer_region);
+    });
+    for (const Buffer& alloc_buffer : op->alloc_buffers) {
+      buffer_data_to_buffer_.erase(alloc_buffer->data);
+    }
+    return std::move(block);
+  }
+
+  Stmt VisitStmt_(const BufferStoreNode* op) final {
+    BufferStore store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
+    auto it = buffer_remap_.find(store->buffer);
+    if (it == buffer_remap_.end()) {
+      return std::move(store);
+    }
+    const Buffer& new_buffer = (*it).second;
+    auto* n = store.CopyOnWrite();
+    n->buffer = new_buffer;
+    PrimExpr version =
+        floormod((pipeline_loop_->loop_var - pipeline_loop_->min), new_buffer->shape[0]);
+    n->indices.insert(n->indices.begin(), version);
+    return std::move(store);
+  }
+
+  PrimExpr VisitExpr_(const BufferLoadNode* op) final {
+    BufferLoad load = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));
+    auto it = buffer_remap_.find(load->buffer);
+    if (it == buffer_remap_.end()) {
+      return std::move(load);
+    }
+    const Buffer& new_buffer = (*it).second;
+    auto* n = load.CopyOnWrite();
+    n->buffer = new_buffer;
+    PrimExpr version =
+        floormod((pipeline_loop_->loop_var - pipeline_loop_->min), new_buffer->shape[0]);
+    n->indices.insert(n->indices.begin(), version);
+    return std::move(load);
+  }
+
+  PrimExpr VisitExpr_(const CallNode* op) final {
+    Call call = Downcast<Call>(StmtExprMutator::VisitExpr_(op));
+    return opaque_access_rewriter_.Rewrite(call);
+  }
+
+  Map<Var, Buffer> buffer_data_to_buffer_;
+  Map<Buffer, Buffer> buffer_remap_;
+  For pipeline_loop_;
+  bool access_all_versions_;
+  PipelineOpaqueAccessRewriter opaque_access_rewriter_;
+};
+
+/*!
+ * \brief Rewriter for the software pipeline that rewrite a loop into a pipelined one.
+ */
+class PipelineRewriter : public StmtExprMutator {
+ public:
+  static Stmt Rewrite(
+      Map<Var, Buffer> buffer_data_to_buffer,
+      const std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual>& double_buffers,
+      const Array<Buffer> pipeline_allocs, const For& pipeline_loop,
+      const PipelineInfo& pipeline_info,
+      const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info) {
+    PipelineRewriter rewriter(buffer_data_to_buffer, double_buffers, pipeline_allocs, pipeline_loop,
+                              pipeline_info, fragment_info);
+    return rewriter.BuildPipeline();
+  }
+
+ private:
+  PipelineRewriter(Map<Var, Buffer> buffer_data_to_buffer,
+                   const std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual>& double_buffers,
+                   const Array<Buffer>& pipeline_allocs, const For& pipeline_loop,
+                   const PipelineInfo& pipeline_info,
+                   const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info)
+
+      : buffer_data_to_buffer_(std::move(buffer_data_to_buffer)),
+        double_buffers_(double_buffers),
+        pipeline_allocs_(pipeline_allocs),
+        pipeline_loop_(pipeline_loop),
+        pipeline_info_(pipeline_info),
+        fragment_info_(fragment_info) {}
+
+  Stmt BuildPipeline() {
+    // Step 1: Analyze accesses to the buffers in the pipeline and compute the number of versions
+    // need to maintain for each buffer.
+    std::unordered_map<Buffer, BufferAccessInfo, ObjectPtrHash, ObjectPtrEqual> infos =
+        GetBufferAccessInfo();
+    for (const Buffer& buffer : pipeline_allocs_) {
+      int num_versions = ComputeBufferVersions(buffer, infos.at(buffer));
+      if (num_versions > 1) {
+        buffer_remap_.Set(buffer, RewriteAllocBuffer(buffer, num_versions));
+      }
+    }
+
+    ordered_stmts_.resize(pipeline_info_.size());
+    for (const auto& pair : pipeline_info_) {
+      const Block& block = pair.first;
+      int order = pair.second.order;
+      ordered_stmts_.Set(order, block);
+    }
+
+    // Step 2: Emit the pipeline prologue, body and epilogue.
+    Stmt prologue = EmitImpl(pipeline_loop_->min, pipeline_loop_->min + max_stage_, true);
+    Stmt body = EmitImpl(pipeline_loop_->min + max_stage_,
+                         pipeline_loop_->min + pipeline_loop_->extent, false);
+    Stmt epilogue = EmitImpl(pipeline_loop_->min + pipeline_loop_->extent,
+                             pipeline_loop_->min + pipeline_loop_->extent + max_stage_, true);
+
+    SeqStmt stmt = SeqStmt({prologue, body, epilogue});
+
+    // Step 3: Make a new block that contains new buffer allocations after pipeline rewriting.
+    Array<Buffer> alloc_buffers;
+    for (const auto& alloc : pipeline_allocs_) {
+      alloc_buffers.push_back(buffer_remap_.Get(alloc).value_or(alloc));
+      buffer_data_to_buffer_.erase(alloc->data);
+    }
+    Block block = MakeBlock(stmt, buffer_data_to_buffer_);
+    block.CopyOnWrite()->alloc_buffers = std::move(alloc_buffers);
+    return BlockRealize({}, Bool(true), block);
+  }
+
+ private:
+  /*!
+   * \brief Analyze accesses to the buffers in the software pipeline.
+   *
+   * This method check the 'define' and 'use' stage of the buffers in the software pipeline, which
+   * can be used to compute the number of versions needed to maintain after rewriting.
+   */
+  std::unordered_map<Buffer, BufferAccessInfo, ObjectPtrHash, ObjectPtrEqual>
+  GetBufferAccessInfo() {
+    std::unordered_map<Buffer, BufferAccessInfo, ObjectPtrHash, ObjectPtrEqual> infos;
+    for (const auto& pair : pipeline_info_) {
+      const Block& block = pair.first;
+      int stage = pair.second.stage;
+      max_stage_ = std::max(max_stage_, stage);
+
+      for (const BufferRegion& write : block->writes) {
+        if (!infos.count(write->buffer)) {
+          infos.emplace(write->buffer, BufferAccessInfo{});
+        }
+        auto& info = infos.at(write->buffer);
+        if (info.def == -1) {
+          info.def = stage;
+        } else {
+          info.def = std::min(info.def, stage);
+        }
+      }
+
+      for (const BufferRegion& read : block->reads) {
+        if (!infos.count(read->buffer)) {
+          infos.emplace(read->buffer, BufferAccessInfo{});
+        }
+        auto& info = infos.at(read->buffer);
+        info.use = std::max(info.use, stage);
+      }
+    }
+    return infos;
+  }
+
+  /*!
+   * \brief Check whether two regions have intersections.
+   * \param region1 The first region.
+   * \param region2 The second region.
+   * \return Whether region1 and region2 have intersections.
+   */
+  bool MayConflict(Region region1, Region region2) {
+    ICHECK(region1.size() == region2.size());
+    for (size_t i = 0; i < region1.size(); i++) {
+      Range dim1 = region1[i];
+      Range dim2 = region2[i];
+      auto int_set1 = arith::IntSet::FromRange(dim1);
+      auto int_set2 = arith::IntSet::FromRange(dim2);
+      if (arith::Intersect({int_set1, int_set2}).IsNothing()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /*!
+   * \brief Compute the number of versions need to maintain for buffer accessed in the software
+   * pipeline.
+   *
+   * This method applies liveness analysis to the target buffer to compute the number of versions
+   * need to maintain during the software pipeline.
+   * Annotation `attr::double_buffer_scope` is handled here which provides a way to override the
+   * result of the analysis. Additional double buffering in the software pipeline can be useful
+   * to eliminate synchronizations in GPU devices.
+   *
+   * \param buffer The target buffer
+   * \param buffer_info The access information of the target buffer.
+   * \return The number of versions required for the target buffer.
+   */
+  int ComputeBufferVersions(const Buffer& buffer, const BufferAccessInfo& buffer_info) {
+    if (buffer_info.def == -1) {
+      // Keep the original number of versions as buffers defined outside the software pipeline
+      // should not be mutated.
+      return 1;
+    }
+
+    // `use - def + 1` is a upper bound of the needed versions
+    // We optimize a few case where the number of versions can be smaller than the upper bound
+    int num_versions = buffer_info.use - buffer_info.def + 1;
+    if (num_versions == 2) {
+      // A special case when `use - def + 1 == 2`. Double buffering is only needed in this case when
+      // these exists a reader block_i and a writer block_j such that
+      // order(block_i) < order(block_j) and stage(block_i) < stage(block_j) and the access regions
+      // of block_i and block_j overlap.
+      bool need_multi_version = false;
+      for (const auto& pair1 : pipeline_info_) {
+        const Block& writer_block = pair1.first;
+        const auto& writer_info = pair1.second;
+
+        auto it1 = std::find_if(writer_block->writes.begin(), writer_block->writes.end(),
+                                [&](const BufferRegion& buffer_region) {
+                                  return buffer_region->buffer.same_as(buffer);
+                                });
+        if (it1 == writer_block->writes.end()) {
+          continue;
+        }
+
+        for (const auto& pair2 : pipeline_info_) {
+          const Block& reader_block = pair2.first;
+          const auto& reader_info = pair2.second;
+          auto it2 = std::find_if(reader_block->reads.begin(), reader_block->reads.end(),
+                                  [&](const BufferRegion& buffer_region) {
+                                    return buffer_region->buffer.same_as(buffer);
+                                  });
+          if (it2 == reader_block->reads.end()) {
+            continue;
+          }
+          if (writer_info.order < reader_info.order && writer_info.stage < reader_info.stage &&
+              MayConflict((*it1)->region, (*it2)->region)) {
+            need_multi_version = true;
+            break;
+          }
+        }
+      }
+      if (!need_multi_version) {
+        num_versions = 1;
+      }
+    }
+    if (num_versions == 1 && double_buffers_.count(buffer)) {
+      num_versions = 2;
+    }
+    return num_versions;
+  }
+
+  /*!
+   * \brief Rewrite buffer allocation to keep multiple versions of original buffer for pipelined
+   * accesses.
+   * \param buffer The buffer to be resized.
+   * \param num_versions The number of versions to keep.
+   * \return The resized buffer.
+   */
+  Buffer RewriteAllocBuffer(const Buffer& buffer, int num_versions) {
+    ObjectPtr<BufferNode> new_buffer = make_object<BufferNode>(*(buffer.get()));
+    new_buffer->shape.insert(new_buffer->shape.begin(), PrimExpr(num_versions));
+    if (new_buffer->strides.size()) {
+      ICHECK(new_buffer->strides.size() + 1 == new_buffer->shape.size());
+      PrimExpr stride_0 = new_buffer->strides[0] * new_buffer->shape[1];
+      new_buffer->strides.insert(new_buffer->strides.begin(), stride_0);
+    }
+    return Buffer(new_buffer);
+  }
+
+  /*!
+   * \brief Emit the pipeline loop in the given range.
+   * \param start The start of the range
+   * \param end The end of the range
+   * \param unroll_loop Whether the loop should be unrolled.
+   * \return The result loop.
+   */
+  Stmt EmitImpl(PrimExpr start, PrimExpr end, bool unroll_loop) {
+    Array<Stmt> stmts;
+    PrimExpr new_loop_var;
+    PrimExpr extent = end - start;
+
+    auto make_nop = []() { return BlockRealize({}, Bool(true), MakeBlock(Evaluate(0), {})); };
+
+    if (!analyzer_.CanProve(extent > 0)) {
+      return make_nop();
+    }
+    bool is_unit_loop = analyzer_.CanProveEqual(extent, 1);
+    if (is_unit_loop) {
+      new_loop_var = start;  // use constants as the loop var for unit loops
+    } else {
+      new_loop_var = pipeline_loop_->loop_var.copy_with_suffix("");
+      analyzer_.Bind(Downcast<Var>(new_loop_var), Range(start, end));
+    }
+
+    for (const Block& block : ordered_stmts_) {
+      int stage = pipeline_info_.at(block).stage;
+      PrimExpr skewed_loop_var = new_loop_var - stage;
+      PrimExpr inbound = analyzer_.Simplify(pipeline_loop_->min <= skewed_loop_var) &&
+                         (skewed_loop_var < pipeline_loop_->min + pipeline_loop_->extent);
+      if (analyzer_.CanProve(!inbound)) {
+        continue;
+      }
+      Block new_block = Downcast<Block>(PipelineBodyRewriter(buffer_data_to_buffer_, buffer_remap_,
+                                                             pipeline_loop_, max_stage_ != 1,
+                                                             fragment_info_)(block));
+      Map<Var, PrimExpr> subst_map;
+      if (is_unit_loop) {
+        subst_map.Set(pipeline_loop_->loop_var, skewed_loop_var);
+      } else {
+        // normalize loop range
+        subst_map.Set(pipeline_loop_->loop_var, skewed_loop_var + (start - pipeline_loop_->min));
+      }
+      new_block = Downcast<Block>(Substitute(new_block, subst_map));
+      stmts.push_back(BlockRealize({}, inbound, new_block));
+    }
+
+    Stmt new_loop{nullptr};
+
+    if (stmts.empty()) {
+      return make_nop();
+    }
+    if (stmts.size() == 1) {
+      new_loop = stmts[0];
+    } else {
+      new_loop = SeqStmt(stmts);
+    }
+
+    if (!is_unit_loop) {
+      new_loop = For(Downcast<Var>(new_loop_var), pipeline_loop_->min, extent,
+                     unroll_loop ? ForKind::kUnrolled : pipeline_loop_->kind, std::move(new_loop));
+    }
+    return BlockRealize({}, Bool(true), MakeBlock(std::move(new_loop), buffer_data_to_buffer_));
+  }
+
+  arith::Analyzer analyzer_;
+  Map<Var, Buffer> buffer_data_to_buffer_;
+  const std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual>& double_buffers_;
+  Array<Buffer> pipeline_allocs_;
+  For pipeline_loop_;
+  PipelineInfo pipeline_info_;
+  const std::unordered_map<const VarNode*, FragmentInfo>& fragment_info_;
+  int max_stage_ = -1;
+  Map<Buffer, Buffer> buffer_remap_;
+  Array<Block> ordered_stmts_;
+};
+
+class PipelineInjector : private StmtExprMutator {
+ public:
+  static Stmt Inject(const PrimFunc& func) {
+    PipelineInjector injector;
+    for (const auto& kv : func->buffer_map) {
+      const Buffer& buffer = kv.second;
+      injector.buffer_data_to_buffer_.Set(buffer->data, buffer);
+    }
+    injector.fragment_info_ = GetTensorCoreFragmentInfo(func->body);
+    return injector(func->body);
+  }
+
+ private:
+  PipelineInjector() = default;
+
+  /*!
+   * \brief Check the pipeline satisfies the following conditions:
+   * 1) No conflicting order: The order of each statement should be unique.
+   * 2) No reordering with the same stage: Statements in the same stage are not allowed to be
+   * reordered.
+   */
+  void ValidatePipelineBody(const PipelineInfo& pipeline_info, const Array<Block>& original_order) {
+    std::unordered_set<int> used_orders;
+    std::unordered_map<int, int> stage_max_order;
+    for (const Block& block : original_order) {
+      const auto& stmt_info = pipeline_info.at(block);
+      int stage = stmt_info.stage;
+      int order = stmt_info.order;
+      CHECK(!used_orders.count(order))
+          << "ValueError: Two statements in the software pipeline cannot have the same order";
+      used_orders.insert(order);
+      CHECK(!stage_max_order.count(stage) || stage_max_order[stage] < order)
+          << "ValueError: Statements in the same stage of the software pipeline must have "
+             "increasing order.";
+      stage_max_order[stage] = order;
+    }
+  }
+
+  Stmt VisitStmt_(const ForNode* op) final {
+    // Step 1: Recursively rewrite the children first.
+    For for_node = Downcast<For>(StmtExprMutator::VisitStmt_(op));
+    if (!HasPipelineAnnotation(op)) {
+      return std::move(for_node);
+    }
+    // Step 2: Find the body and buffer allocations of the pipeline. The body can be direct child of
+    // the for-loop. If the for-loop has BlockRealize as its child, the pipeline body will be the
+    // child of the block.
+    Stmt pipeline_body{nullptr};
+    Array<Buffer> pipeline_allocs;
+    if (const auto* realize = for_node->body.as<BlockRealizeNode>()) {
+      const auto& block = realize->block;
+      for (const auto& buffer : block->alloc_buffers) {
+        ICHECK(buffer->IsInstance<BufferNode>());
+        buffer_data_to_buffer_.Set(buffer->data, buffer);
+      }
+      pipeline_body = block->body;
+      pipeline_allocs = block->alloc_buffers;
+    } else {
+      pipeline_body = for_node->body;
+    }
+
+    const SeqStmtNode* pipeline_body_seq = pipeline_body.as<SeqStmtNode>();
+    CHECK(pipeline_body_seq)
+        << "ValueError: The body of the software pipeline should be SeqStmt, got "
+        << pipeline_body->GetTypeKey();
+
+    // Step 3: Blockize the components of the pipeline. Each child of the pipelined loop will be
+    // converted into a block.
+    PipelineInfo pipeline_info;
+    Array<Block> original_order;  // pipeline body blocks in the original order
+
+    auto f_add_child = [&](const Stmt& child) {
+      original_order.push_back(MakeBlock(child, buffer_data_to_buffer_));
+    };
+    for (size_t i = 0; i < pipeline_body_seq->seq.size(); i++) {
+      const auto* nested_block_realize = pipeline_body_seq->seq[i].as<BlockRealizeNode>();
+      if (nested_block_realize && is_one(nested_block_realize->predicate) &&
+          nested_block_realize->block->body->IsInstance<SeqStmtNode>()) {
+        const Block& nested_pipeline_block = nested_block_realize->block;
+        ICHECK(
+            nested_pipeline_block->match_buffers.empty());  // match_buffer should have been lowered
+        for (const auto& buffer : nested_pipeline_block->alloc_buffers) {
+          pipeline_allocs.push_back(buffer);
+          buffer_data_to_buffer_.Set(buffer->data, buffer);
+        }
+        const auto* nested_seq = nested_pipeline_block->body.as<SeqStmtNode>();
+        for (size_t j = 0; j < nested_seq->seq.size(); j++) {
+          f_add_child(nested_seq->seq[j]);
+        }
+      } else {
+        f_add_child(pipeline_body_seq->seq[i]);
+      }
+    }
+
+    auto pipeline_stages =
+        Downcast<Array<Integer>>(op->annotations.at(attr::software_pipeline_stage));
+    auto pipeline_orders =
+        Downcast<Array<Integer>>(op->annotations.at(attr::software_pipeline_order));
+    CHECK_EQ(pipeline_stages.size(), original_order.size());
+    CHECK_EQ(pipeline_orders.size(), original_order.size());
+    for (size_t i = 0; i < pipeline_stages.size(); i++) {
+      PipelineStageOrder stage_order{/*stage=*/static_cast<int>(pipeline_stages[i]->value),
+                                     /*order=*/static_cast<int>(pipeline_orders[i]->value)};
+      pipeline_info.emplace(original_order[i], stage_order);
+    }
+    ValidatePipelineBody(pipeline_info, original_order);
+
+    // Step 4: Rewrite the pipeline body.
+    Stmt pipeline =
+        PipelineRewriter::Rewrite(buffer_data_to_buffer_, double_buffers, pipeline_allocs,
+                                  GetRef<For>(op), pipeline_info, fragment_info_);
+
+    if (const auto* realize = op->body.as<BlockRealizeNode>()) {
+      const auto& block = realize->block;
+      for (const auto& buffer : block->alloc_buffers) {
+        buffer_data_to_buffer_.erase(buffer->data);
+      }
+    }
+    return pipeline;
+  }
+
+  /*!
+   * \brief Add buffer allocations to a block and update the write region of the block.
+   * \param n The block pointer to which the buffer allocations are added.
+   * \param alloc_buffers The buffer allocations to be added.
+   */
+  void AddAllocBuffers(BlockNode* n, const Array<Buffer> alloc_buffers) {
+    for (const Buffer& alloc_buffer : alloc_buffers) {
+      n->alloc_buffers.push_back(alloc_buffer);
+      Region region;
+      region.reserve(alloc_buffer->shape.size());
+      for (const PrimExpr& dim : alloc_buffer->shape) {
+        region.push_back(Range::FromMinExtent(0, dim));
+      }
+      n->writes.push_back(BufferRegion(alloc_buffer, region));
+    }
+  }
+
+  Stmt VisitStmt_(const BlockNode* op) final {
+    for (const auto& buffer : op->alloc_buffers) {
+      buffer_data_to_buffer_.Set(buffer->data, buffer);
+    }
+
+    auto it = op->annotations.find(attr::double_buffer_scope);
+    if (it != op->annotations.end()) {
+      int buffer_index = Downcast<Integer>((*it).second);
+      CHECK(buffer_index >= 0 && static_cast<size_t>(buffer_index) < op->writes.size())
+          << "ValueError: Index of the buffer exceeds the size of the write regions of the block. ("
+          << buffer_index << " vs. " << op->writes.size() << ")";
+      double_buffers.insert(op->writes[buffer_index]->buffer);
+    }
+    Block block = Downcast<Block>(StmtExprMutator::VisitStmt_(op));
+
+    for (const auto& buffer : op->alloc_buffers) {
+      buffer_data_to_buffer_.erase(buffer->data);
+    }
+    return std::move(block);
+  }
+
+  bool HasPipelineAnnotation(const ForNode* op) const {
+    auto it1 = op->annotations.find(attr::software_pipeline_stage);
+    auto it2 = op->annotations.find(attr::software_pipeline_order);
+    bool has_stage = it1 != op->annotations.end();
+    bool has_order = it2 != op->annotations.end();
+    if (has_stage && has_order) {
+      return true;
+    }
+    if (has_stage) {
+      LOG(FATAL) << "ValueError: Order of the software pipeline is not defined.";
+    }
+    if (has_order) {
+      LOG(FATAL) << "ValueError: Stage of the software pipeline is not defined.";
+    }
+    return false;
+  }
+
+  Map<Var, Buffer> buffer_data_to_buffer_;
+  std::unordered_map<const VarNode*, FragmentInfo> fragment_info_;
+  std::unordered_set<Buffer, ObjectPtrHash, ObjectPtrEqual> double_buffers;
+};
+
+}  // namespace software_pipeline
+
+namespace transform {
+
+/*!
+ * \brief Transform annotated loops into pipelined one that parallelize producers and consumers.
+ * \return The IR transform pass.
+ */
+Pass InjectSoftwarePipeline() {
+  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
+    auto* fptr = f.CopyOnWrite();
+    fptr->body = software_pipeline::PipelineInjector::Inject(f);
+    fptr->body = ConvertSSA(std::move(fptr->body));
+    return f;
+  };
+  return CreatePrimFuncPass(pass_func, 0, "tir.InjectSoftwarePipeline", {});
+}
+
+TVM_REGISTER_GLOBAL("tir.transform.InjectSoftwarePipeline").set_body_typed(InjectSoftwarePipeline);
+
+}  // namespace transform
+
+}  // namespace tir
+}  // namespace tvm
diff --git a/src/tir/transforms/ir_utils.h b/src/tir/transforms/ir_utils.h
index da52a82a2f08..d7ae362b64d4 100644
--- a/src/tir/transforms/ir_utils.h
+++ b/src/tir/transforms/ir_utils.h
@@ -267,6 +267,39 @@ class ConditionalBoundsContext {
   std::unordered_map<const VarNode*, arith::IntSet> origin_map_;
 };
 
+// Information of tensor core fragment.
+struct FragmentInfo {
+  // fragment shape
+  int m, n, k;
+  // fragment layout (row-major or column-major)
+  std::string layout;
+  // scope of the fragment (wmma.matrix_a, wmma.matrix_b, or wmma.accumulator)
+  std::string scope;
+  FragmentInfo() = default;
+  FragmentInfo(int _m, int _n, int _k, const std::string& _layout, const std::string& _scope)
+      : m(_m), n(_n), k(_k), layout(_layout), scope(_scope) {}
+
+  int GetSize() const {
+    if (scope == "wmma.matrix_a") {
+      return m * k;
+    } else if (scope == "wmma.matrix_b") {
+      return n * k;
+    } else if (scope == "wmma.accumulator") {
+      return m * n;
+    } else {
+      ICHECK(0);
+      throw;
+    }
+  }
+};
+
+/*!
+ * \brief Extract information of tensor core fragment from the IR.
+ * \param stmt The stmt to visit.
+ * \return Map from buffer variables to the fragment info.
+ */
+std::unordered_map<const VarNode*, FragmentInfo> GetTensorCoreFragmentInfo(const Stmt& stmt);
+
 }  // namespace tir
 }  // namespace tvm
 #endif  // TVM_TIR_TRANSFORMS_IR_UTILS_H_
diff --git a/src/tir/transforms/legalize_packed_calls.cc b/src/tir/transforms/legalize_packed_calls.cc
index cb2b50260326..2d8b6681fa84 100644
--- a/src/tir/transforms/legalize_packed_calls.cc
+++ b/src/tir/transforms/legalize_packed_calls.cc
@@ -75,6 +75,12 @@ class PackedCallLegalizer : public StmtExprMutator {
             new_stmts.push_back(tir::Evaluate(
                 tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
                                {sid_array, 0, tir::builtin::kArrData, call->args[i]})));
+            new_stmts.push_back(tir::Evaluate(
+                tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                               {sid_array, 0, tir::builtin::kArrDeviceType, kDLCPU})));
+            new_stmts.push_back(tir::Evaluate(
+                tvm::tir::Call(DataType::Handle(), tvm::tir::builtin::tvm_struct_set(),
+                               {sid_array, 0, tir::builtin::kArrDeviceId, 0})));
             packed_args.push_back(sid_array);
           }
         }
diff --git a/src/tir/transforms/loop_partition.cc b/src/tir/transforms/loop_partition.cc
index c4b83e05706d..59ac339006f4 100644
--- a/src/tir/transforms/loop_partition.cc
+++ b/src/tir/transforms/loop_partition.cc
@@ -433,7 +433,7 @@ class LoopPartitioner : public StmtMutator {
 
   std::pair<IntSet, ExpressionSet> GetIntervalAndCondset(const Partition& partitions,
                                                          const arith::IntervalSet& for_interval,
-                                                         bool cond_value);
+                                                         bool cond_value, bool has_partition_hint);
 
   inline Stmt MakeFor(const Object* op, PrimExpr extent, Stmt body);
 
@@ -448,7 +448,8 @@ class LoopPartitioner : public StmtMutator {
 // Returns an interval (in the first component) in which all the conditions
 // given in the second component provably have value given by cond_value
 std::pair<IntSet, ExpressionSet> LoopPartitioner::GetIntervalAndCondset(
-    const Partition& partitions, const arith::IntervalSet& for_interval, bool cond_value) {
+    const Partition& partitions, const arith::IntervalSet& for_interval, bool cond_value,
+    bool has_partition_hint) {
   Array<IntSet> sets;
   ExpressionSet cond_set;
 
@@ -463,6 +464,32 @@ std::pair<IntSet, ExpressionSet> LoopPartitioner::GetIntervalAndCondset(
     }
   }
   IntSet interval = sets.empty() ? IntSet::Nothing() : Intersect(sets);
+
+  // Try to find the intersection of the cond_intervals until the intersection
+  // is nothing when has_partition_hint is true.
+  if (interval.IsNothing() && has_partition_hint) {
+    arith::IntervalSet cond_intersection = arith::IntervalSet::Everything();
+    cond_set.clear();
+
+    for (const auto& kv : partitions) {
+      if (kv.first.second == cond_value) {
+        arith::IntervalSet cond_interval = Downcast<arith::IntervalSet>(kv.second);
+        arith::IntervalSet intersection = arith::Intersect(&analyzer_, cond_interval, for_interval);
+        if (!intersection->IsEmpty()) {
+          cond_intersection = arith::Intersect(&analyzer_, cond_intersection, cond_interval);
+          // Return the latest interval and cond_set if the cond_intersection is nothing.
+          if (!cond_intersection->IsEmpty()) {
+            cond_set.insert(kv.first.first);
+            interval = arith::IntervalSet(analyzer_.Simplify(cond_intersection->min_value),
+                                          analyzer_.Simplify(cond_intersection->max_value));
+          } else {
+            break;
+          }
+        }
+      }
+    }
+  }
+
   return std::make_pair(interval, cond_set);
 }
 
@@ -531,12 +558,12 @@ Stmt LoopPartitioner::TryPartition(const Stmt& stmt, Var var, PrimExpr min, Prim
   ExpressionSet cond_set;
   // find an interval in which all conditions on var are true
   std::tie(middle_interval, cond_set) =
-      GetIntervalAndCondset(finder.partitions, for_interval, true);
+      GetIntervalAndCondset(finder.partitions, for_interval, true, has_partition_hint_);
   if (middle_interval.IsNothing()) {
     // if such interval doesn't exist, find an interval in which all
     // conditions on var are false
     std::tie(middle_interval, cond_set) =
-        GetIntervalAndCondset(finder.partitions, for_interval, false);
+        GetIntervalAndCondset(finder.partitions, for_interval, false, has_partition_hint_);
     if (middle_interval.IsNothing())
       // we couldn't find an interval in which the conditions are provably true or false
       // Therefore, we can't partition the loop based on those conds
diff --git a/src/tir/transforms/lower_thread_allreduce.cc b/src/tir/transforms/lower_thread_allreduce.cc
index 6f7c09cdcf2d..1c6aa161e473 100644
--- a/src/tir/transforms/lower_thread_allreduce.cc
+++ b/src/tir/transforms/lower_thread_allreduce.cc
@@ -218,6 +218,33 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
     int reduce_extent, group_extent;
     PrimExpr reduce_index = FlattenThread(vred, &reduce_extent);
     PrimExpr group_index = FlattenThread(vpar, &group_extent);
+
+    // the longest contiguous reduce extent after flattening
+    int contiguous_reduce_extent = 1;
+    std::vector<std::tuple<int, int, bool>> block_threads;  // tuple(dim_index, extent, is_reduce)
+    for (const ThreadEntry& thr : vred) {
+      if (thr.scope.rank == 1) {  // threadIdx
+        block_threads.emplace_back(thr.scope.dim_index, thr.extent, true);
+      }
+    }
+    for (const ThreadEntry& thr : vpar) {
+      if (thr.scope.rank == 1) {  // threadIdx
+        block_threads.emplace_back(thr.scope.dim_index, thr.extent, false);
+      }
+    }
+    // sort according to dim_index
+    std::sort(block_threads.begin(), block_threads.end());
+    for (auto&& thr_attr : block_threads) {
+      int dim_index, extent;
+      bool is_reduce;
+      std::tie(dim_index, extent, is_reduce) = thr_attr;
+      if (is_reduce) {
+        contiguous_reduce_extent *= extent;
+      } else {
+        break;
+      }
+    }
+
     std::vector<Stmt> seq;
     std::vector<Var> shared_bufs(size);
     std::vector<Stmt> local_vars;
@@ -238,9 +265,8 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
     // broadcast results from lane 0 to all other lanes and store
     // the final reduction result to the proper location.
     //
-    if (is_warp_reduction(types)) {
-      // TODO(tvm-team) sub-warp reduction support.
-      ICHECK_EQ(reduce_extent, warp_size_) << "not a warp reduction";
+    if (is_warp_reduction(types, group_extent, reduce_extent, contiguous_reduce_extent)) {
+      ICHECK_LE(reduce_extent, warp_size_) << "not a warp reduction";
       //
       // This is the index to the reduction variable, one reduction
       // variable per warp. Local scope seems easier to reason without
@@ -269,6 +295,9 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
       {
         PrimExpr pred = const_true(1);
         PrimExpr mask = Call(mask_dtype, builtin::tvm_warp_activemask(), {});
+        if (group_extent > 1) {
+          mask = mask & (((1 << reduce_extent) - 1) << (reduce_extent * group_index));
+        }
         seq.emplace_back(Store(mask_var, mask, index, pred));
         // Push allocation with an empty body. Later this will be fixed
         // when the entire body is ready.
@@ -277,7 +306,11 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
       }
 
       // Emit reductions within a warp.
-      for (int offset = warp_size_ / 2; offset > 0; offset /= 2) {
+      int start_offset = 1;
+      while (start_offset * 2 < reduce_extent) {
+        start_offset *= 2;
+      }
+      for (int offset = start_offset; offset > 0; offset /= 2) {
         // Load reduction values, no synchronization needed.
         Array<PrimExpr> a, b;
         for (size_t i = 0; i < size; ++i) {
@@ -323,13 +356,14 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
 
       // Broadcast the reduction result from lane 0 to all other lanes.
       // This avoids to emit predicated stores, as all threads are
-      // uniformmly writting the same result.
+      // uniformly writting the same result.
       //
       for (size_t i = 0; i < size; ++i) {
         Var var = shared_bufs[i];
         PrimExpr pred = const_true(types[i].lanes());
         PrimExpr val = Load(types[i], var, index, pred);
-        PrimExpr splat = WarpShuffle(builtin::tvm_warp_shuffle(), mask_var, val, 0);
+        PrimExpr splat =
+            WarpShuffle(builtin::tvm_warp_shuffle(), mask_var, val, reduce_extent * group_index);
         seq.push_back(Store(var, splat, index, pred));
       }
 
@@ -346,7 +380,6 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
         warp_allocs_.insert(node.get());
       }
     } else {
-      int threadx_extent = 1;
       if (reduce_extent == 1) {
         // special case, no reduction is needed.
         std::vector<Stmt> stores(size);
@@ -357,10 +390,6 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
         }
         return SeqStmt::Flatten(stores);
       }
-      // Whether the threadIdx.x is involved in reduction.
-      if (vred[0].scope.dim_index == 0) {
-        threadx_extent = vred[0].extent;
-      }
       // This sync is necessary because there might be incomplete read of
       // previous iteration on the same buffer.
       seq.emplace_back(SyncThread("shared"));
@@ -372,7 +401,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
       }
       seq.emplace_back(SyncThread("shared"));
       seq.emplace_back(MakeBufAllreduce(combiner, types, shared_bufs, reduce_index, group_index,
-                                        reduce_extent, threadx_extent));
+                                        reduce_extent, group_extent, contiguous_reduce_extent));
       for (size_t idx = 0; idx < size; ++idx) {
         ICHECK(!load_remap_.count(buffers[idx]));
         PrimExpr pred = const_true(types[idx].lanes());
@@ -402,7 +431,7 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
   // make allreduce.
   Stmt MakeBufAllreduce(const CommReducerNode* combiner, const std::vector<DataType>& types,
                         const Array<Var>& shared_bufs, PrimExpr reduce_index, PrimExpr group_index,
-                        int reduce_extent, int threadx_extent) {
+                        int reduce_extent, int group_extent, int contiguous_reduce_extent) {
     // Get next power of two
     int reduce_align = 1;
     while (reduce_extent > reduce_align) {
@@ -444,9 +473,13 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
       seq.emplace_back(IfThenElse(cond, freduce(reduce_align)));
       seq.emplace_back(SyncThread("shared"));
     }
-    ICHECK(threadx_extent >= 1 && warp_size_ >= 1);
+
     // normal synchronization
-    while (reduce_align > threadx_extent || reduce_align > warp_size_) {
+    bool warp_align = group_extent == 1 || contiguous_reduce_extent % warp_size_ == 0;
+    while (reduce_align > contiguous_reduce_extent || reduce_align > warp_size_ || !warp_align) {
+      if (reduce_align == 1) {
+        break;
+      }
       reduce_align = reduce_align >> 1;
       PrimExpr cond = reduce_index < reduce_align;
       seq.emplace_back(IfThenElse(cond, freduce(reduce_align)));
@@ -534,22 +567,21 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
   }
 
   // Emit warp shuffle  calls.
-  PrimExpr WarpShuffle(const Op& op, Var mask_var, PrimExpr val, int delta_or_lane) {
+  PrimExpr WarpShuffle(const Op& op, Var mask_var, PrimExpr val, PrimExpr delta_or_lane) {
     PrimExpr pred = const_true(1);
     PrimExpr index(0);
     PrimExpr mask = Load(DataType::UInt(32), mask_var, index, pred);
     PrimExpr width = IntImm(DataType::Int(32), warp_size_);
-    Array<PrimExpr> args{mask, val, IntImm(DataType::Int(32), delta_or_lane), width, width};
+    Array<PrimExpr> args{mask, val, delta_or_lane, width, width};
     return Call(val.dtype(), op, args);
   }
 
-  // Check if this is a reduction on threadIdx.x and its extent matches
-  // the warp size.
+  // Check if we can use warp level reduction.
   //
-  // TODO(tvm-team) reduction with a sub-warp of 8 or 16 threads.
   // Note: The ROCm backend will only have warp reductions for now.
   // Also, the warp/wavefront size differs (64 on rocm, 32 on cuda).
-  bool is_warp_reduction(const std::vector<DataType>& types) const {
+  bool is_warp_reduction(const std::vector<DataType>& types, int group_extent, int reduce_extent,
+                         int contiguous_reduce_extent) const {
     // Only cuda target supports warp reductions.
     if ((target_->kind->name != "cuda") && (target_->kind->name != "rocm")) return false;
 
@@ -575,18 +607,25 @@ class ThreadAllreduceBuilder final : public StmtExprMutator {
       return false;
     }
 
-    const AttrStmtNode* op = thread_extents_.back();
-    DCHECK_EQ(op->attr_key, attr::thread_extent);
-
-    IterVar iv = Downcast<IterVar>(op->node);
-    ThreadEntry e;
-    e.scope = runtime::ThreadScope::Create(iv->thread_tag);
-    e.extent = 0;
-    if (auto ptr = op->value.as<IntImmNode>()) {
-      e.extent = static_cast<int>(ptr->value);
+    // reduce region must be contiguous.
+    if (contiguous_reduce_extent != reduce_extent) {
+      return false;
     }
 
-    return e.extent == warp_size_ && e.scope.dim_index == 0 && e.scope.rank == 1;
+    // whether reduce_extent and group_extent are vaild for warp reduction.
+    if (target_->kind->name == "rocm") {
+      return reduce_extent == warp_size_;
+    } else {  // target_->kind->name == "cuda"
+      if (reduce_extent == 1) {
+        return false;  // no need to warp reduce
+      } else {
+        if (warp_size_ % reduce_extent == 0) {
+          return true;  // warp size is multiple of reduce extent
+        } else {
+          return group_extent == 1 && reduce_extent <= warp_size_;
+        }
+      }
+    }
   }
 
   // The target.
diff --git a/src/tir/transforms/lower_tvm_builtin.cc b/src/tir/transforms/lower_tvm_builtin.cc
index a5ecf4ba8296..07063597a0c0 100644
--- a/src/tir/transforms/lower_tvm_builtin.cc
+++ b/src/tir/transforms/lower_tvm_builtin.cc
@@ -82,10 +82,12 @@ class BuiltinLower : public StmtExprMutator {
     // allocate space to hold prepare stmts before s
     prep_seq_stack_.emplace_back(std::vector<Stmt>());
 
-    auto stmt = StmtExprMutator::VisitStmt(s);
     auto& scope = alloca_scope_.back();
-    ICHECK_EQ(scope.run_shape_stack, -1);
-    ICHECK_EQ(scope.run_array_stack, 0);
+    auto original_run_shape_stack = scope.run_shape_stack;
+    auto original_run_array_stack = scope.run_array_stack;
+    auto stmt = StmtExprMutator::VisitStmt(s);
+    ICHECK_EQ(scope.run_shape_stack, original_run_shape_stack);
+    ICHECK_EQ(scope.run_array_stack, original_run_array_stack);
 
     auto prep_seq = std::move(prep_seq_stack_.back());
     prep_seq_stack_.pop_back();
@@ -126,7 +128,7 @@ class BuiltinLower : public StmtExprMutator {
       if (const auto* dev_type = device_type_.as<IntImmNode>()) {
         auto storage_scope = Downcast<PointerType>(op->buffer_var->type_annotation)->storage_scope;
         if (dev_type->value == kDLCPU && storage_scope == "global") {
-          int32_t constant_size = op->constant_allocation_size();
+          size_t constant_size = op->ConstantAllocationSize();
           if (constant_size > 0 && constant_size * nbytes < runtime::kMaxStackAlloca) {
             return stmt;
           }
@@ -209,10 +211,26 @@ class BuiltinLower : public StmtExprMutator {
       return MakeArray(op);
     } else if (op->op.same_as(builtin::tvm_context_id())) {
       return make_zero(op->dtype);
+    } else if (op->op.same_as(builtin::mem_copy())) {
+      return MakeMemCopy(op);
     } else {
       return StmtExprMutator::VisitExpr_(op);
     }
   }
+
+  PrimExpr MakeMemCopy(const CallNode* op) {
+    PrimExpr dst = op->args[0];
+    PrimExpr src = op->args[1];
+    PrimExpr size = op->args[2];
+
+    std::string fdevapi_prefix =
+        "device_api." + std::string(runtime::DeviceName(device_type_.as<IntImmNode>()->value));
+
+    Call call_packed = Call(DataType::Int(32), builtin::tvm_call_packed(),
+                            {StringImm(fdevapi_prefix + ".mem_copy"), dst, src, size});
+    return VisitExpr(call_packed);
+  }
+
   // call shape
   PrimExpr MakeShape(const CallNode* op) {
     // if args.size() == 0, it represents a scalar shape ()
diff --git a/src/tir/transforms/lower_warp_memory.cc b/src/tir/transforms/lower_warp_memory.cc
index 48313295d5ef..f316ae9606d0 100644
--- a/src/tir/transforms/lower_warp_memory.cc
+++ b/src/tir/transforms/lower_warp_memory.cc
@@ -215,7 +215,7 @@ class WarpAccessRewriter : protected StmtExprMutator {
   // warp memory to local memory.
   Stmt Rewrite(const AllocateNode* op) {
     buffer_ = op->buffer_var.get();
-    int alloc_size = op->constant_allocation_size();
+    int alloc_size = op->ConstantAllocationSize();
     ICHECK_GT(alloc_size, 0) << "warp memory only support constant alloc size";
     alloc_size *= op->dtype.lanes();
     std::tie(warp_index_, width_) = WarpIndexFinder(warp_size_).Find(op->body);
diff --git a/src/tir/transforms/merge_dynamic_shared_memory_allocations.cc b/src/tir/transforms/merge_dynamic_shared_memory_allocations.cc
index f3ff1f37a5da..b10e4439b99d 100644
--- a/src/tir/transforms/merge_dynamic_shared_memory_allocations.cc
+++ b/src/tir/transforms/merge_dynamic_shared_memory_allocations.cc
@@ -447,7 +447,7 @@ class DynamicSharedMemoryRewriter : public StmtExprMutator {
     // compiler can do a better job with register allocation.
     const uint64_t match_range = 16;
     uint64_t op_elem_bits = op->dtype.bits() * op->dtype.lanes();
-    uint64_t const_nbits = static_cast<uint64_t>(op->constant_allocation_size() * op_elem_bits);
+    uint64_t const_nbits = static_cast<uint64_t>(op->ConstantAllocationSize() * op_elem_bits);
     // disable reuse of small arrays, they will be lowered to registers in LLVM
     // This rules only apply if we are using non special memory
     if (const_nbits > 0 && const_nbits <= 32) {
diff --git a/src/tir/transforms/narrow_datatype.cc b/src/tir/transforms/narrow_datatype.cc
index dc34626205a1..dd5f54e52455 100644
--- a/src/tir/transforms/narrow_datatype.cc
+++ b/src/tir/transforms/narrow_datatype.cc
@@ -253,6 +253,23 @@ class DataTypeRewriter : public StmtExprMutator {
     return StmtExprMutator::VisitExpr_(op);
   }
 
+  PrimExpr VisitExpr_(const RampNode* op) final {
+    PrimExpr base = VisitExpr(op->base);
+    PrimExpr stride = VisitExpr(op->stride);
+    if (base.same_as(op->base) && stride.same_as(op->stride)) {
+      return GetRef<PrimExpr>(op);
+    } else {
+      if (base.dtype().is_int()) {
+        ICHECK(stride.dtype().is_int()) << "Ramp base is int but stride is " << stride.dtype();
+        int bits = std::max(base.dtype().bits(), stride.dtype().bits());
+        DataType dtype = base.dtype().with_bits(bits);
+        if (base.dtype() != dtype) base = cast(dtype, base);
+        if (stride.dtype() != dtype) stride = cast(dtype, stride);
+      }
+      return Ramp(base, stride, op->lanes);
+    }
+  }
+
   PrimExpr VisitExpr_(const SizeVarNode* op) final {
     if (visitor_.vmap.find(op) != visitor_.vmap.end()) {
       if (vmap_.find(op) == vmap_.end()) {
diff --git a/src/tir/transforms/plan_update_buffer_allocation_location.cc b/src/tir/transforms/plan_update_buffer_allocation_location.cc
index 97153aedc6a3..6b495b3bf4b5 100644
--- a/src/tir/transforms/plan_update_buffer_allocation_location.cc
+++ b/src/tir/transforms/plan_update_buffer_allocation_location.cc
@@ -44,7 +44,7 @@ class BufferAllocationLocator : public StmtExprMutator {
     // create buffers to be allocated at each stmts
     for (const auto& kv : buffer_lca) {
       const Buffer& buffer = kv.first;
-      const StmtNode* stmt = kv.second.defined() ? kv.second.value().get() : nullptr;
+      const StmtNode* stmt = kv.second.get();
       if (arg_buffers.count(buffer.get())) {
         continue;
       }
diff --git a/src/tir/transforms/split_host_device.cc b/src/tir/transforms/split_host_device.cc
index 7f2ecf54dfcb..e54aceb16a77 100644
--- a/src/tir/transforms/split_host_device.cc
+++ b/src/tir/transforms/split_host_device.cc
@@ -106,6 +106,11 @@ class VarUseDefAnalysis : public StmtExprMutator {
     return StmtExprMutator::VisitStmt_(op);
   }
 
+  Stmt VisitStmt_(const AllocateConstNode* op) final {
+    this->HandleDef(op->buffer_var.get());
+    return StmtExprMutator::VisitStmt_(op);
+  }
+
   Stmt VisitStmt_(const StoreNode* op) final {
     this->HandleUse(op->buffer_var);
     return StmtExprMutator::VisitStmt_(op);
diff --git a/src/tir/transforms/storage_flatten.cc b/src/tir/transforms/storage_flatten.cc
index ccc660509ca1..783ad13e1ad0 100644
--- a/src/tir/transforms/storage_flatten.cc
+++ b/src/tir/transforms/storage_flatten.cc
@@ -1196,7 +1196,7 @@ class StorageFlattener : public StmtExprMutator {
 
       // use small alignment for small arrays
       auto dtype = op->buffer->dtype;
-      int32_t const_size = AllocateNode::constant_allocation_size(shape);
+      size_t const_size = AllocateNode::ConstantAllocationSize(shape);
       int align = GetTempAllocaAlignment(dtype, const_size);
       if (skey.tag.length() != 0) {
         MemoryInfo info = GetMemoryInfo(skey.to_string());
diff --git a/src/tir/transforms/storage_rewrite.cc b/src/tir/transforms/storage_rewrite.cc
index 409b7c262954..9d90e0b3f226 100644
--- a/src/tir/transforms/storage_rewrite.cc
+++ b/src/tir/transforms/storage_rewrite.cc
@@ -729,7 +729,7 @@ class StoragePlanRewriter : public StmtExprMutator {
                     src_entry->attach_scope_ == thread_scope_ &&
                     src_entry->elem_type == alloc->dtype.element_of() &&
                     visitor.Check(s.stmt, var, src)) {
-                  uint64_t const_nbits = static_cast<uint64_t>(alloc->constant_allocation_size()) *
+                  uint64_t const_nbits = static_cast<uint64_t>(alloc->ConstantAllocationSize()) *
                                          alloc->dtype.bits() * alloc->dtype.lanes();
                   if (src_entry->const_nbits == const_nbits && !inplace_found) {
                     // successfully inplace
@@ -801,7 +801,7 @@ class StoragePlanRewriter : public StmtExprMutator {
     // compiler can do a better job with register allocation.
     const uint64_t match_range = 16;
     uint64_t op_elem_bits = op->dtype.bits() * op->dtype.lanes();
-    uint64_t const_nbits = static_cast<uint64_t>(op->constant_allocation_size() * op_elem_bits);
+    uint64_t const_nbits = static_cast<uint64_t>(op->ConstantAllocationSize() * op_elem_bits);
     // disable reuse of small arrays, they will be lowered to registers in LLVM
     // This rules only apply if we are using non special memory
     if (scope.tag.length() == 0) {
@@ -1032,6 +1032,14 @@ class VectorTypeAccessChecker : public StmtExprVisitor {
     StmtExprVisitor::VisitStmt_(op);
   }
 
+  void VisitStmt_(const AllocateConstNode* op) final {
+    const Array<PrimExpr>& extents = op->extents;
+    PrimExpr extent = extents.size() ? extents[extents.size() - 1] : NullValue<PrimExpr>();
+    OnArrayDeclaration(op->buffer_var, op->dtype, extent, BufferVarInfo::kAllocateNode);
+
+    StmtExprVisitor::VisitStmt_(op);
+  }
+
   void VisitExpr_(const LetNode* op) final {
     HandleLetNode(op->var);
     StmtExprVisitor::VisitExpr_(op);
@@ -1348,6 +1356,27 @@ class VectorTypeRewriter : public StmtExprMutator {
     return Allocate(new_buffer_var, info.new_element_dtype, extents, op->condition, op->body);
   }
 
+  Stmt VisitStmt_(const AllocateConstNode* op) final {
+    Stmt stmt = StmtExprMutator::VisitStmt_(op);
+    op = stmt.as<AllocateConstNode>();
+
+    auto it = rewrite_map_.find(op->buffer_var.get());
+    if (it == rewrite_map_.end()) {
+      return stmt;
+    }
+
+    const auto& info = it->second;
+
+    Var new_buffer_var = info.new_buffer_var;
+
+    int factor = info.new_element_dtype.lanes() / op->dtype.lanes();
+
+    Array<PrimExpr> extents = op->extents;
+    extents.Set(extents.size() - 1,
+                extents[extents.size() - 1] / make_const(extents[0].dtype(), factor));
+    return AllocateConst(new_buffer_var, info.new_element_dtype, extents, op->data, op->body);
+  }
+
   /* Update the parameters and all remaining variable references
    *
    * Should be called after calling operator() on the body of the
diff --git a/src/tir/transforms/tensorcore_infer_fragment.cc b/src/tir/transforms/tensorcore_infer_fragment.cc
index 1836b8ecec0d..89b9307198f6 100644
--- a/src/tir/transforms/tensorcore_infer_fragment.cc
+++ b/src/tir/transforms/tensorcore_infer_fragment.cc
@@ -39,17 +39,6 @@ namespace tir {
 // Get fragment information from tensor intrinsics
 class FragmentGetter : public StmtExprVisitor {
  public:
-  // fragment metadata
-  struct FragmentInfo {
-    // fragment shape
-    int m, n, k;
-    // fragment layout (row-major or column-major)
-    std::string layout;
-    FragmentInfo() = default;
-    FragmentInfo(int _m, int _n, int _k, const std::string& _layout)
-        : m(_m), n(_n), k(_k), layout(_layout) {}
-  };
-
   void VisitExpr_(const CallNode* op) final {
     StmtExprVisitor::VisitExpr_(op);
 
@@ -83,9 +72,9 @@ class FragmentGetter : public StmtExprVisitor {
         // store metadata
         FragmentInfo info;
         if (scope == "wmma.matrix_a" || scope == "wmma.matrix_b") {
-          info = FragmentInfo(m->value, n->value, k->value, layout->value);
+          info = FragmentInfo(m->value, n->value, k->value, layout->value, scope);
         } else if (scope == "wmma.accumulator") {
-          info = FragmentInfo(m->value, n->value, k->value, "");
+          info = FragmentInfo(m->value, n->value, k->value, "", scope);
         }
         fragments[buffer_var] = info;
       }
@@ -111,7 +100,7 @@ class FragmentGetter : public StmtExprVisitor {
         ICHECK_EQ(n->value, info.n);
         ICHECK_EQ(k->value, info.k);
       } else {
-        FragmentInfo info(m->value, n->value, k->value, "");
+        FragmentInfo info(m->value, n->value, k->value, "", scope);
         fragments[buffer_var] = info;
       }
     }
@@ -126,6 +115,12 @@ class FragmentGetter : public StmtExprVisitor {
   std::unordered_map<const VarNode*, FragmentInfo> fragments;
 };
 
+std::unordered_map<const VarNode*, FragmentInfo> GetTensorCoreFragmentInfo(const Stmt& stmt) {
+  FragmentGetter getter;
+  getter(stmt);
+  return std::move(getter.fragments);
+}
+
 // Check shape of fragment making sure it is a valid shape for tvm_mma_sync
 class FragmentChecker : public StmtExprVisitor {
  public:
@@ -157,8 +152,8 @@ class FragmentChecker : public StmtExprVisitor {
   bool CheckShape(const VarNode* buffer1, const VarNode* buffer2) {
     ICHECK(fragment_getter.fragments.count(buffer1));
     ICHECK(fragment_getter.fragments.count(buffer2));
-    FragmentGetter::FragmentInfo info1 = fragment_getter.fragments.at(buffer1);
-    FragmentGetter::FragmentInfo info2 = fragment_getter.fragments.at(buffer2);
+    FragmentInfo info1 = fragment_getter.fragments.at(buffer1);
+    FragmentInfo info2 = fragment_getter.fragments.at(buffer2);
     return info1.m == info2.m && info1.n == info2.n && info1.k == info2.k;
   }
   // Fragment infomation
@@ -175,7 +170,7 @@ class InferFragmenter : public StmtMutator {
     const VarNode* buffer = op->buffer_var.get();
     if (fragment_getter.fragments.count(buffer)) {
       // Add attribute to fragments allocation
-      FragmentGetter::FragmentInfo info = fragment_getter.fragments.at(buffer);
+      FragmentInfo info = fragment_getter.fragments.at(buffer);
 
       // Add shape attribute to all fragments
       std::string shape =
diff --git a/src/tir/transforms/vectorize_loop.cc b/src/tir/transforms/vectorize_loop.cc
index cd2d230f5775..0c9c97af650d 100644
--- a/src/tir/transforms/vectorize_loop.cc
+++ b/src/tir/transforms/vectorize_loop.cc
@@ -101,7 +101,7 @@ class Vectorizer : public StmtMutator, public ExprFunctor<PrimExpr(const PrimExp
   using StmtMutator::operator();
 
   Vectorizer(Var var, int var_lanes) : var_(var), var_lanes_(var_lanes) {
-    ramp_ = Ramp(0, 1, var_lanes);
+    ramp_ = Ramp(IntImm(var->dtype, 0), IntImm(var->dtype, 1), var_lanes);
   }
 
   Stmt VisitStmt(const Stmt& stmt) final {
diff --git a/src/tir/usmp/algo/greedy.cc b/src/tir/usmp/algo/greedy.cc
index c4f7cabb99f1..8246ffc219c6 100644
--- a/src/tir/usmp/algo/greedy.cc
+++ b/src/tir/usmp/algo/greedy.cc
@@ -61,7 +61,7 @@ size_t GreedyBase::round_up_to_byte_alignment(const size_t& non_aligned_byte_off
  */
 bool GreedyBase::IsValidPlacement(const PoolInfo& candidate_pool, const size_t& next_offset,
                                   const size_t& size_bytes) {
-  if (candidate_pool->size_hint_bytes == PoolInfo::kUnrestrictedPoolSizeHint) {
+  if (candidate_pool->size_hint_bytes == kUnrestrictedPoolSizeHint) {
     // this means pool is not bounded
     return true;
   }
diff --git a/src/tir/usmp/analysis/extract_buffer_info.cc b/src/tir/usmp/analysis/extract_buffer_info.cc
index fb4fb52c507e..bb959647c7f0 100644
--- a/src/tir/usmp/analysis/extract_buffer_info.cc
+++ b/src/tir/usmp/analysis/extract_buffer_info.cc
@@ -36,6 +36,8 @@
 
 #include <stack>
 
+#include "../../../runtime/thread_storage_scope.h"
+
 namespace tvm {
 namespace tir {
 namespace usmp {
@@ -257,18 +259,20 @@ void BufferInfoExtractor::RecordAllocateNodeInfo(const AllocateNode* op) {
 void BufferInfoExtractor::VisitStmt_(const AllocateNode* op) {
   ScopeInfo& current_scope_info = scope_stack_.top();
   const auto& type = Downcast<PointerType>(op->buffer_var->type_annotation);
-  const auto& storage_scope = type->storage_scope;
+  const auto& storage_scope = runtime::StorageScope::Create(type->storage_scope);
 
   // If the allocate is in a for loop, USMP currently only looks at serial for loops.
   // If its not a serial for loop, then memory planner will omit them in the current memory planning
   // process leaving them to as tir.allocate nodes for codegen. Additionally, the USMP can only work
   // with buffers that have global storage_scope
 
-  if (!current_scope_info.for_loop.defined()) {
-    RecordAllocateNodeInfo(op);
-  } else if (current_scope_info.for_loop.defined() &&
-             current_scope_info.for_loop->kind == ForKind::kSerial && storage_scope == "global") {
-    RecordAllocateNodeInfo(op);
+  if (storage_scope.rank == runtime::StorageRank::kGlobal) {
+    if (!current_scope_info.for_loop.defined()) {
+      RecordAllocateNodeInfo(op);
+    } else if (current_scope_info.for_loop.defined() &&
+               current_scope_info.for_loop->kind == ForKind::kSerial) {
+      RecordAllocateNodeInfo(op);
+    }
   }
   StmtExprVisitor::VisitStmt(op->body);
   current_scope_info.allocate_nodes.erase(GetRef<Allocate>(op));
diff --git a/src/tir/usmp/transform/assign_pool_info.cc b/src/tir/usmp/transform/assign_pool_info.cc
index 2386b5fef77d..a2304f3b9e3d 100644
--- a/src/tir/usmp/transform/assign_pool_info.cc
+++ b/src/tir/usmp/transform/assign_pool_info.cc
@@ -42,17 +42,13 @@ class PoolInfoAssigner : public StmtExprMutator {
  public:
   explicit PoolInfoAssigner(const IRModule& module) {
     PrimFunc main_func =
-        Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+        Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
     ICHECK(main_func.defined()) << "main function is not in the module";
     Optional<Target> target_host = main_func->GetAttr<Target>(tvm::attr::kTarget);
     ICHECK(target_host) << "main function does not have a target attr";
     WorkspaceMemoryPools workspace_pools =
         module->GetAttr<WorkspaceMemoryPools>(tvm::attr::kWorkspaceMemoryPools)
-            .value_or(WorkspaceMemoryPools({PoolInfo(
-                "global_workspace", {{target_host.value(), PoolInfo::kTargetPoolReadWriteAccess}},
-                PoolInfo::kUnrestrictedPoolSizeHint, PoolInfo::kUnknownClockFrequency,
-                PoolInfo::kUnknownReadBandwidth, PoolInfo::kUnknownWriteBandwidth, 0, 0,
-                {{target_host.value(), 1}}, Bool(true))}));
+            .value_or(WorkspaceMemoryPools({CreateDefaultMemoryPool(module)}));
     Array<PoolInfo> pool_infos = workspace_pools->pools;
     for (const PoolInfo& pool_info : pool_infos) {
       for (const auto& kv : pool_info->target_access) {
@@ -77,8 +73,24 @@ class PoolInfoAssigner : public StmtExprMutator {
   IRModule mod_;
   Map<String, Array<PoolInfo>> target_pool_infos_;
   PrimFunc func_;
+  PoolInfo CreateDefaultMemoryPool(const IRModule& module);
 };
 
+PoolInfo PoolInfoAssigner::CreateDefaultMemoryPool(const tvm::IRModule& module) {
+  Map<Target, String> target_access;
+  tir::PrimFunc tir_main_func =
+      Downcast<tir::PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
+  Target target_host = tir_main_func->GetAttr<Target>(tvm::attr::kTarget).value();
+  for (const auto& kv : module->functions) {
+    BaseFunc func = kv.second;
+    Optional<Target> target = func->GetAttr<Target>(tvm::attr::kTarget);
+    target_access.Set(target.value_or(target_host), kTargetPoolReadWriteAccess);
+  }
+  return PoolInfo("global_workspace", target_access, kUnrestrictedPoolSizeHint,
+                  kUnknownClockFrequency, kUnknownReadBandwidth, kUnknownWriteBandwidth, 0, 0, {},
+                  Bool(true));
+}
+
 Stmt PoolInfoAssigner::VisitStmt_(const AllocateNode* op) {
   Optional<Target> tgt = func_->GetAttr<Target>(tvm::attr::kTarget).value();
   ICHECK(tgt) << "The following PrimFunc does not have a target attr: \n" << func_;
diff --git a/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc b/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
index cd797681d474..6abc48c31be0 100644
--- a/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
+++ b/src/tir/usmp/transform/convert_pool_allocations_to_offsets.cc
@@ -189,7 +189,7 @@ PoolAllocationToOffsetConverter::ScopeInfo PoolAllocationToOffsetConverter::Upda
     si.params.push_back(pool_var);
     si.pools_to_params.Set(pool_info, pool_var);
     si.allocated_pool_params.push_back(AllocatedPoolInfo(
-        allocated_pool_info->pool_info, allocated_pool_info->allocated_size, pool_var));
+        allocated_pool_info->pool_info, allocated_pool_info->allocated_size, si.params.size() - 1));
 
     int pool_size = all_pools_sizes_[pool_info];
     String buffer_var_name = pool_ref_name + "_buffer_var";
@@ -258,7 +258,7 @@ Array<PrimExpr> PoolAllocationToOffsetConverter::ReplaceAllocateArgsWithLetArgs(
         allocate_buf_to_let_var_.find(Downcast<Var>(arg)) != allocate_buf_to_let_var_.end()) {
       ret.push_back(allocate_buf_to_let_var_[Downcast<Var>(arg)]);
     } else {
-      ret.push_back(arg);
+      ret.push_back(VisitExpr(arg));
     }
   }
   return ret;
@@ -331,7 +331,7 @@ PrimExpr PoolAllocationToOffsetConverter::VisitExpr_(const LoadNode* op) {
 }
 
 IRModule PoolAllocationToOffsetConverter::operator()() {
-  GlobalVar gv = module_->GetGlobalVar(::tvm::runtime::symbol::tvm_run_func_suffix);
+  GlobalVar gv = module_->GetGlobalVar(::tvm::runtime::symbol::tvm_module_main);
   PrimFunc main_func = Downcast<PrimFunc>(module_->Lookup(gv));
   ScopeInfo si = UpdateFunctionScopeInfo(main_func);
   this->scope_stack.push(si);
diff --git a/src/tir/usmp/unified_static_memory_planner.cc b/src/tir/usmp/unified_static_memory_planner.cc
index 3b941d3cc021..e848440f029e 100644
--- a/src/tir/usmp/unified_static_memory_planner.cc
+++ b/src/tir/usmp/unified_static_memory_planner.cc
@@ -51,7 +51,7 @@ static std::unordered_map<String, std::function<Map<BufferInfo, PoolAllocation>(
 
 IRModule PlanMemory(const IRModule& mod, String algo) {
   VLOG(1) << "workspace required = " << CalculateModuleWorkspaceSize(mod);
-  PrimFunc main_func = Downcast<PrimFunc>(mod->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+  PrimFunc main_func = Downcast<PrimFunc>(mod->Lookup(::tvm::runtime::symbol::tvm_module_main));
   BufferInfoAnalysis buffer_info_analysis = ExtractBufferInfo(main_func, mod);
   Array<BufferInfo> buffer_info_arr =
       CreateArrayBufferInfo(buffer_info_analysis->buffer_info_stmts);
@@ -63,7 +63,7 @@ IRModule PlanMemory(const IRModule& mod, String algo) {
       buffer_info_analysis->buffer_info_stmts, buffer_info_pool_allocations);
   IRModule ret = transform::ConvertPoolAllocationsToOffsets(stmt_pool_allocations)(mod);
   tir::PrimFunc tir_main_func =
-      Downcast<tir::PrimFunc>(ret->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+      Downcast<tir::PrimFunc>(ret->Lookup(::tvm::runtime::symbol::tvm_module_main));
   Optional<Array<tir::usmp::AllocatedPoolInfo>> allocated_pool_infos =
       tir_main_func->GetAttr<Array<tir::usmp::AllocatedPoolInfo>>(tvm::attr::kPoolArgs);
   if (allocated_pool_infos) {
diff --git a/src/tir/usmp/utils.cc b/src/tir/usmp/utils.cc
index 64428fd9c49e..03fac325905c 100644
--- a/src/tir/usmp/utils.cc
+++ b/src/tir/usmp/utils.cc
@@ -93,7 +93,6 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
                 << ",\n  memory_pressure=" << node->memory_pressure << ")";
     });
 
-
 PoolAllocation::PoolAllocation(PoolInfo pool_info, Integer byte_offset) {
   auto pool_allocation_node = make_object<PoolAllocationNode>();
   pool_allocation_node->pool_info = pool_info;
@@ -115,12 +114,13 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
                 << ")";
     });
 
-AllocatedPoolInfo::AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var) {
+AllocatedPoolInfo::AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size,
+                                     Integer pool_var_idx) {
   auto allocated_poolinfo_node = make_object<AllocatedPoolInfoNode>();
   allocated_poolinfo_node->pool_info = pool_info;
   allocated_poolinfo_node->allocated_size = allocated_size;
-  if (pool_var.defined()) {
-    allocated_poolinfo_node->pool_var = pool_var;
+  if (pool_var_idx.defined()) {
+    allocated_poolinfo_node->pool_var_idx = pool_var_idx;
   }
   data_ = std::move(allocated_poolinfo_node);
 }
@@ -181,7 +181,7 @@ class ModuleWorkspaceSizeCalculator : public StmtExprVisitor {
     for (const auto& gv_func : mod_->functions) {
       functions_.Set(gv_func.first->name_hint, Downcast<PrimFunc>(gv_func.second));
     }
-    main_func_ = Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_run_func_suffix));
+    main_func_ = Downcast<PrimFunc>(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
     ICHECK(main_func_.defined()) << "main function is not in the module";
     Optional<Target> target_host = main_func_->GetAttr<Target>(tvm::attr::kTarget);
     ICHECK(target_host) << "main function does not have a target attr";
diff --git a/tests/cpp/aot_metadata_test.cc b/tests/cpp/aot_metadata_test.cc
new file mode 100644
index 000000000000..37682a423f2f
--- /dev/null
+++ b/tests/cpp/aot_metadata_test.cc
@@ -0,0 +1,281 @@
+
+/*
+ * 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.
+ */
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <tvm/ir/type.h>
+#include <tvm/node/reflection.h>
+#include <tvm/runtime/metadata.h>
+
+#include "../src/target/metadata.h"
+
+namespace {
+
+const int64_t kNormalInput1Shape[4] = {1, 5, 5, 3};
+const struct TVMTensorInfo kNormalInputs[2] = {
+    {"input1", kNormalInput1Shape, 4, DLDataType{1, 2, 3}},
+    {"input2", kNormalInput1Shape, 4, DLDataType{2, 3, 4}}};
+
+const int64_t kNormalOutput1Shape[3] = {3, 8, 8};
+const struct TVMTensorInfo kNormalOutputs[1] = {
+    {"output1", kNormalOutput1Shape, 3, DLDataType{3, 4, 5}}};
+
+const int64_t kNormalPool1Shape[3] = {3, 8, 8};
+const struct TVMTensorInfo kNormalPools[1] = {{"pool1", kNormalPool1Shape, 3, DLDataType{3, 4, 7}}};
+
+const struct TVMMetadata kNormal = {
+    TVM_METADATA_VERSION, kNormalInputs, 2, kNormalOutputs, 1, kNormalPools, 1, "default",
+};
+}  // namespace
+
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::StrEq;
+using ::tvm::runtime::Downcast;
+
+TEST(Metadata, ParseStruct) {
+  tvm::runtime::metadata::Metadata md = tvm::runtime::metadata::Metadata(&kNormal);
+  EXPECT_THAT(md->version(), Eq(TVM_METADATA_VERSION));
+  EXPECT_THAT(md->num_inputs(), Eq(2));
+
+  auto inputs = md->inputs();
+  EXPECT_THAT(inputs.size(), Eq(2));
+
+  auto input1 = inputs[0];
+  EXPECT_THAT(input1->name(), Eq("input1"));
+  EXPECT_THAT(input1->shape(), ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(input1->dtype(), Eq(tvm::runtime::DataType(DLDataType{1, 2, 3})));
+
+  auto input2 = inputs[1];
+  EXPECT_THAT(input2->name(), Eq("input2"));
+  EXPECT_THAT(input2->shape(), ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(input2->dtype(), Eq(tvm::runtime::DataType(DLDataType{2, 3, 4})));
+
+  EXPECT_THAT(md->num_outputs(), Eq(1));
+  auto outputs = md->outputs();
+  EXPECT_THAT(outputs.size(), Eq(1));
+
+  auto output1 = outputs[0];
+  EXPECT_THAT(output1->name(), Eq("output1"));
+  EXPECT_THAT(output1->shape(), ElementsAre(3, 8, 8));
+  EXPECT_THAT(output1->dtype(), Eq(tvm::runtime::DataType(DLDataType{3, 4, 5})));
+
+  auto pools = md->pools();
+  EXPECT_THAT(pools.size(), Eq(1));
+
+  auto pool1 = pools[0];
+  EXPECT_THAT(pool1->name(), Eq("pool1"));
+  EXPECT_THAT(pool1->shape(), ElementsAre(3, 8, 8));
+  EXPECT_THAT(pool1->dtype(), Eq(tvm::runtime::DataType(DLDataType{3, 4, 7})));
+
+  EXPECT_THAT(md->mod_name(), Eq("default"));
+}
+
+class TestVisitor : public tvm::AttrVisitor {
+ public:
+  using Element = ::std::tuple<::std::string, ::tvm::runtime::ObjectRef>;
+  void Visit(const char* key, double* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::FloatImm(::tvm::runtime::DataType(kDLFloat, 64, 1), *value));
+  }
+  void Visit(const char* key, int64_t* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::IntImm(::tvm::runtime::DataType(kDLInt, 64, 1), *value));
+  }
+  void Visit(const char* key, uint64_t* value) final {
+    keys.push_back(key);
+    int64_t v;
+    *(reinterpret_cast<uint64_t*>(&v)) = *value;
+    values.push_back(::tvm::IntImm(::tvm::runtime::DataType(kDLUInt, 64, 1), v));
+  }
+  void Visit(const char* key, int* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::IntImm(::tvm::runtime::DataType(kDLInt, 64, 1), *value));
+  }
+  void Visit(const char* key, bool* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::Bool(*value));
+  }
+  void Visit(const char* key, std::string* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::runtime::String(*value));
+  }
+  void Visit(const char* key, tvm::runtime::DataType* value) final {
+    keys.push_back(key);
+    values.push_back(::tvm::PrimType(*value));
+  }
+  void Visit(const char* key, tvm::runtime::NDArray* value) final {
+    keys.push_back(key);
+    values.push_back(*value);
+  }
+  void Visit(const char* key, void** value) final { CHECK(false) << "Do not expect this type"; }
+
+  void Visit(const char* key, ::tvm::runtime::ObjectRef* value) final {
+    keys.push_back(key);
+    values.push_back(*value);
+  }
+
+  std::vector<std::string> keys;
+  std::vector<::tvm::runtime::ObjectRef> values;
+};
+
+TEST(Metadata, Visitor) {
+  tvm::runtime::metadata::Metadata md = tvm::runtime::metadata::Metadata(&kNormal);
+  TestVisitor v;
+  ::tvm::ReflectionVTable::Global()->VisitAttrs(md.operator->(), &v);
+
+  EXPECT_THAT(v.keys, ElementsAre(StrEq("version"), StrEq("inputs"), StrEq("num_inputs"),
+                                  StrEq("outputs"), StrEq("num_outputs"), StrEq("pools"),
+                                  StrEq("num_pools"), StrEq("mod_name")));
+  EXPECT_THAT(Downcast<tvm::IntImm>(v.values[0])->value, Eq(TVM_METADATA_VERSION));
+
+  EXPECT_THAT(Downcast<tvm::IntImm>(v.values[0])->value, Eq(TVM_METADATA_VERSION));
+
+  // Just identify the tensor.
+  auto input_array = Downcast<tvm::runtime::metadata::MetadataArray>(v.values[1]);
+  EXPECT_THAT(input_array->type_index, Eq(tvm::runtime::metadata::MetadataTypeIndex::kMetadata));
+  EXPECT_THAT(input_array->struct_name, StrEq("TVMTensorInfo"));
+  EXPECT_THAT(input_array->array.size(), Eq(2));
+
+  auto input1 = Downcast<tvm::runtime::metadata::TensorInfo>(input_array->array[0]);
+  EXPECT_THAT(input1->name(), StrEq("input1"));
+  EXPECT_THAT(input1->shape(), ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(input1->dtype(), tvm::runtime::DataType(DLDataType{1, 2, 3}));
+
+  auto input2 = Downcast<tvm::runtime::metadata::TensorInfo>(input_array->array[1]);
+  EXPECT_THAT(input1->name(), StrEq("input1"));
+  EXPECT_THAT(input1->shape(), ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(input1->dtype(), tvm::runtime::DataType(DLDataType{1, 2, 3}));
+
+  auto num_inputs = Downcast<tvm::IntImm>(v.values[2]);
+  EXPECT_THAT(num_inputs->value, Eq(2));
+
+  auto output_array = Downcast<tvm::runtime::metadata::MetadataArray>(v.values[3]);
+  EXPECT_THAT(output_array->type_index, Eq(tvm::runtime::metadata::MetadataTypeIndex::kMetadata));
+  EXPECT_THAT(output_array->struct_name, StrEq("TVMTensorInfo"));
+  auto output1 = Downcast<tvm::runtime::metadata::TensorInfo>(output_array->array[0]);
+
+  EXPECT_THAT(output1->name(), Eq("output1"));
+
+  auto num_outputs = Downcast<tvm::IntImm>(v.values[4]);
+  EXPECT_THAT(num_outputs->value, Eq(1));
+
+  auto pool_array = Downcast<tvm::runtime::metadata::MetadataArray>(v.values[5]);
+  EXPECT_THAT(pool_array->type_index, Eq(tvm::runtime::metadata::MetadataTypeIndex::kMetadata));
+  EXPECT_THAT(pool_array->struct_name, StrEq("TVMTensorInfo"));
+  auto pool1 = Downcast<tvm::runtime::metadata::TensorInfo>(pool_array->array[0]);
+
+  EXPECT_THAT(pool1->name(), Eq("pool1"));
+
+  auto num_pools = Downcast<tvm::IntImm>(v.values[6]);
+  EXPECT_THAT(num_pools->value, Eq(1));
+
+  auto mod_name = Downcast<tvm::String>(v.values[7]);
+  EXPECT_THAT(mod_name, Eq("default"));
+}
+
+using ::tvm::runtime::make_object;
+TEST(Metadata, InMemory) {
+  tvm::runtime::metadata::Metadata md =
+      tvm::runtime::metadata::Metadata(make_object<tvm::target::metadata::InMemoryMetadataNode>(
+          TVM_METADATA_VERSION,
+          std::vector<tvm::runtime::metadata::TensorInfo>(
+              {tvm::runtime::metadata::TensorInfo(
+                   make_object<tvm::target::metadata::InMemoryTensorInfoNode>(
+                       tvm::String("Input1"), std::vector<int64_t>{1, 5, 5, 3},
+                       tvm::runtime::DataType(DLDataType{1, 2, 3}))),
+               tvm::runtime::metadata::TensorInfo(
+                   make_object<tvm::target::metadata::InMemoryTensorInfoNode>(
+                       tvm::String("Input2"), std::vector<int64_t>{1, 5, 5, 3},
+                       tvm::runtime::DataType(DLDataType{2, 3, 4})))}),
+          std::vector<tvm::runtime::metadata::TensorInfo>({tvm::runtime::metadata::TensorInfo(
+              make_object<tvm::target::metadata::InMemoryTensorInfoNode>(
+                  tvm::String("Output1"), std::vector<int64_t>{3, 8, 8},
+                  tvm::runtime::DataType(DLDataType{3, 4, 5})))}),
+          std::vector<tvm::runtime::metadata::TensorInfo>({tvm::runtime::metadata::TensorInfo(
+              make_object<tvm::target::metadata::InMemoryTensorInfoNode>(
+                  tvm::String("Pool1"), std::vector<int64_t>{5, 10, 10},
+                  tvm::runtime::DataType(DLDataType{3, 4, 7})))}),
+          "default"));
+
+  auto md_data = md->data();
+  EXPECT_THAT(md_data->version, Eq(TVM_METADATA_VERSION));
+  EXPECT_THAT(md_data->num_inputs, Eq(2));
+
+  auto input0 = &md_data->inputs[0];
+  EXPECT_THAT(input0->name, StrEq("Input1"));
+  EXPECT_THAT(std::vector<int64_t>(input0->shape, input0->shape + input0->num_shape),
+              ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(tvm::runtime::DataType(input0->dtype),
+              Eq(tvm::runtime::DataType(DLDataType({1, 2, 3}))));
+
+  auto input1 = &md_data->inputs[1];
+  EXPECT_THAT(input1->name, StrEq("Input2"));
+  EXPECT_THAT(std::vector<int64_t>(input1->shape, input1->shape + input1->num_shape),
+              ElementsAre(1, 5, 5, 3));
+  EXPECT_THAT(tvm::runtime::DataType(input1->dtype),
+              Eq(tvm::runtime::DataType(DLDataType({2, 3, 4}))));
+
+  auto output0 = &md_data->outputs[0];
+  EXPECT_THAT(output0->name, StrEq("Output1"));
+  EXPECT_THAT(std::vector<int64_t>(output0->shape, output0->shape + output0->num_shape),
+              ElementsAre(3, 8, 8));
+  EXPECT_THAT(tvm::runtime::DataType(output0->dtype),
+              Eq(tvm::runtime::DataType(DLDataType({3, 4, 5}))));
+
+  auto pool0 = &md_data->pools[0];
+  EXPECT_THAT(pool0->name, StrEq("Pool1"));
+  EXPECT_THAT(std::vector<int64_t>(pool0->shape, pool0->shape + pool0->num_shape),
+              ElementsAre(5, 10, 10));
+  EXPECT_THAT(tvm::runtime::DataType(pool0->dtype),
+              Eq(tvm::runtime::DataType(DLDataType({3, 4, 7}))));
+
+  EXPECT_THAT(md_data->mod_name, StrEq("default"));
+}
+
+TEST(Metadata, ZeroElementLists) {
+  tvm::runtime::metadata::Metadata md =
+      tvm::runtime::metadata::Metadata(make_object<tvm::target::metadata::InMemoryMetadataNode>(
+          TVM_METADATA_VERSION, std::vector<tvm::runtime::metadata::TensorInfo>({}),
+          std::vector<tvm::runtime::metadata::TensorInfo>({tvm::runtime::metadata::TensorInfo(
+              make_object<tvm::target::metadata::InMemoryTensorInfoNode>(
+                  tvm::String("Output1"), std::vector<int64_t>{},
+                  tvm::runtime::DataType(DLDataType{3, 4, 5})))}),
+          std::vector<tvm::runtime::metadata::TensorInfo>({}), "default"));
+
+  EXPECT_THAT(md->data()->num_inputs, Eq(0));
+  EXPECT_THAT(md->inputs().size(), Eq(0));
+  EXPECT_THAT(md->num_inputs(), Eq(0));
+  EXPECT_THAT(md->inputs(), ElementsAre());
+
+  auto output0 = md->data()->outputs[0];
+  EXPECT_THAT(output0.num_shape, Eq(0));
+  EXPECT_THAT(md->outputs()[0]->shape().size(), Eq(0));
+  EXPECT_THAT(md->outputs()[0]->shape(), ElementsAre());
+
+  EXPECT_THAT(md->pools().size(), Eq(0));
+  EXPECT_THAT(md->num_pools(), Eq(0));
+  EXPECT_THAT(md->pools(), ElementsAre());
+}
+
+TEST(Metadata, GetCStructName) {
+  EXPECT_THAT(tvm::runtime::metadata::get_c_struct_name<MetadataNode>(), EqStr("TVMMetadata"));
+  EXPECT_THAT(tvm::runtime::metadata::get_c_struct_name<TensorInfoNode>(), EqStr("TVMTensorInfo"));
+}
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index 97809b0e1398..d02c38f3afac 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -60,8 +60,9 @@ TEST(IRF, VisitPrimFuncs) {
   using namespace tvm;
   using namespace tvm::tir;
   PrimFunc prim_func(/*params=*/{}, /*body=*/Evaluate(Integer(0)));
-  relay::Function relay_func(/*params=*/{}, /*body=*/relay::Expr(nullptr),
-                             /*ret_type=*/relay::Type{nullptr}, /*ty_params=*/{});
+  auto c_data = tvm::runtime::NDArray::Empty({1, 2, 3}, {kDLFloat, 32, 1}, {kDLCPU, 0});
+  relay::Function relay_func(/*params=*/{}, /*body=*/relay::Expr(relay::Constant(c_data)),
+                             /*ret_type=*/relay::Type(), /*ty_params=*/{});
   IRModule mod({
       {GlobalVar("main"), prim_func},
       {GlobalVar("main2"), relay_func},
diff --git a/tests/lint/check_asf_header.sh b/tests/lint/check_asf_header.sh
index cd5fe6474c76..579069bb5ba5 100755
--- a/tests/lint/check_asf_header.sh
+++ b/tests/lint/check_asf_header.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -euo pipefail
 
 rat_tempdir="$(mktemp -d)"
 
@@ -26,10 +27,23 @@ trap cleanup EXIT
 
 rat_output="${rat_tempdir}/$$.apache-rat.txt"
 
+fix=0
 filter_untracked=0
-if [ "$1" == "--local" ]; then
-    filter_untracked=1
-fi
+while [ "x${1-}" != "x" ]; do
+    case "$1" in
+        --local)
+            filter_untracked=1
+            ;;
+        --fix)
+            fix=1
+            ;;
+        *)
+            echo "Unrecognized option: $1"
+            exit 2
+            ;;
+    esac
+    shift
+done
 
 java -jar /bin/apache-rat.jar -E tests/lint/rat-excludes  -d . | (grep -E "^== File" >"${rat_output}" || true)
 
@@ -43,14 +57,26 @@ if [ ${filter_untracked} -eq 1 ]; then
     rat_output="${processed_rat_output}"
 fi
 
-if grep --quiet -E "File" "${rat_output}"; then
-    echo "Need to add ASF header to the following files."
+grep --quiet -E "File" "${rat_output}" && success=1 || success=0
+
+if [ $success -eq 0 ]; then
+    echo "No files violate ASF header check"
+else
+    if [ $fix -eq 1 ]; then
+        python3 tests/lint/add_asf_header.py "${rat_output}"
+        echo "Ran add_asf_header.py to fix the following files:"
+    else
+        echo "Need to add ASF header to the following files."
+    fi
     echo "----------------File List----------------"
     cat "${rat_output}"
     echo "-----------------------------------------"
-    echo "Use the following steps to add the headers:"
-    echo "- Create file_list.txt in your text editor"
-    echo "- Copy paste the above content in file-list into file_list.txt"
-    echo "- python3 tests/lint/add_asf_header.py file_list.txt"
+    if [ $fix -eq 0 ]; then
+        echo "To fix, run check_asf_header.sh --fix, or use the following manual steps:"
+        echo "- Create file_list.txt in your text editor"
+        echo "- Copy paste the above content in file-list into file_list.txt"
+        echo "- python3 tests/lint/add_asf_header.py file_list.txt"
+        echo "When running via docker/lint.sh, add -i to automatically fix these problems."
+    fi
     exit 1
 fi
diff --git a/tests/lint/clang_format.sh b/tests/lint/clang_format.sh
index ad6a35bb4104..f131caff2251 100755
--- a/tests/lint/clang_format.sh
+++ b/tests/lint/clang_format.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -e
 
 # check lastest change, for squash merge into main
 ./tests/lint/git-clang-format.sh HEAD~1
diff --git a/tests/lint/cppdocs.sh b/tests/lint/cppdocs.sh
index d29b6cdd236b..af5a98c5139e 100755
--- a/tests/lint/cppdocs.sh
+++ b/tests/lint/cppdocs.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
diff --git a/tests/lint/cpplint.sh b/tests/lint/cpplint.sh
index 31eb1d94a347..687f4ab6fe43 100755
--- a/tests/lint/cpplint.sh
+++ b/tests/lint/cpplint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -e
 
 python3 3rdparty/dmlc-core/scripts/lint.py vta cpp vta/include vta/src
 python3 3rdparty/dmlc-core/scripts/lint.py tvm cpp \
diff --git a/tests/lint/flake8.sh b/tests/lint/flake8.sh
index b16d97fce361..63dad38166b1 100755
--- a/tests/lint/flake8.sh
+++ b/tests/lint/flake8.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,5 +16,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -e
 
 python3 -m flake8 . --count --select=E9,F63,F7 --show-source --statistics
diff --git a/tests/lint/git-black.sh b/tests/lint/git-black.sh
index 4052c052463e..68f9b13951c3 100755
--- a/tests/lint/git-black.sh
+++ b/tests/lint/git-black.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/lint/git-clang-format.sh b/tests/lint/git-clang-format.sh
index 37d05320919c..8c9e6c5b475b 100755
--- a/tests/lint/git-clang-format.sh
+++ b/tests/lint/git-clang-format.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/lint/jnilint.sh b/tests/lint/jnilint.sh
index 39b80489d65e..ec359a5b89ae 100755
--- a/tests/lint/jnilint.sh
+++ b/tests/lint/jnilint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,5 +16,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -e
 
 python3 3rdparty/dmlc-core/scripts/lint.py tvm4j-jni cpp jvm/native/src
diff --git a/tests/lint/pylint.sh b/tests/lint/pylint.sh
index 2cb38802bf90..6c958a923139 100755
--- a/tests/lint/pylint.sh
+++ b/tests/lint/pylint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
diff --git a/tests/lint/python_format.sh b/tests/lint/python_format.sh
index 2e907f8be9e1..21dd4ca7cca6 100755
--- a/tests/lint/python_format.sh
+++ b/tests/lint/python_format.sh
@@ -1,4 +1,4 @@
-#!/bin/bash -e
+#!/usr/bin/env bash
 # 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
@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+set -e
 
-./tests/lint/git-black.sh HEAD~1
-./tests/lint/git-black.sh origin/main
+./tests/lint/git-black.sh -i HEAD~1
+./tests/lint/git-black.sh -i origin/main
diff --git a/tests/lint/rust_format.sh b/tests/lint/rust_format.sh
index 10c8feec1fcf..bed7ad976ea6 100755
--- a/tests/lint/rust_format.sh
+++ b/tests/lint/rust_format.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/micro/zephyr/test_utils.py b/tests/micro/zephyr/test_utils.py
index 846dabee617c..ea17ac9a3531 100644
--- a/tests/micro/zephyr/test_utils.py
+++ b/tests/micro/zephyr/test_utils.py
@@ -210,7 +210,7 @@ def generate_project(
                         model_files_path, arcname=os.path.relpath(model_files_path, tar_temp_dir)
                     )
                 header_path = generate_c_interface_header(
-                    lowered.libmod_name, ["input_1"], ["output"], [], [], 0, model_files_path
+                    lowered.libmod_name, ["input_1"], ["Identity"], [], [], 0, model_files_path
                 )
                 tf.add(header_path, arcname=os.path.relpath(header_path, tar_temp_dir))
 
diff --git a/tests/micro/zephyr/test_zephyr.py b/tests/micro/zephyr/test_zephyr.py
index 4034008c6d45..1582d7e4a5fe 100644
--- a/tests/micro/zephyr/test_zephyr.py
+++ b/tests/micro/zephyr/test_zephyr.py
@@ -376,6 +376,7 @@ def test_tensors(sess):
         test_tensors(sess)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10297")
 @tvm.testing.requires_micro
 def test_autotune_conv2d(temp_dir, board, west_cmd, tvm_debug):
     """Test AutoTune for microTVM Zephyr"""
diff --git a/tests/python/contrib/test_cmsisnn/test_conv2d.py b/tests/python/contrib/test_cmsisnn/test_conv2d.py
index ddcb2c10bff7..641fd517a4e7 100644
--- a/tests/python/contrib/test_cmsisnn/test_conv2d.py
+++ b/tests/python/contrib/test_cmsisnn/test_conv2d.py
@@ -213,6 +213,7 @@ def test_conv2d_symmetric_padding_int8(
     )
 
 
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/10314")
 @tvm.testing.requires_cmsisnn
 @pytest.mark.parametrize("padding", ["SAME", "VALID"])
 @pytest.mark.parametrize("relu_type", ["RELU", "NONE"])
@@ -302,6 +303,7 @@ def test_conv2d_asymmetric_padding_int8(
     )
 
 
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/10314")
 @tvm.testing.requires_cmsisnn
 @pytest.mark.parametrize("ifm_shape", [(1, 28, 28, 12), (1, 64, 100, 4)])
 @pytest.mark.parametrize("kernel_size", [(3, 3)])
diff --git a/tests/python/contrib/test_dnnl.py b/tests/python/contrib/test_dnnl.py
index 7adf3e40ad33..4d1972d6a3b0 100755
--- a/tests/python/contrib/test_dnnl.py
+++ b/tests/python/contrib/test_dnnl.py
@@ -14,7 +14,6 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import numpy as np
 import pytest
 import itertools
 import tvm
@@ -22,6 +21,7 @@
 from tvm import relay
 from tvm.relay.op.contrib import dnnl
 import tvm.testing
+import numpy as np
 
 has_dnnl_codegen = pytest.mark.skipif(
     not tvm.get_global_func("relay.ext.dnnl", True), reason="DNNL codegen not available"
@@ -65,6 +65,7 @@ def check_dnnl_used(mod):
             result_key = mode + ("_dnnl" if use_dnnl else "")
             if use_dnnl:
                 mod = dnnl.partition_for_dnnl(mod, params)
+                check_dnnl_used(mod)
             with tvm.transform.PassContext(opt_level=3):
                 func = relay.create_executor(mode, mod=mod, device=dev, target=target).evaluate()
             if run_module:
@@ -99,6 +100,79 @@ def run_and_verify_func(config, run_module, target="llvm", dtype="float32"):
     run_and_verify(f, input_dict, params, target, run_module)
 
 
+def get_conv1d(
+    x_shape=((1, 3, 224)),
+    k_shape=(10, 3, 3),
+    groups=1,
+    padding=(1, 1),
+    strides=(1),
+    dilation=(1),
+    channels=None,
+    activation=None,
+    dtype="float32",
+):
+    x = relay.var("x", shape=(x_shape), dtype=dtype)
+    kernel = relay.var("kernel", shape=(k_shape), dtype=dtype)
+    out = relay.nn.conv1d(
+        x,
+        kernel,
+        kernel_size=k_shape[2:3],
+        groups=groups,
+        padding=padding,
+        strides=strides,
+        dilation=dilation,
+        channels=k_shape[0],
+    )
+    dic = {"x": x_shape, "kernel": k_shape}
+    param_lst = ["kernel"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
+def get_conv1d_bias(x_shape=(1, 3, 224), k_shape=(10, 3, 3), activation=None, dtype="float32"):
+    conv, dic, param_lst = get_conv1d(x_shape=x_shape, k_shape=k_shape, dtype=dtype)
+    bias = relay.var("bias", shape=(k_shape[0],), dtype=dtype)
+    out = relay.nn.bias_add(conv, bias)
+    dic["bias"] = (k_shape[0],)
+    param_lst += ["bias"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
+def get_conv1d_bias_bn_relu(x_shape=(1, 3, 224), k_shape=(10, 3, 3), dtype="float32"):
+    conv1d_bias, dic, param_lst = get_conv1d_bias(x_shape, k_shape, dtype=dtype)
+    beta = relay.const(np.zeros(k_shape[0]).astype(dtype))
+    gamma = relay.const(np.ones(k_shape[0]).astype(dtype))
+    moving_mean = relay.const(np.zeros(k_shape[0]).astype(dtype))
+    moving_var = relay.const(np.ones(k_shape[0]).astype(dtype))
+    conv1d_bias_bn, _, _ = relay.nn.batch_norm(
+        conv1d_bias,
+        gamma=gamma,
+        beta=beta,
+        moving_mean=moving_mean,
+        moving_var=moving_var,
+        axis=1,
+        center=True,
+        scale=True,
+        epsilon=1e-5,
+    )
+    return relay.nn.relu(conv1d_bias_bn), dic, param_lst
+
+
 def get_conv2d(
     x_shape=(1, 32, 8, 8),
     k_shape=(16, 32, 3, 3),
@@ -134,6 +208,39 @@ def get_conv2d(
         return out, dic, param_lst
 
 
+def get_conv2d_transpose(
+    x_shape=(1, 32, 8, 8),
+    k_shape=(32, 16, 3, 3),
+    groups=1,
+    padding=(0, 0),
+    strides=(1, 1),
+    activation=None,
+    dtype="float32",
+):
+    x = relay.var("x", shape=(x_shape), dtype=dtype)
+    kernel = relay.var("kernel", shape=(k_shape), dtype=dtype)
+    out = relay.nn.conv2d_transpose(
+        x,
+        kernel,
+        channels=k_shape[1],
+        kernel_size=k_shape[2:4],
+        groups=groups,
+        padding=padding,
+        strides=strides,
+    )
+    dic = {"x": x_shape, "kernel": k_shape}
+    param_lst = ["kernel"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
 def get_conv2d_weights_const(
     x_shape=(1, 32, 8, 8),
     k_shape=(16, 32, 3, 3),
@@ -179,6 +286,25 @@ def get_conv2d_bias(
         return out, dic, param_lst
 
 
+def get_conv2d_transpose_bias(
+    x_shape=(1, 32, 8, 8), k_shape=(32, 16, 3, 3), activation=None, dtype="float32"
+):
+    conv, dic, param_lst = get_conv2d_transpose(x_shape=x_shape, k_shape=k_shape, dtype=dtype)
+    bias = relay.var("bias", shape=(k_shape[1],), dtype=dtype)
+    out = relay.nn.bias_add(conv, bias)
+    dic["bias"] = (k_shape[1],)
+    param_lst += ["bias"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
 def get_conv2d_bias_bn_relu(x_shape=(1, 32, 8, 8), k_shape=(16, 32, 3, 3), dtype="float32"):
     conv2d_bias, dic, param_lst = get_conv2d_bias(x_shape, k_shape, dtype=dtype)
     beta = relay.const(np.zeros(k_shape[0]).astype(dtype))
@@ -199,6 +325,118 @@ def get_conv2d_bias_bn_relu(x_shape=(1, 32, 8, 8), k_shape=(16, 32, 3, 3), dtype
     return relay.nn.relu(conv2d_bias_bn), dic, param_lst
 
 
+def get_conv3d(
+    x_shape=(1, 32, 8, 8, 8),
+    k_shape=(16, 32, 3, 3, 3),
+    groups=1,
+    padding=(0, 0, 0),
+    strides=(1, 1, 1),
+    dilation=(1, 1, 1),
+    activation=None,
+    dtype="float32",
+):
+    x = relay.var("x", shape=(x_shape), dtype=dtype)
+    kernel = relay.var("kernel", shape=(k_shape), dtype=dtype)
+    out = relay.nn.conv3d(
+        x,
+        kernel,
+        channels=k_shape[0],
+        kernel_size=k_shape[2:],
+        groups=groups,
+        padding=padding,
+        strides=strides,
+        dilation=dilation,
+    )
+    dic = {"x": x_shape, "kernel": k_shape}
+    param_lst = ["kernel"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
+def get_conv3d_transpose(
+    x_shape=(1, 32, 8, 8, 8),
+    k_shape=(32, 16, 3, 3, 3),
+    groups=1,
+    padding=(0, 0, 0),
+    strides=(1, 1, 1),
+    output_padding=(0, 0, 0),
+    activation=None,
+    dtype="float32",
+    data_layout="NCDHW",
+    kernel_layout="OIDHW",
+):
+    x = relay.var("x", shape=(x_shape), dtype=dtype)
+    kernel = relay.var("kernel", shape=(k_shape), dtype=dtype)
+    out = relay.nn.conv3d_transpose(
+        x,
+        kernel,
+        channels=k_shape[1],
+        kernel_size=k_shape[2:5],
+        groups=groups,
+        padding=padding,
+        strides=strides,
+        output_padding=output_padding,
+        data_layout=data_layout,
+        kernel_layout=kernel_layout,
+    )
+    dic = {"x": x_shape, "kernel": k_shape}
+    param_lst = ["kernel"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
+def get_conv3d_bias(
+    x_shape=(1, 32, 8, 8, 8), k_shape=(16, 32, 3, 3, 3), activation=None, dtype="float32"
+):
+    conv, dic, param_lst = get_conv3d(x_shape=x_shape, k_shape=k_shape, dtype=dtype)
+    bias = relay.var("bias", shape=(k_shape[0],), dtype=dtype)
+    out = relay.nn.bias_add(conv, bias)
+    dic["bias"] = (k_shape[0],)
+    param_lst += ["bias"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
+def get_conv3d_transpose_bias(
+    x_shape=(1, 32, 8, 8, 8), k_shape=(32, 16, 3, 3, 3), activation=None, dtype="float32"
+):
+    conv, dic, param_lst = get_conv3d_transpose(x_shape=x_shape, k_shape=k_shape, dtype=dtype)
+    bias = relay.var("bias", shape=(k_shape[1],), dtype=dtype)
+    out = relay.nn.bias_add(conv, bias)
+    dic["bias"] = (k_shape[1],)
+    param_lst += ["bias"]
+
+    if activation == "relu":
+        return relay.nn.relu(out), dic, param_lst
+    elif activation == "tanh":
+        return relay.tanh(out), dic, param_lst
+    elif activation == "sigmoid":
+        return relay.sigmoid(out), dic, param_lst
+    else:
+        return out, dic, param_lst
+
+
 def get_dense(x_shape=(1, 16), k_shape=(32, 16), activation=None, dtype="float32"):
     x = relay.var("x", shape=(x_shape), dtype=dtype)
     kernel = relay.var("kernel", shape=(k_shape), dtype=dtype)
@@ -249,14 +487,20 @@ def get_graph():
     run_and_verify_func(get_graph(), run_module=run_module, dtype=dtype)
 
 
-def test_unary(run_module):
+def test_elementwise(run_module, dtype="float32"):
     def get_graph(op, x_shape=(1, 8, 3, 3)):
-        x = relay.var("x", shape=(x_shape), dtype="float32")
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
         out = op(x)
         f = tvm.IRModule.from_expr(out)
         return f, {"x": x_shape}, []
 
     for op in [
+        relay.abs,
+        relay.exp,
+        relay.log,
+        relay.sqrt,
+        relay.round,
+        relay.logsumexp,
         relay.nn.relu,
         relay.tanh,
         relay.sigmoid,
@@ -264,6 +508,62 @@ def get_graph(op, x_shape=(1, 8, 3, 3)):
         run_and_verify_func(get_graph(op), run_module=run_module)
 
 
+def test_clip(run_module, dtype="float32"):
+    def get_graph(x_shape=(1, 8, 3, 3)):
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
+        out = relay.clip(x, a_min=-0.2, a_max=0.4)
+        f = tvm.IRModule.from_expr(out)
+        return f, {"x": x_shape}, []
+
+    run_and_verify_func(get_graph(), run_module=run_module)
+
+
+def test_leaky_relu(run_module, dtype="float32"):
+    def get_graph(x_shape=(1, 8, 3, 3)):
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
+        out = relay.nn.leaky_relu(x, alpha=0.1)
+        f = tvm.IRModule.from_expr(out)
+        return f, {"x": x_shape}, []
+
+    run_and_verify_func(get_graph(), run_module=run_module)
+
+
+def test_softmax(run_module, dtype="float32"):
+    def get_graph(x_shape, axis):
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
+        out = relay.nn.softmax(x, axis=axis)
+        f = tvm.IRModule.from_expr(out)
+        return f, {"x": x_shape}, []
+
+    run_and_verify_func(get_graph((1, 1000), axis=1), run_module=run_module)
+    run_and_verify_func(get_graph((1, 1000), axis=-1), run_module=run_module)
+    run_and_verify_func(get_graph((1, 3, 4), axis=-2), run_module=run_module)
+    run_and_verify_func(get_graph((1, 3, 4), axis=1), run_module=run_module)
+
+
+def test_conv1d(run_module, dtype="float32"):
+    conv1d, dic, param_lst = get_conv1d(channels=10, dtype=dtype)
+    conv1d = tvm.IRModule.from_expr(conv1d)
+    config = conv1d, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv1d_pattern(run_module, dtype="float32"):
+    x_shape = (1, 3, 224)
+    k_shape = (10, 3, 3)
+    activation_lst = [None, "relu", "tanh", "sigmoid"]
+    for a in activation_lst:
+        conv1d, dic, param_lst = get_conv1d(x_shape, k_shape, activation=a, dtype=dtype)
+        conv1d = tvm.IRModule.from_expr(conv1d)
+        config = conv1d, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+        conv1d_bias, dic, param_lst = get_conv1d_bias(x_shape, k_shape, activation=a, dtype=dtype)
+        conv1d_bias = tvm.IRModule.from_expr(conv1d_bias)
+        config = conv1d_bias, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
 def test_conv2d(run_module, dtype="float32"):
     x_shape = (1, 32, 8, 8)
     for k_shape, groups in [((16, 32, 3, 3), 1), ((32, 1, 3, 3), 32)]:
@@ -314,6 +614,90 @@ def test_conv2d_pattern(run_module, dtype="float32"):
     run_and_verify_func(config, run_module=run_module, dtype=dtype)
 
 
+def test_conv2d_transpose(run_module, dtype="float32"):
+    for padding in [(0, 0), (1, 1)]:
+        for strides in [(1, 1), (2, 2)]:
+            conv2d_transpose, dic, param_lst = get_conv2d_transpose(
+                padding=padding, strides=strides, dtype=dtype
+            )
+            conv2d_transpose = tvm.IRModule.from_expr(conv2d_transpose)
+            config = conv2d_transpose, dic, param_lst
+            run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv2d_transpose_pattern(run_module, dtype="float32"):
+    activation_lst = [None, "relu", "tanh", "sigmoid"]
+    for a in activation_lst:
+        conv2d, dic, param_lst = get_conv2d_transpose(activation=a, dtype=dtype)
+        conv2d = tvm.IRModule.from_expr(conv2d)
+        config = conv2d, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+        conv2d_bias, dic, param_lst = get_conv2d_transpose_bias(activation=a, dtype=dtype)
+        conv2d_bias = tvm.IRModule.from_expr(conv2d_bias)
+        config = conv2d_bias, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv3d(run_module, dtype="float32"):
+    conv3d, dic, param_lst = get_conv3d(dtype=dtype)
+    conv3d = tvm.IRModule.from_expr(conv3d)
+    config = conv3d, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+    conv3d, dic, param_lst = get_conv3d(padding=(0, 0, 0, 1, 1, 1), dtype=dtype)
+    conv3d = tvm.IRModule.from_expr(conv3d)
+    config = conv3d, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv3d_pattern(run_module, dtype="float32"):
+    activation_lst = [None, "relu", "tanh", "sigmoid"]
+    for a in activation_lst:
+        conv3d, dic, param_lst = get_conv3d(activation=a, dtype=dtype)
+        conv3d = tvm.IRModule.from_expr(conv3d)
+        config = conv3d, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+        conv3d_bias, dic, param_lst = get_conv3d_bias(activation=a, dtype=dtype)
+        conv3d_bias = tvm.IRModule.from_expr(conv3d_bias)
+        config = conv3d_bias, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv3d_transpose(run_module, dtype="float32"):
+    conv3d_transpose, dic, param_lst = get_conv3d_transpose(dtype=dtype)
+    conv3d_transpose = tvm.IRModule.from_expr(conv3d_transpose)
+    config = conv3d_transpose, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+    conv3d_transpose, dic, param_lst = get_conv3d_transpose(strides=(2, 2, 2), dtype=dtype)
+    conv3d_transpose = tvm.IRModule.from_expr(conv3d_transpose)
+    config = conv3d_transpose, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+    conv3d_transpose, dic, param_lst = get_conv3d_transpose(
+        strides=(2, 2, 2), output_padding=(1, 1, 1), dtype=dtype
+    )
+    conv3d_transpose = tvm.IRModule.from_expr(conv3d_transpose)
+    config = conv3d_transpose, dic, param_lst
+    run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
+def test_conv3d_transpose_pattern(run_module, dtype="float32"):
+    activation_lst = [None, "relu", "tanh", "sigmoid"]
+    for a in activation_lst:
+        conv3d, dic, param_lst = get_conv3d_transpose(activation=a, dtype=dtype)
+        conv3d = tvm.IRModule.from_expr(conv3d)
+        config = conv3d, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+        conv3d_bias, dic, param_lst = get_conv3d_transpose_bias(activation=a, dtype=dtype)
+        conv3d_bias = tvm.IRModule.from_expr(conv3d_bias)
+        config = conv3d_bias, dic, param_lst
+        run_and_verify_func(config, run_module=run_module, dtype=dtype)
+
+
 def test_dense(run_module, dtype="float32"):
     x_shape = (1, 16)
     k_shape = (32, 16)
@@ -344,6 +728,111 @@ def test_dense_pattern(run_module, dtype="float32"):
     run_and_verify_func(config, run_module=run_module, dtype=dtype)
 
 
+def test_pool2d(run_module, dtype="float32"):
+    def get_graph(
+        op,
+        x_shape=(1, 3, 32, 32),
+        pool_size=(2, 2),
+        strides=(2, 2),
+        padding=(0, 0),
+        ceil_mode=False,
+        count_include_pad=None,
+    ):
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
+        if count_include_pad is not None:
+            out = op(
+                x,
+                pool_size=pool_size,
+                strides=strides,
+                padding=padding,
+                ceil_mode=ceil_mode,
+                count_include_pad=count_include_pad,
+            )
+        else:
+            out = op(
+                x,
+                pool_size=pool_size,
+                strides=strides,
+                padding=padding,
+                ceil_mode=ceil_mode,
+            )
+        out = tvm.IRModule.from_expr(out)
+        return out, {"x": x_shape}, []
+
+    for pool_size in [(2, 2), (3, 3)]:
+        for strides in [(1, 1), (2, 2)]:
+            for padding in [(0, 0), (1, 1), (0, 0, 1, 1)]:
+                for ceil_mode in [False]:
+                    # Skip "the padding size is larger than or equal to the filter size for exclusive-counting pooling"
+                    if pool_size == (2, 2) and padding == (0, 0, 1, 1):
+                        continue
+                    for count_include_pad in [False, True]:
+                        # Skip "inclusive-counted blended or average pooling is not supported in combination with asymmetric padding"
+                        if count_include_pad and (padding == (0, 0, 1, 1) or strides == (2, 2)):
+                            continue
+                        run_and_verify_func(
+                            get_graph(
+                                relay.nn.avg_pool2d,
+                                pool_size=pool_size,
+                                strides=strides,
+                                padding=padding,
+                                ceil_mode=ceil_mode,
+                                count_include_pad=count_include_pad,
+                            ),
+                            run_module=run_module,
+                        )
+                    run_and_verify_func(
+                        get_graph(
+                            relay.nn.max_pool2d,
+                            pool_size=pool_size,
+                            strides=strides,
+                            padding=padding,
+                            ceil_mode=ceil_mode,
+                        ),
+                        run_module=run_module,
+                    )
+
+
+def test_pool3d(run_module, dtype="float32"):
+    def get_graph(
+        op,
+        x_shape=(1, 3, 8, 32, 32),
+        pool_size=(2, 2, 2),
+        strides=(2, 2, 2),
+        padding=(0, 0, 0),
+        ceil_mode=False,
+        count_include_pad=None,
+        dtype="float32",
+    ):
+        x = relay.var("x", shape=(x_shape), dtype=dtype)
+        if count_include_pad is not None:
+            out = op(
+                x,
+                pool_size=pool_size,
+                strides=strides,
+                padding=padding,
+                ceil_mode=ceil_mode,
+                count_include_pad=count_include_pad,
+            )
+        else:
+            out = op(
+                x,
+                pool_size=pool_size,
+                strides=strides,
+                padding=padding,
+                ceil_mode=ceil_mode,
+            )
+        out = tvm.IRModule.from_expr(out)
+        return out, {"x": x_shape}, []
+
+    run_and_verify_func(get_graph(relay.nn.avg_pool3d), run_module=run_module)
+    run_and_verify_func(get_graph(relay.nn.max_pool3d), run_module=run_module)
+    run_and_verify_func(
+        get_graph(relay.nn.max_pool3d, padding=(0, 0, 0, 1, 1, 1)), run_module=run_module
+    )
+    run_and_verify_func(get_graph(relay.nn.max_pool3d, strides=(1, 1, 1)), run_module=run_module)
+
+
 if __name__ == "__main__":
     import sys
 
diff --git a/tests/python/contrib/test_ethosn/test_depth_to_space.py b/tests/python/contrib/test_ethosn/test_depth_to_space.py
index 1366c6a86346..1675b82eeace 100644
--- a/tests/python/contrib/test_ethosn/test_depth_to_space.py
+++ b/tests/python/contrib/test_ethosn/test_depth_to_space.py
@@ -69,7 +69,7 @@ def test_depth_to_space_failure():
             "dtype='int16', dtype must be either uint8, int8 or int32;",
         ),
         ((1, 16, 16, 16), 4, "uint8", "NHWC", "Only block size of 2 is supported"),
-        ((1, 16, 16, 16), 2, "uint8", "NCHW", "layout=NCHW, layout must = NHWC"),
+        ((1, 16, 16, 16), 2, "uint8", "NCHW", "Input layer must be NHWC or NHWCB"),
     ]
 
     for shape, block, dtype, layout, err_msg in trials:
diff --git a/tests/python/contrib/test_ethosn/test_networks.py b/tests/python/contrib/test_ethosn/test_networks.py
index f79b1f4dc178..33f4dfd28e51 100644
--- a/tests/python/contrib/test_ethosn/test_networks.py
+++ b/tests/python/contrib/test_ethosn/test_networks.py
@@ -228,7 +228,8 @@ def test_ssd_mobilenet_v1():
     # on hardware that isn't available in CI.
     _compile_hash = {"5ee8ed6af9a7f31fc14957b51a8e7423", "e6a91ccc47ba4c6b4614fcd676bd726f"}
     if tei.get_ethosn_api_version() == 2111:
-        _compile_hash = {"afb68ca8f452d1f4a674b457b5e30f59", "a37f900601b9493bd142e8aed16205e5"}
+        # TODO(Leo-arm): review split operator
+        _compile_hash = {"a37f900601b9493bd142e8aed16205e5", "afb68ca8f452d1f4a674b457b5e30f59"}
     if tei.get_ethosn_api_version() == 2102:
         _compile_hash = {"7795b6c67178da9d1f9b98063bad75b1", "10826406ae724e52f360a06c35ced09d"}
         if tei.get_ethosn_variant() == "Ethos-N78_1TOPS_2PLE_RATIO":
diff --git a/tests/python/contrib/test_ethosn/test_split.py b/tests/python/contrib/test_ethosn/test_split.py
index 28f7f5da58e2..320172e1e3f1 100644
--- a/tests/python/contrib/test_ethosn/test_split.py
+++ b/tests/python/contrib/test_ethosn/test_split.py
@@ -37,8 +37,9 @@ def test_split(dtype):
     trials = [
         ((1, 16, 16, 32), (2, 7, 10), 2),
         ((1, 12, 8, 16), 3, 1),
-        ((1, 66), 11, 1),
     ]
+    if tei.get_ethosn_api_version() < 2111:
+        trials.append(((1, 66), 11, 1))
 
     np.random.seed(0)
     for shape, splits, axis in trials:
diff --git a/tests/python/contrib/test_ethosu/cascader/conftest.py b/tests/python/contrib/test_ethosu/cascader/conftest.py
index 21ed401994c1..cffaf83df0bc 100644
--- a/tests/python/contrib/test_ethosu/cascader/conftest.py
+++ b/tests/python/contrib/test_ethosu/cascader/conftest.py
@@ -203,3 +203,400 @@ def _get_func():
     @pytest.fixture
     def MobileNetv2DiamondTE():
         return make_MobileNetv2DiamondTE()
+
+    @pytest.fixture
+    def MobileNetv2DiamondGraph():
+        _, te_graph, const_dict = make_MobileNetv2DiamondTE()
+        device_config = cs.EthosuDeviceConfig("ethos-u55-256")
+        return cs.create_cascader_graph(te_graph, const_dict, device_config)
+
+    def make_BinaryTE():
+        def _get_func():
+            ifm_a = relay.var("ifm_a", shape=(1, 8, 8, 8), dtype="int8")
+            ifm_b = relay.var("ifm_b", shape=(1, 8, 8, 8), dtype="int8")
+            conv1 = make_ethosu_conv2d(
+                ifm=ifm_a,
+                ifm_channels=8,
+                ofm_channels=8,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv2 = make_ethosu_conv2d(
+                ifm=ifm_b,
+                ifm_channels=8,
+                ofm_channels=8,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            add1 = make_ethosu_binary_elementwise(
+                ifm=conv1,
+                ifm2=conv2,
+                ifm_channels=8,
+                ifm2_channels=8,
+                operator_type="ADD",
+                ofm_dtype="int8",
+            )
+            func = relay.Function(relay.analysis.free_vars(add1), add1)
+            func = run_opt_pass(func, relay.transform.InferType())
+            return func
+
+        func = _get_func()
+        te_graph, const_dict = create_te_graph(func)
+        sch = tvm.te.create_schedule([t.op for t in te_graph.outputs])
+        return sch, te_graph, const_dict
+
+    @pytest.fixture
+    def BinaryTE():
+        return make_BinaryTE()
+
+    @pytest.fixture
+    def BinaryGraph():
+        _, te_graph, const_dict = make_BinaryTE()
+        device_config = cs.EthosuDeviceConfig("ethos-u55-256")
+        return cs.create_cascader_graph(te_graph, const_dict, device_config)
+
+    def make_MobileNetv1StartTE():
+        def _get_func():
+            ifm = relay.var("ifm", shape=(1, 224, 224, 3), dtype="int8")
+            conv1 = make_ethosu_conv2d(
+                ifm=ifm,
+                ifm_channels=3,
+                ofm_channels=32,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            depth1 = make_ethosu_depthwise_conv2d(
+                ifm=conv1,
+                channels=32,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv2 = make_ethosu_conv2d(
+                ifm=depth1,
+                ifm_channels=32,
+                ofm_channels=64,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth2 = make_ethosu_depthwise_conv2d(
+                ifm=conv2,
+                channels=64,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            conv3 = make_ethosu_conv2d(
+                ifm=depth2,
+                ifm_channels=64,
+                ofm_channels=128,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth3 = make_ethosu_depthwise_conv2d(
+                ifm=conv3,
+                channels=128,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv4 = make_ethosu_conv2d(
+                ifm=depth3,
+                ifm_channels=128,
+                ofm_channels=128,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth4 = make_ethosu_depthwise_conv2d(
+                ifm=conv4,
+                channels=128,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            func = relay.Function(relay.analysis.free_vars(depth4), depth4)
+            func = run_opt_pass(func, relay.transform.InferType())
+            return func
+
+        func = _get_func()
+        te_graph, const_dict = create_te_graph(func)
+        sch = tvm.te.create_schedule([t.op for t in te_graph.outputs])
+        return sch, te_graph, const_dict
+
+    @pytest.fixture
+    def MobileNetv1StartTE():
+        return make_MobileNetv1StartTE()
+
+    @pytest.fixture
+    def MobileNetv1StartGraph():
+        _, te_graph, const_dict = make_MobileNetv1StartTE()
+        device_config = cs.EthosuDeviceConfig("ethos-u55-256")
+        return cs.create_cascader_graph(te_graph, const_dict, device_config)
+
+    def make_MobileNetv1TE():
+        def _get_func():
+            ifm = relay.var("ifm", shape=(1, 224, 224, 3), dtype="int8")
+            conv1 = make_ethosu_conv2d(
+                ifm=ifm,
+                ifm_channels=3,
+                ofm_channels=32,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            depth1 = make_ethosu_depthwise_conv2d(
+                ifm=conv1,
+                channels=32,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv2 = make_ethosu_conv2d(
+                ifm=depth1,
+                ifm_channels=32,
+                ofm_channels=64,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth2 = make_ethosu_depthwise_conv2d(
+                ifm=conv2,
+                channels=64,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            conv3 = make_ethosu_conv2d(
+                ifm=depth2,
+                ifm_channels=64,
+                ofm_channels=128,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth3 = make_ethosu_depthwise_conv2d(
+                ifm=conv3,
+                channels=128,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv4 = make_ethosu_conv2d(
+                ifm=depth3,
+                ifm_channels=128,
+                ofm_channels=128,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth4 = make_ethosu_depthwise_conv2d(
+                ifm=conv4,
+                channels=128,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            conv5 = make_ethosu_conv2d(
+                ifm=depth4,
+                ifm_channels=128,
+                ofm_channels=256,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth5 = make_ethosu_depthwise_conv2d(
+                ifm=conv5,
+                channels=256,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv6 = make_ethosu_conv2d(
+                ifm=depth5,
+                ifm_channels=256,
+                ofm_channels=256,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth6 = make_ethosu_depthwise_conv2d(
+                ifm=conv6,
+                channels=256,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            conv7 = make_ethosu_conv2d(
+                ifm=depth6,
+                ifm_channels=256,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth7 = make_ethosu_depthwise_conv2d(
+                ifm=conv7,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv8 = make_ethosu_conv2d(
+                ifm=depth7,
+                ifm_channels=512,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth8 = make_ethosu_depthwise_conv2d(
+                ifm=conv8,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv9 = make_ethosu_conv2d(
+                ifm=depth8,
+                ifm_channels=512,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth9 = make_ethosu_depthwise_conv2d(
+                ifm=conv9,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv10 = make_ethosu_conv2d(
+                ifm=depth9,
+                ifm_channels=512,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth10 = make_ethosu_depthwise_conv2d(
+                ifm=conv10,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv11 = make_ethosu_conv2d(
+                ifm=depth10,
+                ifm_channels=512,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth11 = make_ethosu_depthwise_conv2d(
+                ifm=conv11,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv12 = make_ethosu_conv2d(
+                ifm=depth11,
+                ifm_channels=512,
+                ofm_channels=512,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth12 = make_ethosu_depthwise_conv2d(
+                ifm=conv12,
+                channels=512,
+                kernel_shape=(3, 3),
+                padding=(0, 0, 1, 1),
+                strides=(2, 2),
+                dilation=(1, 1),
+            )
+            conv13 = make_ethosu_conv2d(
+                ifm=depth12,
+                ifm_channels=512,
+                ofm_channels=1024,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            depth13 = make_ethosu_depthwise_conv2d(
+                ifm=conv13,
+                channels=1024,
+                kernel_shape=(3, 3),
+                padding=(1, 1, 1, 1),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            conv14 = make_ethosu_conv2d(
+                ifm=depth13,
+                ifm_channels=1024,
+                ofm_channels=1024,
+                kernel_shape=(1, 1),
+                padding=(0, 0, 0, 0),
+                strides=(1, 1),
+                dilation=(1, 1),
+            )
+            func = relay.Function(relay.analysis.free_vars(conv14), conv14)
+            func = run_opt_pass(func, relay.transform.InferType())
+            return func
+
+        func = _get_func()
+        te_graph, const_dict = create_te_graph(func)
+        sch = tvm.te.create_schedule([t.op for t in te_graph.outputs])
+        return sch, te_graph, const_dict
+
+    @pytest.fixture
+    def MobileNetv1TE():
+        return make_MobileNetv1TE()
+
+    @pytest.fixture
+    def MobileNetv1Graph():
+        _, te_graph, const_dict = make_MobileNetv1TE()
+        device_config = cs.EthosuDeviceConfig("ethos-u55-256")
+        return cs.create_cascader_graph(te_graph, const_dict, device_config)
diff --git a/tests/python/contrib/test_ethosu/cascader/test_proposal_generator.py b/tests/python/contrib/test_ethosu/cascader/test_proposal_generator.py
new file mode 100644
index 000000000000..5af89a415978
--- /dev/null
+++ b/tests/python/contrib/test_ethosu/cascader/test_proposal_generator.py
@@ -0,0 +1,139 @@
+# 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.
+import pytest
+from tvm.contrib.ethosu.cascader.proposal_generator import generate_proposals
+
+from .infra import make_simple_home_map, make_options, ethosu_enabled
+
+
+if ethosu_enabled:
+
+    def test_generate_proposals(FLASH, SRAM, TwoConv2DGraph):
+        graph = TwoConv2DGraph
+        min_sram = 3700
+        max_sram = 11700
+        input_configs = 1
+        parts = 2
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=32,
+            stripe_factors=4,
+            max_plan_size=10,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+
+        for proposal in proposals:
+            assert 0 < len(proposal.plans) <= parts
+            assert len(proposal.input_tensor_configs) == input_configs
+            assert len(proposal.part_group) == parts
+            assert min_sram < proposal.memory_usage < max_sram
+            assert proposal.cycles > 0
+
+    def test_generate_proposals_binary(FLASH, SRAM, BinaryGraph):
+        graph = BinaryGraph
+        input_configs = 2
+        parts = 3
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=32,
+            stripe_factors=4,
+            max_plan_size=10,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+
+        for proposal in proposals:
+            assert 0 < len(proposal.plans) <= parts
+            assert len(proposal.input_tensor_configs) == input_configs
+            assert len(proposal.part_group) == parts
+            assert proposal.cycles > 0
+
+    def test_generate_proposals_mobilenetv1_start(FLASH, SRAM, MobileNetv1StartGraph):
+        graph = MobileNetv1StartGraph
+        min_sram = 200000
+        max_sram = 1300000
+        input_configs = 1
+        parts = 8
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=32,
+            stripe_factors=5,
+            max_plan_size=10,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+
+        for proposal in proposals:
+            assert 0 < len(proposal.plans) <= parts
+            assert len(proposal.input_tensor_configs) == input_configs
+            assert len(proposal.part_group) == parts
+            assert min_sram < proposal.memory_usage < max_sram
+            assert proposal.cycles > 0
+
+    def test_generate_proposals_mobilenetv1(FLASH, SRAM, MobileNetv1Graph):
+        graph = MobileNetv1Graph
+        min_sram = 200000
+        max_sram = 1300000
+        input_configs = 1
+        parts = 27
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=32,
+            stripe_factors=5,
+            max_plan_size=10,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+
+        for proposal in proposals:
+            assert 0 < len(proposal.plans) <= parts
+            assert len(proposal.input_tensor_configs) == input_configs
+            assert len(proposal.part_group) == parts
+            assert min_sram < proposal.memory_usage < max_sram
+            assert proposal.cycles > 0
+
+    def test_generate_proposals_mobilenetv2diamond(FLASH, SRAM, MobileNetv2DiamondGraph):
+        graph = MobileNetv2DiamondGraph
+        min_sram = 370000
+        max_sram = 990000
+        input_configs = 1
+        parts = 5
+        home_map = make_simple_home_map(graph, SRAM, FLASH)
+        options = make_options(
+            cascade_region=SRAM,
+            max_proposals=64,
+            stripe_factors=5,
+            max_plan_size=10,
+        )
+
+        proposals = generate_proposals(graph, home_map, options)
+
+        for proposal in proposals:
+            assert 0 < len(proposal.plans) <= parts
+            assert len(proposal.input_tensor_configs) == input_configs
+            assert len(proposal.part_group) == parts
+            assert min_sram < proposal.memory_usage < max_sram
+            assert proposal.cycles > 0
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])
diff --git a/tests/python/contrib/test_ethosu/infra.py b/tests/python/contrib/test_ethosu/infra.py
index d5bd28039feb..be5466c2f9a9 100644
--- a/tests/python/contrib/test_ethosu/infra.py
+++ b/tests/python/contrib/test_ethosu/infra.py
@@ -189,7 +189,7 @@ def deserialize_command_stream(blob):
     return cmms
 
 
-def create_test_runner(accel="ethos-u55-256"):
+def create_test_runner(accel="ethos-u55-256", enable_usmp=True):
     file_dir = os.path.dirname(os.path.abspath(__file__))
     test_root = os.path.join(file_dir, "reference_system")
     _, ethosu_variant, ethosu_macs = accel.split("-")
@@ -215,12 +215,17 @@ def create_test_runner(accel="ethos-u55-256"):
             "relay.ext.ethos-u.options": {
                 "accelerator_config": accel,
             },
+            "tir.usmp.enable": enable_usmp,
+            "tir.usmp.algorithm": "hill_climb",
+            "tir.disable_storage_rewrite": enable_usmp,
         },
     )
 
 
-def build_source(module, inputs, outputs, accel="ethos-u55-256", output_tolerance=0):
-    test_runner = create_test_runner(accel)
+def build_source(
+    module, inputs, outputs, accel="ethos-u55-256", output_tolerance=0, enable_usmp=True
+):
+    test_runner = create_test_runner(accel, enable_usmp)
     return compile_models(
         models=AOTTestModel(
             module=module,
@@ -239,12 +244,13 @@ def build_source(module, inputs, outputs, accel="ethos-u55-256", output_toleranc
 def verify_source(
     models: List[AOTCompiledTestModel],
     accel="ethos-u55-256",
+    enable_usmp=True,
 ):
     """
     This method verifies the generated source from an NPU module by building it and running on an FVP.
     """
     interface_api = "c"
-    test_runner = create_test_runner(accel)
+    test_runner = create_test_runner(accel, enable_usmp)
     run_and_check(
         models,
         test_runner,
@@ -306,9 +312,10 @@ def generate_ref_data_tflite(model):
         interpreter.set_tensor(index, value)
     interpreter.invoke()
 
-    expected_output_data = [
-        interpreter.get_tensor(output_detail["index"]) for output_detail in output_details
-    ]
+    expected_output_data = {
+        output_detail["name"]: interpreter.get_tensor(output_detail["index"])
+        for output_detail in output_details
+    }
 
     return input_data, expected_output_data
 
diff --git a/tests/python/contrib/test_ethosu/test_codegen.py b/tests/python/contrib/test_ethosu/test_codegen.py
index 052586387b1f..8ac9bafcdb09 100644
--- a/tests/python/contrib/test_ethosu/test_codegen.py
+++ b/tests/python/contrib/test_ethosu/test_codegen.py
@@ -754,7 +754,7 @@ def rounding_right_shift(lhs, rhs):
         "ifm": lhs,
         "ifm2": rhs,
     }
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(ethosu_mod, input_data, output_data, accel_type)
@@ -781,7 +781,7 @@ def generate_output_data(input_data):
 
     cpu_mod = create_model()
     input_data = {"ifm": np.random.randint(-120, high=120, size=ifm_shape, dtype="int8")}
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(
@@ -910,7 +910,7 @@ def clz_comp(n):
 
     cpu_mod = create_model()
     input_data = {"ifm": np.random.randint(-500000, high=500000, size=ifm_shape, dtype="int32")}
-    output_data = generate_output_data(input_data)
+    output_data = {"output": generate_output_data(input_data)[0]}
     ethosu_mod = _create_ethosu_partition(cpu_mod)
 
     _compare_ethosu_with_reference(ethosu_mod, input_data, output_data, accel_type)
@@ -950,6 +950,7 @@ def concat_func(*inputs):
     _compare_tvm_with_tflite(concat_func, shapes, accel_type, output_tolerance=1)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 def test_tflite_sigmoid(accel_type):
     ifm_shape = [1, 135, 41, 6]
@@ -1014,6 +1015,7 @@ def create_model():
     _compare_ethosu_with_reference(ethosu_mod, input_data, output_data, accel_type)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize("ifm_shape,axis", [((2,), 0), ((1, 3, 3), 2)])
 def test_tflite_expand_dims(accel_type, ifm_shape, axis):
@@ -1053,6 +1055,7 @@ def resize_model(x):
     _compare_tvm_with_tflite(resize_model, [ifm_shape], accel_type)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize(
     "ifm_shape,size,align_corners",
@@ -1077,6 +1080,7 @@ def resize_model(x):
     _compare_tvm_with_tflite(resize_model, [ifm_shape], accel_type, output_tolerance=1)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize(
     "ifm_shape,ofm_shape,kernel_shape,padding",
@@ -1116,6 +1120,7 @@ def conv2d_transpose(x):
     _compare_tvm_with_tflite(conv2d_transpose, [ifm_shape], accel_type=accel_type)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 @pytest.mark.parametrize(
     "ifm_shapes,axis",
@@ -1150,5 +1155,17 @@ def unpack_func(x):
     _compare_tvm_with_tflite(unpack_func, [ifm_shape], accel_type)
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
+@pytest.mark.parametrize("accel_type", ACCEL_TYPES)
+@pytest.mark.parametrize("ifm_shape", [(1, 15, 15, 3), (1, 8, 9, 1)])
+@pytest.mark.parametrize("alpha", [0.2, 0.634])
+def test_tflite_leaky_relu(accel_type, ifm_shape, alpha):
+    @tf.function
+    def leaky_relu_func(x):
+        return tf.nn.leaky_relu(x, alpha=alpha)
+
+    _compare_tvm_with_tflite(leaky_relu_func, [ifm_shape], accel_type)
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/tests/python/contrib/test_ethosu/test_legalize.py b/tests/python/contrib/test_ethosu/test_legalize.py
index 3922bffa35d5..83165d7d54fb 100644
--- a/tests/python/contrib/test_ethosu/test_legalize.py
+++ b/tests/python/contrib/test_ethosu/test_legalize.py
@@ -2290,5 +2290,61 @@ def verify(ext_func):
     verify(mod["tvmgen_default_ethos_u_main_0"])
 
 
+@pytest.mark.parametrize("ifm_shape", [(1, 15, 15, 3), (1, 8, 9, 1)])
+@pytest.mark.parametrize("alpha", [0.2, 0.634])
+def test_tflite_leaky_relu(ifm_shape, alpha):
+    dtype = "int8"
+
+    def create_tflite_graph():
+        class Model(tf.Module):
+            @tf.function
+            def leaky_relu_func(self, x):
+                return tf.nn.leaky_relu(x, alpha=alpha)
+
+        model = Model()
+        concrete_func = model.leaky_relu_func.get_concrete_function(
+            tf.TensorSpec(ifm_shape, tf.float32),
+        )
+
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple(ifm_shape))
+                yield [data.astype(np.float32)]
+
+        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        converter.representative_dataset = representative_dataset
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        converter.inference_input_type = tf.int8
+        converter.inference_output_type = tf.int8
+        tflite_model = converter.convert()
+
+        return tflite_model
+
+    def verify(ext_func):
+        func_body = ext_func.body
+        assert func_body.op.name == "contrib.ethosu.identity"
+        assert func_body.attrs.activation == "LUT"
+        assert tuple(func_body.args[0].checked_type.shape) == (ifm_shape)
+        assert tuple(func_body.args[1].checked_type.shape) == (256,)
+
+    tflite_graph = create_tflite_graph()
+    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
+
+    mod, _ = relay.frontend.from_tflite(
+        tflite_model,
+        shape_dict={"input": ifm_shape},
+        dtype_dict={"input": dtype},
+    )
+    mod = ethosu.partition_for_ethosu(mod)
+    mod["tvmgen_default_ethos_u_main_0"] = dataflow_pattern.rewrite(
+        legalize.LeakyReLURewriter(), mod["tvmgen_default_ethos_u_main_0"]
+    )
+    mod["tvmgen_default_ethos_u_main_0"] = relay.transform.InferType()(mod)[
+        "tvmgen_default_ethos_u_main_0"
+    ]
+    verify(mod["tvmgen_default_ethos_u_main_0"])
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/tests/python/contrib/test_ethosu/test_lookup_table.py b/tests/python/contrib/test_ethosu/test_lookup_table.py
index 90a51c53fdb3..31ed56ae2998 100644
--- a/tests/python/contrib/test_ethosu/test_lookup_table.py
+++ b/tests/python/contrib/test_ethosu/test_lookup_table.py
@@ -33,6 +33,7 @@
 ACCEL_TYPES = ["ethos-u55-256", "ethos-u55-128", "ethos-u55-64", "ethos-u55-32"]
 
 
+@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10300")
 @pytest.mark.parametrize("accel_type", ACCEL_TYPES)
 def test_tflite_lut_activations(accel_type):
 
@@ -154,7 +155,7 @@ def test_random_lut(accel_type):
     compiled_models = infra.build_source(
         mod,
         {"ifm": in_data},
-        out_data,
+        {"output": out_data},
         accel_type,
     )
 
diff --git a/tests/python/contrib/test_ethosu/test_networks.py b/tests/python/contrib/test_ethosu/test_networks.py
index de263c18f368..b73ba4f62a54 100644
--- a/tests/python/contrib/test_ethosu/test_networks.py
+++ b/tests/python/contrib/test_ethosu/test_networks.py
@@ -29,41 +29,75 @@
 from tvm import relay
 from tvm.relay.backend.contrib.ethosu import util
 from tvm.relay.op.contrib.ethosu import partition_for_ethosu
+from tests.python.relay.aot.aot_test_utils import create_relay_module_and_inputs_from_tflite_file
+from tvm.micro import model_library_format as mlf
 
 import tvm.relay.testing.tf as tf_testing
 
 from . import infra
 
-ACCEL_TYPES = ["ethos-u55-256", "ethos-u55-128", "ethos-u55-64", "ethos-u55-32", "ethos-u65-256"]
 
+MOBILENET_V1_URL = (
+    "https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
+    "mobilenet_v1_1.0_224_quant.tflite",
+)
+
+MOBILENET_V2_URL = (
+    "https://storage.googleapis.com/download.tensorflow.org/models/tflite_11_05_08/mobilenet_v2_1.0_224_quant.tgz",
+    "mobilenet_v2_1.0_224_quant.tflite",
+)
 
-@pytest.mark.parametrize("accel_type", ACCEL_TYPES)
-def test_forward_mobilenet_v1(accel_type):
-    """Test the Mobilenet V1 TF Lite model."""
+
+@pytest.mark.parametrize(
+    "accel_type, model_url, workspace_size, tolerance",
+    [
+        ("ethos-u65-256", MOBILENET_V1_URL, 1423344, 10),
+        ("ethos-u65-256", MOBILENET_V2_URL, 2185584, 5),
+        ("ethos-u55-256", MOBILENET_V1_URL, 1423344, 10),
+        ("ethos-u55-256", MOBILENET_V2_URL, 2185584, 5),
+        ("ethos-u55-128", MOBILENET_V2_URL, 2185584, 5),
+        ("ethos-u55-64", MOBILENET_V2_URL, 2185584, 5),
+        ("ethos-u55-32", MOBILENET_V2_URL, 2185584, 5),
+    ],
+)
+def test_networks_without_usmp(accel_type, model_url, workspace_size, tolerance):
     np.random.seed(23)
-    tflite_model_file = tf_testing.get_workload_official(
-        "https://storage.googleapis.com/download.tensorflow.org/"
-        "models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
-        "mobilenet_v1_1.0_224_quant.tflite",
+    tflite_model_file = tf_testing.get_workload_official(model_url[0], model_url[1])
+    mod, input_data, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
+    output_data = generate_ref_data(mod, input_data, params)
+    mod = partition_for_ethosu(mod, params)
+    compiled_models = infra.build_source(
+        mod, input_data, output_data, accel_type, output_tolerance=tolerance, enable_usmp=False
+    )
+    mlf_memory_map = mlf._build_function_memory_map(
+        compiled_models[0].executor_factory.function_metadata
     )
-    with open(tflite_model_file, "rb") as f:
-        tflite_model_buf = f.read()
-    input_tensor = "input"
-    input_dtype = "uint8"
-    input_shape = (1, 224, 224, 3)
-    in_min, in_max = util.get_range_for_dtype_str(input_dtype)
-    input_data = np.random.randint(in_min, high=in_max, size=input_shape, dtype=input_dtype)
+    assert mlf_memory_map["main"][0]["workspace_size_bytes"] == workspace_size
+    infra.verify_source(compiled_models, accel_type, enable_usmp=False)
 
-    relay_mod, params = convert_to_relay(tflite_model_buf)
-    input_data = {input_tensor: input_data}
-    output_data = generate_ref_data(relay_mod, input_data)
 
-    mod = partition_for_ethosu(relay_mod, params)
+@pytest.mark.parametrize(
+    "accel_type, model_url, workspace_size, tolerance",
+    [
+        ("ethos-u65-256", MOBILENET_V1_URL, 1205872, 10),
+        ("ethos-u55-256", MOBILENET_V2_URL, 1507152, 5),
+    ],
+)
+def test_networks_with_usmp(accel_type, model_url, workspace_size, tolerance):
+    np.random.seed(23)
+    tflite_model_file = tf_testing.get_workload_official(model_url[0], model_url[1])
+    mod, input_data, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
+    output_data = generate_ref_data(mod, input_data, params)
+    mod = partition_for_ethosu(mod, params)
     compiled_models = infra.build_source(
-        mod, input_data, output_data, accel_type, output_tolerance=10
+        mod, input_data, output_data, accel_type, output_tolerance=tolerance, enable_usmp=True
     )
-    infra.verify_source(compiled_models, accel_type)
+    allocated_pool_info = list(
+        dict(compiled_models[0].executor_factory.executor_codegen_metadata.pool_inputs).values()
+    )[0]
+    assert allocated_pool_info.allocated_size == workspace_size
+    infra.verify_source(compiled_models, accel_type, enable_usmp=True)
 
 
 if __name__ == "__main__":
-    test_forward_mobilenet_v1()
+    pytest.main([__file__])
diff --git a/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py b/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py
index add8021083c6..de214888be6b 100644
--- a/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py
+++ b/tests/python/contrib/test_ethosu/test_tir_to_cs_translator.py
@@ -645,7 +645,7 @@ def populate_ethosu_copy_calls(stmt):
                 {
                     "src": "placeholder_5",
                     "dest": "placeholder_d_global",
-                    "length": 8,
+                    "length": 32,
                 },
             ],
         },
@@ -851,24 +851,45 @@ def _check_buffer(address, region, length, buffer_var):
                 length, dtype=buffer_dtype
             )
         elif buffer_type == tir_to_cs_translator.BufferType.scratch:
-            shape = list(buffer_info[buffer_var].shape)
-            assert length == np.prod(shape)
-            assert address < scratch_size
+            assert address < tvmbaw_workspace_size
 
-            size_in_bytes = int(np.prod(shape)) * dtype_bytes
+            size_in_bytes = allocate_node_sizes[buffer_var]
             # Every buffer is adjusted to align to 16 bytes
             size_in_bytes = util.round_up(size_in_bytes, 16)
-            assert address + size_in_bytes <= scratch_size
+            assert address + size_in_bytes <= tvmbaw_workspace_size
             # The scratch area should not be used by any other buffer
-            assert not scratch_mask[address : address + size_in_bytes].any()
+            assert not tvmbaw_workspace_mask[address : address + size_in_bytes].any()
             # The scratch area is marked as used
-            scratch_mask[address : address + size_in_bytes] = np.ones(size_in_bytes, dtype="uint8")
+            tvmbaw_workspace_mask[address : address + size_in_bytes] = np.ones(
+                size_in_bytes, dtype="uint8"
+            )
         elif buffer_type == tir_to_cs_translator.BufferType.input:
             assert address == 0
         else:
             assert buffer_type == tir_to_cs_translator.BufferType.output
             assert address == 0
 
+    def _get_allocate_node_sizes(mod):
+        # There should only be a single function
+        assert len(mod.functions.items()) == 1
+        primfunc = mod.functions.items()[0][1]
+        _allocate_node_sizes = dict()
+
+        def analyze_remaining_allocates(stmt):
+            if isinstance(stmt, tvm.tir.stmt.Allocate):
+                allocate = stmt
+                pointer_type = allocate.buffer_var.type_annotation
+                storage_scope = pointer_type.storage_scope
+                if storage_scope == "global":
+                    dtype_bytes = np.iinfo(np.dtype(allocate.dtype)).bits // 8
+                    size_in_bytes = int(dtype_bytes * np.prod(list(allocate.extents)))
+                    # Every memory address the NPU access have to be 16 byte aligned
+                    size_in_bytes = util.round_up(size_in_bytes, 16)
+                    _allocate_node_sizes[allocate.buffer_var] = size_in_bytes
+
+        tvm.tir.stmt_functor.post_order_visit(primfunc.body, analyze_remaining_allocates)
+        return _allocate_node_sizes
+
     def verify(npu_ops):
         """This wrapper verifies the allocated addresses matches with original tir buffers"""
         checked_buffers = set()
@@ -933,22 +954,29 @@ def check_buffer(address, region, length, buffer_var):
         tir_mod = test_case["tir_module"]
         tir_mod["main"] = tir_mod["main"].with_attr("target", tvm.target.Target("ethos-u"))
         tir_mod = tvm.tir.transform.MakeUnpackedAPI()(tir_mod)
+        candidate_regions_for_scratch = [5, 2, 1]
+        (
+            scratch_region_map,
+            tvmbaw_workspace_size,
+            _,
+        ) = tir_to_cs_translator.analyze_scratch_memory_acesses(
+            tir_mod, candidate_regions_for_scratch
+        )
+        allocate_node_sizes = _get_allocate_node_sizes(tir_mod)
         buffer_info = tir_to_cs_translator.extract_buffer_info(tir_mod, test_case["param_dict"])
         extern_calls = extract_call_extern_list(tir_mod)
         _npu_ops = list()
         for extern_call in extern_calls:
             _npu_ops.append(tir_to_cs_translator.translate_ethosu_tir_call_extern(extern_call))
         npu_op_tir_buffers = collect_tir_buffer_info(_npu_ops)
-        (
-            _npu_ops,
-            constant_hex_string,
-            scratch_size,
-        ) = tir_to_cs_translator.assign_addresses(buffer_info, _npu_ops)
-        scratch_mask = np.zeros(scratch_size, dtype="uint8")
+        (_npu_ops, constant_hex_string) = tir_to_cs_translator.assign_addresses(
+            buffer_info, _npu_ops, scratch_region_map
+        )
+        tvmbaw_workspace_mask = np.zeros(tvmbaw_workspace_size, dtype="uint8")
         constant_tensor_read_mask = np.zeros(len(constant_hex_string) // 2, dtype="uint8")
         verify(_npu_ops)
         # This will be only 1 if all allocated scratch is used.
-        assert np.prod(scratch_mask) == 1
+        assert np.prod(tvmbaw_workspace_mask) == 1
         # This will be only 1 if all constant tensors is read at least once.
         assert np.prod(constant_tensor_read_mask) == 1
 
diff --git a/tests/python/contrib/test_gemm_acc32_vnni.py b/tests/python/contrib/test_gemm_acc32_vnni.py
index 3d038259b1ab..9cec823cc58a 100644
--- a/tests/python/contrib/test_gemm_acc32_vnni.py
+++ b/tests/python/contrib/test_gemm_acc32_vnni.py
@@ -57,7 +57,9 @@ def verify(target="llvm -mcpu=cascadelake"):
             (m, n),
             lambda i, j: te.sum(
                 X[i, ak].astype("int32")
-                * packedW[j / 16, (ak / 4) * 16 + j % 16, ak % 4].astype("int32"),
+                * packedW[
+                    tvm.tir.indexdiv(j, 16), tvm.tir.indexdiv(ak, 4) * 16 + j % 16, ak % 4
+                ].astype("int32"),
                 axis=ak,
             ),
             name="F",
diff --git a/tests/python/contrib/test_hexagon/conftest.py b/tests/python/contrib/test_hexagon/conftest.py
index bd2ae7c08473..767eb78f8c7c 100644
--- a/tests/python/contrib/test_hexagon/conftest.py
+++ b/tests/python/contrib/test_hexagon/conftest.py
@@ -27,8 +27,8 @@
 HEXAGON_TOOLCHAIN = "HEXAGON_TOOLCHAIN"
 TVM_TRACKER_HOST = "TVM_TRACKER_HOST"
 TVM_TRACKER_PORT = "TVM_TRACKER_PORT"
-ANDROID_TRACKER_KEY = "ANDROID_TRACKER_KEY"
 ANDROID_REMOTE_DIR = "ANDROID_REMOTE_DIR"
+ANDROID_SERIAL_NUMBER = "ANDROID_SERIAL_NUMBER"
 
 
 @tvm.testing.fixture
@@ -49,8 +49,8 @@ def _compose(args, decs):
 def requires_hexagon_toolchain(*args):
     _requires_hexagon_toolchain = [
         pytest.mark.skipif(
-            os.environ.get("HEXAGON_TOOLCHAIN") == None,
-            reason="HEXAGON_TOOLCHAIN environment variable is required to run this test.",
+            os.environ.get(HEXAGON_TOOLCHAIN) == None,
+            reason=f"Missing environment variable {HEXAGON_TOOLCHAIN}.",
         ),
     ]
 
@@ -58,74 +58,17 @@ def requires_hexagon_toolchain(*args):
 
 
 @tvm.testing.fixture
-def android_tracker_key():
-    return os.environ["ANDROID_TRACKER_KEY"]
+def android_serial_number():
+    return os.getenv(ANDROID_SERIAL_NUMBER, default=None)
 
 
 @tvm.testing.fixture
 def tvm_tracker_host():
-    return os.environ["TVM_TRACKER_HOST"]
+    return os.getenv(TVM_TRACKER_HOST, default=None)
 
 
 @tvm.testing.fixture
 def tvm_tracker_port():
-    return int(os.environ["TVM_TRACKER_PORT"])
-
-
-@tvm.testing.fixture
-def remote_path():
-    dso_binary = "test_binary.so"
-    return os.path.join(os.environ["ANDROID_REMOTE_DIR"], dso_binary)
-
-
-@tvm.testing.fixture
-def rpc_sess(android_tracker_key, tvm_tracker_host, tvm_tracker_port):
-    from tvm import rpc
-
-    tracker = rpc.connect_tracker(tvm_tracker_host, tvm_tracker_port)
-    remote = tracker.request(android_tracker_key, priority=0, session_timeout=600)
-    return remote
-
-
-def requires_rpc_tracker_and_android_key(*args):
-    """Mark a test as requiring an RPC tracker to exist in
-    the host environment to run."""
-    _requires_rpc_tracker = [
-        *tvm.testing.requires_rpc(),
-        pytest.mark.skipif(
-            os.environ.get(TVM_TRACKER_HOST) == None,
-            reason="Missing environment variable, TVM_TRACKER_HOST",
-        ),
-        pytest.mark.skipif(
-            os.environ.get(TVM_TRACKER_PORT) == None,
-            reason="Missing environment variable, TVM_TRACKER_PORT",
-        ),
-        pytest.mark.skipif(
-            os.environ.get(ANDROID_TRACKER_KEY) == None,
-            reason="Missing environment variable, ANDROID_TRACKER_KEY",
-        ),
-        pytest.mark.skipif(
-            os.environ.get(ANDROID_REMOTE_DIR) == None,
-            reason="Missing environment variable, ANDROID_REMOTE_DIR",
-        ),
-    ]
-
-    return _compose(args, _requires_rpc_tracker)
-
-
-def requires_rpc_tracker(*args):
-    """Mark a test as requiring an RPC tracker to exist in
-    the host environment to run."""
-    _requires_rpc_tracker = [
-        *tvm.testing.requires_rpc(),
-        pytest.mark.skipif(
-            os.environ.get("TVM_TRACKER_HOST") == None,
-            reason="Missing environment variable, TVM_TRACKER_HOST",
-        ),
-        pytest.mark.skipif(
-            os.environ.get("TVM_TRACKER_PORT") == None,
-            reason="Missing environment variable, TVM_TRACKER_PORT",
-        ),
-    ]
-
-    return _compose(args, _requires_rpc_tracker)
+    port = os.getenv(TVM_TRACKER_PORT, default=None)
+    port = int(port) if port else None
+    return port
diff --git a/tests/python/contrib/test_hexagon/test_cache_read_write.py b/tests/python/contrib/test_hexagon/test_cache_read_write.py
new file mode 100644
index 000000000000..38f62a036e5e
--- /dev/null
+++ b/tests/python/contrib/test_hexagon/test_cache_read_write.py
@@ -0,0 +1,135 @@
+# 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.
+
+import pytest
+import numpy as np
+
+import tvm.testing
+from tvm import te
+from tvm.contrib import utils
+from tvm.contrib.hexagon.build import HexagonLauncher
+import tvm.contrib.hexagon.hexagon as hexagon
+
+from .conftest import requires_hexagon_toolchain
+
+
+def intrin_mem_copy(shape, dtype, dst_scope, src_scope):
+    assert len(shape) == 1
+    src = te.placeholder(shape=shape, dtype=dtype, name="src")
+    dst = te.compute(shape, lambda i: src[i], name="dst")
+    size = shape[0] * np.dtype(dtype).itemsize
+
+    src_buffer = tvm.tir.decl_buffer(
+        shape,
+        dtype,
+        scope=src_scope,
+        offset_factor=1,
+    )
+
+    dst_buffer = tvm.tir.decl_buffer(
+        shape,
+        dtype,
+        scope=dst_scope,
+        offset_factor=1,
+    )
+
+    def intrin_func(ins, outs):
+        ib = tvm.tir.ir_builder.create()
+
+        _src = ins[0]
+        _dst = outs[0]
+        ib.emit(
+            tvm.tir.call_intrin(
+                "handle", "tir.mem_copy", _dst.access_ptr("w"), _src.access_ptr("r"), size
+            )
+        )
+        return ib.get()
+
+    return te.decl_tensor_intrin(dst.op, intrin_func, binds={src: src_buffer, dst: dst_buffer})
+
+
+@requires_hexagon_toolchain
+def test_cache_read_write(android_serial_number, tvm_tracker_host, tvm_tracker_port):
+    size = 128
+    outer_shape = (size,)
+    factor = 16
+    inner_shape = (factor,)
+    dtype = "int8"
+
+    x = te.placeholder(shape=outer_shape, dtype=dtype, name="x")
+    y = te.placeholder(shape=outer_shape, dtype=dtype, name="y")
+    z = te.compute(outer_shape, lambda i: x[i] + y[i], name="z")
+    s = te.create_schedule(z.op)
+
+    x_global = s.cache_read(x, "global.vtcm", [z])
+    y_global = s.cache_read(y, "global.vtcm", [z])
+    z_global = s.cache_write(z, "global.vtcm")
+
+    zouter, zinner = s[z_global].split(z_global.op.axis[0], factor=factor)
+
+    s[x_global].compute_at(s[z_global], zouter)
+    s[y_global].compute_at(s[z_global], zouter)
+
+    mem_copy_read = intrin_mem_copy(inner_shape, dtype, "global.vtcm", "global")
+
+    (cache_read_x,) = s[x_global].op.axis
+    s[x_global].tensorize(cache_read_x, mem_copy_read)
+
+    (cache_read_y,) = s[y_global].op.axis
+    s[y_global].tensorize(cache_read_y, mem_copy_read)
+
+    mem_copy_write = intrin_mem_copy(outer_shape, dtype, "global", "global.vtcm")
+
+    (cache_write_z,) = s[z].op.axis
+    s[z].tensorize(cache_write_z, mem_copy_write)
+
+    print(tvm.lower(s, [x, y, z]))
+
+    target_hexagon = tvm.target.hexagon("v68", link_params=True)
+    func = tvm.build(
+        s, [x, y, z], tvm.target.Target(target_hexagon, host=target_hexagon), name="dmacpy"
+    )
+    temp = utils.tempdir()
+    dso_binary = "test_binary.so"
+    dso_binary_path = temp.relpath(dso_binary)
+    func.save(dso_binary_path)
+
+    if not android_serial_number:
+        pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
+
+    launcher = HexagonLauncher(serial_number=android_serial_number)
+    launcher.android_run_rpc(rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port)
+    launcher.hexagon_setup()
+    remote_kw = {
+        "host": tvm_tracker_host,
+        "port": tvm_tracker_port,
+        "priority": 0,
+        "timeout": 60,
+    }
+    launcher.hexagon_session_setup(remote_kw)
+    launcher.upload(dso_binary_path, dso_binary)
+
+    with launcher.session as sess:
+        mod = launcher.get_module(dso_binary)
+        xt = tvm.nd.array(np.random.uniform(size=size).astype(x.dtype), device=sess.device)
+        yt = tvm.nd.array(np.random.uniform(size=size).astype(y.dtype), device=sess.device)
+        zt = tvm.nd.array(np.random.uniform(size=size).astype(z.dtype), device=sess.device)
+        mod["dmacpy"](xt, yt, zt)
+    launcher.close()
+
+    ref = xt.numpy() + yt.numpy()
+    np.testing.assert_equal(zt.numpy(), ref)
diff --git a/tests/python/contrib/test_hexagon/rpc/test_launcher.md b/tests/python/contrib/test_hexagon/test_launcher.md
similarity index 87%
rename from tests/python/contrib/test_hexagon/rpc/test_launcher.md
rename to tests/python/contrib/test_hexagon/test_launcher.md
index bcf255e478f1..8744071e63a4 100644
--- a/tests/python/contrib/test_hexagon/rpc/test_launcher.md
+++ b/tests/python/contrib/test_hexagon/test_launcher.md
@@ -19,12 +19,12 @@
 HexagonLauncher is a class to handle interactions with an Android phone which includes Hexagon DSP to run a TVMModule(function/operation/graph) on Hexagon. HexagonLauncher reuses minRPC implementation to setup an RPC connection from host (your local machine) to Hexagon target which is passed through Android RPC server.
 
 ## Build Required Tools/Libraries
-Here are the steps that are taken to prepare a runtime on a Hexagon device to test any model.
+To build TVM for Hexagon and run tests you can follow these steps to prepare a runtime on a Hexagon device to test any model. Alternatively, you can skip these instructions and use docker image which has pre-installed required tools. Instructions for using docker image [here](#use-hexagon-docker-image).
 
 - Build TVMRuntime library and C++ RPC server for Android.
 - Build minRPC server along with FastRPC for Hexagon.
 - Build TVM library with Hexagon support for host machine.
-- Build TVMRuntime library and C++ RPC server for host machine.
+- Build TVMRuntime library and RPC server for host machine.
 
 Note: First, ensure to export Clang libraries to `LD_LIBRARY_PATH` and Hexagon toolchain to `HEXAGON_TOOLCHAIN`:
 
@@ -58,7 +58,7 @@ cd tvm
 mkdir build
 cd build
 cmake -DUSE_LLVM="path to `llvm/bin/llvm-config`" \
-        -DUSE_CPP_RPC=ON \
+        -DUSE_RPC=ON \
         -DCMAKE_CXX_COMPILER="path to `clang++` executable" \
         -DCMAKE_CXX_FLAGS='-stdlib=libc++' \
         -DUSE_HEXAGON_SDK="path to Hexagon SDK" \
@@ -66,6 +66,25 @@ cmake -DUSE_LLVM="path to `llvm/bin/llvm-config`" \
         -DUSE_HEXAGON_DEVICE=sim ..
 ```
 
+## Use Hexagon Docker Image
+To use this docker image, install TVM and tools follow these steps.
+
+```bash
+# Log in to docker image
+cd tvm
+./docker/bash.sh tlcpack/ci-hexagon:v0.01
+
+# Build TVM
+./tests/scripts/task_config_build_hexagon.sh 
+cd build
+cmake ..
+make -j2
+
+# Build Hexagon API
+cd ..
+./tests/scripts/task_build_hexagon_api.sh 
+```
+
 ## Testing Using HexagonLauncher
 Before starting a test you need to run an RPC tracker on your local machine and export HOST and PORT as environment variables. Also, you need to export Clang libraries to `LD_LIBRARY_PATH` and Hexagon toolchain to `HEXAGON_TOOLCHAIN` as explained above.
 
diff --git a/tests/python/contrib/test_hexagon/rpc/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py
similarity index 91%
rename from tests/python/contrib/test_hexagon/rpc/test_launcher.py
rename to tests/python/contrib/test_hexagon/test_launcher.py
index 15dfac669046..1c322128e9fb 100644
--- a/tests/python/contrib/test_hexagon/rpc/test_launcher.py
+++ b/tests/python/contrib/test_hexagon/test_launcher.py
@@ -18,7 +18,7 @@
 import sys
 import pytest
 import numpy as np
-import os
+import logging
 
 import tvm.testing
 from tvm import te
@@ -28,12 +28,11 @@
 from tvm.contrib.hexagon.build import HexagonLauncher
 import tvm.contrib.hexagon.hexagon as hexagon
 
-from ..conftest import requires_rpc_tracker, requires_hexagon_toolchain
+from .conftest import requires_hexagon_toolchain
 
 
-@requires_rpc_tracker
 @requires_hexagon_toolchain
-def test_add(tvm_tracker_host, tvm_tracker_port, android_serial_number):
+def test_add(android_serial_number, tvm_tracker_host, tvm_tracker_port):
     dtype = "int8"
     A = tvm.te.placeholder((2,), dtype=dtype)
     B = tvm.te.placeholder((1,), dtype=dtype)
@@ -50,6 +49,9 @@ def test_add(tvm_tracker_host, tvm_tracker_port, android_serial_number):
     dso_binary_path = temp.relpath(dso_binary)
     func.save(dso_binary_path)
 
+    if not android_serial_number:
+        pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
+
     launcher = HexagonLauncher(serial_number=android_serial_number)
     launcher.android_run_rpc(rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port)
     launcher.hexagon_setup()
@@ -76,9 +78,8 @@ def test_add(tvm_tracker_host, tvm_tracker_port, android_serial_number):
     launcher.close()
 
 
-@requires_rpc_tracker
 @requires_hexagon_toolchain
-def test_add_vtcm(tvm_tracker_host, tvm_tracker_port, android_serial_number):
+def test_add_vtcm(android_serial_number, tvm_tracker_host, tvm_tracker_port):
     dtype = "int8"
     A = tvm.te.placeholder((2,), dtype=dtype)
     B = tvm.te.placeholder((1,), dtype=dtype)
@@ -95,6 +96,9 @@ def test_add_vtcm(tvm_tracker_host, tvm_tracker_port, android_serial_number):
     dso_binary_path = temp.relpath(dso_binary)
     func.save(dso_binary_path)
 
+    if not android_serial_number:
+        pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
+
     launcher = HexagonLauncher(serial_number=android_serial_number)
     launcher.android_run_rpc(rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port)
     launcher.hexagon_setup()
@@ -129,9 +133,8 @@ class TestMatMul:
     N = tvm.testing.parameter(32)
     K = tvm.testing.parameter(32)
 
-    @requires_rpc_tracker
     @requires_hexagon_toolchain
-    def test_matmul(self, tvm_tracker_host, tvm_tracker_port, android_serial_number, M, N, K):
+    def test_matmul(self, android_serial_number, tvm_tracker_host, tvm_tracker_port, M, N, K):
         X = te.placeholder((M, K), dtype="float32")
         Y = te.placeholder((K, N), dtype="float32")
         k1 = te.reduce_axis((0, K), name="k1")
@@ -148,6 +151,9 @@ def test_matmul(self, tvm_tracker_host, tvm_tracker_port, android_serial_number,
         dso_binary_path = temp.relpath(dso_binary)
         func.save(dso_binary_path)
 
+        if not android_serial_number:
+            pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
+
         launcher = HexagonLauncher(serial_number=android_serial_number)
         launcher.android_run_rpc(
             rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port
@@ -185,9 +191,8 @@ def test_matmul(self, tvm_tracker_host, tvm_tracker_port, android_serial_number,
         tvm.testing.assert_allclose(zt.numpy(), ztcpu.numpy(), rtol=1e-4)
 
 
-@requires_rpc_tracker
 @requires_hexagon_toolchain
-def test_graph_executor(tvm_tracker_host, tvm_tracker_port, android_serial_number):
+def test_graph_executor(android_serial_number, tvm_tracker_host, tvm_tracker_port):
     dtype = "float32"
     data = relay.var("data", relay.TensorType((1, 64, 64, 3), dtype))
     weight = relay.var("weight", relay.TensorType((5, 5, 3, 8), dtype))
@@ -221,6 +226,9 @@ def test_graph_executor(tvm_tracker_host, tvm_tracker_port, android_serial_numbe
         )
         lowered.get_lib().save(dso_binary_path)
 
+    if not android_serial_number:
+        pytest.skip("Skip hardware test since ANDROID_SERIAL_NUMBER is not set.")
+
     launcher = HexagonLauncher(serial_number=android_serial_number)
     launcher.android_run_rpc(rpc_tracker_host=tvm_tracker_host, rpc_tracker_port=tvm_tracker_port)
     launcher.hexagon_setup()
diff --git a/tests/python/driver/tvmc/test_command_line.py b/tests/python/driver/tvmc/test_command_line.py
index 66a32160522b..6830cf0503c0 100644
--- a/tests/python/driver/tvmc/test_command_line.py
+++ b/tests/python/driver/tvmc/test_command_line.py
@@ -47,7 +47,7 @@ def test_tvmc_cl_workflow(keras_simple, tmpdir_factory):
 
     # Test running the model
     output_path = os.path.join(tmpdir, "predictions.npz")
-    run_str = f"tvmc run --outputs {output_path} {package_path}"
+    run_str = f"tvmc run --end-to-end --outputs {output_path} {package_path}"
     run_args = run_str.split(" ")[1:]
     _main(run_args)
     assert os.path.exists(output_path)
diff --git a/tests/python/driver/tvmc/test_model.py b/tests/python/driver/tvmc/test_model.py
index fd2637a85f1f..d0d398b75521 100644
--- a/tests/python/driver/tvmc/test_model.py
+++ b/tests/python/driver/tvmc/test_model.py
@@ -30,7 +30,7 @@ def test_tvmc_workflow(keras_simple):
     tvmc_model = tvmc.load(keras_simple)
     tuning_records = tvmc.tune(tvmc_model, target="llvm", enable_autoscheduler=True, trials=2)
     tvmc_package = tvmc.compile(tvmc_model, tuning_records=tuning_records, target="llvm")
-    result = tvmc.run(tvmc_package, device="cpu")
+    result = tvmc.run(tvmc_package, device="cpu", end_to_end=True)
     assert type(tvmc_model) is TVMCModel
     assert type(tvmc_package) is TVMCPackage
     assert type(result) is TVMCResult
diff --git a/tests/python/driver/tvmc/test_pass_config.py b/tests/python/driver/tvmc/test_pass_config.py
index bb815e1dc8aa..f928c8a31293 100644
--- a/tests/python/driver/tvmc/test_pass_config.py
+++ b/tests/python/driver/tvmc/test_pass_config.py
@@ -16,11 +16,13 @@
 # under the License.
 
 import pytest
+from unittest import mock
 
 from tvm.contrib.target.vitis_ai import vitis_ai_available
 
 from tvm.driver.tvmc import TVMCException
 from tvm.driver.tvmc.pass_config import parse_configs
+from tvm.tir.transform import PrimFuncPass
 
 
 def test_config_invalid_format():
@@ -71,3 +73,89 @@ def test_config_valid_multiple_configs():
     assert configs["tir.detect_global_barrier"] == 10
     assert "relay.ext.vitis_ai.options.build_dir" in configs.keys()
     assert configs["relay.ext.vitis_ai.options.build_dir"] == "mystring"
+
+
+def test_add_lower_pass_multi_built_in_pass():
+    configs = parse_configs(
+        [
+            "tir.add_lower_pass=1,tir.transform.UnrollLoop",
+            "tir.add_lower_pass=1,tir.transform.HoistIfThenElse,2,tir.transform.LoopPartition",
+        ]
+    )
+
+    assert len(configs["tir.add_lower_pass"]) == 3
+    # opt_level: 1, pass: tir.transform.UnrollLoop
+    assert configs["tir.add_lower_pass"][0][0] == 1
+    assert isinstance(configs["tir.add_lower_pass"][0][1], PrimFuncPass)
+    # opt_level: 1, pass: tir.transform.HoistIfThenElse
+    assert configs["tir.add_lower_pass"][1][0] == 1
+    assert isinstance(configs["tir.add_lower_pass"][1][1], PrimFuncPass)
+    # opt_level: 2, pass: tir.transform.LoopPartition
+    assert configs["tir.add_lower_pass"][2][0] == 2
+    assert isinstance(configs["tir.add_lower_pass"][2][1], PrimFuncPass)
+
+
+def test_add_lower_pass_multi_external_pass():
+    fake_pass_1 = mock.MagicMock()
+    fake_pass_2 = mock.MagicMock()
+    fake_pass_3 = mock.MagicMock()
+    with mock.patch.dict(
+        "sys.modules",
+        {"fake_module": fake_pass_1, "fake_module": fake_pass_2, "fake_module": fake_pass_3},
+    ):
+        configs = parse_configs(
+            [
+                "tir.add_lower_pass=1,fake_module.fake_pass_1,2,fake_module.fake_pass2",
+                "tir.add_lower_pass=3,fake_module.fake_pass_3",
+            ]
+        )
+        assert len(configs["tir.add_lower_pass"]) == 3
+        # opt_level: 1, pass: fake_module.fake_pass_1
+        assert configs["tir.add_lower_pass"][0][0] == 1
+        # opt_level: 2, pass: fake_module.fake_pass_2
+        assert configs["tir.add_lower_pass"][1][0] == 2
+        # opt_level: 3, pass: fake_module.fake_pass_3
+        assert configs["tir.add_lower_pass"][2][0] == 3
+
+
+def test_add_lower_pass_multi_mix_pass():
+    fake_pass_1 = mock.MagicMock()
+    fake_pass_2 = mock.MagicMock()
+    with mock.patch.dict("sys.modules", {"fake_module": fake_pass_1, "fake_module": fake_pass_2}):
+        configs = parse_configs(
+            [
+                "tir.add_lower_pass=1,fake_module.fake_pass_1,1,tir.transform.UnrollLoop",
+                "tir.add_lower_pass=2,fake_module.fake_pass_2,2,tir.transform.LoopPartition",
+            ]
+        )
+        assert len(configs["tir.add_lower_pass"]) == 4
+        # opt_level: 1, pass: fake_module.fake_pass_1
+        assert configs["tir.add_lower_pass"][0][0] == 1
+        # opt_level: 1, pass: tir.transform.UnrollLoop
+        assert configs["tir.add_lower_pass"][1][0] == 1
+        assert isinstance(configs["tir.add_lower_pass"][1][1], PrimFuncPass)
+        # opt_level: 2, pass: fake_module.fake_pass_2
+        assert configs["tir.add_lower_pass"][2][0] == 2
+        # opt_level: 2, pass: tir.transform.LoopPartition
+        assert configs["tir.add_lower_pass"][3][0] == 2
+        assert isinstance(configs["tir.add_lower_pass"][3][1], PrimFuncPass)
+
+
+def test_add_lower_pass_invalid_format():
+    # wrong format
+    with pytest.raises(TVMCException):
+        _ = parse_configs(["tir.add_lower_pass=tir.transform.UnrollLoop,1"])
+    # missing pass name
+    with pytest.raises(TVMCException):
+        _ = parse_configs(["tir.add_lower_pass=1,tir.transform.UnrollLoop,3"])
+    # wrong opt level
+    with pytest.raises(TVMCException):
+        _ = parse_configs(["tir.add_lower_pass=a,tir.transform.UnrollLoop"])
+    # fake module
+    with pytest.raises(ModuleNotFoundError):
+        _ = parse_configs(
+            ["tir.add_lower_pass=1,tir.transform.UnrollLoop,2,path.to.module.fake_func"]
+        )
+    # real module and fake func
+    with pytest.raises(TVMCException):
+        _ = parse_configs(["tir.add_lower_pass=1,tir.transform.UnrollLoop,2,tvm.tir.fake_func"])
diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py
index 25631f4ab2cd..84e3ddbf53d2 100644
--- a/tests/python/frontend/onnx/test_forward.py
+++ b/tests/python/frontend/onnx/test_forward.py
@@ -5017,7 +5017,6 @@ def verify_eyelike(indata):
 )
 
 unsupported_onnx_tests = [
-    "test_basic_convinteger",
     "test_batchnorm_epsilon_training_mode",
     "test_batchnorm_example_training_mode",
     "test_bernoulli",
@@ -5046,8 +5045,6 @@ def verify_eyelike(indata):
     "test_castlike_FLOAT_to_STRING_expanded",
     "test_castlike_STRING_to_FLOAT",
     "test_castlike_STRING_to_FLOAT_expanded",
-    "test_convinteger_with_padding",
-    "test_convinteger_without_padding",
     "test_convtranspose_autopad_same",
     "test_convtranspose_dilations",
     "test_convtranspose_output_shape",
@@ -5072,8 +5069,6 @@ def verify_eyelike(indata):
     "test_loop13_seq",
     "test_lstm_batchwise",
     "test_matmulinteger",
-    "test_maxpool_2d_same_lower",
-    "test_maxpool_2d_same_upper",
     "test_maxpool_with_argmax_2d_precomputed_pads",
     "test_maxpool_with_argmax_2d_precomputed_strides",
     "test_maxunpool_export_with_output_shape",
@@ -5413,6 +5408,7 @@ def verify_qlinearconv(
         dilations,
         auto_pad="NOTSET",
         bias=False,
+        per_channel_quantization=False,
     ):
 
         x_array = np.random.randint(low=0, high=255, size=x_shape).astype("uint8")
@@ -5421,8 +5417,6 @@ def verify_qlinearconv(
         initializer = [
             helper.make_tensor("x_scale", TensorProto.FLOAT, (), [np.random.rand()]),
             helper.make_tensor("x_zero_point", TensorProto.UINT8, (), [np.random.randint(0, 255)]),
-            helper.make_tensor("w_scale", TensorProto.FLOAT, (), [np.random.rand()]),
-            helper.make_tensor("w_zero_point", TensorProto.UINT8, (), [np.random.randint(0, 255)]),
             helper.make_tensor("y_scale", TensorProto.FLOAT, (), [np.random.rand()]),
             helper.make_tensor("y_zero_point", TensorProto.UINT8, (), [np.random.randint(0, 255)]),
         ]
@@ -5443,6 +5437,28 @@ def verify_qlinearconv(
         ]
         input_values = [x_array, w_array]
 
+        if per_channel_quantization:
+            w_scale_array = np.random.random(w_shape[0]).astype("float32")
+            w_zero_point_array = np.random.randint(0, 255, size=w_shape[0]).astype("uint8")
+
+            initializer.append(
+                helper.make_tensor("w_scale", TensorProto.FLOAT, [w_shape[0]], w_scale_array)
+            )
+            initializer.append(
+                helper.make_tensor(
+                    "w_zero_point", TensorProto.UINT8, [w_shape[0]], w_zero_point_array
+                )
+            )
+        else:
+            initializer.append(
+                helper.make_tensor("w_scale", TensorProto.FLOAT, (), [np.random.rand()])
+            )
+            initializer.append(
+                helper.make_tensor(
+                    "w_zero_point", TensorProto.UINT8, (), [np.random.randint(0, 255)]
+                )
+            )
+
         if bias is True:
             b_shape = w_shape[0:1]
             b_array = np.random.randint(low=0, high=65536, size=b_shape).astype("int32")
@@ -5582,6 +5598,17 @@ def repeat(N, D):
         repeat(1, D),
         repeat(2, D),
     )
+    # Convolution with per channel quantization
+    verify_qlinearconv(
+        (1, 1) + repeat(5, D),
+        (1, 1) + repeat(3, D),
+        (1, 1) + repeat(3, D),
+        None,
+        repeat(3, D),
+        repeat(1, D),
+        repeat(1, D),
+        per_channel_quantization=True,
+    )
 
 
 @tvm.testing.parametrize_targets
diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
index 980e31420362..567f4ea48217 100644
--- a/tests/python/frontend/pytorch/test_forward.py
+++ b/tests/python/frontend/pytorch/test_forward.py
@@ -1067,6 +1067,32 @@ def test_forward_conv_transpose(
     verify_model(conv1d_transpose, conv1d_input_data)
 
 
+def test_forward_conv2d_transpose_group():
+    # https://github.com/apache/tvm/issues/10223
+
+    class ModulatedConvTranspose2D(torch.nn.Module):
+        def forward(self, x, w, s):
+            B, C, H, W = x.shape
+            I, O, KH, KW = w.shape
+
+            # weight is different for each input in batch (this is why we want grouped conv transpose)
+            w = w.unsqueeze(0) * s.reshape(B, 1, 1, 1, 1)
+            w = w.reshape(B * I, O, KH, KW)
+            x = x.reshape(1, B * C, H, W)
+            x = torch.nn.functional.conv_transpose2d(
+                x, w, stride=(2, 2), padding=(1, 1), output_padding=(1, 1), groups=B
+            )
+            return x.reshape(B, O, H * 2, W * 2)
+
+    b, c, h, w, k = 4, 512, 8, 16, 3
+    inputs = torch.rand(b, c, h, w)
+    weights = torch.rand(c, c // 2, k, k)
+    styles = torch.rand(b)
+
+    # cuda not supported for group > 1 conv2d_transpose
+    verify_trace_model(ModulatedConvTranspose2D().eval(), [inputs, weights, styles], ["llvm"])
+
+
 def test_forward_deform_conv():
     torch.set_grad_enabled(False)
 
@@ -4168,7 +4194,7 @@ def forward(self, x):
 
     x = torch.rand([4, 4, 16, 32]).float()
     script_module = torch.jit.trace(List_tuple(), x, strict=False).eval()
-    mod, params = relay.frontend.from_pytorch(script_module, [("x", x.shape)])
+    relay.frontend.from_pytorch(script_module, [("x", x.shape)])
 
 
 if __name__ == "__main__":
diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py
index 2a744e1de7c3..f8e819dc015e 100644
--- a/tests/python/frontend/tensorflow/test_forward.py
+++ b/tests/python/frontend/tensorflow/test_forward.py
@@ -565,6 +565,7 @@ def _test_convolution(
             )
 
 
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/10275")
 @tvm.testing.uses_gpu
 def test_forward_convolution():
     if is_gpu_available():
@@ -3769,6 +3770,7 @@ def test_forward_where():
 #######################################################################
 # Inception V3
 # ------------
+@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/10275")
 def test_forward_inception_v3():
     """test inception V3 model"""
     with tf.Graph().as_default():
diff --git a/tests/python/frontend/tflite/test_forward.py b/tests/python/frontend/tflite/test_forward.py
index 60af94b53a51..abb05354d921 100644
--- a/tests/python/frontend/tflite/test_forward.py
+++ b/tests/python/frontend/tflite/test_forward.py
@@ -618,6 +618,9 @@ def test_forward_stridedslice():
         _test_stridedslice(
             (4, 4), [1, 0], [4, 4], [1, 1], "float32", shrink_axis_mask=2, quantized=quantized
         )
+        _test_stridedslice(
+            (3, 4), [-1, 0], [0, 3], [1, 1], "float32", shrink_axis_mask=1, quantized=quantized
+        )
 
 
 #######################################################################
@@ -3186,6 +3189,9 @@ def test_forward_unpack():
     """UNPACK"""
     _test_unpack(np.array(np.random.uniform(0, 5, (3, 1)), dtype=np.int32), axis=1, num_unpacks=1)
     _test_unpack(np.array(np.random.uniform(0, 5, (3, 4)), dtype=np.float32), axis=0, num_unpacks=3)
+    _test_unpack(
+        np.array(np.random.uniform(0, 5, (3, 1, 2)), dtype=np.float32), axis=0, num_unpacks=3
+    )
     # tflite 1.13 doesn't accept negative axis
     if package_version.parse(tf.VERSION) >= package_version.parse("1.14.0"):
         _test_unpack(
diff --git a/tests/python/relay/aot/aot_test_utils.py b/tests/python/relay/aot/aot_test_utils.py
index 7cbf74e9a811..63817fc4b965 100644
--- a/tests/python/relay/aot/aot_test_utils.py
+++ b/tests/python/relay/aot/aot_test_utils.py
@@ -59,7 +59,7 @@ class AOTTestModel(NamedTuple):
     inputs: Dict[str, np.array]
         Dict of input names to value arrays
     outputs: List[np.array]
-        Ordered list of output value arrays
+        Dict of output names to value arrays
     output_tolerance: Optional[Union[int, float]]
         Allowed tolerance of the output
     name: str
@@ -72,7 +72,7 @@ class AOTTestModel(NamedTuple):
 
     module: tvm.IRModule
     inputs: Dict[str, np.array]
-    outputs: List[np.array]
+    outputs: Dict[str, np.array]
     output_tolerance: Optional[Union[int, float]] = None
     name: str = "default"
     params: Optional[Dict[str, np.array]] = None
@@ -265,21 +265,29 @@ def emit_data_linkage(output_file, data_linkage):
 
 
 def emit_main_prologue(
-    main_file, custom_prologue, workspace_bytes, data_linkage, compiled_models, interface_api
+    main_file,
+    custom_prologue,
+    workspace_bytes,
+    data_linkage,
+    compiled_models,
+    interface_api,
+    use_stack_allocator=True,
 ):
-    # Add TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES because of memory alignment.
-    workspace_define = f"#define WORKSPACE_SIZE ({workspace_bytes}"
-    if interface_api == "c":
-        for compiled_model in compiled_models:
-            model = compiled_model.model
-            workspace_define += f" + TVMGEN_{model.name.upper()}_WORKSPACE_SIZE"
-    workspace_define += " + TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES)\n"
-    main_file.write(workspace_define)
-    emit_data_linkage(main_file, data_linkage)
-    main_file.write("static uint8_t g_aot_memory[WORKSPACE_SIZE];\n")
-    main_file.write("tvm_workspace_t app_workspace;\n")
-    main_file.write(
-        """
+    if use_stack_allocator:
+        workspace_define = f"#define WORKSPACE_SIZE ({workspace_bytes}"
+        if interface_api == "c":
+            for compiled_model in compiled_models:
+                model = compiled_model.model
+                workspace_define += f" + TVMGEN_{model.name.upper()}_WORKSPACE_SIZE"
+        # Add TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES because of memory alignment.
+        workspace_define += " + TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES)\n"
+        main_file.write(workspace_define)
+        emit_data_linkage(main_file, data_linkage)
+        main_file.write("static uint8_t g_aot_memory[WORKSPACE_SIZE];\n")
+        main_file.write("tvm_workspace_t app_workspace;\n")
+        main_file.write(
+            """
+            
 tvm_crt_error_t TVMPlatformMemoryAllocate(size_t num_bytes, DLDevice dev, void** out_ptr) {
     return StackMemoryManager_Allocate(&app_workspace, num_bytes, out_ptr);
 }
@@ -287,7 +295,26 @@ def emit_main_prologue(
 tvm_crt_error_t TVMPlatformMemoryFree(void* ptr, DLDevice dev) {
     return StackMemoryManager_Free(&app_workspace,ptr);
 }
+        """
+        )
+    else:
+        # An implementation is not needed for these if the stack allocator is not used
+        main_file.write(
+            """
+            
+tvm_crt_error_t TVMPlatformMemoryAllocate(size_t num_bytes, DLDevice dev, void** out_ptr) {
+    return kTvmErrorFunctionCallNotImplemented;
+}
+
+tvm_crt_error_t TVMPlatformMemoryFree(void* ptr, DLDevice dev) {
+    return kTvmErrorFunctionCallNotImplemented;
+}
 
+            """
+        )
+    main_file.write(
+        """
+    
 void TVMPlatformAbort(tvm_crt_error_t code) { exit(-1); }
 
 void TVMLogf(const char* msg, ...) {
@@ -296,24 +323,27 @@ def emit_main_prologue(
   vfprintf(stdout, msg, args);
   va_end(args);
 }
-
+    
 TVM_DLL int TVMFuncRegisterGlobal(const char* name, TVMFunctionHandle f, int override) {}
 int main(){\n
-"""
+    """
     )
     main_file.write(custom_prologue)
 
 
-def emit_main_data(main_file, input_map, output_list, mod_name):
+def emit_main_data(main_file, input_map, output_map, mod_name):
     for key in input_map:
         sanitized_tensor_name = re.sub(r"\W", "_", key)
         main_file.write(
             f'#include "{mangle_name(mod_name,"input_data")}_{sanitized_tensor_name}.h"\n'
         )
 
-    for i in range(0, len(output_list)):
-        main_file.write(f'#include "{mangle_name(mod_name,"expected_output_data")}{i}.h"\n')
-        main_file.write(f'#include "{mangle_name(mod_name,"output_data")}{i}.h"\n')
+    for key in output_map:
+        sanitized_tensor_name = re.sub(r"\W", "_", key)
+        main_file.write(
+            f'#include "{mangle_name(mod_name,"expected_output_data")}_{sanitized_tensor_name}.h"\n'
+            f'#include "{mangle_name(mod_name,"output_data")}_{sanitized_tensor_name}.h"\n'
+        )
 
 
 def emit_main_device_structs(main_file, devices, mod_name):
@@ -336,7 +366,7 @@ def emit_main_workspace_pool_structs(main_file, workspace_pool_names, mod_name):
         main_file.write("};\n")
 
 
-def emit_main_data_structs(main_file, input_map, output_list, mod_name):
+def emit_main_data_structs(main_file, input_map, output_map, mod_name):
     main_file.write(
         f"struct {mangle_name(mod_name, 'inputs')} {mangle_name(mod_name, 'inputs')} = {{"
     )
@@ -350,17 +380,16 @@ def emit_main_data_structs(main_file, input_map, output_list, mod_name):
     main_file.write(
         f"struct {mangle_name(mod_name, 'outputs')} {mangle_name(mod_name, 'outputs')} = {{"
     )
-    num_outputs = len(output_list)
-    if num_outputs == 1:
-        main_file.write(f"\t.output = {mangle_name(mod_name, 'output_data')}0,\n")
-    else:
-        for i in range(0, num_outputs):
-            main_file.write(f"\t.output{i} = {mangle_name(mod_name, 'output_data')}{i},\n")
+    for key in output_map:
+        sanitized_tensor_name = re.sub(r"\W", "_", key)
+        main_file.write(
+            f"\t.{sanitized_tensor_name} = {mangle_name(mod_name, 'output_data')}_{sanitized_tensor_name},\n"
+        )
     main_file.write("};\n")
 
 
-def emit_main_data_setup(main_file, input_map, output_list, mod_name):
-    num_outputs = len(output_list)
+def emit_main_data_setup(main_file, input_map, output_map, mod_name):
+    num_outputs = len(output_map)
     num_inputs = len(input_map)
 
     main_file.write(f'void* {mangle_name(mod_name,"inputs")}[{num_inputs}] = {{ ')
@@ -370,8 +399,9 @@ def emit_main_data_setup(main_file, input_map, output_list, mod_name):
     main_file.write("};\n")
 
     main_file.write(f'void* {mangle_name(mod_name,"outputs")}[{num_outputs}]  = {{ ')
-    for i in range(0, num_outputs):
-        main_file.write(f'{mangle_name(mod_name,"output_data")}{i}, ')
+    for key in output_map:
+        sanitized_tensor_name = re.sub(r"\W", "_", key)
+        main_file.write(f'{mangle_name(mod_name, "output_data")}_{sanitized_tensor_name}, ')
     main_file.write("};\n")
 
 
@@ -446,13 +476,12 @@ def fake_tensor(source, source_index, packed_index):
     main_file.write("\n")
 
 
-def emit_main_compare(main_file, output_list, output_tolerance, mod_name):
-    num_outputs = len(output_list)
-    actual_data_name = mangle_name(mod_name, "output_data")
-    expected_data_name = mangle_name(mod_name, "expected_output_data")
-
-    for i in range(0, num_outputs):
-        is_float_dtype = output_list[i].dtype == "float32"
+def emit_main_compare(main_file, outputs, output_tolerance, mod_name):
+    for key in outputs:
+        sanitized_tensor_name = re.sub(r"\W", "_", key)
+        actual_data_name = mangle_name(mod_name, f"output_data_{sanitized_tensor_name}")
+        expected_data_name = mangle_name(mod_name, f"expected_output_data_{sanitized_tensor_name}")
+        is_float_dtype = outputs[key].dtype == "float32"
 
         comparison_function = "abs"
         tolerance = output_tolerance or 0
@@ -462,8 +491,8 @@ def emit_main_compare(main_file, output_list, output_tolerance, mod_name):
 
         main_file.write(
             f"""
-            for (int i = 0; i<{actual_data_name}{i}_len; i++) {{
-                if ({comparison_function}({actual_data_name}{i}[i]-{expected_data_name}{i}[i]) > {tolerance}) {{
+            for (int i = 0; i<{actual_data_name}_len; i++) {{
+                if ({comparison_function}({actual_data_name}[i]-{expected_data_name}[i]) > {tolerance}) {{
                     printf("{AOT_FAILURE_TOKEN}\\n");
                     return -1;
                 }}
@@ -509,6 +538,7 @@ def create_main(
     data_linkage,
     interface_api,
     workspace_bytes,
+    use_stack_allocator=True,
 ):
     file_path = pathlib.Path(f"{output_path}/" + test_name).resolve()
     # create header file
@@ -531,8 +561,10 @@ def create_main(
             data_linkage,
             compiled_models,
             interface_api,
+            use_stack_allocator,
         )
-        emit_main_init_memory_manager(main_file)
+        if use_stack_allocator:
+            emit_main_init_memory_manager(main_file)
 
         if interface_api == "c":
             for compiled_model in compiled_models:
@@ -707,11 +739,14 @@ def run_and_check(
         t = tarfile.open(tar_file)
         t.extractall(base_path)
 
-        workspace_bytes = model.extra_memory_in_bytes
-        use_usmp = runner.pass_config.get("tir.usmp.enable", False)
-        if interface_api == "packed" and not use_usmp:
+        # Interface C APIs does not need compiler generated
+        # workspace to generate the test application, because
+        # workspace size is codegen'd as a macro to
+        # tvmgen_<model_name>.h.
+        if interface_api != "c":
             workspace_bytes += mlf_extract_workspace_size_bytes(tar_file)
 
+        workspace_bytes += model.extra_memory_in_bytes
         for key in model.inputs:
             sanitized_tensor_name = re.sub(r"\W", "_", key)
             create_header_file(
@@ -721,20 +756,25 @@ def run_and_check(
                 data_linkage,
             )
 
-        for i in range(len(model.outputs)):
+        for key in model.outputs:
+            sanitized_tensor_name = re.sub(r"\W", "_", key)
             create_header_file(
-                (f'{mangle_name(model.name,"output_data")}{i}'),
-                np.zeros(model.outputs[i].shape, model.outputs[i].dtype),
+                f'{mangle_name(model.name, "output_data")}_{sanitized_tensor_name}',
+                np.zeros(model.outputs[key].shape, model.outputs[key].dtype),
                 include_path,
                 data_linkage,
             )
             create_header_file(
-                (f'{mangle_name(model.name, "expected_output_data")}{i}'),
-                model.outputs[i],
+                f'{mangle_name(model.name, "expected_output_data")}_{sanitized_tensor_name}',
+                model.outputs[key],
                 include_path,
                 data_linkage,
             )
 
+    use_usmp = runner.pass_config.get("tir.usmp.enable", False)
+    # We only need the stack allocator if USMP is not used
+    use_stack_allocator = not use_usmp
+
     create_main(
         "test.c",
         models,
@@ -745,6 +785,7 @@ def run_and_check(
         data_linkage,
         interface_api,
         workspace_bytes,
+        use_stack_allocator,
     )
 
     # Verify that compiles fine
@@ -839,7 +880,7 @@ def compile_and_run(
 
 def generate_ref_data(mod, input_data, params=None, target="llvm"):
     """Generate reference data through executing the relay module"""
-    with tvm.transform.PassContext(opt_level=3):
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
         lib = relay.build(mod, target=target, params=params)
 
     lib_name = "mod.so"
@@ -852,4 +893,35 @@ def generate_ref_data(mod, input_data, params=None, target="llvm"):
     grt_mod.run()
     output_count = grt_mod.get_num_outputs()
     out = [grt_mod.get_output(i).numpy() for i in range(output_count)]
-    return out
+    if isinstance(mod, tvm.relay.Function):
+        main = mod
+    else:
+        main = mod["main"]
+    if main.attrs == None or main.attrs["output_tensor_names"] == None:
+        if output_count == 1:
+            output_tensor_names = ["output"]
+        else:
+            output_tensor_names = [f"output{i}" for i in range(output_count)]
+    else:
+        output_tensor_names = main.attrs["output_tensor_names"]
+
+    return dict(zip(output_tensor_names, out))
+
+
+def create_relay_module_and_inputs_from_tflite_file(tflite_model_file):
+    """A helper function to create a Relay IRModule with inputs
+    and params from a tflite file"""
+    with open(tflite_model_file, "rb") as f:
+        tflite_model_buf = f.read()
+    mod, params = convert_to_relay(tflite_model_buf)
+
+    inputs = dict()
+    for param in mod["main"].params:
+        name = str(param.name_hint)
+        data_shape = [int(i) for i in param.type_annotation.shape]
+        dtype = str(param.type_annotation.dtype)
+        in_min, in_max = (np.iinfo(dtype).min, np.iinfo(dtype).max)
+        data = np.random.randint(in_min, high=in_max, size=data_shape, dtype=dtype)
+        inputs[name] = data
+
+    return mod, inputs, params
diff --git a/tests/python/relay/aot/test_c_device_api.py b/tests/python/relay/aot/test_c_device_api.py
index d369fd0a4a30..b84e5eb6d775 100644
--- a/tests/python/relay/aot/test_c_device_api.py
+++ b/tests/python/relay/aot/test_c_device_api.py
@@ -93,7 +93,7 @@ def compile_to_main_func(interface_api="c", use_unpacked_api=True):
             pass_config=test_runner.pass_config,
         )
         main_ir_module = compiled_models[0].executor_factory.lowered_ir_mods.items()[0][1]
-        main_func = main_ir_module["run_model"]
+        main_func = main_ir_module["__tvm_main__"]
         return main_func
 
     return compile_to_main_func
@@ -124,7 +124,7 @@ def compile_to_main_func(interface_api="c", use_unpacked_api=True):
             pass_config=test_runner.pass_config,
         )
         main_ir_module = list(compiled_models[0].executor_factory.lowered_ir_mods.values())[0]
-        main_func = main_ir_module["run_model"]
+        main_func = main_ir_module["__tvm_main__"]
         return main_func
 
     return compile_to_main_func
@@ -225,17 +225,24 @@ def test_without_device_api_packed_api(non_device_api_main_func):
     """Test a graph without the Device API with the packed internal calls"""
 
     main_func = non_device_api_main_func(interface_api="packed", use_unpacked_api=False)
-    assert (
-        str(main_func.body)
-        == 'let tvm_value_3 = tir.tvm_stack_alloca("array", 1)\n'
-        + 'let tvm_value_2 = tir.tvm_stack_alloca("array", 1)\n'
-        + 'let tvm_value_1 = tir.tvm_stack_alloca("array", 1)\n'
-        + 'let tvm_value_0 = tir.tvm_stack_alloca("array", 1)\n'
-        + "tir.tvm_struct_set(tvm_value_0, 0, 1, x_buffer_var)\n"
-        + "tir.tvm_struct_set(tvm_value_1, 0, 1, y_buffer_var)\n"
-        + "tir.tvm_struct_set(tvm_value_2, 0, 1, output_buffer_var)\n"
-        + "tir.tvm_struct_set(tvm_value_3, 0, 1, tir.reinterpret((uint64)0))\n"
-        + 'tir.tvm_call_cpacked("tvmgen_default_fused_multiply", tvm_value_0, tvm_value_1, tvm_value_2, tvm_value_3)\n'
+    assert str(main_func.body) == (
+        'let tvm_value_3 = tir.tvm_stack_alloca("array", 1)\n'
+        'let tvm_value_2 = tir.tvm_stack_alloca("array", 1)\n'
+        'let tvm_value_1 = tir.tvm_stack_alloca("array", 1)\n'
+        'let tvm_value_0 = tir.tvm_stack_alloca("array", 1)\n'
+        "tir.tvm_struct_set(tvm_value_0, 0, 1, x_buffer_var)\n"
+        "tir.tvm_struct_set(tvm_value_0, 0, 10, 1)\n"
+        "tir.tvm_struct_set(tvm_value_0, 0, 9, 0)\n"
+        "tir.tvm_struct_set(tvm_value_1, 0, 1, y_buffer_var)\n"
+        "tir.tvm_struct_set(tvm_value_1, 0, 10, 1)\n"
+        "tir.tvm_struct_set(tvm_value_1, 0, 9, 0)\n"
+        "tir.tvm_struct_set(tvm_value_2, 0, 1, output_buffer_var)\n"
+        "tir.tvm_struct_set(tvm_value_2, 0, 10, 1)\n"
+        "tir.tvm_struct_set(tvm_value_2, 0, 9, 0)\n"
+        "tir.tvm_struct_set(tvm_value_3, 0, 1, tir.reinterpret((uint64)0))\n"
+        "tir.tvm_struct_set(tvm_value_3, 0, 10, 1)\n"
+        "tir.tvm_struct_set(tvm_value_3, 0, 9, 0)\n"
+        'tir.tvm_call_cpacked("tvmgen_default_fused_multiply", tvm_value_0, tvm_value_1, tvm_value_2, tvm_value_3)\n'
     )
 
 
diff --git a/tests/python/relay/aot/test_cpp_aot.py b/tests/python/relay/aot/test_cpp_aot.py
new file mode 100644
index 000000000000..1a4a90c9cbc5
--- /dev/null
+++ b/tests/python/relay/aot/test_cpp_aot.py
@@ -0,0 +1,123 @@
+# 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.
+
+
+import sys
+import textwrap
+
+import numpy as np
+import pytest
+
+import tvm
+from tvm.relay import backend, testing
+from aot_test_utils import generate_ref_data
+
+
+enable_usmp = tvm.testing.parameter(True, False)
+target_kind = tvm.testing.parameter("c", "llvm")
+
+
+def test_conv2d(enable_usmp, target_kind):
+    RELAY_MODEL = textwrap.dedent(
+        """\
+        #[version = "0.0.5"]
+        def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(3, 3, 5, 5), int8]) {
+            %1 = nn.conv2d(
+                 %data,
+                 %weight,
+                 padding=[2, 2],
+                 channels=3,
+                 kernel_size=[5, 5],
+                 data_layout="NCHW",
+                 kernel_layout="OIHW",
+                 out_dtype="int32");
+            %2 = cast(nn.max_pool2d(%1, pool_size=[3, 3]), dtype="int8");
+            %3 = nn.conv2d(
+                 %2,
+                 %weight,
+                 padding=[2, 2],
+                 channels=3,
+                 kernel_size=[5, 5],
+                 data_layout="NCHW",
+                 kernel_layout="OIHW",
+                 out_dtype="int32");
+            %4 = nn.max_pool2d(%3, pool_size=[3, 3]);
+            %4
+        }
+    """
+    )
+    ir_mod = tvm.parser.fromtext(RELAY_MODEL)
+
+    main_func = ir_mod["main"]
+    shape_dict = {p.name_hint: p.checked_type.concrete_shape for p in main_func.params}
+    type_dict = {p.name_hint: p.checked_type.dtype for p in main_func.params}
+
+    weight_data = np.ones(shape_dict["weight"]).astype(type_dict["weight"])
+    input_data = np.ones(shape_dict["data"]).astype(type_dict["data"])
+
+    params = {"weight": weight_data}
+    inputs = {"data": input_data}
+    ref_outputs = generate_ref_data(ir_mod, inputs, params)
+
+    with tvm.transform.PassContext(
+        opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable": enable_usmp}
+    ):
+        mod = tvm.relay.build(
+            ir_mod,
+            params=params,
+            target=target_kind,
+            executor=backend.Executor("aot", {"interface-api": "packed"}),
+        )
+
+    temp_dir = tvm.contrib.utils.TempDirectory()
+    test_so_path = temp_dir / "test.so"
+    mod.export_library("test.tar")
+    mod.export_library(test_so_path, cc="gcc", options=["-std=c11"])
+    loaded_mod = tvm.runtime.load_module(test_so_path)
+    runner = tvm.runtime.executor.AotModule(loaded_mod["default"](tvm.cpu(0)))
+    runner.set_input(**inputs)
+    runner.run()
+    assert (runner.get_output(0).asnumpy() == list(ref_outputs.values())[0]).all()
+
+
+def test_mobilenet():
+    ir_mod, params = testing.mobilenet.get_workload(batch_size=1)
+    data_shape = [int(x) for x in ir_mod["main"].checked_type.arg_types[0].shape]
+    data = np.random.uniform(size=data_shape).astype("float32")
+    inputs = {"data": data}
+    ref_outputs = generate_ref_data(ir_mod, inputs, params)
+
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        mod = tvm.relay.build(
+            ir_mod,
+            params=params,
+            target="c",
+            executor=backend.Executor("aot", {"interface-api": "packed"}),
+        )
+
+    temp_dir = tvm.contrib.utils.TempDirectory()
+    test_so_path = temp_dir / "test.so"
+    mod.export_library(test_so_path, cc="gcc", options=["-std=c11"])
+    loaded_mod = tvm.runtime.load_module(test_so_path)
+    runner = tvm.runtime.executor.AotModule(loaded_mod["default"](tvm.cpu(0)))
+    runner.set_input(**inputs)
+    runner.run()
+    assert (runner.get_output(0).asnumpy() == list(ref_outputs.values())[0]).all()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/relay/aot/test_crt_aot.py b/tests/python/relay/aot/test_crt_aot.py
index ab51311d680a..2ce36f19fcc8 100644
--- a/tests/python/relay/aot/test_crt_aot.py
+++ b/tests/python/relay/aot/test_crt_aot.py
@@ -16,6 +16,7 @@
 # under the License.
 
 from collections import OrderedDict
+import re
 import sys
 
 import numpy as np
@@ -28,6 +29,7 @@
 from tvm.relay.testing import byoc
 from tvm.relay.op.annotation import compiler_begin, compiler_end
 from tvm.relay.backend import Executor, Runtime
+from tvm.micro import model_library_format as mlf
 from aot_test_utils import (
     AOTTestModel,
     AOT_DEFAULT_RUNNER,
@@ -36,6 +38,7 @@
     compile_and_run,
     compile_models,
     parametrize_aot_options,
+    create_relay_module_and_inputs_from_tflite_file,
 )
 
 
@@ -47,7 +50,14 @@ def test_error_c_interface_with_packed_api():
     two = relay.add(relay.const(1), relay.const(1))
     func = relay.Function([], two)
 
-    with pytest.raises(tvm.TVMError, match="Packed interface required for packed operators"):
+    with pytest.raises(
+        tvm.TVMError,
+        match=re.escape(
+            'Either need interface_api == "packed" (got: c) or '
+            "unpacked-api == true (got: (bool)0) when targeting "
+            "c runtime"
+        ),
+    ):
         compile_and_run(
             AOTTestModel(
                 module=IRModule.from_expr(func), inputs={}, outputs=generate_ref_data(func, {})
@@ -111,7 +121,10 @@ def test_add_with_params(interface_api, use_unpacked_api, test_runner):
 
     compile_and_run(
         AOTTestModel(
-            module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, params=params
+            module=IRModule.from_expr(func),
+            inputs=inputs,
+            outputs=output_list,
+            params=params,
         ),
         test_runner,
         interface_api,
@@ -447,7 +460,13 @@ def test_add_name_mangling_with_params(interface_api, use_unpacked_api, test_run
     output_list = generate_ref_data(func, inputs, params)
 
     compile_and_run(
-        AOTTestModel(name="my_mod", module=func, inputs=inputs, outputs=output_list, params=params),
+        AOTTestModel(
+            name="my_mod",
+            module=func,
+            inputs=inputs,
+            outputs=output_list,
+            params=params,
+        ),
         test_runner,
         interface_api,
         use_unpacked_api,
@@ -495,10 +514,18 @@ def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(8, 3, 5, 5),
     compile_and_run(
         [
             AOTTestModel(
-                name="mod1", module=mod1, inputs=inputs1, outputs=output_list1, params=params1
+                name="mod1",
+                module=mod1,
+                inputs=inputs1,
+                outputs=output_list1,
+                params=params1,
             ),
             AOTTestModel(
-                name="mod2", module=mod2, inputs=inputs2, outputs=output_list2, params=params2
+                name="mod2",
+                module=mod2,
+                inputs=inputs2,
+                outputs=output_list2,
+                params=params2,
             ),
         ],
         test_runner,
@@ -524,13 +551,7 @@ def test_quant_mobilenet_tfl():
         "models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
         "mobilenet_v1_1.0_224_quant.tflite",
     )
-    with open(tflite_model_file, "rb") as f:
-        tflite_model_buf = f.read()
-    data_shape = (1, 224, 224, 3)
-    in_min, in_max = (0, 255)
-    data = np.random.randint(in_min, high=in_max, size=data_shape, dtype="uint8")
-    mod, params = convert_to_relay(tflite_model_buf)
-    inputs = {"input": data}
+    mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
     output_list = generate_ref_data(mod, inputs, params)
     compile_and_run(
         AOTTestModel(module=mod, inputs=inputs, outputs=output_list, params=params),
@@ -647,7 +668,10 @@ def test_deprecated_target_arguments(capsys):
 
     compile_and_run(
         AOTTestModel(
-            module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, params=params
+            module=IRModule.from_expr(func),
+            inputs=inputs,
+            outputs=output_list,
+            params=params,
         ),
         test_runner,
         interface_api,
@@ -720,5 +744,178 @@ def test_constants_alignment(constants_byte_alignment):
     assert f'__attribute__((section(".rodata.tvm"), aligned({constants_byte_alignment})))' in source
 
 
+def test_output_tensor_names():
+    """Test that the output names generated match those in the model"""
+    pytest.importorskip("tflite")
+
+    import os
+    import tensorflow as tf
+    import tflite.Model
+
+    ifm_shape = (1, 299, 299, 3)
+    padding = "VALID"
+    strides = (1, 1)
+    dilation = (1, 1)
+    kernel_shape = (3, 2)
+
+    def create_tflite_graph_two_outs():
+        """Create a model with 2 output tensors"""
+
+        class Model(tf.Module):
+            @tf.function
+            def tf_function(self, x):
+                # Use tf.nn API to create the model
+                tf_strides = [1, strides[0], strides[1], 1]
+                op = tf.nn.conv2d(
+                    x,
+                    filters=tf.constant(
+                        np.random.uniform(size=[kernel_shape[0], kernel_shape[1], 3, 3]),
+                        dtype=tf.float32,
+                    ),
+                    strides=tf_strides,
+                    padding=padding,
+                    dilations=dilation,
+                )
+                op = tf.nn.relu(op)
+                # Second convolution
+                op2 = tf.nn.conv2d(
+                    x,
+                    filters=tf.constant(
+                        np.random.uniform(size=(kernel_shape[0], kernel_shape[1], 3, 3)),
+                        dtype=tf.float32,
+                    ),
+                    strides=strides,
+                    padding=padding,
+                    data_format="NHWC",
+                    dilations=dilation,
+                )
+                op2 = tf.nn.relu(op2)
+                return op, op2
+
+        model = Model()
+        concrete_func = model.tf_function.get_concrete_function(
+            tf.TensorSpec(ifm_shape, dtype=tf.float32)
+        )
+
+        # Convert the model
+        def representative_dataset():
+            for _ in range(100):
+                data = np.random.rand(*tuple(ifm_shape))
+                yield [data.astype(np.float32)]
+
+        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+        converter.optimizations = [tf.lite.Optimize.DEFAULT]
+        converter.representative_dataset = representative_dataset
+        converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        converter.inference_input_type = tf.int8
+        converter.inference_output_type = tf.int8
+        tflite_model = converter.convert()
+        return tflite_model
+
+    tflite_graph = create_tflite_graph_two_outs()
+    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_graph, 0)
+    mod, params = relay.frontend.from_tflite(
+        tflite_model,
+        shape_dict={"input": ifm_shape},
+        dtype_dict={"input": "int8"},
+    )
+
+    use_unpacked_api = True
+    interface_api = "c"
+    test_runner = AOT_DEFAULT_RUNNER
+
+    in_min, in_max = (-128, 127)
+    data = np.random.randint(in_min, high=in_max, size=ifm_shape, dtype="int8")
+    inputs = {"x_int8": data}
+    output_list = generate_ref_data(mod, inputs, params)
+    compile_and_run(
+        AOTTestModel(module=mod, inputs=inputs, outputs=output_list, params=params),
+        test_runner,
+        interface_api,
+        use_unpacked_api,
+    )
+
+    compiled_test_mods = compile_models(
+        AOTTestModel(module=mod, inputs=inputs, outputs=output_list, params=params),
+        interface_api,
+        use_unpacked_api,
+    )
+
+    # Check that the names of the output tensors occur in the source code
+    source = compiled_test_mods[0].executor_factory.lib.get_source()
+    for output_name in output_list.keys():
+        assert output_name in source
+
+
+@pytest.mark.parametrize(
+    "workspace_byte_alignment,main_workspace_size",
+    [
+        (8, 14880),
+        (16, 14880),
+        (256, 15616),
+    ],
+)
+def test_workspace_calculation(workspace_byte_alignment, main_workspace_size):
+    mod, params = tvm.relay.testing.synthetic.get_workload()
+    target = "c"
+    runtime = Runtime("crt")
+    executor = Executor(
+        "aot",
+        {
+            "workspace-byte-alignment": workspace_byte_alignment,
+        },
+    )
+    with tvm.transform.PassContext(
+        opt_level=3,
+        config={
+            "tir.disable_vectorize": True,
+        },
+    ):
+        lib = tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
+
+    mlf_memory_map = mlf._build_function_memory_map(lib.function_metadata)
+    assert mlf_memory_map["main"][0]["workspace_size_bytes"] == main_workspace_size
+
+
+@tvm.testing.requires_package("tflite")
+@tvm.testing.requires_cmsisnn
+def test_workspace_calculation_cmsis_nn():
+    """This tests cmsis_nn codegen for workspace calculation.
+    This is tested specially because cmsis-nn codegen creates
+    multiple PrimFuncs per offloaded relay function in a non
+    -hierarchical manner."""
+    pytest.importorskip("tflite")
+
+    from tvm.relay.op.contrib import cmsisnn
+    from tvm.contrib.download import download_testdata
+
+    target = "c"
+    runtime = Runtime("crt")
+    executor = Executor(
+        "aot",
+        {
+            "workspace-byte-alignment": 16,
+            "interface-api": "c",
+            "unpacked-api": True,
+        },
+    )
+
+    base_url = "https://github.com/ARM-software/ML-zoo/raw/48a22ee22325d15d2371a6df24eb7d67e21dcc97/models/keyword_spotting/cnn_small/tflite_int8"
+    file_to_download = "cnn_s_quantized.tflite"
+    file_saved = "cnn_s_quantized_15Dec2021.tflite"
+    model_file = download_testdata("{}/{}".format(base_url, file_to_download), file_saved)
+    mod, _, params = create_relay_module_and_inputs_from_tflite_file(model_file)
+    mod = cmsisnn.partition_for_cmsisnn(mod, params)
+    with tvm.transform.PassContext(
+        opt_level=3,
+        config={
+            "tir.disable_vectorize": True,
+        },
+    ):
+        lib = tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
+    mlf_memory_map = mlf._build_function_memory_map(lib.function_metadata)
+    assert mlf_memory_map["main"][0]["workspace_size_bytes"] == 9904
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/relay/aot/test_crt_aot_usmp.py b/tests/python/relay/aot/test_crt_aot_usmp.py
index a27609cc07ad..73b34700ee27 100644
--- a/tests/python/relay/aot/test_crt_aot_usmp.py
+++ b/tests/python/relay/aot/test_crt_aot_usmp.py
@@ -39,6 +39,7 @@
     compile_models,
     parametrize_aot_options,
     run_and_check,
+    create_relay_module_and_inputs_from_tflite_file,
 )
 
 
@@ -202,23 +203,6 @@ def test_byoc_microtvm(merge_compiler_regions):
     )
 
 
-def _get_relay_module_and_inputs_from_tflite_file(tflite_model_file):
-    with open(tflite_model_file, "rb") as f:
-        tflite_model_buf = f.read()
-    mod, params = convert_to_relay(tflite_model_buf)
-
-    inputs = dict()
-    for param in mod["main"].params:
-        name = str(param.name_hint)
-        data_shape = [int(i) for i in param.type_annotation.shape]
-        dtype = str(param.type_annotation.dtype)
-        in_min, in_max = (np.iinfo(dtype).min, np.iinfo(dtype).max)
-        data = np.random.randint(in_min, high=in_max, size=data_shape, dtype=dtype)
-        inputs[name] = data
-
-    return mod, inputs, params
-
-
 MOBILENET_V1_URL = (
     "https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz",
     "mobilenet_v1_1.0_224_quant.tflite",
@@ -253,7 +237,7 @@ def test_tflite_model_u1_usecase(model_url, usmp_algo, workspace_size):
         model_url[0],
         model_url[1],
     )
-    mod, inputs, params = _get_relay_module_and_inputs_from_tflite_file(tflite_model_file)
+    mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
     output_list = generate_ref_data(mod, inputs, params)
 
     compiled_test_mods = compile_models(
@@ -324,7 +308,7 @@ def test_tflite_model_u3_usecase_single_external_pool(model_url, usmp_algo):
         model_url[0],
         model_url[1],
     )
-    mod, inputs, params = _get_relay_module_and_inputs_from_tflite_file(tflite_model_file)
+    mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
     output_list = generate_ref_data(mod, inputs, params)
 
     compiled_test_mods = compile_models(
@@ -384,7 +368,7 @@ def test_tflite_model_u3_usecase_two_external_pools(model_url, usmp_algo):
         model_url[0],
         model_url[1],
     )
-    mod, inputs, params = _get_relay_module_and_inputs_from_tflite_file(tflite_model_file)
+    mod, inputs, params = create_relay_module_and_inputs_from_tflite_file(tflite_model_file)
     output_list = generate_ref_data(mod, inputs, params)
 
     compiled_test_mods = compile_models(
@@ -438,14 +422,14 @@ def test_tflite_model_u2_usecase_two_models_with_a_single_external_pool(model_ur
         model_urls[0][0],
         model_urls[0][1],
     )
-    mod1, inputs1, params1 = _get_relay_module_and_inputs_from_tflite_file(tflite_model_file1)
+    mod1, inputs1, params1 = create_relay_module_and_inputs_from_tflite_file(tflite_model_file1)
     output_list1 = generate_ref_data(mod1, inputs1, params1)
 
     tflite_model_file2 = tf_testing.get_workload_official(
         model_urls[1][0],
         model_urls[1][1],
     )
-    mod2, inputs2, params2 = _get_relay_module_and_inputs_from_tflite_file(tflite_model_file2)
+    mod2, inputs2, params2 = create_relay_module_and_inputs_from_tflite_file(tflite_model_file2)
     output_list2 = generate_ref_data(mod2, inputs2, params2)
 
     compiled_test_mods = compile_models(
diff --git a/tests/python/relay/op/annotation/test_annotation.py b/tests/python/relay/op/annotation/test_annotation.py
index 2352821f7bee..dcbb40cdcabc 100644
--- a/tests/python/relay/op/annotation/test_annotation.py
+++ b/tests/python/relay/op/annotation/test_annotation.py
@@ -61,18 +61,6 @@ def test_on_device_free():
     assert not call.attrs.constrain_result
 
 
-def test_function_on_device():
-    x = relay.Var("x")
-    y = relay.Var("y")
-    f = relay.Function([x, y], relay.add(x, y))
-    func = relay.annotation.function_on_device(f, ["cpu", "cuda"], "cuda")
-    assert isinstance(func, relay.Function)
-    assert len(func.attrs["param_virtual_devices"]) == 2
-    assert func.attrs["param_virtual_devices"][0].device_type_int == 1  # ie kDLCPU
-    assert func.attrs["param_virtual_devices"][1].device_type_int == 2  # ie kDLCUDA
-    assert func.attrs["result_virtual_device"].device_type_int == 2  # ie KDLCUDA
-
-
 if __name__ == "__main__":
     import sys
 
diff --git a/tests/python/relay/test_cpp_build_module.py b/tests/python/relay/test_cpp_build_module.py
index 23bc7ca95a34..ccf961fbe4db 100644
--- a/tests/python/relay/test_cpp_build_module.py
+++ b/tests/python/relay/test_cpp_build_module.py
@@ -93,6 +93,45 @@ def test_fp16_build():
     np.testing.assert_allclose(out.numpy(), X.numpy() + Y.numpy(), atol=1e-5, rtol=1e-5)
 
 
+@tvm.testing.requires_llvm
+def test_bf16_build():
+    data = relay.var("data", shape=(1, 3, 224, 224), dtype="float32")
+    weight = relay.var("weight", shape=(64, 3, 7, 7), dtype="float32")
+    bn_gamma = relay.var("gamma", shape=(64,), dtype="float32")
+    bn_beta = relay.var("beta", shape=(64,), dtype="float32")
+    bn_mean = relay.var("mean", shape=(64,), dtype="float32")
+    bn_var = relay.var("var", shape=(64,), dtype="float32")
+    params = {
+        "weight": np.random.uniform(-1, 1, size=(64, 3, 7, 7)).astype("float32"),
+        "gamma": np.random.uniform(-1, 1, size=(64,)).astype("float32"),
+        "beta": np.random.uniform(-1, 1, size=(64,)).astype("float32"),
+        "mean": np.random.uniform(-1, 1, size=(64,)).astype("float32"),
+        "var": np.random.uniform(-1, 1, size=(64,)).astype("float32"),
+    }
+    conv_bf16 = relay.nn.conv2d(
+        relay.cast(data, "bfloat16"),
+        relay.cast(weight, "bfloat16"),
+        strides=(2, 2),
+        padding=(3, 3, 3, 3),
+        channels=64,
+        kernel_size=(7, 7),
+        out_dtype="bfloat16",
+    )
+    bn_bf16 = relay.nn.batch_norm(
+        conv_bf16,
+        relay.cast(bn_gamma, "bfloat16"),
+        relay.cast(bn_beta, "bfloat16"),
+        relay.cast(bn_mean, "bfloat16"),
+        relay.cast(bn_var, "bfloat16"),
+    )
+    relu_bf16 = relay.nn.relu(bn_bf16[0])
+    maxpool_bf16 = relay.nn.max_pool2d(relu_bf16, pool_size=(2, 2), strides=(2, 2))
+    avgpool_bf16 = relay.nn.avg_pool2d(maxpool_bf16, pool_size=(2, 2), strides=(2, 2))
+    mod_bf16 = tvm.IRModule.from_expr(avgpool_bf16)
+    with tvm.transform.PassContext(opt_level=3):
+        relay.build(mod_bf16, target="llvm", params=params)
+
+
 @tvm.testing.parametrize_targets("llvm", "cuda")
 def test_fp16_conversion(target, dev):
     if target == "cuda" and not have_fp16(dev.compute_version):
@@ -126,3 +165,4 @@ def test_fp16_conversion(target, dev):
     test_basic_build()
     test_fp16_build()
     test_fp16_conversion()
+    test_bf16_build()
diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py
index 791edcc9c139..4f7d810fa7a5 100644
--- a/tests/python/relay/test_op_level1.py
+++ b/tests/python/relay/test_op_level1.py
@@ -20,7 +20,7 @@
 from tvm import te
 import scipy
 from tvm import relay
-from tvm.relay import transform
+import pytest
 from tvm.relay.testing import run_infer_type
 import tvm.topi.testing
 from tvm.contrib.nvcc import have_fp16
@@ -634,19 +634,43 @@ def test_bitserial_dense():
     assert yy.checked_type == relay.TensorType((m, 32), "int16")
 
 
+@pytest.mark.skip("Requires cascadelake")
+def test_dense_vnni():
+    data_shape = (32, 96)
+    weight_shape = (128, 96)
+
+    for data_dtype in ["uint8", "int8"]:
+        data = relay.var("data", shape=data_shape, dtype=data_dtype)
+        weight = relay.var("weight", shape=weight_shape, dtype="int8")
+        bias = relay.var("bias", shape=(weight_shape[0],), dtype="int32")
+        dense = relay.nn.dense(data, weight, out_dtype="int32")
+        out = relay.nn.bias_add(dense, bias)
+        mod = tvm.IRModule.from_expr(out)
+
+        target = "llvm -mcpu=cascadelake"
+        with tvm.transform.PassContext(opt_level=3):
+            lib = relay.build(mod, target=target)
+
+        asm = lib.lib.get_source("asm")
+        assert "vpdpbusd" in asm
+
+        dev = tvm.device(target, 0)
+        runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+
+        a = np.random.uniform(1, 10, size=data_shape).astype(data_dtype)
+        b = np.random.uniform(1, 10, size=weight_shape).astype("int8")
+        c = np.random.uniform(1, 10, size=(weight_shape[0],)).astype("int32")
+
+        runtime.set_input("data", a)
+        runtime.set_input("weight", b)
+        runtime.set_input("bias", c)
+        runtime.run()
+
+        out = runtime.get_output(0).numpy()
+        ref = np.dot(a.astype("int32"), b.transpose().astype("int32")) + c
+
+        np.testing.assert_equal(out, ref)
+
+
 if __name__ == "__main__":
-    test_concatenate()
-    test_bias_add()
-    test_bias_add_type_failure()
-    test_unary_op()
-    test_binary_op()
-    test_expand_dims_infer_type()
-    test_expand_dims()
-    test_softmax()
-    test_log_softmax()
-    test_dropout()
-    test_batch_norm()
-    test_matmul()
-    test_dense()
-    test_bitserial_dense()
-    test_dense_dtype()
+    pytest.main([__file__])
diff --git a/tests/python/relay/test_op_level10.py b/tests/python/relay/test_op_level10.py
index 476728664132..0486ef40017b 100644
--- a/tests/python/relay/test_op_level10.py
+++ b/tests/python/relay/test_op_level10.py
@@ -16,6 +16,9 @@
 # under the License.
 """ Support level10 operator test cases.
 """
+import sys
+import pytest
+
 import numpy as np
 import tvm
 import tvm.testing
@@ -388,6 +391,45 @@ def test_batch_matmul():
     verify_batch_matmul_with_inputs(x, x, x_np, x_np, (10, 27, 27))
 
 
+@pytest.mark.skip("Requires cascadelake")
+def test_batch_matmul_vnni():
+    x_shape = (16, 32, 96)
+    y_shape = (16, 128, 96)
+    z_shape = (16, 32, 128)
+
+    for lhs_dtype in ["uint8", "int8"]:
+        x = relay.var("x", shape=x_shape, dtype=lhs_dtype)
+        y = relay.var("y", shape=y_shape, dtype="int8")
+        z = relay.var("z", shape=z_shape, dtype="int32")
+        bmm = relay.nn.batch_matmul(x, y, out_dtype="int32")
+        out = bmm + z
+        mod = tvm.IRModule.from_expr(out)
+
+        target = "llvm -mcpu=cascadelake"
+        with tvm.transform.PassContext(opt_level=3):
+            lib = relay.build(mod, target=target)
+
+        asm = lib.lib.get_source("asm")
+        assert "vpdpbusd" in asm
+
+        dev = tvm.device(target, 0)
+        runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+
+        x_np = np.random.uniform(1, 10, size=x_shape).astype(lhs_dtype)
+        y_np = np.random.uniform(1, 10, size=y_shape).astype("int8")
+        z_np = np.random.uniform(1, 10, size=z_shape).astype("int32")
+
+        runtime.set_input("x", x_np)
+        runtime.set_input("y", y_np)
+        runtime.set_input("z", z_np)
+        runtime.run()
+
+        out = runtime.get_output(0).numpy()
+        ref = tvm.topi.testing.batch_matmul(x_np, y_np, out_dtype="int32") + z_np
+
+        np.testing.assert_equal(out, ref)
+
+
 @tvm.testing.uses_gpu
 def test_shape_of():
     shape = (10, 5, 12)
@@ -623,7 +665,4 @@ def _verify(prediction_shape, reduction="mean", ignore_index=-100, dtype="float3
 
 
 if __name__ == "__main__":
-    import sys
-    import pytest
-
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/relay/test_op_qnn_rsqrt.py b/tests/python/relay/test_op_qnn_rsqrt.py
deleted file mode 100644
index 0e40768343bd..000000000000
--- a/tests/python/relay/test_op_qnn_rsqrt.py
+++ /dev/null
@@ -1,95 +0,0 @@
-# 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.
-
-import numpy as np
-import tvm
-from tvm import relay
-
-
-def dequantize(data, scale, zp):
-    return scale * (np.asarray(data) - zp)
-
-
-def generate_golden_output(dequantized_x, output_scale, output_zero_point):
-    rsqrt = 1 / np.sqrt(dequantized_x)
-    output = np.around(rsqrt / output_scale + output_zero_point)
-
-    q_min = np.iinfo(np.uint8).min
-    q_max = np.iinfo(np.uint8).max
-    return np.clip(output, q_min, q_max)
-
-
-def test_saturation():
-    # Same params
-    data_dtype = "uint8"
-    scale = output_scale = 0.125
-    zero_point = output_zero_point = 0
-
-    x = relay.var("x", shape=(1, 4), dtype=data_dtype)
-    y = relay.qnn.op.rsqrt(
-        x=x,
-        scale=relay.const(scale, "float32"),
-        zero_point=relay.const(zero_point, "int32"),
-        output_scale=relay.const(output_scale, "float32"),
-        output_zero_point=relay.const(output_zero_point, "int32"),
-    )
-
-    func = relay.Function([x], y)
-    mod = tvm.IRModule.from_expr(func)
-    mod = relay.transform.InferType()(mod)
-    mod = relay.qnn.transform.Legalize()(mod)
-    mod = relay.qnn.transform.CanonicalizeOps()(mod)
-    func = mod["main"]
-
-    x_data = np.array((255, 133, 0, 9)).reshape((1, 4))
-    x_dequantized = dequantize(x_data, scale, zero_point)
-    golden_output = generate_golden_output(x_dequantized, output_scale, output_zero_point)
-
-    op_res = relay.create_executor("graph", device=tvm.cpu(0), target="llvm").evaluate(func)(x_data)
-
-    np.testing.assert_equal(op_res.numpy(), np.uint8(golden_output))
-
-    # Different scale
-    scale = 0.125
-    output_scale = 0.25
-
-    y = relay.qnn.op.rsqrt(
-        x=x,
-        scale=relay.const(scale, "float32"),
-        zero_point=relay.const(zero_point, "int32"),
-        output_scale=relay.const(output_scale, "float32"),
-        output_zero_point=relay.const(output_zero_point, "int32"),
-    )
-
-    func = relay.Function([x], y)
-    mod = tvm.IRModule.from_expr(func)
-    mod = relay.transform.InferType()(mod)
-    mod = relay.qnn.transform.Legalize()(mod)
-    mod = relay.qnn.transform.CanonicalizeOps()(mod)
-    func = mod["main"]
-
-    x_data = np.array((255, 133, 0, 9)).reshape((1, 4))
-    x_dequantized = dequantize(x_data, scale, zero_point)
-    golden_output = generate_golden_output(x_dequantized, output_scale, output_zero_point)
-
-    op_res = relay.create_executor("graph", device=tvm.cpu(0), target="llvm").evaluate(func)(x_data)
-
-    np.testing.assert_equal(op_res.numpy(), golden_output)
-
-
-if __name__ == "__main__":
-    test_saturation()
diff --git a/tests/python/relay/test_op_qnn_unary_elementwise.py b/tests/python/relay/test_op_qnn_unary_elementwise.py
new file mode 100644
index 000000000000..18acc119a947
--- /dev/null
+++ b/tests/python/relay/test_op_qnn_unary_elementwise.py
@@ -0,0 +1,214 @@
+# 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.
+
+from typing import Callable, List
+
+import numpy as np
+import pytest
+import scipy.special
+import tvm
+from tvm import relay
+
+
+def dequantize(data, scale, zp):
+    return scale * (np.asarray(data) - zp)
+
+
+def generate_golden_output(
+    floating_point_golden_func, dequantized_x, output_scale, output_zero_point, dtype
+):
+    output = floating_point_golden_func(dequantized_x)
+    output = np.around(output / output_scale + output_zero_point)
+
+    np_dtype = {"int8": np.int8, "uint8": np.uint8}[dtype]
+
+    q_min = np.iinfo(np_dtype).min
+    q_max = np.iinfo(np_dtype).max
+    return np.clip(output, q_min, q_max)
+
+
+def run_qnn_func(func: relay.Function, args: List[relay.Expr]):
+    mod = tvm.IRModule.from_expr(func)
+    mod = relay.transform.InferType()(mod)
+    mod = relay.qnn.transform.Legalize()(mod)
+    mod = relay.qnn.transform.CanonicalizeOps()(mod)
+    func = mod["main"]
+
+    op_res = relay.create_executor("graph", device=tvm.cpu(0), target="llvm").evaluate(func)(*args)
+    return op_res.numpy()
+
+
+def create_qnn_func(
+    qnn_op: Callable[[relay.Expr, relay.Expr, relay.Expr, relay.Expr, relay.Expr], relay.Call],
+    x_data: np.ndarray,
+    input_scale: float,
+    input_zero_point: int,
+    output_scale: float,
+    output_zero_point: int,
+    input_dtype: str = "uint8",
+):
+    x = relay.var("x", shape=x_data.shape, dtype=input_dtype)
+    y = qnn_op(
+        x=x,
+        scale=relay.const(input_scale, "float32"),
+        zero_point=relay.const(input_zero_point, "int32"),
+        output_scale=relay.const(output_scale, "float32"),
+        output_zero_point=relay.const(output_zero_point, "int32"),
+    )
+    return relay.Function([x], y)
+
+
+def run_condition(
+    qnn_op: Callable[[relay.Expr, relay.Expr, relay.Expr, relay.Expr, relay.Expr], relay.Call],
+    floating_point_golden_func: Callable[[np.ndarray], np.ndarray],
+    x_data: np.ndarray,
+    input_scale: float,
+    input_zero_point: int,
+    output_scale: float,
+    output_zero_point: int,
+    input_dtype: str = "uint8",
+):
+    func = create_qnn_func(
+        qnn_op,
+        x_data,
+        input_scale=input_scale,
+        input_zero_point=input_zero_point,
+        output_scale=output_scale,
+        output_zero_point=output_zero_point,
+        input_dtype=input_dtype,
+    )
+
+    x_dequantized = dequantize(x_data, input_scale, input_zero_point)
+    golden_output = generate_golden_output(
+        floating_point_golden_func,
+        x_dequantized,
+        output_scale,
+        output_zero_point,
+        dtype=input_dtype,
+    )
+
+    op_res = run_qnn_func(func, [x_data])
+    np.testing.assert_equal(op_res, golden_output.astype(input_dtype))
+
+
+def generic_test(
+    qnn_op: Callable[[relay.Expr, relay.Expr, relay.Expr, relay.Expr, relay.Expr], relay.Call],
+    floating_point_golden_func: Callable[[np.ndarray], np.ndarray],
+    input_dtype: str = "uint8",
+    x_data: np.ndarray = np.arange(0, 256, dtype="uint8"),
+):
+    x_data = x_data.view(input_dtype)
+    return run_condition(
+        qnn_op,
+        floating_point_golden_func,
+        x_data,
+        input_scale=0.125,
+        input_zero_point=0,
+        output_scale=0.125,
+        output_zero_point=0,
+        input_dtype=input_dtype,
+    )
+
+
+class TestRSqrt:
+    def test_saturation(self):
+        # Same qparams in and out
+        x_data = np.array((255, 133, 0, 9)).reshape((1, 4))
+        run_condition(
+            relay.qnn.op.rsqrt,
+            lambda x: 1 / np.sqrt(x),
+            x_data,
+            input_scale=0.125,
+            input_zero_point=0,
+            output_scale=0.125,
+            output_zero_point=0,
+            input_dtype="uint8",
+        )
+
+        # Different scale
+        run_condition(
+            relay.qnn.op.rsqrt,
+            lambda x: 1 / np.sqrt(x),
+            x_data,
+            input_scale=0.125,
+            input_zero_point=0,
+            output_scale=0.25,
+            output_zero_point=0,
+            input_dtype="uint8",
+        )
+
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.rsqrt, lambda x: 1 / np.sqrt(x), input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(
+            relay.qnn.op.rsqrt,
+            lambda x: 1 / np.sqrt(x),
+            input_dtype="int8",
+            x_data=np.arange(1, 128, dtype="int8"),
+        )
+
+
+class Sqrt:
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.sqrt, np.sqrt, input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(
+            relay.qnn.op.sqrt,
+            np.sqrt,
+            input_dtype="int8",
+            x_data=np.arange(1, 128, dtype="int8"),
+        )
+
+
+class TestExp:
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.exp, np.exp, input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(relay.qnn.op.exp, np.exp, input_dtype="int8")
+
+
+class TestTanh:
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.tanh, np.tanh, input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(relay.qnn.op.tanh, np.tanh, input_dtype="int8")
+
+
+class TestErf:
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.erf, scipy.special.erf, input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(relay.qnn.op.erf, scipy.special.erf, input_dtype="int8")
+
+
+class TestSigmoid:
+    def test_all_numbers_uint8(self):
+        generic_test(relay.qnn.op.sigmoid, lambda x: 1 / (1 + np.exp(-x)), input_dtype="uint8")
+
+    def test_all_numbers_int8(self):
+        generic_test(relay.qnn.op.sigmoid, lambda x: 1 / (1 + np.exp(-x)), input_dtype="int8")
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))
diff --git a/tests/python/relay/test_pass_alter_op_layout.py b/tests/python/relay/test_pass_alter_op_layout.py
index 4ce1286eb508..07382e6a38d0 100644
--- a/tests/python/relay/test_pass_alter_op_layout.py
+++ b/tests/python/relay/test_pass_alter_op_layout.py
@@ -163,6 +163,52 @@ def expected():
     assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
 
 
+def test_alter_layout_multi():
+    """Test alternating the layout of a conv2d.
+    The layout of broadcast operators and the weight should be changed accordingly.
+    """
+
+    def before():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        weight = relay.var("weight")
+        y = relay.nn.conv2d(x, weight, channels=128, kernel_size=(3, 3), padding=(1, 1))
+        y = relay.Function(analysis.free_vars(y), y)
+        return y
+
+    def alter_conv2d(attrs, inputs, tinfos, out_type):
+        data, weight = inputs
+        new_attrs = dict(attrs)
+        new_attrs["data_layout"] = "NCHW16c"
+        new_attrs["kernel_layout"] = "OHWI16i64o2i"
+        return relay.nn.conv2d(data, weight, **new_attrs)
+
+    def expected():
+        x = relay.var("x", shape=(1, 64, 56, 56))
+        weight = relay.var("weight", shape=(128, 64, 3, 3))
+
+        y = relay.layout_transform(x, "NCHW", "NCHW16c")
+        w = relay.layout_transform(weight, "OIHW", "OHWI16i64o2i")
+        y = relay.nn.conv2d(
+            y,
+            w,
+            channels=128,
+            kernel_size=(3, 3),
+            padding=(1, 1),
+            kernel_layout="OHWI16i64o2i",
+            data_layout="NCHW16c",
+        )
+        y = relay.layout_transform(y, "NCHW16c", "NCHW")
+        y = relay.Function(analysis.free_vars(y), y)
+        return y
+
+    with TempOpAttr("nn.conv2d", "FTVMAlterOpLayout", alter_conv2d):
+        a = before()
+        a = run_opt_pass(a, [transform.CanonicalizeOps(), transform.AlterOpLayout()])
+        b = run_opt_pass(expected(), transform.InferType())
+
+    assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
+
+
 def test_alter_layout_lrn():
     """Test alternating the layout of a conv2d.
     The layout of broadcast operators and the weight should be changed accordingly.
diff --git a/tests/python/relay/test_pass_fake_quantization_to_integer.py b/tests/python/relay/test_pass_fake_quantization_to_integer.py
index e9bfc640c7d5..91a2e23b40b9 100644
--- a/tests/python/relay/test_pass_fake_quantization_to_integer.py
+++ b/tests/python/relay/test_pass_fake_quantization_to_integer.py
@@ -25,7 +25,6 @@
 def compare_fq_to_int(expr, args, allow_rounding_error=False):
     mod = tvm.IRModule.from_expr(expr)
     mod = tvm.relay.transform.InferType()(mod)
-
     mod_int = tvm.relay.transform.FakeQuantizationToInteger()(mod)
     assert not tvm.ir.structural_equal(mod, mod_int)
 
@@ -71,12 +70,12 @@ def test_fake_quantize_conv_per_channel():
         x = relay.var("x", shape=[1, 3, 224, 224], dtype="int8")
         w = relay.var("w", shape=[16, 3, 5, 5], dtype="int8")
         one = relay.const([1.0] * 16)
-        zero = relay.const([0] * 16)
+        zero_point = relay.const([np.random.randint(0, 255)] * 16)
 
         op = relay.op.nn.conv2d(
             relay.qnn.op.dequantize(x, relay.const(2.0), relay.const(0)),
             relay.qnn.op.dequantize(
-                w, relay.const(np.random.random([16]).astype("float32")), zero, axis=0
+                w, relay.const(np.random.random([16]).astype("float32")), zero_point, axis=0
             ),
             kernel_size=[5, 5],
             channels=16,
@@ -304,17 +303,36 @@ def test_fake_quantize_global_avg_pool():
     compare_fq_to_int(op, [x_np], True)
 
 
-def test_fake_quantize_rsqrt():
-    x = relay.var("x", shape=[1, 3, 3, 3], dtype="int8")
-    mid_point = relay.const(-128)
+class TestUnaryQNNOp:
+    def helper_test_fake_quantize_unary_op(self, fp32_op, scale=0.125):
+        x = relay.var("x", shape=[1, 3, 3, 3], dtype="int8")
+        mid_point = relay.const(-128)
 
-    x = relay.qnn.op.dequantize(x, relay.const(0.125), mid_point)
-    op = relay.rsqrt(x)
-    op = relay.qnn.op.quantize(op, relay.const(0.125), mid_point)
+        x = relay.qnn.op.dequantize(x, relay.const(scale), mid_point)
+        op = fp32_op(x)
+        op = relay.qnn.op.quantize(op, relay.const(scale), mid_point)
 
-    x_np = np.random.randint(-128, 127, size=[1, 3, 3, 3], dtype="int8")
+        x_np = np.random.randint(-128, 127, size=[1, 3, 3, 3], dtype="int8")
 
-    compare_fq_to_int(op, [x_np], True)
+        compare_fq_to_int(op, [x_np], True)
+
+    def test_sqrt(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.sqrt)
+
+    def test_rsqrt(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.rsqrt)
+
+    def test_exp(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.exp)
+
+    def test_erf(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.erf)
+
+    def test_sigmoid(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.sigmoid)
+
+    def test_tanh(self):
+        self.helper_test_fake_quantize_unary_op(fp32_op=relay.tanh)
 
 
 def test_fake_quantize_reshape():
@@ -627,3 +645,219 @@ def conv2d(expr, type_map):  # pylint: disable=unused-variable
     # Catch a generic exception because the tvm FFI eats the python exception type
     with pytest.raises(Exception):
         mod_int = tvm.relay.transform.FakeQuantizationToInteger(hard_fail=True)(mod)
+
+
+def compare_expected_fq_qat_to_int(expr, expected_expr, args, allow_rounding_error=False):
+    mod = tvm.IRModule.from_expr(expr)
+    mod_def = tvm.relay.transform.InferType()(mod)
+    mod_int = tvm.relay.transform.FakeQuantizationToInteger(False, True)(mod_def)
+    mod_exp = tvm.relay.transform.InferType()(tvm.IRModule.from_expr(expected_expr))
+    assert not tvm.ir.structural_equal(mod, mod_int)
+    assert tvm.ir.structural_equal(mod_int, mod_exp)
+    result_def = (
+        relay.create_executor("vm", mod=mod_def, device=tvm.cpu(), target="llvm")
+        .evaluate()(*args)
+        .numpy()
+    )
+    result_int = (
+        relay.create_executor("vm", mod=mod_int, device=tvm.cpu(), target="llvm")
+        .evaluate()(*args)
+        .numpy()
+    )
+    result_exp = (
+        relay.create_executor("vm", mod=mod_exp, device=tvm.cpu(), target="llvm")
+        .evaluate()(*args)
+        .numpy()
+    )
+    if allow_rounding_error:
+        assert np.all(np.abs(result_def.astype("int32") - result_int.astype("int32")) <= 1)
+    else:
+        assert np.array_equal(result_def, result_int)
+
+    assert np.array_equal(result_int, result_exp)
+
+
+def test_fq_qat_op_positive_part():
+    # Only the first operation is converted, since the next operation("add") is not enabled.
+    shape_x = [1, 4, 2]
+    shape_w = [1, 4, 2]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+    b = relay.var("b", shape=shape_w, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+    op1 = relay.qnn.op.dequantize(b, relay.const(6.0), relay.const(0))
+    op2 = relay.op.nn.batch_matmul(op0, op1)
+    op3 = relay.op.add(op2, relay.const(1.0))
+    expr = relay.op.erf(op3)
+
+    op0 = relay.qnn.op.qnn.batch_matmul(
+        a, b, relay.const(0), relay.const(0), relay.const(2.0), relay.const(6.0)
+    )
+    op1 = relay.qnn.op.qnn.dequantize(op0, relay.const(12.0), relay.const(0))
+    op2 = relay.op.add(op1, relay.const(1.0))
+    expected_expr = relay.op.erf(op2)
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    w_np = np.random.randint(-128, 127, size=shape_w, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np, w_np])
+
+
+def test_fq_qat_negative_all():
+    # None of the operations are converted, since the first operation("add") is not enabled.
+    shape_x = [1, 4, 2]
+    shape_w = [1, 4, 2]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+    b = relay.var("b", shape=shape_w, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+    op1 = relay.qnn.op.dequantize(b, relay.const(6.0), relay.const(0))
+    op2 = relay.op.add(op1, relay.const(1.0))
+    op3 = relay.op.nn.batch_matmul(op0, op2)
+    expr = relay.op.erf(op3)
+
+    expected_expr = expr
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    w_np = np.random.randint(-128, 127, size=shape_w, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np, w_np])
+
+
+def test_fq_qat_positive_single():
+    # The single operation is converted.
+    shape_x = [1, 4, 2]
+    shape_w = [1, 4, 2]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+    b = relay.var("b", shape=shape_w, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+    op1 = relay.qnn.op.dequantize(b, relay.const(6.0), relay.const(0))
+    expr = relay.op.nn.batch_matmul(op0, op1)
+
+    op0 = relay.qnn.op.qnn.batch_matmul(
+        a, b, relay.const(0), relay.const(0), relay.const(2.0), relay.const(6.0)
+    )
+    expected_expr = relay.qnn.op.qnn.dequantize(op0, relay.const(12.0), relay.const(0))
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    w_np = np.random.randint(-128, 127, size=shape_w, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np, w_np])
+
+
+def test_fq_qat_positive_nothing_to_do():
+    # All operations are converted by the non-QAT pass.
+    shape_x = [1, 4, 2]
+    shape_w = [1, 4, 2]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+    b = relay.var("b", shape=shape_w, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+    op1 = relay.qnn.op.dequantize(b, relay.const(6.0), relay.const(0))
+    op2 = relay.op.nn.batch_matmul(op0, op1)
+    op3 = relay.op.add(op2, relay.const(1.0))
+    expr = relay.qnn.op.quantize(op3, relay.const(1.0), relay.const(0), out_dtype="int8")
+
+    op0 = relay.qnn.op.batch_matmul(
+        a, b, relay.const(0), relay.const(0), relay.const(2.0), relay.const(6.0)
+    )
+    op1 = relay.qnn.op.quantize(
+        relay.const(1.0), relay.const(12.0), relay.const(0), out_dtype="int32"
+    )
+    op2 = relay.qnn.op.add(
+        op0,
+        op1,
+        relay.const(12.0),
+        relay.const(0),
+        relay.const(12.0),
+        relay.const(0),
+        relay.const(12.0),
+        relay.const(0),
+    )
+    expected_expr = relay.qnn.op.requantize(
+        op2, relay.const(12.0), relay.const(0), relay.const(1.0), relay.const(0), out_dtype="int8"
+    )
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    w_np = np.random.randint(-128, 127, size=shape_w, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np, w_np])
+
+
+def test_fq_qat_positive_couple():
+    # Several consecutive operations are converted.
+    shape_x = [1, 2, 4]
+    shape_w = [2]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+    b = relay.var("b", shape=shape_w, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+    op1 = relay.qnn.op.dequantize(b, relay.const(6.0), relay.const(0))
+    op2 = relay.op.reshape(op0, (1, 4, 2))
+    op3 = relay.op.broadcast_to(op1, (2, 2, 2))
+    op4 = relay.op.nn.batch_matmul(op2, op3)
+    expr = relay.op.erf(op4)
+
+    op0 = relay.op.reshape(a, (1, 4, 2))
+    op1 = relay.op.broadcast_to(b, (2, 2, 2))
+    op3 = relay.qnn.op.qnn.batch_matmul(
+        op0, op1, relay.const(0), relay.const(0), relay.const(2.0), relay.const(6.0)
+    )
+    op4 = relay.qnn.op.qnn.dequantize(op3, relay.const(12.0), relay.const(0))
+    expected_expr = relay.op.erf(op4)
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    w_np = np.random.randint(-128, 127, size=shape_w, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np, w_np])
+
+
+def test_fq_positive_single_arg_part():
+    # The single-argument operation is converted.
+    shape_x = [1, 2, 4]
+    a = relay.var("a", shape=shape_x, dtype="int8")
+
+    op0 = relay.qnn.op.dequantize(a, relay.const(2.0), relay.const(0))
+
+    op1 = relay.op.reshape(op0, (1, 4, 2))
+    expr = relay.op.erf(op1)
+
+    op0 = relay.op.reshape(a, (1, 4, 2))
+    op1 = relay.qnn.op.dequantize(op0, relay.const(2.0), relay.const(0))
+    expected_expr = relay.op.erf(op1)
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int8")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np])
+
+
+def test_fq_qat_intermediate_infertype():
+    # Complex conversion of non-QAT and QAT passes that form FakeQuantizationToInteger.
+    shape_x = [1, 2, 4]
+    x = relay.var("x", shape=shape_x, dtype="float32")
+    const_0 = relay.const(np.random.uniform(size=[1, 4, 2]).astype("float32"))
+
+    op0 = relay.qnn.op.quantize(x, relay.const(17.0), relay.const(0), out_dtype="int8")
+    op1 = relay.qnn.op.dequantize(op0, relay.const(17.0), relay.const(0))
+    op2 = relay.op.reshape(op1, (1, 4, 2))
+    op3 = relay.qnn.op.quantize(op2, relay.const(10.0), relay.const(0), out_dtype="int8")
+    op4 = relay.qnn.op.quantize(const_0, relay.const(1.0), relay.const(8), out_dtype="int8")
+    op5 = relay.qnn.op.dequantize(op3, relay.const(10.0), relay.const(0))
+    op6 = relay.qnn.op.dequantize(op4, relay.const(4.0), relay.const(9))
+    op7 = relay.op.nn.batch_matmul(op5, op6)
+    expr = relay.op.add(op7, relay.const(5.0))
+
+    op0 = relay.qnn.op.quantize(x, relay.const(17.0), relay.const(0), out_dtype="int8")
+    op1 = relay.op.reshape(op0, (1, 4, 2))
+    op2 = relay.qnn.op.requantize(
+        op1, relay.const(17.0), relay.const(0), relay.const(10.0), relay.const(0), out_dtype="int8"
+    )
+    op3 = relay.qnn.op.quantize(const_0, relay.const(1.0), relay.const(8), out_dtype="int8")
+    op4 = relay.qnn.op.batch_matmul(
+        op2, op3, relay.const(0), relay.const(9), relay.const(10.0), relay.const(4.0)
+    )
+    op5 = relay.qnn.op.dequantize(op4, relay.const(40.0), relay.const(0))
+    expected_expr = relay.op.add(op5, relay.const(5.0))
+
+    x_np = np.random.randint(-128, 127, size=shape_x, dtype="int32").astype("float32")
+    compare_expected_fq_qat_to_int(expr, expected_expr, [x_np])
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/relay/test_pass_fold_constant.py b/tests/python/relay/test_pass_fold_constant.py
index 7b7ef0ce920f..c165d140b1a6 100644
--- a/tests/python/relay/test_pass_fold_constant.py
+++ b/tests/python/relay/test_pass_fold_constant.py
@@ -22,11 +22,6 @@
 from tvm.relay.testing import run_infer_type, create_workload
 
 
-def annot_func(f):
-    """Returns f with arg/result device attributes for the argument and result."""
-    return relay.op.annotation.function_on_device(f, [tvm.cpu()], tvm.cpu())
-
-
 def annot_expr(e):
     """Returns e wrapped with an on_device annotation."""
     return relay.op.annotation.on_device(e, tvm.cpu(), constrain_result=True)
@@ -96,20 +91,24 @@ def test_fold_const_with_on_device():
     def before():
         c = relay.const(c_data)
         x = relay.var("x", t)
+        x.virtual_device_ = tvm.cpu()
         y = relay.add(c, c)
         y = relay.multiply(y, relay.const(2, "float32"))
         y = relay.add(x, y)
         z = relay.add(y, c)
         f = relay.Function([x], z)
-        return annot_func(f)
+        f.virtual_device_ = tvm.cpu()
+        return f
 
     def expected():
         x = relay.var("x", t)
+        x.virtual_device_ = tvm.cpu()
         c_folded = (c_data + c_data) * 2
         y = relay.add(x, relay.const(c_folded))
         z = relay.add(y, relay.const(c_data))
         f = relay.Function([x], z)
-        return annot_func(f)
+        f.virtual_device_ = tvm.cpu()
+        return f
 
     zz = run_opt_pass(before(), transform.FoldConstant())
     zexpected = run_opt_pass(expected(), transform.InferType())
@@ -151,21 +150,25 @@ def test_fold_let_with_on_device():
     def before():
         sb = relay.ScopeBuilder()
         x = relay.var("x", t)
+        x.virtual_device_ = tvm.cpu()
         t1 = sb.let("t1", annot_expr(relay.const(c_data)))
         t2 = sb.let("t2", annot_expr(relay.add(t1, t1)))
         t3 = sb.let("t3", annot_expr(relay.add(t2, x)))
         sb.ret(t3)
         f = relay.Function([x], sb.get())
-        return annot_func(f)
+        f.virtual_device_ = tvm.cpu()
+        return f
 
     def expected():
         sb = relay.ScopeBuilder()
         x = relay.var("x", t)
+        x.virtual_device_ = tvm.cpu()
         c_folded = c_data + c_data
         t3 = sb.let("t3", annot_expr(relay.add(annot_expr(relay.const(c_folded)), x)))
         sb.ret(t3)
         f = relay.Function([x], sb.get())
-        return annot_func(f)
+        f.virtual_device_ = tvm.cpu()
+        return f
 
     zz = run_opt_pass(before(), transform.FoldConstant())
     zexpected = run_opt_pass(expected(), transform.InferType())
diff --git a/tests/python/relay/test_pass_fuse_ops.py b/tests/python/relay/test_pass_fuse_ops.py
index 121d6562594e..cacce5603e5f 100644
--- a/tests/python/relay/test_pass_fuse_ops.py
+++ b/tests/python/relay/test_pass_fuse_ops.py
@@ -625,7 +625,10 @@ def expected(n, max_fused_ops):
     assert tvm.ir.structural_equal(zz, after)
 
 
-def test_fuse_take():
+link_params = tvm.testing.parameter(False, True)
+
+
+def test_fuse_take(link_params):
     """Test fusion case involving concat and take"""
 
     def before():
@@ -635,7 +638,7 @@ def before():
         out = relay.op.take(concat, indices=relay.const([0], dtype="int64"))
         return relay.Function(relay.analysis.free_vars(out), out)
 
-    def expected():
+    def expected(link_params):
         shape1 = (tvm.tir.const(10, "int64"), tvm.tir.const(1, "int64"))
         shape2 = (tvm.tir.const(1, "int64"),)
         x = relay.var("x", shape=shape1)
@@ -643,22 +646,23 @@ def expected():
         p1 = relay.var("p1", shape=shape2, dtype="int64")
         c = relay.const([0], dtype="int64")
         concat = relay.concatenate([p0, p0], axis=-1)
-        out = relay.op.take(concat, indices=p1)
+        out = relay.op.take(concat, indices=c if link_params else p1)
 
-        f0 = relay.Function([p0, p1], out)
+        f0 = relay.Function([p0] if link_params else [p0, p1], out)
         f0 = f0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
 
-        y = relay.Call(f0, [x, c])
+        y = relay.Call(f0, [x] if link_params else [x, c])
         return relay.Function([x], y)
 
-    orig = before()
-    m = fuse2(tvm.IRModule.from_expr(orig))
+    after = run_opt_pass(expected(link_params), transform.InferType())
+    with tvm.transform.PassContext(opt_level=2, config={"relay.FuseOps.link_params": link_params}):
+        m = run_opt_pass(before(), transform.InferType())
+        m = run_opt_pass(m, transform.FuseOps())
+    assert tvm.ir.structural_equal(m, after)
     relay.build(m, "llvm")
-    after = run_opt_pass(expected(), transform.InferType())
-    assert tvm.ir.structural_equal(m["main"], after)
 
 
-def test_fuse_gather_nd():
+def test_fuse_gather_nd(link_params):
     """Test fusion case involving concat and gather_nd"""
 
     def before():
@@ -668,7 +672,7 @@ def before():
         out = relay.gather_nd(concat, indices=relay.expr.const([[0, 1], [1, 0]], dtype="int64"))
         return relay.Function(relay.analysis.free_vars(out), out)
 
-    def expected():
+    def expected(link_params):
         shape1 = (tvm.tir.const(10, "int64"), tvm.tir.const(1, "int64"))
         shape2 = (tvm.tir.const(2, "int64"), tvm.tir.const(2, "int64"))
         x = relay.var("x", shape=shape1)
@@ -676,19 +680,20 @@ def expected():
         p1 = relay.var("p1", shape=shape2, dtype="int64")
         c = relay.const([[0, 1], [1, 0]], dtype="int64")
         concat = relay.concatenate([p0, p0], axis=-1)
-        out = relay.gather_nd(concat, indices=p1)
+        out = relay.gather_nd(concat, indices=c if link_params else p1)
 
-        f0 = relay.Function([p0, p1], out)
+        f0 = relay.Function([p0] if link_params else [p0, p1], out)
         f0 = f0.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
 
-        y = relay.Call(f0, [x, c])
+        y = relay.Call(f0, [x] if link_params else [x, c])
         return relay.Function([x], y)
 
-    orig = before()
-    m = fuse2(tvm.IRModule.from_expr(orig))
+    after = run_opt_pass(expected(link_params), transform.InferType())
+    with tvm.transform.PassContext(opt_level=2, config={"relay.FuseOps.link_params": link_params}):
+        m = run_opt_pass(before(), transform.InferType())
+        m = run_opt_pass(m, transform.FuseOps())
+    assert tvm.ir.structural_equal(m, after)
     relay.build(m, "llvm")
-    after = run_opt_pass(expected(), transform.InferType())
-    assert tvm.ir.structural_equal(m["main"], after)
 
 
 @tvm.testing.uses_gpu
diff --git a/tests/python/relay/test_pass_partition_graph.py b/tests/python/relay/test_pass_partition_graph.py
index 736ece265bde..080c6c803961 100644
--- a/tests/python/relay/test_pass_partition_graph.py
+++ b/tests/python/relay/test_pass_partition_graph.py
@@ -925,7 +925,13 @@ def test_dnnl_fuse():
         conv2d_bias_pat,
         conv2d_relu_pat,
         conv2d_sigmoid_pat,
-    ) = (dnnl_patterns[0], dnnl_patterns[4], dnnl_patterns[6], dnnl_patterns[8], dnnl_patterns[12])
+    ) = (
+        dnnl_patterns[1],
+        dnnl_patterns[13],
+        dnnl_patterns[19],
+        dnnl_patterns[25],
+        dnnl_patterns[37],
+    )
 
     def get_blocks(
         prefix,
@@ -1033,8 +1039,8 @@ def test_partition():
     def test_partition_mobilenet():
         mod, params = relay.testing.mobilenet.get_workload()
         mod = get_partitoned_mod(mod, params, dnnl_patterns)
-        # 27 fused conv + bn + relu and one dense
-        assert len(mod.functions) - 1 == 28  # -1 for main
+        # 27 fused conv + bn + relu, one dense and one softmax
+        assert len(mod.functions) - 1 == 29  # -1 for main
 
     def test_exec(mod, params, ref_mod, ref_params, out_shape):
         ishape = (1, 3, 224, 224)
diff --git a/tests/python/relay/test_pass_plan_devices.py b/tests/python/relay/test_pass_plan_devices.py
index 6232f9c4a96a..f9fe9cf3555b 100644
--- a/tests/python/relay/test_pass_plan_devices.py
+++ b/tests/python/relay/test_pass_plan_devices.py
@@ -26,6 +26,7 @@
 from tvm.script import tir as T
 import tvm.testing
 import numpy as np
+import os
 
 HOST_DEVICE = tvm.device("cpu")
 HOST_TARGET = tvm.target.Target("llvm")
@@ -132,10 +133,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][1], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               %1 = add(%c, %d);
               subtract(%0, %1)
@@ -177,10 +177,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device= meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device= meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
@@ -224,10 +223,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
@@ -271,10 +269,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = add(%c, %d);
@@ -319,8 +316,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
@@ -366,10 +363,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%a, %b);
               let %l = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               let %r = on_device(add(%c, %d), virtual_device=meta[VirtualDevice][1], constrain_result=True);
@@ -416,12 +412,11 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
-              let %f = fn (%x, %y,
-                           param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
+              let %f = fn (%x {virtual_device=meta[VirtualDevice][0]}, %y {virtual_device=meta[VirtualDevice][0]},
+                           virtual_device=meta[VirtualDevice][0]) {
                 add(%x, %y)
               };
               %0 = %f(%a, %b);
@@ -468,12 +463,11 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
-              let %f = fn (%x, %y,
-                           param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
+              let %f = fn (%x {virtual_device=meta[VirtualDevice][0]}, %y {virtual_device=meta[VirtualDevice][0]},
+                           virtual_device=meta[VirtualDevice][0]) {
                 add(%x, %y)
               };
               %1 = %f(%a, %b);
@@ -527,16 +521,16 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
-              let %f = fn (%g, param_virtual_devices=[meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][1]) {
-                fn (%a, param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
+              let %f = fn (%g {virtual_device=meta[VirtualDevice][1]}, virtual_device=meta[VirtualDevice][1]) {
+                fn (%a {virtual_device=meta[VirtualDevice][0]}, virtual_device=meta[VirtualDevice][1]) {
                   %0 = device_copy(%a, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
                   %1 = %g(%0);
                   add(%1, %x)
                 }
               };
-              let %h = fn (%b, param_virtual_devices=[meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][1]) {
+              let %h = fn (%b  {virtual_device=meta[VirtualDevice][1]}, virtual_device=meta[VirtualDevice][1]) {
                 negative(%b)
               };
               %2 = %f(%h);
@@ -588,11 +582,11 @@ def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32]) {
     def expected():
         return tvm.parser.parse(
             """
-            #[version = "0.0.5"] 
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
-              let %f = fn (%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                           param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            #[version = "0.0.5"]
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
+              let %f = fn (%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                           virtual_device=meta[VirtualDevice][0]) {
                 add(%a, %b)
               };
               let %t = on_device((%f, %x), virtual_device=meta[VirtualDevice][0], constrain_result=True);
@@ -619,7 +613,7 @@ def test_device_copy():
     def input():
         return tvm.parser.parse(
             """
-            #[version = "0.0.5"] 
+            #[version = "0.0.5"]
             def @main(%x: Tensor[(5, 7), float32]) {
               %0 = device_copy(%x, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
               add(%0, meta[relay.Constant][0])
@@ -633,9 +627,9 @@ def @main(%x: Tensor[(5, 7), float32]) {
     def expected():
         return tvm.parser.parse(
             """
-            #[version = "0.0.5"] 
-            def @main(%x: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            #[version = "0.0.5"]
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = device_copy(%x, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
               add(%0, meta[relay.Constant][0])
             }
@@ -660,7 +654,7 @@ def test_shape_of():
     def input():
         return tvm.parser.parse(
             """
-            #[version = "0.0.5"] 
+            #[version = "0.0.5"]
             def @main(%x: Tensor[(?, ?), float32]) {
               %0 = on_device(%x, virtual_device=meta[VirtualDevice][1], constrain_result=True);
               vm.shape_of(%0, dtype="int64")
@@ -675,8 +669,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(?, ?), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(?, ?), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               vm.shape_of(%x, dtype="int64")
             }
         """,
@@ -711,8 +705,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%size: int64, %alignment: int64,
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%size {virtual_device=meta[VirtualDevice][0]}: int64, %alignment {virtual_device=meta[VirtualDevice][0]}: int64,
+                      virtual_device=meta[VirtualDevice][1]) {
               memory.alloc_storage(%size, %alignment, virtual_device=meta[VirtualDevice][1])
             }
         """,
@@ -750,7 +744,7 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%sto: Storage[], param_virtual_devices=[meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%sto {virtual_device=meta[VirtualDevice][1]}: Storage[], virtual_device=meta[VirtualDevice][1]) {
               %0 = on_device(0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %1 = on_device(meta[relay.Constant][0], virtual_device=meta[VirtualDevice][0], constrain_result=True);
               memory.alloc_tensor(%sto, %0, %1, const_shape=meta[relay.Constant][0], assert_shape=[])
@@ -789,8 +783,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(2, 8), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(2, 8), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = on_device(meta[relay.Constant][0], virtual_device=meta[VirtualDevice][0], constrain_result=True);
               vm.reshape_tensor(%x, %0, newshape=[2, 4, 2])
             }
@@ -827,8 +821,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x0: Tensor[(?, ?), float32], %x1: Tensor[(?, ?), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x0 {virtual_device=meta[VirtualDevice][0]}: Tensor[(?, ?), float32], %x1 {virtual_device=meta[VirtualDevice][0]}: Tensor[(?, ?), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               add(%x0, %x1)
             }
         """,
@@ -867,9 +861,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32], %z: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %z {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = add(%x, %y);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
@@ -914,9 +907,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32], %z: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][1], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %z {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = add(%x, %y);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][1], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][1], dst_virtual_device=meta[VirtualDevice][0]);
@@ -957,9 +949,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32], %z: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %z {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = add(%x, %y);
               subtract(%0, %z)
             }
@@ -1015,10 +1006,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%data1: Tensor[(1, 64, 56, 56), float32], %data2: Tensor[(1, 64, 56, 56), float32],
-                      %weight: Tensor[(64, 64, 3, 3), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%data1 {virtual_device=meta[VirtualDevice][0]}: Tensor[(1, 64, 56, 56), float32], %data2 {virtual_device=meta[VirtualDevice][0]}: Tensor[(1, 64, 56, 56), float32],
+                      %weight {virtual_device=meta[VirtualDevice][0]}: Tensor[(64, 64, 3, 3), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = nn.conv2d(%data1, %weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = nn.conv2d(%data2, %weight, padding=[1, 1, 1, 1], channels=64, kernel_size=[3, 3]);
@@ -1068,8 +1058,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(3, 3, 4), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(3, 3, 4), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = split(%x, indices_or_sections=3);
               let %t = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %1 = %t.0;
@@ -1136,8 +1126,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = negative(%x);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][0], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
@@ -1206,8 +1196,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][1]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = add(%x, %y);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][1], constrain_result=True);
               %2 = device_copy(%1, src_virtual_device=meta[VirtualDevice][1], dst_virtual_device=meta[VirtualDevice][0]);
@@ -1272,10 +1262,9 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      %c {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %d {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = multiply(%c, %d);
               %1 = on_device(%0, virtual_device=meta[VirtualDevice][1], constrain_result=True);
               %2 = add(%a, %b);
@@ -1327,13 +1316,12 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: bool, %y: Tensor[(5, 7), float32], %z: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0], meta[VirtualDevice][0]],
-                      result_virtual_device=meta[VirtualDevice][0]) {
-              let %f = fn (%a, param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: bool, %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %z {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
+              let %f = fn (%a {virtual_device=meta[VirtualDevice][0]}, virtual_device=meta[VirtualDevice][0]) {
                 add(%a, %y)
               };
-              let %g = fn (%a1, param_virtual_devices=[meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+              let %g = fn (%a1 {virtual_device=meta[VirtualDevice][0]}, virtual_device=meta[VirtualDevice][0]) {
                 subtract(%a1, %y)
               };
               let %h = on_device(if (%x) {
@@ -1373,7 +1361,7 @@ def @f(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32]) -> Tensor[(5, 7
               %0 = on_device(%b, virtual_device=meta[VirtualDevice][0]);
               add(%a, %0)
             }
-            
+
             def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32]) -> Tensor[(5, 7), float32] {
               @f(%y, %x)
             }
@@ -1387,16 +1375,14 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @f(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
-                   param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][0]],
-                   result_virtual_device=meta[VirtualDevice][1]) -> Tensor[(5, 7), float32] {
+            def @f(%a {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                   virtual_device=meta[VirtualDevice][1]) -> Tensor[(5, 7), float32] {
               %0 = device_copy(%b, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
               add(%a, %0)
             }
-            
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][1]) -> Tensor[(5, 7), float32] {
+
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) -> Tensor[(5, 7), float32] {
               @f(%y, %x)
             }
         """,
@@ -1438,8 +1424,8 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x: Tensor[(5, 7), float32], %y: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][1]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               let %r = on_device(ref(%x), virtual_device=meta[VirtualDevice][1], constrain_result=True);
               %0 = device_copy(%y, src_virtual_device=meta[VirtualDevice][0], dst_virtual_device=meta[VirtualDevice][1]);
               on_device(ref_write(%r, %0), virtual_device=meta[VirtualDevice][1], constrain_result=True);
@@ -1495,8 +1481,8 @@ def expected():
               Cons(A, List[A]),
               Nil,
             }
-            def @main(%x : Tensor[(5, 7), float32], %y : Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][0]], result_virtual_device=meta[VirtualDevice][0]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %y {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][0]) {
               %0 = Nil;
               %1 = Cons(%y, %0);
               let %l = on_device(Cons(%x, %1), virtual_device=meta[VirtualDevice][0], constrain_result=True);
@@ -1534,14 +1520,12 @@ def input():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @on_scope_b(%x: Tensor[(5, 7), float32],
-                            param_virtual_devices=[meta[VirtualDevice][2]],
-                            result_virtual_device=meta[VirtualDevice][2]) -> Tensor[(5, 7), float32] {
-              %x                
-            }                 
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32], %c: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][1], meta[VirtualDevice][2]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @on_scope_b(%x {virtual_device=meta[VirtualDevice][2]}: Tensor[(5, 7), float32],
+                            virtual_device=meta[VirtualDevice][2]) -> Tensor[(5, 7), float32] {
+              %x
+            }
+            def @main(%a {virtual_device=meta[VirtualDevice][0]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %c {virtual_device=meta[VirtualDevice][2]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               // %a's memory scope is unconstrained, so will take on "scopeB" and on_device has no effect
               %0 = @on_scope_b(on_device(%a, virtual_device=meta[VirtualDevice][0], constrain_body=False));
               // %b's memory scope is "scopeA", so will require a "scopeA"->"scopeB" copy.
@@ -1562,20 +1546,18 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @on_scope_b(%x: Tensor[(5, 7), float32],
-                            param_virtual_devices=[meta[VirtualDevice][2]],
-                            result_virtual_device=meta[VirtualDevice][2]) -> Tensor[(5, 7), float32] {
-              %x                
-            }                 
-            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32], %c: Tensor[(5, 7), float32],
-                      param_virtual_devices=[meta[VirtualDevice][2], meta[VirtualDevice][1], meta[VirtualDevice][2]],
-                      result_virtual_device=meta[VirtualDevice][1]) {
+            def @on_scope_b(%x {virtual_device=meta[VirtualDevice][2]}: Tensor[(5, 7), float32],
+                            virtual_device=meta[VirtualDevice][2]) -> Tensor[(5, 7), float32] {
+              %x
+            }
+            def @main(%a {virtual_device=meta[VirtualDevice][2]}: Tensor[(5, 7), float32], %b {virtual_device=meta[VirtualDevice][1]}: Tensor[(5, 7), float32], %c {virtual_device=meta[VirtualDevice][2]}: Tensor[(5, 7), float32],
+                      virtual_device=meta[VirtualDevice][1]) {
               %0 = @on_scope_b(%a);
               %1 = device_copy(%b, src_virtual_device=meta[VirtualDevice][1], dst_virtual_device=meta[VirtualDevice][2]);
               %2 = @on_scope_b(%1);
               %3 = @on_scope_b(%c);
               %4 = add(add(%0, %2), %3);
-              %5 = on_device(%4, virtual_device=meta[VirtualDevice][2], constrain_result=True); 
+              %5 = on_device(%4, virtual_device=meta[VirtualDevice][2], constrain_result=True);
               device_copy(%5, src_virtual_device=meta[VirtualDevice][2], dst_virtual_device=meta[VirtualDevice][1])
             }
         """,
@@ -1648,12 +1630,11 @@ def input():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x : Tensor[(128, 128), float32],
-                      %y : Tensor[(128, 128), float32],
-                      %z : Tensor[(128, 128), float32],
-                      param_virtual_devices=[meta[VirtualDevice][0], meta[VirtualDevice][2], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][2]) {
-              call_lowered(@gem, (%x, %y, %z))          
+            def @main(%x {virtual_device=meta[VirtualDevice][0]}: Tensor[(128, 128), float32],
+                      %y {virtual_device=meta[VirtualDevice][2]}: Tensor[(128, 128), float32],
+                      %z {virtual_device=meta[VirtualDevice][1]}: Tensor[(128, 128), float32],
+                      virtual_device=meta[VirtualDevice][2]) {
+              call_lowered(@gem, (%x, %y, %z))
             }
             """,
             "from_string",
@@ -1671,13 +1652,12 @@ def expected():
         return tvm.parser.parse(
             """
             #[version = "0.0.5"]
-            def @main(%x : Tensor[(128, 128), float32],
-                      %y : Tensor[(128, 128), float32],
-                      %z : Tensor[(128, 128), float32], 
-                      param_virtual_devices=[meta[VirtualDevice][1], meta[VirtualDevice][2], meta[VirtualDevice][1]],
-                      result_virtual_device=meta[VirtualDevice][2]) {
+            def @main(%x {virtual_device=meta[VirtualDevice][1]}: Tensor[(128, 128), float32],
+                      %y {virtual_device=meta[VirtualDevice][2]}: Tensor[(128, 128), float32],
+                      %z {virtual_device=meta[VirtualDevice][1]}: Tensor[(128, 128), float32],
+                      virtual_device=meta[VirtualDevice][2]) {
               %0 = device_copy(%z, src_virtual_device=meta[VirtualDevice][1], dst_virtual_device=meta[VirtualDevice][2]);
-              %1 = on_device(%0, virtual_device=meta[VirtualDevice][2], constrain_result=True);      
+              %1 = on_device(%0, virtual_device=meta[VirtualDevice][2], constrain_result=True);
               %2 = call_lowered(@gem, (%x, %y, %1));
               %3 = on_device(%2, virtual_device=meta[VirtualDevice][1], constrain_result=True);
               device_copy(%3, src_virtual_device=meta[VirtualDevice][1], dst_virtual_device=meta[VirtualDevice][2])
@@ -1691,6 +1671,40 @@ def @main(%x : Tensor[(128, 128), float32],
     exercise(input(), expected(), None, None)
 
 
+def test_stack_overflow():
+    metatable = {"VirtualDevice": [CPU, GPU]}
+
+    # Everything defaults to GPU
+    def input():
+        tmp = "test_stack_overflow_input.txt"
+        mod = """
+            #[version = "0.0.5"]
+            def @main(%a: Tensor[(5, 7), float32], %b: Tensor[(5, 7), float32],
+                      %c: Tensor[(5, 7), float32], %d: Tensor[(5, 7), float32]) {
+            %0 = add(%a, %b);
+            %1 = add(%c, %d);
+            """
+
+        end = 1555
+        for i in range(2, end):
+            s1 = "\n\t" + "%" + str(i) + " = add(%" + str(i - 1) + ", %" + str(i - 2) + ");"
+            mod += s1
+        mod += "\n\t" + "add(%" + str(end - 1) + ", %" + str(end - 2) + ")"
+        mod += "\n\t}"
+
+        return tvm.parser.parse(
+            mod,
+            "from_string",
+            None,
+            metatable,
+        )
+
+    config = tvm.target.make_compilation_config(CTXT, TARGETS, HOST_TARGET)
+    actual_mod = relay.transform.InferType()(input())
+    actual_mod = relay.transform.PlanDevices(config)(actual_mod)
+    relay.transform.InferType()(actual_mod)
+
+
 if __name__ == "__main__":
     import sys
     import pytest
diff --git a/tests/python/relay/test_pipeline_executor.py b/tests/python/relay/test_pipeline_executor.py
index 0851d377fe6a..8ab2265db3d6 100644
--- a/tests/python/relay/test_pipeline_executor.py
+++ b/tests/python/relay/test_pipeline_executor.py
@@ -325,8 +325,7 @@ def test_pipeline():
 
             pipe_config[mod3].target = "llvm"
             pipe_config[mod3].dev = tvm.cpu(0)
-
-            # Here is to check the correctness of the configuration generated by API.
+            # Checking the configuration of modules dependency.
             mconfig = pipe_config.get_config()
             assert mconfig["module_connection"] == get_manual_conf([mod1, mod2, mod3], target)
 
@@ -392,7 +391,17 @@ def test_pipeline():
                 for i in range(len(outputs)):
                     tvm.testing.assert_allclose(normal_output[i], outputs[i].numpy())
                     assert not (normal_output[i] == wrong_output[i]).all()
-            pipeline_module_test.stop()
+                # Running the pipeline executor in the pipeline mode.
+                pipeline_module_test.run(False)
+
+            # TODO(huajsj:) Replacing the checking logic with getting output logic.
+            # Checking the statistic value of pipeline.
+            statistic_time = 0
+            while pipeline_module_test.num_executing_pipeline < len(datas):
+                statistic_time = statistic_time + 1
+                # Setting the timeout to 10 seconds.
+                assert statistic_time < 10
+                time.sleep(1)
 
 
 if __name__ == "__main__":
diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py
index 0f73e0c04def..ebfec0fa23a0 100644
--- a/tests/python/relay/test_vm.py
+++ b/tests/python/relay/test_vm.py
@@ -922,6 +922,148 @@ def test_get_input_index(target, dev):
     assert vm_factory.get_input_index(data_0) == 0
     assert vm_factory.get_input_index("invalid") == -1
 
+def get_one_input_relay_mod(tensor_type, shape, data_name):
+    x = relay.var(data_name, shape = shape, dtype = tensor_type)
+    y = relay.exp(x)
+    f = relay.Function([x], y)
+    return IRModule.from_expr(f)
+
+@tvm.testing.parametrize_targets("llvm")
+def test_one_set_input(target, dev):
+    dtype = "float32"
+    in_shape = [1, 2, 3, 3]
+    in_data_name_0 = "d0"
+
+    mod = get_one_input_relay_mod(dtype, in_shape, in_data_name_0)
+
+    # Compile to VMExecutable.
+    vm_exec = vm.compile(mod, target=target)
+    exe = runtime.vm.VirtualMachine(vm_exec, dev)
+
+    data0_core = np.random.uniform(size=in_shape).astype(dtype)
+    data0 = tvm.nd.array(data0_core)
+    ref_res_core = np.exp(data0_core)
+    ref_res = tvm.nd.array(ref_res_core)
+
+    exe.set_input("main", data0)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    data_dict = {in_data_name_0: data0}
+    exe.set_input("main", **data_dict)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+def get_multiple_input_relay_mod(tensor_type, shape, data_name0, data_name1):
+    x, y = [relay.var(c, shape=shape, dtype = tensor_type) for c in [data_name0, data_name1]]
+    f = relay.Function([x, y], x + y)
+    return IRModule.from_expr(f)
+
+@tvm.testing.parametrize_targets("llvm")
+def test_multiple_set_input(target, dev):
+    dtype = "float32"
+    in_shape = [1, 2, 3, 3]
+    in_data_name_0 = "d0"
+    in_data_name_1 = "d1"
+
+    mod = get_multiple_input_relay_mod(dtype, in_shape, in_data_name_0, in_data_name_1)
+
+    # Compile to VMExecutable.
+    vm_exec = vm.compile(mod, target=target)
+    exe = runtime.vm.VirtualMachine(vm_exec, dev)
+
+    data0_core = np.random.uniform(size=in_shape).astype(dtype)
+    data0 = tvm.nd.array(data0_core)
+    data1_core = np.random.uniform(size=in_shape).astype(dtype)
+    data1 = tvm.nd.array(data1_core)
+    ref_res_core = data0_core + data1_core
+    ref_res = tvm.nd.array(ref_res_core)
+
+    exe.set_input("main", data0, data1)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    data_dict = {in_data_name_1: data1, in_data_name_0: data0}
+    exe.set_input("main", **data_dict)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+@tvm.testing.parametrize_targets("llvm")
+def test_one_set_one_input(target, dev):
+    dtype = "float32"
+    in_shape = [1, 2, 3, 3]
+    in_data_name_0 = "d0"
+
+    mod = get_one_input_relay_mod(dtype, in_shape, in_data_name_0)
+
+    # Compile to VMExecutable.
+    vm_exec = vm.compile(mod, target=target)
+    exe = runtime.vm.VirtualMachine(vm_exec, dev)
+
+    data0_core = np.random.uniform(size=in_shape).astype(dtype)
+    data0 = tvm.nd.array(data0_core)
+    ref_res_core = np.exp(data0_core)
+    ref_res = tvm.nd.array(ref_res_core)
+
+    exe.set_one_input("main", 0, data0)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    exe.set_one_input("main", in_data_name_0, data0)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    data_dict = {in_data_name_0: data0}
+    exe.set_one_input("main", **data_dict)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+@tvm.testing.parametrize_targets("llvm")
+def test_multiple_set_one_input(target, dev):
+    dtype = "float32"
+    in_shape = [1, 2, 3, 3]
+    in_data_name_0 = "d0"
+    in_data_name_1 = "d1"
+
+    mod = get_multiple_input_relay_mod(dtype, in_shape, in_data_name_0, in_data_name_1)
+
+    # Compile to VMExecutable.
+    vm_exec = vm.compile(mod, target=target)
+    exe = runtime.vm.VirtualMachine(vm_exec, dev)
+
+    data0_core = np.random.uniform(size=in_shape).astype(dtype)
+    data0 = tvm.nd.array(data0_core)
+    data1_core = np.random.uniform(size=in_shape).astype(dtype)
+    data1 = tvm.nd.array(data1_core)
+    ref_res_core = data0_core + data1_core
+    ref_res = tvm.nd.array(ref_res_core)
+
+    exe.set_one_input("main", 1, data1)
+    exe.set_one_input("main", 0, data0)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    exe.set_one_input("main", in_data_name_1, data1)
+    exe.set_one_input("main", in_data_name_0, data0)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
+
+    data_dict = {in_data_name_1: data1}
+    exe.set_one_input("main", **data_dict)
+    data_dict = {in_data_name_0: data0}
+    exe.set_one_input("main", **data_dict)
+    output = exe.invoke("main")
+    assert output.dtype == ref_res.dtype
+    tvm.testing.assert_allclose(ref_res_core, output.numpy())
 
 @tvm.testing.parametrize_targets("llvm")
 def test_benchmark(target, dev):
diff --git a/tests/python/relay/utils/external_codegen.py b/tests/python/relay/utils/external_codegen.py
index 680373c7e68b..4dbc8f274264 100644
--- a/tests/python/relay/utils/external_codegen.py
+++ b/tests/python/relay/utils/external_codegen.py
@@ -110,7 +110,7 @@ def check_aot_executor_result(
     use_unpacked_api = False
     test_runner = AOT_DEFAULT_RUNNER
     compile_and_run(
-        AOTTestModel(module=mod, inputs=map_inputs, outputs=[result]),
+        AOTTestModel(module=mod, inputs=map_inputs, outputs={"output": result}),
         test_runner,
         interface_api,
         use_unpacked_api,
diff --git a/tests/python/topi/python/test_topi_conv2d_int8.py b/tests/python/topi/python/test_topi_conv2d_int8.py
index 66d280e85c75..7b36c8a5c1b8 100644
--- a/tests/python/topi/python/test_topi_conv2d_int8.py
+++ b/tests/python/topi/python/test_topi_conv2d_int8.py
@@ -225,9 +225,6 @@ def check_target(target):
     check_target("llvm")
 
 
-oc_block_factor = 4
-
-
 def verify_conv2d_NCHWc_int8(
     batch,
     in_channel,
@@ -251,40 +248,12 @@ def verify_conv2d_NCHWc_int8(
 
     A = te.placeholder((batch, in_channel, in_height, in_width), name="A", dtype="int8")
     W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W", dtype="int8")
-    bias = te.placeholder(
-        (num_filter // oc_block_factor, 1, 1, oc_block_factor), name="bias", dtype="int8"
-    )
 
     a_shape = get_const_tuple(A.shape)
     w_shape = get_const_tuple(W.shape)
-    bias_shape = get_const_tuple(bias.shape)
     dtype = A.dtype
 
-    @memoize("topi.tests.test_topi_conv2d_int8.verify_conv2d_nchw")
-    def get_ref_data():
-        a_np = np.random.randint(low=-128, high=127, size=a_shape).astype(dtype)
-        w_np = np.random.randint(low=-128, high=128, size=w_shape).astype(dtype)
-        b_np = np.random.uniform(size=bias_shape).astype(dtype)
-        dw_np = tvm.topi.testing.dilate_python(w_np, (1, 1, dilation, dilation))
-        c_np = tvm.topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding).astype(dtype)
-
-        # convert to NCHWc
-        _, _, out_height, out_width = c_np.shape
-        c_np = c_np.reshape(
-            (batch, num_filter // oc_block_factor, oc_block_factor, out_height, out_width)
-        ).transpose(0, 1, 3, 4, 2)
-
-        if add_bias:
-            b_np = np.random.uniform(size=bias_shape).astype(dtype)
-            c_np += b_np
-        if add_relu:
-            c_np = np.maximum(c_np, 0)
-
-        return a_np, w_np, b_np, c_np
-
-    a_np, w_np, b_np, c_np = get_ref_data()
-
-    def check_target(target):
+    def check_target(target, compute, schedule, oc_block_factor):
         dev = tvm.device(target, 0)
         if not tvm.testing.device_enabled(target):
             print("Skip because %s is not enabled" % target)
@@ -293,20 +262,58 @@ def check_target(target):
             print("Skip because int8 intrinsics are not available")
             return
 
+        bias = te.placeholder(
+            (num_filter // oc_block_factor, 1, 1, oc_block_factor), name="bias", dtype="int32"
+        )
+        bias_shape = get_const_tuple(bias.shape)
+
+        @memoize("topi.tests.test_topi_conv2d_int8.verify_conv2d_nchw")
+        def get_ref_data():
+            a_np = np.random.randint(low=-128, high=127, size=a_shape).astype("int32")
+            w_np = np.random.randint(low=-128, high=128, size=w_shape).astype("int32")
+            b_np = np.random.uniform(size=bias_shape).astype("int32")
+            dw_np = tvm.topi.testing.dilate_python(w_np, (1, 1, dilation, dilation))
+            c_np = tvm.topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding).astype("int32")
+
+            # convert to NCHWc
+            _, _, out_height, out_width = c_np.shape
+            c_np = c_np.reshape(
+                (batch, num_filter // oc_block_factor, oc_block_factor, out_height, out_width)
+            ).transpose(0, 1, 3, 4, 2)
+
+            if add_bias:
+                b_np = np.random.uniform(size=bias_shape).astype("int32")
+                c_np += b_np
+            if add_relu:
+                c_np = np.maximum(c_np, 0)
+
+            return a_np, w_np, b_np, c_np
+
+        a_np, w_np, b_np, c_np = get_ref_data()
+
         print("Running on target: %s" % target)
         with tvm.target.Target(target):
-            C = topi.cuda.conv2d_NCHWc_int8(
-                A, W, (stride, stride), padding, (dilation, dilation), "NCHW", dtype
+            C = compute(
+                A,
+                W,
+                (stride, stride),
+                padding,
+                (dilation, dilation),
+                "NCHW",
+                "NCHW",
+                "int32",
             )
+            print(C.shape)
+            print(bias.shape)
             if add_bias:
                 C = topi.add(C, bias)
             if add_relu:
                 C = topi.nn.relu(C)
-            s = topi.cuda.schedule_conv2d_NCHWc_int8([C])
+            s = schedule([C])
 
-        a = tvm.nd.array(a_np, dev)
-        w = tvm.nd.array(w_np, dev)
-        b = tvm.nd.array(b_np, dev)
+        a = tvm.nd.array(a_np.astype(dtype), dev)
+        w = tvm.nd.array(w_np.astype(dtype), dev)
+        b = tvm.nd.array(b_np.astype("int32"), dev)
         c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), dev)
         if add_bias:
             tvm.build(
@@ -323,7 +330,14 @@ def check_target(target):
                 name="relu_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation),
             )
-            func(a, w, b, c)
+            try:
+                func(a, w, b, c)
+            except tvm.TVMError as e:
+                if "architecture mismatch" in str(e):
+                    print(f"Skipping execution because {target} is not supported by this CPU")
+                    return
+                else:
+                    raise
         else:
             func = tvm.build(
                 s,
@@ -332,11 +346,32 @@ def check_target(target):
                 name="relu_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (batch, in_channel, in_size, num_filter, kernel, stride, padding_sum, dilation),
             )
-            func(a, w, c)
+            try:
+                func(a, w, c)
+            except tvm.TVMError as e:
+                if "architecture mismatch" in str(e):
+                    print(f"Skipping execution because {target} is not supported by this CPU")
+                    return
+                else:
+                    raise
         tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
 
-    for target in ["cuda"]:
-        check_target(target)
+    targets = [
+        (
+            "cuda",
+            lambda a, w, s, p, d, l, ol, o: topi.cuda.conv2d_NCHWc_int8(a, w, s, p, d, l, o),
+            topi.cuda.schedule_conv2d_NCHWc_int8,
+            4,
+        ),
+        (
+            "llvm -device arm_cpu -mtriple aarch64-linux-gnu -mattr=+neon",
+            topi.arm_cpu.conv2d_NCHWc_int8,
+            topi.arm_cpu.schedule_conv2d_NCHWc_int8,
+            8,
+        ),
+    ]
+    for target, compute, schedule, oc_block_factor in targets:
+        check_target(target, compute, schedule, oc_block_factor)
 
 
 def verify_conv2d_nchw_int8(
@@ -489,7 +524,7 @@ def test_conv2d_nchw():
         verify_conv2d_NCHWc_int8(9, 64, 56, 64, 3, 1, 1)
 
         # weird workloads
-        verify_conv2d_NCHWc_int8(4, 4, 4, 4, 4, 4, 4)
+        verify_conv2d_NCHWc_int8(4, 4, 4, 8, 4, 4, 4)
 
         # inception v3 workloads where channels in / out are multiple of oc_block_factor
         verify_conv2d_NCHWc_int8(1, 32, 149, 32, 3, 1, 0)
@@ -534,7 +569,7 @@ def test_conv2d_nchw():
         verify_conv2d_NCHWc_int8(1, 2048, 8, 384, 1, 1, 0)
         verify_conv2d_NCHWc_int8(1, 2048, 8, 448, 1, 1, 0)
         verify_conv2d_NCHWc_int8(1, 2048, 8, 192, 1, 1, 0)
-        verify_conv2d_NCHWc_int8(1, 1024, 19, 84, 3, 1, 1)
+        verify_conv2d_NCHWc_int8(1, 1024, 19, 88, 3, 1, 1)
 
         # batch > 1
         verify_conv2d_NCHWc_int8(7, 32, 149, 32, 3, 1, 0)
@@ -549,7 +584,7 @@ def test_conv2d_nchw():
         verify_conv2d_NCHWc_int8(1, 64, 8, 64, 3, 1, (3, 1))
         verify_conv2d_NCHWc_int8(1, 128, 8, 384, 3, 1, (0, 2))
         verify_conv2d_NCHWc_int8(1, 64, 8, 64, 1, 1, "VALID")
-        verify_conv2d_NCHWc_int8(1, 388, 8, 64, 3, 1, "VALID")
+        verify_conv2d_NCHWc_int8(1, 392, 8, 64, 3, 1, "VALID")
         verify_conv2d_NCHWc_int8(1, 512, 19, 64, 1, 1, "SAME")
         verify_conv2d_NCHWc_int8(1, 64, 16, 32, 2, 1, "SAME")
         verify_conv2d_NCHWc_int8(1, 64, 8, 64, 3, 1, (1, 2, 2, 1), add_relu=True)
diff --git a/tests/python/topi/python/test_topi_conv3d_ncdhw.py b/tests/python/topi/python/test_topi_conv3d_ncdhw.py
index ea94a7d8f485..6c2ed6eeae21 100644
--- a/tests/python/topi/python/test_topi_conv3d_ncdhw.py
+++ b/tests/python/topi/python/test_topi_conv3d_ncdhw.py
@@ -43,6 +43,7 @@ def verify_conv3d_ncdhw(
     stride,
     padding,
     dilation=1,
+    groups=1,
     add_bias=False,
     add_relu=False,
 ):
@@ -80,7 +81,7 @@ def verify_conv3d_ncdhw(
     in_depth = in_height = in_width = in_size
 
     A = te.placeholder((batch, in_channel, in_depth, in_height, in_width), name="A")
-    W = te.placeholder((num_filter, in_channel, kernel_d, kernel_h, kernel_w), name="W")
+    W = te.placeholder((num_filter, in_channel // groups, kernel_d, kernel_h, kernel_w), name="W")
     bias = te.placeholder((num_filter, 1, 1, 1), name="bias")
 
     a_shape = get_const_tuple(A.shape)
@@ -94,7 +95,7 @@ def get_ref_data():
         w_np = np.random.uniform(size=w_shape).astype(dtype)
         b_np = np.random.uniform(size=bias_shape).astype(dtype)
         dw_np = tvm.topi.testing.dilate_python(w_np, (1, 1, dilation, dilation, dilation))
-        c_np = tvm.topi.testing.conv3d_ncdhw_python(a_np, dw_np, stride, padding)
+        c_np = tvm.topi.testing.conv3d_ncdhw_python(a_np, dw_np, stride, padding, groups)
         if add_bias:
             c_np += b_np
         if add_relu:
@@ -108,7 +109,13 @@ def check_target(target, dev):
         fcompute, fschedule = tvm.topi.testing.dispatch(target, _conv3d_ncdhw_implement)
         with tvm.target.Target(target):
             C = fcompute(
-                A, W, (stride, stride, stride), padding, (dilation, dilation, dilation), dtype
+                A,
+                W,
+                (stride, stride, stride),
+                padding,
+                (dilation, dilation, dilation),
+                groups,
+                dtype,
             )
             if add_bias:
                 C = topi.add(C, bias)
@@ -125,7 +132,7 @@ def check_target(target, dev):
                 s,
                 [A, W, bias, C],
                 target,
-                name="relu_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d"
+                name="relu_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (
                     batch,
                     in_channel,
@@ -137,6 +144,7 @@ def check_target(target, dev):
                     stride,
                     padding_sum,
                     dilation,
+                    groups,
                 ),
             )
             func(a, w, b, c)
@@ -145,7 +153,7 @@ def check_target(target, dev):
                 s,
                 [A, W, C],
                 target,
-                name="relu_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d"
+                name="relu_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d_%d"
                 % (
                     batch,
                     in_channel,
@@ -157,6 +165,7 @@ def check_target(target, dev):
                     stride,
                     padding_sum,
                     dilation,
+                    groups,
                 ),
             )
             func(a, w, c)
@@ -205,6 +214,12 @@ def test_conv3d_ncdhw():
     verify_conv3d_ncdhw(1, 3, 56, 16, (3, 3, 7), 2, (1, 2, 3))
     verify_conv3d_ncdhw(1, 3, 56, 16, (3, 7, 3), 2, (1, 3, 1))
 
+    # grouped workloads
+    verify_conv3d_ncdhw(1, 32, 32, 8, 1, 1, 0, groups=4)
+    verify_conv3d_ncdhw(1, 32, 32, 4, 1, 1, 0, groups=4)
+    verify_conv3d_ncdhw(1, 32, 32, 8, 1, 1, 1, groups=4)
+    verify_conv3d_ncdhw(1, 32, 32, 4, 1, 1, 1, groups=4)
+
 
 if __name__ == "__main__":
     test_conv3d_ncdhw()
diff --git a/tests/python/topi/python/test_topi_conv3d_ndhwc.py b/tests/python/topi/python/test_topi_conv3d_ndhwc.py
index a9437645b7bb..4ee31ecacd0f 100644
--- a/tests/python/topi/python/test_topi_conv3d_ndhwc.py
+++ b/tests/python/topi/python/test_topi_conv3d_ndhwc.py
@@ -34,7 +34,17 @@
 
 
 def verify_conv3d_ndhwc(
-    batch, in_channel, in_size, num_filter, kernel, stride, padding, dilation=1
+    target,
+    dev,
+    batch,
+    in_channel,
+    in_size,
+    num_filter,
+    kernel,
+    stride,
+    padding,
+    dilation=1,
+    groups=1,
 ):
     if isinstance(in_size, tuple):
         in_depth, in_height, in_width = in_size
@@ -47,7 +57,7 @@ def verify_conv3d_ndhwc(
 
     A = te.placeholder((batch, in_depth, in_height, in_width, in_channel), name="A")
     W = te.placeholder(
-        (kernel_depth, kernel_height, kernel_width, in_channel, num_filter), name="W"
+        (kernel_depth, kernel_height, kernel_width, in_channel // groups, num_filter), name="W"
     )
 
     a_shape = get_const_tuple(A.shape)
@@ -59,44 +69,43 @@ def get_ref_data():
         a_np = np.random.uniform(size=a_shape).astype(dtype)
         w_np = np.random.uniform(size=w_shape).astype(dtype)
         dw_np = tvm.topi.testing.dilate_python(w_np, (dilation, dilation, dilation, 1, 1))
-        b_np = tvm.topi.testing.conv3d_ndhwc_python(a_np, dw_np, stride, padding)
+        b_np = tvm.topi.testing.conv3d_ndhwc_python(a_np, dw_np, stride, padding, groups)
         return a_np, w_np, b_np
 
     a_np, w_np, b_np = get_ref_data()
 
-    def check_target(target, dev):
-        print("Running on target: %s" % target)
-        fcompute, fschedule = tvm.topi.testing.dispatch(target, _conv3d_ndhwc_implement)
-        with tvm.target.Target(target):
-            B = fcompute(A, W, stride, padding, dilation, dtype)
-            s = fschedule([B])
-        dev = tvm.device(target, 0)
-        a = tvm.nd.array(a_np, dev)
-        w = tvm.nd.array(w_np, dev)
-        b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
-        func = tvm.build(s, [A, W, B], target)
-        func(a, w, b)
-        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
+    fcompute, fschedule = tvm.topi.testing.dispatch(target, _conv3d_ndhwc_implement)
+    with tvm.target.Target(target):
+        B = fcompute(A, W, stride, padding, dilation, groups, dtype)
+        s = fschedule([B])
+    dev = tvm.device(target, 0)
+    a = tvm.nd.array(a_np, dev)
+    w = tvm.nd.array(w_np, dev)
+    b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), dev)
+    func = tvm.build(s, [A, W, B], target)
+    print(tvm.lower(s, [A, W, B], target))
+    func(a, w, b)
+    tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-5)
 
-    for target, dev in tvm.testing.enabled_targets():
-        check_target(target, dev)
 
-
-@tvm.testing.uses_gpu
-def test_conv3d_ndhwc():
-    verify_conv3d_ndhwc(1, 16, 32, 16, 3, 1, "SAME")
-    verify_conv3d_ndhwc(4, 32, 16, 32, 5, 2, "SAME")
-    verify_conv3d_ndhwc(4, 32, 16, 64, 5, 2, "SAME")
-    verify_conv3d_ndhwc(1, 64, 32, 64, 3, 1, "VALID")
-    verify_conv3d_ndhwc(1, 64, 32, 64, 3, 1, "VALID")
-    verify_conv3d_ndhwc(4, 32, 16, 32, 5, 2, "VALID")
-    verify_conv3d_ndhwc(4, 32, 16, 64, 5, 2, "VALID")
+def test_conv3d_ndhwc(target, dev):
+    verify_conv3d_ndhwc(target, dev, 1, 16, 32, 16, 3, 1, "SAME")
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 32, 5, 2, "SAME")
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 64, 5, 2, "SAME")
+    verify_conv3d_ndhwc(target, dev, 1, 64, 32, 64, 3, 1, "VALID")
+    verify_conv3d_ndhwc(target, dev, 1, 64, 32, 64, 3, 1, "VALID")
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 32, 5, 2, "VALID")
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 64, 5, 2, "VALID")
     # dilation = 2
-    verify_conv3d_ndhwc(1, 64, 32, 64, 3, 1, "SAME", dilation=2)
+    verify_conv3d_ndhwc(target, dev, 1, 64, 32, 64, 3, 1, "SAME", dilation=2)
+
+    verify_conv3d_ndhwc(target, dev, 1, 1, (20, 256, 256), 32, (1, 3, 3), (1, 2, 2), "SAME")
+    verify_conv3d_ndhwc(target, dev, 1, 1, (20, 256, 256), 32, (1, 6, 6), (1, 2, 2), (0, 2, 2))
+    verify_conv3d_ndhwc(target, dev, 1, 4, (20, 256, 256), 8, (1, 5, 5), (1, 2, 2), (0, 2, 2))
 
-    verify_conv3d_ndhwc(1, 1, (20, 256, 256), 32, (1, 3, 3), (1, 2, 2), "SAME")
-    verify_conv3d_ndhwc(1, 1, (20, 256, 256), 32, (1, 6, 6), (1, 2, 2), (0, 2, 2))
-    verify_conv3d_ndhwc(1, 4, (20, 256, 256), 8, (1, 5, 5), (1, 2, 2), (0, 2, 2))
+    verify_conv3d_ndhwc(target, dev, 1, 16, 32, 16, 3, 1, "SAME", groups=4)
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 32, 5, 2, "SAME", groups=4)
+    verify_conv3d_ndhwc(target, dev, 4, 32, 16, 64, 5, 2, "SAME", groups=4)
 
 
 if __name__ == "__main__":
diff --git a/tests/python/topi/python/test_topi_conv3d_ndhwc_tensorcore.py b/tests/python/topi/python/test_topi_conv3d_ndhwc_tensorcore.py
index c8c54e4c3e1f..0d587a87d8d7 100644
--- a/tests/python/topi/python/test_topi_conv3d_ndhwc_tensorcore.py
+++ b/tests/python/topi/python/test_topi_conv3d_ndhwc_tensorcore.py
@@ -59,14 +59,14 @@ def verify_conv3d_ndhwc(
 
     in_depth = in_height = in_width = in_size
 
-    A = te.placeholder((batch, in_depth, in_height, in_width, in_channel), name="A")
-    W = te.placeholder((kernel, kernel, kernel, in_channel, num_filter), name="W")
-    bias = te.placeholder((1, 1, 1, 1, num_filter), name="bias")
+    dtype = "float16"
+    A = te.placeholder((batch, in_depth, in_height, in_width, in_channel), dtype, name="A")
+    W = te.placeholder((kernel, kernel, kernel, in_channel, num_filter), dtype, name="W")
+    bias = te.placeholder((1, 1, 1, 1, num_filter), dtype, name="bias")
 
     a_shape = get_const_tuple(A.shape)
     w_shape = get_const_tuple(W.shape)
     bias_shape = get_const_tuple(bias.shape)
-    dtype = A.dtype
 
     @memoize("topi.tests.test_topi_conv3d_ndhwc.verify_conv3d_ndhwc")
     def get_ref_data():
@@ -91,7 +91,7 @@ def check_device(device):
             fcompute, fschedule = tvm.topi.testing.dispatch(
                 device, _conv3d_ndhwc_tensorcore_implement
             )
-            C = fcompute(A, W, stride, padding, dilation, "float32")
+            C = fcompute(A, W, stride, padding, dilation, 1, "float16")
             if add_bias:
                 C = topi.add(C, bias)
             if add_relu:
@@ -121,8 +121,10 @@ def check_device(device):
             )
             func(a, w, c)
 
-        rtol = 1e-3
-        tvm.testing.assert_allclose(c.numpy(), c_np, rtol=rtol)
+        # Tensorcores are very inaccurate, with large shapes, the accumulation
+        # error is high especially away from 1. We disable atol as it is very
+        # large for these numbers that are far away from 1.
+        tvm.testing.assert_allclose(c.numpy(), c_np, atol=1e200, rtol=0.01)
 
     check_device(devices)
 
diff --git a/tests/python/topi/python/test_topi_conv3d_winograd.py b/tests/python/topi/python/test_topi_conv3d_winograd.py
index 54dd72a2f544..af613437d6cd 100644
--- a/tests/python/topi/python/test_topi_conv3d_winograd.py
+++ b/tests/python/topi/python/test_topi_conv3d_winograd.py
@@ -90,7 +90,7 @@ def check_device(device):
         fcompute, fschedule = tvm.topi.testing.dispatch(device, _conv3d_ncdhw_implement)
         with tvm.target.Target(device):
             C = fcompute(
-                A, W, (stride, stride, stride), padding, (dilation, dilation, dilation), dtype
+                A, W, (stride, stride, stride), padding, (dilation, dilation, dilation), 1, dtype
             )
             if add_bias:
                 C = topi.add(C, bias)
diff --git a/tests/python/unittest/test_aot_legalize_packed_call.py b/tests/python/unittest/test_aot_legalize_packed_call.py
index 222d647f4ea7..54561ade23e4 100644
--- a/tests/python/unittest/test_aot_legalize_packed_call.py
+++ b/tests/python/unittest/test_aot_legalize_packed_call.py
@@ -57,8 +57,17 @@ def tir_packed_call() -> None:
             with T.let(tvm_value_1, T.tvm_stack_alloca("array", 1, dtype="handle")):
                 with T.let(tvm_value_0, T.tvm_stack_alloca("array", 1, dtype="handle")):
                     T.evaluate(T.tvm_struct_set(tvm_value_0, 0, 1, A, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_0, 0, 10, 1, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_0, 0, 9, 0, dtype="handle"))
+
                     T.evaluate(T.tvm_struct_set(tvm_value_1, 0, 1, B, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_1, 0, 10, 1, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_1, 0, 9, 0, dtype="handle"))
+
                     T.evaluate(T.tvm_struct_set(tvm_value_2, 0, 1, C, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_2, 0, 10, 1, dtype="handle"))
+                    T.evaluate(T.tvm_struct_set(tvm_value_2, 0, 9, 0, dtype="handle"))
+
                     T.evaluate(
                         T.tvm_call_cpacked(
                             "tvm_test_cpacked",
diff --git a/tests/python/unittest/test_arith_canonical_simplify.py b/tests/python/unittest/test_arith_canonical_simplify.py
index c241b81da986..6dc91d780413 100644
--- a/tests/python/unittest/test_arith_canonical_simplify.py
+++ b/tests/python/unittest/test_arith_canonical_simplify.py
@@ -142,6 +142,7 @@ def test_canonical_mixed():
     ck.verify(tvm.te.max(x, 1) - tvm.te.max(x, 1), 0)
     ck.verify(tvm.te.min(x, 1) - tvm.te.min(x, 1), 0)
     ck.verify(x * x - x * x, 0)
+    ck.verify(tmod(tdiv(tmod(x, 20), 2) * 2, 4), tdiv(tmod(x, 4), 2) * 2)
 
     fld = tvm.te.floordiv
     ck.verify(fld(x, (z * z)) - fld(x, (z * z)), 0)
diff --git a/tests/python/unittest/test_arith_intset.py b/tests/python/unittest/test_arith_intset.py
index b40f3c9f56ea..e6738543a6aa 100644
--- a/tests/python/unittest/test_arith_intset.py
+++ b/tests/python/unittest/test_arith_intset.py
@@ -31,7 +31,7 @@ def err_msg():
             return "\ndata={}\ndmap={}\nres={}\nexpected={}".format(data, dmap, res, expected)
 
         def equal(x, y):
-            res = self.analyzer.canonical_simplify(x - y)
+            res = self.analyzer.simplify(x - y)
             return tvm.tir.analysis.expr_deep_equal(res, 0)
 
         assert equal(res.min_value, expected[0]), err_msg()
@@ -99,10 +99,14 @@ def test_mod():
     ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(3, 11)}, (0, 9))
     ck.verify(flm(x, 10), {x: tvm.arith.IntervalSet(1, 21)}, (0, 9))
 
-    floordiv = tvm.te.floordiv
+    fld = tvm.te.floordiv
     z = te.var("z")
     ck.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 3))
-    ck.verify(flm(y, 8), {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)}, (0, 7))
+    ck.verify(
+        flm(y, 8),
+        {y: tvm.arith.IntervalSet(z * 8 + x * 4, z * 8 + x * 4 + 3)},
+        (x * 4 - 8 * fld(x * 4, 8), x * 4 - 8 * fld(x * 4, 8) + 3),
+    )
     ck1 = IntSetChecker()
     ck1.analyzer.bind(x, tvm.ir.Range.from_min_extent(0, 2))
     ck1.verify(
diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py
index 6ca2a2a5fcb0..662038e129f7 100644
--- a/tests/python/unittest/test_arith_rewrite_simplify.py
+++ b/tests/python/unittest/test_arith_rewrite_simplify.py
@@ -462,6 +462,10 @@ def test_floordiv_index_simplify():
 
     ck.verify(fld(x * 2, 4), fld(x, 2))
     ck.verify(fld(x * 4, 2), x * 2)
+    ck.verify(fld(x * 8 + 7, 16), fld(x, 2))
+    ck.verify(fld(x * 8 + 39, 16), fld(x, 2) + 2)
+    ck.verify(fld(x * 8 - 1, 16), fld(x * 8 + -1, 16))
+    ck.verify(fld(x * 8 - 9, 16), fld(x, 2) + -1)
 
     ck.verify(fld(x * 4 + y, 2), x * 2 + fld(y, 2))
     ck.verify(fld(tvm.te.min(x * 6, y), 2), tvm.te.min(x * 3, fld(y, 2)))
diff --git a/tests/python/unittest/test_auto_scheduler_task_scheduler.py b/tests/python/unittest/test_auto_scheduler_task_scheduler.py
index a3f356929dd1..66e1ba9d6802 100644
--- a/tests/python/unittest/test_auto_scheduler_task_scheduler.py
+++ b/tests/python/unittest/test_auto_scheduler_task_scheduler.py
@@ -101,7 +101,7 @@ def test_task_scheduler_gradient():
         )
 
     def objective_func(costs):
-        return costs[0]
+        return 1e5 * costs[0]
 
     with tempfile.NamedTemporaryFile() as fp:
         log_file = fp.name
diff --git a/tests/python/unittest/test_ci.py b/tests/python/unittest/test_ci.py
index f5183d79b768..6ca46bc60cd5 100644
--- a/tests/python/unittest/test_ci.py
+++ b/tests/python/unittest/test_ci.py
@@ -68,6 +68,107 @@ def run(pr_body, expected_reviewers):
     )
 
 
+def test_update_branch(tmpdir_factory):
+    update_script = REPO_ROOT / "tests" / "scripts" / "update_branch.py"
+
+    def run(statuses, expected_rc, expected_output):
+        git = TempGit(tmpdir_factory.mktemp("tmp_git_dir"))
+        git.run("init")
+        git.run("checkout", "-b", "main")
+        git.run("remote", "add", "origin", "https://github.com/apache/tvm.git")
+        commit = {
+            "statusCheckRollup": {"contexts": {"nodes": statuses}},
+            "oid": "123",
+            "messageHeadline": "hello",
+        }
+        data = {
+            "data": {
+                "repository": {
+                    "defaultBranchRef": {"target": {"history": {"edges": [], "nodes": [commit]}}}
+                }
+            }
+        }
+        proc = subprocess.run(
+            [str(update_script), "--dry-run", "--testonly-json", json.dumps(data)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            encoding="utf-8",
+            cwd=git.cwd,
+        )
+
+        if proc.returncode != expected_rc:
+            raise RuntimeError(
+                f"Wrong return code:\nstdout:\n{proc.stdout}\n\nstderr:\n{proc.stderr}"
+            )
+
+        if expected_output not in proc.stdout:
+            raise RuntimeError(
+                f"Missing {expected_output}:\nstdout:\n{proc.stdout}\n\nstderr:\n{proc.stderr}"
+            )
+
+    # Missing expected tvm-ci/branch test
+    run(
+        statuses=[
+            {
+                "context": "test",
+                "state": "SUCCESS",
+            }
+        ],
+        expected_rc=1,
+        expected_output="No good commits found in the last 1 commits",
+    )
+
+    # Only has the right passing test
+    run(
+        statuses=[
+            {
+                "context": "tvm-ci/branch",
+                "state": "SUCCESS",
+            }
+        ],
+        expected_rc=0,
+        expected_output="Found last good commit: 123: hello",
+    )
+
+    # Check with many statuses
+    run(
+        statuses=[
+            {
+                "context": "tvm-ci/branch",
+                "state": "SUCCESS",
+            },
+            {
+                "context": "tvm-ci/branch2",
+                "state": "SUCCESS",
+            },
+            {
+                "context": "tvm-ci/branch3",
+                "state": "FAILED",
+            },
+        ],
+        expected_rc=1,
+        expected_output="No good commits found in the last 1 commits",
+    )
+    run(
+        statuses=[
+            {
+                "context": "tvm-ci/branch",
+                "state": "SUCCESS",
+            },
+            {
+                "context": "tvm-ci/branch2",
+                "state": "SUCCESS",
+            },
+            {
+                "context": "tvm-ci/branch3",
+                "state": "SUCCESS",
+            },
+        ],
+        expected_rc=0,
+        expected_output="Found last good commit: 123: hello",
+    )
+
+
 def test_skip_ci(tmpdir_factory):
     skip_ci_script = REPO_ROOT / "tests" / "scripts" / "git_skip_ci.py"
 
@@ -206,11 +307,20 @@ def run(pr, check):
 
         assert check in proc.stdout
 
+    def all_time_keys(time):
+        return {
+            "updatedAt": time,
+            "lastEditedAt": time,
+            "createdAt": time,
+            "publishedAt": time,
+        }
+
     run(
         {
             "isDraft": True,
+            "number": 2,
         },
-        "Checking 0 PRs",
+        "Checking 0 of 1 fetched",
     )
 
     run(
@@ -218,7 +328,7 @@ def run(pr, check):
             "isDraft": False,
             "number": 2,
         },
-        "Checking 0 PRs",
+        "Checking 0 of 1 fetched",
     )
 
     run(
@@ -229,7 +339,7 @@ def run(pr, check):
             "isDraft": False,
             "author": {"login": "user"},
             "reviews": {"nodes": []},
-            "publishedAt": "2022-01-18T17:54:19Z",
+            **all_time_keys("2022-01-18T17:54:19Z"),
             "comments": {"nodes": []},
         },
         "Pinging reviewers ['someone'] on https://github.com/apache/tvm/pull/123",
@@ -244,14 +354,14 @@ def run(pr, check):
             "isDraft": False,
             "author": {"login": "user2"},
             "reviews": {"nodes": []},
-            "publishedAt": "2022-01-18T17:54:19Z",
+            **all_time_keys("2022-01-18T17:54:19Z"),
             "comments": {
                 "nodes": [
-                    {"updatedAt": "2022-01-19T17:54:19Z", "bodyText": "abc"},
+                    {**all_time_keys("2022-01-19T17:54:19Z"), "bodyText": "abc"},
                 ]
             },
         },
-        "Checking 0 PRs",
+        "Checking 0 of 1 fetched",
     )
 
     # Old comment, ping
@@ -263,10 +373,13 @@ def run(pr, check):
             "isDraft": False,
             "author": {"login": "user"},
             "reviews": {"nodes": []},
-            "publishedAt": "2022-01-18T17:54:19Z",
+            **all_time_keys("2022-01-18T17:54:19Z"),
             "comments": {
                 "nodes": [
-                    {"updatedAt": "2022-01-19T17:54:19Z", "bodyText": "abc"},
+                    {
+                        **all_time_keys("2022-01-18T17:54:19Z"),
+                        "bodyText": "abc",
+                    },
                 ]
             },
         },
@@ -282,10 +395,10 @@ def run(pr, check):
             "isDraft": False,
             "author": {"login": "user"},
             "reviews": {"nodes": []},
-            "publishedAt": "2022-01-18T17:54:19Z",
+            **all_time_keys("2022-01-18T17:54:19Z"),
             "comments": {
                 "nodes": [
-                    {"updatedAt": "2022-01-27T17:54:19Z", "bodyText": "abc"},
+                    {**all_time_keys("2022-01-27T17:54:19Z"), "bodyText": "abc"},
                 ]
             },
         },
diff --git a/tests/python/unittest/test_link_params.py b/tests/python/unittest/test_link_params.py
index 0e56f6c00ae2..ea4f4ff975d7 100644
--- a/tests/python/unittest/test_link_params.py
+++ b/tests/python/unittest/test_link_params.py
@@ -20,7 +20,6 @@
 import os
 import re
 import sys
-import tempfile
 
 import numpy as np
 import pytest
@@ -34,20 +33,21 @@
 
 INPUT_SHAPE = (1, 3, 16, 16)
 
-
 KERNEL_SHAPE = (3, 3, 3, 3)
 
 
 # The data types that are linkable.
-LINKABLE_DTYPES = (
-    [f"uint{b}" for b in (8, 16, 32, 64)]
-    + [f"int{b}" for b in (8, 16, 32, 64)]
-    + ["float32", "float64"]
+linkable_dtype = tvm.testing.parameter(
+    *(
+        [f"uint{b}" for b in (8, 16, 32, 64)]
+        + [f"int{b}" for b in (8, 16, 32, 64)]
+        + ["float32", "float64"]
+    )
 )
 
 
 def dtype_info(dtype):
-    """Lookup numpy type info for the given string dtype (of LINKABLE_DTYPES above)."""
+    """Lookup numpy type info for the given string dtype (of linkable_dtype params above)."""
     if "int" in dtype:
         return np.iinfo(getattr(np, dtype))
     else:
@@ -181,58 +181,56 @@ def _add_decl(name, dtype):
 
 
 @tvm.testing.requires_llvm
-def test_llvm_link_params():
-    for dtype in LINKABLE_DTYPES:
-        ir_mod, param_init = _make_mod_and_params(dtype)
-        rand_input = _make_random_tensor(dtype, INPUT_SHAPE)
-        target = "llvm"
-        runtime = Runtime("crt", {"system-lib": True})
-        executor = Executor("graph", {"link-params": True})
-        with tvm.transform.PassContext(opt_level=3):
-            lib = tvm.relay.build(
-                ir_mod, target, runtime=runtime, executor=executor, params=param_init
-            )
-
-            # NOTE: Need to export_library() and load_library() to link all the Module(llvm, ...)
-            # against one another.
-            temp_dir = tempfile.mkdtemp()
-            export_file = os.path.join(temp_dir, "lib.so")
-            lib.lib.export_library(export_file)
-            mod = tvm.runtime.load_module(export_file)
-            assert set(lib.params.keys()) == {"p0", "p1"}  # NOTE: op folded
-            assert mod.get_function("TVMSystemLibEntryPoint") != None
-
-            graph = json.loads(lib.graph_json)
-            for p in lib.params:
-                _verify_linked_param(dtype, lib, mod, graph, p) or found_one
-
-            # Wrap in function to explicitly deallocate the runtime.
-            def _run_linked(lib, mod):
-                graph_json, _, _ = lib
-                graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
-                graph_rt.set_input("rand_input", rand_input)  # NOTE: params not required.
-                graph_rt.run()
-                return graph_rt.get_output(0)
-
-            linked_output = _run_linked(lib, mod)
-
-        runtime = Runtime("cpp", {"system-lib": True})
-        with tvm.transform.PassContext(opt_level=3):
-            lib = tvm.relay.build(ir_mod, "llvm", runtime=runtime, params=param_init)
-
-            def _run_unlinked(lib):
-                graph_json, mod, lowered_params = lib
-                graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
-                graph_rt.set_input("rand_input", rand_input, **lowered_params)
-                graph_rt.run()
-                return graph_rt.get_output(0)
-
-            unlinked_output = _run_unlinked(lib)
-
-        if "int" in dtype:
-            np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy())
-        else:
-            np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy())
+def test_llvm_link_params(linkable_dtype):
+    ir_mod, param_init = _make_mod_and_params(linkable_dtype)
+    rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE)
+    main_func = ir_mod["main"]
+    target = "llvm"
+    runtime = Runtime("crt", {"system-lib": True})
+    executor = Executor("graph", {"link-params": True})
+    with tvm.transform.PassContext(opt_level=3):
+        lib = tvm.relay.build(ir_mod, target, runtime=runtime, executor=executor, params=param_init)
+
+        # NOTE: Need to export_library() and load_library() to link all the Module(llvm, ...)
+        # against one another.
+        temp_dir = utils.TempDirectory()
+        export_file = temp_dir / "lib.so"
+        lib.lib.export_library(export_file)
+        mod = tvm.runtime.load_module(export_file)
+        assert len(lib.params.keys()) == 0  # NOTE: params became tir.constants
+        assert mod.get_function("TVMSystemLibEntryPoint") != None
+
+        graph = json.loads(lib.graph_json)
+        for p in lib.params:
+            _verify_linked_param(linkable_dtype, lib, mod, graph, p) or found_one
+
+        # Wrap in function to explicitly deallocate the runtime.
+        def _run_linked(lib, mod):
+            graph_json, _, _ = lib
+            graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
+            graph_rt.set_input("rand_input", rand_input)  # NOTE: params not required.
+            graph_rt.run()
+            return graph_rt.get_output(0)
+
+        linked_output = _run_linked(lib, mod)
+
+    runtime = Runtime("cpp", {"system-lib": True})
+    with tvm.transform.PassContext(opt_level=3):
+        lib = tvm.relay.build(ir_mod, "llvm", runtime=runtime, params=param_init)
+
+        def _run_unlinked(lib):
+            graph_json, mod, lowered_params = lib
+            graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
+            graph_rt.set_input("rand_input", rand_input, **lowered_params)
+            graph_rt.run()
+            return graph_rt.get_output(0)
+
+        unlinked_output = _run_unlinked(lib)
+
+    if "int" in linkable_dtype:
+        np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy())
+    else:
+        np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy())
 
 
 def _get_c_datatype(dtype):
@@ -247,161 +245,140 @@ def _get_c_datatype(dtype):
         assert False, f"unknown dtype {dtype}"
 
 
-def _format_c_value(dtype, width, x):
-    if "int" in dtype:
-        hex_formatstr = f'{{:{"+" if dtype.startswith("int") else ""}#0{width}x}}'
-        return hex_formatstr.format(x)
-    elif "float" in dtype:
-        to_ret = float(x).hex()
-        if "inf" in to_ret:
-            return ("-" if x < 0 else "") + "INFINITY"
-        elif "nan" in to_ret:
-            return "NAN"
-
-        before, after = to_ret.split("p")
-        return f'{before.rstrip("0")}p{after}'
-    else:
-        assert False, f"don't know dtype {dtype}"
-
-
 HEX_NUM_RE = re.compile(r"[+\-]?(?:(?:0x[0-9A-Fa-f.p+-]+)|(?:INFINITY)|(?:NAN))")
 
 
-def test_c_link_params():
+def test_c_link_params(linkable_dtype):
     temp_dir = utils.tempdir()
-    for dtype in LINKABLE_DTYPES:
-        mod, param_init = _make_mod_and_params(dtype)
-        rand_input = _make_random_tensor(dtype, INPUT_SHAPE)
-        target = "c"
-        executor = Executor("graph", {"link-params": True})
-        with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
-            lib = tvm.relay.build(mod, target, executor=executor, params=param_init)
-            assert set(lib.params.keys()) == {"p0", "p1"}  # NOTE: op folded
-
-            src = lib.lib.get_source()
-            lib.lib.save(temp_dir.relpath("test.c"), "c")
-            c_dtype = _get_c_datatype(dtype)
-            src_lines = src.split("\n")
-            param = lib.params["p0"].numpy().reshape(np.prod(KERNEL_SHAPE))
-            param_def = f'static const {c_dtype} __attribute__((section(".rodata.tvm"), aligned(16))) __tvm_param__p0[{np.prod(param.shape)}] = {{'
-            for i, line in enumerate(src_lines):
-                if line == param_def:
-                    i += 1
-                    break
-            else:
-                assert False, f'did not find parameter definition "{param_def}":\n{src}'
-
-            cursor = 0
-            width = dtype_info(dtype).bits // 4 + 2
-            if dtype.startswith("int"):
-                width += 1  # Account for sign
-
-            while "};" not in src_lines[i]:
-                for match in HEX_NUM_RE.finditer(src_lines[i]):
-                    assert match.group() == _format_c_value(dtype, width, param[cursor]), (
-                        f'p0 byte {cursor}: want "{_format_c_value(dtype, width, param[cursor])}" got '
-                        f'"{match.group(0)}"; full p0 follows:\n{src}'
-                    )
-                    cursor += 1
+    mod, param_init = _make_mod_and_params(linkable_dtype)
+    rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE)
+    main_func = mod["main"]
+    target = "c"
+    executor = Executor("graph", {"link-params": True})
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        lib = tvm.relay.build(mod, target, executor=executor, params=param_init)
+        assert len(lib.params.keys()) == 0  # NOTE: params became tir.constants
+
+        src = lib.lib.get_source()
+        lib.lib.save(temp_dir.relpath("test.c"), "c")
+        c_dtype = _get_c_datatype(linkable_dtype)
+        src_lines = src.split("\n")
+        param = param_init[f"{linkable_dtype}_a"].reshape(np.prod(KERNEL_SHAPE))
+        param_def = rf"^static const {c_dtype} __attribute__\(\(section\(\".rodata.tvm\"\), aligned\(16\)\)\) constant_\d+\[{np.prod(param.shape)}\] = {{$"
+        for i, line in enumerate(src_lines):
+            if re.match(param_def, line):
                 i += 1
+                break
+        else:
+            assert False, f'did not find parameter definition "{param_def}":\n{src}'
+
+        cursor = 0
+        width = dtype_info(linkable_dtype).bits // 4 + 2
+        if linkable_dtype.startswith("int"):
+            width += 1  # Account for sign
+
+        while "};" not in src_lines[i]:
+            for match in HEX_NUM_RE.finditer(src_lines[i]):
+                cursor += 1
+            i += 1
+
+        assert cursor == np.prod(param.shape)
+
+        # Need a unique name per library to avoid dlopen caching the lib load.
+        lib_path = temp_dir.relpath(f"test-{linkable_dtype}-linked.so")
+        lib["remove_params"]().export_library(lib_path)
+        lib_mod = tvm.runtime.load_module(lib_path)
+
+        #            lib_mod = lib_factory['default']()
+        graph = json.loads(lib.graph_json)
+        for p in lib.params:
+            _verify_linked_param(linkable_dtype, lib, lib_mod, graph, p)
+
+        # Wrap in function to explicitly deallocate the runtime.
+        def _run_linked(lib_mod):
+            graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0)))
+            graph_rt.set_input("rand_input", rand_input)  # NOTE: params not required.
+            graph_rt.run()
+
+            return graph_rt.get_output(0)
+
+        linked_output = _run_linked(lib_mod)
+
+    linked_params = lib.params
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        lib = tvm.relay.build(mod, "c", params=param_init)
+        _, _, params = lib
+        # Need a unique name per library to avoid dlopen caching the lib load.
+        lib_path = temp_dir.relpath(f"test-{linkable_dtype}-unlinked.so")
+        lib.export_library(lib_path)
+        lib_mod = tvm.runtime.load_module(lib_path)
+
+        def _run_unlinked(lib_mod):
+            graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0)))
+            graph_rt.set_input("rand_input", rand_input, **params)
+            graph_rt.run()
+            return graph_rt.get_output(0)
+
+        unlinked_output = _run_unlinked(lib_mod)
+
+    if "int" in linkable_dtype:
+        np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy())
+    else:
+        np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy())
 
-            assert cursor == np.prod(param.shape)
-
-            # Need a unique name per library to avoid dlopen caching the lib load.
-            lib_path = temp_dir.relpath(f"test-{dtype}-linked.so")
-            lib["remove_params"]().export_library(lib_path)
-            lib_mod = tvm.runtime.load_module(lib_path)
-
-            #            lib_mod = lib_factory['default']()
-            graph = json.loads(lib.graph_json)
-            for p in lib.params:
-                _verify_linked_param(dtype, lib, lib_mod, graph, p)
-
-            # Wrap in function to explicitly deallocate the runtime.
-            def _run_linked(lib_mod):
-                graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0)))
-                graph_rt.set_input("rand_input", rand_input)  # NOTE: params not required.
-                graph_rt.run()
-
-                return graph_rt.get_output(0)
-
-            linked_output = _run_linked(lib_mod)
 
-        linked_params = lib.params
-        with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
-            lib = tvm.relay.build(mod, "c", params=param_init)
-            _, _, params = lib
-            # Need a unique name per library to avoid dlopen caching the lib load.
-            lib_path = temp_dir.relpath(f"test-{dtype}-unlinked.so")
-            lib.export_library(lib_path)
-            lib_mod = tvm.runtime.load_module(lib_path)
+@tvm.testing.requires_micro
+def test_crt_link_params(linkable_dtype):
+    from tvm import micro
 
-            def _run_unlinked(lib_mod):
-                graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0)))
-                graph_rt.set_input("rand_input", rand_input, **params)
-                graph_rt.run()
-                return graph_rt.get_output(0)
+    mod, param_init = _make_mod_and_params(linkable_dtype)
+    rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE)
+    main_func = mod["main"]
+    target = "c"
+    runtime = Runtime("crt", {"system-lib": True})
+    executor = Executor("graph", {"link-params": True})
+    with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
+        factory = tvm.relay.build(
+            mod, target, runtime=runtime, executor=executor, params=param_init
+        )
+        assert len(factory.get_params().keys()) == 0  # NOTE: params became tir.constants
+
+        temp_dir = tvm.contrib.utils.tempdir()
+        template_project_dir = os.path.join(tvm.micro.get_standalone_crt_dir(), "template", "host")
+        project = tvm.micro.generate_project(
+            template_project_dir, factory, temp_dir / "project", {"verbose": 1}
+        )
+        project.build()
+        project.flash()
+        with tvm.micro.Session(project.transport()) as sess:
+            graph_rt = tvm.micro.session.create_local_graph_executor(
+                factory.get_graph_json(), sess.get_system_lib(), sess.device
+            )
 
-            unlinked_output = _run_unlinked(lib_mod)
+            assert len(factory.params.keys()) == 0  # NOTE: params became tir.constants
 
-        if "int" in dtype:
-            np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy())
-        else:
-            np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy())
+            # NOTE: not setting params here.
+            graph_rt.set_input("rand_input", rand_input)
+            graph_rt.run()
+            linked_output = graph_rt.get_output(0).numpy()
 
+    runtime = Runtime("cpp", {"system-lib": True})
+    with tvm.transform.PassContext(opt_level=3):
+        lib = tvm.relay.build(mod, "llvm", runtime=runtime, params=param_init)
 
-@tvm.testing.requires_micro
-def test_crt_link_params():
-    from tvm import micro
+        def _run_unlinked(lib):
+            graph_json, mod, lowered_params = lib
+            graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
+            graph_rt.set_input("rand_input", rand_input, **lowered_params)
+            graph_rt.run()
+            return graph_rt.get_output(0).numpy()
 
-    for dtype in LINKABLE_DTYPES:
-        mod, param_init = _make_mod_and_params(dtype)
-        rand_input = _make_random_tensor(dtype, INPUT_SHAPE)
-        target = "c"
-        runtime = Runtime("crt", {"system-lib": True})
-        executor = Executor("graph", {"link-params": True})
-        with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
-            factory = tvm.relay.build(
-                mod, target, runtime=runtime, executor=executor, params=param_init
-            )
-            assert set(factory.get_params().keys()) == {"p0", "p1"}  # NOTE: op folded
+        unlinked_output = _run_unlinked(lib)
 
-            temp_dir = tvm.contrib.utils.tempdir()
-            template_project_dir = os.path.join(
-                tvm.micro.get_standalone_crt_dir(), "template", "host"
-            )
-            project = tvm.micro.generate_project(
-                template_project_dir, factory, temp_dir / "project", {"verbose": 1}
-            )
-            project.build()
-            project.flash()
-            with tvm.micro.Session(project.transport()) as sess:
-                graph_rt = tvm.micro.session.create_local_graph_executor(
-                    factory.get_graph_json(), sess.get_system_lib(), sess.device
-                )
-
-                # NOTE: not setting params here.
-                graph_rt.set_input("rand_input", rand_input)
-                graph_rt.run()
-                linked_output = graph_rt.get_output(0).numpy()
-
-        runtime = Runtime("cpp", {"system-lib": True})
-        with tvm.transform.PassContext(opt_level=3):
-            lib = tvm.relay.build(mod, "llvm", runtime=runtime, params=param_init)
-
-            def _run_unlinked(lib):
-                graph_json, mod, lowered_params = lib
-                graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0))
-                graph_rt.set_input("rand_input", rand_input, **lowered_params)
-                graph_rt.run()
-                return graph_rt.get_output(0).numpy()
-
-            unlinked_output = _run_unlinked(lib)
-
-        if "int" in dtype:
-            np.testing.assert_equal(unlinked_output, linked_output)
-        else:
-            np.testing.assert_allclose(unlinked_output, linked_output)
+    if "int" in linkable_dtype:
+        np.testing.assert_equal(unlinked_output, linked_output)
+    else:
+        np.testing.assert_allclose(unlinked_output, linked_output)
 
 
 if __name__ == "__main__":
diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_cross_thread_reduction.py b/tests/python/unittest/test_meta_schedule_schedule_rule_cross_thread_reduction.py
index 7bed18b0f9ea..47f405842c98 100644
--- a/tests/python/unittest/test_meta_schedule_schedule_rule_cross_thread_reduction.py
+++ b/tests/python/unittest/test_meta_schedule_schedule_rule_cross_thread_reduction.py
@@ -84,7 +84,7 @@ def test_gpu_softmax_mn():
             "v4 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l5, l6 = sch.split(loop=l3, factors=[None, v4])",
             'sch.bind(loop=l6, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l7, l8, l9 = sch.get_loops(block=b0)",
             "l10, l11 = sch.split(loop=l9, factors=[None, v4])",
@@ -97,7 +97,7 @@ def test_gpu_softmax_mn():
             "v4 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l5, l6 = sch.split(loop=l3, factors=[None, v4])",
             'sch.bind(loop=l6, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l7, l8, l9 = sch.get_loops(block=b0)",
             "l10, l11 = sch.split(loop=l9, factors=[None, v4])",
@@ -111,7 +111,7 @@ def test_gpu_softmax_mn():
             "v5 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l6, l7 = sch.split(loop=l4, factors=[None, v5])",
             'sch.bind(loop=l7, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b1, loop=l3, preserve_unit_loops=1)",
+            "sch.compute_at(block=b1, loop=l3, preserve_unit_loops=True)",
             'sch.set_scope(block=b1, buffer_index=0, storage_scope="shared")',
             "l8, l9, l10 = sch.get_loops(block=b1)",
             "l11, l12 = sch.split(loop=l10, factors=[None, v5])",
@@ -121,7 +121,7 @@ def test_gpu_softmax_mn():
             "v16 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l17, l18 = sch.split(loop=l15, factors=[None, v16])",
             'sch.bind(loop=l18, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b0, loop=l14, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l14, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l19, l20, l21 = sch.get_loops(block=b0)",
             "l22, l23 = sch.split(loop=l21, factors=[None, v16])",
@@ -161,7 +161,7 @@ def test_gpu_softmax_mn_after_inline():
             "v4 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l5, l6 = sch.split(loop=l3, factors=[None, v4])",
             'sch.bind(loop=l6, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l2, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l7, l8, l9 = sch.get_loops(block=b0)",
             "l10, l11 = sch.split(loop=l9, factors=[None, v4])",
@@ -175,14 +175,14 @@ def test_gpu_softmax_mn_after_inline():
             "v5 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l6, l7 = sch.split(loop=l4, factors=[None, v5])",
             'sch.bind(loop=l7, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b1, loop=l3, preserve_unit_loops=1)",
+            "sch.compute_at(block=b1, loop=l3, preserve_unit_loops=True)",
             'sch.set_scope(block=b1, buffer_index=0, storage_scope="shared")',
             "l8, l9, l10 = sch.get_loops(block=b1)",
             "l11, l12 = sch.split(loop=l10, factors=[None, v5])",
             'sch.bind(loop=l12, thread_axis="threadIdx.x")',
             "b13, b14 = sch.get_consumers(block=b0)",
             "l15, l16, l17, l18 = sch.get_loops(block=b13)",
-            "sch.compute_at(block=b0, loop=l15, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l15, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l19, l20, l21 = sch.get_loops(block=b0)",
             "l22, l23 = sch.split(loop=l21, factors=[None, v5])",
@@ -210,7 +210,7 @@ def test_gpu_batch_norm_bmn():
             "v3 = sch.sample_categorical(candidates=[4, 8, 16, 32, 64, 128, 256, 512], probs=[0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125, 0.125])",
             "l4, l5 = sch.split(loop=l2, factors=[None, v3])",
             'sch.bind(loop=l5, thread_axis="threadIdx.x")',
-            "sch.compute_at(block=b0, loop=l4, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l4, preserve_unit_loops=True)",
             'sch.set_scope(block=b0, buffer_index=0, storage_scope="shared")',
             "l6, l7, l8, l9 = sch.get_loops(block=b0)",
             "l10 = sch.fuse(l8, l9)",
diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py b/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py
index c6a63aae7427..8882ed625bf1 100644
--- a/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py
+++ b/tests/python/unittest/test_meta_schedule_schedule_rule_multi_level_tiling.py
@@ -55,7 +55,7 @@ def test_cpu_matmul():
             "l22, l23 = sch.split(loop=l3, factors=[v20, v21])",
             "sch.reorder(l8, l16, l9, l17, l22, l10, l18, l23, l11, l19)",
             'b24 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")',
-            "sch.reverse_compute_at(block=b24, loop=l17, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b24, loop=l17, preserve_unit_loops=True)",
         ],
         [
             'b0 = sch.get_block(name="C", func_name="main")',
@@ -69,7 +69,7 @@ def test_cpu_matmul():
             "l22, l23 = sch.split(loop=l3, factors=[v20, v21])",
             "sch.reorder(l8, l16, l9, l17, l22, l10, l18, l23, l11, l19)",
             'b24 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")',
-            "sch.reverse_compute_at(block=b24, loop=l16, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b24, loop=l16, preserve_unit_loops=True)",
         ],
         [
             'b0 = sch.get_block(name="C", func_name="main")',
@@ -116,7 +116,7 @@ def test_cpu_matmul_relu():
             "l22, l23 = sch.split(loop=l3, factors=[v20, v21])",
             "sch.reorder(l8, l16, l9, l17, l22, l10, l18, l23, l11, l19)",
             "b24, = sch.get_consumers(block=b0)",
-            "sch.reverse_compute_at(block=b24, loop=l17, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b24, loop=l17, preserve_unit_loops=True)",
         ],
         [
             'b0 = sch.get_block(name="C", func_name="main")',
@@ -130,7 +130,7 @@ def test_cpu_matmul_relu():
             "l22, l23 = sch.split(loop=l3, factors=[v20, v21])",
             "sch.reorder(l8, l16, l9, l17, l22, l10, l18, l23, l11, l19)",
             "b24, = sch.get_consumers(block=b0)",
-            "sch.reverse_compute_at(block=b24, loop=l16, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b24, loop=l16, preserve_unit_loops=True)",
         ],
         [
             'b0 = sch.get_block(name="C", func_name="main")',
@@ -186,15 +186,15 @@ def test_cuda_matmul():
             'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32)',
             'sch.annotate(block_or_loop=b0, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)',
             'b33 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")',
-            "sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=True)",
             'b34 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
-            "sch.compute_at(block=b34, loop=l27, preserve_unit_loops=1)",
+            "sch.compute_at(block=b34, loop=l27, preserve_unit_loops=True)",
             "l35, l36, l37, l38, l39, l40 = sch.get_loops(block=b34)",
             "l41 = sch.fuse(l39, l40)",
             "v42 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
             'sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v42)',
             'b43 = sch.cache_read(block=b0, read_buffer_index=2, storage_scope="shared")',
-            "sch.compute_at(block=b43, loop=l27, preserve_unit_loops=1)",
+            "sch.compute_at(block=b43, loop=l27, preserve_unit_loops=True)",
             "l44, l45, l46, l47, l48, l49 = sch.get_loops(block=b43)",
             "l50 = sch.fuse(l48, l49)",
             "v51 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
@@ -240,15 +240,15 @@ def test_cuda_matmul_relu():
             "l32 = sch.fuse(l11, l21)",
             'sch.bind(loop=l32, thread_axis="threadIdx.x")',
             'b33 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")',
-            "sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=1)",
+            "sch.reverse_compute_at(block=b33, loop=l32, preserve_unit_loops=True)",
             'b34 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")',
-            "sch.compute_at(block=b34, loop=l27, preserve_unit_loops=1)",
+            "sch.compute_at(block=b34, loop=l27, preserve_unit_loops=True)",
             "l35, l36, l37, l38, l39, l40 = sch.get_loops(block=b34)",
             "l41 = sch.fuse(l39, l40)",
             "v42 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
             'sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v42)',
             'b43 = sch.cache_read(block=b0, read_buffer_index=2, storage_scope="shared")',
-            "sch.compute_at(block=b43, loop=l27, preserve_unit_loops=1)",
+            "sch.compute_at(block=b43, loop=l27, preserve_unit_loops=True)",
             "l44, l45, l46, l47, l48, l49 = sch.get_loops(block=b43)",
             "l50 = sch.fuse(l48, l49)",
             "v51 = sch.sample_categorical(candidates=[1, 2, 3, 4], probs=[0.25, 0.25, 0.25, 0.25])",
diff --git a/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py b/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py
index 92c7da922c39..18db006c6ca8 100644
--- a/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py
+++ b/tests/python/unittest/test_meta_schedule_schedule_rule_random_compute_location.py
@@ -74,7 +74,7 @@ def test_random_compute_location():
         [
             'b0 = sch.get_block(name="move", func_name="main")',
             "l1 = sch.sample_compute_location(block=b0)",
-            "sch.compute_at(block=b0, loop=l1, preserve_unit_loops=1)",
+            "sch.compute_at(block=b0, loop=l1, preserve_unit_loops=True)",
         ]
     ]
     mod = Add
diff --git a/tests/python/unittest/test_meta_schedule_tune_relay.py b/tests/python/unittest/test_meta_schedule_tune_relay.py
index 7e6f89dfb149..56b6feae067b 100644
--- a/tests/python/unittest/test_meta_schedule_tune_relay.py
+++ b/tests/python/unittest/test_meta_schedule_tune_relay.py
@@ -21,6 +21,11 @@
 import numpy as np
 from typing import Tuple, List
 
+try:
+    import torch
+except ModuleNotFoundError:
+    pass
+
 import tvm
 from tvm import relay
 from tvm.ir import IRModule
diff --git a/tests/python/unittest/test_micro_model_library_format.py b/tests/python/unittest/test_micro_model_library_format.py
index 48437eaf58d7..76e95c960482 100644
--- a/tests/python/unittest/test_micro_model_library_format.py
+++ b/tests/python/unittest/test_micro_model_library_format.py
@@ -108,14 +108,21 @@ def validate_graph_json(extract_dir, factory):
 
 @tvm.testing.requires_micro
 @pytest.mark.parametrize(
-    "executor,runtime,should_generate_interface",
+    "executor,runtime,should_generate_interface,json_constants_size_bytes",
     [
-        (Executor("graph"), Runtime("crt", {"system-lib": True}), False),
-        (Executor("aot"), Runtime("crt"), False),
-        (Executor("aot", {"unpacked-api": True, "interface-api": "c"}), Runtime("crt"), True),
+        (Executor("graph"), Runtime("crt", {"system-lib": True}), False, 8),
+        (Executor("aot", {"link-params": True}), Runtime("crt"), False, 0),
+        (
+            Executor("aot", {"unpacked-api": True, "interface-api": "c"}),
+            Runtime("crt"),
+            True,
+            0,
+        ),
     ],
 )
-def test_export_model_library_format_c(executor, runtime, should_generate_interface):
+def test_export_model_library_format_c(
+    executor, runtime, should_generate_interface, json_constants_size_bytes
+):
     target = tvm.target.target.micro("host")
     with utils.TempDirectory.set_keep_for_debug(True):
         with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}):
@@ -165,7 +172,7 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[
                 ]
             assert metadata["memory"]["functions"]["main"] == [
                 {
-                    "constants_size_bytes": 8,
+                    "constants_size_bytes": json_constants_size_bytes,
                     "device": 1,
                     "io_size_bytes": 18,
                     "workspace_size_bytes": 0,
@@ -193,7 +200,10 @@ def @main(%a : Tensor[(1, 2), uint8], %b : Tensor[(1, 2), float32], %c : Tensor[
 
         with open(os.path.join(extract_dir, "parameters", "add.params"), "rb") as params_f:
             params = tvm.relay.load_param_dict(params_f.read())
-            assert "p0" in params
+            if json_constants_size_bytes != 0:
+                assert "p0" in params
+            else:
+                assert len(params) == 0
 
 
 @tvm.testing.requires_micro
diff --git a/tests/python/unittest/test_subwarp_reduction_cuda.py b/tests/python/unittest/test_subwarp_reduction_cuda.py
new file mode 100644
index 000000000000..8778c75f5699
--- /dev/null
+++ b/tests/python/unittest/test_subwarp_reduction_cuda.py
@@ -0,0 +1,68 @@
+# 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.
+import tvm
+import tvm.testing
+import numpy as np
+from tvm.script import tir as T
+
+
+@T.prim_func
+def reduce(a: T.handle, b: T.handle, d1: T.int32, d2: T.int32, d3: T.int32) -> None:
+    A = T.match_buffer(a, [1, d1, d2, d3])
+    B = T.match_buffer(b, [1, d1, d2])
+
+    for i, j, k, l in T.grid(1, d1, d2, d3):
+        with T.block("reduce"):
+            vi, vj, vk, vl = T.axis.remap("SSSR", [i, j, k, l])
+            with T.init():
+                B[vi, vj, vk] = 0.0
+            B[vi, vj, vk] = B[vi, vj, vk] + A[vi, vj, vk, vl]
+
+
+@tvm.testing.requires_gpu
+@tvm.testing.requires_cuda
+def test_cuda_subwarp_reduction():
+    def check(d1: int, d2: int, d3: int):
+        _, _, _d1, _d2, _d3 = reduce.params
+        mod = reduce.specialize({_d1: d1, _d2: d2, _d3: d3})
+        sch = tvm.tir.Schedule(mod)
+        blk = sch.get_block("reduce")
+        i, j, k, l = sch.get_loops(blk)
+        sch.bind(i, "blockIdx.x")
+        sch.bind(j, "threadIdx.z")
+        sch.bind(k, "threadIdx.y")
+        sch.bind(l, "threadIdx.x")
+        f = tvm.build(sch.mod["main"], target="cuda")
+
+        # prepare input and output array
+        a_np = np.random.rand(1, d1, d2, d3).astype("float32")
+        b_np = a_np.sum(axis=-1).astype("float32")
+        a = tvm.nd.array(a_np, tvm.cuda(0))
+        b = tvm.nd.array(np.zeros_like(b_np), tvm.cuda(0))
+
+        # launch kernel
+        f(a, b)
+        tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-6, atol=1e-6)
+
+    for d1 in range(1, 5):
+        for d2 in range(1, 5):
+            for d3 in range(2, 33):
+                check(d1, d2, d3)
+
+
+if __name__ == "__main__":
+    test_cuda_subwarp_reduction()
diff --git a/tests/python/unittest/test_te_schedule_ops.py b/tests/python/unittest/test_te_schedule_ops.py
index ca3ab3aade98..f85cdc619687 100644
--- a/tests/python/unittest/test_te_schedule_ops.py
+++ b/tests/python/unittest/test_te_schedule_ops.py
@@ -21,6 +21,12 @@
 from tvm.driver.build_module import schedule_to_module
 
 
+def test_const():
+    x = tvm.te.const(1, "int32")
+    assert x.dtype == "int32"
+    assert isinstance(x, tvm.tir.IntImm)
+
+
 def test_schedule0():
     m = te.var("m")
     l = te.var("l")
diff --git a/tests/python/unittest/test_te_tensor.py b/tests/python/unittest/test_te_tensor.py
index 6958888e9bb6..6b09410af6c4 100644
--- a/tests/python/unittest/test_te_tensor.py
+++ b/tests/python/unittest/test_te_tensor.py
@@ -37,6 +37,8 @@ def test_tensor():
     assert T.op.output(0).__hash__() == T.__hash__()
     d = {T.op.output(0): 1}
     assert d[T] == 1
+    assert T == T.op.output(0)
+    assert T.same_as(T.op.output(0))
     assert T[0][0][0].astype("float16").dtype == "float16"
 
 
@@ -49,6 +51,8 @@ def test_rank_zero():
     print(T)
     print(T.op.body)
     assert tuple(T.shape) == ()
+    assert T == T.op.output(0)
+    assert T.same_as(T.op.output(0))
 
 
 def test_conv1d():
diff --git a/tests/python/unittest/test_tir_nodes.py b/tests/python/unittest/test_tir_nodes.py
index fe719ee99693..96224ef6fe55 100644
--- a/tests/python/unittest/test_tir_nodes.py
+++ b/tests/python/unittest/test_tir_nodes.py
@@ -22,7 +22,12 @@
 
 def test_const():
     x = tvm.tir.const(1, "int32")
-    print(x.dtype)
+    assert x.dtype == "int32"
+    assert isinstance(x, tvm.tir.IntImm)
+
+
+def test_te_const():
+    x = tvm.te.const(1, "int32")
     assert x.dtype == "int32"
     assert isinstance(x, tvm.tir.IntImm)
 
diff --git a/tests/python/unittest/test_tir_ptx_mma.py b/tests/python/unittest/test_tir_ptx_mma.py
index 4b8e3fcaffef..c304e818ef05 100644
--- a/tests/python/unittest/test_tir_ptx_mma.py
+++ b/tests/python/unittest/test_tir_ptx_mma.py
@@ -310,6 +310,10 @@ def gemm_mma_m8n8k16_row_col_s8s8s32(a: T.handle, b: T.handle, c: T.handle):
         )
 
 
+# This test uses mma instructions that are not available on NVCC 10.1.
+# Failure occurs during the external call to nvcc, when attempting to
+# generate the .fatbin file.
+@tvm.testing.requires_nvcc_version(11)
 @tvm.testing.requires_cuda
 def test_gemm_mma_m8n8k16_row_col_s8s8s32():
     sch = tvm.tir.Schedule(gemm_mma_m8n8k16_row_col_s8s8s32)
@@ -384,6 +388,10 @@ def gemm_mma_m8n8k16_row_col_s8u8s32(a: T.handle, b: T.handle, c: T.handle):
         )
 
 
+# This test uses mma instructions that are not available on NVCC 10.1.
+# Failure occurs during the external call to nvcc, when attempting to
+# generate the .fatbin file.
+@tvm.testing.requires_nvcc_version(11)
 @tvm.testing.requires_cuda
 def test_gemm_mma_m8n8k16_row_col_s8u8s32():
     sch = tvm.tir.Schedule(gemm_mma_m8n8k16_row_col_s8u8s32)
@@ -458,6 +466,10 @@ def gemm_mma_m8n8k32_row_col_s4s4s32(a: T.handle, b: T.handle, c: T.handle):
         )
 
 
+# This test uses mma instructions that are not available on NVCC 10.1.
+# Failure occurs during the external call to nvcc, when attempting to
+# generate the .fatbin file.
+@tvm.testing.requires_nvcc_version(11)
 @tvm.testing.requires_cuda
 def test_gemm_mma_m8n8k32_row_col_s4s4s32():
     sch = tvm.tir.Schedule(gemm_mma_m8n8k32_row_col_s4s4s32)
@@ -524,6 +536,10 @@ def gemm_mma_m8n8k32_row_col_s4u4s32(a: T.handle, b: T.handle, c: T.handle):
         )
 
 
+# This test uses mma instructions that are not available on NVCC 10.1.
+# Failure occurs during the external call to nvcc, when attempting to
+# generate the .fatbin file.
+@tvm.testing.requires_nvcc_version(11)
 @tvm.testing.requires_cuda
 def test_gemm_mma_m8n8k32_row_col_s4u4s32():
     sch = tvm.tir.Schedule(gemm_mma_m8n8k32_row_col_s4u4s32)
diff --git a/tests/python/unittest/test_tir_schedule_compute_at.py b/tests/python/unittest/test_tir_schedule_compute_at.py
index dc0812780503..e1cf399d49a1 100644
--- a/tests/python/unittest/test_tir_schedule_compute_at.py
+++ b/tests/python/unittest/test_tir_schedule_compute_at.py
@@ -912,7 +912,7 @@ def concat_two_elemwise(x: T.Buffer[(16,), "float32"],
     for i in T.serial(24):
         with T.block("T_concat"):
             ax = T.axis.spatial(24, i)
-            T_concat[ax] = T.if_then_else(16 <= ax, T_add_1[ax - 16], T_add_2[ax], dtype="float32")
+            T_concat[ax] = T.if_then_else(16 <= ax, T_add_2[ax - 16], T_add_1[ax], dtype="float32")
 
 @T.prim_func
 def concat_two_elemwise_after_compute_at(x: T.Buffer[(16,), "float32"],
@@ -922,16 +922,16 @@ def concat_two_elemwise_after_compute_at(x: T.Buffer[(16,), "float32"],
     T_add_2 = T.alloc_buffer([8], dtype="float32")
     for i in T.serial(24):
         with T.block("T_add_1"):
-            ax = T.axis.spatial(16, i - 16)
-            T.where(16 <= i)
+            ax = T.axis.spatial(16, i)
+            T.where(i < 16)
             T_add_1[ax] = x[ax] + T.float32(1)
         with T.block("T_add_2"):
-            ax = T.axis.spatial(8, i)
-            T.where(i < 8)
+            ax = T.axis.spatial(8, i - 16)
+            T.where(16 <= i)
             T_add_2[ax] = y[ax] + T.float32(2)
         with T.block("T_concat"):
             ax = T.axis.spatial(24, i)
-            T_concat[ax] = T.if_then_else(16 <= ax, T_add_1[ax - 16], T_add_2[ax], dtype="float32")
+            T_concat[ax] = T.if_then_else(16 <= ax, T_add_2[ax - 16], T_add_1[ax], dtype="float32")
 
 @T.prim_func
 def floordiv_and_floormod_indices(a: T.handle, b: T.handle) -> None:
@@ -962,6 +962,61 @@ def floordiv_and_floormod_indices_after_reverse_compute_at(a: T.handle, b: T.han
                 v_i = T.axis.spatial(256, i * 16 + ax0)
                 Y[v_i] = temp[v_i // 16, v_i % 16]
 
+
+@T.prim_func
+def tiled_repeat_op(x: T.Buffer[(4,), "float32"], T_repeat: T.Buffer[(64,), "float32"]) -> None:
+    T_add = T.alloc_buffer([4], dtype="float32")
+    for i0 in T.serial(4):
+        with T.block("T_add"):
+            ax0 = T.axis.spatial(4, i0)
+            T_add[ax0] = x[ax0] + 1.0
+    for i0_0, i0_1 in T.grid(8, 8):
+        with T.block("T_repeat"):
+            ax0 = T.axis.spatial(64, i0_0 * 8 + i0_1)
+            T_repeat[ax0] = T_add[ax0 // 16]
+
+@T.prim_func
+def tiled_repeat_op_after_compute_at(x: T.Buffer[(4,), "float32"], T_repeat: T.Buffer[(64,), "float32"]) -> None:
+    T_add = T.alloc_buffer([4], dtype="float32")
+    for i0_0 in T.serial(8):
+        with T.block("T_add"):
+            ax0 = T.axis.spatial(4, i0_0 // 2)
+            T_add[ax0] = x[ax0] + T.float32(1)
+        for i0_1 in T.serial(8):
+            with T.block("T_repeat"):
+                ax0 = T.axis.spatial(64, i0_0 * 8 + i0_1)
+                T_repeat[ax0] = T_add[ax0 // 16]
+
+@T.prim_func
+def static_bound(A: T.Buffer[(32, 1), "float32"], C: T.Buffer[(32, 1), "float32"]) -> None:
+    B = T.alloc_buffer((32, 1), "float32")
+    for i, j in T.grid(32, 1):
+        with T.block("B"):
+            vi = T.axis.spatial(32, i)
+            vj = T.axis.spatial(1, j)
+            B[vi, vj] = A[vi, vj] * 2.0
+    for i, j in T.grid(32, 32):
+        with T.block("C"):
+            vi = T.axis.spatial(32, i)
+            vj = T.axis.spatial(1, j)
+            T.where(j < 1)
+            C[vi, vj] = B[vi, vj] + 1.0
+
+@T.prim_func
+def static_bound_after_compute_at(A: T.Buffer[(32, 1), "float32"], C: T.Buffer[(32, 1), "float32"]) -> None:
+    B = T.alloc_buffer((32, 1), "float32")
+    for i in range(32):
+        for ax0, ax1 in T.grid(1, 1):
+            with T.block("B"):
+                vi = T.axis.spatial(32, i + ax0)
+                vj = T.axis.spatial(1, ax1)
+                B[vi, vj] = A[vi, vj] * 2.0
+        for j in range(32):
+            with T.block("C"):
+                vi = T.axis.spatial(32, i)
+                vj = T.axis.spatial(1, j)
+                T.where(j < 1)
+                C[vi, vj] = B[vi, vj] + 1.0
 # pylint: enable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
 # fmt: on
 
@@ -1077,6 +1132,14 @@ def test_compute_at_concat():
     verify_trace_roundtrip(sch=sch, mod=concat_two_elemwise)
 
 
+def test_compute_at_tiled_repeat_op():
+    sch = tir.Schedule(tiled_repeat_op, debug_mask="all")
+    outer_ax, _ = sch.get_loops(sch.get_block("T_repeat"))
+    sch.compute_at(sch.get_block("T_add"), outer_ax)
+    tvm.ir.assert_structural_equal(tiled_repeat_op_after_compute_at, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=tiled_repeat_op)
+
+
 def test_reverse_compute_at_tiled():
     sch = tir.Schedule(tiled, debug_mask="all")
     block = sch.get_block("C")
@@ -1133,6 +1196,15 @@ def test_compact_dataflow():
     verify_trace_roundtrip(sch=sch, mod=not_all_compact_data_flow)
 
 
+def test_compute_at_simplify_static_bound():
+    sch = tir.Schedule(static_bound, debug_mask="all")
+    block = sch.get_block("B")
+    loop, _ = sch.get_loops(sch.get_block("C"))
+    sch.compute_at(block, loop, preserve_unit_loops=True)
+    tvm.ir.assert_structural_equal(static_bound_after_compute_at, sch.mod["main"])
+    verify_trace_roundtrip(sch=sch, mod=static_bound)
+
+
 def test_fail_subtree_complete_block():
     sch = tir.Schedule(fail_subtree_compact_dataflow, debug_mask="all")
     block = sch.get_block("B_0")
diff --git a/tests/python/unittest/test_tir_schedule_trace.py b/tests/python/unittest/test_tir_schedule_trace.py
index f1c97c57b2ff..1923eb23af5b 100644
--- a/tests/python/unittest/test_tir_schedule_trace.py
+++ b/tests/python/unittest/test_tir_schedule_trace.py
@@ -82,6 +82,15 @@ def _make_compute_inline(input):  # pylint: disable=redefined-builtin
     )
 
 
+def _make_split(inputs, outputs):  # pylint: disable=redefined-builtin
+    return Instruction(
+        kind=InstructionKind.get("Split"),
+        inputs=inputs,
+        attrs=[],
+        outputs=outputs,
+    )
+
+
 def _make_enter_postproc():
     return Instruction(
         kind=InstructionKind.get("EnterPostproc"),
@@ -129,6 +138,17 @@ def _make_trace_3(b0, b1, add_postproc):  # pylint: disable=invalid-name
     return Trace(insts=insts, decisions={})
 
 
+def _make_trace_4(b0, l1, l2, l3):  # pylint: disable=invalid-name
+    return Trace(
+        insts=[
+            _make_get_block(name="B", output=b0),
+            _make_get_loops(input=b0, outputs=[l1]),
+            _make_split([l1, None, 32], [l2, l3]),
+        ],
+        decisions={},
+    )
+
+
 def test_trace_construct_1():
     trace = _make_trace_1(BlockRV(), LoopRV(), LoopRV())
     assert str(trace) == "\n".join(
@@ -235,6 +255,17 @@ def test_trace_simplified_2():
     )
 
 
+def test_trace_simplified_3():
+    trace = _make_trace_4(BlockRV(), LoopRV(), LoopRV(), LoopRV()).simplified(remove_postproc=False)
+    assert str(trace) == "\n".join(
+        (
+            'b0 = sch.get_block(name="B", func_name="main")',
+            "l1, = sch.get_loops(block=b0)",
+            "l2, l3 = sch.split(loop=l1, factors=[None, 32])",
+        )
+    )
+
+
 def test_apply_json_to_schedule_1():
     trace = _make_trace_2(BlockRV())
     json_obj = trace.as_json()
diff --git a/tests/python/unittest/test_tir_transform_compact_buffer_region.py b/tests/python/unittest/test_tir_transform_compact_buffer_region.py
index 9b844853f243..6c42ed7d9280 100644
--- a/tests/python/unittest/test_tir_transform_compact_buffer_region.py
+++ b/tests/python/unittest/test_tir_transform_compact_buffer_region.py
@@ -505,6 +505,72 @@ def compacted_opaque_access_annotated_func(a: T.handle) -> None:
                 T.store(C.data, i, T.load("float32", B.data, i))
 
 
+@T.prim_func
+def sparse_read_cache(
+    A_data: T.Buffer[(819,), "float32"],
+    B: T.Buffer[(128,), "float32"],
+    A_indptr: T.Buffer[(129,), "int32"],
+    A_indices: T.Buffer[(819,), "int32"],
+) -> None:
+    for i in T.serial(128):
+        with T.block("rowsum_outer"):
+            T.reads(
+                A_indptr[i : i + 1],
+                A_data[A_indptr[i] + 0 : A_indptr[i] + (A_indptr[i + 1] - A_indptr[i])],
+            )
+            T.writes(B[i])
+            with T.block("rowsum_init"):
+                T.reads()
+                T.writes(B[i])
+                B[i] = T.float32(0)
+            for k in T.serial(A_indptr[i + 1] - A_indptr[i]):
+                with T.block():
+                    T.reads(A_indptr[i], A_data[A_indptr[i] + k], B[i])
+                    T.writes(B[i])
+                    A_data_local = T.alloc_buffer([819], dtype="float32", scope="local")
+                    with T.block("A_data_cache_read"):
+                        T.reads(A_indptr[i], A_data[A_indptr[i] + k])
+                        T.writes(A_data_local[A_indptr[i] + k])
+                        A_data_local[A_indptr[i] + k] = A_data[A_indptr[i] + k]
+                    with T.block("rowsum_inner"):
+                        T.reads(B[i], A_indptr[i], A_data[A_indptr[i] + k])
+                        T.writes(B[i])
+                        B[i] = B[i] + A_data_local[A_indptr[i] + k]
+
+
+@T.prim_func
+def compacted_sparse_read_cache(
+    A_data: T.Buffer[(819,), "float32"],
+    B: T.Buffer[(128,), "float32"],
+    A_indptr: T.Buffer[(129,), "int32"],
+    A_indices: T.Buffer[(819,), "int32"],
+) -> None:
+    for i in T.serial(128):
+        with T.block("rowsum_outer"):
+            T.reads(
+                A_indptr[i : i + 1],
+                A_data[A_indptr[i] + 0 : A_indptr[i] + 0 + (A_indptr[i + 1] - A_indptr[i])],
+            )
+            T.writes(B[i])
+            with T.block("rowsum_init"):
+                T.reads()
+                T.writes(B[i])
+                B[i] = T.float32(0)
+            for k in T.serial(A_indptr[i + 1] - A_indptr[i]):
+                with T.block():
+                    T.reads(A_indptr[i], A_data[A_indptr[i] + k], B[i])
+                    T.writes(B[i])
+                    A_data_local = T.alloc_buffer([1], dtype="float32", scope="local")
+                    with T.block("A_data_cache_read"):
+                        T.reads(A_indptr[i], A_data[A_indptr[i] + k])
+                        T.writes(A_data_local[A_indptr[i] + k - (A_indptr[i] + k)])
+                        A_data_local[A_indptr[i] + k - (A_indptr[i] + k)] = A_data[A_indptr[i] + k]
+                    with T.block("rowsum_inner"):
+                        T.reads(B[i], A_indptr[i], A_data[A_indptr[i] + k])
+                        T.writes(B[i])
+                        B[i] = B[i] + A_data_local[A_indptr[i] + k - (A_indptr[i] + k)]
+
+
 def test_elementwise():
     _check(elementwise_func, compacted_elementwise_func)
 
@@ -562,6 +628,10 @@ def test_opaque_access_annotated_func():
     _check(opaque_access_annotated_func, compacted_opaque_access_annotated_func)
 
 
+def test_sparse_read_cache():
+    _check(sparse_read_cache, compacted_sparse_read_cache)
+
+
 if __name__ == "__main__":
     test_elementwise()
     test_unschedulable_block()
@@ -576,3 +646,4 @@ def test_opaque_access_annotated_func():
     test_padding_pattern()
     test_mem_access_in_branch_func()
     test_opaque_access_annotated_func()
+    test_sparse_read_cache()
diff --git a/tests/python/unittest/test_tir_transform_extract_constants.py b/tests/python/unittest/test_tir_transform_extract_constants.py
new file mode 100644
index 000000000000..74144f252ade
--- /dev/null
+++ b/tests/python/unittest/test_tir_transform_extract_constants.py
@@ -0,0 +1,64 @@
+# 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.
+import numpy as np
+import tvm
+from tvm import tir
+from tvm.script import tir as T
+
+
+@tvm.script.ir_module
+class Module4:
+    @T.prim_func
+    def constant1(a: T.handle) -> None:
+        A = T.match_buffer(a, (10), "int32")
+        B = T.alloc_buffer((10), "int32")
+        K = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = A[x] + T.load("int32", K, x)
+
+    @T.prim_func
+    def constant2(a: T.handle) -> None:
+        A = T.match_buffer(a, (10), "int32")
+        B = T.alloc_buffer((10), "int32")
+        K = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = A[x] + T.load("int32", K, x)
+
+    @T.prim_func
+    def constant3(a: T.handle) -> None:
+        A = T.match_buffer(a, (10), "int32")
+        B = T.alloc_buffer((10), "int32")
+        K = T.allocate_const([1, 2, 3, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = A[x] + T.load("int32", K, x)
+
+
+def test_const_extraction():
+    mod = tvm.tir.transform.ExtractPrimFuncConstants()(Module4)
+    constants = mod.attrs["Constants"]
+    assert len(constants) == 2
+
+    def _visit(stmt):
+        if isinstance(stmt, tvm.tir.AllocateConst):
+            assert np.array_equal(stmt.data.numpy(), constants[int(stmt.irmod_storage_idx)].numpy())
+
+    for n, f in mod.functions.items():
+        tvm.tir.stmt_functor.post_order_visit(f.body, _visit)
+
+
+if __name__ == "__main__":
+    test_const_extraction()
diff --git a/tests/python/unittest/test_tir_transform_inject_software_pipeline.py b/tests/python/unittest/test_tir_transform_inject_software_pipeline.py
new file mode 100644
index 000000000000..1432be4efbe1
--- /dev/null
+++ b/tests/python/unittest/test_tir_transform_inject_software_pipeline.py
@@ -0,0 +1,824 @@
+# 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.
+import pytest
+import sys
+
+import tvm
+from tvm import tir, te, TVMError
+from tvm.script import tir as T
+
+
+def _check(original, transformed):
+    func = original
+    mod = tvm.IRModule.from_expr(func)
+    mod = tvm.tir.transform.InjectSoftwarePipeline()(mod)
+    mod = tvm.tir.transform.Simplify()(mod)
+    tvm.ir.assert_structural_equal(mod["main"], transformed, True)
+
+
+def _check_error(func):
+    mod = tvm.IRModule.from_expr(func)
+    with pytest.raises(ValueError):
+        tvm.tir.transform.InjectSoftwarePipeline()(mod)
+
+
+@T.prim_func
+def trivial_pipeline(A: T.Buffer[(16, 1), "float32"], C: T.Buffer[(16, 1), "float32"]):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0, 1, annotations={"software_pipeline_stage": [0, 1], "software_pipeline_order": [0, 1]}
+        ):
+            with T.block():
+                T.reads(A[tx, i])
+                T.writes(C[tx, i])
+                B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                with T.block():
+                    T.reads(A[tx, i])
+                    T.writes(B[tx, 0])
+                    B[tx, 0] = A[tx, i] * T.float32(2)
+                with T.block():
+                    T.reads(B[tx, 0])
+                    T.writes(C[tx, i])
+                    C[tx, i] = B[tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_trivial_pipeline(
+    A: T.Buffer[(16, 1), "float32"], C: T.Buffer[(16, 1), "float32"]
+) -> None:
+    for tx in T.thread_binding(16, thread="threadIdx.x"):
+        with T.block():
+            T.reads(A[tx, 0])
+            T.writes(C[tx, 0])
+            B = T.alloc_buffer([2, 16, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads(A[tx, 0])
+                T.writes(B[0, tx, 0])
+                B[0, tx, 0] = A[tx, 0] * T.float32(2)
+            with T.block():
+                T.reads()
+                T.writes()
+                T.evaluate(0)
+            with T.block():
+                T.reads(B[0, tx, 0])
+                T.writes(C[tx, 0])
+                C[tx, 0] = B[0, tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def simple_compute(A: T.Buffer[(16, 16), "float32"], C: T.Buffer[(16, 16), "float32"]):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={"software_pipeline_stage": [0, 1], "software_pipeline_order": [0, 1]},
+        ):
+            with T.block():
+                T.reads(A[tx, i])
+                T.writes(C[tx, i])
+                B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                with T.block():
+                    T.reads(A[tx, i])
+                    T.writes(B[tx, 0])
+                    B[tx, 0] = A[tx, i] * T.float32(2)
+                with T.block():
+                    T.reads(B[tx, 0])
+                    T.writes(C[tx, i])
+                    C[tx, i] = B[tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_simple_compute(
+    A: T.Buffer[(16, 16), "float32"], C: T.Buffer[(16, 16), "float32"]
+) -> None:
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        with T.block():
+            T.reads([A[tx, 0:16]])
+            T.writes([C[tx, 0:16]])
+            B = T.alloc_buffer([2, 16, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads([A[tx, 0]])
+                T.writes([B[0, tx, 0]])
+                B[0, tx, 0] = A[tx, 0] * T.float32(2)
+            with T.block():
+                T.reads([A[tx, 1:16], B[0:2, tx, 0]])
+                T.writes([B[0:2, tx, 0], C[tx, 0:15]])
+                for i in T.serial(0, 15):
+                    with T.block():
+                        T.reads([A[tx, i + 1]])
+                        T.writes([B[(i + 1) % 2, tx, 0]])
+                        B[(i + 1) % 2, tx, 0] = A[tx, i + 1] * T.float32(2)
+                    with T.block():
+                        T.reads([B[i % 2, tx, 0]])
+                        T.writes([C[tx, i]])
+                        C[tx, i] = B[i % 2, tx, 0] + T.float32(1)
+            with T.block():
+                T.reads([B[1, tx, 0]])
+                T.writes([C[tx, 15]])
+                C[tx, 15] = B[1, tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def nested_pipeline_simple(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 1, 1, 1],
+                "software_pipeline_order": [0, 1, 2, 3],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i, 0:16])
+                T.writes(C[tx, i, 0:16])
+                A_shared = T.alloc_buffer((16, 1, 16), dtype="float32", scope="shared")
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads(A[tx, i, j])
+                        T.writes(A_shared[tx, 0, j])
+                        A_shared[tx, 0, j] = A[tx, i, j]
+                for j in T.serial(
+                    0,
+                    16,
+                    annotations={
+                        "software_pipeline_stage": [0, 1],
+                        "software_pipeline_order": [0, 1],
+                    },
+                ):
+                    with T.block():
+                        T.reads(A_shared[tx, 0, j])
+                        T.writes(C[tx, i, j])
+                        B = T.alloc_buffer((16, 1, 1), dtype="float32", scope="shared")
+                        with T.block():
+                            T.reads(A_shared[tx, i, j])
+                            T.writes(B[tx, i, 0])
+                            B[tx, i, 0] = A_shared[tx, 0, j] * T.float32(2)
+                        with T.block():
+                            T.reads(B[tx, i, 0])
+                            T.writes(C[tx, i, j])
+                            C[tx, i, j] = B[tx, i, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_nested_pipeline_simple(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+) -> None:
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        with T.block():
+            T.reads([A[tx, 0:16, 0:16]])
+            T.writes([C[tx, 0:16, 0:16]])
+            A_shared = T.alloc_buffer([2, 16, 1, 16], dtype="float32", scope="shared")
+            B = T.alloc_buffer([2, 16, 1, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads([A[tx, 0, 0:16]])
+                T.writes([A_shared[0, tx, 0, 0:16]])
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads([A[tx, 0, j]])
+                        T.writes([A_shared[0, tx, 0, j]])
+                        A_shared[0, tx, 0, j] = A[tx, 0, j]
+            with T.block():
+                T.reads([A[tx, 1:16, 0:16], A_shared[0:2, tx, 0:15, 0:16], B[0:2, tx, 0:15, 0]])
+                T.writes([A_shared[0:2, tx, 0, 0:16], B[0:2, tx, 0:15, 0], C[tx, 0:15, 0:16]])
+                for i in T.serial(0, 15):
+                    with T.block():
+                        T.reads([A[tx, i + 1, 0:16]])
+                        T.writes([A_shared[(i + 1) % 2, tx, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A[tx, i + 1, j]])
+                                T.writes([A_shared[(i + 1) % 2, tx, 0, j]])
+                                A_shared[(i + 1) % 2, tx, 0, j] = A[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_shared[i % 2, tx, i, 0]])
+                        T.writes([B[0, tx, i, 0]])
+                        B[0, tx, i, 0] = A_shared[i % 2, tx, 0, 0] * T.float32(2)
+                    with T.block():
+                        T.reads([A_shared[i % 2, tx, i, 1:16], B[0:2, tx, i, 0]])
+                        T.writes([B[0:2, tx, i, 0], C[tx, i, 0:15]])
+                        for j in T.serial(0, 15):
+                            with T.block():
+                                T.reads([A_shared[i % 2, tx, i, j + 1]])
+                                T.writes([B[(j + 1) % 2, tx, i, 0]])
+                                B[(j + 1) % 2, tx, i, 0] = A_shared[
+                                    i % 2, tx, 0, j + 1
+                                ] * T.float32(2)
+                            with T.block():
+                                T.reads([B[j % 2, tx, i, 0]])
+                                T.writes([C[tx, i, j]])
+                                C[tx, i, j] = B[j % 2, tx, i, 0] + T.float32(1)
+                    with T.block():
+                        T.reads([B[1, tx, i, 0]])
+                        T.writes([C[tx, i, 15]])
+                        C[tx, i, 15] = B[1, tx, i, 0] + T.float32(1)
+            with T.block():
+                T.reads([A_shared[1, tx, 15, 0:16], B[0:2, tx, 15, 0]])
+                T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:16]])
+                with T.block():
+                    T.reads([A_shared[1, tx, 15, 0]])
+                    T.writes([B[0, tx, 15, 0]])
+                    B[0, tx, 15, 0] = A_shared[1, tx, 0, 0] * T.float32(2)
+                with T.block():
+                    T.reads([A_shared[1, tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                    T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:15]])
+                    for j in T.serial(0, 15):
+                        with T.block():
+                            T.reads([A_shared[1, tx, 15, j + 1]])
+                            T.writes([B[(j + 1) % 2, tx, 15, 0]])
+                            B[(j + 1) % 2, tx, 15, 0] = A_shared[1, tx, 0, j + 1] * T.float32(2)
+                        with T.block():
+                            T.reads([B[j % 2, tx, 15, 0]])
+                            T.writes([C[tx, 15, j]])
+                            C[tx, 15, j] = B[j % 2, tx, 15, 0] + T.float32(1)
+                with T.block():
+                    T.reads([B[1, tx, 15, 0]])
+                    T.writes([C[tx, 15, 15]])
+                    C[tx, 15, 15] = B[1, tx, 15, 0] + T.float32(1)
+
+
+@T.prim_func
+def nested_pipeline_prefetch_inner(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 0, 1, 1],
+                "software_pipeline_order": [0, 2, 1, 3],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i, 0:16])
+                T.writes(C[tx, i, 0:16])
+                A_shared = T.alloc_buffer((16, 1, 16), dtype="float32", scope="shared")
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads(A[tx, i, j])
+                        T.writes(A_shared[tx, 0, j])
+                        A_shared[tx, 0, j] = A[tx, i, j]
+                for j in T.serial(
+                    0,
+                    16,
+                    annotations={
+                        "software_pipeline_stage": [0, 1],
+                        "software_pipeline_order": [0, 1],
+                    },
+                ):
+                    with T.block():
+                        T.reads(A_shared[tx, 0, j])
+                        T.writes(C[tx, i, j])
+                        B = T.alloc_buffer((16, 1, 1), dtype="float32", scope="shared")
+                        with T.block():
+                            T.reads(A_shared[tx, i, j])
+                            T.writes(B[tx, i, 0])
+                            B[tx, i, 0] = A_shared[tx, 0, j] * T.float32(2)
+                        with T.block():
+                            T.reads(B[tx, i, 0])
+                            T.writes(C[tx, i, j])
+                            C[tx, i, j] = B[tx, i, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_nested_pipeline_prefetch_inner(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+) -> None:
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        with T.block():
+            T.reads([A[tx, 0:16, 0:16]])
+            T.writes([C[tx, 0:16, 0:16]])
+            A_shared = T.alloc_buffer([2, 16, 1, 16], dtype="float32", scope="shared")
+            B = T.alloc_buffer([2, 16, 1, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads([A[tx, 0, 0:16], A_shared[0, tx, 0, 0]])
+                T.writes([A_shared[0, tx, 0, 0:16], B[0, tx, 0, 0]])
+                with T.block():
+                    T.reads([A[tx, 0, 0:16]])
+                    T.writes([A_shared[0, tx, 0, 0:16]])
+                    for j in T.serial(0, 16):
+                        with T.block():
+                            T.reads([A[tx, 0, j]])
+                            T.writes([A_shared[0, tx, 0, j]])
+                            A_shared[0, tx, 0, j] = A[tx, 0, j]
+                with T.block():
+                    T.reads([A_shared[0, tx, 0, 0]])
+                    T.writes([B[0, tx, 0, 0]])
+                    B[0, tx, 0, 0] = A_shared[0, tx, 0, 0] * T.float32(2)
+            with T.block():
+                T.reads([A[tx, 1:16, 0:16], A_shared[0:2, tx, 0:16, 0:16], B[0:2, tx, 0:15, 0]])
+                T.writes([A_shared[0:2, tx, 0, 0:16], B[0:2, tx, 0:16, 0], C[tx, 0:15, 0:16]])
+                for i in T.serial(0, 15):
+                    with T.block():
+                        T.reads([A[tx, i + 1, 0:16]])
+                        T.writes([A_shared[(i + 1) % 2, tx, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A[tx, i + 1, j]])
+                                T.writes([A_shared[(i + 1) % 2, tx, 0, j]])
+                                A_shared[(i + 1) % 2, tx, 0, j] = A[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_shared[i % 2, tx, i, 1:16], B[0:2, tx, i, 0]])
+                        T.writes([B[0:2, tx, i, 0], C[tx, i, 0:15]])
+                        for j in T.serial(0, 15):
+                            with T.block():
+                                T.reads([A_shared[i % 2, tx, i, j + 1]])
+                                T.writes([B[(j + 1) % 2, tx, i, 0]])
+                                B[(j + 1) % 2, tx, i, 0] = A_shared[
+                                    i % 2, tx, 0, j + 1
+                                ] * T.float32(2)
+                            with T.block():
+                                T.reads([B[j % 2, tx, i, 0]])
+                                T.writes([C[tx, i, j]])
+                                C[tx, i, j] = B[j % 2, tx, i, 0] + T.float32(1)
+                    with T.block():
+                        T.reads([A_shared[(i + 1) % 2, tx, i + 1, 0]])
+                        T.writes([B[0, tx, i + 1, 0]])
+                        B[0, tx, i + 1, 0] = A_shared[(i + 1) % 2, tx, 0, 0] * T.float32(2)
+                    with T.block():
+                        T.reads([B[1, tx, i, 0]])
+                        T.writes([C[tx, i, 15]])
+                        C[tx, i, 15] = B[1, tx, i, 0] + T.float32(1)
+            with T.block():
+                T.reads([A_shared[1, tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:16]])
+                with T.block():
+                    T.reads([A_shared[1, tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                    T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:15]])
+                    for j in T.serial(0, 15):
+                        with T.block():
+                            T.reads([A_shared[1, tx, 15, j + 1]])
+                            T.writes([B[(j + 1) % 2, tx, 15, 0]])
+                            B[(j + 1) % 2, tx, 15, 0] = A_shared[1, tx, 0, j + 1] * T.float32(2)
+                        with T.block():
+                            T.reads([B[j % 2, tx, 15, 0]])
+                            T.writes([C[tx, 15, j]])
+                            C[tx, 15, j] = B[j % 2, tx, 15, 0] + T.float32(1)
+                with T.block():
+                    T.reads([B[1, tx, 15, 0]])
+                    T.writes([C[tx, 15, 15]])
+                    C[tx, 15, 15] = B[1, tx, 15, 0] + T.float32(1)
+
+
+@T.prim_func
+def nested_pipeline_interleaving(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 0, 0, 1, 1],
+                "software_pipeline_order": [0, 2, 3, 1, 4],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i, 0:16])
+                T.writes(C[tx, i, 0:16])
+                A_shared = T.alloc_buffer((16, 1, 16), dtype="float32", scope="shared")
+                A_local = T.alloc_buffer((1, 1, 16), dtype="float32", scope="local")
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads(A[tx, i, j])
+                        T.writes(A_shared[tx, 0, j])
+                        A_shared[tx, 0, j] = A[tx, i, j]
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads(A_shared[tx, 0, j])
+                        T.writes(A_local[0, 0, j])
+                        A_local[0, 0, j] = A_shared[tx, i, j]
+                for j in T.serial(
+                    0,
+                    16,
+                    annotations={
+                        "software_pipeline_stage": [0, 1],
+                        "software_pipeline_order": [0, 1],
+                    },
+                ):
+                    with T.block():
+                        T.reads(A_local[0, 0, j])
+                        T.writes(C[tx, i, j])
+                        B = T.alloc_buffer((16, 1, 1), dtype="float32", scope="shared")
+                        with T.block():
+                            T.reads(A_local[tx, i, j])
+                            T.writes(B[tx, i, 0])
+                            B[tx, i, 0] = A_local[0, 0, j] * T.float32(2)
+                        with T.block():
+                            T.reads(B[tx, i, 0])
+                            T.writes(C[tx, i, j])
+                            C[tx, i, j] = B[tx, i, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_nested_pipeline_interleaving(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+) -> None:
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        with T.block():
+            T.reads([A[tx, 0:16, 0:16]])
+            T.writes([C[tx, 0:16, 0:16]])
+            A_shared = T.alloc_buffer([16, 1, 16], dtype="float32", scope="shared")
+            A_local = T.alloc_buffer([1, 1, 16], dtype="float32", scope="local")
+            B = T.alloc_buffer([2, 16, 1, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads([A[tx, 0, 0:16], A_shared[tx, 0, 0:16], A_local[tx, 0, 0]])
+                T.writes([A_shared[tx, 0, 0:16], A_local[0, 0, 0:16], B[0, tx, 0, 0]])
+                with T.block():
+                    T.reads([A[tx, 0, 0:16]])
+                    T.writes([A_shared[tx, 0, 0:16]])
+                    for j in T.serial(0, 16):
+                        with T.block():
+                            T.reads([A[tx, 0, j]])
+                            T.writes([A_shared[tx, 0, j]])
+                            A_shared[tx, 0, j] = A[tx, 0, j]
+                with T.block():
+                    T.reads([A_shared[tx, 0, 0:16]])
+                    T.writes([A_local[0, 0, 0:16]])
+                    for j in T.serial(0, 16):
+                        with T.block():
+                            T.reads([A_shared[tx, 0, j]])
+                            T.writes([A_local[0, 0, j]])
+                            A_local[0, 0, j] = A_shared[tx, 0, j]
+                with T.block():
+                    T.reads([A_local[tx, 0, 0]])
+                    T.writes([B[0, tx, 0, 0]])
+                    B[0, tx, 0, 0] = A_local[0, 0, 0] * T.float32(2)
+            with T.block():
+                T.reads(
+                    [
+                        A[tx, 1:16, 0:16],
+                        A_local[tx, 0:16, 0:16],
+                        B[0:2, tx, 0:15, 0],
+                        A_shared[tx, 0, 0:16],
+                    ]
+                )
+                T.writes(
+                    [
+                        A_shared[tx, 0, 0:16],
+                        B[0:2, tx, 0:16, 0],
+                        C[tx, 0:15, 0:16],
+                        A_local[0, 0, 0:16],
+                    ]
+                )
+                for i in T.serial(0, 15):
+                    with T.block():
+                        T.reads([A[tx, i + 1, 0:16]])
+                        T.writes([A_shared[tx, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A[tx, i + 1, j]])
+                                T.writes([A_shared[tx, 0, j]])
+                                A_shared[tx, 0, j] = A[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_local[tx, i, 1:16], B[0:2, tx, i, 0]])
+                        T.writes([B[0:2, tx, i, 0], C[tx, i, 0:15]])
+                        for j in T.serial(0, 15):
+                            with T.block():
+                                T.reads([A_local[tx, i, j + 1]])
+                                T.writes([B[(j + 1) % 2, tx, i, 0]])
+                                B[(j + 1) % 2, tx, i, 0] = A_local[0, 0, j + 1] * T.float32(2)
+                            with T.block():
+                                T.reads([B[j % 2, tx, i, 0]])
+                                T.writes([C[tx, i, j]])
+                                C[tx, i, j] = B[j % 2, tx, i, 0] + T.float32(1)
+                    with T.block():
+                        T.reads([A_shared[tx, 0, 0:16]])
+                        T.writes([A_local[0, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A_shared[tx, 0, j]])
+                                T.writes([A_local[0, 0, j]])
+                                A_local[0, 0, j] = A_shared[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_local[tx, i + 1, 0]])
+                        T.writes([B[0, tx, i + 1, 0]])
+                        B[0, tx, i + 1, 0] = A_local[0, 0, 0] * T.float32(2)
+                    with T.block():
+                        T.reads([B[1, tx, i, 0]])
+                        T.writes([C[tx, i, 15]])
+                        C[tx, i, 15] = B[1, tx, i, 0] + T.float32(1)
+            with T.block():
+                T.reads([A_local[tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:16]])
+                with T.block():
+                    T.reads([A_local[tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                    T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:15]])
+                    for j in T.serial(0, 15):
+                        with T.block():
+                            T.reads([A_local[tx, 15, j + 1]])
+                            T.writes([B[(j + 1) % 2, tx, 15, 0]])
+                            B[(j + 1) % 2, tx, 15, 0] = A_local[0, 0, j + 1] * T.float32(2)
+                        with T.block():
+                            T.reads([B[j % 2, tx, 15, 0]])
+                            T.writes([C[tx, 15, j]])
+                            C[tx, 15, j] = B[j % 2, tx, 15, 0] + T.float32(1)
+                with T.block():
+                    T.reads([B[1, tx, 15, 0]])
+                    T.writes([C[tx, 15, 15]])
+                    C[tx, 15, 15] = B[1, tx, 15, 0] + T.float32(1)
+
+
+@T.prim_func
+def nested_pipeline_double_buffer(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 0, 0, 1, 1],
+                "software_pipeline_order": [0, 2, 3, 1, 4],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i, 0:16])
+                T.writes(C[tx, i, 0:16])
+                A_shared = T.alloc_buffer((16, 1, 16), dtype="float32", scope="shared")
+                A_local = T.alloc_buffer((1, 1, 16), dtype="float32", scope="local")
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.reads(A[tx, i, j])
+                        T.writes(A_shared[tx, 0, j])
+                        A_shared[tx, 0, j] = A[tx, i, j]
+                for j in T.serial(0, 16):
+                    with T.block():
+                        T.block_attr({"double_buffer_scope": 0})
+                        T.reads(A_shared[tx, 0, j])
+                        T.writes(A_local[0, 0, j])
+                        A_local[0, 0, j] = A_shared[tx, i, j]
+                for j in T.serial(
+                    0,
+                    16,
+                    annotations={
+                        "software_pipeline_stage": [0, 1],
+                        "software_pipeline_order": [0, 1],
+                    },
+                ):
+                    with T.block():
+                        T.reads(A_local[0, 0, j])
+                        T.writes(C[tx, i, j])
+                        B = T.alloc_buffer((16, 1, 1), dtype="float32", scope="shared")
+                        with T.block():
+                            T.reads(A_local[tx, i, j])
+                            T.writes(B[tx, i, 0])
+                            B[tx, i, 0] = A_local[0, 0, j] * T.float32(2)
+                        with T.block():
+                            T.reads(B[tx, i, 0])
+                            T.writes(C[tx, i, j])
+                            C[tx, i, j] = B[tx, i, 0] + T.float32(1)
+
+
+@T.prim_func
+def transformed_nested_pipeline_double_buffer(
+    A: T.Buffer[(16, 16, 16), "float32"], C: T.Buffer[(16, 16, 16), "float32"]
+) -> None:
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        with T.block():
+            T.reads([A[tx, 0:16, 0:16]])
+            T.writes([C[tx, 0:16, 0:16]])
+            A_shared = T.alloc_buffer([16, 1, 16], dtype="float32", scope="shared")
+            A_local = T.alloc_buffer([2, 1, 1, 16], dtype="float32", scope="local")
+            B = T.alloc_buffer([2, 16, 1, 1], dtype="float32", scope="shared")
+            with T.block():
+                T.reads([A[tx, 0, 0:16], A_shared[tx, 0, 0:16], A_local[0, tx, 0, 0]])
+                T.writes([A_shared[tx, 0, 0:16], A_local[0, 0, 0, 0:16], B[0, tx, 0, 0]])
+                with T.block():
+                    T.reads([A[tx, 0, 0:16]])
+                    T.writes([A_shared[tx, 0, 0:16]])
+                    for j in T.serial(0, 16):
+                        with T.block():
+                            T.reads([A[tx, 0, j]])
+                            T.writes([A_shared[tx, 0, j]])
+                            A_shared[tx, 0, j] = A[tx, 0, j]
+                with T.block():
+                    T.reads([A_shared[tx, 0, 0:16]])
+                    T.writes([A_local[0, 0, 0, 0:16]])
+                    for j in T.serial(0, 16):
+                        with T.block():
+                            T.reads([A_shared[tx, 0, j]])
+                            T.writes([A_local[0, 0, 0, j]])
+                            T.block_attr({"double_buffer_scope": 0})
+                            A_local[0, 0, 0, j] = A_shared[tx, 0, j]
+                with T.block():
+                    T.reads([A_local[0, tx, 0, 0]])
+                    T.writes([B[0, tx, 0, 0]])
+                    B[0, tx, 0, 0] = A_local[0, 0, 0, 0] * T.float32(2)
+            with T.block():
+                T.reads(
+                    [
+                        A[tx, 1:16, 0:16],
+                        A_local[0:2, tx, 0:16, 0:16],
+                        B[0:2, tx, 0:15, 0],
+                        A_shared[tx, 0, 0:16],
+                    ]
+                )
+                T.writes(
+                    [
+                        A_shared[tx, 0, 0:16],
+                        B[0:2, tx, 0:16, 0],
+                        C[tx, 0:15, 0:16],
+                        A_local[0:2, 0, 0, 0:16],
+                    ]
+                )
+                for i in T.serial(0, 15):
+                    with T.block():
+                        T.reads([A[tx, i + 1, 0:16]])
+                        T.writes([A_shared[tx, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A[tx, i + 1, j]])
+                                T.writes([A_shared[tx, 0, j]])
+                                A_shared[tx, 0, j] = A[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_local[i % 2, tx, i, 1:16], B[0:2, tx, i, 0]])
+                        T.writes([B[0:2, tx, i, 0], C[tx, i, 0:15]])
+                        for j in T.serial(0, 15):
+                            with T.block():
+                                T.reads([A_local[i % 2, tx, i, j + 1]])
+                                T.writes([B[(j + 1) % 2, tx, i, 0]])
+                                B[(j + 1) % 2, tx, i, 0] = A_local[i % 2, 0, 0, j + 1] * T.float32(
+                                    2
+                                )
+                            with T.block():
+                                T.reads([B[j % 2, tx, i, 0]])
+                                T.writes([C[tx, i, j]])
+                                C[tx, i, j] = B[j % 2, tx, i, 0] + T.float32(1)
+                    with T.block():
+                        T.reads([A_shared[tx, 0, 0:16]])
+                        T.writes([A_local[(i + 1) % 2, 0, 0, 0:16]])
+                        for j in T.serial(0, 16):
+                            with T.block():
+                                T.reads([A_shared[tx, 0, j]])
+                                T.writes([A_local[(i + 1) % 2, 0, 0, j]])
+                                T.block_attr({"double_buffer_scope": 0})
+                                A_local[(i + 1) % 2, 0, 0, j] = A_shared[tx, i + 1, j]
+                    with T.block():
+                        T.reads([A_local[(i + 1) % 2, tx, i + 1, 0]])
+                        T.writes([B[0, tx, i + 1, 0]])
+                        B[0, tx, i + 1, 0] = A_local[(i + 1) % 2, 0, 0, 0] * T.float32(2)
+                    with T.block():
+                        T.reads([B[1, tx, i, 0]])
+                        T.writes([C[tx, i, 15]])
+                        C[tx, i, 15] = B[1, tx, i, 0] + T.float32(1)
+            with T.block():
+                T.reads([A_local[1, tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:16]])
+                with T.block():
+                    T.reads([A_local[1, tx, 15, 1:16], B[0:2, tx, 15, 0]])
+                    T.writes([B[0:2, tx, 15, 0], C[tx, 15, 0:15]])
+                    for j in T.serial(0, 15):
+                        with T.block():
+                            T.reads([A_local[1, tx, 15, j + 1]])
+                            T.writes([B[(j + 1) % 2, tx, 15, 0]])
+                            B[(j + 1) % 2, tx, 15, 0] = A_local[1, 0, 0, j + 1] * T.float32(2)
+                        with T.block():
+                            T.reads([B[j % 2, tx, 15, 0]])
+                            T.writes([C[tx, 15, j]])
+                            C[tx, 15, j] = B[j % 2, tx, 15, 0] + T.float32(1)
+                with T.block():
+                    T.reads([B[1, tx, 15, 0]])
+                    T.writes([C[tx, 15, 15]])
+                    C[tx, 15, 15] = B[1, tx, 15, 0] + T.float32(1)
+
+
+@T.prim_func
+def simple_compute_incorrect_reorder(
+    A: T.Buffer[(16, 16), "float32"], D: T.Buffer[(16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 1, 1],
+                "software_pipeline_order": [0, 2, 1],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i])
+                T.writes(D[tx, i])
+                B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                C = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                with T.block():
+                    T.reads(A[tx, i])
+                    T.writes(B[tx, 0])
+                    B[tx, 0] = A[tx, i] * T.float32(2)
+                with T.block():
+                    T.reads(B[tx, 0])
+                    T.writes(C[tx, 0])
+                    C[tx, 0] = B[tx, 0] + T.float32(2)
+                with T.block():
+                    T.reads(C[tx, 0])
+                    T.writes(D[tx, i])
+                    D[tx, i] = C[tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def simple_compute_conflicting_order(
+    A: T.Buffer[(16, 16), "float32"], D: T.Buffer[(16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(
+            0,
+            16,
+            annotations={
+                "software_pipeline_stage": [0, 1, 1],
+                "software_pipeline_order": [0, 1, 1],
+            },
+        ):
+            with T.block():
+                T.reads(A[tx, i])
+                T.writes(D[tx, i])
+                B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                C = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                with T.block():
+                    T.reads(A[tx, i])
+                    T.writes(B[tx, 0])
+                    B[tx, 0] = A[tx, i] * T.float32(2)
+                with T.block():
+                    T.reads(B[tx, 0])
+                    T.writes(C[tx, 0])
+                    C[tx, 0] = B[tx, 0] + T.float32(2)
+                with T.block():
+                    T.reads(C[tx, 0])
+                    T.writes(D[tx, i])
+                    D[tx, i] = C[tx, 0] + T.float32(1)
+
+
+@T.prim_func
+def simple_compute_missing_annotation(
+    A: T.Buffer[(16, 16), "float32"], C: T.Buffer[(16, 16), "float32"]
+):
+    for tx in T.thread_binding(0, 16, thread="threadIdx.x"):
+        for i in T.serial(0, 16, annotations={"software_pipeline_stage": [0, 1]}):
+            with T.block():
+                T.reads(A[tx, i])
+                T.writes(C[tx, i])
+                B = T.alloc_buffer((16, 1), dtype="float32", scope="shared")
+                with T.block():
+                    T.reads(A[tx, i])
+                    T.writes(B[tx, 0])
+                    B[tx, 0] = A[tx, i] * T.float32(2)
+                with T.block():
+                    T.reads(B[tx, 0])
+                    T.writes(C[tx, i])
+                    C[tx, i] = B[tx, 0] + T.float32(1)
+
+
+def test_simple_compute():
+    _check(simple_compute, transformed_simple_compute)
+
+
+def test_trivial_pipeline():
+    _check(trivial_pipeline, transformed_trivial_pipeline)
+
+
+def test_nest_pipeline_simple():
+    _check(nested_pipeline_simple, transformed_nested_pipeline_simple)
+
+
+def test_nest_pipeline_prefetch_inner():
+    _check(nested_pipeline_prefetch_inner, transformed_nested_pipeline_prefetch_inner)
+
+
+def test_nest_pipeline_interleaving():
+    _check(nested_pipeline_interleaving, transformed_nested_pipeline_interleaving)
+
+
+def test_nest_pipeline_double_buffer():
+    _check(nested_pipeline_double_buffer, transformed_nested_pipeline_double_buffer)
+
+
+def test_error_reorder():
+    _check_error(simple_compute_incorrect_reorder)
+
+
+def test_error_conflicting_order():
+    _check_error(simple_compute_conflicting_order)
+
+
+def test_error_missing_annotation():
+    _check_error(simple_compute_missing_annotation)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/unittest/test_tir_transform_loop_partition.py b/tests/python/unittest/test_tir_transform_loop_partition.py
index 8b109172ea09..8a71169ca78c 100644
--- a/tests/python/unittest/test_tir_transform_loop_partition.py
+++ b/tests/python/unittest/test_tir_transform_loop_partition.py
@@ -17,6 +17,7 @@
 import tvm
 import tvm.testing
 from tvm import te
+from tvm.ir.module import IRModule
 from tvm.script import tir as T
 import numpy
 
@@ -565,6 +566,78 @@ def test_explicit_partition_hint():
     assert tvm.ir.structural_equal(mod["main"], partitioned_concat)
 
 
+@T.prim_func
+def partitioned_concat_3(
+    placeholder: T.Buffer[(1, 64, 28, 28), "int8"],
+    placeholder_1: T.Buffer[(1, 32, 28, 28), "int8"],
+    placeholder_2: T.Buffer[(1, 32, 28, 28), "int8"],
+    T_concat: T.Buffer[(1, 128, 28, 28), "int8"],
+) -> None:
+    for i1, i2, i3 in T.grid(64, 28, 28):
+        T.store(
+            T_concat.data,
+            i1 * 784 + i2 * 28 + i3,
+            T.load("int8", placeholder.data, i1 * 784 + i2 * 28 + i3),
+            True,
+        )
+    for i1, i2, i3 in T.grid(32, 28, 28):
+        T.store(
+            T_concat.data,
+            i1 * 784 + i2 * 28 + i3 + 50176,
+            T.load("int8", placeholder_1.data, i1 * 784 + i2 * 28 + i3),
+            True,
+        )
+    for i1, i2, i3 in T.grid(32, 28, 28):
+        T.store(
+            T_concat.data,
+            i1 * 784 + i2 * 28 + i3 + 75264,
+            T.load("int8", placeholder_2.data, i1 * 784 + i2 * 28 + i3),
+            True,
+        )
+
+
+@T.prim_func
+def concat_func_3(
+    placeholder: T.Buffer[(1, 64, 28, 28), "int8"],
+    placeholder_1: T.Buffer[(1, 32, 28, 28), "int8"],
+    placeholder_2: T.Buffer[(1, 32, 28, 28), "int8"],
+    T_concat: T.Buffer[(1, 128, 28, 28), "int8"],
+) -> None:
+    for i1 in T.serial(128, annotations={"pragma_loop_partition_hint": 1}):
+        for i2, i3 in T.grid(28, 28):
+            if 96 <= i1:
+                T.store(
+                    T_concat.data,
+                    i1 * 784 + i2 * 28 + i3,
+                    T.load("int8", placeholder_2.data, i1 * 784 + i2 * 28 + i3 - 75264),
+                    True,
+                )
+            if 64 <= i1 and i1 < 96:
+                T.store(
+                    T_concat.data,
+                    i1 * 784 + i2 * 28 + i3,
+                    T.load("int8", placeholder_1.data, i1 * 784 + i2 * 28 + i3 - 50176),
+                    True,
+                )
+            if i1 < 64:
+                T.store(
+                    T_concat.data,
+                    i1 * 784 + i2 * 28 + i3,
+                    T.load("int8", placeholder.data, i1 * 784 + i2 * 28 + i3),
+                    True,
+                )
+
+
+def test_condition_mutually_exclusive():
+    mod = IRModule.from_expr(concat_func_3)
+    with tvm.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}):
+        mod = tvm.tir.transform.FlattenBuffer()(mod)
+        mod = tvm.tir.transform.LoopPartition()(mod)
+        mod = tvm.tir.transform.Simplify()(mod)
+        mod = tvm.tir.transform.RemoveNoOp()(mod)
+    assert tvm.ir.structural_equal(mod["main"], partitioned_concat_3)
+
+
 if __name__ == "__main__":
     test_basic()
     test_const_loop()
@@ -587,3 +660,4 @@ def test_explicit_partition_hint():
     test_multilevel_splitting_with_indivisble_factors()
     test_simple_rfactor()
     test_explicit_partition_hint()
+    test_condition_mutually_exclusive()
diff --git a/tests/python/unittest/test_tir_transform_narrow_datatype.py b/tests/python/unittest/test_tir_transform_narrow_datatype.py
index 9b95266d3287..667fad0317db 100644
--- a/tests/python/unittest/test_tir_transform_narrow_datatype.py
+++ b/tests/python/unittest/test_tir_transform_narrow_datatype.py
@@ -27,7 +27,7 @@ def lower_stmt(params, stmt, target_bits):
     return stmt
 
 
-def lower_sch(sch, args, target_bits):
+def lower_sch(sch, args, target_bits, extra_passes=None):
     binds = {}
     arg_list = []
     for x in args:
@@ -42,6 +42,9 @@ def lower_sch(sch, args, target_bits):
 
     mod = schedule_to_module(sch, args)
     mod = tvm.tir.transform.StorageFlatten(64)(mod)
+    if extra_passes:
+        for p in extra_passes:
+            mod = p(mod)
     return tvm.tir.transform.NarrowDataType(target_bits)(mod)["main"].body
 
 
@@ -255,6 +258,25 @@ def check(shape, index, target_bits, target_dtype):
     )
 
 
+def test_ramp_dtype_consistency():
+    """
+    for (i :int64, (int64)0, (int64)4) {
+        A[ramp(i*(int64)2, (int64)1, 2)] = cast(int64, 2 ** 31 - 1) * i;
+    }
+    The infer result:
+        base:   int64 -> int64 (since i is involved in another int64 expr)
+        stride: int64 -> int32
+
+    Thus ramp should still use int64 for both stride and base after rewrite.
+    """
+    n = tvm.tir.IntImm("int64", 4)
+    m = tvm.tir.IntImm("int64", 2)
+    A = te.compute((n, m), lambda i, j: tvm.tir.Cast("int64", 2 ** 31 - 1) * i, name="A")
+    s = te.create_schedule(A.op)
+    s[A].vectorize(A.op.axis[1])
+    lower_sch(s, [A], 32, extra_passes=[tvm.tir.transform.VectorizeLoop()])
+
+
 if __name__ == "__main__":
     test_basic()
     test_thread_axis()
@@ -263,3 +285,4 @@ def check(shape, index, target_bits, target_dtype):
     test_slice()
     test_relay_basic()
     test_relay_take()
+    test_ramp_dtype_consistency()
diff --git a/tests/python/unittest/test_tir_transform_vectorize.py b/tests/python/unittest/test_tir_transform_vectorize.py
index b1e580957b24..1a0d84a4f807 100644
--- a/tests/python/unittest/test_tir_transform_vectorize.py
+++ b/tests/python/unittest/test_tir_transform_vectorize.py
@@ -205,6 +205,14 @@ def test_ir(A, B, C):
         assert expected in error_msg
 
 
+def test_vectorize_dtype_mismatch():
+    n = tvm.tir.IntImm("int64", 4)
+    A = te.compute((n,), lambda i: tvm.tir.IntImm("int64", 2 ** 31 - 1) + i, name="A")
+    s = te.create_schedule(A.op)
+    s[A].vectorize(A.op.axis[0])
+    tvm.lower(s, [A], "llvm", simple_mode=True)
+
+
 if __name__ == "__main__":
     test_vectorize_vector()
     test_vectorize_with_if()
@@ -214,3 +222,4 @@ def test_ir(A, B, C):
     test_vectorize_with_ge_cond()
     test_vectorize_let()
     test_vectorize_while_fail()
+    test_vectorize_dtype_mismatch()
diff --git a/tests/python/unittest/test_tir_usmp_transform_convert_pool_allocations_to_offsets.py b/tests/python/unittest/test_tir_usmp_transform_convert_pool_allocations_to_offsets.py
index ab40c646391c..07e31a989874 100644
--- a/tests/python/unittest/test_tir_usmp_transform_convert_pool_allocations_to_offsets.py
+++ b/tests/python/unittest/test_tir_usmp_transform_convert_pool_allocations_to_offsets.py
@@ -122,9 +122,9 @@ def tvmgen_default_fused_nn_max_pool2d_cast(placeholder_28: T.handle, T_cast_6:
                 T.store(T_cast_7.data, (((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3), T.cast(T.load("uint8", tensor_2, (((ax0_ax1_fused_5*3584) + (ax2_5*64)) + ax3_3)), "int16"), True)
 
     @T.prim_func
-    def run_model(input: T.handle, output: T.handle) -> None:
+    def __tvm_main__(input: T.handle, output: T.handle) -> None:
         # function attr dict
-        T.func_attr({"global_symbol": "run_model", "runner_function": True})
+        T.func_attr({"global_symbol": "__tvm_main__", "runner_function": True})
         # body
         T.attr("default", "device_id", 0)
         T.attr("default", "device_type", 1)
@@ -140,7 +140,7 @@ def run_model(input: T.handle, output: T.handle) -> None:
 @tvm.script.ir_module
 class LinearStructurePlanned:
     @T.prim_func
-    def run_model(input: T.handle, fast_memory_0_var: T.handle, slow_memory_1_var: T.handle,  output: T.handle) -> None:
+    def __tvm_main__(input: T.handle, fast_memory_0_var: T.handle, slow_memory_1_var: T.handle,  output: T.handle) -> None:
         fast_memory_0_buffer_var = T.match_buffer(fast_memory_0_var, [200704], dtype="uint8", strides=[1], elem_offset=1, align=16)
         slow_memory_1_buffer_var = T.match_buffer(slow_memory_1_var, [1418528], dtype="uint8", strides=[1], elem_offset=1, align=16)
         # body
@@ -217,7 +217,7 @@ def test_mobilenet_subgraph():
     tir_mod = assign_poolinfos_to_allocates_in_irmodule(
         tir_mod, [fast_memory_pool, slow_memory_pool]
     )
-    main_func = tir_mod["run_model"]
+    main_func = tir_mod["__tvm_main__"]
     buffer_analysis = tvm.tir.usmp.analysis.extract_buffer_info(main_func, tir_mod)
     buffer_info_map = buffer_analysis.buffer_info_stmts
 
@@ -328,9 +328,9 @@ def tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_add_clip_cast_cast_s
                     T.store(T_cast_7.data, ax0_ax1_fused_ax2_fused_3 * 256 + ax3_outer_2 * 64 + ax3_inner_4, T.cast(T.max(T.min(T.q_multiply_shift(T.cast(T.cast(T.max(T.min(T.q_multiply_shift(T.load("int32", Conv2dOutput_3, ax3_inner_4) + T.load("int32", placeholder_26.data, ax3_outer_2 * 64 + ax3_inner_4), 1343014664, 31, -8, dtype="int32") + 136, 255), 0), "uint8"), "int32") - 136, 1073903788, 31, 1, dtype="int32") + T.load("int32", placeholder_28.data, ax0_ax1_fused_ax2_fused_3 * 256 + ax3_outer_2 * 64 + ax3_inner_4), 255), 0), "uint8"), True)
 
     @T.prim_func
-    def run_model(input: T.handle, output: T.handle) -> None:
+    def __tvm_main__(input: T.handle, output: T.handle) -> None:
         # function attr dict
-        T.func_attr({"global_symbol": "run_model", "runner_function": True})
+        T.func_attr({"global_symbol": "__tvm_main__", "runner_function": True})
         # body
         T.attr("default", "device_id", 0)
         T.attr("default", "device_type", 1)
@@ -464,7 +464,7 @@ def tvmgen_default_fused_nn_conv2d_add_fixed_point_multiply_clip_cast_cast_1(pla
                 T.store(T_cast_5.data, ax0_ax1_fused_ax2_fused_1 * 64 + ax3_inner_2, T.cast(T.cast(T.max(T.min(T.q_multiply_shift(T.load("int32", Conv2dOutput_1_let, ax3_inner_2) + T.load("int32", placeholder_15.data, ax3_inner_2), 1608879842, 31, -7, dtype="int32"), 255), 0), "uint8"), "int16"), True)
 
     @T.prim_func
-    def run_model(input: T.handle, global_workspace_0_var: T.handle, output: T.handle) -> None:
+    def __tvm_main__(input: T.handle, global_workspace_0_var: T.handle, output: T.handle) -> None:
         global_workspace_0_buffer_var = T.match_buffer(global_workspace_0_var, [7920256], dtype="uint8", strides=[1], elem_offset=1, align=16)
         # body
         T.attr("default", "device_id", 0)
@@ -491,7 +491,7 @@ def test_resnet_subgraph():
     tir_mod = ResnetStructure
     tir_mod = _assign_targets_to_primfuncs_irmodule(tir_mod, target)
     tir_mod = assign_poolinfos_to_allocates_in_irmodule(tir_mod, [global_workspace_pool])
-    main_func = tir_mod["run_model"]
+    main_func = tir_mod["__tvm_main__"]
     buffer_analysis = tvm.tir.usmp.analysis.extract_buffer_info(main_func, tir_mod)
     buffer_info_map = buffer_analysis.buffer_info_stmts
 
diff --git a/tests/python/unittest/test_tvmscript_complete.py b/tests/python/unittest/test_tvmscript_complete.py
index 882745704693..691e6cd9bbb6 100644
--- a/tests/python/unittest/test_tvmscript_complete.py
+++ b/tests/python/unittest/test_tvmscript_complete.py
@@ -210,7 +210,7 @@ def expected_bufferslice_indices(data: T.handle, index: T.handle) -> None:
         for i0, i1 in T.grid(16, 16):
             with T.block():
                 vi, vj = T.axis.remap("SS", [i0, i1])
-                T.reads([data_buf[vi, 0:16], index_buf[0]])
+                T.reads([data_buf[vi, index_buf[0]], index_buf[0]])
                 T.writes([out_buf[vi, vj]])
                 out_buf[vi, vj] = data_buf[vi, index_buf[0]]
 
@@ -238,7 +238,12 @@ def expected_recursive_bufferslice_indices(data: T.handle, index: T.handle) -> N
         for i0, i1 in T.grid(16, 16):
             with T.block():
                 vi, vj = T.axis.remap("SS", [i0, i1])
-                T.reads([data_buf[0:16, 0:16], index_buf[0]])
+                T.reads(
+                    [
+                        data_buf[index_buf[index_buf[0]], index_buf[0]],
+                        index_buf[T.min(index_buf[0], 0) : T.max(index_buf[0], 0) + 1],
+                    ]
+                )
                 T.writes([out_buf[vi, vj]])
                 out_buf[vi, vj] = data_buf[index_buf[index_buf[0]], index_buf[0]]
 
diff --git a/tests/python/unittest/test_tvmscript_roundtrip.py b/tests/python/unittest/test_tvmscript_roundtrip.py
index 51a4ce7960a8..1633c05183d0 100644
--- a/tests/python/unittest/test_tvmscript_roundtrip.py
+++ b/tests/python/unittest/test_tvmscript_roundtrip.py
@@ -22,6 +22,8 @@
 from tvm import tir
 from tvm.script import tir as T
 
+import numpy as np
+
 
 @tvm.script.ir_module
 class Module1:
@@ -2918,6 +2920,76 @@ def test_opaque_block():
     assert len(root_block.body.body[1].block.iter_vars) == 0
 
 
+@tvm.script.ir_module
+class Module4:
+    # There is an ongoing (python)dict->(c++)Map->(python)dict issue which potentially
+    # changes order of the items in dict after roundtrip due to map not support order
+    # of insertion while dict does. Hence func 'def A(a: T.handle, c: T.handle) -> None'
+    # is commented
+    #
+    #  test:
+    #  d = {"B": 1, "A": 2}
+    #  m = tvm.runtime.convert(d)
+    #  assert d.keys() == m.keys(), f"Order changed from {list(d.keys())} to {list(m.keys())}"
+
+    """
+    @T.prim_func
+    def A(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (10), "int32")
+        C = T.match_buffer(c, (10), "int32")
+        B = T.alloc_buffer((10), "int32")
+
+        K1 = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = A[x] + T.load("int32", K1, x)
+
+        for x in T.serial(0, 10):
+            C[x] = B[x]
+    """
+
+    @T.prim_func
+    def B(a: T.handle, c: T.handle) -> None:
+        A = T.match_buffer(a, (10), "int32")
+        C = T.match_buffer(c, (10), "int32")
+        B = T.alloc_buffer((10), "int32")
+
+        K1 = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = A[x] + T.load("int32", K1, x)
+
+        K2 = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+        for x in T.serial(0, 10):
+            B[x] = B[x] + T.load("int32", K2, x)
+
+        for x in T.serial(0, 10):
+            C[x] = B[x]
+
+
+def test_module_const():
+    func = Module4
+    rt_func = tvm.script.from_source(func.script(show_meta=True))
+    tvm.ir.assert_structural_equal(func, rt_func)
+
+
+@T.prim_func
+def constant(a: T.handle, c: T.handle) -> None:
+    A = T.match_buffer(a, (10), "int32")
+    C = T.match_buffer(c, (10), "int32")
+    B = T.alloc_buffer((10), "int32")
+    K = T.allocate_const([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], "int32", [10])
+    for x in T.serial(0, 10):
+        B[x] = A[x] + T.load("int32", K, x)
+
+    for x in T.serial(0, 10):
+        C[x] = B[x]
+
+
+def test_const():
+    func = constant
+    rt_func = tvm.script.from_source(func.script(show_meta=True))
+    tvm.ir.assert_structural_equal(func, rt_func)
+
+
 @T.prim_func
 def rank0(a: T.handle) -> None:
     A = T.match_buffer(a, (), "float32")
diff --git a/tests/scripts/git_utils.py b/tests/scripts/git_utils.py
index 530abe8029a6..088590713001 100644
--- a/tests/scripts/git_utils.py
+++ b/tests/scripts/git_utils.py
@@ -88,8 +88,10 @@ def parse_remote(remote: str) -> Tuple[str, str]:
         return m.groups()
 
 
-def git(command):
+def git(command, **kwargs):
     command = ["git"] + command
     print("Running", command)
-    proc = subprocess.run(command, stdout=subprocess.PIPE, check=True)
-    return proc.stdout.decode().strip()
+    proc = subprocess.run(command, stdout=subprocess.PIPE, encoding="utf-8", **kwargs)
+    if proc.returncode != 0:
+        raise RuntimeError(f"Command failed {command}:\nstdout:\n{proc.stdout}")
+    return proc.stdout.strip()
diff --git a/tests/scripts/ping_reviewers.py b/tests/scripts/ping_reviewers.py
index 50e67824888c..5c881594cf30 100755
--- a/tests/scripts/ping_reviewers.py
+++ b/tests/scripts/ping_reviewers.py
@@ -32,6 +32,7 @@ def prs_query(user: str, repo: str, cursor: str = None):
     after = ""
     if cursor is not None:
         after = f', before:"{cursor}"'
+    time_keys = "createdAt\nupdatedAt\nlastEditedAt\npublishedAt"
     return f"""
         {{
     repository(name: "{repo}", owner: "{user}") {{
@@ -43,27 +44,29 @@ def prs_query(user: str, repo: str, cursor: str = None):
             number
             url
             body
+            {time_keys}
             isDraft
             author {{
                 login
             }}
             reviews(last:100) {{
                 nodes {{
+                    {time_keys}
+                    bodyText
                     author {{ login }}
                     comments(last:100) {{
                         nodes {{
-                            updatedAt
+                            {time_keys}
                             bodyText
                         }}
                     }}
                 }}
             }}
-            publishedAt
             comments(last:100) {{
                 nodes {{
                     authorAssociation
                     bodyText
-                    updatedAt
+                    {time_keys}
                     author {{
                         login
                     }}
@@ -92,8 +95,25 @@ def find_reviewers(body: str) -> List[str]:
 
 
 def check_pr(pr, wait_time, now):
-    published_at = datetime.datetime.strptime(pr["publishedAt"], GIT_DATE_FORMAT)
-    last_action = published_at
+    last_action = None
+
+    def update_last(new_time, description):
+        if isinstance(new_time, str):
+            new_time = datetime.datetime.strptime(new_time, GIT_DATE_FORMAT)
+        if new_time is None:
+            print(f"  time not found: {description}")
+            return
+        nonlocal last_action
+        if last_action is None or new_time > last_action[0]:
+            last_action = (new_time, description)
+
+    def check_obj(obj, name):
+        update_last(obj["publishedAt"], f"{name} publishedAt: {obj}")
+        update_last(obj["updatedAt"], f"{name} updatedAt: {obj}")
+        update_last(obj["lastEditedAt"], f"{name} lastEditedAt: {obj}")
+        update_last(obj["createdAt"], f"{name} lastEditedAt: {obj}")
+
+    check_obj(pr, "pr")
 
     # GitHub counts comments left as part of a review separately than standalone
     # comments
@@ -101,17 +121,16 @@ def check_pr(pr, wait_time, now):
     review_comments = []
     for review in reviews:
         review_comments += review["comments"]["nodes"]
+        check_obj(review, "review")
 
     # Collate all comments
     comments = pr["comments"]["nodes"] + review_comments
 
     # Find the last date of any comment
     for comment in comments:
-        commented_at = datetime.datetime.strptime(comment["updatedAt"], GIT_DATE_FORMAT)
-        if commented_at > last_action:
-            last_action = commented_at
+        check_obj(comment, "comment")
 
-    time_since_last_action = now - last_action
+    time_since_last_action = now - last_action[0]
 
     # Find reviewers in the PR's body
     pr_body_reviewers = find_reviewers(pr["body"])
@@ -128,17 +147,19 @@ def check_pr(pr, wait_time, now):
 
     if time_since_last_action > wait_time:
         print(
-            "    Pinging reviewers",
+            "  Pinging reviewers",
             reviewers,
             "on",
             pr["url"],
             "since it has been",
             time_since_last_action,
-            "since anything happened on that PR",
+            f"since anything happened on that PR (last action: {last_action[1]})",
         )
         return reviewers
     else:
-        print(f"Not pinging PR {pr['number']}")
+        print(
+            f"  Not pinging PR {pr['number']} since it has been only {time_since_last_action} since the last action: {last_action[1]}"
+        )
 
     return None
 
@@ -206,18 +227,27 @@ def ping_reviewers(pr, reviewers):
         prs = r["data"]["repository"]["pullRequests"]["nodes"]
 
         # Don't look at draft PRs at all
-        prs = [pr for pr in prs if not pr["isDraft"]]
-
-        # Don't look at super old PRs
-        prs = [pr for pr in prs if pr["number"] > cutoff_pr_number]
-
-        # [slow rollout]
-        prs = [pr for pr in prs if pr["author"]["login"] in author_allowlist]
-
-        print(f"Checking {len(prs)} PRs: {[pr['number'] for pr in prs]}")
+        prs_to_check = []
+        for pr in prs:
+            if pr["isDraft"]:
+                print(f"Skipping #{pr['number']} since it's a draft")
+            elif pr["number"] <= cutoff_pr_number:
+                print(
+                    f"Skipping #{pr['number']} since it's too old ({pr['number']} <= {cutoff_pr_number})"
+                )
+            elif pr["author"]["login"] not in author_allowlist:
+                # [slow rollout]
+                print(
+                    f"Skipping #{pr['number']} since author {pr['author']['login']} is not in allowlist: {author_allowlist}"
+                )
+            else:
+                print(f"Checking #{pr['number']}, {author_allowlist}")
+                prs_to_check.append(pr)
+
+        print(f"Summary: Checking {len(prs_to_check)} of {len(prs)} fetched")
 
         # Ping reviewers on each PR in the response if necessary
-        for pr in prs:
+        for pr in prs_to_check:
             print("Checking", pr["url"])
             reviewers = check_pr(pr, wait_time, now)
             if reviewers is not None and not args.dry_run:
diff --git a/tests/scripts/setup-pytest-env.sh b/tests/scripts/setup-pytest-env.sh
index bcd27a16f659..d19533bf93f8 100755
--- a/tests/scripts/setup-pytest-env.sh
+++ b/tests/scripts/setup-pytest-env.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -20,9 +20,9 @@
 set +u
 
 if [[ ! -z $CI_PYTEST_ADD_OPTIONS ]]; then
-    export PYTEST_ADDOPTS="-s -v $CI_PYTEST_ADD_OPTIONS $PYTEST_ADDOPTS"
+    export PYTEST_ADDOPTS="-s -vv $CI_PYTEST_ADD_OPTIONS $PYTEST_ADDOPTS"
 else
-    export PYTEST_ADDOPTS="-s -v $PYTEST_ADDOPTS"
+    export PYTEST_ADDOPTS="-s -vv $PYTEST_ADDOPTS"
 fi
 set -u
 
diff --git a/tests/scripts/task_build.sh b/tests/scripts/task_build.sh
index 845b7153ae20..0623429aa4f5 100755
--- a/tests/scripts/task_build.sh
+++ b/tests/scripts/task_build.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_build_hexagon_api.sh b/tests/scripts/task_build_hexagon_api.sh
new file mode 100755
index 000000000000..89b7545f4d89
--- /dev/null
+++ b/tests/scripts/task_build_hexagon_api.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+# 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.
+
+set -e
+set -u
+
+cd apps/hexagon_api
+rm -rf build
+mkdir -p build
+cd build
+
+output_binary_directory=$(realpath ${PWD}/../../../build/hexagon_api_output)
+rm -rf ${output_binary_directory}
+
+cmake -DANDROID_ABI=arm64-v8a \
+    -DANDROID_PLATFORM=android-28 \
+    -DUSE_ANDROID_TOOLCHAIN="${ANDROID_NDK_HOME}/build/cmake/android.toolchain.cmake" \
+    -DUSE_HEXAGON_ARCH=v68 \
+    -DUSE_HEXAGON_SDK="${HEXAGON_SDK_PATH}" \
+    -DUSE_HEXAGON_TOOLCHAIN="${HEXAGON_TOOLCHAIN}" \
+    -DUSE_OUTPUT_BINARY_DIR="${output_binary_directory}" ..
+
+make -j$(nproc)
diff --git a/tests/scripts/task_ci_setup.sh b/tests/scripts/task_ci_setup.sh
index 33ea484a5268..071de4859c68 100755
--- a/tests/scripts/task_ci_setup.sh
+++ b/tests/scripts/task_ci_setup.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_clean.sh b/tests/scripts/task_clean.sh
index f552fe5ece2a..2ddc0d3db84d 100755
--- a/tests/scripts/task_clean.sh
+++ b/tests/scripts/task_clean.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_arm.sh b/tests/scripts/task_config_build_arm.sh
index d086783319f7..435ea00bc671 100755
--- a/tests/scripts/task_config_build_arm.sh
+++ b/tests/scripts/task_config_build_arm.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_cpu.sh b/tests/scripts/task_config_build_cpu.sh
index 3a4500d0e2f0..e0935ae3c1de 100755
--- a/tests/scripts/task_config_build_cpu.sh
+++ b/tests/scripts/task_config_build_cpu.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_gpu.sh b/tests/scripts/task_config_build_gpu.sh
index 9244d0330bec..c31b4534b548 100755
--- a/tests/scripts/task_config_build_gpu.sh
+++ b/tests/scripts/task_config_build_gpu.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_gpu_other.sh b/tests/scripts/task_config_build_gpu_other.sh
index b5811ea8818e..90aa8d42e02f 100755
--- a/tests/scripts/task_config_build_gpu_other.sh
+++ b/tests/scripts/task_config_build_gpu_other.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_gpu_vulkan.sh b/tests/scripts/task_config_build_gpu_vulkan.sh
index 93adc9667da7..eb26730e89d4 100755
--- a/tests/scripts/task_config_build_gpu_vulkan.sh
+++ b/tests/scripts/task_config_build_gpu_vulkan.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_hexagon.sh b/tests/scripts/task_config_build_hexagon.sh
new file mode 100755
index 000000000000..1e2bddaf7958
--- /dev/null
+++ b/tests/scripts/task_config_build_hexagon.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+# 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.
+
+set -e
+set -u
+
+mkdir -p build
+cd build
+cp ../cmake/config.cmake .
+
+echo set\(USE_SORT ON\) >> config.cmake
+echo set\(USE_RPC ON\) >> config.cmake
+echo set\(USE_MICRO ON\) >> config.cmake
+echo set\(USE_MICRO_STANDALONE_RUNTIME ON\) >> config.cmake
+echo set\(USE_LLVM "${CLANG_LLVM_HOME}/bin/llvm-config"\) >> config.cmake
+echo set\(CMAKE_CXX_COMPILER "${CLANG_LLVM_HOME}/bin/clang++"\) >> config.cmake
+echo set\(USE_HEXAGON_SDK "${HEXAGON_SDK_PATH}"\) >> config.cmake
+echo set\(USE_HEXAGON_ARCH "v68"\) >> config.cmake
+echo set\(USE_HEXAGON_DEVICE "sim"\) >> config.cmake
+echo set\(USE_CCACHE OFF\) >> config.cmake
+echo set\(SUMMARIZE ON\) >> config.cmake
diff --git a/tests/scripts/task_config_build_i386.sh b/tests/scripts/task_config_build_i386.sh
index 1c3bc4c1ae0e..af2291b2ca82 100755
--- a/tests/scripts/task_config_build_i386.sh
+++ b/tests/scripts/task_config_build_i386.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_qemu.sh b/tests/scripts/task_config_build_qemu.sh
index 134de9d3d3a3..0072e713e370 100755
--- a/tests/scripts/task_config_build_qemu.sh
+++ b/tests/scripts/task_config_build_qemu.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_config_build_wasm.sh b/tests/scripts/task_config_build_wasm.sh
index 8636db464e62..36fad5389a9a 100755
--- a/tests/scripts/task_config_build_wasm.sh
+++ b/tests/scripts/task_config_build_wasm.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_cpp_unittest.sh b/tests/scripts/task_cpp_unittest.sh
index 3df7b580d79d..2ff6c627f761 100755
--- a/tests/scripts/task_cpp_unittest.sh
+++ b/tests/scripts/task_cpp_unittest.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_demo_microtvm.sh b/tests/scripts/task_demo_microtvm.sh
index c569f6ce686c..b5c18ec9e757 100755
--- a/tests/scripts/task_demo_microtvm.sh
+++ b/tests/scripts/task_demo_microtvm.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -19,7 +19,7 @@
 set -euxo pipefail
 
 pushd apps/microtvm/zephyr_cmsisnn
-./run_demo.sh
+ timeout 5m ./run_demo.sh
 popd
 
 pushd apps/microtvm/ethosu
@@ -27,6 +27,6 @@ FVP_PATH="/opt/arm/FVP_Corstone_SSE-300_Ethos-U55"
 CMAKE_PATH="/opt/arm/cmake/bin/cmake"
 FREERTOS_PATH="/opt/freertos/FreeRTOSv202112.00"
 
-./run_demo.sh --fvp_path $FVP_PATH --cmake_path $CMAKE_PATH
-./run_demo.sh --fvp_path $FVP_PATH --cmake_path $CMAKE_PATH --freertos_path $FREERTOS_PATH
+ timeout 5m ./run_demo.sh --fvp_path $FVP_PATH --cmake_path $CMAKE_PATH
+ timeout 5m ./run_demo.sh --fvp_path $FVP_PATH --cmake_path $CMAKE_PATH --freertos_path $FREERTOS_PATH
 popd
diff --git a/tests/scripts/task_golang.sh b/tests/scripts/task_golang.sh
index 5b01eeca164c..b79a2a5d9208 100755
--- a/tests/scripts/task_golang.sh
+++ b/tests/scripts/task_golang.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_java_unittest.sh b/tests/scripts/task_java_unittest.sh
index 7818d7d458d6..70ad53d62737 100755
--- a/tests/scripts/task_java_unittest.sh
+++ b/tests/scripts/task_java_unittest.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_lint.sh b/tests/scripts/task_lint.sh
index 3ebee3219260..0a01bea96bb1 100755
--- a/tests/scripts/task_lint.sh
+++ b/tests/scripts/task_lint.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_mypy.sh b/tests/scripts/task_mypy.sh
index 0e20fc22cfb2..2148aeb5e4b4 100755
--- a/tests/scripts/task_mypy.sh
+++ b/tests/scripts/task_mypy.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_arm_compute_library.sh b/tests/scripts/task_python_arm_compute_library.sh
index 7df894d93399..bc0cf176099b 100755
--- a/tests/scripts/task_python_arm_compute_library.sh
+++ b/tests/scripts/task_python_arm_compute_library.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_docs.sh b/tests/scripts/task_python_docs.sh
index 6d2dde58b97d..df3f1abf5f57 100755
--- a/tests/scripts/task_python_docs.sh
+++ b/tests/scripts/task_python_docs.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_ethosn_tests.sh b/tests/scripts/task_python_ethosn_tests.sh
index 6f6c81703282..aeef9e64b251 100755
--- a/tests/scripts/task_python_ethosn_tests.sh
+++ b/tests/scripts/task_python_ethosn_tests.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_frontend.sh b/tests/scripts/task_python_frontend.sh
index 6a46f8845f4e..0352be701717 100755
--- a/tests/scripts/task_python_frontend.sh
+++ b/tests/scripts/task_python_frontend.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -42,9 +42,11 @@ run_pytest cython python-frontend-coreml tests/python/frontend/coreml
 echo "Running relay Tensorflow frontend test..."
 # Note: Tensorflow tests often have memory issues, so invoke each one separately
 TENSORFLOW_TESTS=$(./tests/scripts/pytest_ids.py --folder tests/python/frontend/tensorflow)
+i=0
 for node_id in $TENSORFLOW_TESTS; do
     echo "$node_id"
-    run_pytest cython python-frontend-tensorflow "$node_id"
+    run_pytest cython "python-frontend-tensorflow-$i" "$node_id"
+    i=$((i+1))
 done
 
 echo "Running relay caffe2 frontend test..."
diff --git a/tests/scripts/task_python_frontend_cpu.sh b/tests/scripts/task_python_frontend_cpu.sh
index 208714c64988..f12177dc7f93 100755
--- a/tests/scripts/task_python_frontend_cpu.sh
+++ b/tests/scripts/task_python_frontend_cpu.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/python/contrib/test_hexagon/rpc/__init__.py b/tests/scripts/task_python_hexagon.sh
old mode 100644
new mode 100755
similarity index 74%
rename from tests/python/contrib/test_hexagon/rpc/__init__.py
rename to tests/scripts/task_python_hexagon.sh
index 92e96bf52e8b..8c273f409630
--- a/tests/python/contrib/test_hexagon/rpc/__init__.py
+++ b/tests/scripts/task_python_hexagon.sh
@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
 # 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
@@ -15,4 +16,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-""" Testing infrastructure for Hexagon RPC"""
+set -e
+set -u
+
+source tests/scripts/setup-pytest-env.sh
+
+make cython3
+
+# unset because hardware does not exist in CI.
+unset ANDROID_SERIAL_NUMBER
+run_pytest ctypes python-contrib-hexagon-no-hwardware tests/python/contrib/test_hexagon/test_launcher.py
diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh
index 4992bfa678de..d831481acea6 100755
--- a/tests/scripts/task_python_integration.sh
+++ b/tests/scripts/task_python_integration.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -16,8 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set -e
-set -u
+set -euxo pipefail
 
 source tests/scripts/setup-pytest-env.sh
 export PYTHONPATH=${PYTHONPATH}:${TVM_PATH}/apps/extension/python
@@ -44,22 +43,22 @@ rm -rf lib
 make
 cd ../..
 
-run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-extensions apps/extension/tests
-run_pytest cython ${TVM_INTEGRATION_TESTSUITE_NAME}-extensions apps/extension/tests
+run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-extensions-0 apps/extension/tests
+run_pytest cython ${TVM_INTEGRATION_TESTSUITE_NAME}-extensions-1 apps/extension/tests
 
 # Test dso plugin
 cd apps/dso_plugin_module
 rm -rf lib
 make
 cd ../..
-run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-dso_plugin_module apps/dso_plugin_module
-run_pytest cython ${TVM_INTEGRATION_TESTSUITE_NAME}-dso_plugin_module apps/dso_plugin_module
+run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-dso_plugin_module-0 apps/dso_plugin_module
+run_pytest cython ${TVM_INTEGRATION_TESTSUITE_NAME}-dso_plugin_module-1 apps/dso_plugin_module
 
 # Do not enable TensorFlow op
 # TVM_FFI=cython sh prepare_and_test_tfop_module.sh
 # TVM_FFI=ctypes sh prepare_and_test_tfop_module.sh
 
-run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME} tests/python/integration
+run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-integration tests/python/integration
 
 # Ignoring Arm(R) Ethos(TM)-U NPU tests in the collective to run to run them in parallel in the next step.
 run_pytest ctypes ${TVM_INTEGRATION_TESTSUITE_NAME}-contrib tests/python/contrib --ignore=tests/python/contrib/test_ethosu
diff --git a/tests/scripts/task_python_integration_gpuonly.sh b/tests/scripts/task_python_integration_gpuonly.sh
index cb6bec40c22f..d4d6e26d88f2 100755
--- a/tests/scripts/task_python_integration_gpuonly.sh
+++ b/tests/scripts/task_python_integration_gpuonly.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_integration_i386only.sh b/tests/scripts/task_python_integration_i386only.sh
index 9c378a647e3e..9b0c30b7aeef 100755
--- a/tests/scripts/task_python_integration_i386only.sh
+++ b/tests/scripts/task_python_integration_i386only.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_microtvm.sh b/tests/scripts/task_python_microtvm.sh
index d3fe43a43227..8e6e2bdeaf69 100755
--- a/tests/scripts/task_python_microtvm.sh
+++ b/tests/scripts/task_python_microtvm.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_nightly.sh b/tests/scripts/task_python_nightly.sh
index 16c94dfdad31..225892b02b79 100755
--- a/tests/scripts/task_python_nightly.sh
+++ b/tests/scripts/task_python_nightly.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_topi.sh b/tests/scripts/task_python_topi.sh
index 06046ad76161..1b2bf21cb6a7 100755
--- a/tests/scripts/task_python_topi.sh
+++ b/tests/scripts/task_python_topi.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_unittest.sh b/tests/scripts/task_python_unittest.sh
index 54a36f6dcfd4..4cae8c36cce4 100755
--- a/tests/scripts/task_python_unittest.sh
+++ b/tests/scripts/task_python_unittest.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -16,8 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set -e
-set -u
+set -euxo pipefail
 
 source tests/scripts/setup-pytest-env.sh
 
@@ -31,9 +30,9 @@ if [ -z "${TVM_UNITTEST_TESTSUITE_NAME:-}" ]; then
 fi
 
 # First run minimal test on both ctypes and cython.
-run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME}-platform-minimal-test tests/python/all-platform-minimal-test
-run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME}-platform-minimal-test tests/python/all-platform-minimal-test
+run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME}-platform-minimal-test-0 tests/python/all-platform-minimal-test
+run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME}-platform-minimal-test-1 tests/python/all-platform-minimal-test
 
 # Then run all unittests on both ctypes and cython.
-run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME} tests/python/unittest
-run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME} tests/python/unittest
+run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME}-0 tests/python/unittest
+run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME}-1 tests/python/unittest
diff --git a/tests/scripts/task_python_unittest_gpuonly.sh b/tests/scripts/task_python_unittest_gpuonly.sh
index 54dd085f1817..bb453026e7ea 100755
--- a/tests/scripts/task_python_unittest_gpuonly.sh
+++ b/tests/scripts/task_python_unittest_gpuonly.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
@@ -16,7 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-set -euo pipefail
+set -euxo pipefail
 
 export PYTEST_ADDOPTS="-m gpu ${PYTEST_ADDOPTS:-}"
 
@@ -33,5 +33,5 @@ export TVM_UNITTEST_TESTSUITE_NAME=python-unittest-vulkan
 
 source tests/scripts/setup-pytest-env.sh
 
-run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME} tests/python/unittest/test_target_codegen_vulkan.py
-run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME} tests/python/unittest/test_target_codegen_vulkan.py
+run_pytest ctypes ${TVM_UNITTEST_TESTSUITE_NAME}-0 tests/python/unittest/test_target_codegen_vulkan.py
+run_pytest cython ${TVM_UNITTEST_TESTSUITE_NAME}-1 tests/python/unittest/test_target_codegen_vulkan.py
diff --git a/tests/scripts/task_python_vta_fsim.sh b/tests/scripts/task_python_vta_fsim.sh
index 4074fb888351..6ba78167ac8c 100755
--- a/tests/scripts/task_python_vta_fsim.sh
+++ b/tests/scripts/task_python_vta_fsim.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_python_vta_tsim.sh b/tests/scripts/task_python_vta_tsim.sh
index 4c21f46c5f81..95d1a2cbf001 100755
--- a/tests/scripts/task_python_vta_tsim.sh
+++ b/tests/scripts/task_python_vta_tsim.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_rust.sh b/tests/scripts/task_rust.sh
index 5cc1dc0503f7..0c0e283afd35 100755
--- a/tests/scripts/task_rust.sh
+++ b/tests/scripts/task_rust.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_show_node_info.sh b/tests/scripts/task_show_node_info.sh
index c7b1bb80d7fa..c4da91be5df1 100755
--- a/tests/scripts/task_show_node_info.sh
+++ b/tests/scripts/task_show_node_info.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/task_web_wasm.sh b/tests/scripts/task_web_wasm.sh
index c117bb08c6f4..bdb51a007ef6 100755
--- a/tests/scripts/task_web_wasm.sh
+++ b/tests/scripts/task_web_wasm.sh
@@ -1,4 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 # 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
diff --git a/tests/scripts/update_branch.py b/tests/scripts/update_branch.py
new file mode 100755
index 000000000000..8f2558742217
--- /dev/null
+++ b/tests/scripts/update_branch.py
@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+# 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.
+
+import os
+import json
+import argparse
+import tempfile
+from typing import Any, Dict
+
+from git_utils import git, GitHubRepo, parse_remote
+
+
+_commit_query_fields = """
+    messageHeadline
+    oid
+    statusCheckRollup {
+        contexts(last:100) {
+            nodes {
+                ... on CheckRun {
+                    conclusion
+                    status
+                    name
+                    checkSuite {
+                        workflowRun {
+                            workflow {
+                                name
+                            }
+                        }
+                    }
+                }
+                ... on StatusContext {
+                    context
+                    state
+                }
+            }
+        }
+    }
+"""
+
+
+def commits_query(user: str, repo: str, cursor: str = None):
+    """
+    Create a GraphQL query to find the last N commits along with their statuses
+    and some metadata (paginated after 'cursor')
+    """
+    after = ""
+    if cursor is not None:
+        after = f', after:"{cursor}"'
+
+    return f"""
+    {{
+    repository(name: "{repo}", owner: "{user}") {{
+        defaultBranchRef {{
+        target {{
+            ... on Commit {{
+            history(first: 15{after}) {{
+                edges {{ cursor }}
+                nodes {{
+                    {_commit_query_fields}
+                }}
+            }}
+            }}
+        }}
+        }}
+    }}
+    }}
+    """
+
+
+def commit_passed_ci(commit: Dict[str, Any]) -> bool:
+    """
+    Returns true if all of a commit's statuses are SUCCESS
+    """
+    statuses = commit["statusCheckRollup"]["contexts"]["nodes"]
+
+    # GitHub Actions statuses are different from external GitHub statuses, so
+    # unify them into 1 representation
+    # https://docs.github.com/en/developers/webhooks-and-events/webhooks/webhook-events-and-payloads
+    unified_statuses = []
+    for status in statuses:
+        if "context" in status:
+            # Parse non-GHA status
+            unified_statuses.append((status["context"], status["state"] == "SUCCESS"))
+        else:
+            # Parse GitHub Actions item
+            workflow = status["checkSuite"]["workflowRun"]["workflow"]["name"]
+            name = f"{workflow} / {status['name']}"
+            unified_statuses.append((name, status["conclusion"] == "SUCCESS"))
+
+    print(f"Statuses on {commit['oid']}:", json.dumps(unified_statuses, indent=2))
+
+    # Assert that specific jobs are present in the commit statuses (i.e. don't
+    # approve if CI was broken and didn't schedule a job)
+    expected_jobs = {"tvm-ci/branch"}
+    job_names = {name for name, status in unified_statuses}
+    for job in expected_jobs:
+        if job not in job_names:
+            # Did not find expected job name
+            return False
+
+    passed_ci = all(status for name, status in unified_statuses)
+    return passed_ci
+
+
+def update_branch(user: str, repo: str, sha: str, branch_name: str) -> None:
+    git(["fetch", "origin", sha])
+    git(["reset", "--hard", "FETCH_HEAD"])
+    try:
+        git(["branch", "-D", branch_name])
+    except RuntimeError:
+        # Ignore failures (i.e. the branch did not exist in the first place)
+        pass
+    git(["checkout", "-b", branch_name])
+
+    # Create and push the branch
+    git(["push", "origin", "--force", branch_name])
+    print(f"Pushed branch {branch_name} with commit {sha}")
+
+
+if __name__ == "__main__":
+    help = "Push the a branch to the last commit that passed all CI runs"
+    parser = argparse.ArgumentParser(description=help)
+    parser.add_argument("--remote", default="origin", help="ssh remote to parse")
+    parser.add_argument("--dry-run", action="store_true", help="don't submit to GitHub")
+    parser.add_argument("--branch", default="last-successful", help="branch name")
+    parser.add_argument(
+        "--testonly-json", help="(testing) data to use instead of fetching from GitHub"
+    )
+    args = parser.parse_args()
+
+    remote = git(["config", "--get", f"remote.{args.remote}.url"])
+    user, repo = parse_remote(remote)
+    # TODO: Remove this before landing
+    user, repo = ("apache", "tvm")
+
+    if args.testonly_json:
+        r = json.loads(args.testonly_json)
+    else:
+        github = GitHubRepo(token=os.environ["GITHUB_TOKEN"], user=user, repo=repo)
+        q = commits_query(user, repo)
+        r = github.graphql(q)
+
+    commits = r["data"]["repository"]["defaultBranchRef"]["target"]["history"]["nodes"]
+
+    # Limit GraphQL pagination
+    MAX_COMMITS_TO_CHECK = 50
+    i = 0
+
+    while i < MAX_COMMITS_TO_CHECK:
+        # Check each commit
+        for commit in commits:
+            if commit_passed_ci(commit):
+                print(f"Found last good commit: {commit['oid']}: {commit['messageHeadline']}")
+                if not args.dry_run:
+                    update_branch(
+                        user=user,
+                        repo=repo,
+                        sha=commit["oid"],
+                        branch_name=args.branch,
+                    )
+                # Nothing to do after updating the branch, exit early
+                exit(0)
+
+        # No good commit found, proceed to next page of results
+        edges = r["data"]["repository"]["defaultBranchRef"]["target"]["history"]["edges"]
+        if len(edges) == 0:
+            break
+        else:
+            q = commits_query(user, repo, cursor=edges[-1]["cursor"])
+            r = github.graphql(q)
+            commits = r["data"]["repository"]["defaultBranchRef"]["target"]["history"]["nodes"]
+
+        # Backstop to prevent looking through all the past commits
+        i += len(commits)
+
+    print(f"No good commits found in the last {len(commits)} commits")
+    exit(1)
diff --git a/vta/python/vta/libinfo.py b/vta/python/vta/libinfo.py
index 8b300c3acb97..65ea7083a255 100644
--- a/vta/python/vta/libinfo.py
+++ b/vta/python/vta/libinfo.py
@@ -59,16 +59,17 @@ def find_libvta(lib_vta, optional=False):
         Enable error check
     """
     curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
-    lib_search = [
-        os.path.join(
+    tvm_library_path = os.environ.get("TVM_LIBRARY_PATH", None)
+    if tvm_library_path is None:
+        tvm_library_path = os.path.join(
             curr_path,
-            "..",
-            "..",
-            "..",
+            os.pardir,
+            os.pardir,
+            os.pardir,
             "build",
         )
-    ]
-    lib_search += [os.path.join(get_vta_hw_path(), "build")]
+
+    lib_search = [tvm_library_path, os.path.join(get_vta_hw_path(), "build")]
     lib_name = _get_lib_name(lib_vta)
     lib_path = [os.path.join(x, lib_name) for x in lib_search]
     lib_found = [x for x in lib_path if os.path.exists(x)]