[graphite] Support uniform buffers in ComputeSteps

* Introduced the `ComptueStep::prepareUniformBuffer()` method which is
  invoked with a `UniformManager` pointer to handle the field encoding with
  the correct layout constraints.

* Unlike RenderSteps, ComputeSteps do not use a PipelineDataGatherer nor
  to they need to declare their uniforms using the `Uniform` POD type,
  since ComputeSteps are not assembled from snippets. However, they
  still need to call `UniformManager::setExpectedUniforms` to satisfy
  UniformManager's debug validation assertions.

* Renamed `ComputeStep::prepareBuffer()` to `prepareStorageBuffer()` as
  this is now exclusively used for internally allocated storage buffers.

Change-Id: I7e30a4402a5aeb6cd4eafd5abde2f999e431eb76
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/665584
Commit-Queue: Arman Uguray <[email protected]>
Reviewed-by: Jim Van Verth <[email protected]>
Reviewed-by: Michael Ludwig <[email protected]>

src/gpu/graphite/compute/ComputeStep.cpp

@@ -77,9 +77,16 @@ ComputeStep::ComputeStep(std::string_view name,
     }
 }
 
-void ComputeStep::prepareBuffer(
+void ComputeStep::prepareStorageBuffer(
         const DrawParams&, int, int, const ResourceDesc&, void*, size_t) const {
-    SK_ABORT("ComputeSteps using a mapped resource must override prepareBuffer()");
+    SK_ABORT("ComputeSteps using a mapped storage buffer must override prepareStorageBuffer()");
+}
+
+void ComputeStep::prepareUniformBuffer(const DrawParams&,
+                                       int,
+                                       const ResourceDesc&,
+                                       UniformManager*) const {
+    SK_ABORT("ComputeSteps using a uniform buffer must override prepareUniformBuffer()");
 }
 
 size_t ComputeStep::calculateBufferSize(const DrawParams&,

src/gpu/graphite/compute/ComputeStep.h

@@ -23,6 +23,7 @@
 namespace skgpu::graphite {
 
 class DrawParams;
+class UniformManager;
 struct ResourceBindingRequirements;
 
 /**
@@ -98,14 +99,16 @@ class ComputeStep {
         kClear,
 
         // The ComputeStep will be asked to initialize the memory on the CPU via
-        // `ComputeStep::prepareBuffer` prior to pipeline execution. This may incur a transfer cost
-        // on platforms that do not allow buffers to be mapped in shared memory.
+        // `ComputeStep::prepareStorageBuffer` or `ComputeStep::prepareUniformBuffer` prior to
+        // pipeline execution. This may incur a transfer cost on platforms that do not allow buffers
+        // to be mapped in shared memory.
         //
         // If multiple ComputeSteps in a DispatchGroup declare a mapped resource with the same
-        // shared slot number, only the first ComputeStep in the series will receive a call to
-        // `ComputeStep::prepareBuffer`.
+        // shared slot number, only the first ComputeStep in the group will receive a call to
+        // prepare the buffer.
         //
-        // This only has meaning for buffer resources.
+        // This only has meaning for buffer resources. A resource with the `kUniformBuffer` resource
+        // type must specify the `kMapped` resource policy.
         kMapped,
     };
 
@@ -167,18 +170,37 @@ class ComputeStep {
         return WorkgroupSize();
     }
 
-    // Populates a buffer resource which was specified as "mapped". This method will only be called
-    // once for a resource right after its allocation and before pipeline execution. For shared
-    // resources, only the first ComputeStep in a DispatchGroup will be asked to prepare the buffer.
+    // Populates a storage buffer resource which was specified as "mapped". This method will only be
+    // called once for a resource right after its allocation and before pipeline execution. For
+    // shared resources, only the first ComputeStep in a DispatchGroup will be asked to prepare the
+    // buffer.
     //
     // `resourceIndex` matches the order in which `resource` was enumerated by
     // `ComputeStep::resources()`.
-    virtual void prepareBuffer(const DrawParams&,
-                               int ssboIndex,
-                               int resourceIndex,
-                               const ResourceDesc& resource,
-                               void* buffer,
-                               size_t bufferSize) const;
+    virtual void prepareStorageBuffer(const DrawParams&,
+                                      int ssboIndex,
+                                      int resourceIndex,
+                                      const ResourceDesc& resource,
+                                      void* buffer,
+                                      size_t bufferSize) const;
+
+    // Populates a uniform buffer resource. This method will be called once for a resource right
+    // after its allocation and before pipeline execution. For shared resources, only the first
+    // ComputeStep in a DispatchGroup will be asked to prepare the buffer.
+    //
+    // `resourceIndex` matches the order in which `resource` was enumerated by
+    // `ComputeStep::resources()`.
+    //
+    // The implementation must use the provided `UniformManager` to populate the buffer. On debug
+    // builds, the implementation must validate the buffer layout by setting up an expectation, for
+    // example:
+    //
+    //     SkDEBUGCODE(mgr->setExpectedUniforms({{"foo", SkSLType::kFloat}}));
+    //
+    virtual void prepareUniformBuffer(const DrawParams&,
+                                      int resourceIndex,
+                                      const ResourceDesc&,
+                                      UniformManager*) const;
 
     SkSpan<const ResourceDesc> resources() const { return SkSpan(fResources); }
 

src/gpu/graphite/compute/DispatchGroup.cpp

@@ -9,11 +9,14 @@
 
 #include "include/gpu/graphite/Recorder.h"
 #include "src/gpu/graphite/BufferManager.h"
+#include "src/gpu/graphite/Caps.h"
 #include "src/gpu/graphite/CommandBuffer.h"
 #include "src/gpu/graphite/ComputePipeline.h"
 #include "src/gpu/graphite/Log.h"
+#include "src/gpu/graphite/PipelineData.h"
 #include "src/gpu/graphite/RecorderPriv.h"
 #include "src/gpu/graphite/ResourceProvider.h"
+#include "src/gpu/graphite/UniformManager.h"
 
 namespace skgpu::graphite {
 
@@ -263,9 +266,9 @@ DispatchResourceOptional Builder::allocateResource(const ComputeStep* step,
             SkASSERT(bufferSize);
             if (resource.fPolicy == ResourcePolicy::kMapped) {
                 auto [ptr, bufInfo] = bufferMgr->getStoragePointer(bufferSize);
-                // Allocation failures are handled below.
                 if (ptr) {
-                    step->prepareBuffer(params, ssboIdx, resourceIdx, resource, ptr, bufferSize);
+                    step->prepareStorageBuffer(
+                            params, ssboIdx, resourceIdx, resource, ptr, bufferSize);
                     result = bufInfo;
                 }
             } else {
@@ -279,6 +282,23 @@ DispatchResourceOptional Builder::allocateResource(const ComputeStep* step,
             }
             break;
         }
+        case Type::kUniformBuffer: {
+            SkASSERT(resource.fPolicy == ResourcePolicy::kMapped);
+
+            const auto& resourceReqs = fRecorder->priv().caps()->resourceBindingRequirements();
+            UniformManager uboMgr(resourceReqs.fUniformBufferLayout);
+            step->prepareUniformBuffer(params, resourceIdx, resource, &uboMgr);
+
+            auto dataBlock = uboMgr.finishUniformDataBlock();
+            SkASSERT(dataBlock.size());
+
+            auto [writer, bufInfo] = bufferMgr->getUniformWriter(dataBlock.size());
+            if (bufInfo) {
+                writer.write(dataBlock.data(), dataBlock.size());
+                result = bufInfo;
+            }
+            break;
+        }
         case Type::kStorageTexture: {
             auto [size, colorType] =
                     step->calculateTextureParameters(params, resourceIdx, resource);
@@ -302,10 +322,6 @@ DispatchResourceOptional Builder::allocateResource(const ComputeStep* step,
             // slot by calling the Builder::assignSharedTexture() method.
             SK_ABORT("a sampled texture must be externally assigned to a ComputeStep");
             break;
-        case Type::kUniformBuffer:
-            // TODO(b/259564970): Support uniform buffers
-            SkUNREACHABLE;
-            break;
     }
     return result;
 }