[aot] Revert "Load AOT module from memory (taichi-dev#6692)" (taichi-…

…dev#6704)

This reverts commit 60a5eee.

Issue: #

### Brief Summary

c_api/include/taichi/cpp/taichi.hpp

@@ -797,14 +797,6 @@ class Runtime {
     return load_aot_module(path.c_str());
   }
 
-  AotModule create_aot_module(const void *tcm, size_t size) {
-    TiAotModule aot_module = ti_create_aot_module(runtime_, tcm, size);
-    return AotModule(runtime_, aot_module, true);
-  }
-  AotModule create_aot_module(const std::vector<uint8_t> &tcm) {
-    return create_aot_module(tcm.data(), tcm.size());
-  }
-
   void copy_memory_device_to_device(const TiMemorySlice &dst_memory,
                                     const TiMemorySlice &src_memory) {
     ti_copy_memory_device_to_device(runtime_, &dst_memory, &src_memory);

c_api/include/taichi/taichi_core.h

@@ -995,11 +995,6 @@ TI_DLL_EXPORT void TI_API_CALL ti_wait(TiRuntime runtime);
 TI_DLL_EXPORT TiAotModule TI_API_CALL
 ti_load_aot_module(TiRuntime runtime, const char *module_path);
 
-// Function `ti_create_aot_module`
-TI_DLL_EXPORT TiAotModule TI_API_CALL ti_create_aot_module(TiRuntime runtime,
-                                                           const void *tcm,
-                                                           uint64_t size);
-
 // Function `ti_destroy_aot_module`
 //
 // Destroys a loaded AOT module and releases all related resources.

c_api/src/taichi_core_impl.cpp

@@ -4,7 +4,6 @@
 #include "taichi_llvm_impl.h"
 #include "taichi/program/ndarray.h"
 #include "taichi/program/texture.h"
-#include "taichi/common/virtual_dir.h"
 
 struct ErrorCache {
   TiError error{TI_ERROR_SUCCESS};
@@ -22,7 +21,7 @@ const char *describe_error(TiError error) {
     case TI_ERROR_NOT_SUPPORTED:
       return "not supported";
     case TI_ERROR_CORRUPTED_DATA:
-      return "corrupted data";
+      return "path not found";
     case TI_ERROR_NAME_NOT_FOUND:
       return "name not found";
     case TI_ERROR_INVALID_ARGUMENT:
@@ -493,8 +492,6 @@ TiAotModule ti_load_aot_module(TiRuntime runtime, const char *module_path) {
   TI_CAPI_ARGUMENT_NULL_RV(runtime);
   TI_CAPI_ARGUMENT_NULL_RV(module_path);
 
-  // (penguinliong) Should call `create_aot_module` directly after all backends
-  // adapted to it.
   TiAotModule aot_module = ((Runtime *)runtime)->load_aot_module(module_path);
 
   if (aot_module == TI_NULL_HANDLE) {
@@ -505,26 +502,6 @@ TiAotModule ti_load_aot_module(TiRuntime runtime, const char *module_path) {
   TI_CAPI_TRY_CATCH_END();
   return out;
 }
-TiAotModule ti_create_aot_module(TiRuntime runtime,
-                                 const void *tcm,
-                                 uint64_t size) {
-  TiAotModule out = TI_NULL_HANDLE;
-  TI_CAPI_TRY_CATCH_BEGIN();
-  TI_CAPI_ARGUMENT_NULL_RV(runtime);
-  TI_CAPI_ARGUMENT_NULL_RV(tcm);
-
-  auto dir = taichi::io::VirtualDir::from_zip(tcm, size);
-  if (dir == TI_NULL_HANDLE) {
-    ti_set_last_error(TI_ERROR_CORRUPTED_DATA, "tcm");
-    return TI_NULL_HANDLE;
-  }
-
-  Error err = ((Runtime *)runtime)->create_aot_module(dir.get(), out);
-  err.set_last_error();
-
-  TI_CAPI_TRY_CATCH_END();
-  return out;
-}
 void ti_destroy_aot_module(TiAotModule aot_module) {
   TI_CAPI_TRY_CATCH_BEGIN();
   TI_CAPI_ARGUMENT_NULL(aot_module);

c_api/src/taichi_core_impl.h

@@ -31,7 +31,6 @@
 #include "taichi/rhi/device.h"
 #include "taichi/aot/graph_data.h"
 #include "taichi/aot/module_loader.h"
-#include "taichi/common/virtual_dir.h"
 
 #define TI_CAPI_NOT_SUPPORTED(x) ti_set_last_error(TI_ERROR_NOT_SUPPORTED, #x);
 #define TI_CAPI_NOT_SUPPORTED_IF(x)                \
@@ -90,28 +89,6 @@
     ti_set_last_error(TI_ERROR_INVALID_STATE, "c++ exception"); \
   }
 
-struct Error {
-  TiError error;
-  std::string message;
-
-  Error(TiError error, const std::string &message)
-      : error(error), message(message) {
-  }
-  Error() : error(TI_ERROR_SUCCESS), message() {
-  }
-  Error(const Error &) = delete;
-  Error(Error &&) = default;
-  Error &operator=(const Error &) = delete;
-  Error &operator=(Error &&) = default;
-
-  // Set this error as the last error if it's not `TI_ERROR_SUCCESS`.
-  inline void set_last_error() const {
-    if (error != TI_ERROR_SUCCESS) {
-      ti_set_last_error(error, message.c_str());
-    }
-  }
-};
-
 class Runtime {
  protected:
   // 32 is a magic number in `taichi/inc/constants.h`.
@@ -127,18 +104,7 @@ class Runtime {
 
   virtual taichi::lang::Device &get() = 0;
 
-  [[deprecated("create_aot_module")]] virtual TiAotModule load_aot_module(
-      const char *module_path) {
-    auto dir = taichi::io::VirtualDir::open(module_path);
-    TiAotModule aot_module = TI_NULL_HANDLE;
-    Error err = create_aot_module(dir.get(), aot_module);
-    err.set_last_error();
-    return aot_module;
-  }
-  virtual Error create_aot_module(const taichi::io::VirtualDir *dir,
-                                  TiAotModule &out) {
-    TI_NOT_IMPLEMENTED
-  }
+  virtual TiAotModule load_aot_module(const char *module_path) = 0;
   virtual TiMemory allocate_memory(
       const taichi::lang::Device::AllocParams &params);
   virtual void free_memory(TiMemory devmem);

c_api/src/taichi_gfx_impl.cpp

@@ -4,15 +4,14 @@
 GfxRuntime::GfxRuntime(taichi::Arch arch) : Runtime(arch) {
 }
 
-Error GfxRuntime::create_aot_module(const taichi::io::VirtualDir *dir,
-                                    TiAotModule &out) {
+TiAotModule GfxRuntime::load_aot_module(const char *module_path) {
   taichi::lang::gfx::AotModuleParams params{};
-  params.dir = dir;
+  params.module_path = module_path;
   params.runtime = &get_gfx_runtime();
   std::unique_ptr<taichi::lang::aot::Module> aot_module =
       taichi::lang::aot::Module::load(arch, params);
   if (aot_module->is_corrupted()) {
-    return Error(TI_ERROR_CORRUPTED_DATA, "aot_module");
+    return TI_NULL_HANDLE;
   }
 
   const taichi::lang::DeviceCapabilityConfig &current_devcaps =
@@ -22,16 +21,15 @@ Error GfxRuntime::create_aot_module(const taichi::io::VirtualDir *dir,
   for (const auto &pair : required_devcaps.devcaps) {
     uint32_t current_version = current_devcaps.get(pair.first);
     uint32_t required_version = pair.second;
-    if (current_version != required_version) {
-      return Error(TI_ERROR_INCOMPATIBLE_MODULE,
-                   taichi::lang::to_string(pair.first).c_str());
+    if (current_version == required_version) {
+      ti_set_last_error(TI_ERROR_INCOMPATIBLE_MODULE,
+                        taichi::lang::to_string(pair.first).c_str());
     }
   }
 
   size_t root_size = aot_module->get_root_size();
   params.runtime->add_root_buffer(root_size);
-  out = (TiAotModule) new AotModule(*this, std::move(aot_module));
-  return Error();
+  return (TiAotModule)(new AotModule(*this, std::move(aot_module)));
 }
 void GfxRuntime::buffer_copy(const taichi::lang::DevicePtr &dst,
                              const taichi::lang::DevicePtr &src,

c_api/src/taichi_gfx_impl.h

@@ -2,7 +2,6 @@
 
 #include "taichi_core_impl.h"
 #include "taichi/runtime/gfx/runtime.h"
-#include "taichi/common/virtual_dir.h"
 
 class GfxRuntime;
 
@@ -11,8 +10,7 @@ class GfxRuntime : public Runtime {
   GfxRuntime(taichi::Arch arch);
   virtual taichi::lang::gfx::GfxRuntime &get_gfx_runtime() = 0;
 
-  virtual Error create_aot_module(const taichi::io::VirtualDir *dir,
-                                  TiAotModule &out) override final;
+  virtual TiAotModule load_aot_module(const char *module_path) override final;
   virtual void buffer_copy(const taichi::lang::DevicePtr &dst,
                            const taichi::lang::DevicePtr &src,
                            size_t size) override final;

c_api/taichi.json

@@ -866,27 +866,6 @@
                         }
                     ]
                 },
-                {
-                    "name": "create_aot_module",
-                    "type": "function",
-                    "parameters": [
-                        {
-                            "name": "@return",
-                            "type": "handle.aot_module"
-                        },
-                        {
-                            "type": "handle.runtime"
-                        },
-                        {
-                            "name": "tcm",
-                            "type": "const void*"
-                        },
-                        {
-                            "name": "size",
-                            "type": "uint64_t"
-                        }
-                    ]
-                },
                 {
                     "name": "destroy_aot_module",
                     "type": "function",

c_api/tests/c_api_interface_test.cpp

@@ -165,40 +165,3 @@ TEST_F(CapiTest, TestLoadTcmAotModule) {
     }
   }
 }
-
-TEST_F(CapiTest, TestCreateTcmAotModule) {
-  if (capi::utils::is_vulkan_available()) {
-    const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
-
-    std::stringstream aot_mod_ss;
-    aot_mod_ss << folder_dir << "/module.tcm";
-
-    std::vector<uint8_t> tcm;
-    {
-      std::fstream f(aot_mod_ss.str(),
-                     std::ios::in | std::ios::binary | std::ios::ate);
-      TI_ASSERT(f.is_open());
-      tcm.resize(f.tellg());
-      f.seekg(std::ios::beg);
-      f.read((char *)tcm.data(), tcm.size());
-    }
-
-    {
-      // Vulkan Runtime
-      TiArch arch = TiArch::TI_ARCH_VULKAN;
-      ti::Runtime runtime(arch);
-      ti::AotModule aot_mod = runtime.create_aot_module(tcm);
-      ti::Kernel run = aot_mod.get_kernel("run");
-      ti::NdArray<int32_t> arr =
-          runtime.allocate_ndarray<int32_t>({16}, {}, true);
-      run[0] = arr;
-      run.launch();
-      runtime.wait();
-      std::vector<int32_t> data(16);
-      arr.read(data);
-      for (int32_t i = 0; i < 16; ++i) {
-        TI_ASSERT(data.at(i) == i);
-      }
-    }
-  }
-}

taichi/common/virtual_dir.cpp

@@ -60,11 +60,9 @@ struct ZipArchiveVirtualDir : public VirtualDir {
     f.seekg(std::ios::beg);
     f.read((char *)archive_data.data(), archive_data.size());
 
-    return from_zip(archive_data.data(), archive_data.size());
-  }
-  static std::unique_ptr<VirtualDir> from_zip(const void *data, size_t size) {
     zip::ZipArchive archive;
-    bool succ = zip::ZipArchive::try_from_bytes(data, size, archive);
+    bool succ = zip::ZipArchive::try_from_bytes(archive_data.data(),
+                                                archive_data.size(), archive);
     if (!succ) {
       return nullptr;
     }
@@ -123,14 +121,5 @@ std::unique_ptr<VirtualDir> VirtualDir::open(const std::string &path) {
   }
 }
 
-std::unique_ptr<VirtualDir> VirtualDir::from_zip(const void *data,
-                                                 size_t size) {
-  return ZipArchiveVirtualDir::from_zip(data, size);
-}
-std::unique_ptr<VirtualDir> VirtualDir::from_fs_dir(
-    const std::string &base_dir) {
-  return FilesystemVirtualDir::create(base_dir);
-}
-
 }  // namespace io
 }  // namespace taichi

taichi/common/virtual_dir.h

@@ -12,8 +12,6 @@ struct TI_DLL_EXPORT VirtualDir {
   // Open a virtual directory based on what `path` points to. Zip files and
   // filesystem directories are supported.
   static std::unique_ptr<VirtualDir> open(const std::string &path);
-  static std::unique_ptr<VirtualDir> from_zip(const void *data, size_t size);
-  static std::unique_ptr<VirtualDir> from_fs_dir(const std::string &base_dir);
 
   // Get the `size` of the file at `path` in the virtual directory. Returns
   // false when the file doesn't exist.

taichi/runtime/gfx/aot_module_loader_impl.cpp

@@ -5,6 +5,7 @@
 
 #include "taichi/runtime/gfx/runtime.h"
 #include "taichi/aot/graph_data.h"
+#include "taichi/common/virtual_dir.h"
 
 namespace taichi::lang {
 namespace gfx {
@@ -26,10 +27,9 @@ class AotModuleImpl : public aot::Module {
       : module_path_(params.module_path),
         runtime_(params.runtime),
         device_api_backend_(device_api_backend) {
-    std::unique_ptr<io::VirtualDir> dir_alt =
-        io::VirtualDir::from_fs_dir(module_path_);
-    const io::VirtualDir *dir =
-        params.dir == nullptr ? dir_alt.get() : params.dir;
+    auto dir = io::VirtualDir::open(params.module_path);
+    TI_ERROR_IF(dir == nullptr, "cannot open aot module '{}'",
+                params.module_path);
 
     bool succ = true;
 

taichi/runtime/gfx/aot_module_loader_impl.h

@@ -11,25 +11,20 @@
 #include "taichi/codegen/spirv/kernel_utils.h"
 #include "taichi/aot/module_builder.h"
 #include "taichi/aot/module_loader.h"
-#include "taichi/common/virtual_dir.h"
 
 namespace taichi::lang {
 namespace gfx {
 
 struct TI_DLL_EXPORT AotModuleParams {
-  std::string module_path{};
-  const io::VirtualDir *dir{nullptr};
+  std::string module_path;
   GfxRuntime *runtime{nullptr};
   bool enable_lazy_loading{false};
 
   AotModuleParams() = default;
 
-  [[deprecated]] AotModuleParams(const std::string &path, GfxRuntime *rt)
+  AotModuleParams(const std::string &path, GfxRuntime *rt)
       : module_path(path), runtime(rt) {
   }
-  AotModuleParams(const io::VirtualDir *dir, GfxRuntime *rt)
-      : dir(dir), runtime(rt) {
-  }
 };
 
 TI_DLL_EXPORT std::unique_ptr<aot::Module> make_aot_module(

tests/test_config.json

@@ -198,10 +198,6 @@
     "CapiTest.TestLoadTcmAotModule": [
         ["cpp", "aot", "python_scripts", "tcm_test.py"],
         "--arch=vulkan"
-    ],
-    "CapiTest.TestCreateTcmAotModule": [
-        ["cpp", "aot", "python_scripts", "tcm_test.py"],
-        "--arch=vulkan"
     ]
 }
 }