Skip to content

Commit 10ac6d2

Browse files
mthrokfacebook-github-bot
authored andcommitted
Move helper functions out of common utility for better locality (#2512)
Summary: This commits move helper functions/definitions around so that better locality of logics are achieved. ## Detail `ffmpeg.[h|cpp]` implements classes that convert FFmpeg structures into RAII semantics. Initially it these classes included the construction logic in their constructors, but such logics were extracted to factory functions in #2373. Now the reason why the factory functions stayed in `ffmpeg.[h|cpp]` was because the logic for the initialization and clean-up of AVDictionary class was only available in `ffmpeg.cpp`. Now AVDictionary class handling is properly defined in #2507, the factory functions, which are not that reusable better stay with the implementation that use them. This makes `ffmpeg.h` lean and clean, makes it easier to see what can be reused. Pull Request resolved: #2512 Reviewed By: hwangjeff Differential Revision: D37477592 Pulled By: mthrok fbshipit-source-id: 8c1b5059ea5f44649cc0eb1f82d1a92877ef186e
1 parent 515fd01 commit 10ac6d2

File tree

5 files changed

+175
-209
lines changed

5 files changed

+175
-209
lines changed

torchaudio/csrc/ffmpeg/decoder.cpp

Lines changed: 120 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,19 +6,133 @@ namespace ffmpeg {
66
////////////////////////////////////////////////////////////////////////////////
77
// Decoder
88
////////////////////////////////////////////////////////////////////////////////
9+
namespace {
10+
AVCodecContextPtr get_decode_context(
11+
enum AVCodecID codec_id,
12+
const c10::optional<std::string>& decoder_name) {
13+
const AVCodec* pCodec = !decoder_name.has_value()
14+
? avcodec_find_decoder(codec_id)
15+
: avcodec_find_decoder_by_name(decoder_name.value().c_str());
16+
17+
if (!pCodec) {
18+
std::stringstream ss;
19+
if (!decoder_name.has_value()) {
20+
ss << "Unsupported codec: \"" << avcodec_get_name(codec_id) << "\", ("
21+
<< codec_id << ").";
22+
} else {
23+
ss << "Unsupported codec: \"" << decoder_name.value() << "\".";
24+
}
25+
throw std::runtime_error(ss.str());
26+
}
27+
28+
AVCodecContext* pCodecContext = avcodec_alloc_context3(pCodec);
29+
if (!pCodecContext) {
30+
throw std::runtime_error("Failed to allocate CodecContext.");
31+
}
32+
return AVCodecContextPtr(pCodecContext);
33+
}
34+
35+
#ifdef USE_CUDA
36+
enum AVPixelFormat get_hw_format(
37+
AVCodecContext* ctx,
38+
const enum AVPixelFormat* pix_fmts) {
39+
const enum AVPixelFormat* p = nullptr;
40+
AVPixelFormat pix_fmt = *static_cast<AVPixelFormat*>(ctx->opaque);
41+
for (p = pix_fmts; *p != -1; p++) {
42+
if (*p == pix_fmt) {
43+
return *p;
44+
}
45+
}
46+
TORCH_WARN("Failed to get HW surface format.");
47+
return AV_PIX_FMT_NONE;
48+
}
49+
50+
const AVCodecHWConfig* get_cuda_config(const AVCodec* pCodec) {
51+
for (int i = 0;; ++i) {
52+
const AVCodecHWConfig* config = avcodec_get_hw_config(pCodec, i);
53+
if (!config) {
54+
break;
55+
}
56+
if (config->device_type == AV_HWDEVICE_TYPE_CUDA &&
57+
config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) {
58+
return config;
59+
}
60+
}
61+
std::stringstream ss;
62+
ss << "CUDA device was requested, but the codec \"" << pCodec->name
63+
<< "\" is not supported.";
64+
throw std::runtime_error(ss.str());
65+
}
66+
#endif
67+
68+
void init_codec_context(
69+
AVCodecContext* pCodecContext,
70+
AVCodecParameters* pParams,
71+
const OptionDict& decoder_option,
72+
const torch::Device& device,
73+
AVBufferRefPtr& pHWBufferRef) {
74+
int ret = avcodec_parameters_to_context(pCodecContext, pParams);
75+
if (ret < 0) {
76+
throw std::runtime_error(
77+
"Failed to set CodecContext parameter: " + av_err2string(ret));
78+
}
79+
80+
#ifdef USE_CUDA
81+
// Enable HW Acceleration
82+
if (device.type() == c10::DeviceType::CUDA) {
83+
const AVCodecHWConfig* config = get_cuda_config(pCodecContext->codec);
84+
// TODO: check how to log
85+
// C10_LOG << "Decoder " << pCodec->name << " supports device " <<
86+
// av_hwdevice_get_type_name(config->device_type);
87+
88+
// https://www.ffmpeg.org/doxygen/trunk/hw__decode_8c_source.html#l00221
89+
// 1. Set HW pixel format (config->pix_fmt) to opaue pointer.
90+
static thread_local AVPixelFormat pix_fmt = config->pix_fmt;
91+
pCodecContext->opaque = static_cast<void*>(&pix_fmt);
92+
// 2. Set pCodecContext->get_format call back function which
93+
// will retrieve the HW pixel format from opaque pointer.
94+
pCodecContext->get_format = get_hw_format;
95+
// 3. Create HW device context and set to pCodecContext.
96+
AVBufferRef* hw_device_ctx = nullptr;
97+
ret = av_hwdevice_ctx_create(
98+
&hw_device_ctx,
99+
AV_HWDEVICE_TYPE_CUDA,
100+
std::to_string(device.index()).c_str(),
101+
nullptr,
102+
0);
103+
if (ret < 0) {
104+
throw std::runtime_error(
105+
"Failed to create CUDA device context: " + av_err2string(ret));
106+
}
107+
assert(hw_device_ctx);
108+
pCodecContext->hw_device_ctx = av_buffer_ref(hw_device_ctx);
109+
pHWBufferRef.reset(hw_device_ctx);
110+
}
111+
#endif
112+
113+
AVDictionary* opts = get_option_dict(decoder_option);
114+
ret = avcodec_open2(pCodecContext, pCodecContext->codec, &opts);
115+
clean_up_dict(opts);
116+
117+
if (ret < 0) {
118+
throw std::runtime_error(
119+
"Failed to initialize CodecContext: " + av_err2string(ret));
120+
}
121+
122+
if (pParams->codec_type == AVMEDIA_TYPE_AUDIO && !pParams->channel_layout)
123+
pParams->channel_layout =
124+
av_get_default_channel_layout(pCodecContext->channels);
125+
}
126+
} // namespace
127+
9128
Decoder::Decoder(
10129
AVCodecParameters* pParam,
11130
const c10::optional<std::string>& decoder_name,
12131
const OptionDict& decoder_option,
13132
const torch::Device& device)
14133
: pCodecContext(get_decode_context(pParam->codec_id, decoder_name)) {
15134
init_codec_context(
16-
pCodecContext,
17-
pParam,
18-
decoder_name,
19-
decoder_option,
20-
device,
21-
pHWBufferRef);
135+
pCodecContext, pParam, decoder_option, device, pHWBufferRef);
22136
}
23137

24138
int Decoder::process_packet(AVPacket* pPacket) {

torchaudio/csrc/ffmpeg/ffmpeg.cpp

Lines changed: 0 additions & 182 deletions
Original file line numberDiff line numberDiff line change
@@ -35,65 +35,13 @@ void clean_up_dict(AVDictionary* p) {
3535
}
3636
}
3737

38-
namespace {
39-
40-
// https://github.com/FFmpeg/FFmpeg/blob/4e6debe1df7d53f3f59b37449b82265d5c08a172/doc/APIchanges#L252-L260
41-
// Starting from libavformat 59 (ffmpeg 5),
42-
// AVInputFormat is const and related functions expect constant.
43-
#if LIBAVFORMAT_VERSION_MAJOR >= 59
44-
#define AVINPUT_FORMAT_CONST const
45-
#else
46-
#define AVINPUT_FORMAT_CONST
47-
#endif
48-
49-
} // namespace
50-
5138
////////////////////////////////////////////////////////////////////////////////
5239
// AVFormatContext
5340
////////////////////////////////////////////////////////////////////////////////
5441
void AVFormatContextDeleter::operator()(AVFormatContext* p) {
5542
avformat_close_input(&p);
5643
};
5744

58-
AVFormatContextPtr get_input_format_context(
59-
const std::string& src,
60-
const c10::optional<std::string>& device,
61-
const OptionDict& option,
62-
AVIOContext* io_ctx) {
63-
AVFormatContext* pFormat = avformat_alloc_context();
64-
if (!pFormat) {
65-
throw std::runtime_error("Failed to allocate AVFormatContext.");
66-
}
67-
if (io_ctx) {
68-
pFormat->pb = io_ctx;
69-
}
70-
71-
auto* pInput = [&]() -> AVINPUT_FORMAT_CONST AVInputFormat* {
72-
if (device.has_value()) {
73-
std::string device_str = device.value();
74-
AVINPUT_FORMAT_CONST AVInputFormat* p =
75-
av_find_input_format(device_str.c_str());
76-
if (!p) {
77-
std::ostringstream msg;
78-
msg << "Unsupported device/format: \"" << device_str << "\"";
79-
throw std::runtime_error(msg.str());
80-
}
81-
return p;
82-
}
83-
return nullptr;
84-
}();
85-
86-
AVDictionary* opt = get_option_dict(option);
87-
int ret = avformat_open_input(&pFormat, src.c_str(), pInput, &opt);
88-
clean_up_dict(opt);
89-
90-
if (ret < 0)
91-
throw std::runtime_error(
92-
"Failed to open the input \"" + src + "\" (" + av_err2string(ret) +
93-
").");
94-
return AVFormatContextPtr(pFormat);
95-
}
96-
9745
AVFormatContextPtr::AVFormatContextPtr(AVFormatContext* p)
9846
: Wrapper<AVFormatContext, AVFormatContextDeleter>(p) {}
9947

@@ -162,136 +110,6 @@ void AVCodecContextDeleter::operator()(AVCodecContext* p) {
162110
avcodec_free_context(&p);
163111
};
164112

165-
namespace {
166-
const AVCodec* get_decode_codec(
167-
enum AVCodecID codec_id,
168-
const c10::optional<std::string>& decoder_name) {
169-
const AVCodec* pCodec = !decoder_name.has_value()
170-
? avcodec_find_decoder(codec_id)
171-
: avcodec_find_decoder_by_name(decoder_name.value().c_str());
172-
173-
if (!pCodec) {
174-
std::stringstream ss;
175-
if (!decoder_name.has_value()) {
176-
ss << "Unsupported codec: \"" << avcodec_get_name(codec_id) << "\", ("
177-
<< codec_id << ").";
178-
} else {
179-
ss << "Unsupported codec: \"" << decoder_name.value() << "\".";
180-
}
181-
throw std::runtime_error(ss.str());
182-
}
183-
return pCodec;
184-
}
185-
186-
} // namespace
187-
188-
AVCodecContextPtr get_decode_context(
189-
enum AVCodecID codec_id,
190-
const c10::optional<std::string>& decoder_name) {
191-
const AVCodec* pCodec = get_decode_codec(codec_id, decoder_name);
192-
193-
AVCodecContext* pCodecContext = avcodec_alloc_context3(pCodec);
194-
if (!pCodecContext) {
195-
throw std::runtime_error("Failed to allocate CodecContext.");
196-
}
197-
return AVCodecContextPtr(pCodecContext);
198-
}
199-
200-
#ifdef USE_CUDA
201-
enum AVPixelFormat get_hw_format(
202-
AVCodecContext* ctx,
203-
const enum AVPixelFormat* pix_fmts) {
204-
const enum AVPixelFormat* p = nullptr;
205-
AVPixelFormat pix_fmt = *static_cast<AVPixelFormat*>(ctx->opaque);
206-
for (p = pix_fmts; *p != -1; p++) {
207-
if (*p == pix_fmt) {
208-
return *p;
209-
}
210-
}
211-
TORCH_WARN("Failed to get HW surface format.");
212-
return AV_PIX_FMT_NONE;
213-
}
214-
215-
const AVCodecHWConfig* get_cuda_config(const AVCodec* pCodec) {
216-
for (int i = 0;; ++i) {
217-
const AVCodecHWConfig* config = avcodec_get_hw_config(pCodec, i);
218-
if (!config) {
219-
break;
220-
}
221-
if (config->device_type == AV_HWDEVICE_TYPE_CUDA &&
222-
config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX) {
223-
return config;
224-
}
225-
}
226-
std::stringstream ss;
227-
ss << "CUDA device was requested, but the codec \"" << pCodec->name
228-
<< "\" is not supported.";
229-
throw std::runtime_error(ss.str());
230-
}
231-
#endif
232-
233-
void init_codec_context(
234-
AVCodecContext* pCodecContext,
235-
AVCodecParameters* pParams,
236-
const c10::optional<std::string>& decoder_name,
237-
const OptionDict& decoder_option,
238-
const torch::Device& device,
239-
AVBufferRefPtr& pHWBufferRef) {
240-
const AVCodec* pCodec = get_decode_codec(pParams->codec_id, decoder_name);
241-
242-
int ret = avcodec_parameters_to_context(pCodecContext, pParams);
243-
if (ret < 0) {
244-
throw std::runtime_error(
245-
"Failed to set CodecContext parameter: " + av_err2string(ret));
246-
}
247-
248-
#ifdef USE_CUDA
249-
// Enable HW Acceleration
250-
if (device.type() == c10::DeviceType::CUDA) {
251-
const AVCodecHWConfig* config = get_cuda_config(pCodec);
252-
// TODO: check how to log
253-
// C10_LOG << "Decoder " << pCodec->name << " supports device " <<
254-
// av_hwdevice_get_type_name(config->device_type);
255-
256-
// https://www.ffmpeg.org/doxygen/trunk/hw__decode_8c_source.html#l00221
257-
// 1. Set HW pixel format (config->pix_fmt) to opaue pointer.
258-
static thread_local AVPixelFormat pix_fmt = config->pix_fmt;
259-
pCodecContext->opaque = static_cast<void*>(&pix_fmt);
260-
// 2. Set pCodecContext->get_format call back function which
261-
// will retrieve the HW pixel format from opaque pointer.
262-
pCodecContext->get_format = get_hw_format;
263-
// 3. Create HW device context and set to pCodecContext.
264-
AVBufferRef* hw_device_ctx = nullptr;
265-
ret = av_hwdevice_ctx_create(
266-
&hw_device_ctx,
267-
AV_HWDEVICE_TYPE_CUDA,
268-
std::to_string(device.index()).c_str(),
269-
nullptr,
270-
0);
271-
if (ret < 0) {
272-
throw std::runtime_error(
273-
"Failed to create CUDA device context: " + av_err2string(ret));
274-
}
275-
assert(hw_device_ctx);
276-
pCodecContext->hw_device_ctx = av_buffer_ref(hw_device_ctx);
277-
pHWBufferRef.reset(hw_device_ctx);
278-
}
279-
#endif
280-
281-
AVDictionary* opts = get_option_dict(decoder_option);
282-
ret = avcodec_open2(pCodecContext, pCodec, &opts);
283-
clean_up_dict(opts);
284-
285-
if (ret < 0) {
286-
throw std::runtime_error(
287-
"Failed to initialize CodecContext: " + av_err2string(ret));
288-
}
289-
290-
if (pParams->codec_type == AVMEDIA_TYPE_AUDIO && !pParams->channel_layout)
291-
pParams->channel_layout =
292-
av_get_default_channel_layout(pCodecContext->channels);
293-
}
294-
295113
AVCodecContextPtr::AVCodecContextPtr(AVCodecContext* p)
296114
: Wrapper<AVCodecContext, AVCodecContextDeleter>(p) {}
297115

torchaudio/csrc/ffmpeg/ffmpeg.h

Lines changed: 9 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,15 @@ namespace ffmpeg {
2727

2828
using OptionDict = std::map<std::string, std::string>;
2929

30+
// https://github.com/FFmpeg/FFmpeg/blob/4e6debe1df7d53f3f59b37449b82265d5c08a172/doc/APIchanges#L252-L260
31+
// Starting from libavformat 59 (ffmpeg 5),
32+
// AVInputFormat is const and related functions expect constant.
33+
#if LIBAVFORMAT_VERSION_MAJOR >= 59
34+
#define AVFORMAT_CONST const
35+
#else
36+
#define AVFORMAT_CONST
37+
#endif
38+
3039
// Replacement of av_err2str, which causes
3140
// `error: taking address of temporary array`
3241
// https://github.com/joncampbell123/composite-video-simulator/issues/5
@@ -84,13 +93,6 @@ struct AVFormatContextPtr
8493
explicit AVFormatContextPtr(AVFormatContext* p);
8594
};
8695

87-
// create format context for reading media
88-
AVFormatContextPtr get_input_format_context(
89-
const std::string& src,
90-
const c10::optional<std::string>& device,
91-
const OptionDict& option,
92-
AVIOContext* io_ctx = nullptr);
93-
9496
////////////////////////////////////////////////////////////////////////////////
9597
// AVIO
9698
////////////////////////////////////////////////////////////////////////////////
@@ -166,20 +168,6 @@ struct AVCodecContextPtr
166168
explicit AVCodecContextPtr(AVCodecContext* p);
167169
};
168170

169-
// Allocate codec context from either decoder name or ID
170-
AVCodecContextPtr get_decode_context(
171-
enum AVCodecID codec_id,
172-
const c10::optional<std::string>& decoder);
173-
174-
// Initialize codec context with the parameters
175-
void init_codec_context(
176-
AVCodecContext* pCodecContext,
177-
AVCodecParameters* pParams,
178-
const c10::optional<std::string>& decoder_name,
179-
const OptionDict& decoder_option,
180-
const torch::Device& device,
181-
AVBufferRefPtr& pHWBufferRef);
182-
183171
////////////////////////////////////////////////////////////////////////////////
184172
// AVFilterGraph
185173
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)