2 * This file is part of FFmpeg.
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.
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.
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
21 #include "hwcontext.h"
22 #include "hwcontext_internal.h"
23 #include "hwcontext_cuda_internal.h"
29 #define CUDA_FRAME_ALIGNMENT 256
31 typedef struct CUDAFramesContext {
32 int shift_width, shift_height;
35 static const enum AVPixelFormat supported_formats[] = {
46 static int cuda_frames_get_constraints(AVHWDeviceContext *ctx,
48 AVHWFramesConstraints *constraints)
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);
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;
61 constraints->valid_hw_formats = av_malloc_array(2, sizeof(*constraints->valid_hw_formats));
62 if (!constraints->valid_hw_formats)
63 return AVERROR(ENOMEM);
65 constraints->valid_hw_formats[0] = AV_PIX_FMT_CUDA;
66 constraints->valid_hw_formats[1] = AV_PIX_FMT_NONE;
71 static void cuda_buffer_free(void *opaque, uint8_t *data)
73 AVHWFramesContext *ctx = opaque;
74 AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx;
75 CudaFunctions *cu = hwctx->internal->cuda_dl;
79 cu->cuCtxPushCurrent(hwctx->cuda_ctx);
81 cu->cuMemFree((CUdeviceptr)data);
83 cu->cuCtxPopCurrent(&dummy);
86 static AVBufferRef *cuda_pool_alloc(void *opaque, int size)
88 AVHWFramesContext *ctx = opaque;
89 AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx;
90 CudaFunctions *cu = hwctx->internal->cuda_dl;
92 AVBufferRef *ret = NULL;
93 CUcontext dummy = NULL;
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");
103 err = cu->cuMemAlloc(&data, size);
104 if (err != CUDA_SUCCESS)
107 ret = av_buffer_create((uint8_t*)data, size, cuda_buffer_free, ctx, 0);
114 cu->cuCtxPopCurrent(&dummy);
118 static int cuda_frames_init(AVHWFramesContext *ctx)
120 CUDAFramesContext *priv = ctx->internal->priv;
123 for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++) {
124 if (ctx->sw_format == supported_formats[i])
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);
133 av_pix_fmt_get_chroma_sub_sample(ctx->sw_format, &priv->shift_width, &priv->shift_height);
136 int size = av_image_get_buffer_size(ctx->sw_format, ctx->width, ctx->height, CUDA_FRAME_ALIGNMENT);
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);
148 static int cuda_get_buffer(AVHWFramesContext *ctx, AVFrame *frame)
152 frame->buf[0] = av_buffer_pool_get(ctx->pool);
154 return AVERROR(ENOMEM);
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);
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;
169 frame->format = AV_PIX_FMT_CUDA;
170 frame->width = ctx->width;
171 frame->height = ctx->height;
176 static int cuda_transfer_get_formats(AVHWFramesContext *ctx,
177 enum AVHWFrameTransferDirection dir,
178 enum AVPixelFormat **formats)
180 enum AVPixelFormat *fmts;
182 fmts = av_malloc_array(2, sizeof(*fmts));
184 return AVERROR(ENOMEM);
186 fmts[0] = ctx->sw_format;
187 fmts[1] = AV_PIX_FMT_NONE;
194 static int cuda_transfer_data_from(AVHWFramesContext *ctx, AVFrame *dst,
197 CUDAFramesContext *priv = ctx->internal->priv;
198 AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx;
199 CudaFunctions *cu = device_hwctx->internal->cuda_dl;
205 err = cu->cuCtxPushCurrent(device_hwctx->cuda_ctx);
206 if (err != CUDA_SUCCESS)
207 return AVERROR_UNKNOWN;
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),
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;
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;
234 cu->cuCtxPopCurrent(&dummy);
239 static int cuda_transfer_data_to(AVHWFramesContext *ctx, AVFrame *dst,
242 CUDAFramesContext *priv = ctx->internal->priv;
243 AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx;
244 CudaFunctions *cu = device_hwctx->internal->cuda_dl;
250 err = cu->cuCtxPushCurrent(device_hwctx->cuda_ctx);
251 if (err != CUDA_SUCCESS)
252 return AVERROR_UNKNOWN;
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),
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;
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;
279 cu->cuCtxPopCurrent(&dummy);
284 static void cuda_device_uninit(AVHWDeviceContext *ctx)
286 AVCUDADeviceContext *hwctx = ctx->hwctx;
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;
293 cuda_free_functions(&hwctx->internal->cuda_dl);
296 av_freep(&hwctx->internal);
299 static int cuda_device_init(AVHWDeviceContext *ctx)
301 AVCUDADeviceContext *hwctx = ctx->hwctx;
304 if (!hwctx->internal) {
305 hwctx->internal = av_mallocz(sizeof(*hwctx->internal));
306 if (!hwctx->internal)
307 return AVERROR(ENOMEM);
310 if (!hwctx->internal->cuda_dl) {
311 ret = cuda_load_functions(&hwctx->internal->cuda_dl, ctx);
313 av_log(ctx, AV_LOG_ERROR, "Could not dynamically load CUDA\n");
321 cuda_device_uninit(ctx);
325 static int cuda_device_create(AVHWDeviceContext *ctx, const char *device,
326 AVDictionary *opts, int flags)
328 AVCUDADeviceContext *hwctx = ctx->hwctx;
336 device_idx = strtol(device, NULL, 0);
338 if (cuda_device_init(ctx) < 0)
341 cu = hwctx->internal->cuda_dl;
344 if (err != CUDA_SUCCESS) {
345 av_log(ctx, AV_LOG_ERROR, "Could not initialize the CUDA driver API\n");
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);
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");
361 // Setting stream to NULL will make functions automatically use the default CUstream
362 hwctx->stream = NULL;
364 cu->cuCtxPopCurrent(&dummy);
366 hwctx->internal->is_allocated = 1;
371 cuda_device_uninit(ctx);
372 return AVERROR_UNKNOWN;
375 const HWContextType ff_hwcontext_type_cuda = {
376 .type = AV_HWDEVICE_TYPE_CUDA,
379 .device_hwctx_size = sizeof(AVCUDADeviceContext),
380 .frames_priv_size = sizeof(CUDAFramesContext),
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,
392 .pix_fmts = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE },