Replace c10::Dict with std::map in StreamReader/Writer (#3092)

Summary:
This commit is kind of clean up and preparation for future development.

We plan to pass around more complicated objects among StreamReader and StreamWriter, and TorchBind is not expressive enough for defining intermediate object, so we want to use PyBind11 for binding StreamReader/Writer.

PyBind11 converts Python dict into std::map, while TorchBind converts it into c10::Dict. Because of this descrepancy, conversion from c10::Dict to std::map have to happen in multiple places, and this makes the binding code thicker as it requires to wrapper methods.

Using std::map reduces the number of wrapper methods / conversions, because the same method can be bound for file-like object and the others.

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

Reviewed By: nateanl

Differential Revision: D43524808

Pulled By: mthrok

fbshipit-source-id: f7467c66ccd37dbf4abc337bbb18ffaac21a0058

torchaudio/csrc/ffmpeg/CMakeLists.txt

@@ -9,6 +9,7 @@ set(
   sources
   ffmpeg.cpp
   filter_graph.cpp
+  binding_utils.cpp
   stream_reader/buffer/common.cpp
   stream_reader/buffer/chunked_buffer.cpp
   stream_reader/buffer/unchunked_buffer.cpp

torchaudio/csrc/ffmpeg/binding_utils.cpp

+#include <torchaudio/csrc/ffmpeg/binding_utils.h>
+
+namespace torchaudio::io {
+
+OptionDictC10 to_c10(const OptionDict& src) {
+  OptionDictC10 ret;
+  for (auto const& [key, value] : src) {
+    ret.insert(key, value);
+  }
+  return ret;
+}
+
+OptionDict from_c10(const OptionDictC10& src) {
+  OptionDict ret;
+  for (const auto& it : src) {
+    ret.emplace(it.key(), it.value());
+  }
+  return ret;
+}
+
+c10::optional<OptionDict> from_c10(const c10::optional<OptionDictC10>& src) {
+  if (src) {
+    return {from_c10(src.value())};
+  }
+  return {};
+}
+
+} // namespace torchaudio::io

torchaudio/csrc/ffmpeg/binding_utils.h

+#pragma once
+#include <torch/types.h>
+
+namespace torchaudio::io {
+
+using OptionDict = std::map<std::string, std::string>;
+using OptionDictC10 = c10::Dict<std::string, std::string>;
+
+OptionDictC10 to_c10(const OptionDict&);
+OptionDict from_c10(const OptionDictC10&);
+c10::optional<OptionDict> from_c10(const c10::optional<OptionDictC10>&);
+
+} // namespace torchaudio::io

torchaudio/csrc/ffmpeg/ffmpeg.cpp

@@ -14,8 +14,8 @@ namespace io {
 AVDictionary* get_option_dict(const c10::optional<OptionDict>& option) {
   AVDictionary* opt = nullptr;
   if (option) {
-    for (const auto& it : option.value()) {
-      av_dict_set(&opt, it.key().c_str(), it.value().c_str(), 0);
+    for (auto const& [key, value] : option.value()) {
+      av_dict_set(&opt, key.c_str(), value.c_str(), 0);
     }
   }
   return opt;

torchaudio/csrc/ffmpeg/ffmpeg.h

@@ -27,7 +27,7 @@ extern "C" {
 namespace torchaudio {
 namespace io {
 
-using OptionDict = c10::Dict<std::string, std::string>;
+using OptionDict = std::map<std::string, std::string>;
 
 // https://github.com/FFmpeg/FFmpeg/blob/4e6debe1df7d53f3f59b37449b82265d5c08a172/doc/APIchanges#L252-L260
 // Starting from libavformat 59 (ffmpeg 5),

torchaudio/csrc/ffmpeg/pybind/pybind.cpp

@@ -24,7 +24,7 @@ PYBIND11_MODULE(_torchaudio_ffmpeg, m) {
       .def(py::init<
            py::object,
            const c10::optional<std::string>&,
-           const c10::optional<OptionMap>&,
+           const c10::optional<OptionDict>&,
            int64_t>())
       .def("num_src_streams", &StreamReaderFileObj::num_src_streams)
       .def("num_out_streams", &StreamReaderFileObj::num_out_streams)

torchaudio/csrc/ffmpeg/pybind/stream_reader.cpp

@@ -13,7 +13,7 @@ SrcInfoPyBind convert_pybind(SrcStreamInfo ssi) {
       ssi.bit_rate,
       ssi.num_frames,
       ssi.bits_per_sample,
-      dict2map(ssi.metadata),
+      ssi.metadata,
       ssi.sample_rate,
       ssi.num_channels,
       ssi.width,
@@ -28,49 +28,11 @@ StreamReaderFileObj::StreamReaderFileObj(
     const c10::optional<std::map<std::string, std::string>>& option,
     int64_t buffer_size)
     : FileObj(fileobj_, static_cast<int>(buffer_size), false),
-      StreamReaderBinding(pAVIO, format, map2dict(option)) {}
-
-std::map<std::string, std::string> StreamReaderFileObj::get_metadata() const {
-  return dict2map(StreamReader::get_metadata());
-};
+      StreamReaderBinding(pAVIO, format, option) {}
 
 SrcInfoPyBind StreamReaderFileObj::get_src_stream_info(int64_t i) {
   return convert_pybind(StreamReader::get_src_stream_info(static_cast<int>(i)));
 }
 
-void StreamReaderFileObj::add_audio_stream(
-    int64_t i,
-    int64_t frames_per_chunk,
-    int64_t num_chunks,
-    const c10::optional<std::string>& filter_desc,
-    const c10::optional<std::string>& decoder,
-    const c10::optional<std::map<std::string, std::string>>& decoder_option) {
-  StreamReader::add_audio_stream(
-      i,
-      frames_per_chunk,
-      num_chunks,
-      filter_desc,
-      decoder,
-      map2dict(decoder_option));
-}
-
-void StreamReaderFileObj::add_video_stream(
-    int64_t i,
-    int64_t frames_per_chunk,
-    int64_t num_chunks,
-    const c10::optional<std::string>& filter_desc,
-    const c10::optional<std::string>& decoder,
-    const c10::optional<std::map<std::string, std::string>>& decoder_option,
-    const c10::optional<std::string>& hw_accel) {
-  StreamReader::add_video_stream(
-      i,
-      frames_per_chunk,
-      num_chunks,
-      filter_desc,
-      decoder,
-      map2dict(decoder_option),
-      hw_accel);
-}
-
 } // namespace io
 } // namespace torchaudio

torchaudio/csrc/ffmpeg/pybind/stream_reader.h

@@ -16,25 +16,7 @@ class StreamReaderFileObj : protected FileObj, public StreamReaderBinding {
       const c10::optional<std::map<std::string, std::string>>& option,
       int64_t buffer_size);
 
-  std::map<std::string, std::string> get_metadata() const;
-
   SrcInfoPyBind get_src_stream_info(int64_t i);
-
-  void add_audio_stream(
-      int64_t i,
-      int64_t frames_per_chunk,
-      int64_t num_chunks,
-      const c10::optional<std::string>& filter_desc,
-      const c10::optional<std::string>& decoder,
-      const c10::optional<std::map<std::string, std::string>>& decoder_option);
-  void add_video_stream(
-      int64_t i,
-      int64_t frames_per_chunk,
-      int64_t num_chunks,
-      const c10::optional<std::string>& filter_desc,
-      const c10::optional<std::string>& decoder,
-      const c10::optional<std::map<std::string, std::string>>& decoder_option,
-      const c10::optional<std::string>& hw_accel);
 };
 
 } // namespace io

torchaudio/csrc/ffmpeg/pybind/stream_writer.cpp

@@ -10,46 +10,5 @@ StreamWriterFileObj::StreamWriterFileObj(
     : FileObj(fileobj_, static_cast<int>(buffer_size), true),
       StreamWriter(pAVIO, format) {}
 
-void StreamWriterFileObj::set_metadata(
-    const std::map<std::string, std::string>& metadata) {
-  StreamWriter::set_metadata(map2dict(metadata));
-}
-
-void StreamWriterFileObj::add_audio_stream(
-    int64_t sample_rate,
-    int64_t num_channels,
-    std::string format,
-    const c10::optional<std::string>& encoder,
-    const c10::optional<std::map<std::string, std::string>>& encoder_option,
-    const c10::optional<std::string>& encoder_format) {
-  StreamWriter::add_audio_stream(
-      sample_rate,
-      num_channels,
-      format,
-      encoder,
-      map2dict(encoder_option),
-      encoder_format);
-}
-
-void StreamWriterFileObj::add_video_stream(
-    double frame_rate,
-    int64_t width,
-    int64_t height,
-    std::string format,
-    const c10::optional<std::string>& encoder,
-    const c10::optional<std::map<std::string, std::string>>& encoder_option,
-    const c10::optional<std::string>& encoder_format,
-    const c10::optional<std::string>& hw_accel) {
-  StreamWriter::add_video_stream(
-      frame_rate,
-      width,
-      height,
-      format,
-      encoder,
-      map2dict(encoder_option),
-      encoder_format,
-      hw_accel);
-}
-
 } // namespace io
 } // namespace torchaudio

torchaudio/csrc/ffmpeg/pybind/stream_writer.h

@@ -11,24 +11,6 @@ class StreamWriterFileObj : private FileObj, public StreamWriter {
       py::object fileobj,
       const c10::optional<std::string>& format,
       int64_t buffer_size);
-
-  void set_metadata(const std::map<std::string, std::string>&);
-  void add_audio_stream(
-      int64_t sample_rate,
-      int64_t num_channels,
-      std::string format,
-      const c10::optional<std::string>& encoder,
-      const c10::optional<std::map<std::string, std::string>>& encoder_option,
-      const c10::optional<std::string>& encoder_format);
-  void add_video_stream(
-      double frame_rate,
-      int64_t width,
-      int64_t height,
-      std::string format,
-      const c10::optional<std::string>& encoder,
-      const c10::optional<std::map<std::string, std::string>>& encoder_option,
-      const c10::optional<std::string>& encoder_format,
-      const c10::optional<std::string>& hw_accel);
 };
 
 } // namespace io

torchaudio/csrc/ffmpeg/pybind/typedefs.cpp

@@ -89,28 +89,5 @@ FileObj::FileObj(py::object fileobj_, int buffer_size, bool writable)
       buffer_size(buffer_size),
       pAVIO(get_io_context(this, buffer_size, writable)) {}
 
-OptionDict map2dict(const OptionMap& src) {
-  OptionDict dict;
-  for (const auto& it : src) {
-    dict.insert(it.first.c_str(), it.second.c_str());
-  }
-  return dict;
-}
-
-c10::optional<OptionDict> map2dict(const c10::optional<OptionMap>& src) {
-  if (src) {
-    return c10::optional<OptionDict>{map2dict(src.value())};
-  }
-  return {};
-}
-
-OptionMap dict2map(const OptionDict& src) {
-  OptionMap ret;
-  for (const auto& it : src) {
-    ret.insert({it.key(), it.value()});
-  }
-  return ret;
-}
-
 } // namespace io
 } // namespace torchaudio

torchaudio/csrc/ffmpeg/pybind/typedefs.h

@@ -12,13 +12,5 @@ struct FileObj {
   FileObj(py::object fileobj, int buffer_size, bool writable);
 };
 
-using OptionMap = std::map<std::string, std::string>;
-
-OptionDict map2dict(const OptionMap& src);
-
-c10::optional<OptionDict> map2dict(const c10::optional<OptionMap>& src);
-
-OptionMap dict2map(const OptionDict& src);
-
 } // namespace io
 } // namespace torchaudio

