Commit 4ba56323 authored by moto's avatar moto Committed by Facebook GitHub Bot
Browse files

Simplify HW acceleration code (#2534)

Summary:
FFmpeg's API provide multiple ways to initialize decoder. This PR simplifies the initialization by delegating the HW device context management to FFmpeg's native code.

Pull Request resolved: https://github.com/pytorch/audio/pull/2534

Reviewed By: hwangjeff

Differential Revision: D37734573

Pulled By: mthrok

fbshipit-source-id: e61736b4d4d2ca6e94d8965abd93b4e9a68e7351
parent 608b8ea6
...@@ -70,7 +70,7 @@ void init_codec_context( ...@@ -70,7 +70,7 @@ void init_codec_context(
AVCodecParameters* pParams, AVCodecParameters* pParams,
const c10::optional<OptionDict>& decoder_option, const c10::optional<OptionDict>& decoder_option,
const torch::Device& device, const torch::Device& device,
AVBufferRefPtr& pHWBufferRef) { enum AVPixelFormat* pHwFmt) {
int ret = avcodec_parameters_to_context(pCodecContext, pParams); int ret = avcodec_parameters_to_context(pCodecContext, pParams);
if (ret < 0) { if (ret < 0) {
throw std::runtime_error( throw std::runtime_error(
...@@ -80,33 +80,13 @@ void init_codec_context( ...@@ -80,33 +80,13 @@ void init_codec_context(
#ifdef USE_CUDA #ifdef USE_CUDA
// Enable HW Acceleration // Enable HW Acceleration
if (device.type() == c10::DeviceType::CUDA) { if (device.type() == c10::DeviceType::CUDA) {
const AVCodecHWConfig* config = get_cuda_config(pCodecContext->codec); *pHwFmt = get_cuda_config(pCodecContext->codec)->pix_fmt;
// TODO: check how to log
// C10_LOG << "Decoder " << pCodec->name << " supports device " <<
// av_hwdevice_get_type_name(config->device_type);
// https://www.ffmpeg.org/doxygen/trunk/hw__decode_8c_source.html#l00221 // https://www.ffmpeg.org/doxygen/trunk/hw__decode_8c_source.html#l00221
// 1. Set HW pixel format (config->pix_fmt) to opaue pointer. // 1. Set HW pixel format (config->pix_fmt) to opaue pointer.
static thread_local AVPixelFormat pix_fmt = config->pix_fmt; pCodecContext->opaque = static_cast<void*>(pHwFmt);
pCodecContext->opaque = static_cast<void*>(&pix_fmt);
// 2. Set pCodecContext->get_format call back function which // 2. Set pCodecContext->get_format call back function which
// will retrieve the HW pixel format from opaque pointer. // will retrieve the HW pixel format from opaque pointer.
pCodecContext->get_format = get_hw_format; pCodecContext->get_format = get_hw_format;
// 3. Create HW device context and set to pCodecContext.
AVBufferRef* hw_device_ctx = nullptr;
ret = av_hwdevice_ctx_create(
&hw_device_ctx,
AV_HWDEVICE_TYPE_CUDA,
std::to_string(device.index()).c_str(),
nullptr,
0);
if (ret < 0) {
throw std::runtime_error(
"Failed to create CUDA device context: " + av_err2string(ret));
}
assert(hw_device_ctx);
pCodecContext->hw_device_ctx = av_buffer_ref(hw_device_ctx);
pHWBufferRef.reset(hw_device_ctx);
} }
#endif #endif
...@@ -131,8 +111,7 @@ Decoder::Decoder( ...@@ -131,8 +111,7 @@ Decoder::Decoder(
const c10::optional<OptionDict>& decoder_option, const c10::optional<OptionDict>& decoder_option,
const torch::Device& device) const torch::Device& device)
: pCodecContext(get_decode_context(pParam->codec_id, decoder_name)) { : pCodecContext(get_decode_context(pParam->codec_id, decoder_name)) {
init_codec_context( init_codec_context(pCodecContext, pParam, decoder_option, device, &pHwFmt);
pCodecContext, pParam, decoder_option, device, pHWBufferRef);
} }
int Decoder::process_packet(AVPacket* pPacket) { int Decoder::process_packet(AVPacket* pPacket) {
......
...@@ -7,7 +7,7 @@ namespace ffmpeg { ...@@ -7,7 +7,7 @@ namespace ffmpeg {
class Decoder { class Decoder {
AVCodecContextPtr pCodecContext; AVCodecContextPtr pCodecContext;
AVBufferRefPtr pHWBufferRef; enum AVPixelFormat pHwFmt = AV_PIX_FMT_NONE;
public: public:
// Default constructable // Default constructable
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment