[opengl] Runtime refactor 1/n (#2965)

* Rename GLSLLauncher to OpenGLRuntime & seperate materialize runtime / materialize snode * Auto Format * oops * fix nits Co-authored-by: Taichi Gardener <[email protected]>
taichi-dev · Sep 22, 2021 · 992e6ce · 992e6ce
1 parent 4ba617c
commit 992e6ce
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 46 deletions.
diff --git a/taichi/backends/opengl/codegen_opengl.cpp b/taichi/backends/opengl/codegen_opengl.cpp
@@ -150,7 +150,7 @@ class KernelGen : public IRVisitor {
 
     // clang-format off
     if (used.print)  // the runtime buffer is only used for print now..
-      line_appender_header_.append_raw(shaders::kOpenGLRuntimeSourceCode);
+      line_appender_header_.append_raw(shaders::kOpenGlRuntimeSourceCode);
     if (used.listman)
       line_appender_header_.append_raw(shaders::kOpenGLListmanSourceCode);
 
@@ -1174,14 +1174,12 @@ class KernelGen : public IRVisitor {
 FunctionType OpenglCodeGen::gen(void) {
 #if defined(TI_WITH_OPENGL)
   KernelGen codegen(kernel_, kernel_name_, struct_compiled_,
-                    kernel_launcher_->device.get());
+                    runtime_->device.get());
   codegen.run();
   auto compiled = codegen.get_compiled_program();
   auto *ptr = compiled.get();
-  kernel_launcher_->keep(std::move(compiled));
-  return [ptr, launcher = kernel_launcher_](Context &ctx) {
-    ptr->launch(ctx, launcher);
-  };
+  runtime_->keep(std::move(compiled));
+  return [ptr, runtime = runtime_](Context &ctx) { ptr->launch(ctx, runtime); };
 #else
   TI_NOT_IMPLEMENTED
 #endif

diff --git a/taichi/backends/opengl/codegen_opengl.h b/taichi/backends/opengl/codegen_opengl.h
@@ -18,10 +18,10 @@ class OpenglCodeGen {
  public:
   OpenglCodeGen(const std::string &kernel_name,
                 StructCompiledResult *struct_compiled,
-                GLSLLauncher *launcher)
+                OpenGlRuntime *launcher)
       : kernel_name_(kernel_name),
         struct_compiled_(struct_compiled),
-        kernel_launcher_(launcher) {
+        runtime_(launcher) {
   }
 
   FunctionType compile(Kernel &kernel);
@@ -34,7 +34,7 @@ class OpenglCodeGen {
 
   Kernel *kernel_;
   [[maybe_unused]] StructCompiledResult *struct_compiled_;
-  [[maybe_unused]] GLSLLauncher *kernel_launcher_;
+  [[maybe_unused]] OpenGlRuntime *runtime_;
 };
 
 }  // namespace opengl

diff --git a/taichi/backends/opengl/opengl_api.cpp b/taichi/backends/opengl/opengl_api.cpp
@@ -47,17 +47,15 @@ static std::string add_line_markers(std::string x) {
   return x;
 }
 
-struct GLSLLauncherImpl {
-  std::unique_ptr<Device> device;
-
+struct OpenGlRuntimeImpl {
   struct {
     DeviceAllocation runtime = kDeviceNullAllocation;
     DeviceAllocation listman = kDeviceNullAllocation;
     DeviceAllocation root = kDeviceNullAllocation;
     DeviceAllocation gtmp = kDeviceNullAllocation;
   } core_bufs;
 
-  GLSLLauncherImpl() {
+  OpenGlRuntimeImpl() {
   }
 
   std::unique_ptr<GLSLRuntime> runtime;
@@ -264,9 +262,8 @@ struct CompiledProgram::Impl {
     return i;
   }
 
-  void dump_message_buffer(GLSLLauncher *launcher) const {
-    auto runtime = launcher->impl->core_bufs.runtime;
-    auto rt_buf = (GLSLRuntime *)device->map(launcher->impl->core_bufs.runtime);
+  void dump_message_buffer(OpenGlRuntime *runtime) const {
+    auto rt_buf = (GLSLRuntime *)device->map(runtime->impl->core_bufs.runtime);
 
     auto msg_count = rt_buf->msg_count;
     if (msg_count > MAX_MESSAGES) {
@@ -300,7 +297,7 @@ struct CompiledProgram::Impl {
       }
     }
     rt_buf->msg_count = 0;
-    device->unmap(launcher->impl->core_bufs.runtime);
+    device->unmap(runtime->impl->core_bufs.runtime);
   }
 
   bool check_ext_arr_read(int i) const {
@@ -335,7 +332,7 @@ struct CompiledProgram::Impl {
     return access;
   }
 
-  void launch(Context &ctx, GLSLLauncher *launcher) const {
+  void launch(Context &ctx, OpenGlRuntime *launcher) const {
     std::array<void *, taichi_max_num_args> ext_arr_host_ptrs;
 
     uint8_t *args_buf_mapped = nullptr;
@@ -428,12 +425,12 @@ struct CompiledProgram::Impl {
   }
 };
 
-GLSLLauncher::GLSLLauncher(size_t root_size) {
+OpenGlRuntime::OpenGlRuntime() {
   initialize_opengl();
 
   device = std::make_unique<GLDevice>();
 
-  impl = std::make_unique<GLSLLauncherImpl>();
+  impl = std::make_unique<OpenGlRuntimeImpl>();
 
   impl->runtime = std::make_unique<GLSLRuntime>();
   impl->core_bufs.runtime = device->allocate_memory(
@@ -442,8 +439,6 @@ GLSLLauncher::GLSLLauncher(size_t root_size) {
   impl->listman = std::make_unique<GLSLListman>();
   impl->core_bufs.listman = device->allocate_memory({sizeof(GLSLListman)});
 
-  impl->core_bufs.root = device->allocate_memory({root_size});
-
   impl->core_bufs.gtmp =
       device->allocate_memory({taichi_global_tmp_buffer_size});
 
@@ -452,16 +447,23 @@ GLSLLauncher::GLSLLauncher(size_t root_size) {
                        0);
   cmdlist->buffer_fill(impl->core_bufs.listman.get_ptr(0), sizeof(GLSLListman),
                        0);
-  cmdlist->buffer_fill(impl->core_bufs.root.get_ptr(0), root_size, 0);
   cmdlist->buffer_fill(impl->core_bufs.gtmp.get_ptr(0),
                        taichi_global_tmp_buffer_size, 0);
   device->get_compute_stream()->submit_synced(cmdlist.get());
 }
 
-void GLSLLauncher::keep(std::unique_ptr<CompiledProgram> program) {
+void OpenGlRuntime::keep(std::unique_ptr<CompiledProgram> program) {
   impl->programs.push_back(std::move(program));
 }
 
+void OpenGlRuntime::add_snode_tree(size_t size) {
+  impl->core_bufs.root = device->allocate_memory({size});
+
+  auto cmdlist = device->get_compute_stream()->new_command_list();
+  cmdlist->buffer_fill(impl->core_bufs.root.get_ptr(0), size, 0);
+  device->get_compute_stream()->submit_synced(cmdlist.get());
+}
+
 bool is_opengl_api_available() {
   if (get_environ_config("TI_ENABLE_OPENGL", 1) == 0)
     return false;
@@ -470,7 +472,7 @@ bool is_opengl_api_available() {
 
 #else
 struct GLProgram {};
-struct GLSLLauncherImpl {};
+struct OpenGlRuntimeImpl {};
 
 struct CompiledProgram::Impl {
   UsedFeature used;
@@ -491,16 +493,20 @@ struct CompiledProgram::Impl {
     TI_NOT_IMPLEMENTED;
   }
 
-  void launch(Context &ctx, GLSLLauncher *launcher) const {
+  void launch(Context &ctx, OpenGlRuntime *launcher) const {
     TI_NOT_IMPLEMENTED;
   }
 };
 
-GLSLLauncher::GLSLLauncher(size_t size) {
+OpenGlRuntime::OpenGlRuntime() {
+  TI_NOT_IMPLEMENTED;
+}
+
+void OpenGlRuntime::keep(std::unique_ptr<CompiledProgram>) {
   TI_NOT_IMPLEMENTED;
 }
 
-void GLSLLauncher::keep(std::unique_ptr<CompiledProgram>) {
+void OpenGlRuntime::add_snode_tree(size_t size) {
   TI_NOT_IMPLEMENTED;
 }
 
@@ -538,11 +544,11 @@ int CompiledProgram::lookup_or_add_string(const std::string &str) {
   return impl->lookup_or_add_string(str);
 }
 
-void CompiledProgram::launch(Context &ctx, GLSLLauncher *launcher) const {
+void CompiledProgram::launch(Context &ctx, OpenGlRuntime *launcher) const {
   impl->launch(ctx, launcher);
 }
 
-GLSLLauncher::~GLSLLauncher() = default;
+OpenGlRuntime::~OpenGlRuntime() = default;
 
 }  // namespace opengl
 TLANG_NAMESPACE_END
diff --git a/taichi/backends/opengl/opengl_api.h b/taichi/backends/opengl/opengl_api.h
@@ -60,7 +60,7 @@ struct CompiledProgram {
                nullptr);
   void set_used(const UsedFeature &used);
   int lookup_or_add_string(const std::string &str);
-  void launch(Context &ctx, GLSLLauncher *launcher) const;
+  void launch(Context &ctx, OpenGlRuntime *launcher) const;
 };
 
 }  // namespace opengl

diff --git a/taichi/backends/opengl/opengl_kernel_launcher.h b/taichi/backends/opengl/opengl_kernel_launcher.h
@@ -10,16 +10,18 @@ TLANG_NAMESPACE_BEGIN
 namespace opengl {
 
 struct CompiledProgram;
-struct GLSLLauncherImpl;
-struct GLSLLauncher;
+struct OpenGlRuntimeImpl;
+struct OpenGlRuntime;
 class GLBuffer;
 
-struct GLSLLauncher {
-  std::unique_ptr<GLSLLauncherImpl> impl;
+struct OpenGlRuntime {
+  std::unique_ptr<OpenGlRuntimeImpl> impl;
   std::unique_ptr<Device> device{nullptr};
-  GLSLLauncher(size_t size);
-  ~GLSLLauncher();
+  OpenGlRuntime();
+  ~OpenGlRuntime();
   void keep(std::unique_ptr<CompiledProgram> program);
+  // FIXME: Currently GLSL codegen only supports single root
+  void add_snode_tree(size_t size);
 
   void *result_buffer;
 };

diff --git a/taichi/backends/opengl/opengl_program.cpp b/taichi/backends/opengl/opengl_program.cpp
@@ -7,7 +7,7 @@ namespace lang {
 FunctionType OpenglProgramImpl::compile(Kernel *kernel,
                                         OffloadedStmt *offloaded) {
   opengl::OpenglCodeGen codegen(kernel->name, &opengl_struct_compiled_.value(),
-                                opengl_kernel_launcher_.get());
+                                opengl_runtime_.get());
   return codegen.compile(*kernel);
 }
 
@@ -16,9 +16,7 @@ void OpenglProgramImpl::materialize_runtime(MemoryPool *memory_pool,
                                             uint64 **result_buffer_ptr) {
   *result_buffer_ptr = (uint64 *)memory_pool->allocate(
       sizeof(uint64) * taichi_result_buffer_entries, 8);
-  // doesn't do anything other than alloc result buffer. runtime is materialized
-  // together with snode tree.
-  // TODO: separate runtime materialization and tree materialization.
+  opengl_runtime_ = std::make_unique<opengl::OpenGlRuntime>();
 }
 
 void OpenglProgramImpl::materialize_snode_tree(
@@ -31,9 +29,8 @@ void OpenglProgramImpl::materialize_snode_tree(
   opengl::OpenglStructCompiler scomp;
   opengl_struct_compiled_ = scomp.run(*root);
   TI_TRACE("OpenGL root buffer size: {} B", opengl_struct_compiled_->root_size);
-  opengl_kernel_launcher_ = std::make_unique<opengl::GLSLLauncher>(
-      opengl_struct_compiled_->root_size);
-  opengl_kernel_launcher_->result_buffer = result_buffer;
+  opengl_runtime_->add_snode_tree(opengl_struct_compiled_->root_size);
+  opengl_runtime_->result_buffer = result_buffer;
 }
 
 }  // namespace lang

diff --git a/taichi/backends/opengl/opengl_program.h b/taichi/backends/opengl/opengl_program.h
@@ -55,7 +55,7 @@ class OpenglProgramImpl : public ProgramImpl {
 
  private:
   std::optional<opengl::StructCompiledResult> opengl_struct_compiled_;
-  std::unique_ptr<opengl::GLSLLauncher> opengl_kernel_launcher_;
+  std::unique_ptr<opengl::OpenGlRuntime> opengl_runtime_;
 };
 }  // namespace lang
 }  // namespace taichi
diff --git a/taichi/backends/opengl/shaders/runtime.h b/taichi/backends/opengl/shaders/runtime.h
@@ -6,7 +6,7 @@
 #ifdef TI_INSIDE_OPENGL_CODEGEN
 
 #ifndef TI_OPENGL_NESTED_INCLUDE
-#define OPENGL_BEGIN_RUNTIME_DEF constexpr auto kOpenGLRuntimeSourceCode =
+#define OPENGL_BEGIN_RUNTIME_DEF constexpr auto kOpenGlRuntimeSourceCode =
 #define OPENGL_END_RUNTIME_DEF ;
 #else
 #define OPENGL_BEGIN_RUNTIME_DEF