git.sesse.net Git - ffmpeg/blob - libavfilter/vf_hwupload_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 "libavutil/buffer.h"
  20 #include "libavutil/hwcontext.h"
  21 #include "libavutil/hwcontext_cuda.h"
  22 #include "libavutil/log.h"
  23 #include "libavutil/opt.h"
  24
  25 #include "avfilter.h"
  26 #include "formats.h"
  27 #include "internal.h"
  28 #include "video.h"
  29
  30 typedef struct CudaUploadContext {
  31     const AVClass *class;
  32     int device_idx;
  33
  34     AVBufferRef *hwdevice;
  35     AVBufferRef *hwframe;
  36 } CudaUploadContext;
  37
  38 static void cudaupload_ctx_free(AVHWDeviceContext *ctx)
  39 {
  40     AVCUDADeviceContext *hwctx = ctx->hwctx;
  41     cuCtxDestroy(hwctx->cuda_ctx);
  42 }
  43
  44 static av_cold int cudaupload_init(AVFilterContext *ctx)
  45 {
  46     CudaUploadContext *s = ctx->priv;
  47
  48     AVHWDeviceContext   *device_ctx;
  49     AVCUDADeviceContext *device_hwctx;
  50     CUdevice device;
  51     CUcontext cuda_ctx = NULL, dummy;
  52     CUresult err;
  53     int ret;
  54
  55     err = cuInit(0);
  56     if (err != CUDA_SUCCESS) {
  57         av_log(ctx, AV_LOG_ERROR, "Could not initialize the CUDA driver API\n");
  58         return AVERROR_UNKNOWN;
  59     }
  60
  61     err = cuDeviceGet(&device, s->device_idx);
  62     if (err != CUDA_SUCCESS) {
  63         av_log(ctx, AV_LOG_ERROR, "Could not get the device number %d\n", s->device_idx);
  64         return AVERROR_UNKNOWN;
  65     }
  66
  67     err = cuCtxCreate(&cuda_ctx, 0, device);
  68     if (err != CUDA_SUCCESS) {
  69         av_log(ctx, AV_LOG_ERROR, "Error creating a CUDA context\n");
  70         return AVERROR_UNKNOWN;
  71     }
  72
  73     cuCtxPopCurrent(&dummy);
  74
  75     s->hwdevice = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_CUDA);
  76     if (!s->hwdevice) {
  77         cuCtxDestroy(cuda_ctx);
  78         return AVERROR(ENOMEM);
  79     }
  80
  81     device_ctx       = (AVHWDeviceContext*)s->hwdevice->data;
  82     device_ctx->free = cudaupload_ctx_free;
  83
  84     device_hwctx = device_ctx->hwctx;
  85     device_hwctx->cuda_ctx = cuda_ctx;
  86
  87     ret = av_hwdevice_ctx_init(s->hwdevice);
  88     if (ret < 0)
  89         return ret;
  90
  91     return 0;
  92 }
  93
  94 static av_cold void cudaupload_uninit(AVFilterContext *ctx)
  95 {
  96     CudaUploadContext *s = ctx->priv;
  97
  98     av_buffer_unref(&s->hwframe);
  99     av_buffer_unref(&s->hwdevice);
 100 }
 101
 102 static int cudaupload_query_formats(AVFilterContext *ctx)
 103 {
 104     static const enum AVPixelFormat input_pix_fmts[] = {
 105         AV_PIX_FMT_NV12, AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV444P,
 106         AV_PIX_FMT_NONE,
 107     };
 108     static const enum AVPixelFormat output_pix_fmts[] = {
 109         AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE,
 110     };
 111     AVFilterFormats *in_fmts  = ff_make_format_list(input_pix_fmts);
 112     AVFilterFormats *out_fmts = ff_make_format_list(output_pix_fmts);
 113
 114     ff_formats_ref(in_fmts,  &ctx->inputs[0]->out_formats);
 115     ff_formats_ref(out_fmts, &ctx->outputs[0]->in_formats);
 116
 117     return 0;
 118 }
 119
 120 static int cudaupload_config_output(AVFilterLink *outlink)
 121 {
 122     AVFilterContext *ctx = outlink->src;
 123     AVFilterLink *inlink = ctx->inputs[0];
 124     CudaUploadContext *s = ctx->priv;
 125
 126     AVHWFramesContext *hwframe_ctx;
 127     int ret;
 128
 129     av_buffer_unref(&s->hwframe);
 130     s->hwframe = av_hwframe_ctx_alloc(s->hwdevice);
 131     if (!s->hwframe)
 132         return AVERROR(ENOMEM);
 133
 134     hwframe_ctx            = (AVHWFramesContext*)s->hwframe->data;
 135     hwframe_ctx->format    = AV_PIX_FMT_CUDA;
 136     hwframe_ctx->sw_format = inlink->format;
 137     hwframe_ctx->width     = FFALIGN(inlink->w, 16);
 138     hwframe_ctx->height    = FFALIGN(inlink->h, 16);
 139
 140     ret = av_hwframe_ctx_init(s->hwframe);
 141     if (ret < 0)
 142         return ret;
 143
 144     outlink->hw_frames_ctx = av_buffer_ref(s->hwframe);
 145     if (!outlink->hw_frames_ctx)
 146         return AVERROR(ENOMEM);
 147
 148     return 0;
 149 }
 150
 151 static int cudaupload_filter_frame(AVFilterLink *link, AVFrame *in)
 152 {
 153     AVFilterContext   *ctx = link->dst;
 154     CudaUploadContext   *s = ctx->priv;
 155
 156     AVFrame *out = NULL;
 157     int ret;
 158
 159     out = av_frame_alloc();
 160     if (!out) {
 161         ret = AVERROR(ENOMEM);
 162         goto fail;
 163     }
 164
 165     ret = av_hwframe_get_buffer(s->hwframe, out, 0);
 166     if (ret < 0)
 167         goto fail;
 168
 169     out->width  = in->width;
 170     out->height = in->height;
 171
 172     ret = av_hwframe_transfer_data(out, in, 0);
 173     if (ret < 0) {
 174         av_log(ctx, AV_LOG_ERROR, "Error transferring data to the GPU\n");
 175         goto fail;
 176     }
 177
 178     ret = av_frame_copy_props(out, in);
 179     if (ret < 0)
 180         goto fail;
 181
 182     av_frame_free(&in);
 183
 184     return ff_filter_frame(ctx->outputs[0], out);
 185 fail:
 186     av_frame_free(&in);
 187     av_frame_free(&out);
 188     return ret;
 189 }
 190
 191 #define OFFSET(x) offsetof(CudaUploadContext, x)
 192 #define FLAGS (AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM)
 193 static const AVOption cudaupload_options[] = {
 194     { "device", "Number of the device to use", OFFSET(device_idx), AV_OPT_TYPE_INT, { .i64 = 0 }, .flags = FLAGS },
 195     { NULL },
 196 };
 197
 198 AVFILTER_DEFINE_CLASS(cudaupload);
 199
 200 static const AVFilterPad cudaupload_inputs[] = {
 201     {
 202         .name         = "default",
 203         .type         = AVMEDIA_TYPE_VIDEO,
 204         .filter_frame = cudaupload_filter_frame,
 205     },
 206     { NULL }
 207 };
 208
 209 static const AVFilterPad cudaupload_outputs[] = {
 210     {
 211         .name         = "default",
 212         .type         = AVMEDIA_TYPE_VIDEO,
 213         .config_props = cudaupload_config_output,
 214     },
 215     { NULL }
 216 };
 217
 218 AVFilter ff_vf_hwupload_cuda = {
 219     .name        = "hwupload_cuda",
 220     .description = NULL_IF_CONFIG_SMALL("Upload a system memory frame to a CUDA device."),
 221
 222     .init      = cudaupload_init,
 223     .uninit    = cudaupload_uninit,
 224
 225     .query_formats = cudaupload_query_formats,
 226
 227     .priv_size  = sizeof(CudaUploadContext),
 228     .priv_class = &cudaupload_class,
 229
 230     .inputs    = cudaupload_inputs,
 231     .outputs   = cudaupload_outputs,
 232 };