git.sesse.net Git - ffmpeg/blob - libavutil/hwcontext_cuda.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg 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  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include "buffer.h"
  20 #include "common.h"
  21 #include "hwcontext.h"
  22 #include "hwcontext_internal.h"
  23 #include "hwcontext_cuda_internal.h"
  24 #include "mem.h"
  25 #include "pixdesc.h"
  26 #include "pixfmt.h"
  27 #include "imgutils.h"
  28
  29 #define CUDA_FRAME_ALIGNMENT 256
  30
  31 typedef struct CUDAFramesContext {
  32     int shift_width, shift_height;
  33 } CUDAFramesContext;
  34
  35 static const enum AVPixelFormat supported_formats[] = {
  36     AV_PIX_FMT_NV12,
  37     AV_PIX_FMT_YUV420P,
  38     AV_PIX_FMT_YUV444P,
  39     AV_PIX_FMT_P010,
  40     AV_PIX_FMT_P016,
  41     AV_PIX_FMT_YUV444P16,
  42     AV_PIX_FMT_0RGB32,
  43     AV_PIX_FMT_0BGR32,
  44 };
  45
  46 static int cuda_frames_get_constraints(AVHWDeviceContext *ctx,
  47                                        const void *hwconfig,
  48                                        AVHWFramesConstraints *constraints)
  49 {
  50     int i;
  51
  52     constraints->valid_sw_formats = av_malloc_array(FF_ARRAY_ELEMS(supported_formats) + 1,
  53                                                     sizeof(*constraints->valid_sw_formats));
  54     if (!constraints->valid_sw_formats)
  55         return AVERROR(ENOMEM);
  56
  57     for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++)
  58         constraints->valid_sw_formats[i] = supported_formats[i];
  59     constraints->valid_sw_formats[FF_ARRAY_ELEMS(supported_formats)] = AV_PIX_FMT_NONE;
  60
  61     constraints->valid_hw_formats = av_malloc_array(2, sizeof(*constraints->valid_hw_formats));
  62     if (!constraints->valid_hw_formats)
  63         return AVERROR(ENOMEM);
  64
  65     constraints->valid_hw_formats[0] = AV_PIX_FMT_CUDA;
  66     constraints->valid_hw_formats[1] = AV_PIX_FMT_NONE;
  67
  68     return 0;
  69 }
  70
  71 static void cuda_buffer_free(void *opaque, uint8_t *data)
  72 {
  73     AVHWFramesContext *ctx = opaque;
  74     AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx;
  75     CudaFunctions *cu = hwctx->internal->cuda_dl;
  76
  77     CUcontext dummy;
  78
  79     cu->cuCtxPushCurrent(hwctx->cuda_ctx);
  80
  81     cu->cuMemFree((CUdeviceptr)data);
  82
  83     cu->cuCtxPopCurrent(&dummy);
  84 }
  85
  86 static AVBufferRef *cuda_pool_alloc(void *opaque, int size)
  87 {
  88     AVHWFramesContext     *ctx = opaque;
  89     AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx;
  90     CudaFunctions          *cu = hwctx->internal->cuda_dl;
  91
  92     AVBufferRef *ret = NULL;
  93     CUcontext dummy = NULL;
  94     CUdeviceptr data;
  95     CUresult err;
  96
  97     err = cu->cuCtxPushCurrent(hwctx->cuda_ctx);
  98     if (err != CUDA_SUCCESS) {
  99         av_log(ctx, AV_LOG_ERROR, "Error setting current CUDA context\n");
 100         return NULL;
 101     }
 102
 103     err = cu->cuMemAlloc(&data, size);
 104     if (err != CUDA_SUCCESS)
 105         goto fail;
 106
 107     ret = av_buffer_create((uint8_t*)data, size, cuda_buffer_free, ctx, 0);
 108     if (!ret) {
 109         cu->cuMemFree(data);
 110         goto fail;
 111     }
 112
 113 fail:
 114     cu->cuCtxPopCurrent(&dummy);
 115     return ret;
 116 }
 117
 118 static int cuda_frames_init(AVHWFramesContext *ctx)
 119 {
 120     CUDAFramesContext *priv = ctx->internal->priv;
 121     int i;
 122
 123     for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++) {
 124         if (ctx->sw_format == supported_formats[i])
 125             break;
 126     }
 127     if (i == FF_ARRAY_ELEMS(supported_formats)) {
 128         av_log(ctx, AV_LOG_ERROR, "Pixel format '%s' is not supported\n",
 129                av_get_pix_fmt_name(ctx->sw_format));
 130         return AVERROR(ENOSYS);
 131     }
 132
 133     av_pix_fmt_get_chroma_sub_sample(ctx->sw_format, &priv->shift_width, &priv->shift_height);
 134
 135     if (!ctx->pool) {
 136         int size = av_image_get_buffer_size(ctx->sw_format, ctx->width, ctx->height, CUDA_FRAME_ALIGNMENT);
 137         if (size < 0)
 138             return size;
 139
 140         ctx->internal->pool_internal = av_buffer_pool_init2(size, ctx, cuda_pool_alloc, NULL);
 141         if (!ctx->internal->pool_internal)
 142             return AVERROR(ENOMEM);
 143     }
 144
 145     return 0;
 146 }
 147
 148 static int cuda_get_buffer(AVHWFramesContext *ctx, AVFrame *frame)
 149 {
 150     int res;
 151
 152     frame->buf[0] = av_buffer_pool_get(ctx->pool);
 153     if (!frame->buf[0])
 154         return AVERROR(ENOMEM);
 155
 156     res = av_image_fill_arrays(frame->data, frame->linesize, frame->buf[0]->data,
 157                                ctx->sw_format, ctx->width, ctx->height, CUDA_FRAME_ALIGNMENT);
 158     if (res < 0)
 159         return res;
 160
 161     // YUV420P is a special case.
 162     // Nvenc expects the U/V planes in swapped order from how ffmpeg expects them, also chroma is half-aligned
 163     if (ctx->sw_format == AV_PIX_FMT_YUV420P) {
 164         frame->linesize[1] = frame->linesize[2] = frame->linesize[0] / 2;
 165         frame->data[2]     = frame->data[1];
 166         frame->data[1]     = frame->data[2] + frame->linesize[2] * ctx->height / 2;
 167     }
 168
 169     frame->format = AV_PIX_FMT_CUDA;
 170     frame->width  = ctx->width;
 171     frame->height = ctx->height;
 172
 173     return 0;
 174 }
 175
 176 static int cuda_transfer_get_formats(AVHWFramesContext *ctx,
 177                                      enum AVHWFrameTransferDirection dir,
 178                                      enum AVPixelFormat **formats)
 179 {
 180     enum AVPixelFormat *fmts;
 181
 182     fmts = av_malloc_array(2, sizeof(*fmts));
 183     if (!fmts)
 184         return AVERROR(ENOMEM);
 185
 186     fmts[0] = ctx->sw_format;
 187     fmts[1] = AV_PIX_FMT_NONE;
 188
 189     *formats = fmts;
 190
 191     return 0;
 192 }
 193
 194 static int cuda_transfer_data_from(AVHWFramesContext *ctx, AVFrame *dst,
 195                                    const AVFrame *src)
 196 {
 197     CUDAFramesContext           *priv = ctx->internal->priv;
 198     AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx;
 199     CudaFunctions                 *cu = device_hwctx->internal->cuda_dl;
 200
 201     CUcontext dummy;
 202     CUresult err;
 203     int i;
 204
 205     err = cu->cuCtxPushCurrent(device_hwctx->cuda_ctx);
 206     if (err != CUDA_SUCCESS)
 207         return AVERROR_UNKNOWN;
 208
 209     for (i = 0; i < FF_ARRAY_ELEMS(src->data) && src->data[i]; i++) {
 210         CUDA_MEMCPY2D cpy = {
 211             .srcMemoryType = CU_MEMORYTYPE_DEVICE,
 212             .dstMemoryType = CU_MEMORYTYPE_HOST,
 213             .srcDevice     = (CUdeviceptr)src->data[i],
 214             .dstHost       = dst->data[i],
 215             .srcPitch      = src->linesize[i],
 216             .dstPitch      = dst->linesize[i],
 217             .WidthInBytes  = FFMIN(src->linesize[i], dst->linesize[i]),
 218             .Height        = src->height >> (i ? priv->shift_height : 0),
 219         };
 220
 221         err = cu->cuMemcpy2DAsync(&cpy, device_hwctx->stream);
 222         if (err != CUDA_SUCCESS) {
 223             av_log(ctx, AV_LOG_ERROR, "Error transferring the data from the CUDA frame\n");
 224             return AVERROR_UNKNOWN;
 225         }
 226     }
 227
 228     err = cu->cuStreamSynchronize(device_hwctx->stream);
 229     if (err != CUDA_SUCCESS) {
 230         av_log(ctx, AV_LOG_ERROR, "Error synchronizing CUDA stream\n");
 231         return AVERROR_UNKNOWN;
 232     }
 233
 234     cu->cuCtxPopCurrent(&dummy);
 235
 236     return 0;
 237 }
 238
 239 static int cuda_transfer_data_to(AVHWFramesContext *ctx, AVFrame *dst,
 240                                  const AVFrame *src)
 241 {
 242     CUDAFramesContext           *priv = ctx->internal->priv;
 243     AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx;
 244     CudaFunctions                 *cu = device_hwctx->internal->cuda_dl;
 245
 246     CUcontext dummy;
 247     CUresult err;
 248     int i;
 249
 250     err = cu->cuCtxPushCurrent(device_hwctx->cuda_ctx);
 251     if (err != CUDA_SUCCESS)
 252         return AVERROR_UNKNOWN;
 253
 254     for (i = 0; i < FF_ARRAY_ELEMS(src->data) && src->data[i]; i++) {
 255         CUDA_MEMCPY2D cpy = {
 256             .srcMemoryType = CU_MEMORYTYPE_HOST,
 257             .dstMemoryType = CU_MEMORYTYPE_DEVICE,
 258             .srcHost       = src->data[i],
 259             .dstDevice     = (CUdeviceptr)dst->data[i],
 260             .srcPitch      = src->linesize[i],
 261             .dstPitch      = dst->linesize[i],
 262             .WidthInBytes  = FFMIN(src->linesize[i], dst->linesize[i]),
 263             .Height        = src->height >> (i ? priv->shift_height : 0),
 264         };
 265
 266         err = cu->cuMemcpy2DAsync(&cpy, device_hwctx->stream);
 267         if (err != CUDA_SUCCESS) {
 268             av_log(ctx, AV_LOG_ERROR, "Error transferring the data to the CUDA frame\n");
 269             return AVERROR_UNKNOWN;
 270         }
 271     }
 272
 273     err = cu->cuStreamSynchronize(device_hwctx->stream);
 274     if (err != CUDA_SUCCESS) {
 275         av_log(ctx, AV_LOG_ERROR, "Error synchronizing CUDA stream\n");
 276         return AVERROR_UNKNOWN;
 277     }
 278
 279     cu->cuCtxPopCurrent(&dummy);
 280
 281     return 0;
 282 }
 283
 284 static void cuda_device_uninit(AVHWDeviceContext *ctx)
 285 {
 286     AVCUDADeviceContext *hwctx = ctx->hwctx;
 287
 288     if (hwctx->internal) {
 289         if (hwctx->internal->is_allocated && hwctx->cuda_ctx) {
 290             hwctx->internal->cuda_dl->cuCtxDestroy(hwctx->cuda_ctx);
 291             hwctx->cuda_ctx = NULL;
 292         }
 293         cuda_free_functions(&hwctx->internal->cuda_dl);
 294     }
 295
 296     av_freep(&hwctx->internal);
 297 }
 298
 299 static int cuda_device_init(AVHWDeviceContext *ctx)
 300 {
 301     AVCUDADeviceContext *hwctx = ctx->hwctx;
 302     int ret;
 303
 304     if (!hwctx->internal) {
 305         hwctx->internal = av_mallocz(sizeof(*hwctx->internal));
 306         if (!hwctx->internal)
 307             return AVERROR(ENOMEM);
 308     }
 309
 310     if (!hwctx->internal->cuda_dl) {
 311         ret = cuda_load_functions(&hwctx->internal->cuda_dl, ctx);
 312         if (ret < 0) {
 313             av_log(ctx, AV_LOG_ERROR, "Could not dynamically load CUDA\n");
 314             goto error;
 315         }
 316     }
 317
 318     return 0;
 319
 320 error:
 321     cuda_device_uninit(ctx);
 322     return ret;
 323 }
 324
 325 static int cuda_device_create(AVHWDeviceContext *ctx, const char *device,
 326                               AVDictionary *opts, int flags)
 327 {
 328     AVCUDADeviceContext *hwctx = ctx->hwctx;
 329     CudaFunctions *cu;
 330     CUdevice cu_device;
 331     CUcontext dummy;
 332     CUresult err;
 333     int device_idx = 0;
 334
 335     if (device)
 336         device_idx = strtol(device, NULL, 0);
 337
 338     if (cuda_device_init(ctx) < 0)
 339         goto error;
 340
 341     cu = hwctx->internal->cuda_dl;
 342
 343     err = cu->cuInit(0);
 344     if (err != CUDA_SUCCESS) {
 345         av_log(ctx, AV_LOG_ERROR, "Could not initialize the CUDA driver API\n");
 346         goto error;
 347     }
 348
 349     err = cu->cuDeviceGet(&cu_device, device_idx);
 350     if (err != CUDA_SUCCESS) {
 351         av_log(ctx, AV_LOG_ERROR, "Could not get the device number %d\n", device_idx);
 352         goto error;
 353     }
 354
 355     err = cu->cuCtxCreate(&hwctx->cuda_ctx, CU_CTX_SCHED_BLOCKING_SYNC, cu_device);
 356     if (err != CUDA_SUCCESS) {
 357         av_log(ctx, AV_LOG_ERROR, "Error creating a CUDA context\n");
 358         goto error;
 359     }
 360
 361     // Setting stream to NULL will make functions automatically use the default CUstream
 362     hwctx->stream = NULL;
 363
 364     cu->cuCtxPopCurrent(&dummy);
 365
 366     hwctx->internal->is_allocated = 1;
 367
 368     return 0;
 369
 370 error:
 371     cuda_device_uninit(ctx);
 372     return AVERROR_UNKNOWN;
 373 }
 374
 375 const HWContextType ff_hwcontext_type_cuda = {
 376     .type                 = AV_HWDEVICE_TYPE_CUDA,
 377     .name                 = "CUDA",
 378
 379     .device_hwctx_size    = sizeof(AVCUDADeviceContext),
 380     .frames_priv_size     = sizeof(CUDAFramesContext),
 381
 382     .device_create        = cuda_device_create,
 383     .device_init          = cuda_device_init,
 384     .device_uninit        = cuda_device_uninit,
 385     .frames_get_constraints = cuda_frames_get_constraints,
 386     .frames_init          = cuda_frames_init,
 387     .frames_get_buffer    = cuda_get_buffer,
 388     .transfer_get_formats = cuda_transfer_get_formats,
 389     .transfer_data_to     = cuda_transfer_data_to,
 390     .transfer_data_from   = cuda_transfer_data_from,
 391
 392     .pix_fmts             = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE },
 393 };