mediafoundation: look into using texture pool for metadata retrieval, e.g SATD, Bitsused map, etc.

frontend/mediafoundation: use texture pool for SATD map and Bitsused map
The usage of texture pool depends on the updated mfplat.dll with a fix related to D3DFMT_INDEX32.
If the mfplat.dll on the machine does not have the fix, it falls back to the original implementation without the texture pool.

Reviewed-by: Sil Vilerino <sivileri@microsoft.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37376>

src/gallium/frontends/mediafoundation/context.h

@@ -46,6 +46,8 @@ typedef class DX12EncodeContext
    pipe_resource *pPipeResourceSATDMapStats = nullptr;
    pipe_resource *pPipeResourceRCBitAllocMapStats = nullptr;
    pipe_resource *pPipeResourcePSNRStats = nullptr;
+   BOOL bUseSATDMapAllocator = FALSE;
+   BOOL bUseBitsusedMapAllocator = FALSE;
 
    // Keep all the media and sync objects until encode is done
    // and then signal EnqueueResourceRelease so the media
@@ -185,10 +187,27 @@ typedef class DX12EncodeContext
 
       if( pPipeResourceQPMapStats )
          pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourceQPMapStats );
-      if( pPipeResourceSATDMapStats )
-         pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourceSATDMapStats );
-      if( pPipeResourceRCBitAllocMapStats )
-         pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourceRCBitAllocMapStats );
+
+      if( bUseSATDMapAllocator )
+      {
+         pPipeResourceSATDMapStats = nullptr;
+      }
+      else
+      {
+         if( pPipeResourceSATDMapStats )
+            pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourceSATDMapStats );
+      }
+
+      if( bUseBitsusedMapAllocator )
+      {
+         pPipeResourceRCBitAllocMapStats = nullptr;
+      }
+      else
+      {
+         if( pPipeResourceRCBitAllocMapStats )
+            pVlScreen->pscreen->resource_destroy( pVlScreen->pscreen, pPipeResourceRCBitAllocMapStats );
+      }
+
       if( pPipeVideoBuffer )
          pPipeVideoBuffer->destroy( pPipeVideoBuffer );
       if( pDownscaledTwoPassPipeVideoBuffer )

src/gallium/frontends/mediafoundation/encode.cpp

@@ -385,10 +385,20 @@ CDX12EncHMFT::PrepareForEncode( IMFSample *pSample, LPDX12EncodeContext *ppDX12E
          templ.format = (enum pipe_format) m_EncoderCapabilities.m_HWSupportStatsSATDMapOutput.bits.pipe_pixel_format;
          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,
-            done );
+         pDX12EncodeContext->bUseSATDMapAllocator = m_bUseSATDMapAllocator;
+         if ( m_bUseSATDMapAllocator )
+         {
+            pDX12EncodeContext->pPipeResourceSATDMapStats =
+               AllocatePipeResourceFromAllocator( m_spSATDMapAllocator.Get(), m_pVlScreen->pscreen, &templ );
+            CHECKNULL_GOTO( pDX12EncodeContext->pPipeResourceSATDMapStats, E_OUTOFMEMORY, done );
+         }
+         else
+         {
+            CHECKNULL_GOTO(
+               pDX12EncodeContext->pPipeResourceSATDMapStats = m_pVlScreen->pscreen->resource_create( m_pVlScreen->pscreen, &templ ),
+               E_OUTOFMEMORY,
+               done );
+         }
       }
 
       if( m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput.bits.supported && m_uiVideoOutputBitsUsedMapBlockSize > 0 )
@@ -397,10 +407,21 @@ CDX12EncHMFT::PrepareForEncode( IMFSample *pSample, LPDX12EncodeContext *ppDX12E
          templ.format = (enum pipe_format) m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput.bits.pipe_pixel_format;
          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 =
+
+         pDX12EncodeContext->bUseBitsusedMapAllocator = m_bUseBitsusedMapAllocator;
+         if( m_bUseBitsusedMapAllocator )
+         {
+            pDX12EncodeContext->pPipeResourceRCBitAllocMapStats =
+               AllocatePipeResourceFromAllocator( m_spBitsusedMapAllocator.Get(), m_pVlScreen->pscreen, &templ );
+            CHECKNULL_GOTO( pDX12EncodeContext->pPipeResourceRCBitAllocMapStats, E_OUTOFMEMORY, done );
+         }
+         else
+         {
+            CHECKNULL_GOTO( pDX12EncodeContext->pPipeResourceRCBitAllocMapStats =
                             m_pVlScreen->pscreen->resource_create( m_pVlScreen->pscreen, &templ ),
-                         E_OUTOFMEMORY,
-                         done );
+                            E_OUTOFMEMORY,
+                            done );
+         }
       }
 
       if( m_EncoderCapabilities.m_PSNRStatsSupport.bits.supports_y_channel && m_bVideoEnableFramePsnrYuv )

