Skip to content

Commit

Permalink
[Internal] Log GPU decoder/encoder usage (#3233)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #3233

Log usage when cuvid/nvenc is called.

Reviewed By: nateanl

Differential Revision: D44569410

fbshipit-source-id: 7c1e9fe70134af29c784fd076d1b128c9d2e0066
  • Loading branch information
mthrok authored and facebook-github-bot committed Apr 3, 2023
1 parent 61c31bc commit 6270e60
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 0 deletions.
9 changes: 9 additions & 0 deletions torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
#include <torchaudio/csrc/ffmpeg/hw_context.h>
#include <torchaudio/csrc/ffmpeg/stream_reader/stream_processor.h>
#include <stdexcept>
#include <string_view>

namespace torchaudio {
namespace io {
Expand Down Expand Up @@ -149,6 +150,11 @@ void open_codec(
ret >= 0, "Failed to initialize CodecContext: ", av_err2string(ret));
}

bool ends_with(std::string_view str, std::string_view suffix) {
return str.size() >= suffix.size() &&
0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix);
}

AVCodecContextPtr get_codec_ctx(
const AVCodecParameters* params,
const c10::optional<std::string>& decoder_name,
Expand All @@ -161,6 +167,9 @@ AVCodecContextPtr get_codec_ctx(
if (codec_ctx->hw_device_ctx) {
codec_ctx->hw_frames_ctx = get_hw_frames_ctx(codec_ctx);
}
if (ends_with(codec_ctx->codec->name, "_cuvid")) {
C10_LOG_API_USAGE_ONCE("torchaudio.io.StreamReaderCUDA");
}
return codec_ctx;
}

Expand Down
13 changes: 13 additions & 0 deletions torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -807,6 +807,15 @@ EncodeProcess get_audio_encode_process(
std::move(codec_ctx)};
}

namespace {

bool ends_with(std::string_view str, std::string_view suffix) {
return str.size() >= suffix.size() &&
0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix);
}

} // namespace

EncodeProcess get_video_encode_process(
AVFormatContext* format_ctx,
double frame_rate,
Expand Down Expand Up @@ -876,6 +885,10 @@ EncodeProcess get_video_encode_process(
}
open_codec(codec_ctx, encoder_option);

if (ends_with(codec_ctx->codec->name, "_nvenc")) {
C10_LOG_API_USAGE_ONCE("torchaudio.io.StreamReaderCUDA");
}

// 5. Build filter graph
FilterGraph filter_graph = get_video_filter_graph(
src_fmt,
Expand Down

0 comments on commit 6270e60

Please sign in to comment.