CUDA support in blocks

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -33,6 +33,10 @@
 #include "nvcuvid_include/nvcuvid.h"
 
 namespace facebook::torchcodec {
+struct StandAloneFrameAttachedData {
+  cudaStream_t producer_stream = nullptr;
+  torch::stable::Tensor storage;
+};
 
 class BetaCudaDeviceInterface : public DeviceInterface {
  public:
@@ -41,9 +45,12 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 
   void initialize(const SharedAVCodecContext& codec_context) override;
 
-  void initialize_video(
+  void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;
@@ -91,6 +98,8 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       unsigned int pitch,
       const CUVIDPARSERDISPINFO& disp_info);
 
+  void make_frame_standalone(UniqueAVFrame& av_frame) override;
+
   UniqueAVFrame transfer_cpu_frame_to_gpu(
       const AVFrame& cpu_frame,
       AVPixelFormat target_pix_fmt);

src/torchcodec/_core/ColorConverter.cpp

@@ -15,10 +15,8 @@
 
 namespace facebook::torchcodec {
 
-ColorConverter::ColorConverter(
-    const StableDevice& device,
-    std::string_view device_variant) {
-  device_interface_ = create_device_interface(device, device_variant);
+ColorConverter::ColorConverter(const StableDevice& device) {
+  device_interface_ = create_device_interface(device);
   STD_TORCH_CHECK(
       device_interface_ != nullptr,
       "Failed to create device interface. This should never happen, please report.");
@@ -27,22 +25,13 @@ ColorConverter::ColorConverter(
   options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
   options.device = device;
 
-  // No user transforms and no stream: the converter is stream-agnostic and
-  // derives everything it needs from each frame.
-  //
-  // TODO_API_BREAKDOWN Need to refac/rethink all this. It seems unnatural that
-  // the color-converter needs its own device_interface_, but at the same time
-  // the color-conversion *must* be third-party aware, and the only way to
-  // achieve that for now is via the interface.
-  // This will become very relevant when we tackle CUDA, so we can defer until
-  // then. For now this is an OK hack.
+  // TODO_API_BREAKDOWN P1 It seems unnatural that the color-converter needs its
+  // own device_interface_, but at the same time the color-conversion *must* be
+  // third-party aware, and the only way to achieve that for now is via the
+  // interface.
   std::vector<std::unique_ptr<Transform>> no_transforms;
-  device_interface_->initialize_video(
-      /*av_stream=*/nullptr,
-      UniqueDecodingAVFormatContext{},
-      options,
-      no_transforms,
-      /*resized_output_dims=*/std::nullopt);
+  device_interface_->initialize_color_conversion(
+      options, no_transforms, /*resized_output_dims=*/std::nullopt);
 }
 
 torch::stable::Tensor ColorConverter::convert(const AVFrame& av_frame) {

src/torchcodec/_core/ColorConverter.h

@@ -18,8 +18,7 @@ namespace facebook::torchcodec {
 class FORCE_PUBLIC_VISIBILITY ColorConverter {
  public:
   explicit ColorConverter(
-      const StableDevice& device = StableDevice(kStableCPU),
-      std::string_view device_variant = "default");
+      const StableDevice& device = StableDevice(kStableCPU));
 
   torch::stable::Tensor convert(const AVFrame& av_frame);
 

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -80,20 +80,18 @@ void CpuDeviceInterface::initialize(const SharedAVCodecContext& codec_context) {
   codec_context_ = codec_context;
 }
 
-void CpuDeviceInterface::initialize_video(
+void CpuDeviceInterface::initialize_video_decoding(
     const AVStream* av_stream,
     [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx,
+    [[maybe_unused]] const VideoStreamOptions& video_stream_options) {
+  STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
+  time_base_ = av_stream->time_base;
+}
+
+void CpuDeviceInterface::initialize_color_conversion(
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
-  // TODO_API_BREAKDOWN this used to be:
-  // STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
-  // time_base_ = av_stream->time_base;
-  // but now that avStrean can be null (to create a standalone color converter)
-  // we need this workaround. This is bad, we need to preserve the previous
-  // check somehow.
-  time_base_ = (av_stream != nullptr) ? av_stream->time_base
-                                      : AVRational{1, AV_TIME_BASE};
   av_media_type_ = AVMEDIA_TYPE_VIDEO;
   video_stream_options_ = video_stream_options;
   resized_output_dims_ = resized_output_dims;

src/torchcodec/_core/CpuDeviceInterface.h

@@ -27,9 +27,12 @@ class CpuDeviceInterface : public DeviceInterface {
 
   virtual void initialize(const SharedAVCodecContext& codec_context) override;
 
-  virtual void initialize_video(
+  virtual void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  virtual void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;
@@ -76,7 +79,11 @@ class CpuDeviceInterface : public DeviceInterface {
       const FrameDims& output_dims) const;
 
   VideoStreamOptions video_stream_options_;
-  AVRational time_base_;
+  // Default used when color conversion runs standalone (no stream to derive it
+  // from, e.g. the ColorConverter block API). initialize_video_decoding()
+  // overrides it from the stream when there is one. Its value doesn't matter on
+  // color-conversion-only mode, but filtergraph still expects it.
+  AVRational time_base_ = {1, AV_TIME_BASE};
   AVPixelFormat output_pixel_format_;
 
   // If the resized output dimensions are present, then we always use those as

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -114,17 +114,13 @@ void CudaDeviceInterface::initialize(
   codec_context_ = codec_context;
 }
 
-void CudaDeviceInterface::initialize_video(
+void CudaDeviceInterface::initialize_video_decoding(
     const AVStream* av_stream,
     const UniqueDecodingAVFormatContext& av_format_ctx,
-    const VideoStreamOptions& video_stream_options,
-    [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
-    [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {
+    [[maybe_unused]] const VideoStreamOptions& video_stream_options) {
   STD_TORCH_CHECK(av_stream != nullptr, "avStream is null");
   time_base_ = av_stream->time_base;
-  video_stream_options_ = video_stream_options;
 
-  // TODO: Ideally, we should keep all interface implementations independent.
   cpu_interface_ = create_device_interface(kStableCPU);
   STD_TORCH_CHECK(
       cpu_interface_ != nullptr, "Failed to create CPU device interface");
@@ -137,6 +133,13 @@ void CudaDeviceInterface::initialize_video(
       /*resizedOutputDims=*/std::nullopt);
 }
 
+void CudaDeviceInterface::initialize_color_conversion(
+    const VideoStreamOptions& video_stream_options,
+    [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
+    [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {
+  video_stream_options_ = video_stream_options;
+}
+
 void CudaDeviceInterface::register_hardware_device_with_codec(
     AVCodecContext* codec_context) {
   STD_TORCH_CHECK(

src/torchcodec/_core/CudaDeviceInterface.h

@@ -28,14 +28,15 @@ class CudaDeviceInterface : public DeviceInterface {
 
   void initialize(const SharedAVCodecContext& codec_context) override;
 
-  void initialize_video(
+  void initialize_video_decoding(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options) override;
+
+  void initialize_color_conversion(
       const VideoStreamOptions& video_stream_options,
-      [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
-          transforms,
-      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims)
-      override;
+      const std::vector<std::unique_ptr<Transform>>& transforms,
+      const std::optional<FrameDims>& resized_output_dims) override;
 
   void register_hardware_device_with_codec(
       AVCodecContext* codec_context) override;

src/torchcodec/_core/Demuxer.cpp

@@ -69,6 +69,8 @@ Demuxer::Demuxer(
 }
 
 UniqueAVPacket Demuxer::next_packet() {
+  // TODO_API_BREAKDOWN P2: Not a fan of the ReferenceAVPacket / AutoAVPacket /
+  // UniqueAVPacket dance here. Can we simplify?
   ReferenceAVPacket packet(auto_packet_);
   int status =
       read_next_packet(format_context_.get(), active_stream_index_, packet);

src/torchcodec/_core/DeviceInterface.h

@@ -56,15 +56,33 @@ class DeviceInterface {
   // default sendPacket/receiveFrame/flush implementations.
   virtual void initialize(const SharedAVCodecContext& codec_context) = 0;
 
-  // Initialize the device with parameters specific to video decoding. There is
-  // a default empty implementation.
-  virtual void initialize_video(
+  // Initialize state needed to decode packets into raw AVFrames.
+  virtual void initialize_video_decoding(
       [[maybe_unused]] const AVStream* av_stream,
       [[maybe_unused]] const UniqueDecodingAVFormatContext& av_format_ctx,
+      [[maybe_unused]] const VideoStreamOptions& video_stream_options) {}
+
+  // Initialize state needed to color-convert decoded AVFrames into output
+  // tensors.
+  virtual void initialize_color_conversion(
       [[maybe_unused]] const VideoStreamOptions& video_stream_options,
       [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
-          transforms,
-      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {}
+          transforms = {},
+      [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims =
+          std::nullopt) {}
+
+  // Convenience for the combined decode + color-convert path, kept for BC as
+  // it's used by SingleStreamDecoder and out-of-tree interfaces rely on it.
+  void initialize_video(
+      const AVStream* av_stream,
+      const UniqueDecodingAVFormatContext& av_format_ctx,
+      const VideoStreamOptions& video_stream_options,
+      const std::vector<std::unique_ptr<Transform>>& transforms,
+      const std::optional<FrameDims>& resized_output_dims) {
+    initialize_video_decoding(av_stream, av_format_ctx, video_stream_options);
+    initialize_color_conversion(
+        video_stream_options, transforms, resized_output_dims);
+  }
 
   // Initialize the device with parameters specific to audio decoding. There is
   // a default empty implementation.
@@ -138,6 +156,9 @@ class DeviceInterface {
     return avcodec_receive_frame(codec_context_.get(), av_frame.get());
   }
 
+  virtual void make_frame_standalone([[maybe_unused]] UniqueAVFrame& av_frame) {
+  };
+
   // Flush remaining frames from decoder
   virtual void flush() {
     STD_TORCH_CHECK(
@@ -200,7 +221,7 @@ TORCHCODEC_THIRD_PARTY_API bool register_device_interface(
 
 FORCE_PUBLIC_VISIBILITY void validate_device_interface(
     const std::string& device,
-    const std::string& variant);
+    const std::string& variant = "default");
 
 TORCHCODEC_THIRD_PARTY_API std::unique_ptr<DeviceInterface>
 create_device_interface(

src/torchcodec/_core/PacketDecoder.cpp

@@ -8,7 +8,7 @@
 
 namespace facebook::torchcodec {
 
-// TODO_API_BREAKDOWN: we should make sure the block APIs can dispatch to
+// TODO_API_BREAKDOWN P1: we should make sure the block APIs can dispatch to
 // third-party extensions - all of them.
 
 SharedAVCodecContext create_and_open_codec_context(
@@ -58,9 +58,8 @@ const AVCodec* find_decoder(
 PacketDecoder::PacketDecoder(
     const Demuxer& demuxer,
     const StableDevice& device,
-    std::string_view device_variant,
     std::optional<int> ffmpeg_thread_count) {
-  device_interface_ = create_device_interface(device, device_variant);
+  device_interface_ = create_device_interface(device);
   STD_TORCH_CHECK(
       device_interface_ != nullptr,
       "Failed to create device interface. This should never happen, please report.");
@@ -71,6 +70,13 @@ PacketDecoder::PacketDecoder(
   codec_context_ = create_and_open_codec_context(
       stream, av_codec, device_interface_.get(), ffmpeg_thread_count);
   device_interface_->initialize(codec_context_);
+
+  VideoStreamOptions options;
+  options.output_dtype = OutputDtype::UINT8; // dtype not exposed yet
+  options.device = device;
+
+  device_interface_->initialize_video_decoding(
+      stream, demuxer.format_context(), options);
 }
 
 int PacketDecoder::send_packet(AVPacket* packet) {
@@ -90,7 +96,11 @@ int PacketDecoder::send_eof() {
 }
 
 int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) {
-  return device_interface_->receive_frame(av_frame);
+  int status = device_interface_->receive_frame(av_frame);
+  if (status == AVSUCCESS) {
+    device_interface_->make_frame_standalone(av_frame);
+  }
+  return status;
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/PacketDecoder.h

@@ -33,7 +33,6 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder {
   explicit PacketDecoder(
       const Demuxer& demuxer,
       const StableDevice& device = StableDevice(kStableCPU),
-      std::string_view device_variant = "default",
       std::optional<int> ffmpeg_thread_count = std::nullopt);
 
   // Feed one packet to the decoder. Borrows `packet` (does not take ownership).

src/torchcodec/_core/color_conversion.cpp

@@ -191,7 +191,7 @@ void compute_rgb_to_yuv_matrix(
 torch::stable::Tensor convert_yuv_frame_to_rgb(
     const AVFrame& av_frame,
     const StableDevice& device,
-    cudaStream_t nvdec_stream,
+    cudaStream_t producer_stream,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor,
     const FrameDims& output_dims,
     bool is_p016,
@@ -221,9 +221,14 @@ torch::stable::Tensor convert_yuv_frame_to_rgb(
         FrameDims(out_height, out_width), device, out_dtype);
   }
 
+  // TODO_API_BREAKDOWN P1: This may not be the semantic that we want: this will
+  // wait for all ongoin work on the producer stream to finish. But maybe the
+  // producer stream produced the frame data a long time ago, and lots of
+  // kernels have been launched on it already. We'd be waiting on those to
+  // finish even though the data we need is already available.
   cudaStream_t stream = get_current_cuda_stream(device.index());
   sync_streams(
-      /*runningStream=*/nvdec_stream, /*waitingStream=*/stream);
+      /*runningStream=*/producer_stream, /*waitingStream=*/stream);
 
   maybe_update_color_matrix(
       cached_color_matrix,

src/torchcodec/_core/color_conversion.h

@@ -86,7 +86,7 @@ void launch_p016_to_rgb16_kernel(
 torch::stable::Tensor convert_yuv_frame_to_rgb(
     const AVFrame& av_frame,
     const StableDevice& device,
-    cudaStream_t nvdec_stream,
+    cudaStream_t producer_stream,
     std::optional<torch::stable::Tensor> pre_allocated_output_tensor,
     const FrameDims& output_dims,
     bool is_p016,