src/gallium/frontends/mediafoundation/hmft_entrypoints.h

@@ -55,6 +55,7 @@
 #include <d3d11_3.h>
 #include <d3d11_4.h>
 #include <dxgi1_2.h>
+#include <d3d9types.h>
 
 #include "context.h"
 #include "encoder_capabilities.h"
@@ -247,6 +248,11 @@ typedef enum IntraRefreshMode
    HMFT_INTRA_REFRESH_MODE_MAX
 } IntraRefreshMode;
 
+// IMFVideoSampleAllocatorEx only works with MFVideoFormat.
+#ifndef MFVideoFormat_L32
+DEFINE_MEDIATYPE_GUID( MFVideoFormat_L32, D3DFMT_INDEX32 );
+#endif
+
 #ifndef CODECAPI_AVEncVideoEnableFramePsnrYuv
 // AVEncVideoEnableFramePsnrYuv (BOOL)
 // Indicates whether to enable or disable reporting frame PSNR of YUV planes for video encoding.
@@ -594,6 +600,8 @@ class __declspec( uuid( HMFT_GUID ) ) CDX12EncHMFT : CMFD3DManager,
    HRESULT OnFlush();
 
    HRESULT ConfigureSampleAllocator();
+   HRESULT ConfigureMapSampleAllocator( IMFVideoSampleAllocatorEx *spAllocator, UINT32 width, UINT32 height, GUID subtype, UINT32 poolSize );
+   void ConfigureMapSampleAllocatorHelper( ComPtr<IMFVideoSampleAllocatorEx> &allocator, const union pipe_enc_cap_gpu_stats_map &outputStatsMap, uint32_t blockSize, BOOL &useAllocatorFlag );
    HRESULT UpdateAvailableInputType();
    HRESULT InternalCheckInputType( IMFMediaType *pType );
    HRESULT InternalCheckOutputType( IMFMediaType *pType );

src/gallium/frontends/mediafoundation/mfbufferhelp.cpp

@@ -343,3 +343,56 @@ MFAttachPipeResourceAsSampleExtension( struct pipe_context *pPipeContext,
 
    return pSample->SetUnknown( guidExtension, spMediaBuffer.Get() );
 }
+
+// Helper to convert the allocated IMFSample to pipe_resource*
+// Returns NULL on failure.
+struct pipe_resource *
+AllocatePipeResourceFromAllocator( IMFVideoSampleAllocatorEx *pAllocator,
+                                   struct pipe_screen *pScreen,
+                                   const struct pipe_resource *templ )
+{
+   if( !pAllocator || !pScreen || !templ )
+      return nullptr;
+
+   // Allocate or get a sample from the pool.
+   ComPtr<IMFSample> spSample;
+   HRESULT hr = pAllocator->AllocateSample( &spSample );
+   if( FAILED( hr ) || !spSample )
+      return nullptr;
+
+   // Get the underlying D3D12 resource from the sample
+   ComPtr<IMFMediaBuffer> spBuffer;
+   hr = spSample->GetBufferByIndex( 0, &spBuffer );
+   if( FAILED( hr ) || !spBuffer )
+      return nullptr;
+
+   ComPtr<IMFDXGIBuffer> spDXGIBuffer;
+   hr = spBuffer.As( &spDXGIBuffer );
+   if( FAILED( hr ) || !spDXGIBuffer )
+      return nullptr;
+
+   ComPtr<ID3D12Resource> spResource;
+   hr = spDXGIBuffer->GetResource( IID_PPV_ARGS( &spResource ) );
+   if( FAILED( hr ) || !spResource )
+      return nullptr;
+
+   // Build winsys_handle
+   struct winsys_handle whandle = {};
+   whandle.type = WINSYS_HANDLE_TYPE_D3D12_RES;
+   whandle.com_obj = spResource.Detach();
+   // templ->format contains the desired pipe_format
+   whandle.format = templ->format;
+
+   // Call resource_from_handle with the same templ for resource_create.
+   struct pipe_resource *pres = pScreen->resource_from_handle( pScreen, templ, &whandle, PIPE_USAGE_DEFAULT );
+
+   if (!pres) {
+      // Release the detached COM object if resource_from_handle fails
+      if (whandle.com_obj) {
+         static_cast<ID3D12Resource*>(whandle.com_obj)->Release();
+      }
+      return nullptr;
+   }
+
+   return pres;
+}

