[cc] The C backend is now capable of running mpm128 (#1553)

* tmp

* fix root buf

* fix local store

* capable of mpm

* support gtmp

* [skip ci] enforce code format

* [skip ci] clean misc/cc_*.py

* RTi -> Ti; Ti -> Tk

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

examples/taichi_logo.py

@@ -55,6 +55,6 @@ def paint():
 paint()
 
 gui = ti.GUI('Logo', (512, 512))
-while not gui.running:
+while gui.running:
     gui.set_image(x.to_numpy())
     gui.show()

taichi/backends/cc/cc_layout.h

@@ -5,19 +5,24 @@
 TLANG_NAMESPACE_BEGIN
 namespace cccp {
 
+class CCProgram;
+
 class CCLayout {
  public:
-  CCLayout() = default;
+  CCLayout(CCProgram *program) : program(program) {
+  }
 
   std::string get_object() {
     return obj_path;
   }
 
-  void compile();
+  size_t compile();
 
   std::string source;
 
  private:
+  CCProgram *program;
+
   std::string src_path;
   std::string obj_path;
 };

taichi/backends/cc/cc_program.cpp

@@ -1,11 +1,15 @@
 #include "taichi/common/core.h"
+#include "taichi/program/kernel.h"
+#include "taichi/program/program.h"
 #include "taichi/system/dynamic_loader.h"
+#include "struct_cc.h"
 #include "cc_program.h"
 #include "cc_configuation.h"
 #include "cc_runtime.h"
 #include "cc_kernel.h"
 #include "cc_layout.h"
 #include "cc_utils.h"
+#include "context.h"
 
 TLANG_NAMESPACE_BEGIN
 namespace cccp {
@@ -16,35 +20,61 @@ void CCKernel::compile() {
   obj_path = fmt::format("{}/{}.o", runtime_tmp_dir, name);
   src_path = fmt::format("{}/{}.c", runtime_tmp_dir, name);
 
-  std::ofstream(src_path) << program->runtime->header << "\n"
-                          << program->layout->source << "\n"
+  std::ofstream(src_path) << program->get_runtime()->header << "\n"
+                          << program->get_layout()->source << "\n"
                           << source;
   TI_DEBUG("[cc] compiling [{}] -> [{}]:\n{}\n", name, obj_path, source);
   execute(cfg.compile_cmd, obj_path, src_path);
 }
 
+CCContext::CCContext(CCProgram *program, Context *ctx)
+    : args(ctx->args), earg((int *)ctx->extra_args) {
+  root = program->get_root_buffer();
+  gtmp = program->get_gtmp_buffer();
+}
+
 void CCKernel::launch(Context *ctx) {
   program->relink();
-  auto entry = program->load_kernel(name);
   TI_TRACE("[cc] entering kernel [{}]", name);
-  (*entry)();
+  auto entry = program->load_kernel(name);
+  TI_ASSERT(entry);
+  CCContext cc_ctx(program, ctx);
+  (*entry)(&cc_ctx);
   TI_TRACE("[cc] leaving kernel [{}]", name);
 }
 
-void CCLayout::compile() {
-  obj_path = fmt::format("{}/_root.o", runtime_tmp_dir);
-  src_path = fmt::format("{}/_root.c", runtime_tmp_dir);
+size_t CCLayout::compile() {
+  obj_path = fmt::format("{}/_rti_root.o", runtime_tmp_dir);
+  src_path = fmt::format("{}/_rti_root.c", runtime_tmp_dir);
+  auto dll_path = fmt::format("{}/libti_roottest.so", runtime_tmp_dir);
+
+  std::ofstream(src_path) << program->get_runtime()->header << "\n"
+                          << source << "\n"
+                          << "void *Ti_get_root_size(void) { \n"
+                          << "  return (void *) sizeof(struct S0root);\n"
+                          << "}\n";
 
-  std::ofstream(src_path)
-      << source << "\n\nstruct S0root *RTi_get_root() {\n"
-      << "\tstatic struct S0root ti_root;\n\treturn &ti_root;\n}\n";
   TI_DEBUG("[cc] compiling root struct -> [{}]:\n{}\n", obj_path, source);
   execute(cfg.compile_cmd, obj_path, src_path);
+
+  TI_DEBUG("[cc] linking root struct object [{}] -> [{}]", obj_path, dll_path);
+  execute(cfg.link_cmd, dll_path, obj_path);
+
+  TI_DEBUG("[cc] loading root struct object: {}", dll_path);
+  DynamicLoader dll(dll_path);
+  TI_ASSERT_INFO(dll.loaded(), "[cc] could not load shared object: {}",
+                 dll_path);
+
+  using FuncGetRootSizeType = size_t();
+  auto get_root_size = reinterpret_cast<FuncGetRootSizeType *>(
+      dll.load_function("Ti_get_root_size"));
+  TI_ASSERT(get_root_size);
+  return (*get_root_size)();
 }
 
 void CCRuntime::compile() {
-  obj_path = fmt::format("{}/_runtime.o", runtime_tmp_dir);
-  src_path = fmt::format("{}/_runtime.c", runtime_tmp_dir);
+  obj_path = fmt::format("{}/_rti_runtime.o", runtime_tmp_dir);
+  src_path = fmt::format("{}/_rti_runtime.c", runtime_tmp_dir);
 
   std::ofstream(src_path) << header << "\n" << source;
   TI_DEBUG("[cc] compiling runtime -> [{}]:\n{}\n", obj_path, source);
@@ -58,7 +88,6 @@ void CCProgram::relink() {
   dll_path = fmt::format("{}/libti_program.so", runtime_tmp_dir);
 
   std::vector<std::string> objects;
-  objects.push_back(layout->get_object());
   objects.push_back(runtime->get_object());
   for (auto const &ker : kernels) {
     objects.push_back(ker->get_object());
@@ -68,6 +97,7 @@ void CCProgram::relink() {
            fmt::join(objects, "] ["));
   execute(cfg.link_cmd, dll_path, fmt::join(objects, "' '"));
 
+  dll = nullptr;
   TI_DEBUG("[cc] loading shared object: {}", dll_path);
   dll = std::make_unique<DynamicLoader>(dll_path);
   TI_ASSERT_INFO(dll->loaded(), "[cc] could not load shared object: {}",
@@ -76,23 +106,32 @@ void CCProgram::relink() {
   need_relink = false;
 }
 
+void CCProgram::compile_layout(SNode *root) {
+  CCLayoutGen gen(this, root);
+  layout = gen.compile();
+  size_t root_size = layout->compile();
+  TI_INFO("[cc] C backend root buffer size: {} B", root_size);
+  root_buf.resize(root_size, 0);
+  gtmp_buf.resize(taichi_global_tmp_buffer_size, 0);
+}
+
 void CCProgram::add_kernel(std::unique_ptr<CCKernel> kernel) {
   kernels.push_back(std::move(kernel));
   need_relink = true;
 }
 
 // TODO: move this to cc_runtime.cpp:
 void CCProgram::init_runtime() {
-  runtime = std::make_unique<CCRuntime>(
+  runtime = std::make_unique<CCRuntime>(this,
 #include "runtime/base.h"
-      ,
+                                        ,
 #include "runtime/base.c"
   );
   runtime->compile();
 }
 
 CCFuncEntryType *CCProgram::load_kernel(std::string const &name) {
-  return reinterpret_cast<CCFuncEntryType *>(dll->load_function("Ti_" + name));
+  return reinterpret_cast<CCFuncEntryType *>(dll->load_function("Tk_" + name));
 }
 
 CCProgram::CCProgram() {

taichi/backends/cc/cc_program.h

@@ -10,13 +10,16 @@ TI_NAMESPACE_END
 
 TLANG_NAMESPACE_BEGIN
 
+class SNode;
+
 namespace cccp {
 
 class CCKernel;
 class CCLayout;
 class CCRuntime;
+struct CCContext;
 
-using CCFuncEntryType = void();
+using CCFuncEntryType = void(CCContext *);
 
 class CCProgram {
   // Launch C compiler to compile generated source code, and run them
@@ -26,9 +29,29 @@ class CCProgram {
 
   void add_kernel(std::unique_ptr<CCKernel> kernel);
   CCFuncEntryType *load_kernel(std::string const &name);
+  void compile_layout(SNode *root);
   void init_runtime();
   void relink();
 
+  CCLayout *get_layout() {
+    return layout.get();
+  }
+
+  CCRuntime *get_runtime() {
+    return runtime.get();
+  }
+
+  void *get_root_buffer() {
+    return root_buf.data();
+  }
+
+  void *get_gtmp_buffer() {
+    return root_buf.data();
+  }
+
+ private:
+  std::vector<char> root_buf;
+  std::vector<char> gtmp_buf;
   std::vector<std::unique_ptr<CCKernel>> kernels;
   std::unique_ptr<CCRuntime> runtime;
   std::unique_ptr<CCLayout> layout;

taichi/backends/cc/cc_runtime.h

@@ -5,10 +5,14 @@
 TLANG_NAMESPACE_BEGIN
 namespace cccp {
 
+class CCProgram;
+
 class CCRuntime {
  public:
-  CCRuntime(std::string const &header, std::string const &source)
-      : header(header), source(source) {
+  CCRuntime(CCProgram *program,
+            std::string const &header,
+            std::string const &source)
+      : header(header), source(source), program(program) {
   }
 
   std::string get_object() {
@@ -21,6 +25,8 @@ class CCRuntime {
   std::string source;
 
  private:
+  [[maybe_unused]] CCProgram *program;
+
   std::string src_path;
   std::string obj_path;
 };

taichi/backends/cc/cc_utils.h

@@ -11,30 +11,7 @@ TLANG_NAMESPACE_BEGIN
 namespace cccp {
 
 inline std::string cc_data_type_name(DataType dt) {
-  switch (dt) {
-    case DataType::i8:
-      return "char";
-    case DataType::i16:
-      return "short";
-    case DataType::i32:
-      return "int";
-    case DataType::i64:
-      return "long long";
-    case DataType::u8:
-      return "unsigned char";
-    case DataType::u16:
-      return "unsigned short";
-    case DataType::u32:
-      return "unsigned int";
-    case DataType::u64:
-      return "unsigned long long";
-    case DataType::f32:
-      return "float";
-    case DataType::f64:
-      return "double";
-    default:
-      TI_ERROR("Unsupported DataType={} on C backend", data_type_name(dt));
-  }
+  return "Ti_" + data_type_short_name(dt);
 }
 
 inline std::string cc_atomic_op_type_symbol(AtomicOpType op) {