torchaudio/csrc/ffmpeg/stream_reader/stream_reader.cpp

@@ -128,7 +128,7 @@ OptionDict parse_metadata(const AVDictionary* metadata) {
   AVDictionaryEntry* tag = nullptr;
   OptionDict ret;
   while ((tag = av_dict_get(metadata, "", tag, AV_DICT_IGNORE_SUFFIX))) {
-    ret.insert(std::string(tag->key), std::string(tag->value));
+    ret.emplace(std::string(tag->key), std::string(tag->value));
   }
   return ret;
 }

torchaudio/csrc/ffmpeg/stream_reader/stream_reader_binding.cpp

 #include <torch/script.h>
+#include <torchaudio/csrc/ffmpeg/binding_utils.h>
 #include <torchaudio/csrc/ffmpeg/stream_reader/stream_reader_wrapper.h>
 #include <stdexcept>
 
@@ -19,13 +20,14 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   });
   m.class_<StreamReaderBinding>("ffmpeg_StreamReader")
       .def(torch::init<>([](const std::string& src,
-                            const c10::optional<std::string>& device,
-                            const c10::optional<OptionDict>& option) {
-        return c10::make_intrusive<StreamReaderBinding>(src, device, option);
+                            const c10::optional<std::string>& format,
+                            const c10::optional<OptionDictC10>& option) {
+        return c10::make_intrusive<StreamReaderBinding>(
+            src, format, from_c10(option));
       }))
       .def("num_src_streams", [](S self) { return self->num_src_streams(); })
       .def("num_out_streams", [](S self) { return self->num_out_streams(); })
