hwcontext_cuda.c 9.3 KB

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