perf(windows-amd): MF texture encode path for AMD AMF

.github/workflows/instant-reliability.yml

@@ -316,8 +316,11 @@ jobs:
         run: |
           set -euo pipefail
           sudo apt-get update -qq
-          # v4l2loopback needs videodev, which is in the extra kernel modules.
-          sudo apt-get install -y -qq "linux-modules-extra-$(uname -r)" v4l2loopback-dkms xvfb x11-utils xdotool \
+          if ! sudo modprobe videodev 2>/dev/null; then
+            sudo apt-get install -y -qq "linux-modules-extra-$(uname -r)"
+            sudo modprobe videodev
+          fi
+          sudo apt-get install -y -qq "linux-headers-$(uname -r)" v4l2loopback-dkms xvfb x11-utils xdotool \
             dbus ffmpeg pulseaudio pulseaudio-utils libasound2-plugins >/dev/null
           chmod +x "$RUNNER_TEMP/build/cap-gpui" "$RUNNER_TEMP/build/real-app-tests"
           sudo modprobe v4l2loopback video_nr=42 card_label="Cap Test Camera" exclusive_caps=1

crates/enc-mediafoundation/src/video/h264.rs

@@ -17,22 +17,26 @@ use windows::{
         Foundation::E_NOTIMPL,
         Graphics::{
             Direct3D11::{ID3D11Device, ID3D11Texture2D},
-            Dxgi::Common::{DXGI_FORMAT, DXGI_FORMAT_NV12},
+            Dxgi::{
+                Common::{DXGI_FORMAT, DXGI_FORMAT_NV12},
+                IDXGIDevice,
+            },
         },
         Media::MediaFoundation::{
             self, IMFAttributes, IMFDXGIDeviceManager, IMFMediaEventGenerator, IMFMediaType,
-            IMFSample, IMFTransform, MF_E_INVALIDMEDIATYPE, MF_E_NO_MORE_TYPES,
-            MF_E_TRANSFORM_TYPE_NOT_SET, MF_EVENT_FLAG_NONE, MF_EVENT_TYPE,
-            MF_MT_ALL_SAMPLES_INDEPENDENT, MF_MT_AVG_BITRATE, MF_MT_FRAME_RATE, MF_MT_FRAME_SIZE,
-            MF_MT_INTERLACE_MODE, MF_MT_MAJOR_TYPE, MF_MT_PIXEL_ASPECT_RATIO, MF_MT_SUBTYPE,
-            MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS, MF_TRANSFORM_ASYNC_UNLOCK,
-            MFCreateDXGIDeviceManager, MFCreateDXGISurfaceBuffer, MFCreateMediaType,
-            MFCreateSample, MFMediaType_Video, MFT_ENUM_FLAG, MFT_ENUM_FLAG_HARDWARE,
-            MFT_ENUM_FLAG_TRANSCODE_ONLY, MFT_MESSAGE_COMMAND_FLUSH,
-            MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, MFT_MESSAGE_NOTIFY_END_OF_STREAM,
-            MFT_MESSAGE_NOTIFY_END_STREAMING, MFT_MESSAGE_NOTIFY_START_OF_STREAM,
-            MFT_MESSAGE_SET_D3D_MANAGER, MFT_OUTPUT_DATA_BUFFER, MFT_SET_TYPE_TEST_ONLY,
-            MFVideoFormat_H264, MFVideoFormat_NV12, MFVideoInterlace_Progressive,
+            IMFSample, IMFTransform, MF_E_INVALIDMEDIATYPE, MF_E_NO_EVENTS_AVAILABLE,
+            MF_E_NO_MORE_TYPES, MF_E_TRANSFORM_TYPE_NOT_SET, MF_EVENT_FLAG_NO_WAIT,
+            MF_EVENT_FLAG_NONE, MF_EVENT_TYPE, MF_MT_ALL_SAMPLES_INDEPENDENT, MF_MT_AVG_BITRATE,
+            MF_MT_FRAME_RATE, MF_MT_FRAME_SIZE, MF_MT_INTERLACE_MODE, MF_MT_MAJOR_TYPE,
+            MF_MT_PIXEL_ASPECT_RATIO, MF_MT_SUBTYPE, MF_READWRITE_ENABLE_HARDWARE_TRANSFORMS,
+            MF_TRANSFORM_ASYNC_UNLOCK, MFCreateDXGIDeviceManager, MFCreateDXGISurfaceBuffer,
+            MFCreateMediaType, MFCreateSample, MFMediaType_Video, MFT_ENUM_FLAG,
+            MFT_ENUM_FLAG_HARDWARE, MFT_ENUM_FLAG_TRANSCODE_ONLY, MFT_MESSAGE_COMMAND_DRAIN,
+            MFT_MESSAGE_COMMAND_FLUSH, MFT_MESSAGE_NOTIFY_BEGIN_STREAMING,
+            MFT_MESSAGE_NOTIFY_END_OF_STREAM, MFT_MESSAGE_NOTIFY_END_STREAMING,
+            MFT_MESSAGE_NOTIFY_START_OF_STREAM, MFT_MESSAGE_SET_D3D_MANAGER,
+            MFT_OUTPUT_DATA_BUFFER, MFT_SET_TYPE_TEST_ONLY, MFVideoFormat_H264, MFVideoFormat_NV12,
+            MFVideoInterlace_Progressive,
         },
     },
     core::{Error, Interface},