-      .def("get_metadata", [](S self) { return self->get_metadata(); })
+      .def("get_metadata", [](S self) { return to_c10(self->get_metadata()); })
       .def(
           "get_src_stream_info",
           [](S s, int64_t i) { return s->get_src_stream_info(i); })
@@ -47,14 +49,14 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
              int64_t num_chunks,
              const c10::optional<std::string>& filter_desc,
              const c10::optional<std::string>& decoder,
-             const c10::optional<OptionDict>& decoder_option) {
+             const c10::optional<OptionDictC10>& decoder_option) {
             s->add_audio_stream(
                 i,
                 frames_per_chunk,
                 num_chunks,
                 filter_desc,
                 decoder,
-                decoder_option);
+                from_c10(decoder_option));
           })
       .def(
           "add_video_stream",
@@ -64,7 +66,7 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
              int64_t num_chunks,
              const c10::optional<std::string>& filter_desc,
              const c10::optional<std::string>& decoder,
-             const c10::optional<OptionDict>& decoder_option,
+             const c10::optional<OptionDictC10>& decoder_option,
              const c10::optional<std::string>& hw_accel) {
             s->add_video_stream(
                 i,
@@ -72,7 +74,7 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
                 num_chunks,
                 filter_desc,
                 decoder,
-                decoder_option,
+                from_c10(decoder_option),
                 hw_accel);
           })
       .def("remove_stream", [](S s, int64_t i) { s->remove_stream(i); })

