| @@ -0,0 +1,270 @@ |
| | 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 "buffer.h" |
| | 20 | #include "common.h" |
| | 21 | #include "hwcontext.h" |
| | 22 | #include "hwcontext_internal.h" |
| | 23 | #include "hwcontext_cuda.h" |
| | 24 | #include "mem.h" |
| | 25 | #include "pixdesc.h" |
| | 26 | #include "pixfmt.h" |
| | 27 | |
| | 28 | typedef struct CUDAFramesContext { |
| | 29 | int shift_width, shift_height; |
| | 30 | } CUDAFramesContext; |
| | 31 | |
| | 32 | static const enum AVPixelFormat supported_formats[] = { |
| | 33 | AV_PIX_FMT_NV12, |
| | 34 | AV_PIX_FMT_YUV420P, |
| | 35 | AV_PIX_FMT_YUV444P, |
| | 36 | }; |
| | 37 | |
| | 38 | static void cuda_buffer_free(void *opaque, uint8_t *data) |
| | 39 | { |
| | 40 | AVHWFramesContext *ctx = opaque; |
| | 41 | AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx; |
| | 42 | |
| | 43 | CUcontext dummy; |
| | 44 | |
| | 45 | cuCtxPushCurrent(hwctx->cuda_ctx); |
| | 46 | |
| | 47 | cuMemFree((CUdeviceptr)data); |
| | 48 | |
| | 49 | cuCtxPopCurrent(&dummy); |
| | 50 | } |
| | 51 | |
| | 52 | static AVBufferRef *cuda_pool_alloc(void *opaque, int size) |
| | 53 | { |
| | 54 | AVHWFramesContext *ctx = opaque; |
| | 55 | AVCUDADeviceContext *hwctx = ctx->device_ctx->hwctx; |
| | 56 | |
| | 57 | AVBufferRef *ret = NULL; |
| | 58 | CUcontext dummy = NULL; |
| | 59 | CUdeviceptr data; |
| | 60 | CUresult err; |
| | 61 | |
| | 62 | err = cuCtxPushCurrent(hwctx->cuda_ctx); |
| | 63 | if (err != CUDA_SUCCESS) { |
| | 64 | av_log(ctx, AV_LOG_ERROR, "Error setting current CUDA context\n"); |
| | 65 | return NULL; |
| | 66 | } |
| | 67 | |
| | 68 | err = cuMemAlloc(&data, size); |
| | 69 | if (err != CUDA_SUCCESS) |
| | 70 | goto fail; |
| | 71 | |
| | 72 | ret = av_buffer_create((uint8_t*)data, size, cuda_buffer_free, ctx, 0); |
| | 73 | if (!ret) { |
| | 74 | cuMemFree(data); |
| | 75 | goto fail; |
| | 76 | } |
| | 77 | |
| | 78 | fail: |
| | 79 | cuCtxPopCurrent(&dummy); |
| | 80 | return ret; |
| | 81 | } |
| | 82 | |
| | 83 | static int cuda_frames_init(AVHWFramesContext *ctx) |
| | 84 | { |
| | 85 | CUDAFramesContext *priv = ctx->internal->priv; |
| | 86 | int i; |
| | 87 | |
| | 88 | for (i = 0; i < FF_ARRAY_ELEMS(supported_formats); i++) { |
| | 89 | if (ctx->sw_format == supported_formats[i]) |
| | 90 | break; |
| | 91 | } |
| | 92 | if (i == FF_ARRAY_ELEMS(supported_formats)) { |
| | 93 | av_log(ctx, AV_LOG_ERROR, "Pixel format '%s' is not supported\n", |
| | 94 | av_get_pix_fmt_name(ctx->sw_format)); |
| | 95 | return AVERROR(ENOSYS); |
| | 96 | } |
| | 97 | |
| | 98 | av_pix_fmt_get_chroma_sub_sample(ctx->sw_format, &priv->shift_width, &priv->shift_height); |
| | 99 | |
| | 100 | if (!ctx->pool) { |
| | 101 | int size; |
| | 102 | |
| | 103 | switch (ctx->sw_format) { |
| | 104 | case AV_PIX_FMT_NV12: |
| | 105 | case AV_PIX_FMT_YUV420P: |
| | 106 | size = ctx->width * ctx->height * 3 / 2; |
| | 107 | break; |
| | 108 | case AV_PIX_FMT_YUV444P: |
| | 109 | size = ctx->width * ctx->height * 3; |
| | 110 | break; |
| | 111 | } |
| | 112 | |
| | 113 | ctx->internal->pool_internal = av_buffer_pool_init2(size, ctx, cuda_pool_alloc, NULL); |
| | 114 | if (!ctx->internal->pool_internal) |
| | 115 | return AVERROR(ENOMEM); |
| | 116 | } |
| | 117 | |
| | 118 | return 0; |
| | 119 | } |
| | 120 | |
| | 121 | static int cuda_get_buffer(AVHWFramesContext *ctx, AVFrame *frame) |
| | 122 | { |
| | 123 | frame->buf[0] = av_buffer_pool_get(ctx->pool); |
| | 124 | if (!frame->buf[0]) |
| | 125 | return AVERROR(ENOMEM); |
| | 126 | |
| | 127 | switch (ctx->sw_format) { |
| | 128 | case AV_PIX_FMT_NV12: |
| | 129 | frame->data[0] = frame->buf[0]->data; |
| | 130 | frame->data[1] = frame->data[0] + ctx->width * ctx->height; |
| | 131 | frame->linesize[0] = ctx->width; |
| | 132 | frame->linesize[1] = ctx->width; |
| | 133 | break; |
| | 134 | case AV_PIX_FMT_YUV420P: |
| | 135 | frame->data[0] = frame->buf[0]->data; |
| | 136 | frame->data[2] = frame->data[0] + ctx->width * ctx->height; |
| | 137 | frame->data[1] = frame->data[2] + ctx->width * ctx->height / 4; |
| | 138 | frame->linesize[0] = ctx->width; |
| | 139 | frame->linesize[1] = ctx->width / 2; |
| | 140 | frame->linesize[2] = ctx->width / 2; |
| | 141 | break; |
| | 142 | case AV_PIX_FMT_YUV444P: |
| | 143 | frame->data[0] = frame->buf[0]->data; |
| | 144 | frame->data[1] = frame->data[0] + ctx->width * ctx->height; |
| | 145 | frame->data[2] = frame->data[1] + ctx->width * ctx->height; |
| | 146 | frame->linesize[0] = ctx->width; |
| | 147 | frame->linesize[1] = ctx->width; |
| | 148 | frame->linesize[2] = ctx->width; |
| | 149 | break; |
| | 150 | default: |
| | 151 | av_frame_unref(frame); |
| | 152 | return AVERROR_BUG; |
| | 153 | } |
| | 154 | |
| | 155 | frame->format = AV_PIX_FMT_CUDA; |
| | 156 | frame->width = ctx->width; |
| | 157 | frame->height = ctx->height; |
| | 158 | |
| | 159 | return 0; |
| | 160 | } |
| | 161 | |
| | 162 | static int cuda_transfer_get_formats(AVHWFramesContext *ctx, |
| | 163 | enum AVHWFrameTransferDirection dir, |
| | 164 | enum AVPixelFormat **formats) |
| | 165 | { |
| | 166 | enum AVPixelFormat *fmts; |
| | 167 | |
| | 168 | fmts = av_malloc_array(2, sizeof(*fmts)); |
| | 169 | if (!fmts) |
| | 170 | return AVERROR(ENOMEM); |
| | 171 | |
| | 172 | fmts[0] = ctx->sw_format; |
| | 173 | fmts[1] = AV_PIX_FMT_NONE; |
| | 174 | |
| | 175 | *formats = fmts; |
| | 176 | |
| | 177 | return 0; |
| | 178 | } |
| | 179 | |
| | 180 | static int cuda_transfer_data_from(AVHWFramesContext *ctx, AVFrame *dst, |
| | 181 | const AVFrame *src) |
| | 182 | { |
| | 183 | CUDAFramesContext *priv = ctx->internal->priv; |
| | 184 | AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx; |
| | 185 | |
| | 186 | CUcontext dummy; |
| | 187 | CUresult err; |
| | 188 | int i; |
| | 189 | |
| | 190 | err = cuCtxPushCurrent(device_hwctx->cuda_ctx); |
| | 191 | if (err != CUDA_SUCCESS) |
| | 192 | return AVERROR_UNKNOWN; |
| | 193 | |
| | 194 | for (i = 0; i < FF_ARRAY_ELEMS(src->data) && src->data[i]; i++) { |
| | 195 | CUDA_MEMCPY2D cpy = { |
| | 196 | .srcMemoryType = CU_MEMORYTYPE_DEVICE, |
| | 197 | .dstMemoryType = CU_MEMORYTYPE_HOST, |
| | 198 | .srcDevice = (CUdeviceptr)src->data[i], |
| | 199 | .dstHost = dst->data[i], |
| | 200 | .srcPitch = src->linesize[i], |
| | 201 | .dstPitch = dst->linesize[i], |
| | 202 | .WidthInBytes = FFMIN(src->linesize[i], dst->linesize[i]), |
| | 203 | .Height = src->height >> (i ? priv->shift_height : 0), |
| | 204 | }; |
| | 205 | |
| | 206 | err = cuMemcpy2D(&cpy); |
| | 207 | if (err != CUDA_SUCCESS) { |
| | 208 | av_log(ctx, AV_LOG_ERROR, "Error transferring the data from the CUDA frame\n"); |
| | 209 | return AVERROR_UNKNOWN; |
| | 210 | } |
| | 211 | } |
| | 212 | |
| | 213 | cuCtxPopCurrent(&dummy); |
| | 214 | |
| | 215 | return 0; |
| | 216 | } |
| | 217 | |
| | 218 | static int cuda_transfer_data_to(AVHWFramesContext *ctx, AVFrame *dst, |
| | 219 | const AVFrame *src) |
| | 220 | { |
| | 221 | CUDAFramesContext *priv = ctx->internal->priv; |
| | 222 | AVCUDADeviceContext *device_hwctx = ctx->device_ctx->hwctx; |
| | 223 | |
| | 224 | CUcontext dummy; |
| | 225 | CUresult err; |
| | 226 | int i; |
| | 227 | |
| | 228 | err = cuCtxPushCurrent(device_hwctx->cuda_ctx); |
| | 229 | if (err != CUDA_SUCCESS) |
| | 230 | return AVERROR_UNKNOWN; |
| | 231 | |
| | 232 | for (i = 0; i < FF_ARRAY_ELEMS(src->data) && src->data[i]; i++) { |
| | 233 | CUDA_MEMCPY2D cpy = { |
| | 234 | .srcMemoryType = CU_MEMORYTYPE_HOST, |
| | 235 | .dstMemoryType = CU_MEMORYTYPE_DEVICE, |
| | 236 | .srcHost = src->data[i], |
| | 237 | .dstDevice = (CUdeviceptr)dst->data[i], |
| | 238 | .srcPitch = src->linesize[i], |
| | 239 | .dstPitch = dst->linesize[i], |
| | 240 | .WidthInBytes = FFMIN(src->linesize[i], dst->linesize[i]), |
| | 241 | .Height = src->height >> (i ? priv->shift_height : 0), |
| | 242 | }; |
| | 243 | |
| | 244 | err = cuMemcpy2D(&cpy); |
| | 245 | if (err != CUDA_SUCCESS) { |
| | 246 | av_log(ctx, AV_LOG_ERROR, "Error transferring the data from the CUDA frame\n"); |
| | 247 | return AVERROR_UNKNOWN; |
| | 248 | } |
| | 249 | } |
| | 250 | |
| | 251 | cuCtxPopCurrent(&dummy); |
| | 252 | |
| | 253 | return 0; |
| | 254 | } |
| | 255 | |
| | 256 | const HWContextType ff_hwcontext_type_cuda = { |
| | 257 | .type = AV_HWDEVICE_TYPE_CUDA, |
| | 258 | .name = "CUDA", |
| | 259 | |
| | 260 | .device_hwctx_size = sizeof(AVCUDADeviceContext), |
| | 261 | .frames_priv_size = sizeof(CUDAFramesContext), |
| | 262 | |
| | 263 | .frames_init = cuda_frames_init, |
| | 264 | .frames_get_buffer = cuda_get_buffer, |
| | 265 | .transfer_get_formats = cuda_transfer_get_formats, |
| | 266 | .transfer_data_to = cuda_transfer_data_to, |
| | 267 | .transfer_data_from = cuda_transfer_data_from, |
| | 268 | |
| | 269 | .pix_fmts = (const enum AVPixelFormat[]){ AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE }, |
| | 270 | }; |