[graphite] Support code for piet render tasks

- New support code lives under graphite/piet, defining
  PietRenderer (which abstracts backend API object dependencies) and
  PietRenderTask (which integrates piet with graphite's task system).

- ResourceProvider owns a single PietRenderer instance. DrawContext now
  has functions for accumulating and snapping a piet task. Device is
  responsible for snapping and recording a piet task.

Bug: b/241725582
Change-Id: Ie2f3bbb31656553c3d9cb6adeef8156874647743
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/566176
Reviewed-by: Greg Daniel <egdaniel@google.com>
Commit-Queue: Arman Uguray <armansito@google.com>
diff --git a/BUILD.gn b/BUILD.gn
index 8355532..e8a8859 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -1153,6 +1153,11 @@
   public_defines = [ "SK_GRAPHITE_ENABLED" ]
   public = skia_graphite_public
   sources = skia_graphite_sources
+
+  if (skia_use_piet) {
+    sources += skia_graphite_piet_sources
+  }
+
   if (skia_use_metal) {
     public_defines += [ "SK_METAL" ]
     sources += skia_graphite_mtl_sources
diff --git a/gn/graphite.gni b/gn/graphite.gni
index 86933f1..45fd98f 100644
--- a/gn/graphite.gni
+++ b/gn/graphite.gni
@@ -188,3 +188,8 @@
   "$_src/mtl/MtlUtils.h",
   "$_src/mtl/MtlUtils.mm",
 ]
+
+skia_graphite_piet_sources = [
+  "$_src/PietRenderTask.cpp",
+  "$_src/PietRenderTask.h",
+]
diff --git a/src/gpu/graphite/CommandBuffer.cpp b/src/gpu/graphite/CommandBuffer.cpp
index 94c6e15..dab27ba 100644
--- a/src/gpu/graphite/CommandBuffer.cpp
+++ b/src/gpu/graphite/CommandBuffer.cpp
@@ -135,4 +135,11 @@
     return true;
 }
 
+#ifdef SK_ENABLE_PIET_GPU
+void CommandBuffer::renderPietScene(const skgpu::piet::Scene& scene, sk_sp<Texture> target) {
+    this->onRenderPietScene(scene, target.get());
+    this->trackResource(std::move(target));
+}
+#endif
+
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/CommandBuffer.h b/src/gpu/graphite/CommandBuffer.h
index 1e76604..1061c60 100644
--- a/src/gpu/graphite/CommandBuffer.h
+++ b/src/gpu/graphite/CommandBuffer.h
@@ -22,6 +22,12 @@
 class RefCntedCallback;
 }
 
+#ifdef SK_ENABLE_PIET_GPU
+namespace skgpu::piet {
+class Scene;
+}
+#endif
+
 namespace skgpu::graphite {
 
 class Buffer;
@@ -72,6 +78,10 @@
                              const BufferTextureCopyData*,
                              int count);
 
+#ifdef SK_ENABLE_PIET_GPU
+    void renderPietScene(const skgpu::piet::Scene& scene, sk_sp<Texture> target);
+#endif
+
 protected:
     CommandBuffer();
 
@@ -96,6 +106,10 @@
                                        const BufferTextureCopyData*,
                                        int count) = 0;
 
+#ifdef SK_ENABLE_PIET_GPU
+    virtual void onRenderPietScene(const skgpu::piet::Scene& scene, const Texture* target) = 0;
+#endif
+
 #ifdef SK_DEBUG
     bool fHasWork = false;
 #endif
diff --git a/src/gpu/graphite/Context.cpp b/src/gpu/graphite/Context.cpp
index 495288f..3faf18e 100644
--- a/src/gpu/graphite/Context.cpp
+++ b/src/gpu/graphite/Context.cpp
@@ -22,6 +22,7 @@
 #include "src/gpu/graphite/GlobalCache.h"
 #include "src/gpu/graphite/GraphicsPipelineDesc.h"
 #include "src/gpu/graphite/QueueManager.h"
+#include "src/gpu/graphite/RecorderPriv.h"
 #include "src/gpu/graphite/RecordingPriv.h"
 #include "src/gpu/graphite/Renderer.h"
 #include "src/gpu/graphite/ResourceProvider.h"
diff --git a/src/gpu/graphite/Device.cpp b/src/gpu/graphite/Device.cpp
index f469393..e92d3ff 100644
--- a/src/gpu/graphite/Device.cpp
+++ b/src/gpu/graphite/Device.cpp
@@ -899,6 +899,13 @@
         fRecorder->priv().add(std::move(uploadTask));
     }
 
