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