@@ -163,6 +167,7 @@ pub struct H264Encoder {
     output_stream_id: u32,
     output_type: IMFMediaType,
     bitrate: u32,
+    drain_on_stop: bool,
 }
 
 #[derive(Clone, Debug, thiserror::Error)]
@@ -251,6 +256,12 @@ impl H264Encoder {
             frame_rate,
             bitrate_multipler,
         );
+        let adapter = d3d_device
+            .cast::<IDXGIDevice>()
+            .and_then(|device| unsafe { device.GetAdapter()?.GetDesc() })
+            .map_err(NewVideoEncoderError::DeviceManager)?;
+        // Match `cap_frame_converter::GpuVendor::from_id`: AMD reports either PCI ID.
+        let is_amd = matches!(adapter.VendorId, 0x1002 | 0x1022);
 
         let transform =
             EncoderDevice::enumerate_with_flags(MFMediaType_Video, MFVideoFormat_H264, flags)
@@ -270,6 +281,9 @@ impl H264Encoder {
             frame_rate,
         )
         .map_err(NewVideoEncoderError::VideoProcessor)?;
+        if is_amd {
+            video_processor.set_bt709_color_space(format);
+        }
 
         let mut device_manager_reset_token: u32 = 0;
         let media_device_manager = {
@@ -367,6 +381,9 @@ impl H264Encoder {
             MFSetAttributeRatio(&attributes, &MF_MT_PIXEL_ASPECT_RATIO, 1, 1)?;
             output_type.SetUINT32(&MF_MT_INTERLACE_MODE, MFVideoInterlace_Progressive.0 as u32)?;
             output_type.SetUINT32(&MF_MT_ALL_SAMPLES_INDEPENDENT, 1)?;
+            if is_amd {
+                set_bt709_attributes(&output_type)?;
+            }
             transform.SetOutputType(output_stream_id, &output_type, 0)?;
             Ok(output_type)
         })()
@@ -386,6 +403,9 @@ impl H264Encoder {
                 let attributes: IMFAttributes = input_type.cast()?;
                 input_type.SetGUID(&MF_MT_MAJOR_TYPE, &MFMediaType_Video)?;
                 input_type.SetGUID(&MF_MT_SUBTYPE, &MFVideoFormat_NV12)?;
+                if is_amd {
+                    set_bt709_attributes(&input_type)?;
+                }
                 MFSetAttributeSize(
                     &attributes,
                     &MF_MT_FRAME_SIZE,
@@ -431,6 +451,7 @@ impl H264Encoder {
             input_stream_id,
             output_stream_id,
             bitrate,
+            drain_on_stop: is_amd,
 
             output_type,
         })
@@ -557,9 +578,25 @@ impl H264Encoder {
 
             let mut consecutive_empty_samples: u8 = 0;
             let mut should_exit = false;
+            let mut drain_started: Option<Instant> = None;
             while !should_exit {
+                if let Some(started) = drain_started
+                    && started.elapsed() > ENCODER_OPERATION_TIMEOUT
+                {
+                    // A slow drain only costs the tail frames; never fail the
+                    // recording over it. Fall through to the normal shutdown.
+                    tracing::warn!(
+                        frames_encoded = health_monitor.total_frames_encoded,
+                        "MediaFoundation H264 drain did not complete within {:?}; finishing stop without remaining output",
+                        ENCODER_OPERATION_TIMEOUT
+                    );
+                    break;
+                }
+                // While draining, the drain deadline above is the only timeout;
+                // the health check would otherwise count pre-stop idle time.
                 let health_status = health_monitor.check_health();
-                if !health_status.is_healthy
+                if drain_started.is_none()
+                    && !health_status.is_healthy
                     && let Some(reason) = health_status.failure_reason
                 {
                     let _ = self.cleanup_encoder();
@@ -571,11 +608,25 @@ impl H264Encoder {
                     });
                 }
 
-                let event = self.event_generator.GetEvent(MF_EVENT_FLAG_NONE)?;
+                let event = if drain_started.is_some() {
+                    match self.event_generator.GetEvent(MF_EVENT_FLAG_NO_WAIT) {
+                        Ok(event) => event,
+                        Err(error) if error.code() == MF_E_NO_EVENTS_AVAILABLE => {
+                            std::thread::sleep(Duration::from_millis(1));
+                            continue;
+                        }
+                        Err(error) => return Err(error.into()),
+                    }
+                } else {
+                    self.event_generator.GetEvent(MF_EVENT_FLAG_NONE)?
+                };
 
                 let event_type = MF_EVENT_TYPE(event.GetType()? as i32);
                 match event_type {
                     MediaFoundation::METransformNeedInput => {
+                        if drain_started.is_some() {
+                            continue;
+                        }
                         health_monitor.record_input();
                         should_exit = true;
                         let get_frame_started = Instant::now();
@@ -631,6 +682,18 @@ impl H264Encoder {
                                     should_exit = false;
                                 }
                             }
+                        } else if self.drain_on_stop {
+                            // AMD retains input at stop; flushing here discards the final frames.
+                            self.transform
+                                .ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0)?;
+                            self.transform
+                                .ProcessMessage(MFT_MESSAGE_COMMAND_DRAIN, 0)?;
+                            // `get_frame` may have parked through a long pause
+                            // before reporting stop; that idle time is not an
+                            // encoder stall, so rebase the timeout clock.
+                            health_monitor.notify_input_stall_recovered();
+                            drain_started = Some(Instant::now());
+                            should_exit = false;
                         }
                     }
                     MediaFoundation::METransformHaveOutput => {
@@ -661,6 +724,9 @@ impl H264Encoder {
                             }
                         }
                     }
+                    MediaFoundation::METransformDrainComplete if drain_started.is_some() => {
+                        should_exit = true;
+                    }
                     _ => {
                         panic!("Unknown media event type: {}", event_type.0);
                     }
@@ -695,3 +761,17 @@ fn calculate_bitrate(width: u32, height: u32, fps: u32, multiplier: f32) -> u32
     let frame_rate_factor = (fps as f32 - 30.0).max(0.0) / 2.0 + 30.0;
     (width as f32 * height as f32 * frame_rate_factor * multiplier) as u32
 }
+
+fn set_bt709_attributes(media_type: &IMFMediaType) -> windows::core::Result<()> {
+    use MediaFoundation::{
+        MF_MT_TRANSFER_FUNCTION, MF_MT_VIDEO_NOMINAL_RANGE, MF_MT_VIDEO_PRIMARIES,
+        MF_MT_YUV_MATRIX, MFNominalRange_16_235, MFVideoPrimaries_BT709, MFVideoTransFunc_709,
+        MFVideoTransferMatrix_BT709,
+    };
+    unsafe {
+        media_type.SetUINT32(&MF_MT_YUV_MATRIX, MFVideoTransferMatrix_BT709.0 as u32)?;
+        media_type.SetUINT32(&MF_MT_VIDEO_PRIMARIES, MFVideoPrimaries_BT709.0 as u32)?;
+        media_type.SetUINT32(&MF_MT_TRANSFER_FUNCTION, MFVideoTransFunc_709.0 as u32)?;
+        media_type.SetUINT32(&MF_MT_VIDEO_NOMINAL_RANGE, MFNominalRange_16_235.0 as u32)
+    }
+}

crates/enc-mediafoundation/src/video/video_processor.rs

@@ -18,7 +18,10 @@ use windows::{
                 ID3D11VideoContext, ID3D11VideoDevice, ID3D11VideoProcessor,
                 ID3D11VideoProcessorInputView, ID3D11VideoProcessorOutputView,
             },
-            Dxgi::Common::{DXGI_FORMAT, DXGI_RATIONAL, DXGI_SAMPLE_DESC},
+            Dxgi::Common::{
+                DXGI_FORMAT, DXGI_FORMAT_B8G8R8A8_UNORM, DXGI_FORMAT_R8G8B8A8_UNORM, DXGI_RATIONAL,
+                DXGI_SAMPLE_DESC,
+            },
         },
     },
     core::Interface,
@@ -218,6 +221,30 @@ impl VideoProcessor {
         &self.video_output_texture
     }
 
+    pub(super) fn set_bt709_color_space(&self, input_format: DXGI_FORMAT) {
+        let input_range = if matches!(
+            input_format,
+            DXGI_FORMAT_B8G8R8A8_UNORM | DXGI_FORMAT_R8G8B8A8_UNORM
+        ) {
+            D3D11_VIDEO_PROCESSOR_NOMINAL_RANGE_0_255
+        } else {
+            D3D11_VIDEO_PROCESSOR_NOMINAL_RANGE_16_235
+        };
+        // Nominal_Range occupies bits 4-5; unshifted values change Usage/RGB_Range.
+        let output = D3D11_VIDEO_PROCESSOR_COLOR_SPACE {
+            _bitfield: 1 | (1 << 2) | ((D3D11_VIDEO_PROCESSOR_NOMINAL_RANGE_16_235.0 as u32) << 4),
+        };
+        let input = D3D11_VIDEO_PROCESSOR_COLOR_SPACE {
+            _bitfield: 1 | (1 << 2) | ((input_range.0 as u32) << 4),
+        };
+        unsafe {
+            self.video_context
+                .VideoProcessorSetOutputColorSpace(&self.video_processor, &output);
+            self.video_context
+                .VideoProcessorSetStreamColorSpace(&self.video_processor, 0, &input);
+        }
+    }
+
     pub fn process_texture(
         &mut self,
         input_texture: &ID3D11Texture2D,