[AOT] Introducing AOT in TVM (apache#7785)

* [AOT] Introducing AOT in TVM

This change adds the code generation and minimal runtime API to use the
Ahead Of Time (AOT) compilation flow. The main logic is contained in:

- src/relay/backend/aot_codegen.cc

Which produces a TIR PrimFunc traversing the Relay graph

The runtime interface (authored by @Mousius) leaves a gap for future
iterations using platform-specific features from RTOS.

Currently AOT runs successfully on x86 in a host OS, running these
tests on micro is coming soon.

This PR is based on the RFC described here: https://discuss.tvm.apache.org/t/implementing-aot-in-tvm/9206

Co-authored-by: Christopher Sidebottom <Christopher.Sidebottom@arm.com>
Change-Id: I9f731c953231f129e1472298915dddc01788efd7

* Rebasing 2

Change-Id: Ia0a533a49960f1cb4bf3c3833511e539cf7c459f

* Applying comments/refactoring

Change-Id: Iea1832355f8b1d4c921d02c6b4ceec7db3a681c1

* Fixing comments + refactoring - 2

Change-Id: I7200cc17b297e42bf67dcdef6f643e86991ca0a8

* fix linting

Change-Id: Iba6544ac7101595696b352b8702345cf916625f6

* fix linting - 2

Change-Id: I7f80d16005f2c621d37a9aae2cbbd61df0277cbe

* fix linting - 3

Change-Id: I7a1ba40afeea46d5f122563a20cd4b2f08751a1e

* fix tests

Change-Id: I1297ccc54dd6d93647f421e0beb226f410bf73f5

* Addressing comments - 3

Change-Id: Id25d1382c30d6d0a0013b5e8986fb8cd886666dc

* Addressing comments - 4

Change-Id: Ibe29676abe3b75161b5a0903e007118a8318d862

* fix tests - 2

Change-Id: I2117f9d4392bfd87102ecbef0993c8b320f479a0

* fix tests - 3

Change-Id: Ic0373543b0f9a54dbd4dc32d428272f7293200ba

* fix tests - 4

Change-Id: I8a6f229c9a3a9e169779c8d49cbfa3f473348b1f

* Addressing comments - 5

Change-Id: Ib9ccd07c87392034a21b2eb70955d0b091b780f1

* fix tests - 5

Change-Id: I4b13c3b548ced414991e83072e9e6fc99b64f939

* fix tests - 6

Change-Id: Id5af1f778ae25bc60849cc054a605181c1b7a765

* addressing comments - 6

Change-Id: Id94a2bbcaae891f9498d41be538f13a952f55b81

* fix linting - 4

Change-Id: I371a0aa5b81824b5a3a1278fac22ace57832027a

* add missing file

Change-Id: If359bef96dd0773ead4f75f0d9f5234276347e2d

* fix build

Change-Id: I73fc1feb6f7b5d454a528e3289228484dc2b07d5

* addressing comments - 7

Change-Id: I7f908f3908ffc77e408391f62edcc06f2600c6c2

* addressing comments - 8

Change-Id: I90bced4e18259a6d42e6a406d93958e204f3859e

* rebasing

Change-Id: Id28751b069bd046f00faee301b2b446b2ea4fab8

* Addressing comments - 9

Change-Id: I06c9f280de0a9bf0ca5545bbbbfcc70cb66831b3

* fix tests - 7

Change-Id: I739f29779862f05def36e5f3e0722019596d17f8

* Addressing comments - 9

Change-Id: Ie736f40a5225f4e56e79006753d7732127da5408

* Applying comments + fixing tests

Change-Id: I83e16068b93aaccc7a86b79d42f13328bc76b53d

* Applying comments - 10

Change-Id: I443d72f53913849f3c28fd6e416162d1ca99e647

* Addressing comments - 11

Change-Id: I7fefbd0076949b9c38d0abbf2759ebf1502de330

* Addressing comments - 11

Change-Id: Iad028144d7b394b2dd2fce41a35ca689d1680200

* fix tests - 7

Change-Id: I14286e665dcdba1e9bc10bb5a27dd6ced50372b0

* fixing tests -8

Change-Id: I7b4c966da9680870ceda1704c749ee3bdc751926

* fixing tests - 9

Change-Id: Icf62128a604998ed1b7d5af4cbeadf7d39196d0b

Co-authored-by: Christopher Sidebottom <Christopher.Sidebottom@arm.com>

apps/bundle_deploy/bundle.c

@@ -23,8 +23,8 @@
 #include <tvm/runtime/c_runtime_api.h>
 #include <tvm/runtime/crt/crt.h>
 #include <tvm/runtime/crt/graph_executor.h>
-#include <tvm/runtime/crt/memory.h>
 #include <tvm/runtime/crt/packed_func.h>
+#include <tvm/runtime/crt/page_allocator.h>
 
 #ifdef ENABLE_TVM_ABORT_BACKTRACE
 #include "backtrace.h"
@@ -64,8 +64,8 @@ TVM_DLL void* tvm_runtime_create(const char* json_data, const char* params_data,
   dev.device_id = device_id;
 
   // declare pointers
-  TVM_CCALL(MemoryManagerCreate(&g_memory_manager, g_crt_memory, sizeof(g_crt_memory),
-                                CRT_MEMORY_PAGE_SIZE_LOG2));
+  TVM_CCALL(PageMemoryManagerCreate(&g_memory_manager, g_crt_memory, sizeof(g_crt_memory),
+                                    CRT_MEMORY_PAGE_SIZE_LOG2));
   TVM_CCALL(TVMInitializeRuntime());
   TVMPackedFunc pf;
   TVMArgs args = TVMArgs_Create(NULL, NULL, 0);

apps/bundle_deploy/bundle_static.c

@@ -22,8 +22,8 @@
 #include <stdlib.h>
 #include <tvm/runtime/crt/crt.h>
 #include <tvm/runtime/crt/graph_executor.h>
-#include <tvm/runtime/crt/memory.h>
 #include <tvm/runtime/crt/packed_func.h>
+#include <tvm/runtime/crt/page_allocator.h>
 #include <unistd.h>
 
 #ifdef ENABLE_TVM_PLATFORM_ABORT_BACKTRACE
@@ -64,8 +64,8 @@ TVM_DLL void* tvm_runtime_create(const char* json_data, const char* params_data,
   dev.device_id = device_id;
 
   // get pointers
-  TVM_CCALL(MemoryManagerCreate(&g_memory_manager, g_crt_memory, sizeof(g_crt_memory),
-                                CRT_MEMORY_PAGE_SIZE_LOG2));
+  TVM_CCALL(PageMemoryManagerCreate(&g_memory_manager, g_crt_memory, sizeof(g_crt_memory),
+                                    CRT_MEMORY_PAGE_SIZE_LOG2));
   TVM_CCALL(TVMInitializeRuntime());
   TVMPackedFunc pf;
   TVMArgs args = TVMArgs_Create(NULL, NULL, 0);

cmake/modules/StandaloneCrt.cmake

@@ -44,6 +44,7 @@ if(USE_MICRO)
          "src/runtime/crt/include *.h -> include"
          "src/runtime/crt/common *.c -> src/runtime/crt/common"
          "src/runtime/crt/graph_executor *.c -> src/runtime/crt/graph_executor"
+         "src/runtime/crt/aot_executor *.c -> src/runtime/crt/aot_executor"
          "src/runtime/crt/graph_executor_module *.c -> src/runtime/crt/graph_executor_module"
          "src/runtime/crt/host crt_config.h -> template/host"
          "src/runtime/crt/host *.cc -> template/host"
@@ -97,7 +98,7 @@ if(USE_MICRO)
     set(make_quiet )
     endif(${VERBOSE})
 
-    list(APPEND crt_libraries memory graph_executor utvm_rpc_server utvm_rpc_common common)  # NOTE: listed in link order.
+    list(APPEND crt_libraries memory graph_executor aot_executor utvm_rpc_server utvm_rpc_common common)  # NOTE: listed in link order.
     foreach(crt_lib_name IN LISTS crt_libraries)
       list(APPEND crt_library_paths "host_standalone_crt/lib${crt_lib_name}.a")
     endforeach()

include/tvm/runtime/crt/memory.h → include/tvm/runtime/crt/page_allocator.h

@@ -18,12 +18,12 @@
  */
 
 /*!
- * \file tvm/runtime/crt/memory.h
+ * \file tvm/runtime/crt/page_allocator.h
  * \brief An implementation of a dynamic memory allocator for microcontrollers.
  */
 
-#ifndef TVM_RUNTIME_CRT_MEMORY_H_
-#define TVM_RUNTIME_CRT_MEMORY_H_
+#ifndef TVM_RUNTIME_CRT_PAGE_ALLOCATOR_H_
+#define TVM_RUNTIME_CRT_PAGE_ALLOCATOR_H_
 
 #ifdef __cplusplus
 extern "C" {
@@ -72,11 +72,11 @@ struct MemoryManagerInterface {
  * \param page_size_bytes_log2 log2 of the page size, in bytes.
  * \return kTvmErrorNoError on success.
  */
-tvm_crt_error_t MemoryManagerCreate(MemoryManagerInterface** manager, uint8_t* memory_pool,
-                                    size_t memory_pool_size_bytes, size_t page_size_bytes_log2);
+tvm_crt_error_t PageMemoryManagerCreate(MemoryManagerInterface** manager, uint8_t* memory_pool,
+                                        size_t memory_pool_size_bytes, size_t page_size_bytes_log2);
 
 #ifdef __cplusplus
 }  // extern "C"
 #endif
 
-#endif  // TVM_RUNTIME_CRT_MEMORY_H_
+#endif  // TVM_RUNTIME_CRT_PAGE_ALLOCATOR_H_

include/tvm/runtime/crt/stack_allocator.h

@@ -0,0 +1,60 @@
+/*
+ * 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.
+ */
+
+// LINT_C_FILE
+#ifndef TVM_RUNTIME_CRT_STACK_ALLOCATOR_H_
+#define TVM_RUNTIME_CRT_STACK_ALLOCATOR_H_
+#include <stddef.h>
+#include <stdint.h>
+
+#include "crt_config.h"
+#include "error_codes.h"
+
+#define STACK_ALLOCATOR_TAG 0xabcd1234
+#define STACK_ALLOCATOR_TAG_SIZE_BYTES 4
+
+/*! Memory alignment for allocator */
+
+#ifndef TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES
+#define TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES 16
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef struct {
+  uint8_t* next_alloc;    // Pointer to the next block of TVM_RUNTIME_ALLOC_ALIGNMENT_BYTES
+  uint8_t* workspace;     // Pointer to start of the workspace
+  size_t workspace_size;  // Total number of bytes in the workspace
+} tvm_workspace_t;
+
+tvm_crt_error_t StackMemoryManager_Init(tvm_workspace_t* tvm_runtime_workspace,
+                                        uint8_t* g_aot_memory, size_t workspace_size);
+
+tvm_crt_error_t StackMemoryManager_Allocate(tvm_workspace_t* tvm_runtime_workspace, int32_t nbytes,
+                                            void**);
+
+tvm_crt_error_t StackMemoryManager_Free(tvm_workspace_t* tvm_runtime_workspace, void* ptr);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif  // TVM_RUNTIME_CRT_STACK_ALLOCATOR_H_

include/tvm/runtime/executor_info.h

@@ -0,0 +1,39 @@
+/*
+ * 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 executor_info.h
+ * \brief Executor information
+ */
+#ifndef TVM_RUNTIME_EXECUTOR_INFO_H_
+#define TVM_RUNTIME_EXECUTOR_INFO_H_
+
+namespace tvm {
+namespace runtime {
+
+/*! \brief Value used to indicate the graph executor. */
+static constexpr const char* kTvmExecutorGraph = "graph";
+
+/*! \brief Value used to indicate the aot executor. */
+static constexpr const char* kTvmExecutorAot = "aot";
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_EXECUTOR_INFO_H_

include/tvm/runtime/module.h

@@ -230,6 +230,8 @@ 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 */
+constexpr const char* tvm_run_func_prefix = "tvm__run_func";
 }  // namespace symbol
 
 // implementations of inline functions.

include/tvm/tir/builtin.h

@@ -346,6 +346,18 @@ TVM_DLL const Op& tvm_stack_make_array();
  */
 TVM_DLL const Op& tvm_call_packed();
 
+/*!
+ * \brief See pesudo code
+ *
+ * return_type tvm_call_packed(fname, TVMValue* args) {
+ * 	   int ret_code;
+ *     TVMValue ret_value;
+ *     (*fname)(args, type_code_of(args), len(args), &ret_value, &ret_code);
+ *     return cast(return_type, ret_value.v_return_type);
+ *  }
+ */
+TVM_DLL const Op& tvm_call_cpacked();
+
 /*!
  * \brief See pesudo code
  *
@@ -392,6 +404,21 @@ TVM_DLL const Op& tvm_thread_context();
  */
 TVM_DLL const Op& tvm_call_packed_lowered();
 
+/*!
+ * \brief Lowered version of call c-packed, the space of value and
+ *  type codes are explicitly allocated.
+ *
+ *  int tvm_call_packed_lowered(fname,
+ *                              TVMValue* value_stack,
+ *                              int* tcode_stack,
+ *                              int begin,
+ *                              int end) {
+ *     fname(TVMArgs(value_stack[begin:end], tcode_stack[begin:end]),
+ *                   TVMRetValue(value_stack + end, tcode_stack + end));
+ *  }
+ */
+TVM_DLL const Op& tvm_call_cpacked_lowered();
+
 /*!
  * \brief Lowered version of trace intrinsic, the space of value and
  *  type codes are explicitly allocated. The return value is the

python/tvm/driver/tvmc/model.py

@@ -52,7 +52,7 @@
 import tvm.contrib.cc
 from tvm import relay
 from tvm.contrib import utils
-from tvm.relay.backend.graph_executor_factory import GraphExecutorFactoryModule
+from tvm.relay.backend.executor_factory import GraphExecutorFactoryModule
 
 from .common import TVMCException
 
@@ -220,7 +220,7 @@ def export_package(
         self.lib_path = path_lib
 
         with open(temp.relpath(graph_name), "w") as graph_file:
-            graph_file.write(executor_factory.get_json())
+            graph_file.write(executor_factory.get_graph_json())
 
         with open(temp.relpath(param_name), "wb") as params_file:
             params_file.write(relay.save_param_dict(executor_factory.get_params()))

python/tvm/micro/model_library_format.py

@@ -24,7 +24,7 @@
 import tarfile
 
 from ..contrib import utils
-from ..relay.backend import graph_executor_factory
+from ..relay.backend import executor_factory
 from ..relay import param_dict
 
 
@@ -117,7 +117,7 @@ def _build_memory_map(graph_json):
     return memory_map
 
 
-def export_model_library_format(mod: graph_executor_factory.GraphExecutorFactoryModule, file_name):
+def export_model_library_format(mod: executor_factory.ExecutorFactoryModule, file_name):
     """Export the build artifact in Model Library Format.
 
     This function creates a .tar archive containing the build artifacts in a standardized
@@ -126,20 +126,25 @@ def export_model_library_format(mod: graph_executor_factory.GraphExecutorFactory
 
     Parameters
     ----------
-    mod : tvm.relay.backend.graph_executor_factory.GraphExecutorFactoryModule
+    mod : tvm.relay.backend.executor_factory.ExecutorFactoryModule
         The return value of tvm.relay.build, which will be exported into Model Library Format.
     file_name : str
         Path to the .tar archive to generate.
     """
     tempdir = utils.tempdir()
+    is_aot = isinstance(mod, executor_factory.AOTExecutorFactoryModule)
+    memory_map = [] if is_aot else _build_memory_map(mod.get_executor_config())
+    runtime = ["aot"] if is_aot else ["graph"]
+
     metadata = {
         "version": 1,
         "model_name": mod.libmod_name,
         "export_datetime": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%SZ"),
-        "memory": _build_memory_map(mod.graph_json),
+        "memory": memory_map,
         "target": {int(k): str(v) for k, v in mod.target.items()},
-        "runtimes": ["graph"],
+        "runtimes": runtime,
     }
+
     with open(tempdir.relpath("metadata.json"), "w") as json_f:
         json.dump(metadata, json_f, indent=2, sort_keys=True)
 
@@ -156,10 +161,11 @@ def export_model_library_format(mod: graph_executor_factory.GraphExecutorFactory
     with open(tempdir.relpath("relay.txt"), "w") as f:
         f.write(str(mod.ir_mod))
 
-    graph_config_dir_path = tempdir.relpath(os.path.join("runtime-config", "graph"))
-    os.makedirs(graph_config_dir_path)
-    with open(os.path.join(graph_config_dir_path, "graph.json"), "w") as f:
-        f.write(mod.graph_json)
+    if not is_aot:
+        graph_config_dir_path = tempdir.relpath(os.path.join("runtime-config", "graph"))
+        os.makedirs(graph_config_dir_path)
+        with open(os.path.join(graph_config_dir_path, "graph.json"), "w") as f:
+            f.write(mod.get_executor_config())
 
     with tarfile.open(file_name, "w") as tar_f: