Support capturing a gainmapped screenshot

This uses the gainmap concept from UltraHDR and ISO 21496-1 to produce
screenshots that (a) can be rendered in HDR and (b) that are perfectly
backwards-compatible on SDR displays are scenes where displaying SDR is
preferable.

The technical details of the screenshot process at a high level are:
* A client requests that they want a gainmap attached to the screenshot
  result
* SurfaceFlinger asks RenderEngine to perform three render passes
    * First, render the SDR base image, tonemapping using MouriMap
    * Second, render an HDR rendition. This is *also* tonemapped using
      MouriMap to the current display conditions. For HDR UI, this is
      effectively a no-op. For PQ and HLG content, this allows for
      approximating what is on-screen as closely as possible. Note that
      to preserve precision as much as possible, we require an FP16
      framebuffer.
    * Finally, combine the SDR and HDR renditions into a gainmap
* The client now receives the base image with the gainmap, as well as an
  HDR/SDR ratio. The client can use this information to generate gainmap
  metadata for encoding or round-tripping to the display.

This MVP is sufficient for generating screenshots through adb screencap.
Eventually, this can be used to integrate with the applications that
want HDR screenshots, or to remove hacks that disable tonemapping during
certain system animations.

Bug: 329470026
Flag: com.android.graphics.surfaceflinger.flags.true_hdr_screenshots
Test: builds
Test: adb screencap
Change-Id: I0434059d957bb6cc38d019743619d72dda72a269

5 files changed, 117 insertions, 37 deletions

