subprojects/gst-plugins-bad/sys/nvcodec/gstnvencoder.cpp

   1 /* GStreamer
   2  * Copyright (C) 2022 Seungha Yang <seungha@centricular.com>
   3  *
   4  * This library is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Library General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2 of the License, or (at your option) any later version.
   8  *
   9  * This library is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Library General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Library General Public
  15  * License along with this library; if not, write to the
  16  * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
  17  * Boston, MA 02110-1301, USA.
  18  */
  19
  20 #ifdef HAVE_CONFIG_H
  21 #include "config.h"
  22 #endif
  23
  24 #include "gstnvencoder.h"
  25
  26 #include <gst/cuda/gstcudautils.h>
  27 #include <gst/cuda/gstcudamemory.h>
  28 #include <gst/cuda/gstcudabufferpool.h>
  29 #include <string.h>
  30
  31 #ifdef GST_CUDA_HAS_D3D
  32 #include <gst/d3d11/gstd3d11.h>
  33 #endif
  34
  35 #ifdef G_OS_WIN32
  36 #include <wrl.h>
  37
  38 /* *INDENT-OFF* */
  39 using namespace Microsoft::WRL;
  40 /* *INDENT-ON* */
  41 #endif
  42
  43 GST_DEBUG_CATEGORY_STATIC (gst_nv_encoder_debug);
  44 #define GST_CAT_DEFAULT gst_nv_encoder_debug
  45
  46 #define GET_LOCK(e) (&(GST_NV_ENCODER_CAST(e)->priv->lock))
  47 #define GST_NV_ENCODER_LOCK(e) G_STMT_START { \
  48   GST_TRACE_OBJECT (e, "Locking from thread %p", g_thread_self ()); \
  49   g_mutex_lock(GET_LOCK(e)); \
  50   GST_TRACE_OBJECT (e, "Locked from thread %p", g_thread_self ()); \
  51 } G_STMT_END
  52
  53 #define GST_NV_ENCODER_UNLOCK(e) G_STMT_START { \
  54   GST_TRACE_OBJECT (e, "Unlocking from thread %p", g_thread_self ()); \
  55   g_mutex_unlock(GET_LOCK(e)); \
  56 } G_STMT_END
  57
  58 struct _GstNvEncoderPrivate
  59 {
  60   GstCudaContext *context;
  61 #ifdef GST_CUDA_HAS_D3D
  62   GstD3D11Device *device;
  63   GstD3D11Fence *fence;
  64 #endif
  65
  66   GstNvEncoderDeviceMode subclass_device_mode;
  67   GstNvEncoderDeviceMode selected_device_mode;
  68   gint64 dxgi_adapter_luid;
  69   guint cuda_device_id;
  70
  71   NV_ENC_INITIALIZE_PARAMS init_params;
  72   NV_ENC_CONFIG config;
  73   gpointer session;
  74
  75   GstVideoCodecState *input_state;
  76
  77   GstBufferPool *internal_pool;
  78
  79   GstClockTime dts_offset;
  80
  81   /* Array of GstNvEncoderTask, holding ownership */
  82   GArray *task_pool;
  83
  84   GQueue free_tasks;
  85   GQueue output_tasks;
  86
  87   GMutex lock;
  88   GCond cond;
  89
  90   GRecMutex context_lock;
  91
  92   GThread *encoding_thread;
  93
  94   GstFlowReturn last_flow;
  95 };
  96
  97 /**
  98  * GstNvEncoder:
  99  *
 100  * Since: 1.22
 101  */
 102 #define gst_nv_encoder_parent_class parent_class
 103 G_DEFINE_ABSTRACT_TYPE_WITH_PRIVATE (GstNvEncoder, gst_nv_encoder,
 104     GST_TYPE_VIDEO_ENCODER);
 105
 106 static void gst_nv_encoder_finalize (GObject * object);
 107 static void gst_nv_encoder_set_context (GstElement * element,
 108     GstContext * context);
 109 static gboolean gst_nv_encoder_open (GstVideoEncoder * encoder);
 110 static gboolean gst_nv_encoder_close (GstVideoEncoder * encoder);
 111 static gboolean gst_nv_encoder_stop (GstVideoEncoder * encoder);
 112 static gboolean gst_nv_encoder_sink_query (GstVideoEncoder * encoder,
 113     GstQuery * query);
 114 static gboolean gst_nv_encoder_src_query (GstVideoEncoder * encoder,
 115     GstQuery * query);
 116 static gboolean gst_nv_encoder_propose_allocation (GstVideoEncoder *
 117     encoder, GstQuery * query);
 118 static gboolean gst_nv_encoder_set_format (GstVideoEncoder * encoder,
 119     GstVideoCodecState * state);
 120 static GstFlowReturn gst_nv_encoder_handle_frame (GstVideoEncoder *
 121     encoder, GstVideoCodecFrame * frame);
 122 static GstFlowReturn gst_nv_encoder_finish (GstVideoEncoder * encoder);
 123 static gboolean gst_nv_encoder_flush (GstVideoEncoder * encoder);
 124 static void gst_nv_encoder_task_clear (GstNvEncoderTask * task);
 125
 126 static void
 127 gst_nv_encoder_class_init (GstNvEncoderClass * klass)
 128 {
 129   GObjectClass *object_class = G_OBJECT_CLASS (klass);
 130   GstElementClass *element_class = GST_ELEMENT_CLASS (klass);
 131   GstVideoEncoderClass *videoenc_class = GST_VIDEO_ENCODER_CLASS (klass);
 132
 133   object_class->finalize = gst_nv_encoder_finalize;
 134
 135   element_class->set_context = GST_DEBUG_FUNCPTR (gst_nv_encoder_set_context);
 136
 137   videoenc_class->open = GST_DEBUG_FUNCPTR (gst_nv_encoder_open);
 138   videoenc_class->close = GST_DEBUG_FUNCPTR (gst_nv_encoder_close);
 139   videoenc_class->stop = GST_DEBUG_FUNCPTR (gst_nv_encoder_stop);
 140   videoenc_class->sink_query = GST_DEBUG_FUNCPTR (gst_nv_encoder_sink_query);
 141   videoenc_class->src_query = GST_DEBUG_FUNCPTR (gst_nv_encoder_src_query);
 142   videoenc_class->propose_allocation =
 143       GST_DEBUG_FUNCPTR (gst_nv_encoder_propose_allocation);
 144   videoenc_class->set_format = GST_DEBUG_FUNCPTR (gst_nv_encoder_set_format);
 145   videoenc_class->handle_frame =
 146       GST_DEBUG_FUNCPTR (gst_nv_encoder_handle_frame);
 147   videoenc_class->finish = GST_DEBUG_FUNCPTR (gst_nv_encoder_finish);
 148   videoenc_class->flush = GST_DEBUG_FUNCPTR (gst_nv_encoder_flush);
 149
 150   GST_DEBUG_CATEGORY_INIT (gst_nv_encoder_debug, "nvencoder", 0, "nvencoder");
 151
 152   gst_type_mark_as_plugin_api (GST_TYPE_NV_ENCODER, (GstPluginAPIFlags) 0);
 153   gst_type_mark_as_plugin_api (GST_TYPE_NV_ENCODER_PRESET,
 154       (GstPluginAPIFlags) 0);
 155   gst_type_mark_as_plugin_api (GST_TYPE_NV_ENCODER_RC_MODE,
 156       (GstPluginAPIFlags) 0);
 157 }
 158
 159 static void
 160 gst_nv_encoder_init (GstNvEncoder * self)
 161 {
 162   GstNvEncoderPrivate *priv;
 163
 164   self->priv = priv = (GstNvEncoderPrivate *)
 165       gst_nv_encoder_get_instance_private (self);
 166
 167   priv->task_pool = g_array_new (FALSE, TRUE, sizeof (GstNvEncoderTask));
 168   g_array_set_clear_func (priv->task_pool,
 169       (GDestroyNotify) gst_nv_encoder_task_clear);
 170
 171   g_queue_init (&priv->free_tasks);
 172   g_queue_init (&priv->output_tasks);
 173
 174   g_mutex_init (&priv->lock);
 175   g_cond_init (&priv->cond);
 176
 177   g_rec_mutex_init (&priv->context_lock);
 178
 179   gst_video_encoder_set_min_pts (GST_VIDEO_ENCODER (self),
 180       GST_SECOND * 60 * 60 * 1000);
 181 }
 182
 183 static void
 184 gst_nv_encoder_finalize (GObject * object)
 185 {
 186   GstNvEncoder *self = GST_NV_ENCODER (object);
 187   GstNvEncoderPrivate *priv = self->priv;
 188
 189   g_array_unref (priv->task_pool);
 190
 191   g_mutex_clear (&priv->lock);
 192   g_cond_clear (&priv->cond);
 193
 194   g_rec_mutex_clear (&priv->context_lock);
 195
 196   G_OBJECT_CLASS (parent_class)->finalize (object);
 197 }
 198
 199 static void
 200 gst_nv_encoder_set_context (GstElement * element, GstContext * context)
 201 {
 202   GstNvEncoder *self = GST_NV_ENCODER (element);
 203   GstNvEncoderPrivate *priv = self->priv;
 204
 205   g_rec_mutex_lock (&priv->context_lock);
 206
 207   switch (priv->selected_device_mode) {
 208 #ifdef GST_CUDA_HAS_D3D
 209     case GST_NV_ENCODER_DEVICE_D3D11:
 210       gst_d3d11_handle_set_context_for_adapter_luid (element,
 211           context, priv->dxgi_adapter_luid, &priv->device);
 212       break;
 213 #endif
 214     case GST_NV_ENCODER_DEVICE_CUDA:
 215       gst_cuda_handle_set_context (element, context, priv->cuda_device_id,
 216           &priv->context);
 217       break;
 218     default:
 219       break;
 220   }
 221
 222   g_rec_mutex_unlock (&priv->context_lock);
 223
 224   GST_ELEMENT_CLASS (parent_class)->set_context (element, context);
 225 }
 226
 227 static gboolean
 228 gst_nv_encoder_reset (GstNvEncoder * self)
 229 {
 230   GstNvEncoderPrivate *priv = self->priv;
 231
 232   GST_LOG_OBJECT (self, "Reset");
 233
 234   g_array_set_size (priv->task_pool, 0);
 235
 236   if (priv->internal_pool) {
 237     gst_buffer_pool_set_active (priv->internal_pool, FALSE);
 238     gst_clear_object (&priv->internal_pool);
 239   }
 240
 241   if (priv->session) {
 242     NvEncDestroyEncoder (priv->session);
 243     priv->session = NULL;
 244   }
 245
 246   g_queue_clear (&priv->free_tasks);
 247   g_queue_clear (&priv->output_tasks);
 248
 249   priv->last_flow = GST_FLOW_OK;
 250
 251   return TRUE;
 252 }
 253
 254 static gboolean
 255 gst_nv_encoder_device_lock (GstNvEncoder * self)
 256 {
 257   GstNvEncoderPrivate *priv = self->priv;
 258   gboolean ret = TRUE;
 259
 260   switch (priv->selected_device_mode) {
 261 #ifdef GST_CUDA_HAS_D3D
 262     case GST_NV_ENCODER_DEVICE_D3D11:
 263       gst_d3d11_device_lock (priv->device);
 264       break;
 265 #endif
 266     case GST_NV_ENCODER_DEVICE_CUDA:
 267       ret = gst_cuda_context_push (priv->context);
 268       break;
 269     default:
 270       break;
 271   }
 272
 273   return ret;
 274 }
 275
 276 static gboolean
 277 gst_nv_encoder_device_unlock (GstNvEncoder * self)
 278 {
 279   GstNvEncoderPrivate *priv = self->priv;
 280   gboolean ret = TRUE;
 281
 282   switch (priv->selected_device_mode) {
 283 #ifdef GST_CUDA_HAS_D3D
 284     case GST_NV_ENCODER_DEVICE_D3D11:
 285       gst_d3d11_device_unlock (priv->device);
 286       break;
 287 #endif
 288     case GST_NV_ENCODER_DEVICE_CUDA:
 289       ret = gst_cuda_context_pop (nullptr);
 290       break;
 291     default:
 292       break;
 293   }
 294
 295   return ret;
 296 }
 297
 298 static GstFlowReturn
 299 gst_nv_encoder_get_free_task (GstNvEncoder * self, GstNvEncoderTask ** task,
 300     gboolean check_last_flow)
 301 {
 302   GstNvEncoderPrivate *priv = self->priv;
 303   GstFlowReturn ret = GST_FLOW_OK;
 304   GstNvEncoderTask *free_task = NULL;
 305
 306   GST_NV_ENCODER_LOCK (self);
 307   if (check_last_flow) {
 308     if (priv->last_flow != GST_FLOW_OK) {
 309       ret = priv->last_flow;
 310       GST_NV_ENCODER_UNLOCK (self);
 311       return ret;
 312     }
 313
 314     while (priv->last_flow == GST_FLOW_OK && (free_task = (GstNvEncoderTask *)
 315             g_queue_pop_head (&priv->free_tasks)) == NULL) {
 316       g_cond_wait (&priv->cond, &priv->lock);
 317     }
 318
 319     ret = priv->last_flow;
 320     if (ret != GST_FLOW_OK && free_task) {
 321       g_queue_push_tail (&priv->free_tasks, free_task);
 322       free_task = NULL;
 323     }
 324   } else {
 325     while ((free_task = (GstNvEncoderTask *)
 326             g_queue_pop_head (&priv->free_tasks)) == NULL)
 327       g_cond_wait (&priv->cond, &priv->lock);
 328   }
 329   GST_NV_ENCODER_UNLOCK (self);
 330
 331   *task = free_task;
 332
 333   return ret;
 334 }
 335
 336 static gboolean
 337 gst_nv_encoder_drain (GstNvEncoder * self, gboolean locked)
 338 {
 339   GstNvEncoderPrivate *priv = self->priv;
 340   NV_ENC_PIC_PARAMS pic_params = { 0, };
 341   NVENCSTATUS status;
 342   GstNvEncoderTask *task;
 343
 344   if (!priv->session || !priv->encoding_thread)
 345     return TRUE;
 346
 347   GST_DEBUG_OBJECT (self, "Drain");
 348
 349   if (locked)
 350     GST_VIDEO_ENCODER_STREAM_UNLOCK (self);
 351
 352   gst_nv_encoder_get_free_task (self, &task, FALSE);
 353
 354   task->is_eos = TRUE;
 355
 356   pic_params.version = gst_nvenc_get_pic_params_version ();
 357   pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
 358   pic_params.completionEvent = task->event_handle;
 359
 360   gst_nv_encoder_device_lock (self);
 361   status = NvEncEncodePicture (priv->session, &pic_params);
 362   if (status != NV_ENC_SUCCESS) {
 363     GST_DEBUG_OBJECT (self, "Drain returned status %" GST_NVENC_STATUS_FORMAT,
 364         GST_NVENC_STATUS_ARGS (status));
 365 #ifdef G_OS_WIN32
 366     if (task->event_handle) {
 367       SetEvent (task->event_handle);
 368     }
 369 #endif
 370   }
 371   gst_nv_encoder_device_unlock (self);
 372
 373   GST_NV_ENCODER_LOCK (self);
 374   g_queue_push_tail (&priv->output_tasks, task);
 375   g_cond_broadcast (&priv->cond);
 376   GST_NV_ENCODER_UNLOCK (self);
 377
 378   g_clear_pointer (&priv->encoding_thread, g_thread_join);
 379   gst_nv_encoder_reset (self);
 380
 381   if (locked)
 382     GST_VIDEO_ENCODER_STREAM_LOCK (self);
 383
 384   return TRUE;
 385 }
 386
 387 #ifdef GST_CUDA_HAS_D3D
 388 static gboolean
 389 gst_nv_encoder_open_d3d11_device (GstNvEncoder * self)
 390 {
 391   GstNvEncoderPrivate *priv = self->priv;
 392   ComPtr < ID3D10Multithread > multi_thread;
 393   ID3D11Device *device_handle;
 394   HRESULT hr;
 395
 396   if (!gst_d3d11_ensure_element_data_for_adapter_luid (GST_ELEMENT (self),
 397           priv->dxgi_adapter_luid, &priv->device)) {
 398     GST_ERROR_OBJECT (self, "Cannot create d3d11device");
 399     return FALSE;
 400   }
 401
 402   device_handle = gst_d3d11_device_get_device_handle (priv->device);
 403   hr = device_handle->QueryInterface (IID_PPV_ARGS (&multi_thread));
 404   if (!gst_d3d11_result (hr, priv->device)) {
 405     GST_ERROR_OBJECT (self, "ID3D10Multithread interface is unavailable");
 406     gst_clear_object (&priv->device);
 407
 408     return FALSE;
 409   }
 410
 411   multi_thread->SetMultithreadProtected (TRUE);
 412
 413   return TRUE;
 414 }
 415 #endif
 416
 417 static gboolean
 418 gst_nv_encoder_open (GstVideoEncoder * encoder)
 419 {
 420   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 421   GstNvEncoderPrivate *priv = self->priv;
 422
 423   switch (priv->selected_device_mode) {
 424     case GST_NV_ENCODER_DEVICE_AUTO_SELECT:
 425       /* Will open GPU later */
 426       return TRUE;
 427 #ifdef GST_CUDA_HAS_D3D
 428     case GST_NV_ENCODER_DEVICE_D3D11:
 429       return gst_nv_encoder_open_d3d11_device (self);
 430 #endif
 431     case GST_NV_ENCODER_DEVICE_CUDA:
 432       if (!gst_cuda_ensure_element_context (GST_ELEMENT_CAST (encoder),
 433               priv->cuda_device_id, &priv->context)) {
 434         GST_ERROR_OBJECT (self, "failed to create CUDA context");
 435         return FALSE;
 436       }
 437       break;
 438     default:
 439       g_assert_not_reached ();
 440       return FALSE;
 441   }
 442
 443   return TRUE;
 444 }
 445
 446 static gboolean
 447 gst_nv_encoder_close (GstVideoEncoder * encoder)
 448 {
 449   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 450   GstNvEncoderPrivate *priv = self->priv;
 451
 452   gst_clear_object (&priv->context);
 453 #ifdef GST_CUDA_HAS_D3D
 454   gst_clear_d3d11_fence (&priv->fence);
 455   gst_clear_object (&priv->device);
 456 #endif
 457
 458   return TRUE;
 459 }
 460
 461 static gboolean
 462 gst_nv_encoder_stop (GstVideoEncoder * encoder)
 463 {
 464   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 465   GstNvEncoderPrivate *priv = self->priv;
 466
 467   GST_DEBUG_OBJECT (self, "Stop");
 468
 469   gst_nv_encoder_drain (self, FALSE);
 470
 471   if (priv->subclass_device_mode == GST_NV_ENCODER_DEVICE_AUTO_SELECT) {
 472     gst_clear_object (&priv->context);
 473 #ifdef GST_CUDA_HAS_D3D
 474     gst_clear_object (&priv->device);
 475 #endif
 476     priv->selected_device_mode = GST_NV_ENCODER_DEVICE_AUTO_SELECT;
 477   }
 478
 479   g_clear_pointer (&priv->input_state, gst_video_codec_state_unref);
 480
 481   return TRUE;
 482 }
 483
 484 static gboolean
 485 gst_nv_encoder_handle_context_query (GstNvEncoder * self, GstQuery * query)
 486 {
 487   GstNvEncoderPrivate *priv = self->priv;
 488   gboolean ret = FALSE;
 489
 490   g_rec_mutex_lock (&priv->context_lock);
 491
 492   switch (priv->selected_device_mode) {
 493 #ifdef GST_CUDA_HAS_D3D
 494     case GST_NV_ENCODER_DEVICE_D3D11:
 495       ret = gst_d3d11_handle_context_query (GST_ELEMENT (self),
 496           query, priv->device);
 497       break;
 498 #endif
 499     case GST_NV_ENCODER_DEVICE_CUDA:
 500       ret = gst_cuda_handle_context_query (GST_ELEMENT (self),
 501           query, priv->context);
 502       break;
 503     default:
 504       break;
 505   }
 506
 507   g_rec_mutex_unlock (&priv->context_lock);
 508
 509   return ret;
 510 }
 511
 512 static gboolean
 513 gst_nv_encoder_sink_query (GstVideoEncoder * encoder, GstQuery * query)
 514 {
 515   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 516
 517   switch (GST_QUERY_TYPE (query)) {
 518     case GST_QUERY_CONTEXT:
 519       if (gst_nv_encoder_handle_context_query (self, query))
 520         return TRUE;
 521       break;
 522     default:
 523       break;
 524   }
 525
 526   return GST_VIDEO_ENCODER_CLASS (parent_class)->sink_query (encoder, query);
 527 }
 528
 529 static gboolean
 530 gst_nv_encoder_src_query (GstVideoEncoder * encoder, GstQuery * query)
 531 {
 532   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 533
 534   switch (GST_QUERY_TYPE (query)) {
 535     case GST_QUERY_CONTEXT:
 536       if (gst_nv_encoder_handle_context_query (self, query))
 537         return TRUE;
 538       break;
 539     default:
 540       break;
 541   }
 542
 543   return GST_VIDEO_ENCODER_CLASS (parent_class)->src_query (encoder, query);
 544 }
 545
 546 static gboolean
 547 gst_nv_encoder_propose_allocation (GstVideoEncoder * encoder, GstQuery * query)
 548 {
 549   GstNvEncoder *self = GST_NV_ENCODER (encoder);
 550   GstNvEncoderPrivate *priv = self->priv;
 551   GstVideoInfo info;
 552   GstBufferPool *pool = NULL;
 553   GstCaps *caps;
 554   guint size;
 555   GstStructure *config;
 556   GstCapsFeatures *features;
 557   guint min_buffers;
 558
 559   gst_query_parse_allocation (query, &caps, NULL);
 560   if (!caps) {
 561     GST_WARNING_OBJECT (self, "null caps in query");
 562     return FALSE;
 563   }
 564
 565   if (!gst_video_info_from_caps (&info, caps)) {
 566     GST_WARNING_OBJECT (self, "Failed to convert caps into info");
 567     return FALSE;
 568   }
 569
 570   features = gst_caps_get_features (caps, 0);
 571   min_buffers = gst_nv_encoder_get_task_size (self);
 572
 573   switch (priv->subclass_device_mode) {
 574     case GST_NV_ENCODER_DEVICE_AUTO_SELECT:
 575       /* Use upstream pool in case of auto select mode. We don't know which
 576        * GPU to use at this moment */
 577       gst_query_add_allocation_meta (query, GST_VIDEO_META_API_TYPE, nullptr);
 578       gst_query_add_allocation_pool (query, nullptr, info.size, min_buffers, 0);
 579       return TRUE;
 580 #ifdef GST_CUDA_HAS_D3D
 581     case GST_NV_ENCODER_DEVICE_D3D11:
 582       if (features && gst_caps_features_contains (features,
 583               GST_CAPS_FEATURE_MEMORY_D3D11_MEMORY)) {
 584         GST_DEBUG_OBJECT (self, "upstream support d3d11 memory");
 585         pool = gst_d3d11_buffer_pool_new (priv->device);
 586       }
 587       break;
 588 #endif
 589     case GST_NV_ENCODER_DEVICE_CUDA:
 590       if (features && gst_caps_features_contains (features,
 591               GST_CAPS_FEATURE_MEMORY_CUDA_MEMORY)) {
 592         GST_DEBUG_OBJECT (self, "upstream support CUDA memory");
 593         pool = gst_cuda_buffer_pool_new (priv->context);
 594       }
 595       break;
 596     default:
 597       g_assert_not_reached ();
 598       return FALSE;
 599   }
 600
 601   if (!pool)
 602     pool = gst_video_buffer_pool_new ();
 603
 604   config = gst_buffer_pool_get_config (pool);
 605   gst_buffer_pool_config_add_option (config, GST_BUFFER_POOL_OPTION_VIDEO_META);
 606
 607
 608   size = GST_VIDEO_INFO_SIZE (&info);
 609   gst_buffer_pool_config_set_params (config, caps, size, min_buffers, 0);
 610
 611   if (!gst_buffer_pool_set_config (pool, config)) {
 612     GST_WARNING_OBJECT (self, "Failed to set pool config");
 613     gst_object_unref (pool);
 614     return FALSE;
 615   }
 616
 617   config = gst_buffer_pool_get_config (pool);
 618   gst_buffer_pool_config_get_params (config, NULL, &size, NULL, NULL);
 619   gst_structure_free (config);
 620
 621   gst_query_add_allocation_pool (query, pool, size, min_buffers, 0);
 622   gst_query_add_allocation_meta (query, GST_VIDEO_META_API_TYPE, NULL);
 623   gst_object_unref (pool);
 624
 625   return TRUE;
 626 }
 627
 628 /* called with lock */
 629 static void
 630 gst_nv_encoder_task_reset (GstNvEncoder * self, GstNvEncoderTask * task)
 631 {
 632   GstNvEncoderPrivate *priv = self->priv;
 633
 634   if (!task)
 635     return;
 636
 637   if (task->buffer) {
 638     gst_nv_encoder_device_lock (self);
 639     if (priv->session) {
 640       NvEncUnmapInputResource (priv->session,
 641           task->mapped_resource.mappedResource);
 642       NvEncUnregisterResource (priv->session,
 643           task->register_resource.registeredResource);
 644     }
 645     gst_nv_encoder_device_unlock (self);
 646
 647     gst_buffer_unmap (task->buffer, &task->map_info);
 648     gst_clear_buffer (&task->buffer);
 649   }
 650 #ifdef G_OS_WIN32
 651   if (task->event_handle)
 652     ResetEvent (task->event_handle);
 653 #endif
 654
 655   task->is_eos = FALSE;
 656
 657   g_queue_push_head (&priv->free_tasks, task);
 658 }
 659
 660 static gboolean
 661 gst_nv_encoder_create_event_handle (GstNvEncoder * self, gpointer session,
 662     gpointer * event_handle)
 663 {
 664 #ifdef G_OS_WIN32
 665   NV_ENC_EVENT_PARAMS event_params = { 0, };
 666   NVENCSTATUS status;
 667
 668   event_params.version = gst_nvenc_get_event_params_version ();
 669   event_params.completionEvent = CreateEvent (NULL, FALSE, FALSE, NULL);
 670   status = NvEncRegisterAsyncEvent (session, &event_params);
 671
 672   if (status != NV_ENC_SUCCESS) {
 673     GST_ERROR_OBJECT (self,
 674         "Failed to register async event handle, status %"
 675         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
 676     CloseHandle (event_params.completionEvent);
 677     return FALSE;
 678   }
 679
 680   *event_handle = event_params.completionEvent;
 681 #endif
 682
 683   return TRUE;
 684 }
 685
 686 static gboolean
 687 gst_d3d11_encoder_wait_for_event_handle (GstNvEncoder * self,
 688     gpointer event_handle)
 689 {
 690 #ifdef G_OS_WIN32
 691   /* NVCODEC SDK uses 20s */
 692   if (WaitForSingleObject (event_handle, 20000) == WAIT_FAILED) {
 693     GST_ERROR_OBJECT (self, "Failed to wait for completion event");
 694     return FALSE;
 695   }
 696 #endif
 697
 698   return TRUE;
 699 }
 700
 701 static void
 702 gst_nv_encoder_destroy_event_handle (GstNvEncoder * self, gpointer session,
 703     gpointer event_handle)
 704 {
 705 #ifdef G_OS_WIN32
 706   NV_ENC_EVENT_PARAMS event_params = { 0, };
 707   NVENCSTATUS status;
 708
 709   event_params.version = gst_nvenc_get_event_params_version ();
 710   event_params.completionEvent = event_handle;
 711   status = NvEncUnregisterAsyncEvent (session, &event_params);
 712   CloseHandle (event_handle);
 713
 714   if (status != NV_ENC_SUCCESS) {
 715     GST_ERROR_OBJECT (self,
 716         "Failed to unregister async event handle, status %"
 717         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
 718   }
 719 #endif
 720 }
 721
 722 static void
 723 gst_nv_encoder_task_clear (GstNvEncoderTask * task)
 724 {
 725   GstNvEncoder *self;
 726   GstNvEncoderPrivate *priv;
 727
 728   if (!task)
 729     return;
 730
 731   self = task->encoder;
 732   priv = self->priv;
 733
 734   if (priv->session) {
 735     gst_nv_encoder_device_lock (self);
 736     if (task->buffer) {
 737       NvEncUnmapInputResource (priv->session,
 738           task->mapped_resource.mappedResource);
 739       NvEncUnregisterResource (priv->session,
 740           task->register_resource.registeredResource);
 741     }
 742     if (task->output_ptr)
 743       NvEncDestroyBitstreamBuffer (priv->session, task->output_ptr);
 744     if (task->input_buffer.inputBuffer)
 745       NvEncDestroyInputBuffer (priv->session, task->input_buffer.inputBuffer);
 746     if (task->event_handle) {
 747       gst_nv_encoder_destroy_event_handle (self, priv->session,
 748           task->event_handle);
 749     }
 750
 751     gst_nv_encoder_device_unlock (self);
 752   }
 753
 754   if (task->buffer) {
 755     gst_buffer_unmap (task->buffer, &task->map_info);
 756     gst_clear_buffer (&task->buffer);
 757   }
 758
 759   memset (task, 0, sizeof (GstNvEncoderTask));
 760 }
 761
 762 static NV_ENC_PIC_STRUCT
 763 gst_nv_encoder_get_pic_struct (GstNvEncoder * self, GstBuffer * buffer)
 764 {
 765   GstNvEncoderPrivate *priv = self->priv;
 766   GstVideoInfo *info = &priv->input_state->info;
 767
 768   if (!GST_VIDEO_INFO_IS_INTERLACED (info))
 769     return NV_ENC_PIC_STRUCT_FRAME;
 770
 771   if (GST_VIDEO_INFO_INTERLACE_MODE (info) == GST_VIDEO_INTERLACE_MODE_MIXED) {
 772     if (!GST_BUFFER_FLAG_IS_SET (buffer, GST_VIDEO_BUFFER_FLAG_INTERLACED)) {
 773       return NV_ENC_PIC_STRUCT_FRAME;
 774     }
 775
 776     if (GST_BUFFER_FLAG_IS_SET (buffer, GST_VIDEO_BUFFER_FLAG_TFF))
 777       return NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
 778
 779     return NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
 780   }
 781
 782   switch (GST_VIDEO_INFO_FIELD_ORDER (info)) {
 783     case GST_VIDEO_FIELD_ORDER_TOP_FIELD_FIRST:
 784       return NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
 785       break;
 786     case GST_VIDEO_FIELD_ORDER_BOTTOM_FIELD_FIRST:
 787       return NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
 788       break;
 789     default:
 790       break;
 791   }
 792
 793   if (GST_BUFFER_FLAG_IS_SET (buffer, GST_VIDEO_BUFFER_FLAG_TFF))
 794     return NV_ENC_PIC_STRUCT_FIELD_TOP_BOTTOM;
 795
 796   return NV_ENC_PIC_STRUCT_FIELD_BOTTOM_TOP;
 797 }
 798
 799 static GstFlowReturn
 800 gst_nv_encoder_encode_frame (GstNvEncoder * self,
 801     GstVideoCodecFrame * frame, GstNvEncoderTask * task)
 802 {
 803   GstNvEncoderPrivate *priv = self->priv;
 804   NV_ENC_PIC_PARAMS pic_params = { 0, };
 805   NVENCSTATUS status;
 806   guint retry_count = 0;
 807   const guint retry_threshold = 100;
 808
 809   pic_params.version = gst_nvenc_get_pic_params_version ();
 810   if (task->buffer) {
 811     pic_params.inputWidth = task->register_resource.width;
 812     pic_params.inputHeight = task->register_resource.height;
 813     pic_params.inputPitch = task->register_resource.pitch;
 814     pic_params.inputBuffer = task->mapped_resource.mappedResource;
 815     pic_params.bufferFmt = task->mapped_resource.mappedBufferFmt;
 816   } else {
 817     pic_params.inputWidth = task->input_buffer.width;
 818     pic_params.inputHeight = task->input_buffer.height;
 819     pic_params.inputPitch = task->lk_input_buffer.pitch;
 820     pic_params.inputBuffer = task->input_buffer.inputBuffer;
 821     pic_params.bufferFmt = task->input_buffer.bufferFmt;
 822   }
 823
 824   pic_params.frameIdx = frame->system_frame_number;
 825   pic_params.inputTimeStamp = frame->pts;
 826   pic_params.inputDuration = frame->duration;
 827   pic_params.outputBitstream = task->output_ptr;
 828   pic_params.completionEvent = task->event_handle;
 829   pic_params.pictureStruct = gst_nv_encoder_get_pic_struct (self, task->buffer);
 830   if (GST_VIDEO_CODEC_FRAME_IS_FORCE_KEYFRAME (frame))
 831     pic_params.encodePicFlags = NV_ENC_PIC_FLAG_FORCEIDR;
 832
 833   do {
 834     gst_nv_encoder_device_lock (self);
 835     status = NvEncEncodePicture (priv->session, &pic_params);
 836     gst_nv_encoder_device_unlock (self);
 837
 838     if (status == NV_ENC_ERR_ENCODER_BUSY) {
 839       if (retry_count < 100) {
 840         GST_DEBUG_OBJECT (self, "GPU is busy, retry count (%d/%d)", retry_count,
 841             retry_threshold);
 842         retry_count++;
 843
 844         /* Magic number 1ms */
 845         g_usleep (1000);
 846         continue;
 847       } else {
 848         GST_ERROR_OBJECT (self, "GPU is keep busy, give up");
 849         break;
 850       }
 851     }
 852
 853     break;
 854   } while (TRUE);
 855
 856   GST_NV_ENCODER_LOCK (self);
 857   if (status != NV_ENC_SUCCESS && status != NV_ENC_ERR_NEED_MORE_INPUT) {
 858     GST_ERROR_OBJECT (self, "Encode return %" GST_NVENC_STATUS_FORMAT,
 859         GST_NVENC_STATUS_ARGS (status));
 860     gst_nv_encoder_task_reset (self, task);
 861     GST_NV_ENCODER_UNLOCK (self);
 862
 863     return GST_FLOW_ERROR;
 864   }
 865
 866   gst_video_codec_frame_set_user_data (frame, task, NULL);
 867   g_queue_push_tail (&priv->output_tasks, task);
 868   g_cond_broadcast (&priv->cond);
 869   GST_NV_ENCODER_UNLOCK (self);
 870
 871   return GST_FLOW_OK;
 872 }
 873
 874 static GstVideoCodecFrame *
 875 gst_nv_encoder_find_output_frame (GstVideoEncoder * self,
 876     GstNvEncoderTask * task)
 877 {
 878   GList *frames, *iter;
 879   GstVideoCodecFrame *ret = NULL;
 880
 881   frames = gst_video_encoder_get_frames (self);
 882
 883   for (iter = frames; iter; iter = g_list_next (iter)) {
 884     GstVideoCodecFrame *frame = (GstVideoCodecFrame *) iter->data;
 885     GstNvEncoderTask *other = (GstNvEncoderTask *)
 886         gst_video_codec_frame_get_user_data (frame);
 887
 888     if (!other)
 889       continue;
 890
 891     if (other == task) {
 892       ret = frame;
 893       break;
 894     }
 895   }
 896
 897   if (ret)
 898     gst_video_codec_frame_ref (ret);
 899
 900   if (frames)
 901     g_list_free_full (frames, (GDestroyNotify) gst_video_codec_frame_unref);
 902
 903   return ret;
 904 }
 905
 906 static gpointer
 907 gst_nv_encoder_thread_func (GstNvEncoder * self)
 908 {
 909   GstVideoEncoder *encoder = GST_VIDEO_ENCODER (self);
 910   GstNvEncoderClass *klass = GST_NV_ENCODER_GET_CLASS (self);
 911   GstNvEncoderPrivate *priv = self->priv;
 912   GstNvEncoderTask *task = NULL;
 913
 914   do {
 915     NV_ENC_LOCK_BITSTREAM bitstream = { 0, };
 916     NVENCSTATUS status;
 917     GstVideoCodecFrame *frame;
 918     GstFlowReturn ret;
 919
 920     GST_NV_ENCODER_LOCK (self);
 921     while ((task = (GstNvEncoderTask *)
 922             g_queue_pop_head (&priv->output_tasks)) == NULL) {
 923       g_cond_wait (&priv->cond, &priv->lock);
 924     }
 925     GST_NV_ENCODER_UNLOCK (self);
 926
 927     if (task->event_handle) {
 928       if (!gst_d3d11_encoder_wait_for_event_handle (self, task->event_handle)) {
 929         GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
 930             ("Failed to wait for event signal"));
 931         goto error;
 932       }
 933     }
 934
 935     if (task->is_eos) {
 936       GST_INFO_OBJECT (self, "Got EOS packet");
 937
 938       GST_NV_ENCODER_LOCK (self);
 939       gst_nv_encoder_task_reset (self, task);
 940       g_cond_broadcast (&priv->cond);
 941       GST_NV_ENCODER_UNLOCK (self);
 942
 943       goto exit_thread;
 944     }
 945
 946     frame = gst_nv_encoder_find_output_frame (encoder, task);
 947     if (!frame) {
 948       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
 949           ("Failed to find associated codec frame"));
 950       goto error;
 951     }
 952
 953     if (!gst_nv_encoder_device_lock (self)) {
 954       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
 955           ("Failed to lock device"));
 956       goto error;
 957     }
 958
 959     bitstream.version = gst_nvenc_get_lock_bitstream_version ();
 960     bitstream.outputBitstream = task->output_ptr;
 961
 962     status = NvEncLockBitstream (priv->session, &bitstream);
 963     if (status != NV_ENC_SUCCESS) {
 964       gst_nv_encoder_device_unlock (self);
 965       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
 966           ("Failed to lock bitstream, status: %" GST_NVENC_STATUS_FORMAT,
 967               GST_NVENC_STATUS_ARGS (status)));
 968       goto error;
 969     }
 970
 971     if (klass->create_output_buffer) {
 972       frame->output_buffer = klass->create_output_buffer (self, &bitstream);
 973     } else {
 974       frame->output_buffer =
 975           gst_buffer_new_memdup (bitstream.bitstreamBufferPtr,
 976           bitstream.bitstreamSizeInBytes);
 977     }
 978
 979     GST_BUFFER_FLAG_SET (frame->output_buffer, GST_BUFFER_FLAG_MARKER);
 980
 981     if (bitstream.pictureType == NV_ENC_PIC_TYPE_IDR)
 982       GST_VIDEO_CODEC_FRAME_SET_SYNC_POINT (frame);
 983
 984     NvEncUnlockBitstream (priv->session, task->output_ptr);
 985     gst_nv_encoder_device_unlock (self);
 986
 987     frame->dts = frame->pts - priv->dts_offset;
 988     frame->pts = bitstream.outputTimeStamp;
 989     frame->duration = bitstream.outputDuration;
 990
 991     ret = gst_video_encoder_finish_frame (encoder, frame);
 992     if (ret != GST_FLOW_OK) {
 993       GST_INFO_OBJECT (self,
 994           "Finish frame returned %s", gst_flow_get_name (ret));
 995     }
 996
 997     GST_NV_ENCODER_LOCK (self);
 998     gst_nv_encoder_task_reset (self, task);
 999     priv->last_flow = ret;
