[skip ci] use compute shader

taichi-dev · Feb 18, 2020 · 4f8b696 · 4f8b696
1 parent 4909e32
commit 4f8b696
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 33 deletions.
diff --git a/cmake/TaichiCore.cmake b/cmake/TaichiCore.cmake
@@ -88,9 +88,8 @@ if (TI_WITH_CUDA)
     target_link_libraries(${LIBRARY_NAME} ${llvm_ptx_libs})
 endif()
 
-target_link_libraries(${LIBRARY_NAME} /usr/lib/libGLESv2.so GLESv2)
-target_link_libraries(${LIBRARY_NAME} /usr/lib/libGLEW.so GLEW)
-target_link_libraries(${LIBRARY_NAME} /usr/lib/libglut.so glut)
+target_link_libraries(${LIBRARY_NAME} /usr/lib/libGLESv2_nvidia.so GLESv2_nvidia)
+target_link_libraries(${LIBRARY_NAME} /usr/lib/libglfw.so glfw)
 
 # add_executable(runtime runtime/runtime.cpp)
 

diff --git a/taichi/backends/codegen_opengl.cpp b/taichi/backends/codegen_opengl.cpp
@@ -70,9 +70,11 @@ class KernelGen : public IRVisitor
 
   void generate_header()
   {
-    emit("#version 330 core");
-    emit("layout (location = 0) in float tmp7;"); // XXX
-    emit("layout (location = 1) out float output;"); // XXX
+    emit("#version 430 core");
+    //emit("#version 310 es");
+    emit("#extension GL_ARB_compute_shader: enable");
+    emit("layout (location = 0) in float input;");
+    emit("layout (location = 1) out float output;");
     emit("");
   }
 

diff --git a/taichi/platform/opengl/opengl_api.cpp b/taichi/platform/opengl/opengl_api.cpp
@@ -1,10 +1,10 @@
+//#define USE_GLEW
 #ifdef USE_GLEW
 #include <GL/glew.h>
 #else
-#include <GLES3/gl3.h>
+#include <GLES3/gl32.h>
 #endif
-#include <GL/glut.h>
-#include <GL/freeglut_ext.h>
+#include <GLFW/glfw3.h>
 #include "opengl_api.h"
 
 TLANG_NAMESPACE_BEGIN
@@ -19,12 +19,12 @@ struct GLShader
 {
   GLuint id_;
 
-  GLShader(GLuint type = GL_FRAGMENT_SHADER)
+  GLShader(GLuint type = GL_COMPUTE_SHADER)
   {
     id_ = glCreateShader(type);
   }
 
-  GLShader(std::string source, GLuint type = GL_FRAGMENT_SHADER)
+  GLShader(std::string source, GLuint type = GL_COMPUTE_SHADER)
     : GLShader(type)
   {
     this->compile(source);
@@ -159,31 +159,23 @@ void launch_glsl_kernel(std::string source)
 {
   static bool gl_inited = false;
   if (!gl_inited) {
-    TI_INFO("[glsl] initializing GLUT for GL version 3.3");
-    // TODO: make context buffer here
-    glutInitContextVersion(3, 3);
-    glutInitContextProfile(GLUT_CORE_PROFILE);
-    glutInitContextFlags(GLUT_FORWARD_COMPATIBLE);
+    TI_INFO("[glsl] initializing GLFW (with hint GL 4.3)");
+    glfwInit();
+    glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 4);
+    glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
+    GLFWwindow *window = glfwCreateWindow(200, 200, "GLFW Window", nullptr, nullptr);
+    glfwMakeContextCurrent(window);
 #ifdef USE_GLEW
-    TI_INFO("[glsl] initializing GLEW {}", glewGetString(GLEW_VERSION));
+    TI_INFO("[glsl] initializing GLEW");
     GLint status = glewInit();
     if (status != GLEW_OK) {
       TI_ERROR("[glsl] cannot initialize GLEW: {}", glewGetErrorString(status));
     }
 #endif
     gl_inited = true;
   }
-  GLShader frag_shader(source, GL_FRAGMENT_SHADER);
-  GLShader vert_shader("#version 330 core\n\
-      void main()\n\
-      {\n\
-        gl_Position = vec4(0.0, 0.0, 0.0, 1.0);\n\
-      }\n\
-      ", GL_VERTEX_SHADER);
-  GLProgram program;
-  glViewport(0, 0, 2, 2);
-  program.attach(frag_shader);
-  program.attach(vert_shader);
+  GLShader shader(source);
+  GLProgram program(shader);
   program.link();
   program.use();
 
@@ -192,11 +184,11 @@ void launch_glsl_kernel(std::string source)
   vbo.bind_data(data, sizeof(data));
   vbo.set_attrib(GL_FLOAT, 0, 0, 1, 1 * sizeof(float));
   vbo.set_attrib(GL_FLOAT, 1, 1, 1, 1 * sizeof(float));
-  TI_INFO("[glsl] input data[0] = {}", data[0]);
-  TI_INFO("[glsl] output data[0] = {}", data[1]);
-  TI_INFO("[glsl] drawing arrays");
-  glDrawArrays(GL_POINTS, 0, 1);
-  TI_INFO("[glsl] output data[0] = {}", data[1]);
+  TI_INFO("[glsl] output before: {}", data[1]);
+  TI_INFO("[glsl] dispatching compute...");
+  // https://www.khronos.org/opengl/wiki/Compute_Shader
+  glDispatchCompute(1, 1, 1);
+  TI_INFO("[glsl] output after: {}", data[1]);
 }
 
 bool is_opengl_api_available()