src/gallium/frontends/mediafoundation/mfbufferhelp.h

@@ -54,3 +54,8 @@ MFAttachPipeResourceAsSampleExtension( struct pipe_context *pPipeContext,
                                        ID3D12CommandQueue *pSyncObjectQueue,
                                        REFGUID guidExtension,
                                        IMFSample *pSample );
+
+struct pipe_resource *
+AllocatePipeResourceFromAllocator( IMFVideoSampleAllocatorEx *pAllocator,
+                                   struct pipe_screen *pScreen,
+                                   const struct pipe_resource *templ );

src/gallium/frontends/mediafoundation/mfd3dmanager.cpp

@@ -72,6 +72,18 @@ CMFD3DManager::Shutdown( bool bReleaseDeviceManager )
       m_spVideoSampleAllocator = nullptr;
    }
 
+   if( m_spSATDMapAllocator )
+   {
+      m_spSATDMapAllocator->UninitializeSampleAllocator();
+      m_spSATDMapAllocator = nullptr;
+   }
+
+   if( m_spBitsusedMapAllocator )
+   {
+      m_spBitsusedMapAllocator->UninitializeSampleAllocator();
+      m_spBitsusedMapAllocator = nullptr;
+   }
+
    if( m_spDeviceManager != nullptr )
    {
       if( m_hDevice != NULL )
@@ -363,6 +375,8 @@ CMFD3DManager::xOnSetD3DManager( ULONG_PTR ulParam )
    }
 #endif // ( USE_D3D12_PREVIEW_HEADERS && ( D3D12_PREVIEW_SDK_VERSION >= 717 ) )
    CHECKHR_GOTO( MFCreateVideoSampleAllocatorEx( IID_PPV_ARGS( &m_spVideoSampleAllocator ) ), done );
+   CHECKHR_GOTO( MFCreateVideoSampleAllocatorEx( IID_PPV_ARGS( &m_spSATDMapAllocator ) ), done );
+   CHECKHR_GOTO( MFCreateVideoSampleAllocatorEx( IID_PPV_ARGS( &m_spBitsusedMapAllocator ) ), done );
 
    CHECKHR_GOTO( GetDeviceInfo(), done );
 

src/gallium/frontends/mediafoundation/mfd3dmanager.h

@@ -100,6 +100,10 @@ class CMFD3DManager
    ComPtr<ID3D12VideoDevice> m_spVideoDevice;
    ComPtr<ID3D12CommandQueue> m_spStagingQueue;
    ComPtr<IMFVideoSampleAllocatorEx> m_spVideoSampleAllocator;   // Used for software input samples that need to be copied
+   ComPtr<IMFVideoSampleAllocatorEx> m_spSATDMapAllocator;
+   ComPtr<IMFVideoSampleAllocatorEx> m_spBitsusedMapAllocator;
+   BOOL m_bUseSATDMapAllocator = FALSE;
+   BOOL m_bUseBitsusedMapAllocator = FALSE;
    UINT32 m_uiResetToken = 0;
    HANDLE m_hDevice = NULL;
    struct vl_screen *m_pVlScreen = nullptr;

src/gallium/frontends/mediafoundation/mftransform.cpp

@@ -544,6 +544,19 @@ CDX12EncHMFT::OnOutputTypeChanged()
    }
    CHECKHR_GOTO( InitializeEncoder( m_outputPipeProfile, m_uiOutputWidth, m_uiOutputHeight ), done );
 
+   if( bResolutionChange )
+   {
+      ConfigureMapSampleAllocatorHelper( m_spSATDMapAllocator,
+                                         m_EncoderCapabilities.m_HWSupportStatsSATDMapOutput,
+                                         m_uiVideoSatdMapBlockSize,
+                                         m_bUseSATDMapAllocator );
+
+      ConfigureMapSampleAllocatorHelper( m_spBitsusedMapAllocator,
+                                         m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput,
+                                         m_uiVideoOutputBitsUsedMapBlockSize,
+                                         m_bUseBitsusedMapAllocator );
+   }
+
    if( m_gpuFeatureFlags.m_bDisableAsync )
    {
       MFE_INFO( "[dx12 hmft 0x%p] Async is disabled due to lack of GPU support.", this );
@@ -987,6 +1000,68 @@ done:
    return hr;
 }
 
