[opengl] Reset opengl context when taichi program resets (#6987)

taichi-dev · Dec 27, 2022 · 5914971 · 5914971
1 parent fc6931f
commit 5914971
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 22 deletions.
diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp
@@ -417,6 +417,8 @@ void Program::finalize() {
   finalized_ = true;
   num_instances_ -= 1;
   program_impl_->dump_cache_data_to_disk();
+  configs.clear();
+  configs[main_thread_id_] = default_compile_config;
   TI_TRACE("Program ({}) finalized_.", fmt::ptr(this));
 }
 

diff --git a/taichi/program/program.h b/taichi/program/program.h
@@ -137,6 +137,10 @@ class TI_DLL_EXPORT Program {
     return configs[thread_id];
   }
 
+  const CompileConfig &config() {
+    return configs[main_thread_id_];
+  }
+
   struct KernelProfilerQueryResult {
     int counter{0};
     double min{0.0};

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -329,8 +329,7 @@ void export_lang(py::module &m) {
 
   py::class_<Program>(m, "Program")
       .def(py::init<>())
-      .def("config", &Program::this_thread_config,
-           py::return_value_policy::reference)
+      .def("config", &Program::config)
       .def("sync_kernel_profiler",
            [](Program *program) { program->profiler->sync(); })
       .def("update_kernel_profiler",

diff --git a/taichi/rhi/opengl/opengl_api.cpp b/taichi/rhi/opengl/opengl_api.cpp
@@ -23,14 +23,13 @@ int opengl_max_grid_dim = 1024;
 // TODO: Properly support setting GLES/GLSL in opengl backend
 // without this global static boolean.
 static bool kUseGles = false;
+static std::optional<bool> supported;  // std::nullopt
 
 static void glfw_error_callback(int code, const char *description) {
   TI_WARN("GLFW Error {}: {}", code, description);
 }
 
 bool initialize_opengl(bool use_gles, bool error_tolerance) {
-  static std::optional<bool> supported;  // std::nullopt
-
   TI_TRACE("initialize_opengl({}, {}) called", use_gles, error_tolerance);
 
   if (supported.has_value()) {  // this function has been called before
@@ -205,6 +204,12 @@ bool is_gles() {
   return kUseGles;
 }
 
+void reset_opengl() {
+  supported = std::nullopt;
+  kUseGles = false;
+  glfwTerminate();
+}
+
 std::shared_ptr<Device> make_opengl_device() {
   std::shared_ptr<Device> dev = std::make_shared<GLDevice>();
   return dev;

diff --git a/taichi/rhi/opengl/opengl_api.h b/taichi/rhi/opengl/opengl_api.h
@@ -9,6 +9,7 @@ namespace opengl {
 bool initialize_opengl(bool use_gles = false, bool error_tolerance = false);
 bool is_opengl_api_available(bool use_gles = false);
 bool is_gles();
+void reset_opengl();
 
 std::shared_ptr<Device> make_opengl_device();
 

diff --git a/taichi/rhi/opengl/opengl_device.cpp b/taichi/rhi/opengl/opengl_device.cpp
@@ -597,26 +597,14 @@ RhiResult GLDevice::map_range(DevicePtr ptr, uint64_t size, void **mapped_ptr) {
 }
 
 RhiResult GLDevice::map(DeviceAllocation alloc, void **mapped_ptr) {
-  TI_ASSERT_INFO(
-      buffer_to_access_.find(alloc.alloc_id) != buffer_to_access_.end(),
-      "Buffer not created with host_read or write");
   glBindBuffer(GL_SHADER_STORAGE_BUFFER, alloc.alloc_id);
   check_opengl_error("glBindBuffer");
-  // This is pure stupidity.
-  // Why does `glMapBuffer` and `glMapBufferRange` uses two TOTALLY different
-  // enums? Whoever came up with the API is DRUNK!
-  GLbitfield access = buffer_to_access_.at(alloc.alloc_id);
-  GLenum access_oldapi;
-  if (bool(access & GL_MAP_READ_BIT) && bool(access & GL_MAP_WRITE_BIT)) {
-    access_oldapi = GL_READ_WRITE;
-  } else if (access & GL_MAP_WRITE_BIT) {
-    access_oldapi = GL_WRITE_ONLY;
-  } else {
-    access_oldapi = GL_READ_ONLY;
-  }
-  *mapped_ptr = glMapBuffer(GL_SHADER_STORAGE_BUFFER, access_oldapi);
-  check_opengl_error("glMapBuffer");
-  return RhiResult::success;
+
+  int size = 0;
+  glGetBufferParameteriv(GL_SHADER_STORAGE_BUFFER, GL_BUFFER_SIZE, &size);
+  check_opengl_error("glGetBufferParameteriv");
+
+  return map_range(alloc.get_ptr(0), size, mapped_ptr);
 }
 
 void GLDevice::unmap(DevicePtr ptr) {

diff --git a/taichi/runtime/program_impls/opengl/opengl_program.cpp b/taichi/runtime/program_impls/opengl/opengl_program.cpp
@@ -92,6 +92,16 @@ void OpenglProgramImpl::dump_cache_data_to_disk() {
   mgr->dump_with_merging();
 }
 
+void OpenglProgramImpl::finalize() {
+  runtime_.reset();
+  device_.reset();
+  opengl::reset_opengl();
+}
+
+OpenglProgramImpl::~OpenglProgramImpl() {
+  finalize();
+}
+
 const std::unique_ptr<gfx::CacheManager>
     &OpenglProgramImpl::get_cache_manager() {
   if (!cache_manager_) {

diff --git a/taichi/runtime/program_impls/opengl/opengl_program.h b/taichi/runtime/program_impls/opengl/opengl_program.h
@@ -10,6 +10,8 @@ namespace taichi::lang {
 class OpenglProgramImpl : public ProgramImpl {
  public:
   explicit OpenglProgramImpl(CompileConfig &config);
+  ~OpenglProgramImpl() override;
+
   FunctionType compile(Kernel *kernel, OffloadedStmt *offloaded) override;
 
   std::size_t get_snode_num_dynamically_allocated(
@@ -30,6 +32,8 @@ class OpenglProgramImpl : public ProgramImpl {
     runtime_->synchronize();
   }
 
+  void finalize() override;
+
   StreamSemaphore flush() override {
     return runtime_->flush();
   }