torchaudio/csrc/ffmpeg/stream_reader/stream_reader_wrapper.cpp

+#include <torchaudio/csrc/ffmpeg/binding_utils.h>
 #include <torchaudio/csrc/ffmpeg/stream_reader/stream_reader_wrapper.h>
 
 namespace torchaudio {
@@ -13,7 +14,7 @@ SrcInfo convert(SrcStreamInfo ssi) {
       ssi.bit_rate,
       ssi.num_frames,
       ssi.bits_per_sample,
-      ssi.metadata,
+      to_c10(ssi.metadata),
       ssi.sample_rate,
       ssi.num_channels,
       ssi.width,
@@ -27,18 +28,6 @@ OutInfo convert(OutputStreamInfo osi) {
 }
 } // namespace
 
-StreamReaderBinding::StreamReaderBinding(
-    const std::string& src,
-    const c10::optional<std::string>& format,
-    const c10::optional<OptionDict>& option)
-    : StreamReader(src, format, option) {}
-
-StreamReaderBinding::StreamReaderBinding(
-    AVIOContext* io_ctx,
-    const c10::optional<std::string>& format,
-    const c10::optional<OptionDict>& option)
-    : StreamReader(io_ctx, format, option) {}
-
 SrcInfo StreamReaderBinding::get_src_stream_info(int64_t i) {
   return convert(StreamReader::get_src_stream_info(static_cast<int>(i)));
 }

torchaudio/csrc/ffmpeg/stream_reader/stream_reader_wrapper.h

@@ -21,7 +21,7 @@ using SrcInfo = std::tuple<
     int64_t, // bit_rate
     int64_t, // num_frames
     int64_t, // bits_per_sample
-    OptionDict, // metadata
+    c10::Dict<std::string, std::string>, // metadata
     // Audio
     double, // sample_rate
     int64_t, // num_channels
@@ -60,15 +60,7 @@ using ChunkData = std::tuple<torch::Tensor, double>;
 // suitable for Binding the code (i.e. it receives/returns pritimitves)
 struct StreamReaderBinding : public StreamReader,
                              public torch::CustomClassHolder {
-  StreamReaderBinding(
-      const std::string& src,
-      const c10::optional<std::string>& device,
-      const c10::optional<OptionDict>& option);
-
-  StreamReaderBinding(
-      AVIOContext* io_ctx,
-      const c10::optional<std::string>& device,
-      const c10::optional<OptionDict>& option);
+  using StreamReader::StreamReader;
 
   SrcInfo get_src_stream_info(int64_t i);
   OutInfo get_out_stream_info(int64_t i);

torchaudio/csrc/ffmpeg/stream_writer/stream_writer.cpp

@@ -586,9 +586,8 @@ AVStream* StreamWriter::add_stream(AVCodecContextPtr& codec_ctx) {
 
 void StreamWriter::set_metadata(const OptionDict& metadata) {
   av_dict_free(&pFormatContext->metadata);
-  for (const auto& it : metadata) {
-    av_dict_set(
-        &pFormatContext->metadata, it.key().c_str(), it.value().c_str(), 0);
+  for (auto const& [key, value] : metadata) {
+    av_dict_set(&pFormatContext->metadata, key.c_str(), value.c_str(), 0);
   }
 }
 

torchaudio/csrc/ffmpeg/stream_writer/stream_writer_binding.cpp

 #include <torch/script.h>
+#include <torchaudio/csrc/ffmpeg/binding_utils.h>
 #include <torchaudio/csrc/ffmpeg/stream_writer/stream_writer.h>
 
 namespace torchaudio {
 namespace io {
 namespace {
 
-class StreamWriterBinding : public StreamWriter,
+struct StreamWriterBinding : public StreamWriter,
                              public torch::CustomClassHolder {
- public:
-  StreamWriterBinding(
-      const std::string& dst,
-      const c10::optional<std::string>& format)
-      : StreamWriter(dst, format) {}
+  using StreamWriter::StreamWriter;
 };
 
 using S = const c10::intrusive_ptr<StreamWriterBinding>&;
@@ -29,14 +26,14 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
              int64_t num_channels,
              const std::string& format,
              const c10::optional<std::string>& encoder,
-             const c10::optional<OptionDict>& encoder_option,
+             const c10::optional<OptionDictC10>& encoder_option,
              const c10::optional<std::string>& encoder_format) {
             s->add_audio_stream(
                 sample_rate,
                 num_channels,
                 format,
                 encoder,
-                encoder_option,
+                from_c10(encoder_option),
                 encoder_format);
           })
       .def(
@@ -47,7 +44,7 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
              int64_t height,
              const std::string& format,
              const c10::optional<std::string>& encoder,
-             const c10::optional<OptionDict>& encoder_option,
+             const c10::optional<OptionDictC10>& encoder_option,
              const c10::optional<std::string>& encoder_format,
              const c10::optional<std::string>& hw_accel) {
             s->add_video_stream(
@@ -56,17 +53,21 @@ TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
                 height,
                 format,
                 encoder,
-                encoder_option,
+                from_c10(encoder_option),
                 encoder_format,
                 hw_accel);
           })
       .def(
           "set_metadata",
-          [](S s, const OptionDict& metadata) { s->set_metadata(metadata); })
+          [](S s, const OptionDictC10& metadata) {
+            s->set_metadata(from_c10(metadata));
+          })
       .def("dump_format", [](S s, int64_t i) { s->dump_format(i); })
       .def(
           "open",
-          [](S s, const c10::optional<OptionDict>& option) { s->open(option); })
+          [](S s, const c10::optional<OptionDictC10>& option) {
+            s->open(from_c10(option));
+          })
       .def("close", [](S s) { s->close(); })
       .def(
           "write_audio_chunk",