[llvm] [aot] Add LLVM to CAPI part 6: Handle Field initialization in C-API

c_api/src/taichi_llvm_impl.cpp

@@ -5,6 +5,7 @@
 
 #include "taichi/program/compile_config.h"
 #include "taichi/runtime/llvm/llvm_runtime_executor.h"
+#include "taichi/runtime/llvm/llvm_aot_module_loader.h"
 #include "taichi/runtime/cpu/aot_module_loader_impl.h"
 
 #ifdef TI_WITH_CUDA
@@ -89,6 +90,19 @@ TiAotModule LlvmRuntime::load_aot_module(const char *module_path) {
 #endif
   }
 
+  /* TODO(zhanlue): expose allocate/deallocate_snode_tree_type() to C-API
+     Let's initialize SNodeTrees automatically for now since SNodeTreeType isn't
+     ready yet.
+  */
+  auto *llvm_aot_module =
+      dynamic_cast<taichi::lang::LlvmAotModule *>(aot_module.get());
+  TI_ASSERT(llvm_aot_module != nullptr);
+  for (size_t i = 0; i < llvm_aot_module->get_num_snode_trees(); i++) {
+    auto *snode_tree = aot_module->get_field(std::to_string(i));
+    taichi::lang::allocate_aot_snode_tree_type(aot_module.get(), snode_tree,
+                                               this->result_buffer);
+  }
+
   // Insert LLVMRuntime to RuntimeContext
   executor_->prepare_runtime_context(&this->runtime_context_);
   return (TiAotModule)(new AotModule(*this, std::move(aot_module)));

taichi/runtime/llvm/llvm_aot_module_loader.cpp

@@ -54,9 +54,9 @@ std::unique_ptr<aot::CompiledGraph> LlvmAotModule::get_graph(std::string name) {
   return std::make_unique<aot::CompiledGraph>(std::move(graph));
 }
 
-void finalize_aot_field(aot::Module *aot_module,
-                        aot::Field *aot_field,
-                        uint64 *result_buffer) {
+void allocate_aot_snode_tree_type(aot::Module *aot_module,
+                                  aot::Field *aot_field,
+                                  uint64 *result_buffer) {
   auto *llvm_aot_module = dynamic_cast<LlvmAotModule *>(aot_module);
   auto *aot_field_impl = dynamic_cast<llvm_aot::FieldImpl *>(aot_field);
 

taichi/runtime/llvm/llvm_aot_module_loader.h

@@ -6,9 +6,13 @@
 namespace taichi {
 namespace lang {
 
-TI_DLL_EXPORT void finalize_aot_field(aot::Module *aot_module,
-                                      aot::Field *aot_field,
-                                      uint64 *result_buffer);
+/* TODO(zhanlue) refactor this interface once SNodeTreeType is available
+   The "aot::Field" created by "make_new_field()" is a SNodeTree in essense.
+   Therefore we're actually initializing the entire SNodeTree.
+*/
+TI_DLL_EXPORT void allocate_aot_snode_tree_type(aot::Module *aot_module,
+                                                aot::Field *aot_field,
+                                                uint64 *result_buffer);
 
 class LlvmAotModule : public aot::Module {
  public:
@@ -38,6 +42,10 @@ class LlvmAotModule : public aot::Module {
     return executor_;
   }
 
+  size_t get_num_snode_trees() {
+    return cache_reader_->get_num_snode_trees();
+  }
+
   void set_initialized_snode_tree(int snode_tree_id) {
     initialized_snode_tree_ids.insert(snode_tree_id);
   }
@@ -59,6 +67,11 @@ class LlvmAotModule : public aot::Module {
   std::unique_ptr<aot::Kernel> make_new_kernel(
       const std::string &name) override;
 
+  /* TODO(zhanlue): replace "make_new_field()" with "make_snode_tree()" once
+     SNodeTreeType is available Field is not a standalone data structure - it is
+     essentially part of a SNodeTree object. User should always operate on a
+     "SNodeTree" instead of a "Field".
+  */
   std::unique_ptr<aot::Field> make_new_field(const std::string &name) override;
 
   LlvmRuntimeExecutor *const executor_{nullptr};

taichi/runtime/llvm/llvm_offline_cache.cpp

@@ -83,6 +83,10 @@ LlvmOfflineCacheFileReader::LlvmOfflineCacheFileReader(
     : path_(path), data_(std::move(data)), format_(format) {
 }
 
+size_t LlvmOfflineCacheFileReader::get_num_snode_trees() {
+  return data_.fields.size();
+}
+
 bool LlvmOfflineCacheFileReader::get_field_cache(
     LlvmOfflineCache::FieldCacheData &res,
     int snode_tree_id) {

taichi/runtime/llvm/llvm_offline_cache.h

@@ -115,6 +115,8 @@ class LlvmOfflineCacheFileReader {
   bool get_field_cache(LlvmOfflineCache::FieldCacheData &res,
                        int snode_tree_id);
 
+  size_t get_num_snode_trees();
+
   static std::unique_ptr<LlvmOfflineCacheFileReader> make(
       const std::string &path,
       LlvmOfflineCache::Format format = LlvmOfflineCache::Format::LL);

tests/cpp/aot/llvm/field_aot_test.cpp

@@ -33,12 +33,9 @@ void run_field_tests(aot::Module *mod,
   aot::Kernel *k_check_activate_pointer_fields =
       mod->get_kernel("check_activate_pointer_fields");
 
-  // Initialize Fields
-  aot::Field *field_x = mod->get_field("0" /*snode_tree_id*/);
-  aot::Field *field_y = mod->get_field("0" /*snode_tree_id*/);
-
-  finalize_aot_field(mod, field_x, result_buffer);
-  finalize_aot_field(mod, field_y, result_buffer);
+  // Initialize SNodeTree
+  aot::Field *snode_tree_0 = mod->get_field("0" /*snode_tree_id*/);
+  allocate_aot_snode_tree_type(mod, snode_tree_0, result_buffer);
 
   int base_value = 10;
   /* -------- Test Case 1 ------ */