mediafoundation: Request PSNR frame stats

Reviewed-by: Yubo Xie <yuboxie@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/35264>

5 files changed, 37 insertions, 6 deletions

diff --git a/src/gallium/frontends/mediafoundation/context.h b/src/gallium/frontends/mediafoundation/context.h
index b71af769313..e206e25910e 100644
--- a/src/gallium/frontends/mediafoundation/context.h
+++ b/src/gallium/frontends/mediafoundation/context.h
@@ -45,6 +45,7 @@ typedef class DX12EncodeContext
    pipe_resource *pPipeResourceQPMapStats = nullptr;
    pipe_resource *pPipeResourceSATDMapStats = nullptr;
    pipe_resource *pPipeResourceRCBitAllocMapStats = nullptr;
+   pipe_resource *pPipeResourcePSNRStats = nullptr;
 
    // Keep all the media and sync objects until encode is done
    // and then signal EnqueueResourceRelease so the media
@@ -191,5 +192,7 @@ typedef class DX12EncodeContext
          pPipeVideoBuffer->destroy( pPipeVideoBuffer );
       if( pDownscaledTwoPassPipeVideoBuffer )
          pDownscaledTwoPassPipeVideoBuffer->destroy( pDownscaledTwoPassPipeVideoBuffer );
+      if( pPipeResourcePSNRStats )
+         pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourcePSNRStats );
    }
 } *LPDX12EncodeContext;
