avcodec/nvenc: add 4:2:2 encoding and H.264 10-bit support

This commit adds support for 4:2:2 encoding for HEVC and H.264 on
NVIDIA Blackwell GPUs. Additionally, it supports 10-bit encoding
for H.264 on Blackwell GPUs.

Signed-off-by: Diego de Souza <ddesouza@nvidia.com>
Signed-off-by: Timo Rothenpieler <timo@rothenpieler.org>

1 files changed, 62 insertions, 6 deletions

diff --git a/libavcodec/nvenc.c b/libavcodec/nvenc.c
index 06e3fb81a4..f301269dbd 100644
--- a/libavcodec/nvenc.c
+++ b/libavcodec/nvenc.c
@@ -59,6 +59,11 @@ const enum AVPixelFormat ff_nvenc_pix_fmts[] = {
     AV_PIX_FMT_P010,
     AV_PIX_FMT_YUV444P,
     AV_PIX_FMT_P016,      // Truncated to 10bits
+#ifdef NVENC_HAVE_422_SUPPORT
+    AV_PIX_FMT_NV16,
+    AV_PIX_FMT_P210,
+    AV_PIX_FMT_P216,
+#endif
     AV_PIX_FMT_YUV444P16, // Truncated to 10bits
     AV_PIX_FMT_0RGB32,
     AV_PIX_FMT_RGB32,
@@ -87,6 +92,8 @@ const AVCodecHWConfigInternal *const ff_nvenc_hw_configs[] = {
 
 #define IS_10BIT(pix_fmt)  (pix_fmt == AV_PIX_FMT_P010      || \
                             pix_fmt == AV_PIX_FMT_P016      || \
+                            pix_fmt == AV_PIX_FMT_P210      || \
+                            pix_fmt == AV_PIX_FMT_P216      || \
                             pix_fmt == AV_PIX_FMT_YUV444P16 || \
                             pix_fmt == AV_PIX_FMT_X2RGB10   || \
                             pix_fmt == AV_PIX_FMT_X2BGR10   || \
@@ -105,6 +112,10 @@ const AVCodecHWConfigInternal *const ff_nvenc_hw_configs[] = {
                             pix_fmt == AV_PIX_FMT_GBRP16    || \
                             (ctx->rgb_mode == NVENC_RGB_MODE_444 && IS_RGB(pix_fmt)))
 
+#define IS_YUV422(pix_fmt) (pix_fmt == AV_PIX_FMT_NV16 || \
+                            pix_fmt == AV_PIX_FMT_P210 || \
+                            pix_fmt == AV_PIX_FMT_P216)
+
 #define IS_GBRP(pix_fmt) (pix_fmt == AV_PIX_FMT_GBRP || \
                           pix_fmt == AV_PIX_FMT_GBRP16)
 
@@ -477,6 +488,16 @@ static int nvenc_check_capabilities(AVCodecContext *avctx)
         return AVERROR(ENOSYS);
     }
 
+#ifdef NVENC_HAVE_422_SUPPORT
+    ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_YUV422_ENCODE);
+#else
+    ret = 0;
+#endif
+    if (IS_YUV422(ctx->data_pix_fmt) && ret <= 0) {
+        av_log(avctx, AV_LOG_WARNING, "YUV422P not supported\n");
+        return AVERROR(ENOSYS);
+    }
+
     ret = nvenc_check_cap(avctx, NV_ENC_CAPS_SUPPORT_LOSSLESS_ENCODE);
     if (ctx->flags & NVENC_LOSSLESS && ret <= 0) {
         av_log(avctx, AV_LOG_WARNING, "Lossless encoding not supported\n");
@@ -1297,6 +1318,18 @@ static av_cold int nvenc_setup_h264_config(AVCodecContext *avctx)
             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_GUID;
             avctx->profile = AV_PROFILE_H264_HIGH;
             break;
+#ifdef NVENC_HAVE_H264_10BIT_SUPPORT
+        case NV_ENC_H264_PROFILE_HIGH_10:
+            cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_10_GUID;
+            avctx->profile = AV_PROFILE_H264_HIGH_10;
+            break;
+#endif
+#ifdef NVENC_HAVE_422_SUPPORT
+        case NV_ENC_H264_PROFILE_HIGH_422:
+            cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_422_GUID;
+            avctx->profile = AV_PROFILE_H264_HIGH_422;
+            break;
+#endif
         case NV_ENC_H264_PROFILE_HIGH_444P:
             cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
             avctx->profile = AV_PROFILE_H264_HIGH_444_PREDICTIVE;
@@ -1304,21 +1337,37 @@ static av_cold int nvenc_setup_h264_config(AVCodecContext *avctx)
         }
     }
 
