[Impeller] kick off registration and initial PSO compilation of runtime effect earlier. (#52381)

I thought about changing this API so that it blocks on compilation, but I think that isn't necesasary as the workload can happen while the UI thread is building. Plus, even blocking on shader completion does not guarantee that we won't need to compile a variant anyway.

While I was at it I moved the descriptor sets computation to the runtime effect constructor.

With this change the pipleine variant used during the frame is produced much faster than the initial 12ms compilation.

![image](https://github.com/flutter/engine/assets/8975114/42626676-0a71-4503-a3e9-4109c36fbbef)

Fixes flutter/flutter#113719
Fixes flutter/flutter#141222

Author: jonahwilliams

impeller/compiler/reflector.cc

@@ -348,6 +348,8 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
     uniform_description.name = compiler_->get_name(var.self);
     uniform_description.location = compiler_->get_decoration(
         var.self, spv::Decoration::DecorationLocation);
+    uniform_description.binding =
+        compiler_->get_decoration(var.self, spv::Decoration::DecorationBinding);
     uniform_description.type = spir_type.basetype;
     uniform_description.rows = spir_type.vecsize;
     uniform_description.columns = spir_type.columns;
@@ -410,6 +412,7 @@ std::shared_ptr<RuntimeStageData::Shader> Reflector::GenerateRuntimeStageData()
         .name = ubo.name,
         .location = 64,  // Magic constant that must match the descriptor set
                          // location for fragment programs.
+        .binding = 64,
         .type = spirv_cross::SPIRType::Struct,
         .struct_layout = std::move(struct_layout),
         .struct_float_count = float_count,

impeller/compiler/types.h

@@ -48,6 +48,7 @@ enum class SourceLanguage {
 struct UniformDescription {
   std::string name;
   size_t location = 0u;
+  size_t binding = 0u;
   spirv_cross::SPIRType::BaseType type = spirv_cross::SPIRType::BaseType::Float;
   size_t rows = 0u;
   size_t columns = 0u;

impeller/core/runtime_types.h

@@ -39,6 +39,8 @@ struct RuntimeUniformDimensions {
 struct RuntimeUniformDescription {
   std::string name;
   size_t location = 0u;
+  /// Location, but for Vulkan.
+  size_t binding = 0u;
   RuntimeUniformType type = RuntimeUniformType::kFloat;
   RuntimeUniformDimensions dimensions = {};
   size_t bit_width = 0u;

impeller/entity/contents/runtime_effect_contents.cc

@@ -4,6 +4,7 @@
 
 #include "impeller/entity/contents/runtime_effect_contents.h"
 
+#include <algorithm>
 #include <future>
 #include <memory>
 
@@ -19,6 +20,7 @@
 #include "impeller/renderer/pipeline_library.h"
 #include "impeller/renderer/render_pass.h"
 #include "impeller/renderer/shader_function.h"
+#include "impeller/renderer/vertex_descriptor.h"
 
 namespace impeller {
 
@@ -65,18 +67,22 @@ static std::shared_ptr<ShaderMetadata> MakeShaderMetadata(
   return metadata;
 }
 
-bool RuntimeEffectContents::Render(const ContentContext& renderer,
-                                   const Entity& entity,
-                                   RenderPass& pass) const {
+bool RuntimeEffectContents::BootstrapShader(
+    const ContentContext& renderer) const {
+  if (!RegisterShader(renderer)) {
+    return false;
+  }
+  ContentContextOptions options;
+  options.color_attachment_pixel_format =
+      renderer.GetContext()->GetCapabilities()->GetDefaultColorFormat();
+  return !!CreatePipeline(renderer, options);
+}
+
+bool RuntimeEffectContents::RegisterShader(
+    const ContentContext& renderer) const {
   const std::shared_ptr<Context>& context = renderer.GetContext();
   const std::shared_ptr<ShaderLibrary>& library = context->GetShaderLibrary();
 
-  //--------------------------------------------------------------------------
-  /// Get or register shader.
-  ///
-
-  // TODO(113719): Register the shader function earlier.
-
   std::shared_ptr<const ShaderFunction> function = library->GetFunction(
       runtime_stage_->GetEntrypoint(), ShaderStage::kFragment);
 
@@ -123,29 +129,75 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
 
     runtime_stage_->SetClean();
   }
+  return true;
+}
 
-  //--------------------------------------------------------------------------
-  /// Set up the command. Defer setting up the pipeline until the descriptor set
-  /// layouts are known from the uniforms.
-  ///
-
+std::shared_ptr<Pipeline<PipelineDescriptor>>
+RuntimeEffectContents::CreatePipeline(const ContentContext& renderer,
+                                      ContentContextOptions options) const {
+  const std::shared_ptr<Context>& context = renderer.GetContext();
+  const std::shared_ptr<ShaderLibrary>& library = context->GetShaderLibrary();
   const std::shared_ptr<const Capabilities>& caps = context->GetCapabilities();
   const auto color_attachment_format = caps->GetDefaultColorFormat();
   const auto stencil_attachment_format = caps->GetDefaultDepthStencilFormat();
 
   using VS = RuntimeEffectVertexShader;
 
-  //--------------------------------------------------------------------------
-  /// Fragment stage uniforms.
-  ///
+  PipelineDescriptor desc;
+  desc.SetLabel("Runtime Stage");
+  desc.AddStageEntrypoint(
+      library->GetFunction(VS::kEntrypointName, ShaderStage::kVertex));
+  desc.AddStageEntrypoint(library->GetFunction(runtime_stage_->GetEntrypoint(),
+                                               ShaderStage::kFragment));
+
+  std::shared_ptr<VertexDescriptor> vertex_descriptor =
+      std::make_shared<VertexDescriptor>();
+  vertex_descriptor->SetStageInputs(VS::kAllShaderStageInputs,
+                                    VS::kInterleavedBufferLayout);
+  vertex_descriptor->RegisterDescriptorSetLayouts(VS::kDescriptorSetLayouts);
+  vertex_descriptor->RegisterDescriptorSetLayouts(
+      runtime_stage_->GetDescriptorSetLayouts().data(),
+      runtime_stage_->GetDescriptorSetLayouts().size());
+  desc.SetVertexDescriptor(std::move(vertex_descriptor));
+  desc.SetColorAttachmentDescriptor(
+      0u, {.format = color_attachment_format, .blending_enabled = true});
+
+  desc.SetStencilAttachmentDescriptors(StencilAttachmentDescriptor{});
+  desc.SetStencilPixelFormat(stencil_attachment_format);
+
+  desc.SetDepthStencilAttachmentDescriptor(DepthAttachmentDescriptor{});
+  desc.SetDepthPixelFormat(stencil_attachment_format);
+
+  options.ApplyToPipelineDescriptor(desc);
+  auto pipeline = context->GetPipelineLibrary()->GetPipeline(desc).Get();
+  if (!pipeline) {
+    VALIDATION_LOG << "Failed to get or create runtime effect pipeline.";
+    return nullptr;
+  }
 
-  std::vector<DescriptorSetLayout> descriptor_set_layouts;
+  return pipeline;
+}
 
-  BindFragmentCallback bind_callback = [this, &renderer, &context,
-                                        &descriptor_set_layouts](
-                                           RenderPass& pass) {
-    descriptor_set_layouts.clear();
+bool RuntimeEffectContents::Render(const ContentContext& renderer,
+                                   const Entity& entity,
+                                   RenderPass& pass) const {
+  const std::shared_ptr<Context>& context = renderer.GetContext();
+  const std::shared_ptr<ShaderLibrary>& library = context->GetShaderLibrary();
 
+  //--------------------------------------------------------------------------
+  /// Get or register shader. Flutter will do this when the runtime effect
+  /// is first loaded, but this check is added to supporting testing of the
+  /// Aiks API and non-flutter usage of Impeller.
+  ///
+  if (!RegisterShader(renderer)) {
+    return false;
+  }
+
+  //--------------------------------------------------------------------------
+  /// Fragment stage uniforms.
+  ///
+  BindFragmentCallback bind_callback = [this, &renderer,
+                                        &context](RenderPass& pass) {
     size_t minimum_sampler_index = 100000000;
     size_t buffer_index = 0;
     size_t buffer_offset = 0;
@@ -173,9 +225,10 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
                      Context::BackendType::kVulkan)
               << "Uniform " << uniform.name
               << " had unexpected type kFloat for Vulkan backend.";
+
           size_t alignment =
               std::max(uniform.bit_width / 8, DefaultUniformAlignment());
-          auto buffer_view = renderer.GetTransientsBuffer().Emplace(
+          BufferView buffer_view = renderer.GetTransientsBuffer().Emplace(
               uniform_data_->data() + buffer_offset, uniform.GetSize(),
               alignment);
 
@@ -184,23 +237,20 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
           uniform_slot.ext_res_0 = uniform.location;
           pass.BindResource(ShaderStage::kFragment,
                             DescriptorType::kUniformBuffer, uniform_slot,
-                            metadata, buffer_view);
+                            metadata, std::move(buffer_view));
           buffer_index++;
           buffer_offset += uniform.GetSize();
           break;
         }
         case kStruct: {
           FML_DCHECK(renderer.GetContext()->GetBackendType() ==
                      Context::BackendType::kVulkan);
-          descriptor_set_layouts.emplace_back(DescriptorSetLayout{
-              static_cast<uint32_t>(uniform.location),
-              DescriptorType::kUniformBuffer,
-              ShaderStage::kFragment,
-          });
           ShaderUniformSlot uniform_slot;
           uniform_slot.name = uniform.name.c_str();
           uniform_slot.binding = uniform.location;
 
+          // TODO(jonahwilliams): rewrite this to emplace directly into
+          // HostBuffer.
           std::vector<float> uniform_buffer;
           uniform_buffer.reserve(uniform.struct_layout.size());
           size_t uniform_byte_index = 0u;
@@ -214,16 +264,15 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
               FML_UNREACHABLE();
             }
           }
-
           size_t alignment = std::max(sizeof(float) * uniform_buffer.size(),
                                       DefaultUniformAlignment());
 
-          auto buffer_view = renderer.GetTransientsBuffer().Emplace(
+          BufferView buffer_view = renderer.GetTransientsBuffer().Emplace(
               reinterpret_cast<const void*>(uniform_buffer.data()),
               sizeof(float) * uniform_buffer.size(), alignment);
           pass.BindResource(ShaderStage::kFragment,
                             DescriptorType::kUniformBuffer, uniform_slot,
-                            ShaderMetadata{}, buffer_view);
+                            ShaderMetadata{}, std::move(buffer_view));
         }
       }
     }
@@ -243,20 +292,7 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
 
           SampledImageSlot image_slot;
           image_slot.name = uniform.name.c_str();
-
-          uint32_t sampler_binding_location = 0u;
-          if (!descriptor_set_layouts.empty()) {
-            sampler_binding_location =
-                descriptor_set_layouts.back().binding + 1;
-          }
-
-          descriptor_set_layouts.emplace_back(DescriptorSetLayout{
-              sampler_binding_location,
-              DescriptorType::kSampledImage,
-              ShaderStage::kFragment,
-          });
-
-          image_slot.binding = sampler_binding_location;
+          image_slot.binding = uniform.binding;
           image_slot.texture_index = uniform.location - minimum_sampler_index;
           pass.BindResource(ShaderStage::kFragment,
                             DescriptorType::kSampledImage, image_slot,
@@ -273,47 +309,15 @@ bool RuntimeEffectContents::Render(const ContentContext& renderer,
   };
 
   /// Now that the descriptor set layouts are known, get the pipeline.
+  using VS = RuntimeEffectVertexShader;
 
-  PipelineBuilderCallback pipeline_callback = [&](ContentContextOptions
-                                                      options) {
-    // Pipeline creation callback for the cache handler to call.
-    auto create_callback =
-        [&]() -> std::shared_ptr<Pipeline<PipelineDescriptor>> {
-      PipelineDescriptor desc;
-      desc.SetLabel("Runtime Stage");
-      desc.AddStageEntrypoint(
-          library->GetFunction(VS::kEntrypointName, ShaderStage::kVertex));
-      desc.AddStageEntrypoint(library->GetFunction(
-          runtime_stage_->GetEntrypoint(), ShaderStage::kFragment));
-      auto vertex_descriptor = std::make_shared<VertexDescriptor>();
-      vertex_descriptor->SetStageInputs(VS::kAllShaderStageInputs,
-                                        VS::kInterleavedBufferLayout);
-      vertex_descriptor->RegisterDescriptorSetLayouts(
-          VS::kDescriptorSetLayouts);
-      vertex_descriptor->RegisterDescriptorSetLayouts(
-          descriptor_set_layouts.data(), descriptor_set_layouts.size());
-      desc.SetVertexDescriptor(std::move(vertex_descriptor));
-      desc.SetColorAttachmentDescriptor(
-          0u, {.format = color_attachment_format, .blending_enabled = true});
-
-      desc.SetStencilAttachmentDescriptors(StencilAttachmentDescriptor{});
-      desc.SetStencilPixelFormat(stencil_attachment_format);
-
-      desc.SetDepthStencilAttachmentDescriptor(DepthAttachmentDescriptor{});
-      desc.SetDepthPixelFormat(stencil_attachment_format);
-
-      options.ApplyToPipelineDescriptor(desc);
-      auto pipeline = context->GetPipelineLibrary()->GetPipeline(desc).Get();
-      if (!pipeline) {
-        VALIDATION_LOG << "Failed to get or create runtime effect pipeline.";
-        return nullptr;
-      }
-
-      return pipeline;
-    };
-    return renderer.GetCachedRuntimeEffectPipeline(
-        runtime_stage_->GetEntrypoint(), options, create_callback);
-  };
+  PipelineBuilderCallback pipeline_callback =
+      [&](ContentContextOptions options) {
+        // Pipeline creation callback for the cache handler to call.
+        return renderer.GetCachedRuntimeEffectPipeline(
+            runtime_stage_->GetEntrypoint(), options,
+            [&]() { return CreatePipeline(renderer, options); });
+      };
 
   return ColorSourceContents::DrawGeometry<VS>(renderer, entity, pass,
                                                pipeline_callback,

impeller/entity/contents/runtime_effect_contents.h

@@ -35,7 +35,16 @@ class RuntimeEffectContents final : public ColorSourceContents {
               const Entity& entity,
               RenderPass& pass) const override;
 
+  /// Load the runtime effect and ensure a default PSO is initialized.
+  bool BootstrapShader(const ContentContext& renderer) const;
+
  private:
+  bool RegisterShader(const ContentContext& renderer) const;
+
+  std::shared_ptr<Pipeline<PipelineDescriptor>> CreatePipeline(
+      const ContentContext& renderer,
+      ContentContextOptions options) const;
+
   std::shared_ptr<RuntimeStage> runtime_stage_;
   std::shared_ptr<std::vector<uint8_t>> uniform_data_;
   std::vector<TextureInput> texture_inputs_;

impeller/entity/entity_unittests.cc

@@ -2242,6 +2242,20 @@ TEST_P(EntityTest, RuntimeEffectCanSuccessfullyRender) {
                   .has_value());
 }
 
+TEST_P(EntityTest, RuntimeEffectCanPrecache) {
+  auto runtime_stages =
+      OpenAssetAsRuntimeStage("runtime_stage_example.frag.iplr");
+  auto runtime_stage =
+      runtime_stages[PlaygroundBackendToRuntimeStageBackend(GetBackend())];
+  ASSERT_TRUE(runtime_stage);
+  ASSERT_TRUE(runtime_stage->IsDirty());
+
+  auto contents = std::make_shared<RuntimeEffectContents>();
+  contents->SetRuntimeStage(runtime_stage);
+
+  EXPECT_TRUE(contents->BootstrapShader(*GetContentContext()));
+}
+
 TEST_P(EntityTest, RuntimeEffectSetsRightSizeWhenUniformIsStruct) {
   if (GetBackend() != PlaygroundBackend::kVulkan) {
     GTEST_SKIP() << "Test only applies to Vulkan";