1000     g_cond_broadcast (&priv->cond);
1001     GST_NV_ENCODER_UNLOCK (self);
1002
1003     if (ret != GST_FLOW_OK) {
1004       GST_INFO_OBJECT (self, "Push returned %s", gst_flow_get_name (ret));
1005       goto exit_thread;
1006     }
1007   } while (TRUE);
1008
1009 exit_thread:
1010   {
1011     GST_INFO_OBJECT (self, "Exiting thread");
1012
1013     return NULL;
1014
1015   }
1016 error:
1017   {
1018     GST_NV_ENCODER_LOCK (self);
1019     gst_nv_encoder_task_reset (self, task);
1020     priv->last_flow = GST_FLOW_ERROR;
1021     g_cond_broadcast (&priv->cond);
1022     GST_NV_ENCODER_UNLOCK (self);
1023
1024     goto exit_thread;
1025   }
1026 }
1027
1028 static guint
1029 gst_nv_encoder_calculate_task_pool_size (GstNvEncoder * self,
1030     NV_ENC_CONFIG * config)
1031 {
1032   guint num_tasks;
1033
1034   /* At least 4 surfaces are required as documented by Nvidia Encoder guide */
1035   num_tasks = 4;
1036
1037   /* lookahead depth */
1038   num_tasks += config->rcParams.lookaheadDepth;
1039
1040   /* B frames + 1 */
1041   num_tasks += MAX (0, config->frameIntervalP - 1) + 1;
1042
1043   GST_DEBUG_OBJECT (self, "Calculated task pool size: %d "
1044       "(lookahead %d, frameIntervalP %d)",
1045       num_tasks, config->rcParams.lookaheadDepth, config->frameIntervalP);
1046
1047   return num_tasks;
1048 }
1049
1050 static gboolean
1051 gst_nv_encoder_open_encode_session (GstNvEncoder * self, gpointer * session)
1052 {
1053   GstNvEncoderPrivate *priv = self->priv;
1054   NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS session_params = { 0, };
1055   session_params.version =
1056       gst_nvenc_get_open_encode_session_ex_params_version ();
1057   session_params.apiVersion = gst_nvenc_get_api_version ();
1058   NVENCSTATUS status;
1059
1060   switch (priv->selected_device_mode) {
1061 #ifdef GST_CUDA_HAS_D3D
1062     case GST_NV_ENCODER_DEVICE_D3D11:
1063       session_params.deviceType = NV_ENC_DEVICE_TYPE_DIRECTX;
1064       session_params.device = gst_d3d11_device_get_device_handle (priv->device);
1065       break;
1066 #endif
1067     case GST_NV_ENCODER_DEVICE_CUDA:
1068       session_params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
1069       session_params.device = gst_cuda_context_get_handle (priv->context);
1070       break;
1071     default:
1072       g_assert_not_reached ();
1073       return FALSE;
1074   }
1075
1076   status = NvEncOpenEncodeSessionEx (&session_params, session);
1077   if (status != NV_ENC_SUCCESS) {
1078     GST_ERROR_OBJECT (self, "Failed to open session, status: %"
1079         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1080     return FALSE;
1081   }
1082
1083   return TRUE;
1084 }
1085
1086 #ifdef GST_CUDA_HAS_D3D
1087 static GstBufferPool *
1088 gst_nv_encoder_create_d3d11_pool (GstNvEncoder * self,
1089     GstVideoCodecState * state)
1090 {
1091   GstNvEncoderPrivate *priv = self->priv;
1092   GstStructure *config;
1093   GstBufferPool *pool = NULL;
1094   GstD3D11AllocationParams *params;
1095
1096   params = gst_d3d11_allocation_params_new (priv->device, &state->info,
1097       GST_D3D11_ALLOCATION_FLAG_DEFAULT, 0, D3D11_RESOURCE_MISC_SHARED);
1098
1099   pool = gst_d3d11_buffer_pool_new (priv->device);
1100
1101   config = gst_buffer_pool_get_config (pool);
1102   gst_buffer_pool_config_set_d3d11_allocation_params (config, params);
1103   gst_d3d11_allocation_params_free (params);
1104
1105   gst_buffer_pool_config_set_params (config, state->caps,
1106       GST_VIDEO_INFO_SIZE (&state->info), 0, 0);
1107   if (!gst_buffer_pool_set_config (pool, config)) {
1108     GST_ERROR_OBJECT (self, "Failed to set pool config");
1109     gst_object_unref (pool);
1110
1111     return NULL;
1112   }
1113
1114   if (!gst_buffer_pool_set_active (pool, TRUE)) {
1115     GST_ERROR_OBJECT (self, "Failed to set active");
1116     gst_object_unref (pool);
1117     return NULL;
1118   }
1119
1120   return pool;
1121 }
1122 #endif
1123
1124 static GstBufferPool *
1125 gst_nv_encoder_create_pool (GstNvEncoder * self, GstVideoCodecState * state)
1126 {
1127   GstNvEncoderPrivate *priv = self->priv;
1128   GstStructure *config;
1129   GstBufferPool *pool = NULL;
1130
1131   /* At this moment device type must be selected already */
1132   switch (priv->selected_device_mode) {
1133 #ifdef GST_CUDA_HAS_D3D
1134     case GST_NV_ENCODER_DEVICE_D3D11:
1135       return gst_nv_encoder_create_d3d11_pool (self, state);
1136 #endif
1137     case GST_NV_ENCODER_DEVICE_CUDA:
1138       pool = gst_cuda_buffer_pool_new (priv->context);
1139       break;
1140     default:
1141       g_assert_not_reached ();
1142       return FALSE;
1143   }
1144
1145   config = gst_buffer_pool_get_config (pool);
1146   gst_buffer_pool_config_set_params (config, state->caps,
1147       GST_VIDEO_INFO_SIZE (&state->info), 0, 0);
1148   if (!gst_buffer_pool_set_config (pool, config)) {
1149     GST_ERROR_OBJECT (self, "Failed to set pool config");
1150     gst_object_unref (pool);
1151
1152     return NULL;
1153   }
1154
1155   if (!gst_buffer_pool_set_active (pool, TRUE)) {
1156     GST_ERROR_OBJECT (self, "Failed to set active");
1157     gst_object_unref (pool);
1158     return NULL;
1159   }
1160
1161   return pool;
1162 }
1163
1164 static gboolean
1165 gst_nv_encoder_init_session (GstNvEncoder * self, GstBuffer * in_buf)
1166 {
1167   GstNvEncoderPrivate *priv = self->priv;
1168   GstNvEncoderClass *klass = GST_NV_ENCODER_GET_CLASS (self);
1169   GstVideoCodecState *state = priv->input_state;
1170   GstVideoInfo *info = &state->info;
1171   NVENCSTATUS status;
1172   guint task_pool_size;
1173   gint fps_n, fps_d;
1174   GstClockTime frame_duration, min_latency, max_latency;
1175   guint i;
1176
1177   gst_nv_encoder_reset (self);
1178
1179   memset (&priv->init_params, 0, sizeof (NV_ENC_INITIALIZE_PARAMS));
1180   memset (&priv->config, 0, sizeof (NV_ENC_CONFIG));
1181
1182   if (priv->selected_device_mode == GST_NV_ENCODER_DEVICE_AUTO_SELECT) {
1183     GstNvEncoderDeviceData data;
1184     gboolean ret;
1185
1186     if (!in_buf) {
1187       GST_DEBUG_OBJECT (self, "Unknown device mode, open session later");
1188       return TRUE;
1189     }
1190
1191     if (!klass->select_device (self, info, in_buf, &data)) {
1192       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1193           ("Failed to select device mode"));
1194       return FALSE;
1195     }
1196
1197     GST_DEBUG_OBJECT (self,
1198         "Selected device mode: %d, cuda-device-id: %d, adapter-luid %"
1199         G_GINT64_FORMAT, data.device_mode, data.cuda_device_id,
1200         data.adapter_luid);
1201
1202     g_assert (data.device_mode == GST_NV_ENCODER_DEVICE_CUDA ||
1203         data.device_mode == GST_NV_ENCODER_DEVICE_D3D11);
1204
1205     g_rec_mutex_lock (&priv->context_lock);
1206     priv->selected_device_mode = data.device_mode;
1207     priv->cuda_device_id = data.cuda_device_id;
1208     priv->dxgi_adapter_luid = data.adapter_luid;
1209     gst_clear_object (&priv->context);
1210     if (data.device_mode == GST_NV_ENCODER_DEVICE_CUDA)
1211       priv->context = (GstCudaContext *) data.device;
1212 #ifdef GST_CUDA_HAS_D3D
1213     gst_clear_object (&priv->device);
1214     if (data.device_mode == GST_NV_ENCODER_DEVICE_D3D11)
1215       priv->device = (GstD3D11Device *) data.device;
1216 #endif
1217
1218     ret = gst_nv_encoder_open (GST_VIDEO_ENCODER (self));
1219     g_rec_mutex_unlock (&priv->context_lock);
1220
1221     if (!ret) {
1222       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1223           ("Failed to open device"));
1224       return FALSE;
1225     }
1226   }
1227
1228   priv->internal_pool = gst_nv_encoder_create_pool (self, state);
1229   if (!priv->internal_pool) {
1230     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1231         ("Failed to create internal pool"));
1232     return FALSE;
1233   }
1234
1235   if (!gst_nv_encoder_device_lock (self)) {
1236     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL), ("Failed to lock device"));
1237     gst_nv_encoder_reset (self);
1238     return FALSE;
1239   }
1240
1241   if (!gst_nv_encoder_open_encode_session (self, &priv->session)) {
1242     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1243         ("Failed to open session"));
1244     goto error;
1245   }
1246
1247   if (!klass->set_format (self, state, priv->session, &priv->init_params,
1248           &priv->config)) {
1249     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL), ("Failed to set format"));
1250     goto error;
1251   }
1252
1253   priv->init_params.encodeConfig = &priv->config;
1254   status = NvEncInitializeEncoder (priv->session, &priv->init_params);
1255   if (status != NV_ENC_SUCCESS) {
1256     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1257         ("Failed to init encoder, status: %"
1258             GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status)));
1259     goto error;
1260   }
1261
1262   task_pool_size = gst_nv_encoder_calculate_task_pool_size (self,
1263       &priv->config);
1264   g_array_set_size (priv->task_pool, task_pool_size);
1265
1266   for (i = 0; i < task_pool_size; i++) {
1267     NV_ENC_CREATE_BITSTREAM_BUFFER buffer_params = { 0, };
1268     GstNvEncoderTask *task = (GstNvEncoderTask *)
1269         & g_array_index (priv->task_pool, GstNvEncoderTask, i);
1270
1271     task->encoder = self;
1272
1273     buffer_params.version = gst_nvenc_get_create_bitstream_buffer_version ();
1274     status = NvEncCreateBitstreamBuffer (priv->session, &buffer_params);
1275
1276     if (status != NV_ENC_SUCCESS) {
1277       GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1278           ("Failed to create bitstream buffer, status: %"
1279               GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status)));
1280       goto error;
1281     }
1282
1283     task->output_ptr = buffer_params.bitstreamBuffer;
1284
1285     if (priv->init_params.enableEncodeAsync) {
1286       if (!gst_nv_encoder_create_event_handle (self,
1287               priv->session, &task->event_handle)) {
1288         GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1289             ("Failed to create async event handle"));
1290         goto error;
1291       }
1292     }
1293
1294     g_queue_push_tail (&priv->free_tasks, task);
1295   }
1296   gst_nv_encoder_device_unlock (self);
1297
1298   if (!klass->set_output_state (self, priv->input_state, priv->session)) {
1299     GST_ELEMENT_ERROR (self, STREAM, ENCODE, (NULL),
1300         ("Failed to set output state"));
1301     gst_nv_encoder_reset (self);
1302     return FALSE;
1303   }
1304
1305   priv->encoding_thread = g_thread_new ("GstNvEncoderThread",
1306       (GThreadFunc) gst_nv_encoder_thread_func, self);
1307
1308   if (info->fps_n > 0 && info->fps_d > 0) {
1309     fps_n = info->fps_n;
1310     fps_d = info->fps_d;
1311   } else {
1312     fps_n = 25;
1313     fps_d = 1;
1314   }
1315
1316   frame_duration = gst_util_uint64_scale (GST_SECOND, fps_d, fps_n);
1317
1318   priv->dts_offset = 0;
1319   /* Calculate DTS offset for B frame. NVENC does not provide DTS */
1320   if (priv->config.frameIntervalP > 1)
1321     priv->dts_offset = frame_duration * (priv->config.frameIntervalP - 1);
1322
1323   min_latency = priv->dts_offset +
1324       priv->config.rcParams.lookaheadDepth * frame_duration;
1325   max_latency = frame_duration * priv->task_pool->len;
1326   gst_video_encoder_set_latency (GST_VIDEO_ENCODER (self),
1327       min_latency, max_latency);
1328
1329   return TRUE;
1330
1331 error:
1332   gst_nv_encoder_device_unlock (self);
1333
1334   gst_nv_encoder_reset (self);
1335
1336   return FALSE;
1337 }
1338
1339 static gboolean
1340 gst_nv_encoder_reconfigure_session (GstNvEncoder * self)
1341 {
1342   GstNvEncoderPrivate *priv = self->priv;
1343   NV_ENC_RECONFIGURE_PARAMS params = { 0, };
1344   NVENCSTATUS status;
1345
1346   if (!priv->session) {
1347     GST_WARNING_OBJECT (self,
1348         "Encoding session was not configured, open session");
1349     gst_nv_encoder_drain (self, TRUE);
1350
1351     return gst_nv_encoder_init_session (self, nullptr);
1352   }
1353
1354   params.version = gst_nvenc_get_reconfigure_params_version ();
1355   params.reInitEncodeParams = priv->init_params;
1356   params.reInitEncodeParams.encodeConfig = &priv->config;
1357
1358   status = NvEncReconfigureEncoder (priv->session, &params);
1359   if (status != NV_ENC_SUCCESS) {
1360     GST_WARNING_OBJECT (self, "Failed to reconfigure encoder, status %"
1361         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1362     gst_nv_encoder_drain (self, TRUE);
1363
1364     return gst_nv_encoder_init_session (self, nullptr);
1365   }
1366
1367   return TRUE;
1368 }
1369
1370 static gboolean
1371 gst_nv_encoder_set_format (GstVideoEncoder * encoder,
1372     GstVideoCodecState * state)
1373 {
1374   GstNvEncoder *self = GST_NV_ENCODER (encoder);
1375   GstNvEncoderPrivate *priv = self->priv;
1376
1377   gst_nv_encoder_drain (self, TRUE);
1378
1379   g_clear_pointer (&priv->input_state, gst_video_codec_state_unref);
1380   priv->input_state = gst_video_codec_state_ref (state);
1381   priv->last_flow = GST_FLOW_OK;
1382
1383   /* select device again on next buffer */
1384   if (priv->subclass_device_mode == GST_NV_ENCODER_DEVICE_AUTO_SELECT)
1385     priv->selected_device_mode = GST_NV_ENCODER_DEVICE_AUTO_SELECT;
1386
1387   return gst_nv_encoder_init_session (self, nullptr);
1388 }
1389
1390 static NV_ENC_BUFFER_FORMAT
1391 gst_nv_encoder_get_buffer_format (GstNvEncoder * self, GstVideoFormat format)
1392 {
1393   switch (format) {
1394     case GST_VIDEO_FORMAT_NV12:
1395       return NV_ENC_BUFFER_FORMAT_NV12;
1396     case GST_VIDEO_FORMAT_Y444:
1397       return NV_ENC_BUFFER_FORMAT_YUV444;
1398     case GST_VIDEO_FORMAT_P010_10LE:
1399       return NV_ENC_BUFFER_FORMAT_YUV420_10BIT;
1400     case GST_VIDEO_FORMAT_Y444_16LE:
1401       return NV_ENC_BUFFER_FORMAT_YUV444_10BIT;
1402     default:
1403       GST_ERROR_OBJECT (self, "Unexpected format %s",
1404           gst_video_format_to_string (format));
1405       g_assert_not_reached ();
1406       break;
1407   }
1408
1409   return NV_ENC_BUFFER_FORMAT_UNDEFINED;
1410 }
1411
1412 static GstFlowReturn
1413 gst_nv_encoder_copy_system (GstNvEncoder * self, const GstVideoInfo * info,
1414     GstBuffer * buffer, gpointer session, GstNvEncoderTask * task)
1415 {
1416   NVENCSTATUS status;
1417   GstVideoFrame frame;
1418   guint8 *dst_data;
1419   NV_ENC_BUFFER_FORMAT format;
1420
1421   format =
1422       gst_nv_encoder_get_buffer_format (self, GST_VIDEO_INFO_FORMAT (info));
1423   if (format == NV_ENC_BUFFER_FORMAT_UNDEFINED)
1424     return GST_FLOW_ERROR;
1425
1426   if (!gst_video_frame_map (&frame, info, buffer, GST_MAP_READ)) {
1427     GST_ERROR_OBJECT (self, "Failed to map buffer");
1428     return GST_FLOW_ERROR;
1429   }
1430
1431   if (!task->input_buffer.inputBuffer) {
1432     NV_ENC_CREATE_INPUT_BUFFER input_buffer = { 0, };
1433     input_buffer.version = gst_nvenc_get_create_input_buffer_version ();
1434     input_buffer.width = info->width;
1435     input_buffer.height = info->height;
1436     input_buffer.bufferFmt = format;
1437
1438     status = NvEncCreateInputBuffer (session, &input_buffer);
1439     if (status != NV_ENC_SUCCESS) {
1440       GST_ERROR_OBJECT (self, "Failed to create input buffer, status %"
1441           GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1442       gst_video_frame_unmap (&frame);
1443       return GST_FLOW_ERROR;
1444     }
1445
1446     task->input_buffer = input_buffer;
1447   }
1448
1449   task->lk_input_buffer.version = gst_nvenc_get_lock_input_buffer_version ();
1450   task->lk_input_buffer.inputBuffer = task->input_buffer.inputBuffer;
1451   status = NvEncLockInputBuffer (session, &task->lk_input_buffer);
1452   if (status != NV_ENC_SUCCESS) {
1453     GST_ERROR_OBJECT (self, "Failed to lock input buffer, status %"
1454         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1455     gst_video_frame_unmap (&frame);
1456     return GST_FLOW_ERROR;
1457   }
1458
1459   dst_data = (guint8 *) task->lk_input_buffer.bufferDataPtr;
1460
1461   for (guint i = 0; i < GST_VIDEO_FRAME_N_PLANES (&frame); i++) {
1462     guint8 *src_data = (guint8 *) GST_VIDEO_FRAME_PLANE_DATA (&frame, i);
1463     guint width_in_bytes = GST_VIDEO_FRAME_COMP_WIDTH (&frame, i) *
1464         GST_VIDEO_FRAME_COMP_PSTRIDE (&frame, i);
1465     guint stride = GST_VIDEO_FRAME_PLANE_STRIDE (&frame, i);
1466     guint height = GST_VIDEO_FRAME_COMP_HEIGHT (&frame, i);
1467
1468     for (guint j = 0; j < height; j++) {
1469       memcpy (dst_data, src_data, width_in_bytes);
1470       dst_data += task->lk_input_buffer.pitch;
1471       src_data += stride;
1472     }
1473   }
1474
1475   NvEncUnlockInputBuffer (session, task->input_buffer.inputBuffer);
1476   gst_video_frame_unmap (&frame);
1477
1478   return GST_FLOW_OK;
1479 }
1480
1481 static GstFlowReturn
1482 gst_nv_encoder_prepare_task_input_cuda (GstNvEncoder * self,
1483     const GstVideoInfo * info, GstBuffer * buffer, gpointer session,
1484     GstNvEncoderTask * task)
1485 {
1486   GstNvEncoderPrivate *priv = self->priv;
1487   GstMemory *mem;
1488   GstCudaMemory *cmem;
1489   NVENCSTATUS status;
1490
1491   mem = gst_buffer_peek_memory (buffer, 0);
1492   if (!gst_is_cuda_memory (mem)) {
1493     GST_LOG_OBJECT (self, "Not a CUDA buffer, system copy");
1494     return gst_nv_encoder_copy_system (self, info, buffer, session, task);
1495   }
1496
1497   cmem = GST_CUDA_MEMORY_CAST (mem);
1498   if (cmem->context != priv->context) {
1499     GST_LOG_OBJECT (self, "Different context, system copy");
1500     return gst_nv_encoder_copy_system (self, info, buffer, session, task);
1501   }
1502
1503   task->buffer = gst_buffer_ref (buffer);
1504   if (!gst_buffer_map (task->buffer, &task->map_info,
1505           (GstMapFlags) (GST_MAP_READ | GST_MAP_CUDA))) {
1506     GST_ERROR_OBJECT (self, "Failed to map buffer");
1507     gst_clear_buffer (&task->buffer);
1508
1509     return GST_FLOW_ERROR;
1510   }
1511
1512   cmem = (GstCudaMemory *) gst_buffer_peek_memory (task->buffer, 0);
1513
1514   task->register_resource.version = gst_nvenc_get_register_resource_version ();
1515   task->register_resource.resourceType =
1516       NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
1517   task->register_resource.width = cmem->info.width;
1518   task->register_resource.height = cmem->info.height;
1519   task->register_resource.pitch = cmem->info.stride[0];
1520   task->register_resource.resourceToRegister = task->map_info.data;
1521   task->register_resource.bufferFormat =
1522       gst_nv_encoder_get_buffer_format (self, GST_VIDEO_INFO_FORMAT (info));
1523   if (task->register_resource.bufferFormat == NV_ENC_BUFFER_FORMAT_UNDEFINED)
1524     return GST_FLOW_ERROR;
1525
1526   status = NvEncRegisterResource (session, &task->register_resource);
1527   if (status != NV_ENC_SUCCESS) {
1528     GST_ERROR_OBJECT (self, "Failed to register resource, status %"
1529         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1530
1531     gst_buffer_unmap (task->buffer, &task->map_info);
1532     gst_clear_buffer (&task->buffer);
1533
1534     return GST_FLOW_ERROR;
1535   }
1536
1537   task->mapped_resource.version = gst_nvenc_get_map_input_resource_version ();
1538   task->mapped_resource.registeredResource =
1539       task->register_resource.registeredResource;
1540   status = NvEncMapInputResource (session, &task->mapped_resource);
1541   if (status != NV_ENC_SUCCESS) {
1542     GST_ERROR_OBJECT (self, "Failed to map input resource, status %"
1543         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1544     NvEncUnregisterResource (session,
1545         task->register_resource.registeredResource);
1546
1547     gst_buffer_unmap (task->buffer, &task->map_info);
1548     gst_clear_buffer (&task->buffer);
1549
1550     return GST_FLOW_ERROR;
1551   }
1552
1553   return GST_FLOW_OK;
1554 }
1555
1556 #ifdef GST_CUDA_HAS_D3D
1557 static GstBuffer *
1558 gst_nv_encoder_copy_d3d11 (GstNvEncoder * self,
1559     GstBuffer * src_buffer, GstBufferPool * pool, gboolean shared)
1560 {
1561   GstNvEncoderPrivate *priv = self->priv;
1562   D3D11_TEXTURE2D_DESC src_desc, dst_desc;
1563   D3D11_BOX src_box;
1564   guint subresource_idx;
1565   GstMemory *src_mem, *dst_mem;
1566   GstMapInfo src_info, dst_info;
1567   ID3D11Texture2D *src_tex, *dst_tex;
1568   ID3D11Device *device_handle;
1569   ID3D11DeviceContext *device_context;
1570   GstBuffer *dst_buffer;
1571   GstFlowReturn ret;
1572   ComPtr < IDXGIResource > dxgi_resource;
1573   ComPtr < ID3D11Texture2D > shared_texture;
1574   HANDLE shared_handle;
1575   GstD3D11Device *device;
1576   HRESULT hr;
1577
1578   ret = gst_buffer_pool_acquire_buffer (pool, &dst_buffer, NULL);
1579   if (ret != GST_FLOW_OK) {
1580     GST_ERROR_OBJECT (self, "Failed to acquire buffer");
1581     return NULL;
1582   }
1583
1584   src_mem = gst_buffer_peek_memory (src_buffer, 0);
1585   dst_mem = gst_buffer_peek_memory (dst_buffer, 0);
1586
1587   device = GST_D3D11_MEMORY_CAST (src_mem)->device;
1588
1589   device_handle = gst_d3d11_device_get_device_handle (device);
1590   device_context = gst_d3d11_device_get_device_context_handle (device);
1591
1592   if (!gst_memory_map (src_mem, &src_info,
1593           (GstMapFlags) (GST_MAP_READ | GST_MAP_D3D11))) {
1594     GST_WARNING ("Failed to map src memory");
1595     gst_buffer_unref (dst_buffer);
1596     return NULL;
1597   }
1598
1599   if (!gst_memory_map (dst_mem, &dst_info,
1600           (GstMapFlags) (GST_MAP_WRITE | GST_MAP_D3D11))) {
1601     GST_WARNING ("Failed to map dst memory");
1602     gst_memory_unmap (src_mem, &src_info);
1603     gst_buffer_unref (dst_buffer);
1604     return NULL;
1605   }
1606
1607   src_tex = (ID3D11Texture2D *) src_info.data;
1608   dst_tex = (ID3D11Texture2D *) dst_info.data;
1609
1610   gst_d3d11_memory_get_texture_desc (GST_D3D11_MEMORY_CAST (src_mem),
1611       &src_desc);
1612   gst_d3d11_memory_get_texture_desc (GST_D3D11_MEMORY_CAST (dst_mem),
1613       &dst_desc);
1614   subresource_idx =
1615       gst_d3d11_memory_get_subresource_index (GST_D3D11_MEMORY_CAST (src_mem));
1616
1617   if (shared) {
1618     hr = dst_tex->QueryInterface (IID_PPV_ARGS (&dxgi_resource));
1619     if (!gst_d3d11_result (hr, priv->device)) {
1620       GST_ERROR_OBJECT (self,
1621           "IDXGIResource interface is not available, hr: 0x%x", (guint) hr);
1622       goto error;
1623     }
1624
1625     hr = dxgi_resource->GetSharedHandle (&shared_handle);
1626     if (!gst_d3d11_result (hr, priv->device)) {
1627       GST_ERROR_OBJECT (self, "Failed to get shared handle, hr: 0x%x",
1628           (guint) hr);
1629       goto error;
1630     }
1631
1632     hr = device_handle->OpenSharedResource (shared_handle,
1633         IID_PPV_ARGS (&shared_texture));
1634
1635     if (!gst_d3d11_result (hr, device)) {
1636       GST_ERROR_OBJECT (self, "Failed to get shared texture, hr: 0x%x",
1637           (guint) hr);
1638       goto error;
1639     }
1640
1641     dst_tex = shared_texture.Get ();
1642   }
1643
1644   src_box.left = 0;
1645   src_box.top = 0;
1646   src_box.front = 0;
1647   src_box.back = 1;
1648   src_box.right = MIN (src_desc.Width, dst_desc.Width);
1649   src_box.bottom = MIN (src_desc.Height, dst_desc.Height);
1650
1651   if (shared) {
1652     if (priv->fence && priv->fence->device != device)
1653       gst_clear_d3d11_fence (&priv->fence);
1654
1655     if (!priv->fence)
1656       priv->fence = gst_d3d11_device_create_fence (device);
1657
1658     if (!priv->fence) {
1659       GST_ERROR_OBJECT (self, "Couldn't crete fence");
1660       goto error;
1661     }
1662
1663     gst_d3d11_device_lock (device);
1664   }
1665
1666   device_context->CopySubresourceRegion (dst_tex, 0,
1667       0, 0, 0, src_tex, subresource_idx, &src_box);
1668
1669   if (shared) {
1670     if (!gst_d3d11_fence_signal (priv->fence) ||
1671         !gst_d3d11_fence_wait (priv->fence)) {
1672       GST_ERROR_OBJECT (self, "Couldn't sync GPU operation");
1673       gst_d3d11_device_unlock (device);
1674       gst_clear_d3d11_fence (&priv->fence);
1675       goto error;
1676     }
1677
1678     gst_d3d11_device_unlock (device);
1679   }
1680
1681   gst_memory_unmap (dst_mem, &dst_info);
1682   gst_memory_unmap (src_mem, &src_info);
1683
1684   return dst_buffer;
1685
1686 error:
1687   gst_memory_unmap (dst_mem, &dst_info);
1688   gst_memory_unmap (src_mem, &src_info);
1689   gst_buffer_unref (dst_buffer);
1690
1691   return NULL;
1692 }
1693
1694 static GstBuffer *
1695 gst_nv_encoder_upload_d3d11_frame (GstNvEncoder * self,
1696     const GstVideoInfo * info, GstBuffer * buffer, GstBufferPool * pool)
1697 {
1698   GstD3D11Memory *dmem;
1699   D3D11_TEXTURE2D_DESC desc;
1700
1701   dmem = (GstD3D11Memory *) gst_buffer_peek_memory (buffer, 0);
1702
1703   gst_d3d11_memory_get_texture_desc (dmem, &desc);
1704   if (desc.Usage != D3D11_USAGE_DEFAULT) {
1705     GST_TRACE_OBJECT (self, "Not a default usage texture, d3d11 copy");
1706     return gst_nv_encoder_copy_d3d11 (self, buffer, pool, FALSE);
1707   }
1708
1709   GST_TRACE_OBJECT (self, "Use input buffer without copy");
1710
1711   return gst_buffer_ref (buffer);
1712 }
1713
1714 static GstFlowReturn
1715 gst_nv_encoder_prepare_task_input_d3d11 (GstNvEncoder * self,
1716     const GstVideoInfo * info, GstBuffer * buffer, gpointer session,
1717     GstBufferPool * pool, GstNvEncoderTask * task)
1718 {
1719   GstNvEncoderPrivate *priv = self->priv;
1720   GstMemory *mem;
1721   GstD3D11Memory *dmem;
1722   D3D11_TEXTURE2D_DESC desc;
1723   NVENCSTATUS status;
1724
1725   if (gst_buffer_n_memory (buffer) > 1) {
1726     GST_LOG_OBJECT (self, "Not a native DXGI format, system copy");
1727     return gst_nv_encoder_copy_system (self, info, buffer, session, task);
1728   }
1729
1730   mem = gst_buffer_peek_memory (buffer, 0);
1731   if (!gst_is_d3d11_memory (mem)) {
1732     GST_LOG_OBJECT (self, "Not a D3D11 buffer, system copy");
1733     return gst_nv_encoder_copy_system (self, info, buffer, session, task);
1734   }
1735
1736   dmem = GST_D3D11_MEMORY_CAST (mem);
1737   if (dmem->device != priv->device) {
1738     gint64 adapter_luid;
1739
1740     g_object_get (dmem->device, "adapter-luid", &adapter_luid, NULL);
1741     if (adapter_luid == priv->dxgi_adapter_luid) {
1742       GST_LOG_OBJECT (self, "Different device but same GPU, copy d3d11");
1743       task->buffer = gst_nv_encoder_copy_d3d11 (self, buffer, pool, TRUE);
1744     } else {
1745       GST_LOG_OBJECT (self, "Different device, system copy");
1746       return gst_nv_encoder_copy_system (self, info, buffer, session, task);
1747     }
1748   }
1749
1750   if (!task->buffer)
1751     task->buffer = gst_nv_encoder_upload_d3d11_frame (self, info, buffer, pool);
1752
1753   if (!task->buffer) {
1754     GST_ERROR_OBJECT (self, "Failed to upload buffer");
1755     return GST_FLOW_ERROR;
1756   }
1757
1758   if (!gst_buffer_map (task->buffer, &task->map_info,
1759           (GstMapFlags) (GST_MAP_READ | GST_MAP_D3D11))) {
1760     GST_ERROR_OBJECT (self, "Failed to map buffer");
1761     gst_clear_buffer (&task->buffer);
1762
1763     return GST_FLOW_ERROR;
1764   }
1765
1766   dmem = (GstD3D11Memory *) gst_buffer_peek_memory (task->buffer, 0);
1767   gst_d3d11_memory_get_texture_desc (dmem, &desc);
1768
1769   task->register_resource.version = gst_nvenc_get_register_resource_version ();
1770   task->register_resource.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_DIRECTX;
1771   task->register_resource.width = desc.Width;
1772   task->register_resource.height = desc.Height;
1773   switch (desc.Format) {
1774     case DXGI_FORMAT_NV12:
1775       task->register_resource.bufferFormat = NV_ENC_BUFFER_FORMAT_NV12;
1776       break;
1777     case DXGI_FORMAT_P010:
1778       task->register_resource.bufferFormat = NV_ENC_BUFFER_FORMAT_YUV420_10BIT;
1779       break;
1780     default:
1781       GST_ERROR_OBJECT (self, "Unexpected DXGI format %d", desc.Format);
1782       g_assert_not_reached ();
1783       return GST_FLOW_ERROR;
1784   }
1785
1786   task->register_resource.subResourceIndex =
1787       gst_d3d11_memory_get_subresource_index (dmem);
1788   task->register_resource.resourceToRegister =
1789       gst_d3d11_memory_get_resource_handle (dmem);
1790
1791   status = NvEncRegisterResource (session, &task->register_resource);
1792   if (status != NV_ENC_SUCCESS) {
1793     GST_ERROR_OBJECT (self, "Failed to register resource, status %"
1794         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1795
1796     gst_buffer_unmap (task->buffer, &task->map_info);
1797     gst_clear_buffer (&task->buffer);
1798
1799     return GST_FLOW_ERROR;
1800   }
1801
1802   task->mapped_resource.version = gst_nvenc_get_map_input_resource_version ();
1803   task->mapped_resource.registeredResource =
1804       task->register_resource.registeredResource;
1805   status = NvEncMapInputResource (session, &task->mapped_resource);
1806   if (status != NV_ENC_SUCCESS) {
1807     GST_ERROR_OBJECT (self, "Failed to map input resource, status %"
1808         GST_NVENC_STATUS_FORMAT, GST_NVENC_STATUS_ARGS (status));
1809     NvEncUnregisterResource (session,
1810         task->register_resource.registeredResource);
1811
1812     gst_buffer_unmap (task->buffer, &task->map_info);
1813     gst_clear_buffer (&task->buffer);
1814
1815     return GST_FLOW_ERROR;
1816   }
1817
1818   return GST_FLOW_OK;
1819 }
1820 #endif
1821
1822 static GstFlowReturn
1823 gst_nv_encoder_prepare_task_input (GstNvEncoder * self,
1824     const GstVideoInfo * info, GstBuffer * buffer, gpointer session,
1825     GstBufferPool * pool, GstNvEncoderTask * task)
1826 {
1827   GstNvEncoderPrivate *priv = self->priv;
1828   GstFlowReturn ret = GST_FLOW_ERROR;
1829
1830   switch (priv->selected_device_mode) {
1831 #ifdef GST_CUDA_HAS_D3D
1832     case GST_NV_ENCODER_DEVICE_D3D11:
1833       ret = gst_nv_encoder_prepare_task_input_d3d11 (self, info, buffer,
1834           session, pool, task);
1835       break;
1836 #endif
1837     case GST_NV_ENCODER_DEVICE_CUDA:
1838       ret = gst_nv_encoder_prepare_task_input_cuda (self, info, buffer,
1839           session, task);
1840       break;
1841     default:
1842       g_assert_not_reached ();
1843       break;
1844   }
1845
1846   return ret;
1847 }
1848
1849 static GstFlowReturn
1850 gst_nv_encoder_handle_frame (GstVideoEncoder * encoder,
1851     GstVideoCodecFrame * frame)
1852 {
1853   GstNvEncoder *self = GST_NV_ENCODER (encoder);
1854   GstNvEncoderPrivate *priv = self->priv;
1855   GstNvEncoderClass *klass = GST_NV_ENCODER_GET_CLASS (self);
1856   GstFlowReturn ret = GST_FLOW_ERROR;
1857   GstNvEncoderTask *task = NULL;
1858   GstNvEncoderReconfigure reconfig;
1859   GstBuffer *in_buf = frame->input_buffer;
1860
1861   GST_TRACE_OBJECT (self, "Handle frame");
1862
1863   GST_NV_ENCODER_LOCK (self);
1864   ret = priv->last_flow;
1865   GST_NV_ENCODER_UNLOCK (self);
1866
1867   if (ret != GST_FLOW_OK) {
1868     GST_INFO_OBJECT (self, "Last flow was %s", gst_flow_get_name (ret));
1869     gst_video_encoder_finish_frame (encoder, frame);
1870
1871     return ret;
1872   }
1873
1874   if (!priv->session && !gst_nv_encoder_init_session (self, in_buf)) {
1875     GST_ERROR_OBJECT (self, "Encoder object was not configured");
1876     gst_video_encoder_finish_frame (encoder, frame);
1877
1878     return GST_FLOW_NOT_NEGOTIATED;
1879   }
1880
1881   reconfig = klass->check_reconfigure (self, &priv->config);
1882   switch (reconfig) {
1883     case GST_NV_ENCODER_RECONFIGURE_BITRATE:
1884       if (!gst_nv_encoder_reconfigure_session (self)) {
1885         gst_video_encoder_finish_frame (encoder, frame);
1886         return GST_FLOW_NOT_NEGOTIATED;
1887       }
1888       break;
1889     case GST_NV_ENCODER_RECONFIGURE_FULL:
1890     {
1891       gst_nv_encoder_drain (self, TRUE);
1892       if (!gst_nv_encoder_init_session (self, nullptr)) {
1893         gst_video_encoder_finish_frame (encoder, frame);
1894         return GST_FLOW_NOT_NEGOTIATED;
1895       }
1896       break;
1897     }
1898     default:
1899       break;
1900   }
1901
1902   /* Release stream lock temporarily for encoding thread to be able to
1903    * push encoded data */
1904   GST_VIDEO_ENCODER_STREAM_UNLOCK (self);
1905   ret = gst_nv_encoder_get_free_task (self, &task, TRUE);
1906   GST_VIDEO_ENCODER_STREAM_LOCK (self);
1907   if (ret != GST_FLOW_OK) {
1908     GST_DEBUG_OBJECT (self, "Last flow was %s", gst_flow_get_name (ret));
1909     gst_video_encoder_finish_frame (encoder, frame);
1910     return ret;
1911   }
1912
1913   if (!gst_nv_encoder_device_lock (self)) {
1914     GST_ERROR_OBJECT (self, "Failed to lock device");
1915     gst_video_encoder_finish_frame (encoder, frame);
1916
1917     return GST_FLOW_ERROR;
1918   }
1919
1920   g_assert (task->buffer == NULL);
1921   ret = gst_nv_encoder_prepare_task_input (self, &priv->input_state->info,
1922       in_buf, priv->session, priv->internal_pool, task);
1923   gst_nv_encoder_device_unlock (self);
1924
1925   if (ret != GST_FLOW_OK) {
1926     GST_ERROR_OBJECT (self, "Failed to upload frame");
1927     GST_NV_ENCODER_LOCK (self);
1928     gst_nv_encoder_task_reset (self, task);
1929     GST_NV_ENCODER_UNLOCK (self);
1930
1931     gst_video_encoder_finish_frame (encoder, frame);
1932
1933     return ret;
1934   }
1935
1936   ret = gst_nv_encoder_encode_frame (self, frame, task);
1937   if (ret != GST_FLOW_OK) {
1938     GST_ERROR_OBJECT (self, "Failed to encode frame");
1939     gst_video_encoder_finish_frame (encoder, frame);
1940
1941     return ret;
1942   }
1943
1944   gst_video_codec_frame_unref (frame);
1945
1946   return GST_FLOW_OK;
1947 }
1948
1949 static GstFlowReturn
1950 gst_nv_encoder_finish (GstVideoEncoder * encoder)
1951 {
1952   GstNvEncoder *self = GST_NV_ENCODER (encoder);
1953
1954   GST_DEBUG_OBJECT (self, "Finish");
1955
1956   gst_nv_encoder_drain (self, TRUE);
1957
1958   return GST_FLOW_OK;
1959 }
1960
1961 static gboolean
1962 gst_nv_encoder_flush (GstVideoEncoder * encoder)
1963 {
1964   GstNvEncoder *self = GST_NV_ENCODER (encoder);
1965   GstNvEncoderPrivate *priv = self->priv;
1966
1967   GST_DEBUG_OBJECT (self, "Flush");
1968
1969   gst_nv_encoder_drain (self, TRUE);
1970
1971   priv->last_flow = GST_FLOW_OK;
1972
1973   return TRUE;
1974 }
1975
1976 guint
1977 gst_nv_encoder_get_task_size (GstNvEncoder * encoder)
1978 {
1979   g_return_val_if_fail (GST_IS_NV_ENCODER (encoder), 0);
1980
1981   return encoder->priv->task_pool->len;
1982 }
1983
1984 void
1985 gst_nv_encoder_set_device_mode (GstNvEncoder * encoder,
1986     GstNvEncoderDeviceMode mode, guint cuda_device_id, gint64 adapter_luid)
1987 {
1988   GstNvEncoderPrivate *priv = encoder->priv;
1989
1990   priv->subclass_device_mode = mode;
1991   priv->selected_device_mode = mode;
1992   priv->cuda_device_id = cuda_device_id;
1993   priv->dxgi_adapter_luid = adapter_luid;
1994 }
1995
1996 /**
1997  * GstNvEncoderPreset:
1998  *
1999  * Since: 1.22
2000  */
2001 GType
2002 gst_nv_encoder_preset_get_type (void)
2003 {
2004   static GType preset_type = 0;
2005   static const GEnumValue presets[] = {
2006     {GST_NV_ENCODER_PRESET_DEFAULT, "Default", "default"},
2007     {GST_NV_ENCODER_PRESET_HP, "High Performance", "hp"},
2008     {GST_NV_ENCODER_PRESET_HQ, "High Quality", "hq"},
2009     {GST_NV_ENCODER_PRESET_LOW_LATENCY_DEFAULT, "Low Latency", "low-latency"},
2010     {GST_NV_ENCODER_PRESET_LOW_LATENCY_HQ, "Low Latency, High Quality",
2011         "low-latency-hq"},
2012     {GST_NV_ENCODER_PRESET_LOW_LATENCY_HP, "Low Latency, High Performance",
2013         "low-latency-hp"},
2014     {GST_NV_ENCODER_PRESET_LOSSLESS_DEFAULT, "Lossless", "lossless"},
2015     {GST_NV_ENCODER_PRESET_LOSSLESS_HP, "Lossless, High Performance",
2016         "lossless-hp"},
2017     {0, NULL, NULL},
2018   };
2019
2020   if (g_once_init_enter (&preset_type)) {
2021     GType type = g_enum_register_static ("GstNvEncoderPreset", presets);
2022
2023     g_once_init_leave (&preset_type, type);
2024   }
2025
2026   return preset_type;
2027 }
2028
2029 void
2030 gst_nv_encoder_preset_to_guid (GstNvEncoderPreset preset, GUID * guid)
2031 {
2032   switch (preset) {
2033     case GST_NV_ENCODER_PRESET_DEFAULT:
2034       *guid = NV_ENC_PRESET_DEFAULT_GUID;
2035       break;
2036     case GST_NV_ENCODER_PRESET_HP:
2037       *guid = NV_ENC_PRESET_HP_GUID;
2038       break;
2039     case GST_NV_ENCODER_PRESET_HQ:
2040       *guid = NV_ENC_PRESET_HQ_GUID;
2041       break;
2042     case GST_NV_ENCODER_PRESET_LOW_LATENCY_DEFAULT:
2043       *guid = NV_ENC_PRESET_LOW_LATENCY_DEFAULT_GUID;
2044       break;
2045     case GST_NV_ENCODER_PRESET_LOW_LATENCY_HQ:
2046       *guid = NV_ENC_PRESET_LOW_LATENCY_HQ_GUID;
2047       break;
2048     case GST_NV_ENCODER_PRESET_LOW_LATENCY_HP:
2049       *guid = NV_ENC_PRESET_LOW_LATENCY_HP_GUID;
2050       break;
2051     case GST_NV_ENCODER_PRESET_LOSSLESS_DEFAULT:
2052       *guid = NV_ENC_PRESET_LOSSLESS_DEFAULT_GUID;
2053       break;
2054     case GST_NV_ENCODER_PRESET_LOSSLESS_HP:
2055       *guid = NV_ENC_PRESET_LOSSLESS_HP_GUID;
2056       break;
2057     default:
2058       break;
2059   }
2060
2061   *guid = NV_ENC_PRESET_DEFAULT_GUID;
2062 }
2063
2064 /**
2065  * GstNvEncoderRCMode:
2066  *
2067  * Since: 1.22
2068  */
2069 GType
2070 gst_nv_encoder_rc_mode_get_type (void)
2071 {
2072   static GType rc_mode_type = 0;
2073   static const GEnumValue rc_modes[] = {
2074     {GST_NV_ENCODER_RC_MODE_CONSTQP, "Constant Quantization", "cqp"},
2075     {GST_NV_ENCODER_RC_MODE_VBR, "Variable Bit Rate", "vbr"},
2076     {GST_NV_ENCODER_RC_MODE_CBR, "Constant Bit Rate", "cbr"},
2077     {GST_NV_ENCODER_RC_MODE_CBR_LOWDELAY_HQ,
2078         "Low-Delay CBR, High Quality", "cbr-ld-hq"},
2079     {GST_NV_ENCODER_RC_MODE_CBR_HQ, "CBR, High Quality (slower)", "cbr-hq"},
2080     {GST_NV_ENCODER_RC_MODE_VBR_HQ, "VBR, High Quality (slower)", "vbr-hq"},
2081     {0, NULL, NULL},
2082   };
2083
2084   if (g_once_init_enter (&rc_mode_type)) {
2085     GType type = g_enum_register_static ("GstNvEncoderRCMode", rc_modes);
2086
2087     g_once_init_leave (&rc_mode_type, type);
2088   }
2089
2090   return rc_mode_type;
2091 }
2092
2093 NV_ENC_PARAMS_RC_MODE
2094 gst_nv_encoder_rc_mode_to_native (GstNvEncoderRCMode rc_mode)
2095 {
2096   switch (rc_mode) {
2097     case GST_NV_ENCODER_RC_MODE_CONSTQP:
2098       return NV_ENC_PARAMS_RC_CONSTQP;
2099     case GST_NV_ENCODER_RC_MODE_VBR:
2100       return NV_ENC_PARAMS_RC_VBR;
2101     case GST_NV_ENCODER_RC_MODE_CBR:
2102       return NV_ENC_PARAMS_RC_CBR;
2103     case GST_NV_ENCODER_RC_MODE_CBR_LOWDELAY_HQ:
2104       return NV_ENC_PARAMS_RC_CBR_LOWDELAY_HQ;
2105     case GST_NV_ENCODER_RC_MODE_CBR_HQ:
2106       return NV_ENC_PARAMS_RC_CBR_HQ;
2107     case GST_NV_ENCODER_RC_MODE_VBR_HQ:
2108       return NV_ENC_PARAMS_RC_VBR_HQ;
2109     default:
2110       break;
2111   }
2112
2113   return NV_ENC_PARAMS_RC_VBR;
2114 }
2115
2116 const gchar *
2117 gst_nv_encoder_status_to_string (NVENCSTATUS status)
2118 {
2119 #define CASE(err) \
2120     case err: \
2121     return G_STRINGIFY (err);
2122
2123   switch (status) {
2124       CASE (NV_ENC_SUCCESS);
2125       CASE (NV_ENC_ERR_NO_ENCODE_DEVICE);
2126       CASE (NV_ENC_ERR_UNSUPPORTED_DEVICE);
2127       CASE (NV_ENC_ERR_INVALID_ENCODERDEVICE);
2128       CASE (NV_ENC_ERR_INVALID_DEVICE);
2129       CASE (NV_ENC_ERR_DEVICE_NOT_EXIST);
2130       CASE (NV_ENC_ERR_INVALID_PTR);
2131       CASE (NV_ENC_ERR_INVALID_EVENT);
2132       CASE (NV_ENC_ERR_INVALID_PARAM);
2133       CASE (NV_ENC_ERR_INVALID_CALL);
2134       CASE (NV_ENC_ERR_OUT_OF_MEMORY);
2135       CASE (NV_ENC_ERR_ENCODER_NOT_INITIALIZED);
2136       CASE (NV_ENC_ERR_UNSUPPORTED_PARAM);
2137       CASE (NV_ENC_ERR_LOCK_BUSY);
2138       CASE (NV_ENC_ERR_NOT_ENOUGH_BUFFER);
2139       CASE (NV_ENC_ERR_INVALID_VERSION);
2140       CASE (NV_ENC_ERR_MAP_FAILED);
2141       CASE (NV_ENC_ERR_NEED_MORE_INPUT);
2142       CASE (NV_ENC_ERR_ENCODER_BUSY);
2143       CASE (NV_ENC_ERR_EVENT_NOT_REGISTERD);
2144       CASE (NV_ENC_ERR_GENERIC);
2145       CASE (NV_ENC_ERR_INCOMPATIBLE_CLIENT_KEY);
2146       CASE (NV_ENC_ERR_UNIMPLEMENTED);
2147       CASE (NV_ENC_ERR_RESOURCE_REGISTER_FAILED);
2148       CASE (NV_ENC_ERR_RESOURCE_NOT_REGISTERED);
2149       CASE (NV_ENC_ERR_RESOURCE_NOT_MAPPED);
2150     default:
2151       break;
2152   }
2153 #undef CASE
2154
2155   return "Unknown";
2156 }
2157
2158 GstNvEncoderClassData *
2159 gst_nv_encoder_class_data_new (void)
2160 {
2161   GstNvEncoderClassData *data = g_new0 (GstNvEncoderClassData, 1);
2162   data->ref_count = 1;
2163
2164   return data;
2165 }
2166
2167 GstNvEncoderClassData *
2168 gst_nv_encoder_class_data_ref (GstNvEncoderClassData * cdata)
2169 {
2170   g_atomic_int_add (&cdata->ref_count, 1);
2171
2172   return cdata;
2173 }
2174
2175 void
2176 gst_nv_encoder_class_data_unref (GstNvEncoderClassData * cdata)
2177 {
2178   if (g_atomic_int_dec_and_test (&cdata->ref_count)) {
2179     gst_clear_caps (&cdata->sink_caps);
2180     gst_clear_caps (&cdata->src_caps);
2181     if (cdata->formats)
2182       g_list_free_full (cdata->formats, (GDestroyNotify) g_free);
2183     if (cdata->profiles)
2184       g_list_free_full (cdata->profiles, (GDestroyNotify) g_free);
2185     g_free (cdata);
2186   }
2187 }
2188
2189 void
2190 gst_nv_encoder_get_encoder_caps (gpointer session, const GUID * encode_guid,
2191     GstNvEncoderDeviceCaps * device_caps)
2192 {
2193   GstNvEncoderDeviceCaps dev_caps = { 0, };
2194   NV_ENC_CAPS_PARAM caps_param = { 0, };
2195   NVENCSTATUS status;
2196   GUID guid = *encode_guid;
2197
2198   GST_DEBUG_CATEGORY_INIT (gst_nv_encoder_debug, "nvencoder", 0, "nvencoder");
2199
2200   caps_param.version = gst_nvenc_get_caps_param_version ();
2201
2202 #define CHECK_CAPS(to_query,val,default_val) G_STMT_START { \
2203   gint _val; \
2204   caps_param.capsToQuery = to_query; \
2205   status = NvEncGetEncodeCaps (session, guid, &caps_param, \
2206       &_val); \
2207   if (status != NV_ENC_SUCCESS) { \
2208     GST_WARNING ("Unable to query %s, status: %" \
2209         GST_NVENC_STATUS_FORMAT, G_STRINGIFY (to_query), \
2210         GST_NVENC_STATUS_ARGS (status)); \
2211     val = default_val; \
2212   } else { \
2213     GST_DEBUG ("%s: %d", G_STRINGIFY (to_query), _val); \
2214     val = _val; \
2215   } \
2216 } G_STMT_END
2217
2218   CHECK_CAPS (NV_ENC_CAPS_NUM_MAX_BFRAMES, dev_caps.max_bframes, 0);
2219   CHECK_CAPS (NV_ENC_CAPS_SUPPORTED_RATECONTROL_MODES,
2220       dev_caps.ratecontrol_modes, NV_ENC_PARAMS_RC_VBR);
2221   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_FIELD_ENCODING, dev_caps.field_encoding, 0);
2222   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_MONOCHROME, dev_caps.monochrome, 0);
2223   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_FMO, dev_caps.fmo, 0);
2224   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_QPELMV, dev_caps.qpelmv, 0);
2225   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_BDIRECT_MODE, dev_caps.bdirect_mode, 0);
2226   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_CABAC, dev_caps.cabac, 0);
2227   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_ADAPTIVE_TRANSFORM,
2228       dev_caps.adaptive_transform, 0);
2229   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_STEREO_MVC, dev_caps.stereo_mvc, 0);
2230   CHECK_CAPS (NV_ENC_CAPS_NUM_MAX_TEMPORAL_LAYERS, dev_caps.temoral_layers, 0);
2231   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_HIERARCHICAL_PFRAMES,
2232       dev_caps.hierarchical_pframes, 0);
2233   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_HIERARCHICAL_BFRAMES,
2234       dev_caps.hierarchical_bframes, 0);
2235   CHECK_CAPS (NV_ENC_CAPS_LEVEL_MAX, dev_caps.level_max, 0);
2236   CHECK_CAPS (NV_ENC_CAPS_LEVEL_MIN, dev_caps.level_min, 0);
2237   CHECK_CAPS (NV_ENC_CAPS_SEPARATE_COLOUR_PLANE,
2238       dev_caps.separate_colour_plane, 0);
2239   CHECK_CAPS (NV_ENC_CAPS_WIDTH_MAX, dev_caps.width_max, 4096);
2240   CHECK_CAPS (NV_ENC_CAPS_HEIGHT_MAX, dev_caps.height_max, 4096);
2241   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_TEMPORAL_SVC, dev_caps.temporal_svc, 0);
2242   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_DYN_RES_CHANGE, dev_caps.dyn_res_change, 0);
2243   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_DYN_BITRATE_CHANGE,
2244       dev_caps.dyn_bitrate_change, 0);
2245   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_DYN_FORCE_CONSTQP,
2246       dev_caps.dyn_force_constqp, 0);
2247   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_DYN_RCMODE_CHANGE,
2248       dev_caps.dyn_rcmode_change, 0);
2249   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_SUBFRAME_READBACK,
2250       dev_caps.subframe_readback, 0);
2251   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_CONSTRAINED_ENCODING,
2252       dev_caps.constrained_encoding, 0);
2253   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_INTRA_REFRESH, dev_caps.intra_refresh, 0);
2254   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_CUSTOM_VBV_BUF_SIZE,
2255       dev_caps.custom_vbv_buf_size, 0);
2256   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_DYNAMIC_SLICE_MODE,
2257       dev_caps.dynamic_slice_mode, 0);
2258   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_REF_PIC_INVALIDATION,
2259       dev_caps.ref_pic_invalidation, 0);
2260   CHECK_CAPS (NV_ENC_CAPS_PREPROC_SUPPORT, dev_caps.preproc_support, 0);
2261   /* NOTE: Async is Windows only */
2262 #ifdef G_OS_WIN32
2263   CHECK_CAPS (NV_ENC_CAPS_ASYNC_ENCODE_SUPPORT,
2264       dev_caps.async_encoding_support, 0);
2265 #endif
2266   CHECK_CAPS (NV_ENC_CAPS_MB_NUM_MAX, dev_caps.mb_num_max, 0);
2267   CHECK_CAPS (NV_ENC_CAPS_MB_PER_SEC_MAX, dev_caps.mb_per_sec_max, 0);
2268   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_YUV444_ENCODE, dev_caps.yuv444_encode, 0);
2269   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_LOSSLESS_ENCODE, dev_caps.lossless_encode, 0);
2270   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_SAO, dev_caps.sao, 0);
2271   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_MEONLY_MODE, dev_caps.meonly_mode, 0);
2272   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_LOOKAHEAD, dev_caps.lookahead, 0);
2273   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_TEMPORAL_AQ, dev_caps.temporal_aq, 0);
2274   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_10BIT_ENCODE,
2275       dev_caps.supports_10bit_encode, 0);
2276   CHECK_CAPS (NV_ENC_CAPS_NUM_MAX_LTR_FRAMES, dev_caps.num_max_ltr_frames, 0);
2277   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_WEIGHTED_PREDICTION,
2278       dev_caps.weighted_prediction, 0);
2279   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_BFRAME_REF_MODE, dev_caps.bframe_ref_mode, 0);
2280   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_EMPHASIS_LEVEL_MAP,
2281       dev_caps.emphasis_level_map, 0);
2282   CHECK_CAPS (NV_ENC_CAPS_WIDTH_MIN, dev_caps.width_min, 16);
2283   CHECK_CAPS (NV_ENC_CAPS_HEIGHT_MIN, dev_caps.height_min, 16);
2284   CHECK_CAPS (NV_ENC_CAPS_SUPPORT_MULTIPLE_REF_FRAMES,
2285       dev_caps.multiple_ref_frames, 0);
2286 #undef CHECK_CAPS
2287
2288   *device_caps = dev_caps;
2289 }
2290
2291 void
2292 gst_nv_encoder_merge_device_caps (const GstNvEncoderDeviceCaps * a,
2293     const GstNvEncoderDeviceCaps * b, GstNvEncoderDeviceCaps * merged)
2294 {
2295   GstNvEncoderDeviceCaps caps;
2296
2297 #define SELECT_MAX(value) G_STMT_START { \
2298   caps.value = MAX (a->value, b->value); \
2299 } G_STMT_END
2300
2301 #define SELECT_MIN(value) G_STMT_START { \
2302   caps.value = MAX (MIN (a->value, b->value), 1); \
2303 } G_STMT_END
2304
2305   SELECT_MAX (max_bframes);
2306   SELECT_MAX (ratecontrol_modes);
2307   SELECT_MAX (field_encoding);
2308   SELECT_MAX (monochrome);
2309   SELECT_MAX (fmo);
2310   SELECT_MAX (qpelmv);
2311   SELECT_MAX (bdirect_mode);
2312   SELECT_MAX (cabac);
2313   SELECT_MAX (adaptive_transform);
2314   SELECT_MAX (stereo_mvc);
2315   SELECT_MAX (temoral_layers);
2316   SELECT_MAX (hierarchical_pframes);
2317   SELECT_MAX (hierarchical_bframes);
2318   SELECT_MAX (level_max);
2319   SELECT_MAX (level_min);
2320   SELECT_MAX (separate_colour_plane);
2321   SELECT_MAX (width_max);
2322   SELECT_MAX (height_max);
2323   SELECT_MAX (temporal_svc);
2324   SELECT_MAX (dyn_res_change);
2325   SELECT_MAX (dyn_bitrate_change);
2326   SELECT_MAX (dyn_force_constqp);
2327   SELECT_MAX (dyn_rcmode_change);
2328   SELECT_MAX (subframe_readback);
2329   SELECT_MAX (constrained_encoding);
2330   SELECT_MAX (intra_refresh);
2331   SELECT_MAX (custom_vbv_buf_size);
2332   SELECT_MAX (dynamic_slice_mode);
2333   SELECT_MAX (ref_pic_invalidation);
2334   SELECT_MAX (preproc_support);
2335   SELECT_MAX (async_encoding_support);
2336   SELECT_MAX (mb_num_max);
2337   SELECT_MAX (mb_per_sec_max);
2338   SELECT_MAX (yuv444_encode);
2339   SELECT_MAX (lossless_encode);
2340   SELECT_MAX (sao);
2341   SELECT_MAX (meonly_mode);
2342   SELECT_MAX (lookahead);
2343   SELECT_MAX (temporal_aq);
2344   SELECT_MAX (supports_10bit_encode);
2345   SELECT_MAX (num_max_ltr_frames);
2346   SELECT_MAX (weighted_prediction);
2347   SELECT_MAX (bframe_ref_mode);
2348   SELECT_MAX (emphasis_level_map);
2349   SELECT_MIN (width_min);
2350   SELECT_MIN (height_min);
2351   SELECT_MAX (multiple_ref_frames);
2352
2353 #undef SELECT_MAX
2354 #undef SELECT_MIN
2355
2356   *merged = caps;
2357 }