Make the output dtype config the only dtype on VideoStreamOptions

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -71,26 +71,43 @@ static DecoderCapsCache& get_decoder_caps_cache() {
   return cache;
 }
 
+// Which of NVDEC's two surface widths to decode into. The source's own depth is
+// only one of the options: NVDEC will just as happily put a 10-bit source on an
+// 8-bit surface, or an 8-bit source on a 16-bit one.
+enum class SurfaceDepth { EIGHT_BIT, SIXTEEN_BIT, MATCH_SOURCE };
+
+// A surface only has to be wide enough for what the frame will be converted to,
+// so the requested output dtype is what decides its width. Note this is the
+// dtype *config* rather than a resolved OutputDtype: resolving AUTO answers
+// "which dtype for this source", and what we need here is the question it was
+// resolved from, "should the source's own depth decide". A caller that doesn't
+// color-convert at all has no output dtype to honor and wants AUTO too.
+SurfaceDepth surface_depth_for(OutputDtypeConfig output_dtype_config) {
+  switch (output_dtype_config) {
+    case OutputDtypeConfig::UINT8:
+      return SurfaceDepth::EIGHT_BIT;
+    case OutputDtypeConfig::FLOAT32:
+      return SurfaceDepth::SIXTEEN_BIT;
+    case OutputDtypeConfig::AUTO:
+      return SurfaceDepth::MATCH_SOURCE;
+  }
+  STD_TORCH_CHECK(
+      false,
+      "Unexpected output dtype config. This should never happen, please report.");
+}
+
 // NVDEC's output surface formats come in a 4:2:0 and a 4:4:4 flavour, each with