diff --git a/services/surfaceflinger/RegionSamplingThread.cpp b/services/surfaceflinger/RegionSamplingThread.cpp
index 7712d38f43..06c2f26a6d 100644
--- a/services/surfaceflinger/RegionSamplingThread.cpp
+++ b/services/surfaceflinger/RegionSamplingThread.cpp
@@ -346,6 +346,7 @@ void RegionSamplingThread::captureSample() {
     constexpr bool kRegionSampling = true;
     constexpr bool kGrayscale = false;
     constexpr bool kIsProtected = false;
+    constexpr bool kAttachGainmap = false;
 
     SurfaceFlinger::RenderAreaBuilderVariant
             renderAreaBuilder(std::in_place_type<DisplayRenderAreaBuilder>, sampledBounds,
@@ -358,15 +359,15 @@ void RegionSamplingThread::captureSample() {
         std::vector<sp<LayerFE>> layerFEs;
         auto displayState = mFlinger.getSnapshotsFromMainThread(renderAreaBuilder,
                                                                 getLayerSnapshotsFn, layerFEs);
-        fenceResult =
-                mFlinger.captureScreenshot(renderAreaBuilder, buffer, kRegionSampling, kGrayscale,
-                                           kIsProtected, nullptr, displayState, layerFEs)
-                        .get();
+        fenceResult = mFlinger.captureScreenshot(renderAreaBuilder, buffer, kRegionSampling,
+                                                 kGrayscale, kIsProtected, kAttachGainmap, nullptr,
+                                                 displayState, layerFEs)
+                              .get();
     } else {
-        fenceResult =
-                mFlinger.captureScreenshotLegacy(renderAreaBuilder, getLayerSnapshotsFn, buffer,
-                                                 kRegionSampling, kGrayscale, kIsProtected, nullptr)
-                        .get();
+        fenceResult = mFlinger.captureScreenshotLegacy(renderAreaBuilder, getLayerSnapshotsFn,
+                                                       buffer, kRegionSampling, kGrayscale,
+                                                       kIsProtected, kAttachGainmap, nullptr)
+                              .get();
     }
     if (fenceResult.ok()) {
         fenceResult.value()->waitForever(LOG_TAG);
diff --git a/services/surfaceflinger/ScreenCaptureOutput.cpp b/services/surfaceflinger/ScreenCaptureOutput.cpp
index 8bb72b8470..41a9a1bb22 100644
--- a/services/surfaceflinger/ScreenCaptureOutput.cpp
+++ b/services/surfaceflinger/ScreenCaptureOutput.cpp
@@ -93,6 +93,12 @@ renderengine::DisplaySettings ScreenCaptureOutput::generateClientCompositionDisp
     if (mEnableLocalTonemapping) {
         clientCompositionDisplay.tonemapStrategy =
                 renderengine::DisplaySettings::TonemapStrategy::Local;
+        if (static_cast<ui::PixelFormat>(buffer->getPixelFormat()) == ui::PixelFormat::RGBA_FP16) {
+            clientCompositionDisplay.targetHdrSdrRatio =
+                    getState().displayBrightnessNits / getState().sdrWhitePointNits;
+        } else {
+            clientCompositionDisplay.targetHdrSdrRatio = 1.f;
+        }
     }
 
     return clientCompositionDisplay;
diff --git a/services/surfaceflinger/SurfaceFlinger.cpp b/services/surfaceflinger/SurfaceFlinger.cpp
index 0aad2828cb..7590d981d4 100644
--- a/services/surfaceflinger/SurfaceFlinger.cpp
+++ b/services/surfaceflinger/SurfaceFlinger.cpp
@@ -7060,7 +7060,8 @@ void SurfaceFlinger::captureDisplay(const DisplayCaptureArgs& args,
                                                  displayWeak, options),
                         getLayerSnapshotsFn, reqSize,
                         static_cast<ui::PixelFormat>(captureArgs.pixelFormat),
-                        captureArgs.allowProtected, captureArgs.grayscale, captureListener);
+                        captureArgs.allowProtected, captureArgs.grayscale,
+                        captureArgs.attachGainmap, captureListener);
 }
 
 void SurfaceFlinger::captureDisplay(DisplayId displayId, const CaptureArgs& args,
@@ -7117,7 +7118,7 @@ void SurfaceFlinger::captureDisplay(DisplayId displayId, const CaptureArgs& args
                                                  static_cast<ui::Dataspace>(args.dataspace),
                                                  displayWeak, options),
                         getLayerSnapshotsFn, size, static_cast<ui::PixelFormat>(args.pixelFormat),
-                        kAllowProtected, kGrayscale, captureListener);
+                        kAllowProtected, kGrayscale, args.attachGainmap, captureListener);
 }
 
 ScreenCaptureResults SurfaceFlinger::captureLayersSync(const LayerCaptureArgs& args) {
@@ -7228,7 +7229,8 @@ void SurfaceFlinger::captureLayers(const LayerCaptureArgs& args,
                                                  options),
                         getLayerSnapshotsFn, reqSize,
                         static_cast<ui::PixelFormat>(captureArgs.pixelFormat),
-                        captureArgs.allowProtected, captureArgs.grayscale, captureListener);
+                        captureArgs.allowProtected, captureArgs.grayscale,
+                        captureArgs.attachGainmap, captureListener);
 }
 
 // Creates a Future release fence for a layer and keeps track of it in a list to
@@ -7281,7 +7283,7 @@ std::optional<SurfaceFlinger::OutputCompositionState> SurfaceFlinger::getSnapsho
 void SurfaceFlinger::captureScreenCommon(RenderAreaBuilderVariant renderAreaBuilder,
                                          GetLayerSnapshotsFunction getLayerSnapshotsFn,
                                          ui::Size bufferSize, ui::PixelFormat reqPixelFormat,
-                                         bool allowProtected, bool grayscale,
+                                         bool allowProtected, bool grayscale, bool attachGainmap,
                                          const sp<IScreenCaptureListener>& captureListener) {
     SFTRACE_CALL();
 
@@ -7327,9 +7329,9 @@ void SurfaceFlinger::captureScreenCommon(RenderAreaBuilderVariant renderAreaBuil
                 renderengine::impl::ExternalTexture>(buffer, getRenderEngine(),
                                                      renderengine::impl::ExternalTexture::Usage::
                                                              WRITEABLE);
-        auto futureFence =
-                captureScreenshot(renderAreaBuilder, texture, false /* regionSampling */, grayscale,
-                                  isProtected, captureListener, displayState, layerFEs);
+        auto futureFence = captureScreenshot(renderAreaBuilder, texture, false /* regionSampling */,
+                                             grayscale, isProtected, attachGainmap, captureListener,
+                                             displayState, layerFEs);
         futureFence.get();
 
     } else {
@@ -7364,7 +7366,7 @@ void SurfaceFlinger::captureScreenCommon(RenderAreaBuilderVariant renderAreaBuil
                                                              WRITEABLE);
         auto futureFence = captureScreenshotLegacy(renderAreaBuilder, getLayerSnapshotsFn, texture,
                                                    false /* regionSampling */, grayscale,
-                                                   isProtected, captureListener);
+                                                   isProtected, attachGainmap, captureListener);
         futureFence.get();
     }
 }