diff --git a/src/gallium/frontends/mediafoundation/encode.cpp b/src/gallium/frontends/mediafoundation/encode.cpp
index 9ba754623df..48cc8ecaf0c 100644
--- a/src/gallium/frontends/mediafoundation/encode.cpp
+++ b/src/gallium/frontends/mediafoundation/encode.cpp
@@ -335,8 +335,8 @@ CDX12EncHMFT::PrepareForEncode( IMFSample *pSample, LPDX12EncodeContext *ppDX12E
       {
          uint32_t block_size = (1 << m_EncoderCapabilities.m_HWSupportStatsQPMapOutput.bits.log2_values_block_size);
          templ.format = (enum pipe_format) m_EncoderCapabilities.m_HWSupportStatsQPMapOutput.bits.pipe_pixel_format;
-         templ.width0 = static_cast<uint32_t>(std::ceil(alignedWidth / static_cast<float>(block_size)));
-         templ.height0 = static_cast<uint32_t>(std::ceil(alignedHeight / static_cast<float>(block_size)));
+         templ.width0 = static_cast<uint32_t>(std::ceil(m_uiOutputWidth / static_cast<float>(block_size)));
+         templ.height0 = static_cast<uint16_t>(std::ceil(m_uiOutputHeight / static_cast<float>(block_size)));
          CHECKNULL_GOTO(
             pDX12EncodeContext->pPipeResourceQPMapStats = m_pVlScreen->pscreen->resource_create(m_pVlScreen->pscreen, &templ),
             E_OUTOFMEMORY,
@@ -347,8 +347,8 @@ CDX12EncHMFT::PrepareForEncode( IMFSample *pSample, LPDX12EncodeContext *ppDX12E
       {
          uint32_t block_size = (1 << m_EncoderCapabilities.m_HWSupportStatsSATDMapOutput.bits.log2_values_block_size);
          templ.format = (enum pipe_format) m_EncoderCapabilities.m_HWSupportStatsSATDMapOutput.bits.pipe_pixel_format;
-         templ.width0 = static_cast<uint32_t>(std::ceil(alignedWidth / static_cast<float>(block_size)));
-         templ.height0 = static_cast<uint32_t>(std::ceil(alignedHeight / static_cast<float>(block_size)));
+         templ.width0 = static_cast<uint32_t>(std::ceil(m_uiOutputWidth / static_cast<float>(block_size)));
+         templ.height0 = static_cast<uint16_t>(std::ceil(m_uiOutputHeight / static_cast<float>(block_size)));
          CHECKNULL_GOTO(
             pDX12EncodeContext->pPipeResourceSATDMapStats = m_pVlScreen->pscreen->resource_create(m_pVlScreen->pscreen, &templ),
             E_OUTOFMEMORY,
@@ -359,13 +359,34 @@ CDX12EncHMFT::PrepareForEncode( IMFSample *pSample, LPDX12EncodeContext *ppDX12E
       {
          uint32_t block_size = (1 << m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput.bits.log2_values_block_size);
          templ.format = (enum pipe_format) m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput.bits.pipe_pixel_format;
-         templ.width0 = static_cast<uint32_t>(std::ceil(alignedWidth / static_cast<float>(block_size)));
-         templ.height0 = static_cast<uint32_t>(std::ceil(alignedHeight / static_cast<float>(block_size)));
+         templ.width0 = static_cast<uint32_t>(std::ceil(m_uiOutputWidth / static_cast<float>(block_size)));
+         templ.height0 = static_cast<uint16_t>(std::ceil(m_uiOutputHeight / static_cast<float>(block_size)));
          CHECKNULL_GOTO(
             pDX12EncodeContext->pPipeResourceRCBitAllocMapStats = m_pVlScreen->pscreen->resource_create(m_pVlScreen->pscreen, &templ),
             E_OUTOFMEMORY,
             done);
       }
+
+      if (m_EncoderCapabilities.m_PSNRStatsSupport.bits.supports_y_channel)
+      {
+         struct pipe_resource buffer_templ = {};
+         buffer_templ.width0 = 3 * sizeof(float); // Up to 3 float components Y, U, V
+         buffer_templ.target = PIPE_BUFFER;
+         // PIPE_USAGE_STAGING allocates resource in L0 (System Memory) heap
+         // and avoid a bunch of roundtrips for uploading/reading back the bitstream headers
+         // The GPU writes once the slice data (if dGPU over the PCIe bus) and all the other
+         // uploads (e.g bitstream headers from CPU) and readbacks to output MFSamples
+         // happen without moving data between L0/L1 pools
+         buffer_templ.usage = PIPE_USAGE_STAGING;
+         buffer_templ.format = PIPE_FORMAT_R8_UINT;
+         buffer_templ.height0 = 1;
+         buffer_templ.depth0 = 1;
+         buffer_templ.array_size = 1;
+         CHECKNULL_GOTO(
+            pDX12EncodeContext->pPipeResourcePSNRStats = m_pVlScreen->pscreen->resource_create(m_pVlScreen->pscreen, &buffer_templ),
+            E_OUTOFMEMORY,
+            done);
+      }
    }
 
 #endif
diff --git a/src/gallium/frontends/mediafoundation/encode_h264.cpp b/src/gallium/frontends/mediafoundation/encode_h264.cpp
index 4d744afcecc..7635fea84b9 100644
--- a/src/gallium/frontends/mediafoundation/encode_h264.cpp
+++ b/src/gallium/frontends/mediafoundation/encode_h264.cpp
@@ -494,6 +494,7 @@ CDX12EncHMFT::PrepareForEncodeHelper( LPDX12EncodeContext pDX12EncodeContext, bo
    pPicInfo->gpu_stats_qp_map = pDX12EncodeContext->pPipeResourceQPMapStats;
    pPicInfo->gpu_stats_satd_map = pDX12EncodeContext->pPipeResourceSATDMapStats;
    pPicInfo->gpu_stats_rc_bitallocation_map = pDX12EncodeContext->pPipeResourceRCBitAllocMapStats;
+   pPicInfo->gpu_stats_psnr = pDX12EncodeContext->pPipeResourcePSNRStats;
 #endif
 
    // Quality vs speed
diff --git a/src/gallium/frontends/mediafoundation/encoder_capabilities.cpp b/src/gallium/frontends/mediafoundation/encoder_capabilities.cpp
index b0f9be97066..c452d6d9288 100644
--- a/src/gallium/frontends/mediafoundation/encoder_capabilities.cpp
+++ b/src/gallium/frontends/mediafoundation/encoder_capabilities.cpp
@@ -148,4 +148,7 @@ encoder_capabilities::initialize( pipe_screen *pScreen, pipe_video_profile video
 
    m_TwoPassSupport.value =
       pScreen->get_video_param( pScreen, videoProfile, PIPE_VIDEO_ENTRYPOINT_ENCODE, PIPE_VIDEO_CAP_ENC_TWO_PASS );
+
+   m_PSNRStatsSupport.value =
+      pScreen->get_video_param( pScreen, videoProfile, PIPE_VIDEO_ENTRYPOINT_ENCODE, PIPE_VIDEO_CAP_ENC_GPU_STATS_PSNR );
 }
diff --git a/src/gallium/frontends/mediafoundation/encoder_capabilities.h b/src/gallium/frontends/mediafoundation/encoder_capabilities.h
index b598e3d9dd6..a2e8bf89e4b 100644
--- a/src/gallium/frontends/mediafoundation/encoder_capabilities.h
+++ b/src/gallium/frontends/mediafoundation/encoder_capabilities.h
@@ -130,4 +130,7 @@ class encoder_capabilities
 
    // Two pass encode
    union pipe_enc_cap_two_pass m_TwoPassSupport = {};
+
+   // PSNR frame stats
+   union pipe_enc_cap_gpu_stats_psnr m_PSNRStatsSupport = {};
 };