[vulkan] Device API blending (#4815)

* Add actual blending control in device API and disable blending for set_image

* [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>

Author: bobcao3

taichi/backends/device.h

@@ -48,6 +48,21 @@ enum class DeviceCapability : uint32_t {
   wide_lines
 };
 
+enum class BlendOp : uint32_t { add, subtract, reverse_subtract, min, max };
+
+enum class BlendFactor : uint32_t {
+  zero,
+  one,
+  src_color,
+  one_minus_src_color,
+  dst_color,
+  one_minus_dst_color,
+  src_alpha,
+  one_minus_src_alpha,
+  dst_alpha,
+  one_minus_dst_alpha
+};
+
 class Device;
 struct DeviceAllocation;
 struct DevicePtr;
@@ -529,13 +544,26 @@ struct ImageParams {
   bool export_sharing{false};
 };
 
+struct BlendFunc {
+  BlendOp op{BlendOp::add};
+  BlendFactor src_factor{BlendFactor::src_alpha};
+  BlendFactor dst_factor{BlendFactor::one_minus_src_alpha};
+};
+
+struct BlendingParams {
+  bool enable{true};
+  BlendFunc color;
+  BlendFunc alpha;
+};
+
 struct RasterParams {
   TopologyType prim_topology{TopologyType::Triangles};
   PolygonMode polygon_mode{PolygonMode::Fill};
   bool front_face_cull{false};
   bool back_face_cull{false};
   bool depth_test{false};
   bool depth_write{false};
+  std::vector<BlendingParams> blending{};
 };
 
 class GraphicsDevice : public Device {

taichi/backends/vulkan/vulkan_device.cpp

@@ -105,6 +105,40 @@ VkImageLayout image_layout_ti_to_vk(ImageLayout layout) {
   return image_layout_ti_2_vk.at(layout);
 }
 
+const std::unordered_map<BlendOp, VkBlendOp> blend_op_ti_2_vk = {
+    {BlendOp::add, VK_BLEND_OP_ADD},
+    {BlendOp::subtract, VK_BLEND_OP_SUBTRACT},
+    {BlendOp::reverse_subtract, VK_BLEND_OP_REVERSE_SUBTRACT},
+    {BlendOp::min, VK_BLEND_OP_MIN},
+    {BlendOp::max, VK_BLEND_OP_MAX}};
+
+VkBlendOp blend_op_ti_to_vk(BlendOp op) {
+  if (blend_op_ti_2_vk.find(op) == blend_op_ti_2_vk.end()) {
+    TI_ERROR("BlendOp cannot be mapped to vk");
+  }
+  return blend_op_ti_2_vk.at(op);
+}
+
+const std::unordered_map<BlendFactor, VkBlendFactor> blend_factor_ti_2_vk = {
+    {BlendFactor::zero, VK_BLEND_FACTOR_ZERO},
+    {BlendFactor::one, VK_BLEND_FACTOR_ONE},
+    {BlendFactor::src_color, VK_BLEND_FACTOR_SRC_COLOR},
+    {BlendFactor::one_minus_src_color, VK_BLEND_FACTOR_ONE_MINUS_SRC_COLOR},
+    {BlendFactor::dst_color, VK_BLEND_FACTOR_DST_COLOR},
+    {BlendFactor::one_minus_dst_color, VK_BLEND_FACTOR_ONE_MINUS_DST_COLOR},
+    {BlendFactor::src_alpha, VK_BLEND_FACTOR_SRC_ALPHA},
+    {BlendFactor::one_minus_src_alpha, VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA},
+    {BlendFactor::dst_alpha, VK_BLEND_FACTOR_DST_ALPHA},
+    {BlendFactor::one_minus_dst_alpha, VK_BLEND_FACTOR_ONE_MINUS_DST_ALPHA},
+};
+
+VkBlendFactor blend_factor_ti_to_vk(BlendFactor factor) {
+  if (blend_factor_ti_2_vk.find(factor) == blend_factor_ti_2_vk.end()) {
+    TI_ERROR("BlendFactor cannot be mapped to vk");
+  }
+  return blend_factor_ti_2_vk.at(factor);
+}
+
 VulkanPipeline::VulkanPipeline(const Params &params)
     : device_(params.device->vk_device()), name_(params.name) {
   create_descriptor_set_layout(params);
@@ -124,6 +158,9 @@ VulkanPipeline::VulkanPipeline(
     const std::vector<VertexInputBinding> &vertex_inputs,
     const std::vector<VertexInputAttribute> &vertex_attrs)
     : device_(params.device->vk_device()), name_(params.name) {
+  this->graphics_pipeline_template_ =
+      std::make_unique<GraphicsPipelineTemplate>();
+
   create_descriptor_set_layout(params);
   create_shader_stages(params);
   create_pipeline_layout();
@@ -159,38 +196,12 @@ vkapi::IVkPipeline VulkanPipeline::graphics_pipeline(
     return graphics_pipeline_.at(renderpass);
   }
 
-  std::vector<VkPipelineColorBlendAttachmentState> blend_attachments(
-      renderpass_desc.color_attachments.size());
-  for (int i = 0; i < renderpass_desc.color_attachments.size(); i++) {
-    blend_attachments[i].colorWriteMask =
-        VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
-        VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
-    blend_attachments[i].blendEnable = VK_TRUE;
-    blend_attachments[i].srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
-    blend_attachments[i].dstColorBlendFactor =
-        VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
-    blend_attachments[i].colorBlendOp = VK_BLEND_OP_ADD;
-    blend_attachments[i].srcAlphaBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
-    blend_attachments[i].dstAlphaBlendFactor =
-        VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
-    blend_attachments[i].alphaBlendOp = VK_BLEND_OP_ADD;
-  }
-
-  graphics_pipeline_template_->color_blending.attachmentCount =
-      renderpass_desc.color_attachments.size();
-  graphics_pipeline_template_->color_blending.pAttachments =
-      blend_attachments.data();
-
   vkapi::IVkPipeline pipeline = vkapi::create_graphics_pipeline(
       device_, &graphics_pipeline_template_->pipeline_info, renderpass,
       pipeline_layout_);
 
   graphics_pipeline_[renderpass] = pipeline;
 
-  graphics_pipeline_template_->color_blending.attachmentCount = 0;
-  graphics_pipeline_template_->color_blending.pAttachments = nullptr;
-  graphics_pipeline_template_->pipeline_info.renderPass = VK_NULL_HANDLE;
-
   return pipeline;
 }
 
@@ -252,6 +263,39 @@ void VulkanPipeline::create_descriptor_set_layout(const Params &params) {
     //     TI_WARN("attrib {}:{}", location, type->type_name);
     //   }
     // }
+
+    if (code_view.stage == VK_SHADER_STAGE_FRAGMENT_BIT) {
+      uint32_t render_target_count = 0;
+      result = spvReflectEnumerateOutputVariables(&module, &render_target_count,
+                                                  nullptr);
+      TI_ASSERT(result == SPV_REFLECT_RESULT_SUCCESS);
+
+      std::vector<SpvReflectInterfaceVariable *> variables(render_target_count);
+      result = spvReflectEnumerateOutputVariables(&module, &render_target_count,
+                                                  variables.data());
+
+      render_target_count = 0;
+
+      for (auto var : variables) {
+        // We want to remove auxiliary outputs such as frag depth
+        if (var->built_in == -1) {
+          render_target_count++;
+        }
+      }
+
+      graphics_pipeline_template_->blend_attachments.resize(
+          render_target_count);
+
+      VkPipelineColorBlendAttachmentState default_state{};
+      default_state.colorWriteMask =
+          VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
+          VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
+      default_state.blendEnable = VK_FALSE;
+
+      std::fill(graphics_pipeline_template_->blend_attachments.begin(),
+                graphics_pipeline_template_->blend_attachments.end(),
+                default_state);
+    }
   }
 
   for (uint32_t set : sets_used) {
@@ -295,9 +339,6 @@ void VulkanPipeline::create_graphics_pipeline(
     const RasterParams &raster_params,
     const std::vector<VertexInputBinding> &vertex_inputs,
     const std::vector<VertexInputAttribute> &vertex_attrs) {
-  this->graphics_pipeline_template_ =
-      std::make_unique<GraphicsPipelineTemplate>();
-
   // Use dynamic viewport state. These two are just dummies
   VkViewport viewport;
   viewport.width = 1;
@@ -411,13 +452,45 @@ void VulkanPipeline::create_graphics_pipeline(
       VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
   color_blending.logicOpEnable = VK_FALSE;
   color_blending.logicOp = VK_LOGIC_OP_COPY;
-  color_blending.attachmentCount = 0;
-  color_blending.pAttachments = nullptr;  // Filled in later
+  color_blending.attachmentCount =
+      graphics_pipeline_template_->blend_attachments.size();
+  color_blending.pAttachments =
+      graphics_pipeline_template_->blend_attachments.data();
   color_blending.blendConstants[0] = 0.0f;
   color_blending.blendConstants[1] = 0.0f;
   color_blending.blendConstants[2] = 0.0f;
   color_blending.blendConstants[3] = 0.0f;
 
+  if (raster_params.blending.size()) {
+    TI_ASSERT_INFO(raster_params.blending.size() ==
+                       graphics_pipeline_template_->blend_attachments.size(),
+                   "RasterParams::blending (size={}) must either be zero sized "
+                   "or match the number of fragment shader outputs (size={}).",
+                   raster_params.blending.size(),
+                   graphics_pipeline_template_->blend_attachments.size());
+
+    for (int i = 0; i < raster_params.blending.size(); i++) {
+      auto &state = graphics_pipeline_template_->blend_attachments[i];
+      auto &ti_param = raster_params.blending[i];
+      state.blendEnable = ti_param.enable;
+      if (ti_param.enable) {
+        state.colorBlendOp = blend_op_ti_to_vk(ti_param.color.op);
+        state.srcColorBlendFactor =
+            blend_factor_ti_to_vk(ti_param.color.src_factor);
+        state.dstColorBlendFactor =
+            blend_factor_ti_to_vk(ti_param.color.dst_factor);
+        state.alphaBlendOp = blend_op_ti_to_vk(ti_param.alpha.op);
+        state.srcAlphaBlendFactor =
+            blend_factor_ti_to_vk(ti_param.alpha.src_factor);
+        state.dstAlphaBlendFactor =
+            blend_factor_ti_to_vk(ti_param.alpha.dst_factor);
+        state.colorWriteMask =
+            VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
+            VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
+      }
+    }
+  }
+
   VkPipelineDynamicStateCreateInfo &dynamic_state =
       graphics_pipeline_template_->dynamic_state;
   dynamic_state.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;

taichi/backends/vulkan/vulkan_device.h

@@ -313,6 +313,7 @@ class VulkanPipeline : public Pipeline {
     VkPipelineMultisampleStateCreateInfo multisampling{};
     VkPipelineDepthStencilStateCreateInfo depth_stencil{};
     VkPipelineColorBlendStateCreateInfo color_blending{};
+    std::vector<VkPipelineColorBlendAttachmentState> blend_attachments{};
     std::vector<VkDynamicState> dynamic_state_enables = {
         VK_DYNAMIC_STATE_LINE_WIDTH, VK_DYNAMIC_STATE_VIEWPORT,
         VK_DYNAMIC_STATE_SCISSOR};

taichi/ui/backends/vulkan/renderable.cpp

@@ -145,6 +145,10 @@ void Renderable::create_graphics_pipeline() {
   raster_params.depth_test = true;
   raster_params.depth_write = true;
 
+  if (config_.blending) {
+    raster_params.blending.push_back(BlendingParams());
+  }
+
   std::vector<VertexInputBinding> vertex_inputs = {
       {/*binding=*/0, config_.vbo_size(), /*instance=*/false}};
   // TODO: consider using uint8 for colors and normals

taichi/ui/backends/vulkan/renderable.h

@@ -31,6 +31,7 @@ struct RenderableConfig {
   int indices_count{0};
   size_t ubo_size{0};
   size_t ssbo_size{0};
+  bool blending{false};
   std::string vertex_shader_path;
   std::string fragment_shader_path;
   taichi::lang::TopologyType topology_type{

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

@@ -26,6 +26,7 @@ void Circles::init_circles(AppContext *app_context,
       1,
       sizeof(UniformBufferObject),
       0,
+      true,
       app_context->config.package_path + "/shaders/Circles_vk_vert.spv",
       app_context->config.package_path + "/shaders/Circles_vk_frag.spv",
       TopologyType::Points,

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

@@ -29,6 +29,7 @@ void Lines::init_lines(AppContext *app_context,
       indices_count,
       sizeof(UniformBufferObject),
       0,
+      true,
       app_context->config.package_path + "/shaders/Lines_vk_vert.spv",
       app_context->config.package_path + "/shaders/Lines_vk_frag.spv",
       TopologyType::Lines,

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

@@ -53,6 +53,7 @@ void Mesh::init_mesh(AppContext *app_context,
       indices_count,
       sizeof(UniformBufferObject),
       1,
+      true,
       app_context->config.package_path + "/shaders/Mesh_vk_vert.spv",
       app_context->config.package_path + "/shaders/Mesh_vk_frag.spv",
       TopologyType::Triangles,

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

@@ -58,6 +58,7 @@ void Particles::init_particles(AppContext *app_context,
       1,
       sizeof(UniformBufferObject),
       1,
+      true,
       app_context->config.package_path + "/shaders/Particles_vk_vert.spv",
       app_context->config.package_path + "/shaders/Particles_vk_frag.spv",
       TopologyType::Points,

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

@@ -113,6 +113,7 @@ void SetImage::init_set_image(AppContext *app_context,
       6,
       sizeof(UniformBufferObject),
       0,
+      false,
       app_context->config.package_path + "/shaders/SetImage_vk_vert.spv",
       app_context->config.package_path + "/shaders/SetImage_vk_frag.spv",
       TopologyType::Triangles,

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

@@ -25,6 +25,7 @@ void Triangles::init_triangles(AppContext *app_context,
       indices_count,
       sizeof(UniformBufferObject),
       0,
+      true,
       app_context->config.package_path + "/shaders/Triangles_vk_vert.spv",
       app_context->config.package_path + "/shaders/Triangles_vk_frag.spv",
       TopologyType::Triangles,