@@ -7418,7 +7420,8 @@ std::vector<sp<LayerFE>> SurfaceFlinger::extractLayerFEs(
 ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
         const RenderAreaBuilderVariant& renderAreaBuilder,
         const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
-        bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener,
+        bool grayscale, bool isProtected, bool attachGainmap,
+        const sp<IScreenCaptureListener>& captureListener,
         std::optional<OutputCompositionState>& displayState, std::vector<sp<LayerFE>>& layerFEs) {
     SFTRACE_CALL();
 
@@ -7435,19 +7438,87 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
         }
         return ftl::yield<FenceResult>(base::unexpected(NO_ERROR)).share();
     }
+    float displayBrightnessNits = displayState.value().displayBrightnessNits;
+    float sdrWhitePointNits = displayState.value().sdrWhitePointNits;
 
     // Empty vector needed to pass into renderScreenImpl for legacy path
     std::vector<std::pair<Layer*, sp<android::LayerFE>>> layers;
     ftl::SharedFuture<FenceResult> renderFuture =
-            renderScreenImpl(std::move(renderArea), buffer, regionSampling, grayscale, isProtected,
-                             captureResults, displayState, layers, layerFEs);
+            renderScreenImpl(renderArea.get(), buffer, regionSampling, grayscale, isProtected,
+                             attachGainmap, captureResults, displayState, layers, layerFEs);
+
+    if (captureResults.capturedHdrLayers && attachGainmap &&
+        FlagManager::getInstance().true_hdr_screenshots()) {
+        sp<GraphicBuffer> hdrBuffer =
+                getFactory().createGraphicBuffer(buffer->getWidth(), buffer->getHeight(),
+                                                 HAL_PIXEL_FORMAT_RGBA_FP16, 1 /* layerCount */,
+                                                 buffer->getUsage(), "screenshot-hdr");
+        sp<GraphicBuffer> gainmapBuffer =
+                getFactory().createGraphicBuffer(buffer->getWidth(), buffer->getHeight(),
+                                                 buffer->getPixelFormat(), 1 /* layerCount */,
+                                                 buffer->getUsage(), "screenshot-gainmap");
+
+        const status_t bufferStatus = hdrBuffer->initCheck();
+        const status_t gainmapBufferStatus = gainmapBuffer->initCheck();
+
+        if (bufferStatus != OK) {
+            ALOGW("%s: Buffer failed to allocate for hdr: %d. Screenshoting SDR instead.", __func__,
+                  bufferStatus);
+        } else if (gainmapBufferStatus != OK) {
+            ALOGW("%s: Buffer failed to allocate for gainmap: %d. Screenshoting SDR instead.",
+                  __func__, gainmapBufferStatus);
+        } else {
+            captureResults.optionalGainMap = gainmapBuffer;
+            const auto hdrTexture = std::make_shared<
+                    renderengine::impl::ExternalTexture>(hdrBuffer, getRenderEngine(),
+                                                         renderengine::impl::ExternalTexture::
+                                                                 Usage::WRITEABLE);
+            const auto gainmapTexture = std::make_shared<
+                    renderengine::impl::ExternalTexture>(gainmapBuffer, getRenderEngine(),
+                                                         renderengine::impl::ExternalTexture::
+                                                                 Usage::WRITEABLE);
+            ScreenCaptureResults unusedResults;
+            ftl::SharedFuture<FenceResult> hdrRenderFuture =
+                    renderScreenImpl(renderArea.get(), hdrTexture, regionSampling, grayscale,
+                                     isProtected, attachGainmap, unusedResults, displayState,
+                                     layers, layerFEs);
+
+            renderFuture =
+                    ftl::Future(std::move(renderFuture))
+                            .then([&, hdrRenderFuture = std::move(hdrRenderFuture),
+                                   displayBrightnessNits, sdrWhitePointNits,
+                                   dataspace = captureResults.capturedDataspace, buffer, hdrTexture,
+                                   gainmapTexture](FenceResult fenceResult) -> FenceResult {
+                                if (!fenceResult.ok()) {
+                                    return fenceResult;
+                                }
+
+                                auto hdrFenceResult = hdrRenderFuture.get();
+
+                                if (!hdrFenceResult.ok()) {
+                                    return hdrFenceResult;
+                                }
+
+                                return getRenderEngine()
+                                        .drawGainmap(buffer, fenceResult.value()->get(), hdrTexture,
+                                                     hdrFenceResult.value()->get(),
+                                                     displayBrightnessNits / sdrWhitePointNits,
+                                                     static_cast<ui::Dataspace>(dataspace),
+                                                     gainmapTexture)
+                                        .get();
+                            })
+                            .share();
+        };
+    }
 
     if (captureListener) {
         // Defer blocking on renderFuture back to the Binder thread.
         return ftl::Future(std::move(renderFuture))
-                .then([captureListener, captureResults = std::move(captureResults)](
-                              FenceResult fenceResult) mutable -> FenceResult {
+                .then([captureListener, captureResults = std::move(captureResults),
+                       displayBrightnessNits,
+                       sdrWhitePointNits](FenceResult fenceResult) mutable -> FenceResult {
                     captureResults.fenceResult = std::move(fenceResult);
+                    captureResults.hdrSdrRatio = displayBrightnessNits / sdrWhitePointNits;
                     captureListener->onScreenCaptureCompleted(captureResults);
                     return base::unexpected(NO_ERROR);
                 })
@@ -7459,7 +7530,8 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshot(
 ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshotLegacy(
         RenderAreaBuilderVariant renderAreaBuilder, GetLayerSnapshotsFunction getLayerSnapshotsFn,
         const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
-        bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener) {
+        bool grayscale, bool isProtected, bool attachGainmap,
+        const sp<IScreenCaptureListener>& captureListener) {
     SFTRACE_CALL();
 
     auto takeScreenshotFn = [=, this, renderAreaBuilder = std::move(renderAreaBuilder)]() REQUIRES(
@@ -7488,8 +7560,8 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshotLegacy(
 
         auto layerFEs = extractLayerFEs(layers);
         ftl::SharedFuture<FenceResult> renderFuture =
-                renderScreenImpl(std::move(renderArea), buffer, regionSampling, grayscale,
-                                 isProtected, captureResults, displayState, layers, layerFEs);
+                renderScreenImpl(renderArea.get(), buffer, regionSampling, grayscale, isProtected,
+                                 attachGainmap, captureResults, displayState, layers, layerFEs);
 
         if (captureListener) {
             // Defer blocking on renderFuture back to the Binder thread.
@@ -7519,10 +7591,9 @@ ftl::SharedFuture<FenceResult> SurfaceFlinger::captureScreenshotLegacy(
 }
 
 ftl::SharedFuture<FenceResult> SurfaceFlinger::renderScreenImpl(
-        std::unique_ptr<const RenderArea> renderArea,
-        const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
-        bool grayscale, bool isProtected, ScreenCaptureResults& captureResults,
-        std::optional<OutputCompositionState>& displayState,
+        const RenderArea* renderArea, const std::shared_ptr<renderengine::ExternalTexture>& buffer,
+        bool regionSampling, bool grayscale, bool isProtected, bool attachGainmap,
+        ScreenCaptureResults& captureResults, std::optional<OutputCompositionState>& displayState,
         std::vector<std::pair<Layer*, sp<LayerFE>>>& layers, std::vector<sp<LayerFE>>& layerFEs) {
     SFTRACE_CALL();
 
diff --git a/services/surfaceflinger/SurfaceFlinger.h b/services/surfaceflinger/SurfaceFlinger.h
index 282c8cf803..7fb5c36aad 100644
--- a/services/surfaceflinger/SurfaceFlinger.h
+++ b/services/surfaceflinger/SurfaceFlinger.h
@@ -869,7 +869,7 @@ private:
 
     void captureScreenCommon(RenderAreaBuilderVariant, GetLayerSnapshotsFunction,
                              ui::Size bufferSize, ui::PixelFormat, bool allowProtected,
-                             bool grayscale, const sp<IScreenCaptureListener>&);
+                             bool grayscale, bool attachGainmap, const sp<IScreenCaptureListener>&);
 
     std::optional<OutputCompositionState> getDisplayStateFromRenderAreaBuilder(
             RenderAreaBuilderVariant& renderAreaBuilder) REQUIRES(kMainThreadContext);
@@ -883,20 +883,21 @@ private:
     ftl::SharedFuture<FenceResult> captureScreenshot(
             const RenderAreaBuilderVariant& renderAreaBuilder,
             const std::shared_ptr<renderengine::ExternalTexture>& buffer, bool regionSampling,
-            bool grayscale, bool isProtected, const sp<IScreenCaptureListener>& captureListener,
+            bool grayscale, bool isProtected, bool attachGainmap,
+            const sp<IScreenCaptureListener>& captureListener,
             std::optional<OutputCompositionState>& displayState,
             std::vector<sp<LayerFE>>& layerFEs);
 
     ftl::SharedFuture<FenceResult> captureScreenshotLegacy(
             RenderAreaBuilderVariant, GetLayerSnapshotsFunction,
             const std::shared_ptr<renderengine::ExternalTexture>&, bool regionSampling,
-            bool grayscale, bool isProtected, const sp<IScreenCaptureListener>&);
+            bool grayscale, bool isProtected, bool attachGainmap,
+            const sp<IScreenCaptureListener>&);
 
     ftl::SharedFuture<FenceResult> renderScreenImpl(
-            std::unique_ptr<const RenderArea>,
-            const std::shared_ptr<renderengine::ExternalTexture>&, bool regionSampling,
-            bool grayscale, bool isProtected, ScreenCaptureResults&,
-            std::optional<OutputCompositionState>& displayState,
+            const RenderArea*, const std::shared_ptr<renderengine::ExternalTexture>&,
+            bool regionSampling, bool grayscale, bool isProtected, bool attachGainmap,
+            ScreenCaptureResults&, std::optional<OutputCompositionState>& displayState,
             std::vector<std::pair<Layer*, sp<LayerFE>>>& layers,
             std::vector<sp<LayerFE>>& layerFEs);
 
diff --git a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
index a6a275804d..aa9af5d4ca 100644
--- a/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
+++ b/services/surfaceflinger/tests/unittests/TestableSurfaceFlinger.h
@@ -501,9 +501,10 @@ public:
         auto layers = getLayerSnapshotsFn();
         auto layerFEs = mFlinger->extractLayerFEs(layers);
 
-        return mFlinger->renderScreenImpl(std::move(renderArea), buffer, regionSampling,
+        return mFlinger->renderScreenImpl(renderArea.get(), buffer, regionSampling,
                                           false /* grayscale */, false /* isProtected */,
-                                          captureResults, displayState, layers, layerFEs);
+                                          false /* attachGainmap */, captureResults, displayState,
+                                          layers, layerFEs);
     }
 
     auto getLayerSnapshotsForScreenshotsFn(ui::LayerStack layerStack, uint32_t uid) {