[gui] Fix wide line support on macOS (#7205)

### Brief Summary

Changed the wide line rendering feature from rendering line primitives
to expanded quads. Added a "prepare" pass in GGUI to run the compute
shaders that expands & prepares the data.

Screenshot from macOS:

![image](https://user-images.githubusercontent.com/11663476/213070106-35f1d560-7a81-4760-af71-411b529c275a.png)

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

python/taichi/shaders/Lines_vk.frag

@@ -1,17 +1,9 @@
 #version 450
 
-layout(location = 0) in vec2 frag_texcoord;
+layout(location = 0) in vec3 color;
 
 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);
+  out_color = vec4(color, 1);
 }

python/taichi/shaders/Lines_vk.vert

@@ -1,30 +1,11 @@
 #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) in vec2 in_position;
+layout(location = 1) in uint in_color_encoded;
 
-layout(location = 0) out vec2 frag_texcoord;
-
-layout(binding = 0) uniform UniformBufferObject {
-  vec3 color;
-  int use_per_vertex_color;
-}
-ubo;
-
-layout(location = 1) out vec3 selected_color;
+layout(location = 0) out vec3 frag_color;
 
 void main() {
-  float x = in_position.x * 2.0 - 1.0;
-  float y = -(in_position.y * 2.0 - 1.0);
-
-  gl_Position = vec4(x, y, 0.0, 1.0);
-  frag_texcoord = in_texcoord;
-
-  if (ubo.use_per_vertex_color == 0) {
-    selected_color = ubo.color;
-  } else {
-    selected_color = in_color.rgb;
-  }
+  gl_Position = vec4(in_position.x * 2.0 - 1.0, -(in_position.y * 2.0 - 1.0), 0.0, 1.0);
+  frag_color = unpackUnorm4x8(in_color_encoded).rgb;
 }

python/taichi/shaders/SceneLines2quad_vk.comp

@@ -0,0 +1,124 @@
+#version 450
+
+/*
+ * struct Vertex {
+ *   vec3 pos;
+ *   vec3 normal;
+ *   vec2 tex_coord;
+ *   vec4 color;
+ * };
+ *
+ * We will assume pos is always first in struct.
+ */
+
+layout(binding = 0, std430) buffer InputVB {
+  float vb_in[];
+};
+
+layout(binding = 1, std430) buffer InputIB {
+  int ib_in[];
+};
+
+struct VertexOut {
+  vec4 out_pos;
+  vec4 out_color;
+};
+
+layout(binding = 2, std430) buffer OutputVB {
+  VertexOut vb_out[];
+};
+
+layout(binding = 3, std430) buffer OutputIB {
+  int ib_out[];
+};
+
+struct SceneUBO {
+  vec3 camera_pos;
+  mat4 view;
+  mat4 projection;
+  vec3 ambient_light;
+  int point_light_count;
+};
+
+layout(binding = 4) uniform UniformBufferObject {
+  SceneUBO scene;
+  vec3 color;
+  float line_width;
+  int per_vertex_color_offset;
+  int vertex_stride;
+  int start_vertex;
+  int start_index;
+  int num_vertices;
+  int is_indexed;
+  float aspect_ratio;
+}
+ubo;
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+void main() {
+  int thread_idx = int(gl_GlobalInvocationID.x);
+  if (thread_idx * 2 >= ubo.num_vertices) return;
+
+  uint vert0 = ubo.start_vertex;
+  uint vert1 = ubo.start_vertex;
+  if (ubo.is_indexed != 0) {
+    int index = thread_idx * 2 + ubo.start_index;
+    vert0 += ib_in[index];
+    vert1 += ib_in[index + 1];
+  } else {
+    vert0 += thread_idx * 2;
+    vert1 += thread_idx * 2 + 1;
+  }
+
+  vec4 pos0 = vec4(vb_in[vert0 * ubo.vertex_stride    ],
+                   vb_in[vert0 * ubo.vertex_stride + 1],
+                   vb_in[vert0 * ubo.vertex_stride + 2],
+                   1.0);
+  pos0 = ubo.scene.projection * ubo.scene.view * pos0;
+  vec4 color0 = vec4(ubo.color, 1.0);
+  vec4 pos1 = vec4(vb_in[vert1 * ubo.vertex_stride    ],
+                   vb_in[vert1 * ubo.vertex_stride + 1],
+                   vb_in[vert1 * ubo.vertex_stride + 2],
+                   1.0);
+  pos1 = ubo.scene.projection * ubo.scene.view * pos1;
+  vec4 color1 = vec4(ubo.color, 1.0);
+
+  if (ubo.per_vertex_color_offset > 0) {
+    color0 = vec4(vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset],
+                  vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset + 1],
+                  vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset + 2],
+                  1.0);
+    color1 = vec4(vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset],
+                  vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset + 1],
+                  vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset + 2],
+                  1.0);
+  }
+
+  vec2 aspect_adjust = vec2(ubo.aspect_ratio, 1.0);
+  vec2 dir = normalize((pos1.xy / pos1.w - pos0.xy / pos0.w) * aspect_adjust);
+  vec2 tangent = vec2(-dir.y, dir.x) * ubo.line_width / aspect_adjust;
+
+  vec4 out_pos00 = pos0;
+  vec4 out_pos01 = pos0;
+  out_pos00.xy += vec2(tangent * pos0.w);
+  out_pos01.xy -= vec2(tangent * pos0.w);
+
+  vec4 out_pos10 = pos1;
+  vec4 out_pos11 = pos1;
+  out_pos10.xy += vec2(tangent * pos1.w);
+  out_pos11.xy -= vec2(tangent * pos1.w);
+
+  vb_out[thread_idx * 4    ] = VertexOut(out_pos00, color0);
+  vb_out[thread_idx * 4 + 1] = VertexOut(out_pos01, color0);
+
+  vb_out[thread_idx * 4 + 2] = VertexOut(out_pos10, color1);
+  vb_out[thread_idx * 4 + 3] = VertexOut(out_pos11, color1);
+
+  ib_out[thread_idx * 6    ] = thread_idx * 4;
+  ib_out[thread_idx * 6 + 1] = thread_idx * 4 + 1;
+  ib_out[thread_idx * 6 + 2] = thread_idx * 4 + 2;
+
+  ib_out[thread_idx * 6 + 3] = thread_idx * 4 + 2;
+  ib_out[thread_idx * 6 + 4] = thread_idx * 4 + 1;
+  ib_out[thread_idx * 6 + 5] = thread_idx * 4 + 3;
+}

