From 6270e609f2a76b1fd7cf29ab69bdb2832278114f Mon Sep 17 00:00:00 2001
From: Moto Hira <moto@meta.com>
Date: Mon, 3 Apr 2023 15:32:01 -0700
Subject: [PATCH] [Internal] Log GPU decoder/encoder usage (#3233)

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

Log usage when cuvid/nvenc is called.

Reviewed By: nateanl

Differential Revision: D44569410

fbshipit-source-id: 7c1e9fe70134af29c784fd076d1b128c9d2e0066
---
 .../csrc/ffmpeg/stream_reader/stream_processor.cpp  |  9 +++++++++
 .../csrc/ffmpeg/stream_writer/encode_process.cpp    | 13 +++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp b/torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp
index 1306710099..f3b9473761 100644
--- a/torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp
+++ b/torchaudio/csrc/ffmpeg/stream_reader/stream_processor.cpp
@@ -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 {
@@ -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,
@@ -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;
 }
 
diff --git a/torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp b/torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp
index 5e0037bdab..298684abed 100644
--- a/torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp
+++ b/torchaudio/csrc/ffmpeg/stream_writer/encode_process.cpp
@@ -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,
@@ -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,