+#ifdef SK_ENABLE_PIET_GPU
+    auto pietTask = fDC->snapPietRenderTask(fRecorder);
+    if (pietTask) {
+        fRecorder->priv().add(std::move(pietTask));
+    }
+#endif
+
     fClip.recordDeferredClipDraws();
     auto drawTask = fDC->snapRenderPassTask(fRecorder);
     if (drawTask) {
diff --git a/src/gpu/graphite/DrawContext.cpp b/src/gpu/graphite/DrawContext.cpp
index 79a0b1c..4418dc2 100644
--- a/src/gpu/graphite/DrawContext.cpp
+++ b/src/gpu/graphite/DrawContext.cpp
@@ -28,6 +28,10 @@
 #include "src/gpu/graphite/geom/BoundsManager.h"
 #include "src/gpu/graphite/geom/Geometry.h"
 
+#ifdef SK_ENABLE_PIET_GPU
+#include "src/gpu/graphite/PietRenderTask.h"
+#endif
+
 namespace skgpu::graphite {
 
 sk_sp<DrawContext> DrawContext::Make(sk_sp<TextureProxy> target,
@@ -97,6 +101,15 @@
                                          dstRect);
 }
 
+#ifdef SK_ENABLE_PIET_GPU
+bool DrawContext::recordPietSceneRender(Recorder*,
+                                        sk_sp<TextureProxy> targetProxy,
+                                        sk_sp<const skgpu::piet::Scene> scene) {
+    fPendingPietRenders.push_back(PietRenderInstance(std::move(scene), std::move(targetProxy)));
+    return true;
+}
+#endif
+
 void DrawContext::snapDrawPass(Recorder* recorder) {
     if (fPendingDraws->drawCount() == 0) {
         return;
@@ -184,4 +197,13 @@
     return uploadTask;
 }
 
+#ifdef SK_ENABLE_PIET_GPU
+sk_sp<Task> DrawContext::snapPietRenderTask(Recorder* recorder) {
+    if (fPendingPietRenders.empty()) {
+        return nullptr;
+    }
+    return sk_sp<Task>(new PietRenderTask(std::move(fPendingPietRenders)));
+}
+#endif
+
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/DrawContext.h b/src/gpu/graphite/DrawContext.h
index cef0c63..05989b4 100644
--- a/src/gpu/graphite/DrawContext.h
+++ b/src/gpu/graphite/DrawContext.h
@@ -19,6 +19,13 @@
 
 #include <vector>
 
+#ifdef SK_ENABLE_PIET_GPU
+#include "src/gpu/graphite/PietRenderTask.h"
+namespace skgpu::piet {
+class Scene;
+}
+#endif
+
 class SkPixmap;
 
 namespace skgpu::graphite {
@@ -66,6 +73,12 @@
                       const std::vector<MipLevel>& levels,
                       const SkIRect& dstRect);
 
+#ifdef SK_ENABLE_PIET_GPU
+    bool recordPietSceneRender(Recorder* recorder,
+                               sk_sp<TextureProxy> targetProxy,
+                               sk_sp<const skgpu::piet::Scene> pietScene);
+#endif
+
     // Ends the current DrawList being accumulated by the SDC, converting it into an optimized and
     // immutable DrawPass. The DrawPass will be ordered after any other snapped DrawPasses or
     // appended DrawPasses from a child SDC. A new DrawList is started to record subsequent drawing
@@ -96,6 +109,10 @@
     // TODO: see if we can merge transfers into this
     sk_sp<Task> snapUploadTask(Recorder*);
 
+#ifdef SK_ENABLE_PIET_GPU
+    sk_sp<Task> snapPietRenderTask(Recorder*);
+#endif
+
 private:
     DrawContext(sk_sp<TextureProxy>, const SkImageInfo&);
 
@@ -123,6 +140,10 @@
     // Stores the most immediately recorded uploads into Textures. This list is mutable and
     // can be appended to, or have its commands rewritten if they are inlined into a parent DC.
     std::unique_ptr<UploadList> fPendingUploads;
+
+#ifdef SK_ENABLE_PIET_GPU
+    std::vector<PietRenderInstance> fPendingPietRenders;
+#endif
 };
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/PietRenderTask.cpp b/src/gpu/graphite/PietRenderTask.cpp
new file mode 100644
index 0000000..c744d04
--- /dev/null
+++ b/src/gpu/graphite/PietRenderTask.cpp
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "src/gpu/graphite/PietRenderTask.h"
+
+#include "src/gpu/graphite/Log.h"
+#include "src/gpu/graphite/ResourceProvider.h"
+#include "src/gpu/graphite/Texture.h"
+#include "src/gpu/graphite/TextureProxy.h"
+
+namespace skgpu::graphite {
+
+bool PietRenderInstance::prepareResources(ResourceProvider* resourceProvider) {
+    if (!fTargetProxy) {
+        SKGPU_LOG_E("No texture proxy specified for PietRenderTask");
+        return false;
+    }
+    if (!fTargetProxy->instantiate(resourceProvider)) {
+        SKGPU_LOG_E("Could not instantiate texture proxy for PietRenderTask!");
+        return false;
+    }
+    return true;
+}
+
+bool PietRenderInstance::addCommands(CommandBuffer* commandBuffer) {
+    SkASSERT(fTargetProxy && fTargetProxy->isInstantiated());
+    commandBuffer->renderPietScene(*fScene, fTargetProxy->refTexture());
+    return true;
+}
+
+bool PietRenderTask::prepareResources(ResourceProvider* resourceProvider,
+                                      const SkRuntimeEffectDictionary*) {
+    for (PietRenderInstance& instance : fInstances) {
+        if (!instance.prepareResources(resourceProvider)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+bool PietRenderTask::addCommands(ResourceProvider*, CommandBuffer* commandBuffer) {
+    for (PietRenderInstance& instance : fInstances) {
+        instance.addCommands(commandBuffer);
+    }
+    return true;
+}
+
+}  // namespace skgpu::graphite
diff --git a/src/gpu/graphite/PietRenderTask.h b/src/gpu/graphite/PietRenderTask.h
new file mode 100644
index 0000000..15fd009
--- /dev/null
+++ b/src/gpu/graphite/PietRenderTask.h
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#ifndef skgpu_graphite_PietRenderTask_DEFINED
+#define skgpu_graphite_PietRenderTask_DEFINED
+
+#include <vector>
+
+#include "src/gpu/graphite/Task.h"
+#include "src/gpu/piet/Scene.h"
+
+namespace skgpu::graphite {
+
+class TextureProxy;
+
+class PietRenderInstance final {
+public:
+    PietRenderInstance(sk_sp<const skgpu::piet::Scene> scene, sk_sp<TextureProxy> targetProxy)
+            : fScene(std::move(scene))
+            , fTargetProxy(std::move(targetProxy)) {}
+
+    // Make this type move-only.
+    PietRenderInstance(const PietRenderInstance&) = delete;
+    PietRenderInstance(PietRenderInstance&&) = default;
+
+    bool prepareResources(ResourceProvider*);
+
+    bool addCommands(CommandBuffer*);
+
+private:
+    sk_sp<const skgpu::piet::Scene> fScene;
+    sk_sp<TextureProxy> fTargetProxy;
+};
+
+class PietRenderTask final : public Task {
+public:
+    explicit PietRenderTask(std::vector<PietRenderInstance> instances)
+            : fInstances(std::move(instances)) {}
+
+    ~PietRenderTask() override = default;
+
+    bool prepareResources(ResourceProvider*, const SkRuntimeEffectDictionary*) override;
+
+    bool addCommands(ResourceProvider*, CommandBuffer*) override;
+
+private:
+    std::vector<PietRenderInstance> fInstances;
+};
+
+}  // namespace skgpu::graphite
+
+#endif  // skgpu_graphite_PietRenderTask_DEFINED
diff --git a/src/gpu/graphite/mtl/MtlCommandBuffer.h b/src/gpu/graphite/mtl/MtlCommandBuffer.h
index d02e9dd..d095a2e 100644
--- a/src/gpu/graphite/mtl/MtlCommandBuffer.h
+++ b/src/gpu/graphite/mtl/MtlCommandBuffer.h
@@ -18,6 +18,10 @@
 #include "include/core/SkTypes.h"
 #include "include/ports/SkCFObject.h"
 
+#ifdef SK_ENABLE_PIET_GPU
+#include "src/gpu/piet/Render.h"
+#endif
+
 #import <Metal/Metal.h>
 
 namespace skgpu::graphite {
@@ -50,6 +54,10 @@
     }
     bool commit();
 
+#ifdef SK_ENABLE_PIET_GPU
+    void setPietRenderer(const skgpu::piet::MtlRenderer* renderer) { fPietRenderer = renderer; }
+#endif
+
 private:
     MtlCommandBuffer(sk_cfp<id<MTLCommandBuffer>> cmdBuffer, const MtlSharedContext* sharedContext);
 
@@ -117,6 +125,10 @@
                                const BufferTextureCopyData* copyData,
                                int count) override;
 
+#ifdef SK_ENABLE_PIET_GPU
+    void onRenderPietScene(const skgpu::piet::Scene& scene, const Texture* target) override;
+#endif
+
     MtlBlitCommandEncoder* getBlitCommandEncoder();
     void endBlitCommandEncoder();
 
@@ -131,6 +143,10 @@
     size_t fCurrentIndexBufferOffset = 0;
 
     const MtlSharedContext* fSharedContext;
+
+#ifdef SK_ENABLE_PIET_GPU
+    const skgpu::piet::MtlRenderer* fPietRenderer = nullptr;  // owned by MtlQueueManager
+#endif
 };
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/mtl/MtlCommandBuffer.mm b/src/gpu/graphite/mtl/MtlCommandBuffer.mm
index 7ea8369..3a7cf18 100644
--- a/src/gpu/graphite/mtl/MtlCommandBuffer.mm
+++ b/src/gpu/graphite/mtl/MtlCommandBuffer.mm
@@ -637,5 +637,17 @@
     return true;
 }
 
+#ifdef SK_ENABLE_PIET_GPU
+void MtlCommandBuffer::onRenderPietScene(const skgpu::piet::Scene& scene, const Texture* target) {
+    SkASSERT(!fActiveRenderCommandEncoder);
+    SkASSERT(!fActiveComputeCommandEncoder);
+    this->endBlitCommandEncoder();
+
+    SkASSERT(fPietRenderer);
+
+    id<MTLTexture> mtlTexture = static_cast<const MtlTexture*>(target)->mtlTexture();
+    fPietRenderer->render(scene, mtlTexture, fCommandBuffer.get());
+}
+#endif
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/mtl/MtlQueueManager.h b/src/gpu/graphite/mtl/MtlQueueManager.h
index 3259e36..8ef3660 100644
--- a/src/gpu/graphite/mtl/MtlQueueManager.h
+++ b/src/gpu/graphite/mtl/MtlQueueManager.h
@@ -11,6 +11,10 @@
 #include "include/ports/SkCFObject.h"
 #include "src/gpu/graphite/QueueManager.h"
 
+#ifdef SK_ENABLE_PIET_GPU
+#include "src/gpu/piet/Render.h"
+#endif
+
 #import <Metal/Metal.h>
 
 namespace skgpu::graphite {
@@ -35,6 +39,10 @@
 #endif
 
     sk_cfp<id<MTLCommandQueue>> fQueue;
+
+#ifdef SK_ENABLE_PIET_GPU
+    skgpu::piet::MtlRenderer fPietRenderer;
+#endif
 };
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/mtl/MtlQueueManager.mm b/src/gpu/graphite/mtl/MtlQueueManager.mm
index d682bd6..4c86f79 100644
--- a/src/gpu/graphite/mtl/MtlQueueManager.mm
+++ b/src/gpu/graphite/mtl/MtlQueueManager.mm
@@ -15,14 +15,25 @@
 MtlQueueManager::MtlQueueManager(sk_cfp<id<MTLCommandQueue>> queue,
                                  const SharedContext* sharedContext)
         : QueueManager(sharedContext)
-        , fQueue(std::move(queue)) {}
+        , fQueue(std::move(queue))
+#ifdef SK_ENABLE_PIET_GPU
+        , fPietRenderer(this->mtlSharedContext()->device(), fQueue.get())
+#endif
+{
+}
 
 const MtlSharedContext* MtlQueueManager::mtlSharedContext() const {
     return static_cast<const MtlSharedContext*>(fSharedContext);
 }
 
 sk_sp<CommandBuffer> MtlQueueManager::getNewCommandBuffer() {
-    return MtlCommandBuffer::Make(fQueue.get(), this->mtlSharedContext());
+    auto cmdBuffer = MtlCommandBuffer::Make(fQueue.get(), this->mtlSharedContext());
+
+#ifdef SK_ENABLE_PIET_GPU
+    cmdBuffer->setPietRenderer(&fPietRenderer);
+#endif
+
+    return std::move(cmdBuffer);
 }
 
 class WorkSubmission final : public GpuWorkSubmission {
diff --git a/src/gpu/graphite/mtl/MtlTrampoline.h b/src/gpu/graphite/mtl/MtlTrampoline.h
index 831439a..8f5ef29 100644
--- a/src/gpu/graphite/mtl/MtlTrampoline.h
+++ b/src/gpu/graphite/mtl/MtlTrampoline.h
@@ -36,7 +36,6 @@
     static std::unique_ptr<ResourceProvider> MakeResourceProvider(const SharedContext*,
                                                                   sk_sp<GlobalCache>,
                                                                   SingleOwner*);
-
 };
 
 } // namespace skgpu::graphite
diff --git a/src/gpu/graphite/mtl/MtlTrampoline.mm b/src/gpu/graphite/mtl/MtlTrampoline.mm
index 4d2d721..608a5c6 100644
--- a/src/gpu/graphite/mtl/MtlTrampoline.mm
+++ b/src/gpu/graphite/mtl/MtlTrampoline.mm
@@ -13,6 +13,7 @@
 #include "src/gpu/graphite/mtl/MtlSharedContext.h"
 
 namespace skgpu::graphite {
+
 sk_sp<SharedContext> MtlTrampoline::MakeSharedContext(const MtlBackendContext& backendContext,
                                                       const ContextOptions& options) {
     return MtlSharedContext::Make(backendContext, options);