[aot] Use graphs.json instead of TCB (taichi-dev#7392)

Issue: #

### Brief Summary

Removed TCB in GfxRuntime.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

c_api/docs/taichi/taichi_core.h.md

@@ -116,7 +116,7 @@ You can load a Taichi AOT module from the filesystem.
 TiAotModule aot_module = ti_load_aot_module(runtime, "/path/to/aot/module");
 ```
 
-`/path/to/aot/module` should point to the directory that contains a `metadata.tcb`.
+`/path/to/aot/module` should point to the directory that contains a `metadata.json`.
 
 You can destroy an unused AOT module, but please ensure that there is no kernel or compute graph related to it pending to [`ti_flush`](#function-ti_flush).
 

c_api/include/taichi/taichi_core.h

@@ -142,7 +142,7 @@
 // ```
 //
 // `/path/to/aot/module` should point to the directory that contains a
-// `metadata.tcb`.
+// `metadata.json`.
 //
 // You can destroy an unused AOT module, but please ensure that there is no
 // kernel or compute graph related to it pending to

docs/lang/articles/c-api/taichi_core.md

@@ -116,7 +116,7 @@ You can load a Taichi AOT module from the filesystem.
 TiAotModule aot_module = ti_load_aot_module(runtime, "/path/to/aot/module");
 ```
 
-`/path/to/aot/module` should point to the directory that contains a `metadata.tcb`.
+`/path/to/aot/module` should point to the directory that contains a `metadata.json`.
 
 You can destroy an unused AOT module, but please ensure that there is no kernel or compute graph related to it pending to [`ti_flush`](#function-ti_flush).
 

docs/lang/articles/deployment/tutorial.md

@@ -81,8 +81,7 @@ Now that we're done with Kernel compilation, let's take a look at the generated
 ├── demo
 │   ├── add_base_c78_0_k0001_vk_0_t00.spv
 │   ├── init_c76_0_k0000_vk_0_t00.spv
-│   ├── metadata.json
-│   └── metadata.tcb
+│   └── metadata.json
 └── demo.py
 ```
 

taichi/aot/module_builder.h

@@ -74,7 +74,7 @@ class AotModuleBuilder {
 
   static bool all_fields_are_dense_in_container(const SNode *container);
 
- private:
+ protected:
   std::unordered_map<std::string, aot::CompiledGraph> graphs_;
 };
 

taichi/runtime/gfx/aot_module_builder_impl.cpp

@@ -133,30 +133,31 @@ void AotModuleBuilderImpl::dump(const std::string &output_dir,
                                 const std::string &filename) const {
   TI_WARN_IF(!filename.empty(),
              "Filename prefix is ignored on Unified Device API backends.");
-  const std::string bin_path = fmt::format("{}/metadata.tcb", output_dir);
-  write_to_binary_file(ti_aot_data_, bin_path);
 
-  auto converted = AotDataConverter::convert(ti_aot_data_);
   const auto &spirv_codes = ti_aot_data_.spirv_codes;
-  for (int i = 0; i < std::min(ti_aot_data_.kernels.size(), spirv_codes.size());
-       ++i) {
+  for (int i = 0; i < spirv_codes.size(); ++i) {
     auto &k = ti_aot_data_.kernels[i];
-    for (int j = 0; j < std::min(k.tasks_attribs.size(), spirv_codes[i].size());
-         ++j) {
+    for (int j = 0; j < spirv_codes[i].size(); ++j) {
       if (!spirv_codes[i][j].empty()) {
         std::string spv_path =
             write_spv_file(output_dir, k.tasks_attribs[j], spirv_codes[i][j]);
-        converted.kernels[k.name].tasks[j].source_path = spv_path;
       }
     }
   }
 
-  std::string json = liong::json::print(liong::json::serialize(ti_aot_data_));
-  std::fstream f(output_dir + "/metadata.json",
-                 std::ios::trunc | std::ios::out);
-  f.write(json.data(), json.size());
+  {
+    std::string json = liong::json::print(liong::json::serialize(ti_aot_data_));
+    std::fstream f(output_dir + "/metadata.json",
+                   std::ios::trunc | std::ios::out);
+    f.write(json.data(), json.size());
+  }
 
-  dump_graph(output_dir);
+  {
+    std::string json = liong::json::print(liong::json::serialize(graphs_));
+    std::fstream f(output_dir + "/graphs.json",
+                   std::ios::trunc | std::ios::out);
+    f.write(json.data(), json.size());
+  }
 }
 
 void AotModuleBuilderImpl::mangle_aot_data() {

taichi/runtime/gfx/aot_module_loader_impl.cpp

@@ -31,20 +31,20 @@ class AotModuleImpl : public aot::Module {
     const io::VirtualDir *dir =
         params.dir == nullptr ? dir_alt.get() : params.dir;
 
-    bool succ = true;
-
-    std::vector<uint8_t> metadata_json{};
-    succ = dir->load_file("metadata.json", metadata_json) != 0;
-
-    if (!succ) {
-      mark_corrupted();
-      TI_WARN("'metadata.json' cannot be read");
-      return;
+    {
+      std::vector<uint8_t> metadata_json{};
+      bool succ = dir->load_file("metadata.json", metadata_json) != 0;
+
+      if (!succ) {
+        mark_corrupted();
+        TI_WARN("'metadata.json' cannot be read");
+        return;
+      }
+      auto json = liong::json::parse(
+          (const char *)metadata_json.data(),
+          (const char *)(metadata_json.data() + metadata_json.size()));
+      liong::json::deserialize(json, ti_aot_data_);
     }
-    auto json = liong::json::parse(
-        (const char *)metadata_json.data(),
-        (const char *)(metadata_json.data() + metadata_json.size()));
-    liong::json::deserialize(json, ti_aot_data_);
 
     if (!params.enable_lazy_loading) {
       for (int i = 0; i < ti_aot_data_.kernels.size(); ++i) {
@@ -71,14 +71,20 @@ class AotModuleImpl : public aot::Module {
       }
     }
 
-    std::vector<uint8_t> graphs_tcb{};
-    succ = dir->load_file("graphs.tcb", graphs_tcb) &&
-           read_from_binary(graphs_, graphs_tcb.data(), graphs_tcb.size());
+    {
+      std::vector<uint8_t> graphs_json{};
+      bool succ = dir->load_file("graphs.json", graphs_json) != 0;
+
+      if (!succ) {
+        mark_corrupted();
+        TI_WARN("'graphs.json' cannot be read");
+        return;
+      }
 
-    if (!succ) {
-      mark_corrupted();
-      TI_WARN("'graphs.tcb' cannot be read");
-      return;
+      auto json = liong::json::parse(
+          (const char *)graphs_json.data(),
+          (const char *)(graphs_json.data() + graphs_json.size()));
+      liong::json::deserialize(json, graphs_);
     }
   }