-// an 8-bit and a 16-bit variant. We decode on the surface that respects the
-// source chroma, but we don't respect the source bit depth and instead try to
-// honor the user's requested output dtype:
-// - if the user wants uint8 output, we try to decode on a uint8 surface,
-//   including for >8bit sources. It's not always supported by NVDEC, so the
-//   caller must fallback to the >8bit surface in such case.
-// - similarly if the user wants float32 output, we try to decode on a >8bit
-//   surface, including for 8bit sources. The caller must handle a similar
-//   fallback.
-cudaVideoSurfaceFormat get_preferred_surface_format(
+// an 8-bit and a 16-bit variant. We always decode on the surface that respects
+// the source chroma, so those two inputs name exactly one format.
+cudaVideoSurfaceFormat surface_format_for(
     cudaVideoChromaFormat chroma_format,
-    OutputDtype output_dtype) {
-  bool want_uint8 = output_dtype == OutputDtype::UINT8;
+    bool want_8_bits) {
   if (chroma_format == cudaVideoChromaFormat_444) {
-    return want_uint8 ? cudaVideoSurfaceFormat_YUV444
-                      : cudaVideoSurfaceFormat_YUV444_16Bit;
+    return want_8_bits ? cudaVideoSurfaceFormat_YUV444
+                       : cudaVideoSurfaceFormat_YUV444_16Bit;
   } else {
-    return want_uint8 ? cudaVideoSurfaceFormat_NV12
-                      : cudaVideoSurfaceFormat_P016;
+    return want_8_bits ? cudaVideoSurfaceFormat_NV12
+                       : cudaVideoSurfaceFormat_P016;
   }
 }
 
@@ -254,7 +271,7 @@ std::optional<cudaVideoCodec> validate_codec_support(AVCodecID codec_id) {
 std::optional<cudaVideoSurfaceFormat> get_nvdec_surface_format(
     const StableDevice& device,
     const SharedAVCodecContext& codec_context,
-    OutputDtype output_dtype) {
+    SurfaceDepth surface_depth) {
   // Return the surface format to use for NVDEC decoding if the stream is
   // supported, or nullopt to fall back to CPU.
 
@@ -301,39 +318,30 @@ std::optional<cudaVideoSurfaceFormat> get_nvdec_surface_format(
     return std::nullopt;
   }
 
-  auto preferred_format =
-      get_preferred_surface_format(chroma_format.value(), output_dtype);
+  bool source_is_8_bits = bit_depth_minus8 == 0;
+  bool want_8_bits = surface_depth == SurfaceDepth::MATCH_SOURCE
+      ? source_is_8_bits
+      : surface_depth == SurfaceDepth::EIGHT_BIT;
 
   auto is_supported = [&](cudaVideoSurfaceFormat format) {
     return ((caps.nOutputFormatMask >> format) & 1) != 0;
   };
 
+  auto preferred_format =
+      surface_format_for(chroma_format.value(), want_8_bits);
   if (is_supported(preferred_format)) {
     return preferred_format;
   }
 
-  // The preferred_format heuristic tries to take a shortcut that might cause us
-  // to miss valid formats. We fallabck here:
-  // if source is 8bit we can try the 8bit surface.
-  // if surface is 8bit we can try the 16bit surface.
-
-  bool source_is_8_bits = bit_depth_minus8 == 0;
-  if (is_16bit_surface_format(preferred_format) && source_is_8_bits) {
-    auto narrower = preferred_format == cudaVideoSurfaceFormat_YUV444_16Bit
-        ? cudaVideoSurfaceFormat_YUV444
-        : cudaVideoSurfaceFormat_NV12;
-
-    if (is_supported(narrower)) {
-      return narrower;
-    }
-  }
-  if (!is_16bit_surface_format(preferred_format)) {
-    auto wider = preferred_format == cudaVideoSurfaceFormat_YUV444
-        ? cudaVideoSurfaceFormat_YUV444_16Bit
-        : cudaVideoSurfaceFormat_P016;
-
-    if (is_supported(wider)) {
-      return wider;
+  // NVDEC doesn't support the width we asked for. The other one still describes
+  // the samples validly - widening leaves the low bits zeroed, and narrowing
+  // only drops bits an 8-bit source never had - so it is worth a try. Note this
+  // rules out narrowing a source that does have those bits.
+  if (want_8_bits || source_is_8_bits) {
+    auto fallback_format =
+        surface_format_for(chroma_format.value(), !want_8_bits);
+    if (is_supported(fallback_format)) {
+      return fallback_format;
     }
   }
 
@@ -406,10 +414,10 @@ BetaCudaDeviceInterface::Mode BetaCudaDeviceInterface::mode() const {
 }
 
 void BetaCudaDeviceInterface::initialize_color_conversion(
-    const VideoStreamOptions& video_stream_options,
+    [[maybe_unused]] OutputDtype output_dtype,
+    [[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) {
-  output_dtype_ = video_stream_options.output_dtype;
   color_conversion_initialized_ = true;
 }
 
@@ -421,10 +429,12 @@ void BetaCudaDeviceInterface::initialize_video_decoding(
   CudaContextGuard context_guard(device_.index());
   decoding_initialized_ = true;
   rotation_ = rotation_from_degrees(get_rotation_from_stream(av_stream));
-  output_dtype_ = video_stream_options.output_dtype;
+
+  SurfaceDepth surface_depth =
+      surface_depth_for(video_stream_options.output_dtype_config);
 
   auto maybe_surface_format = nvcuvid_available_
-      ? get_nvdec_surface_format(device_, codec_context_, output_dtype_)
+      ? get_nvdec_surface_format(device_, codec_context_, surface_depth)
       : std::nullopt;
 
   if (!maybe_surface_format.has_value()) {

src/torchcodec/_core/BetaCudaDeviceInterface.h

@@ -61,6 +61,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
       const VideoStreamOptions& video_stream_options) override;
 
   void initialize_color_conversion(
+      OutputDtype output_dtype,
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;
@@ -177,7 +178,6 @@ class BetaCudaDeviceInterface : public DeviceInterface {
 
   SwsConfig prev_sws_config_;
   Rotation rotation_ = Rotation::NONE;
-  OutputDtype output_dtype_ = OutputDtype::UINT8;
   cudaVideoSurfaceFormat surface_format_ = cudaVideoSurfaceFormat_NV12;
 
   CachedColorMatrix cached_color_matrix_;

src/torchcodec/_core/ColorConverter.cpp

@@ -47,12 +47,14 @@ void ColorConverter::maybe_initialize_interface(OutputDtype output_dtype) {
   }
 
   VideoStreamOptions options;
-  options.output_dtype = output_dtype;
   options.device = device_;
 
   std::vector<std::unique_ptr<Transform>> no_transforms;
   device_interface_->initialize_color_conversion(
-      options, no_transforms, /*resized_output_dims=*/std::nullopt);
+      output_dtype,
+      options,
+      no_transforms,
+      /*resized_output_dims=*/std::nullopt);
   initialized_output_dtype_ = output_dtype;
 }
 

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -94,14 +94,15 @@ void CpuDeviceInterface::initialize_video_decoding(
 }
 
 void CpuDeviceInterface::initialize_color_conversion(
+    OutputDtype output_dtype,
     const VideoStreamOptions& video_stream_options,
     const std::vector<std::unique_ptr<Transform>>& transforms,
     const std::optional<FrameDims>& resized_output_dims) {
   av_media_type_ = AVMEDIA_TYPE_VIDEO;
   video_stream_options_ = video_stream_options;
+  output_dtype_ = output_dtype;
   resized_output_dims_ = resized_output_dims;
-  output_pixel_format_ =
-      get_output_pixel_format(video_stream_options_.output_dtype);
+  output_pixel_format_ = get_output_pixel_format(output_dtype_);
 
   // We can use swscale when we have a single resize transform.
   // With a single resize, we use swscale twice:
@@ -265,9 +266,8 @@ void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
   torch::stable::Tensor output_tensor;
 
   if (color_conversion_library == ColorConversionLibrary::SWSCALE) {
-    output_tensor =
-        pre_allocated_output_tensor.value_or(allocate_empty_hwc_tensor(
-            output_dims, kStableCPU, video_stream_options_.output_dtype));
+    output_tensor = pre_allocated_output_tensor.value_or(
+        allocate_empty_hwc_tensor(output_dims, kStableCPU, output_dtype_));
 
     auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
     SwsConfig sws_config(

src/torchcodec/_core/CpuDeviceInterface.h

@@ -33,6 +33,7 @@ class CpuDeviceInterface : public DeviceInterface {
       const VideoStreamOptions& video_stream_options) override;
 
   virtual void initialize_color_conversion(
+      OutputDtype output_dtype,
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;
@@ -82,6 +83,9 @@ class CpuDeviceInterface : public DeviceInterface {
       const FrameDims& output_dims) const;
 
   VideoStreamOptions video_stream_options_;
+  // Resolved against the source by whoever set up this conversion, so it isn't
+  // on video_stream_options_.
+  OutputDtype output_dtype_ = OutputDtype::UINT8;
   // 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

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -128,12 +128,14 @@ void CudaDeviceInterface::initialize_video_decoding(
   cpu_interface_->initialize_video(
       av_stream,
       av_format_ctx,
+      OutputDtype::UINT8,
       VideoStreamOptions(),
       {},
       /*resizedOutputDims=*/std::nullopt);
 }
 
 void CudaDeviceInterface::initialize_color_conversion(
+    [[maybe_unused]] OutputDtype output_dtype,
     const VideoStreamOptions& video_stream_options,
     [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>& transforms,
     [[maybe_unused]] const std::optional<FrameDims>& resized_output_dims) {

src/torchcodec/_core/CudaDeviceInterface.h

@@ -34,6 +34,7 @@ class CudaDeviceInterface : public DeviceInterface {
       const VideoStreamOptions& video_stream_options) override;
 
   void initialize_color_conversion(
+      OutputDtype output_dtype,
       const VideoStreamOptions& video_stream_options,
       const std::vector<std::unique_ptr<Transform>>& transforms,
       const std::optional<FrameDims>& resized_output_dims) override;

src/torchcodec/_core/DeviceInterface.h

@@ -63,8 +63,11 @@ class DeviceInterface {
       [[maybe_unused]] const VideoStreamOptions& video_stream_options) {}
 
   // Initialize state needed to color-convert decoded AVFrames into output
-  // tensors.
+  // tensors. `output_dtype` is what the frames are converted to, already
+  // resolved against the source: it is not read off `video_stream_options`,
+  // which only carries the config it was resolved from.
   virtual void initialize_color_conversion(
+      [[maybe_unused]] OutputDtype output_dtype,
       [[maybe_unused]] const VideoStreamOptions& video_stream_options,
       [[maybe_unused]] const std::vector<std::unique_ptr<Transform>>&
           transforms = {},
@@ -76,12 +79,13 @@ class DeviceInterface {
   void initialize_video(
       const AVStream* av_stream,
       const UniqueDecodingAVFormatContext& av_format_ctx,
+      OutputDtype output_dtype,
       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);
+        output_dtype, video_stream_options, transforms, resized_output_dims);
   }
 
   // Initialize the device with parameters specific to audio decoding. There is

src/torchcodec/_core/PacketDecoder.cpp

@@ -105,17 +105,12 @@ PacketDecoder::PacketDecoder(
     return;
   }
 
-  const AVPixFmtDescriptor* stream_desc =
-      av_pix_fmt_desc_get(codec_context_->pix_fmt);
-  int stream_bit_depth = stream_desc ? stream_desc->comp[0].depth : 8;
-
   VideoStreamOptions options;
   options.device = device;
-  // This is ugly: what we actually mean is "let the device interface decode
-  // into the native surface", which matters for NVDEC.
-  // TODO_API_BREAKDOWN CC P1: Find a cleaner way to express this?
-  options.output_dtype =
-      stream_bit_depth > 8 ? OutputDtype::FLOAT32 : OutputDtype::UINT8;
+  // We hand out the decoder's own samples and leave color conversion to a
+  // separate block, so there is no output dtype to size a CUDA decode's surface
+  // for. AUTO is what asks for the source's own depth.
+  options.output_dtype_config = OutputDtypeConfig::AUTO;
 
   device_interface_->initialize_video_decoding(
       stream, demuxer.format_context(), options);

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -436,7 +436,7 @@ void SingleStreamDecoder::add_video_stream(
         active_stream_index_, custom_frame_mappings.value());
   }
 
-  stream_info.video_stream_options.output_dtype = resolve_output_dtype(
+  stream_info.output_dtype = resolve_output_dtype(
       stream_info.video_stream_options.output_dtype_config,
       static_cast<AVPixelFormat>(stream_info.stream->codecpar->format));
 
@@ -485,11 +485,10 @@ void SingleStreamDecoder::add_video_stream(
     transforms_.push_back(std::unique_ptr<Transform>(transform));
   }
 
-  // Pass the resolved options (AUTO -> UINT8/FLOAT32) so the device interface
-  // sees a definite OutputDtype.
   device_interface_->initialize_video(
       stream_info.stream,
       format_context_,
+      stream_info.output_dtype,
       stream_info.video_stream_options,
       transforms_,
       resized_output_dims_);
@@ -628,8 +627,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_at_indices(
       frame_indices.numel(),
       get_output_dims(),
       video_stream_options.device,
-      device_interface_->get_pre_allocation_dtype(
-          video_stream_options.output_dtype));
+      device_interface_->get_pre_allocation_dtype(stream_info.output_dtype));
 
   auto frame_batch_output_pts_seconds =
       mutable_accessor<double, 1>(frame_batch_output.pts_seconds);
@@ -701,8 +699,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_in_range(
       num_output_frames,
       get_output_dims(),
       video_stream_options.device,
-      device_interface_->get_pre_allocation_dtype(
-          video_stream_options.output_dtype));
+      device_interface_->get_pre_allocation_dtype(stream_info.output_dtype));
 
   auto frame_batch_output_pts_seconds =
       mutable_accessor<double, 1>(frame_batch_output.pts_seconds);
@@ -844,8 +841,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
         0,
         get_output_dims(),
         video_stream_options.device,
-        device_interface_->get_pre_allocation_dtype(
-            video_stream_options.output_dtype));
+        device_interface_->get_pre_allocation_dtype(stream_info.output_dtype));
     frame_batch_output.data =
         maybe_permute_and_convert_dtype(frame_batch_output.data);
     return frame_batch_output;
@@ -892,8 +888,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
         num_output_frames,
         get_output_dims(),
         video_stream_options.device,
-        device_interface_->get_pre_allocation_dtype(
-            video_stream_options.output_dtype));
+        device_interface_->get_pre_allocation_dtype(stream_info.output_dtype));
 
     auto frame_batch_output_pts_seconds =
         mutable_accessor<double, 1>(frame_batch_output.pts_seconds);
@@ -945,8 +940,7 @@ FrameBatchOutput SingleStreamDecoder::get_frames_played_in_range(
         num_frames,
         get_output_dims(),
         video_stream_options.device,
-        device_interface_->get_pre_allocation_dtype(
-            video_stream_options.output_dtype));
+        device_interface_->get_pre_allocation_dtype(stream_info.output_dtype));
     auto frame_batch_output_pts_seconds =
         mutable_accessor<double, 1>(frame_batch_output.pts_seconds);
     auto frame_batch_output_duration_seconds =
@@ -1576,8 +1570,7 @@ torch::stable::Tensor SingleStreamDecoder::maybe_permute_and_convert_dtype(
   }
 
   return convert_to_output_dtype(
-      tensor,
-      stream_infos_[active_stream_index_].video_stream_options.output_dtype);
+      tensor, stream_infos_[active_stream_index_].output_dtype);
 }
 
 // --------------------------------------------------------------------------

src/torchcodec/_core/SingleStreamDecoder.h

@@ -245,6 +245,10 @@ class FORCE_PUBLIC_VISIBILITY SingleStreamDecoder {
 
     VideoStreamOptions video_stream_options;
     AudioStreamOptions audio_stream_options;
+    // video_stream_options.output_dtype_config resolved against this stream's
+    // pixel format, which we only know once the stream is added. Decoder state
+    // rather than an option: it is worked out here, not asked for.
+    OutputDtype output_dtype = OutputDtype::UINT8;
   };
 
   // --------------------------------------------------------------------------