[opengl] Texture support in OpenGL (#5296)

* use builtin types

* Texture support in OpenGL

* remove unused stuff

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

python/taichi/examples/rendering/simple_texture.py

@@ -7,7 +7,7 @@
 res = (512, 512)
 pixels = ti.Vector.field(3, dtype=float, shape=res)
 
-tex_format = ti.u8
+tex_format = ti.f32
 texture = ti.Texture(tex_format, 1, (128, 128))
 tex_ndarray = ti.ndarray(tex_format, shape=(128, 128))
 
@@ -16,7 +16,7 @@
 def make_texture(arr: ti.types.ndarray()):
     for i, j in ti.ndrange(128, 128):
         ret = taichi_logo(ti.Vector([i, j]) / 128)
-        ret = ti.cast(ret * 255, ti.u8)
+        # ret = ti.cast(ret * 255, ti.u8)
         arr[i, j] = ret
 
 
@@ -41,13 +41,12 @@ def paint(t: ti.f32, tex: ti.types.texture(num_dimensions=2)):
 
 def main():
 
-    window = ti.ui.Window('UV', res)
-    canvas = window.get_canvas()
+    window = ti.GUI('UV', res)
 
     t = 0.0
     while window.running:
         paint(t, texture)
-        canvas.set_image(pixels)
+        window.set_image(pixels)
         window.show()
         t += 0.03
 

taichi/codegen/spirv/spirv_types.cpp

@@ -384,15 +384,40 @@ class Translate2Spirv : public TypeVisitor {
   }
 
   void visit_int_type(const IntType *type) override {
-    SType vt = spir_builder_->get_null_type();
-    spir_builder_->declare_global(spv::OpTypeInt, vt, type->num_bits(),
-                                  type->is_signed() ? 1 : 0);
+    SType vt;
+    if (type->is_signed()) {
+      if (type->num_bits() == 8) {
+        vt = spir_builder_->i8_type();
+      } else if (type->num_bits() == 16) {
+        vt = spir_builder_->i16_type();
+      } else if (type->num_bits() == 32) {
+        vt = spir_builder_->i32_type();
+      } else if (type->num_bits() == 64) {
+        vt = spir_builder_->i64_type();
+      }
+    } else {
+      if (type->num_bits() == 8) {
+        vt = spir_builder_->u8_type();
+      } else if (type->num_bits() == 16) {
+        vt = spir_builder_->u16_type();
+      } else if (type->num_bits() == 32) {
+        vt = spir_builder_->u32_type();
+      } else if (type->num_bits() == 64) {
+        vt = spir_builder_->u64_type();
+      }
+    }
     ir_node_2_spv_value[type] = vt.id;
   }
 
   void visit_float_type(const FloatType *type) override {
-    SType vt = spir_builder_->get_null_type();
-    spir_builder_->declare_global(spv::OpTypeFloat, vt, type->num_bits());
+    SType vt;
+    if (type->num_bits() == 16) {
+      vt = spir_builder_->f16_type();
+    } else if (type->num_bits() == 32) {
+      vt = spir_builder_->f32_type();
+    } else if (type->num_bits() == 64) {
+      vt = spir_builder_->f64_type();
+    }
     ir_node_2_spv_value[type] = vt.id;
   }
 

taichi/rhi/opengl/opengl_device.cpp

@@ -222,7 +222,9 @@ void GLResourceBinder::image(uint32_t set,
                              uint32_t binding,
                              DeviceAllocation alloc,
                              ImageSamplerConfig sampler_config) {
-  TI_NOT_IMPLEMENTED;
+  TI_ASSERT_INFO(set == 0, "OpenGL only supports set = 0, requested set = {}",
+                 set);
+  texture_binding_map_[binding] = alloc.alloc_id;
 }
 
 void GLResourceBinder::vertex_buffer(DevicePtr ptr, uint32_t binding) {
@@ -330,6 +332,13 @@ void GLCommandList::bind_resources(ResourceBinder *_binder) {
     cmd->target = GL_UNIFORM_BUFFER;
     recorded_commands_.push_back(std::move(cmd));
   }
+  for (auto &[binding, texture] : binder->texture_binding_map()) {
+    auto cmd = std::make_unique<CmdBindTextureToIndex>();
+    cmd->texture = texture;
+    cmd->index = binding;
+    cmd->target = device_->get_image_gl_dims(texture);
+    recorded_commands_.push_back(std::move(cmd));
+  }
 }
 
 void GLCommandList::bind_resources(ResourceBinder *binder,
@@ -432,6 +441,7 @@ void GLCommandList::buffer_to_image(DeviceAllocation dst_img,
   cmd->image = dst_img.alloc_id;
   cmd->buffer = src_buf.alloc_id;
   cmd->offset = src_buf.offset;
+  cmd->device = device_;
   recorded_commands_.push_back(std::move(cmd));
 }
 
@@ -444,6 +454,7 @@ void GLCommandList::image_to_buffer(DevicePtr dst_buf,
   cmd->image = src_img.alloc_id;
   cmd->buffer = dst_buf.alloc_id;
   cmd->offset = dst_buf.offset;
+  cmd->device = device_;
   recorded_commands_.push_back(std::move(cmd));
 }
 
@@ -457,7 +468,7 @@ GLStream::~GLStream() {
 }
 
 std::unique_ptr<CommandList> GLStream::new_command_list() {
-  return std::make_unique<GLCommandList>();
+  return std::make_unique<GLCommandList>(device_);
 }
 
 StreamSemaphore GLStream::submit(
@@ -483,6 +494,9 @@ void GLStream::command_sync() {
   glFinish();
 }
 
+GLDevice::GLDevice() : stream_(this) {
+}
+
 GLDevice::~GLDevice() {
 }
 
@@ -718,6 +732,13 @@ void GLCommandList::CmdBindBufferToIndex::execute() {
   check_opengl_error("glBindBufferBase");
 }
 
+void GLCommandList::CmdBindTextureToIndex::execute() {
+  glActiveTexture(GL_TEXTURE0 + index);
+  check_opengl_error("glActiveTexture");
+  glBindTexture(GL_TEXTURE_2D, texture);
+  check_opengl_error("glBindTexture");
+}
+
 void GLCommandList::CmdBufferBarrier::execute() {
   glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
   check_opengl_error("glMemoryBarrier");
@@ -767,28 +788,31 @@ void GLCommandList::CmdImageTransition::execute() {
 }
 
 void GLCommandList::CmdBufferToImage::execute() {
-  auto image_dims = device->get_image_gl_dims(image);
-  auto image_format = device->get_image_gl_int_dims(image);
-  auto gl_type = gl_internal_format_to_type.at(image_format);
+  GLuint image_dims = device->get_image_gl_dims(image);
+  GLuint image_internal_format = device->get_image_gl_internal_format(image);
+  GLuint image_format = gl_internal_format_to_format.at(image_internal_format);
+  GLuint gl_type = gl_internal_format_to_type.at(image_internal_format);
 
   glBindTexture(image_dims, image);
   check_opengl_error("glBindTexture");
   glBindBuffer(GL_PIXEL_UNPACK_BUFFER, buffer);
   check_opengl_error("glBindBuffer");
   if (image_dims == GL_TEXTURE_1D) {
     glTexSubImage1D(image_dims, /*level=*/0, params.image_offset.x,
-                    params.image_extent.x, image_format, gl_type,
+                    params.image_extent.y, image_format, gl_type,
                     (void *)offset);
   } else if (image_dims == GL_TEXTURE_2D) {
-    glTexSubImage2D(image_dims, /*level=*/0, params.image_offset.x,
-                    params.image_offset.y, params.image_extent.x,
-                    params.image_extent.y, image_format, gl_type,
+    glTexSubImage2D(image_dims, /*level=*/0, /*xoffset=*/params.image_offset.x,
+                    /*yoffset=*/params.image_offset.y,
+                    /*width=*/params.image_extent.x,
+                    /*height=*/params.image_extent.y, image_format, gl_type,
                     (void *)offset);
   } else {
-    glTexSubImage3D(
-        image_dims, /*level=*/0, params.image_offset.x, params.image_offset.y,
-        params.image_offset.z, params.image_extent.x, params.image_extent.y,
-        params.image_extent.z, image_format, gl_type, (void *)offset);
+    glTexSubImage3D(image_dims, /*level=*/0, /*xoffset=*/params.image_offset.x,
+                    params.image_offset.y, params.image_offset.z,
+                    params.image_extent.x, params.image_extent.y,
+                    params.image_extent.z, image_format, gl_type,
+                    (void *)offset);
   }
   check_opengl_error("glTexSubImage");
   glBindTexture(image_dims, /*target=*/0);
@@ -797,7 +821,7 @@ void GLCommandList::CmdBufferToImage::execute() {
 
 void GLCommandList::CmdImageToBuffer::execute() {
   auto image_dims = device->get_image_gl_dims(image);
-  auto image_format = device->get_image_gl_int_dims(image);
+  auto image_format = device->get_image_gl_internal_format(image);
   auto gl_type = gl_internal_format_to_type.at(image_format);
   auto unsized_format = gl_internal_format_to_format.at(image_format);
 

taichi/rhi/opengl/opengl_device.h

@@ -61,9 +61,14 @@ class GLResourceBinder : public ResourceBinder {
     return ubo_binding_map_;
   }
 
+  const std::unordered_map<uint32_t, GLuint> &texture_binding_map() {
+    return texture_binding_map_;
+  }
+
  private:
   std::unordered_map<uint32_t, GLuint> ssbo_binding_map_;
   std::unordered_map<uint32_t, GLuint> ubo_binding_map_;
+  std::unordered_map<uint32_t, GLuint> texture_binding_map_;
 };
 
 class GLPipeline : public Pipeline {
@@ -84,6 +89,8 @@ class GLPipeline : public Pipeline {
 
 class GLCommandList : public CommandList {
  public:
+  GLCommandList(GLDevice *device) : device_(device) {
+  }
   ~GLCommandList() override;
 
   void bind_pipeline(Pipeline *p) override;
@@ -150,6 +157,13 @@ class GLCommandList : public CommandList {
     void execute() override;
   };
 
+  struct CmdBindTextureToIndex : public Cmd {
+    GLuint texture{0};
+    GLuint index{0};
+    GLenum target{GL_TEXTURE_2D};
+    void execute() override;
+  };
+
   struct CmdBufferBarrier : public Cmd {
     void execute() override;
   };
@@ -197,10 +211,13 @@ class GLCommandList : public CommandList {
   };
 
   std::vector<std::unique_ptr<Cmd>> recorded_commands_;
+  GLDevice *device_{nullptr};
 };
 
 class GLStream : public Stream {
  public:
+  GLStream(GLDevice *device) : device_(device) {
+  }
   ~GLStream() override;
 
   std::unique_ptr<CommandList> new_command_list() override;
@@ -212,10 +229,14 @@ class GLStream : public Stream {
       const std::vector<StreamSemaphore> &wait_semaphores = {}) override;
 
   void command_sync() override;
+
+ private:
+  GLDevice *device_{nullptr};
 };
 
 class GLDevice : public GraphicsDevice {
  public:
+  GLDevice();
   ~GLDevice() override;
 
   DeviceAllocation allocate_memory(const AllocParams &params) override;
@@ -269,7 +290,7 @@ class GLDevice : public GraphicsDevice {
     return image_to_dims_.at(image);
   }
 
-  GLuint get_image_gl_int_dims(GLuint image) const {
+  GLuint get_image_gl_internal_format(GLuint image) const {
     return image_to_int_format_.at(image);
   }