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