python/taichi/shaders/SceneLines_vk.frag

@@ -1,11 +1,9 @@
 #version 450
 
-layout(location = 0) in vec2 frag_texcoord;
+layout(location = 0) in vec3 color;
 
 layout(location = 0) out vec4 out_color;
 
-layout(location = 1) in vec3 selected_color;
-
 void main() {
-  out_color = vec4(selected_color, 1);
+  out_color = vec4(color, 1);
 }

python/taichi/shaders/SceneLines_vk.vert

@@ -1,38 +1,12 @@
 #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) in vec4 in_position;
+layout(location = 1) 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;
-};
-
-layout(binding = 0) uniform UBO {
-  SceneUBO scene;
-  vec3 color;
-  int use_per_vertex_color;
-}
-ubo;
-
-layout(location = 1) out vec3 selected_color;
+layout(location = 0) out vec3 frag_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;
-  }
+  gl_Position = in_position;
+  gl_Position.y = -gl_Position.y;
+  frag_color = in_color.rgb;
 }

python/taichi/shaders/lines2quad_vk.comp

@@ -0,0 +1,106 @@
+#version 450
+
+/*
+ * struct Vertex {
+ *   vec3 pos;
+ *   vec3 normal;
+ *   vec2 tex_coord;
+ *   vec4 color;
+ * };
+ *
+ * We will assume pos is always first in struct.
+ */
+
+layout(binding = 0, std430) buffer InputVB {
+  float vb_in[];
+};
+
+layout(binding = 1, std430) buffer InputIB {
+  int ib_in[];
+};
+
+struct VertexOut {
+  vec2 out_pos;
+  uint out_color; // Encoded RGBA8
+  uint pad;
+};
+
+layout(binding = 2, std430) buffer OutputVB {
+  VertexOut vb_out[];
+};
+
+layout(binding = 3, std430) buffer OutputIB {
+  int ib_out[];
+};
+
+layout(binding = 4) uniform UniformBufferObject {
+  vec3 color;
+  float line_width;
+  int per_vertex_color_offset;
+  int vertex_stride;
+  int start_vertex;
+  int start_index;
+  int num_vertices;
+  int is_indexed;
+  float aspect_ratio;
+}
+ubo;
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+void main() {
+  int thread_idx = int(gl_GlobalInvocationID.x);
+  if (thread_idx * 2 >= ubo.num_vertices) return;
+
+  uint vert0 = ubo.start_vertex;
+  uint vert1 = ubo.start_vertex;
+  if (ubo.is_indexed != 0) {
+    int index = thread_idx * 2 + ubo.start_index;
+    vert0 += ib_in[index];
+    vert1 += ib_in[index + 1];
+  } else {
+    vert0 += thread_idx * 2;
+    vert1 += thread_idx * 2 + 1;
+  }
+
+  vec2 pos0 = vec2(vb_in[vert0 * ubo.vertex_stride],
+                   vb_in[vert0 * ubo.vertex_stride + 1]);
+  vec4 color0 = vec4(ubo.color, 1.0);
+  vec2 pos1 = vec2(vb_in[vert1 * ubo.vertex_stride],
+                   vb_in[vert1 * ubo.vertex_stride + 1]);
+  vec4 color1 = vec4(ubo.color, 1.0);
+
+  if (ubo.per_vertex_color_offset > 0) {
+    color0 = vec4(vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset],
+                  vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset + 1],
+                  vb_in[vert0 * ubo.vertex_stride + ubo.per_vertex_color_offset + 2],
+                  1.0);
+    color1 = vec4(vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset],
+                  vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset + 1],
+                  vb_in[vert1 * ubo.vertex_stride + ubo.per_vertex_color_offset + 2],
+                  1.0);
+  }
+
+  vec2 aspect_adjust = vec2(ubo.aspect_ratio, 1.0);
+  vec2 dir = normalize((pos1 - pos0) * aspect_adjust);
+  vec2 tangent = vec2(-dir.y, dir.x) * (ubo.line_width * 0.5) / aspect_adjust;
+
+  vec2 out_pos00 = pos0 + tangent;
+  vec2 out_pos01 = pos0 - tangent;
+
+  vec2 out_pos10 = pos1 + tangent;
+  vec2 out_pos11 = pos1 - tangent;
+
+  vb_out[thread_idx * 4    ] = VertexOut(out_pos00, packUnorm4x8(color0), 0);
+  vb_out[thread_idx * 4 + 1] = VertexOut(out_pos01, packUnorm4x8(color0), 0);
+
+  vb_out[thread_idx * 4 + 2] = VertexOut(out_pos10, packUnorm4x8(color1), 0);
+  vb_out[thread_idx * 4 + 3] = VertexOut(out_pos11, packUnorm4x8(color1), 0);
+
+  ib_out[thread_idx * 6    ] = thread_idx * 4;
+  ib_out[thread_idx * 6 + 1] = thread_idx * 4 + 1;
+  ib_out[thread_idx * 6 + 2] = thread_idx * 4 + 2;
+
+  ib_out[thread_idx * 6 + 3] = thread_idx * 4 + 2;
+  ib_out[thread_idx * 6 + 4] = thread_idx * 4 + 1;
+  ib_out[thread_idx * 6 + 5] = thread_idx * 4 + 3;
+}

taichi/rhi/vulkan/vulkan_device.cpp

@@ -745,6 +745,8 @@ RhiReturn<vkapi::IVkDescriptorSet> VulkanResourceSet::finalize() {
   std::forward_list<VkDescriptorImageInfo> image_infos;
   std::vector<VkWriteDescriptorSet> desc_writes;
 
+  set_->ref_binding_objs.clear();
+
   for (auto &pair : bindings_) {
     uint32_t binding = pair.first;
     VkDescriptorType type = pair.second.type;

taichi/rhi/vulkan/vulkan_device_creator.cpp

@@ -628,12 +628,6 @@ void VulkanDeviceCreator::create_logical_device(bool manual_create) {
   if (device_supported_features.wideLines) {
     device_features.wideLines = true;
     ti_device_->vk_caps().wide_line = true;
-  } else if (params_.is_for_ui) {
-    if (!device_features.wideLines) {
-      RHI_LOG_ERROR(
-          "Taichi GGUI wide lines feature unavailable due to lack of device "
-          "support");
-    }
   }
 
   if (ti_device_->vk_caps().vk_api_version >= VK_API_VERSION_1_1) {

taichi/ui/backends/vulkan/renderable.cpp

@@ -234,7 +234,7 @@ void Renderable::create_vertex_buffer() {
 
   Device::AllocParams vb_params{buffer_size, false, false,
                                 app_context_->requires_export_sharing(),
-                                AllocUsage::Vertex};
+                                AllocUsage::Storage | AllocUsage::Vertex};
   vertex_buffer_ = app_context_->device().allocate_memory(vb_params);
 
   Device::AllocParams staging_vb_params{buffer_size, true, false, false,
@@ -248,7 +248,7 @@ void Renderable::create_index_buffer() {
 
   Device::AllocParams ib_params{buffer_size, false, false,
                                 app_context_->requires_export_sharing(),
-                                AllocUsage::Index};
+                                AllocUsage::Storage | AllocUsage::Index};
   index_buffer_ = app_context_->device().allocate_memory(ib_params);
 
   Device::AllocParams staging_ib_params{buffer_size, true, false, false,