git.sesse.net Git - ffmpeg/blob - libavfilter/vf_hwupload_cuda.c

   1 /*
   2  * This file is part of Libav.
   3  *
   4  * Libav 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  * Libav 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 Libav; 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 = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 160     if (!out) {
 161         ret = AVERROR(ENOMEM);
 162         goto fail;
 163     }
 164
 165     out->width  = in->width;
 166     out->height = in->height;
 167
 168     ret = av_hwframe_transfer_data(out, in, 0);
 169     if (ret < 0) {
 170         av_log(ctx, AV_LOG_ERROR, "Error transferring data to the GPU\n");
 171         goto fail;
 172     }
 173
 174     ret = av_frame_copy_props(out, in);
 175     if (ret < 0)
 176         goto fail;
 177
 178     av_frame_free(&in);
 179
 180     return ff_filter_frame(ctx->outputs[0], out);
 181 fail:
 182     av_frame_free(&in);
 183     av_frame_free(&out);
 184     return ret;
 185 }
 186
 187 #define OFFSET(x) offsetof(CudaUploadContext, x)
 188 #define FLAGS AV_OPT_FLAG_VIDEO_PARAM
 189 static const AVOption options[] = {
 190     { "device", "Number of the device to use", OFFSET(device_idx), AV_OPT_TYPE_INT, { .i64 = 0 }, .flags = FLAGS },
 191     { NULL },
 192 };
 193
 194 static const AVClass cudaupload_class = {
 195     .class_name = "cudaupload",
 196     .item_name  = av_default_item_name,
 197     .option     = options,
 198     .version    = LIBAVUTIL_VERSION_INT,
 199 };
 200
 201 static const AVFilterPad cudaupload_inputs[] = {
 202     {
 203         .name         = "default",
 204         .type         = AVMEDIA_TYPE_VIDEO,
 205         .filter_frame = cudaupload_filter_frame,
 206     },
 207     { NULL }
 208 };
 209
 210 static const AVFilterPad cudaupload_outputs[] = {
 211     {
 212         .name         = "default",
 213         .type         = AVMEDIA_TYPE_VIDEO,
 214         .config_props = cudaupload_config_output,
 215     },
 216     { NULL }
 217 };
 218
 219 AVFilter ff_vf_hwupload_cuda = {
 220     .name        = "hwupload_cuda",
 221     .description = NULL_IF_CONFIG_SMALL("Upload a system memory frame to a CUDA device"),
 222
 223     .init      = cudaupload_init,
 224     .uninit    = cudaupload_uninit,
 225
 226     .query_formats = cudaupload_query_formats,
 227
 228     .priv_size  = sizeof(CudaUploadContext),
 229     .priv_class = &cudaupload_class,
 230
 231     .inputs    = cudaupload_inputs,
 232     .outputs   = cudaupload_outputs,
 233 };