[opengl] Fix GLFW global context issue

cmake/TaichiCore.cmake

@@ -159,6 +159,7 @@ if (TI_WITH_OPENGL OR TI_WITH_VULKAN AND NOT ANDROID)
   endif()
 
   message("Building with GLFW")
+  add_compile_definitions(TI_WITH_GLFW)
   add_subdirectory(external/glfw)
   target_link_libraries(${CORE_LIBRARY_NAME} PRIVATE glfw)
   target_include_directories(${CORE_LIBRARY_NAME} PUBLIC external/glfw/include)

taichi/rhi/opengl/opengl_api.cpp

@@ -4,11 +4,10 @@
 
 #include "glad/gl.h"
 #include "glad/egl.h"
-#ifndef ANDROID
-#include "GLFW/glfw3.h"
-#endif  // ANDROID
 #include "taichi/rhi/opengl/opengl_device.h"
 
+#include "taichi/rhi/window_system.h"
+
 namespace taichi::lang {
 namespace opengl {
 
@@ -28,12 +27,6 @@ static bool kUseGles = false;
 static std::optional<bool> supported;  // std::nullopt
 void *kGetOpenglProcAddr;
 
-#ifndef ANDROID
-static void glfw_error_callback(int code, const char *description) {
-  TI_WARN("GLFW Error {}: {}", code, description);
-}
-#endif  // ANDROID
-
 bool initialize_opengl(bool use_gles, bool error_tolerance) {
   TI_TRACE("initialize_opengl({}, {}) called", use_gles, error_tolerance);
 
@@ -52,8 +45,7 @@ bool initialize_opengl(bool use_gles, bool error_tolerance) {
   void *get_proc_addr = nullptr;
 
 #ifndef ANDROID
-  if (glfwInit()) {
-    glfwSetErrorCallback(glfw_error_callback);
+  if (window_system::glfw_context_acquire()) {
     // Compute Shader requires OpenGL 4.3+ (or OpenGL ES 3.1+)
     if (use_gles) {
       glfwWindowHint(GLFW_CLIENT_API, GLFW_OPENGL_ES_API);
@@ -220,7 +212,7 @@ void reset_opengl() {
   supported = std::nullopt;
   kUseGles = false;
 #ifndef ANDROID
-  glfwTerminate();
+  window_system::glfw_context_release();
 #endif
 }
 

taichi/rhi/window_system.cpp

@@ -0,0 +1,60 @@
+#include "window_system.h"
+#include "taichi/rhi/impl_support.h"
+
+#include <mutex>
+#include <array>
+#include <iostream>
+
+namespace taichi::lang::window_system {
+
+#ifdef TI_WITH_GLFW
+struct GLFWState {
+  std::mutex mutex;
+  int glfw_ref_count = 0;
+};
+
+static GLFWState glfw_state;
+
+static void glfw_error_callback(int code, const char *description) {
+  std::array<char, 1024> buf;
+  snprintf(buf.data(), buf.size(), "GLFW Error %d: %s", code, description);
+  RHI_LOG_ERROR(buf.data());
+}
+
+bool glfw_context_acquire() {
+  std::lock_guard lg(glfw_state.mutex);
+  if (glfw_state.glfw_ref_count == 0) {
+    auto res = glfwInit();
+    if (res != GLFW_TRUE) {
+      return false;
+    }
+
+    glfwSetErrorCallback(glfw_error_callback);
+  }
+  glfw_state.glfw_ref_count++;
+  return true;
+}
+
+void glfw_context_release() {
+  std::lock_guard lg(glfw_state.mutex);
+  glfw_state.glfw_ref_count--;
+  if (glfw_state.glfw_ref_count == 0) {
+    glfwTerminate();
+  } else if (glfw_state.glfw_ref_count < 0) {
+    assert(false && "GLFW context double release?");
+  }
+}
+
+#else
+
+bool glfw_context_acquire() {
+  return false;
+}
+
+void glfw_context_release() {
+  return;
+}
+
+#endif  // TI_WITH_GLFW
+
+}  // namespace taichi::lang::window_system

taichi/rhi/window_system.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#ifdef TI_WITH_GLFW
+#include "GLFW/glfw3.h"
+#endif  // TI_WITH_GLFW
+
+namespace taichi::lang::window_system {
+
+bool glfw_context_acquire();
+void glfw_context_release();
+
+}  // namespace taichi::lang::window_system

taichi/runtime/program_impls/vulkan/vulkan_program.cpp

@@ -7,9 +7,7 @@
 #include "taichi/runtime/gfx/aot_module_loader_impl.h"
 #include "taichi/util/offline_cache.h"
 
-#if !defined(ANDROID)
-#include "GLFW/glfw3.h"
-#endif
+#include "taichi/rhi/window_system.h"
 
 using namespace taichi::lang::vulkan;
 
@@ -85,10 +83,6 @@ FunctionType VulkanProgramImpl::compile(const CompileConfig &compile_config,
       vulkan_runtime_.get());
 }
 
-static void glfw_error_callback(int code, const char *description) {
-  TI_WARN("GLFW Error {}: {}", code, description);
-}
-
 void VulkanProgramImpl::materialize_runtime(MemoryPool *memory_pool,
                                             KernelProfilerBase *profiler,
                                             uint64 **result_buffer_ptr) {
@@ -101,9 +95,7 @@ void VulkanProgramImpl::materialize_runtime(MemoryPool *memory_pool,
 #ifndef ANDROID
   GLFWwindow *glfw_window = nullptr;
 
-  if (glfwInit()) {
-    glfwSetErrorCallback(glfw_error_callback);
-
+  if (window_system::glfw_context_acquire()) {
     // glfw init success
     glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE);
     glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);

taichi/ui/backends/vulkan/window.cpp

@@ -3,6 +3,7 @@
 
 #include "taichi/program/program.h"
 #include "taichi/ui/utils/utils.h"
+#include "taichi/rhi/window_system.h"
 
 using taichi::lang::Program;
 
@@ -91,7 +92,7 @@ Window::~Window() {
   gui_.reset();
   renderer_.reset();
   if (config_.show_window) {
-    glfwTerminate();
+    taichi::lang::window_system::glfw_context_release();
   }
 }
 

taichi/ui/utils/utils.h

@@ -33,9 +33,7 @@
 #endif
 
 #include "taichi/rhi/vulkan/vulkan_common.h"
-#if !defined(ANDROID)
-#include <GLFW/glfw3.h>
-#endif
+#include "taichi/rhi/window_system.h"
 
 #include <stdarg.h>
 
@@ -46,37 +44,27 @@
 
 namespace taichi::ui {
 
-#if !defined(ANDROID)
-inline void initGLFW() {
-  if (!glfwInit()) {
-    printf("cannot initialize GLFW\n");
-    exit(EXIT_FAILURE);
-  }
-}
-
-static void glfw_error_callback(int code, const char *description) {
-  printf("GLFW Error %d: %s\n", code, description);
-}
-
+#ifdef TI_WITH_GLFW
 inline GLFWwindow *create_glfw_window_(const std::string &name,
                                        int screenWidth,
                                        int screenHeight,
                                        int window_pos_x,
                                        int window_pos_y,
                                        bool vsync) {
-  initGLFW();
+  if (!taichi::lang::window_system::glfw_context_acquire()) {
+    printf("cannot initialize GLFW\n");
+    exit(EXIT_FAILURE);
+  }
   GLFWwindow *window;
 
-  glfwSetErrorCallback(glfw_error_callback);
-
   glfwWindowHint(GLFW_VISIBLE, GLFW_FALSE);
   glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
 
   window = glfwCreateWindow(screenWidth, screenHeight, name.c_str(), nullptr,
                             nullptr);
 
   if (!window) {
-    glfwTerminate();
+    taichi::lang::window_system::glfw_context_release();
     exit(EXIT_FAILURE);
   }