Initial implementation of yuv444 encoding

LizardByte · Jun 16, 2024 · dd157c0 · dd157c0
1 parent b87fd42
commit dd157c0
Show file tree

Hide file tree

Showing 31 changed files with 996 additions and 192 deletions.
diff --git a/src/nvenc/nvenc_base.cpp b/src/nvenc/nvenc_base.cpp
@@ -81,9 +81,8 @@ namespace {
 
 namespace nvenc {
 
-  nvenc_base::nvenc_base(NV_ENC_DEVICE_TYPE device_type, void *device):
-      device_type(device_type),
-      device(device) {
+  nvenc_base::nvenc_base(NV_ENC_DEVICE_TYPE device_type):
+      device_type(device_type) {
   }
 
   nvenc_base::~nvenc_base() {
@@ -172,7 +171,7 @@ namespace nvenc {
     };
 
     auto buffer_is_yuv444 = [&]() {
-      return buffer_format == NV_ENC_BUFFER_FORMAT_YUV444 || buffer_format == NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
+      return buffer_format == NV_ENC_BUFFER_FORMAT_AYUV || buffer_format == NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
     };
 
     {
@@ -278,15 +277,15 @@ namespace nvenc {
       }
     };
 
-    auto fill_h264_hevc_vui = [&colorspace](auto &vui_config) {
+    auto fill_h264_hevc_vui = [&](auto &vui_config) {
       vui_config.videoSignalTypePresentFlag = 1;
       vui_config.videoFormat = NV_ENC_VUI_VIDEO_FORMAT_UNSPECIFIED;
       vui_config.videoFullRangeFlag = colorspace.full_range;
       vui_config.colourDescriptionPresentFlag = 1;
       vui_config.colourPrimaries = colorspace.primaries;
       vui_config.transferCharacteristics = colorspace.tranfer_function;
       vui_config.colourMatrix = colorspace.matrix;
-      vui_config.chromaSampleLocationFlag = 1;
+      vui_config.chromaSampleLocationFlag = buffer_is_yuv444() ? 0 : 1;
       vui_config.chromaSampleLocationTop = 0;
       vui_config.chromaSampleLocationBot = 0;
     };
@@ -327,7 +326,9 @@ namespace nvenc {
         auto &format_config = enc_config.encodeCodecConfig.av1Config;
         format_config.repeatSeqHdr = 1;
         format_config.idrPeriod = NVENC_INFINITE_GOPLENGTH;
-        format_config.chromaFormatIDC = 1;  // YUV444 not supported by NVENC yet
+        if (buffer_is_yuv444()) {
+          format_config.chromaFormatIDC = 3;
+        }
         format_config.enableBitstreamPadding = config.insert_filler_data;
         if (buffer_is_10bit()) {
           format_config.inputPixelBitDepthMinus8 = 2;
@@ -337,7 +338,7 @@ namespace nvenc {
         format_config.transferCharacteristics = colorspace.tranfer_function;
         format_config.matrixCoefficients = colorspace.matrix;
         format_config.colorRange = colorspace.full_range;
-        format_config.chromaSamplePosition = 1;
+        format_config.chromaSamplePosition = buffer_is_yuv444() ? 0 : 1;
         set_ref_frames(format_config.maxNumRefFramesInDPB, format_config.numFwdRefs, 8);
         set_minqp_if_enabled(config.min_qp_av1);
 
@@ -386,6 +387,7 @@ namespace nvenc {
     {
       std::string extra;
       if (init_params.enableEncodeAsync) extra += " async";
+      if (buffer_is_yuv444()) extra += " yuv444";
       if (buffer_is_10bit()) extra += " 10-bit";
       if (enc_config.rcParams.multiPass != NV_ENC_MULTI_PASS_DISABLED) extra += " two-pass";
       if (config.vbv_percentage_increase > 0 && get_encoder_cap(NV_ENC_CAPS_SUPPORT_CUSTOM_VBV_BUF_SIZE)) extra += " vbv+" + std::to_string(config.vbv_percentage_increase);
@@ -435,6 +437,11 @@ namespace nvenc {
     assert(registered_input_buffer);
     assert(output_bitstream);
 
+    if (!synchronize_input_buffer()) {
+      BOOST_LOG(error) << "NvEnc: failed to synchronize input buffer";
+      return {};
+    }
+
     NV_ENC_MAP_INPUT_RESOURCE mapped_input_buffer = { min_struct_version(NV_ENC_MAP_INPUT_RESOURCE_VER) };
     mapped_input_buffer.registeredResource = registered_input_buffer;
 
@@ -589,10 +596,12 @@ namespace nvenc {
 
     last_error_string.clear();
     if (status != NV_ENC_SUCCESS) {
+      /* This API function gives broken strings more often than not
       if (nvenc && encoder) {
         last_error_string = nvenc->nvEncGetLastErrorString(encoder);
         if (!last_error_string.empty()) last_error_string += " ";
       }
+      */
       last_error_string += status_string(status);
       return true;
     }

diff --git a/src/nvenc/nvenc_base.h b/src/nvenc/nvenc_base.h
@@ -13,7 +13,7 @@ namespace nvenc {
 
   class nvenc_base {
   public:
-    nvenc_base(NV_ENC_DEVICE_TYPE device_type, void *device);
+    nvenc_base(NV_ENC_DEVICE_TYPE device_type);
     virtual ~nvenc_base();
 
     nvenc_base(const nvenc_base &) = delete;
@@ -39,6 +39,9 @@ namespace nvenc {
     virtual bool
     create_and_register_input_buffer() = 0;
 
+    virtual bool
+    synchronize_input_buffer() { return true; }
+
     virtual bool
     wait_for_async_event(uint32_t timeout_ms) { return false; }
 
@@ -57,9 +60,6 @@ namespace nvenc {
     min_struct_version(uint32_t version, uint32_t v11_struct_version = 0, uint32_t v12_struct_version = 0);
 
     const NV_ENC_DEVICE_TYPE device_type;
-    void *const device;
-
-    std::unique_ptr<NV_ENCODE_API_FUNCTION_LIST> nvenc;
 
     void *encoder = nullptr;
 
@@ -72,6 +72,8 @@ namespace nvenc {
     } encoder_params;
 
     // Derived classes set these variables
+    void *device = nullptr;
+    std::shared_ptr<NV_ENCODE_API_FUNCTION_LIST> nvenc;
     NV_ENC_REGISTERED_PTR registered_input_buffer = nullptr;
     void *async_event_handle = nullptr;
 

diff --git a/src/nvenc/nvenc_d3d11.cpp b/src/nvenc/nvenc_d3d11.cpp
@@ -1,31 +1,15 @@
-#include "src/logging.h"
-
 #ifdef _WIN32
   #include "nvenc_d3d11.h"
 
-  #include "nvenc_utils.h"
-
 namespace nvenc {
 
-  nvenc_d3d11::nvenc_d3d11(ID3D11Device *d3d_device):
-      nvenc_base(NV_ENC_DEVICE_TYPE_DIRECTX, d3d_device),
-      d3d_device(d3d_device) {
-  }
-
   nvenc_d3d11::~nvenc_d3d11() {
-    if (encoder) destroy_encoder();
-
     if (dll) {
       FreeLibrary(dll);
       dll = NULL;
     }
   }
 
-  ID3D11Texture2D *
-  nvenc_d3d11::get_input_texture() {
-    return d3d_input_texture.GetInterfacePtr();
-  }
-
   bool
   nvenc_d3d11::init_library() {
     if (dll) return true;
@@ -64,43 +48,5 @@ namespace nvenc {
     return false;
   }
 
-  bool
-  nvenc_d3d11::create_and_register_input_buffer() {
-    if (!d3d_input_texture) {
-      D3D11_TEXTURE2D_DESC desc = {};
-      desc.Width = encoder_params.width;
-      desc.Height = encoder_params.height;
-      desc.MipLevels = 1;
-      desc.ArraySize = 1;
-      desc.Format = dxgi_format_from_nvenc_format(encoder_params.buffer_format);
-      desc.SampleDesc.Count = 1;
-      desc.Usage = D3D11_USAGE_DEFAULT;
-      desc.BindFlags = D3D11_BIND_RENDER_TARGET;
-      if (d3d_device->CreateTexture2D(&desc, nullptr, &d3d_input_texture) != S_OK) {
-        BOOST_LOG(error) << "NvEnc: couldn't create input texture";
-        return false;
-      }
-    }
-
-    if (!registered_input_buffer) {
-      NV_ENC_REGISTER_RESOURCE register_resource = { min_struct_version(NV_ENC_REGISTER_RESOURCE_VER, 3, 4) };
-      register_resource.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_DIRECTX;
-      register_resource.width = encoder_params.width;
-      register_resource.height = encoder_params.height;
-      register_resource.resourceToRegister = d3d_input_texture.GetInterfacePtr();
-      register_resource.bufferFormat = encoder_params.buffer_format;
-      register_resource.bufferUsage = NV_ENC_INPUT_IMAGE;
-
-      if (nvenc_failed(nvenc->nvEncRegisterResource(encoder, &register_resource))) {
-        BOOST_LOG(error) << "NvEncRegisterResource failed: " << last_error_string;
-        return false;
-      }
-
-      registered_input_buffer = register_resource.registeredResource;
-    }
-
-    return true;
-  }
-
 }  // namespace nvenc
 #endif
diff --git a/src/nvenc/nvenc_d3d11.h b/src/nvenc/nvenc_d3d11.h
@@ -10,25 +10,24 @@ namespace nvenc {
 
   _COM_SMARTPTR_TYPEDEF(ID3D11Device, IID_ID3D11Device);
   _COM_SMARTPTR_TYPEDEF(ID3D11Texture2D, IID_ID3D11Texture2D);
+  _COM_SMARTPTR_TYPEDEF(IDXGIDevice, IID_IDXGIDevice);
+  _COM_SMARTPTR_TYPEDEF(IDXGIAdapter, IID_IDXGIAdapter);
 
-  class nvenc_d3d11 final: public nvenc_base {
+  class nvenc_d3d11: public nvenc_base {
   public:
-    nvenc_d3d11(ID3D11Device *d3d_device);
-    ~nvenc_d3d11();
+    nvenc_d3d11(NV_ENC_DEVICE_TYPE device_type):
+        nvenc_base(device_type) {}
 
-    ID3D11Texture2D *
-    get_input_texture();
+    virtual ~nvenc_d3d11();
 
-  private:
-    bool
-    init_library() override;
+    virtual ID3D11Texture2D *
+    get_input_texture() = 0;
 
+  protected:
     bool
-    create_and_register_input_buffer() override;
+    init_library() override;
 
     HMODULE dll = NULL;
-    const ID3D11DevicePtr d3d_device;
-    ID3D11Texture2DPtr d3d_input_texture;
   };
 
 }  // namespace nvenc

diff --git a/src/nvenc/nvenc_d3d11_native.cpp b/src/nvenc/nvenc_d3d11_native.cpp
@@ -0,0 +1,67 @@
+#ifdef _WIN32
+  #include "nvenc_d3d11_native.h"
+
+  #include "nvenc_utils.h"
+
+namespace nvenc {
+
+  nvenc_d3d11_native::nvenc_d3d11_native(ID3D11Device *d3d_device):
+      nvenc_d3d11(NV_ENC_DEVICE_TYPE_DIRECTX),
+      d3d_device(d3d_device) {
+    device = d3d_device;
+  }
+
+  nvenc_d3d11_native::~nvenc_d3d11_native() {
+    if (encoder) destroy_encoder();
+  }
+
+  ID3D11Texture2D *
+  nvenc_d3d11_native::get_input_texture() {
+    return d3d_input_texture.GetInterfacePtr();
+  }
+
+  bool
+  nvenc_d3d11_native::create_and_register_input_buffer() {
+    if (encoder_params.buffer_format == NV_ENC_BUFFER_FORMAT_YUV444_10BIT) {
+      BOOST_LOG(error) << "NvEnc: 10-bit 4:4:4 encoding is incompatible with D3D11 surface formats, use CUDA interop";
+      return false;
+    }
+
+    if (!d3d_input_texture) {
+      D3D11_TEXTURE2D_DESC desc = {};
+      desc.Width = encoder_params.width;
+      desc.Height = encoder_params.height;
+      desc.MipLevels = 1;
+      desc.ArraySize = 1;
+      desc.Format = dxgi_format_from_nvenc_format(encoder_params.buffer_format);
+      desc.SampleDesc.Count = 1;
+      desc.Usage = D3D11_USAGE_DEFAULT;
+      desc.BindFlags = D3D11_BIND_RENDER_TARGET;
+      if (d3d_device->CreateTexture2D(&desc, nullptr, &d3d_input_texture) != S_OK) {
+        BOOST_LOG(error) << "NvEnc: couldn't create input texture";
+        return false;
+      }
+    }
+
+    if (!registered_input_buffer) {
+      NV_ENC_REGISTER_RESOURCE register_resource = { min_struct_version(NV_ENC_REGISTER_RESOURCE_VER, 3, 4) };
+      register_resource.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_DIRECTX;
+      register_resource.width = encoder_params.width;
+      register_resource.height = encoder_params.height;
+      register_resource.resourceToRegister = d3d_input_texture.GetInterfacePtr();
+      register_resource.bufferFormat = encoder_params.buffer_format;
+      register_resource.bufferUsage = NV_ENC_INPUT_IMAGE;
+
+      if (nvenc_failed(nvenc->nvEncRegisterResource(encoder, &register_resource))) {
+        BOOST_LOG(error) << "NvEncRegisterResource failed: " << last_error_string;
+        return false;
+      }
+
+      registered_input_buffer = register_resource.registeredResource;
+    }
+
+    return true;
+  }
+
+}  // namespace nvenc
+#endif
diff --git a/src/nvenc/nvenc_d3d11_native.h b/src/nvenc/nvenc_d3d11_native.h
@@ -0,0 +1,28 @@
+#pragma once
+#ifdef _WIN32
+
+  #include <comdef.h>
+  #include <d3d11.h>
+
+  #include "nvenc_d3d11.h"
+
+namespace nvenc {
+
+  class nvenc_d3d11_native final: public nvenc_d3d11 {
+  public:
+    nvenc_d3d11_native(ID3D11Device *d3d_device);
+    ~nvenc_d3d11_native();
+
+    ID3D11Texture2D *
+    get_input_texture() override;
+
+  private:
+    bool
+    create_and_register_input_buffer() override;
+
+    const ID3D11DevicePtr d3d_device;
+    ID3D11Texture2DPtr d3d_input_texture;
+  };
+
+}  // namespace nvenc
+#endif