+#ifdef NVENC_HAVE_H264_10BIT_SUPPORT
+    // force setting profile as high10 if input is 10 bit or if it should be encoded as 10 bit
+    if (IS_10BIT(ctx->data_pix_fmt) || ctx->highbitdepth) {
+        cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_10_GUID;
+        avctx->profile = AV_PROFILE_H264_HIGH_10;
+    }
+#endif
+
     // force setting profile as high444p if input is AV_PIX_FMT_YUV444P
     if (IS_YUV444(ctx->data_pix_fmt)) {
         cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_444_GUID;
         avctx->profile = AV_PROFILE_H264_HIGH_444_PREDICTIVE;
     }
 
+#ifdef NVENC_HAVE_422_SUPPORT
+    // force setting profile as high422p if input is AV_PIX_FMT_YUV422P
+    if (IS_YUV422(ctx->data_pix_fmt)) {
+        cc->profileGUID = NV_ENC_H264_PROFILE_HIGH_422_GUID;
+        avctx->profile = AV_PROFILE_H264_HIGH_422;
+    }
+#endif
+
     vui->bitstreamRestrictionFlag = cc->gopLength != 1 || avctx->profile < AV_PROFILE_H264_HIGH;
 
-    h264->chromaFormatIDC = avctx->profile == AV_PROFILE_H264_HIGH_444_PREDICTIVE ? 3 : 1;
+    h264->chromaFormatIDC = IS_YUV444(ctx->data_pix_fmt) ? 3 : IS_YUV422(ctx->data_pix_fmt) ? 2 : 1;
 
     h264->level = ctx->level;
 
 #ifdef NVENC_HAVE_NEW_BIT_DEPTH_API
-    h264->inputBitDepth = h264->outputBitDepth =
-        IS_10BIT(ctx->data_pix_fmt) ? NV_ENC_BIT_DEPTH_10 : NV_ENC_BIT_DEPTH_8;
+    h264->inputBitDepth = IS_10BIT(ctx->data_pix_fmt) ? NV_ENC_BIT_DEPTH_10 : NV_ENC_BIT_DEPTH_8;
+    h264->outputBitDepth = (IS_10BIT(ctx->data_pix_fmt) || ctx->highbitdepth) ? NV_ENC_BIT_DEPTH_10 : NV_ENC_BIT_DEPTH_8;
 #endif
 
     if (ctx->coder >= 0)
@@ -1428,13 +1477,13 @@ static av_cold int nvenc_setup_hevc_config(AVCodecContext *avctx)
         avctx->profile = AV_PROFILE_HEVC_MAIN_10;
     }
 
-    // force setting profile as rext if input is yuv444
-    if (IS_YUV444(ctx->data_pix_fmt)) {
+    // force setting profile as rext if input is yuv444 or yuv422
+    if (IS_YUV444(ctx->data_pix_fmt) || IS_YUV422(ctx->data_pix_fmt)) {
         cc->profileGUID = NV_ENC_HEVC_PROFILE_FREXT_GUID;
         avctx->profile = AV_PROFILE_HEVC_REXT;
     }
 
-    hevc->chromaFormatIDC = IS_YUV444(ctx->data_pix_fmt) ? 3 : 1;
+    hevc->chromaFormatIDC = IS_YUV444(ctx->data_pix_fmt) ? 3 : IS_YUV422(ctx->data_pix_fmt) ? 2 : 1;
 
 #ifdef NVENC_HAVE_NEW_BIT_DEPTH_API
     hevc->inputBitDepth = IS_10BIT(ctx->data_pix_fmt) ? NV_ENC_BIT_DEPTH_10 : NV_ENC_BIT_DEPTH_8;
@@ -1821,6 +1870,13 @@ static NV_ENC_BUFFER_FORMAT nvenc_map_buffer_format(enum AVPixelFormat pix_fmt)
         return NV_ENC_BUFFER_FORMAT_ARGB10;
     case AV_PIX_FMT_X2BGR10:
         return NV_ENC_BUFFER_FORMAT_ABGR10;
+#ifdef NVENC_HAVE_422_SUPPORT
+    case AV_PIX_FMT_NV16:
+        return NV_ENC_BUFFER_FORMAT_NV16;
+    case AV_PIX_FMT_P210:
+    case AV_PIX_FMT_P216:
+        return NV_ENC_BUFFER_FORMAT_P210;
+#endif
     default:
         return NV_ENC_BUFFER_FORMAT_UNDEFINED;
     }