+// Utility function to configure the sample allocator to allocate map samples
+HRESULT
+CDX12EncHMFT::ConfigureMapSampleAllocator( IMFVideoSampleAllocatorEx *spAllocator,
+                                           UINT32 width,
+                                           UINT32 height,
+                                           GUID subtype,
+                                           UINT32 poolSize )
+{
+   if( !spAllocator )
+      return E_POINTER;
+
+   spAllocator->UninitializeSampleAllocator();
+   HRESULT hr = spAllocator->SetDirectXManager( m_spDeviceManager.Get() );
+   if( FAILED( hr ) )
+      return hr;
+
+   ComPtr<IMFAttributes> spAttrs;
+   ComPtr<IMFMediaType> spMapType;
+
+   // Attributes for allocator
+   CHECKHR_GOTO( MFCreateAttributes( &spAttrs, 2 ), done );
+   CHECKHR_GOTO( spAttrs->SetUINT32( MF_SA_BUFFERS_PER_SAMPLE, 1 ), done );
+   CHECKHR_GOTO( spAttrs->SetUINT32( MF_MT_D3D_RESOURCE_VERSION, MF_D3D12_RESOURCE ), done );
+
+   // Media type for the map
+   CHECKHR_GOTO( MFCreateMediaType( &spMapType ), done );
+   CHECKHR_GOTO( spMapType->SetGUID( MF_MT_MAJOR_TYPE, MFMediaType_Video ), done );
+   CHECKHR_GOTO( spMapType->SetGUID( MF_MT_SUBTYPE, subtype ), done );
+   MFSetAttributeSize( spMapType.Get(), MF_MT_FRAME_SIZE, width, height );
+   CHECKHR_GOTO( spMapType->SetUINT32( MF_MT_D3D_RESOURCE_VERSION, MF_D3D12_RESOURCE ), done );
+
+   // Initialize the allocator
+   CHECKHR_GOTO( spAllocator->InitializeSampleAllocatorEx( 1, poolSize, spAttrs.Get(), spMapType.Get() ), done );
+
+done:
+   return hr;
+}
+
+// Helper function to configure map sample allocator.
+void
+CDX12EncHMFT::ConfigureMapSampleAllocatorHelper( ComPtr<IMFVideoSampleAllocatorEx> &allocator,
+                                                 const union pipe_enc_cap_gpu_stats_map &outputStatsMap,
+                                                 uint32_t blockSize,
+                                                 BOOL &useAllocatorFlag )
+{
+   if( allocator != nullptr && outputStatsMap.bits.supported && blockSize > 0 )
+   {
+      uint32_t actualBlockSize = ( 1u << outputStatsMap.bits.log2_values_block_size );
+      uint32_t width = static_cast<uint32_t>( std::ceil( m_uiOutputWidth / static_cast<float>( actualBlockSize ) ) );
+      uint32_t height = static_cast<uint32_t>( std::ceil( m_uiOutputHeight / static_cast<float>( actualBlockSize ) ) );
+
+      if( SUCCEEDED( ConfigureMapSampleAllocator( allocator.Get(), width, height, MFVideoFormat_L32, 10 ) ) )
+      {
+         useAllocatorFlag = TRUE;
+      }
+      else
+      {
+         useAllocatorFlag = FALSE;
+      }
+   }
+}
+
 // internal thread function to handle encoding and output
 void WINAPI
 CDX12EncHMFT::xThreadProc( void *pCtx )
@@ -1930,6 +2005,16 @@ CDX12EncHMFT::ProcessMessage( MFT_MESSAGE_TYPE eMessage, ULONG_PTR ulParam )
          {
             m_EncoderCapabilities.initialize( m_pPipeContext->screen, m_outputPipeProfile );
          }
+
+         ConfigureMapSampleAllocatorHelper( m_spSATDMapAllocator,
+                                            m_EncoderCapabilities.m_HWSupportStatsSATDMapOutput,
+                                            m_uiVideoSatdMapBlockSize,
+                                            m_bUseSATDMapAllocator );
+
+         ConfigureMapSampleAllocatorHelper( m_spBitsusedMapAllocator,
+                                            m_EncoderCapabilities.m_HWSupportStatsRCBitAllocationMapOutput,
+                                            m_uiVideoOutputBitsUsedMapBlockSize,
+                                            m_bUseBitsusedMapAllocator );
          break;
       }
    }