[AOT] Introducing AOT in TVM

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: