[gui] [vulkan] Support 3d-lines rendering

python/taichi/shaders/SceneLines_vk.frag

@@ -0,0 +1,17 @@
+#version 450
+
+layout(location = 0) in vec2 frag_texcoord;
+
+layout(location = 0) out vec4 out_color;
+
+layout(binding = 0) uniform UniformBufferObject {
+  vec3 color;
+  int use_per_vertex_color;
+}
+ubo;
+
+layout(location = 1) in vec3 selected_color;
+
+void main() {
+  out_color = vec4(selected_color, 1);
+}

python/taichi/shaders/SceneLines_vk.vert

@@ -0,0 +1,43 @@
+#version 450
+
+layout(location = 0) in vec3 in_position;
+layout(location = 1) in vec3 in_normal;
+layout(location = 2) in vec2 in_texcoord;
+layout(location = 3) in vec4 in_color;
+
+layout(location = 0) out vec2 frag_texcoord;
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+struct PointLight {
+  vec3 pos;
+  vec3 color;
+};
+
+layout(binding = 0) uniform UBO {
+  SceneUBO scene;
+  vec3 color;
+  int use_per_vertex_color;
+}
+ubo;
+
+layout(location = 1) out vec3 selected_color;
+
+void main() {
+  gl_Position = ubo.scene.projection * ubo.scene.view * vec4(in_position, 1.0);
+  gl_Position.y *= -1.0;
+
+  frag_texcoord = in_texcoord;
+
+  if (ubo.use_per_vertex_color == 0) {
+    selected_color = ubo.color;
+  } else {
+    selected_color = in_color.rgb;
+  }
+}

python/taichi/ui/scene.py

@@ -88,6 +88,35 @@ def set_camera(self, camera):
         """
         self.scene.set_camera(camera.ptr)
 
+    def lines(self,
+              vertices,
+              width,
+              indices=None,
+              color=(0.5, 0.5, 0.5),
+              per_vertex_color=None):
+        """Declare multi-lines inside the scene.
+        Args:
+            vertices: a taichi 3D Vector field, where each element indicate the
+                3D location of points of lines.
+            width: the line width (maybe different on different systems).
+            indices: a taichi int field of shape (2 * #points), which indicate
+                the points indices of the lines. If this is None, then it is
+                assumed that the points are already arranged in lines order.
+            color: a global color of the mesh as 3 floats representing RGB values.
+                If `per_vertex_color` is provided, this is ignored.
+            per_vertex_color (Tuple[float]): a taichi 3D vector field, where each
+                element indicate the RGB color of the line.
+        """
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
+        vbo_info = get_field_info(vbo)
+        indices_info = get_field_info(indices)
+        self.scene.lines(vbo_info, indices_info, has_per_vertex_color, color,
+                         width)
+
     def mesh(self,
              vertices,
              indices=None,

taichi/python/export_ggui.cpp

@@ -136,6 +136,24 @@ struct PyScene {
     scene->set_camera(camera.camera);
   }
 
+  void lines(FieldInfo vbo,
+             FieldInfo indices,
+             bool has_per_vertex_color,
+             py::tuple color_,
+             float width) {
+    RenderableInfo renderable_info;
+    renderable_info.vbo = vbo;
+    renderable_info.indices = indices;
+    renderable_info.has_per_vertex_color = has_per_vertex_color;
+
+    SceneLinesInfo info;
+    info.renderable_info = renderable_info;
+    info.color = tuple_to_vec3(color_);
+    info.width = width;
+
+    return scene->lines(info);
+  }
+
   void mesh(FieldInfo vbo,
             bool has_per_vertex_color,
             FieldInfo indices,
@@ -433,6 +451,7 @@ void export_ggui(py::module &m) {
   py::class_<PyScene>(m, "PyScene")
       .def(py::init<>())
       .def("set_camera", &PyScene::set_camera)
+      .def("lines", &PyScene::lines)
       .def("mesh", &PyScene::mesh)
       .def("particles", &PyScene::particles)
       .def("point_light", &PyScene::point_light)

taichi/ui/backends/vulkan/renderables/scene_lines.cpp

@@ -0,0 +1,84 @@
+#include "scene_lines.h"
+
+#include "taichi/ui/utils/utils.h"
+#include "taichi/rhi/vulkan/vulkan_device.h"
+
+TI_UI_NAMESPACE_BEGIN
+
+namespace vulkan {
+
+using namespace taichi::lang;
+using namespace taichi::lang::vulkan;
+
+void SceneLines::update_data(const SceneLinesInfo &info, const Scene &scene) {
+  Renderable::update_data(info.renderable_info);
+
+  update_ubo(info, scene);
+
+  curr_width_ = info.width;
+}
+
+void SceneLines::init_scene_lines(AppContext *app_context,
+                                  int vertices_count,
+                                  int indices_count) {
+  RenderableConfig config = {
+      vertices_count,
+      indices_count,
+      vertices_count,
+      indices_count,
+      vertices_count,
+      0,
+      indices_count,
+      0,
+      sizeof(UniformBufferObject),
+      0,
+      true,
+      app_context->config.package_path + "/shaders/SceneLines_vk_vert.spv",
+      app_context->config.package_path + "/shaders/SceneLines_vk_frag.spv",
+      TopologyType::Lines,
+  };
+
+  Renderable::init(config, app_context);
+  Renderable::init_render_resources();
+}
+
+SceneLines::SceneLines(AppContext *app_context, VertexAttributes vbo_attrs) {
+  init_scene_lines(app_context, 4, 6);
+}
+
+void SceneLines::update_ubo(const SceneLinesInfo &info, const Scene &scene) {
+  UniformBufferObject ubo{};
+  ubo.scene = scene.current_ubo_;
+  ubo.color = info.color;
+  ubo.use_per_vertex_color = info.renderable_info.has_per_vertex_color;
+  void *mapped = app_context_->device().map(uniform_buffer_);
+  memcpy(mapped, &ubo, sizeof(ubo));
+  app_context_->device().unmap(uniform_buffer_);
+}
+
+void SceneLines::create_bindings() {
+  Renderable::create_bindings();
+  ResourceBinder *binder = pipeline_->resource_binder();
+  binder->buffer(0, 0, uniform_buffer_);
+  binder->rw_buffer(0, 1, storage_buffer_);
+}
+
+void SceneLines::record_this_frame_commands(CommandList *command_list) {
+  command_list->bind_pipeline(pipeline_.get());
+  command_list->bind_resources(pipeline_->resource_binder());
+  command_list->set_line_width(curr_width_);
+
+  if (indexed_) {
+    command_list->draw_indexed(config_.indices_count, 0, 0);
+  } else {
+    command_list->draw(config_.vertices_count, 0);
+  }
+}
+
+void SceneLines::cleanup() {
+  Renderable::cleanup();
+}
+
+}  // namespace vulkan
+
+TI_UI_NAMESPACE_END

taichi/ui/backends/vulkan/renderables/scene_lines.h

@@ -0,0 +1,59 @@
+#pragma once
+
+#include <iostream>
+#include <fstream>
+#include <stdexcept>
+#include <algorithm>
+#include <chrono>
+#include <vector>
+#include <cstring>
+#include <cstdlib>
+#include <cstdint>
+#include <array>
+#include <optional>
+#include <set>
+#include "taichi/ui/utils/utils.h"
+#include "taichi/ui/backends/vulkan/vertex.h"
+
+#include "taichi/ui/backends/vulkan/app_context.h"
+#include "taichi/ui/backends/vulkan/swap_chain.h"
+#include "taichi/ui/backends/vulkan/renderable.h"
+#include "taichi/ui/common/field_info.h"
+#include "taichi/ui/backends/vulkan/scene.h"
+
+TI_UI_NAMESPACE_BEGIN
+
+namespace vulkan {
+
+class SceneLines final : public Renderable {
+ public:
+  SceneLines(AppContext *app_context, VertexAttributes vbo_attrs);
+
+  void update_data(const SceneLinesInfo &info, const Scene &scene);
+
+  virtual void record_this_frame_commands(
+      taichi::lang::CommandList *command_list) override;
+
+ private:
+  struct UniformBufferObject {
+    Scene::SceneUniformBuffer scene;
+    alignas(16) glm::vec3 color;
+    int use_per_vertex_color;
+  };
+
+  void init_scene_lines(AppContext *app_context,
+                        int vertices_count,
+                        int indices_count);
+
+  void update_ubo(const SceneLinesInfo &info, const Scene &scene);
+
+  virtual void cleanup() override;
+
+  virtual void create_bindings() override;
+
+  float curr_width_;
+};
+
+}  // namespace vulkan
+
+TI_UI_NAMESPACE_END

taichi/ui/backends/vulkan/renderer.cpp

@@ -70,6 +70,13 @@ void Renderer::circles(const CirclesInfo &info) {
   next_renderable_ += 1;
 }
 
+void Renderer::scene_lines(const SceneLinesInfo &info, Scene *scene) {
+  SceneLines *scene_lines =
+      get_renderable_of_type<SceneLines>(info.renderable_info.vbo_attrs);
+  scene_lines->update_data(info, *scene);
+  next_renderable_ += 1;
+}
+
 void Renderer::mesh(const MeshInfo &info, Scene *scene) {
   Mesh *mesh = get_renderable_of_type<Mesh>(info.renderable_info.vbo_attrs);
   mesh->update_data(info, *scene);
@@ -90,9 +97,12 @@ void Renderer::scene(Scene *scene) {
   float aspect_ratio = swap_chain_.width() / (float)swap_chain_.height();
   scene->update_ubo(aspect_ratio);
 
-  int object_count = scene->mesh_infos_.size() + scene->particles_infos_.size();
+  int object_count = scene->mesh_infos_.size() +
+                     scene->particles_infos_.size() +
+                     scene->scene_lines_infos_.size();
   int mesh_id = 0;
   int particles_id = 0;
+  int scene_lines_id = 0;
   for (int i = 0; i < object_count; ++i) {
     if (mesh_id < scene->mesh_infos_.size() &&
         scene->mesh_infos_[mesh_id].object_id == i) {
@@ -104,10 +114,17 @@ void Renderer::scene(Scene *scene) {
       particles(scene->particles_infos_[particles_id], scene);
       ++particles_id;
     }
+    // Scene Lines
+    if (scene_lines_id < scene->scene_lines_infos_.size() &&
+        scene->scene_lines_infos_[scene_lines_id].object_id == i) {
+      scene_lines(scene->scene_lines_infos_[scene_lines_id], scene);
+      ++scene_lines_id;
+    }
   }
   scene->next_object_id_ = 0;
   scene->mesh_infos_.clear();
   scene->particles_infos_.clear();
+  scene->scene_lines_infos_.clear();
   scene->point_lights_.clear();
 }
 

taichi/ui/backends/vulkan/renderer.h

@@ -30,6 +30,7 @@
 #include "renderables/particles.h"
 #include "renderables/circles.h"
 #include "renderables/lines.h"
+#include "renderables/scene_lines.h"
 
 namespace taichi {
 namespace lang {
@@ -63,6 +64,8 @@ class TI_DLL_EXPORT Renderer {
 
   void particles(const ParticlesInfo &info, Scene *scene);
 
+  void scene_lines(const SceneLinesInfo &info, Scene *scene);
+
   void scene(Scene *scene);
 
   void draw_frame(Gui *gui);

taichi/ui/backends/vulkan/scene.h

@@ -10,6 +10,7 @@ class Scene final : public SceneBase {
   friend class Renderer;
   friend class Particles;
   friend class Mesh;
+  friend class SceneLines;
 
  private:
   struct SceneUniformBuffer {

taichi/ui/common/scene_base.h

@@ -28,12 +28,22 @@ struct ParticlesInfo {
   int object_id{0};
 };
 
+struct SceneLinesInfo {
+  RenderableInfo renderable_info;
+  glm::vec3 color;
+  float width{0};
+  int object_id{0};
+};
+
 class SceneBase {
  public:
   void set_camera(const Camera &camera) {
     camera_ = camera;
   }
-
+  void lines(const SceneLinesInfo &info) {
+    scene_lines_infos_.push_back(info);
+    scene_lines_infos_.back().object_id = next_object_id_++;
+  }
   void mesh(const MeshInfo &info) {
     mesh_infos_.push_back(info);
     mesh_infos_.back().object_id = next_object_id_++;
@@ -54,6 +64,7 @@ class SceneBase {
   Camera camera_;
   glm::vec3 ambient_light_color_ = glm::vec3(0.1, 0.1, 0.1);
   std::vector<PointLight> point_lights_;
+  std::vector<SceneLinesInfo> scene_lines_infos_;
   std::vector<MeshInfo> mesh_infos_;
   std::vector<ParticlesInfo> particles_infos_;
   int next_object_id_ = 0;