CpuDeviceInterface.cpp

// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#include "CpuDeviceInterface.h"
#include <sstream>

namespace facebook::torchcodec {

namespace {

AVPixelFormat validate_pixel_format(
    const AVCodec& av_codec,
    const std::string& target_pixel_format) {
  AVPixelFormat pixel_format = av_get_pix_fmt(target_pixel_format.c_str());
  const AVPixelFormat* supported_formats =
      get_supported_pixel_formats(av_codec);
  if (supported_formats != nullptr) {
    for (int i = 0; supported_formats[i] != AV_PIX_FMT_NONE; ++i) {
      if (supported_formats[i] == pixel_format) {
        return pixel_format;
      }
    }
  }
  std::stringstream error_msg;
  if (pixel_format == AV_PIX_FMT_NONE) {
    error_msg << "Unknown pixel format: " << target_pixel_format;
  } else {
    error_msg << "Specified pixel format " << target_pixel_format
              << " is not supported by the " << av_codec.name << " encoder.";
  }
  error_msg << "\nSupported pixel formats for " << av_codec.name << ":";
  if (supported_formats != nullptr) {
    for (int i = 0; supported_formats[i] != AV_PIX_FMT_NONE; ++i) {
      error_msg << " " << av_get_pix_fmt_name(supported_formats[i]);
    }
  }
  STD_TORCH_CHECK(false, error_msg.str());
}

AVPixelFormat get_output_pixel_format(OutputDtype output_dtype) {
  return output_dtype == OutputDtype::FLOAT32 ? AV_PIX_FMT_RGB48
                                              : AV_PIX_FMT_RGB24;
}

// Returns the format filter string for the given output pixel format.
std::string get_format_filter_string(AVPixelFormat output_format) {
  return (output_format == AV_PIX_FMT_RGB48) ? "format=rgb48,"
                                             : "format=rgb24,";
}

static bool g_cpu = register_device_interface(
    DeviceInterfaceKey(kStableCPU),
    [](const StableDevice& device) { return new CpuDeviceInterface(device); });

} // namespace

AVPixelFormat CpuDeviceInterface::get_encoding_pixel_format(
    const AVCodec& av_codec,
    const std::optional<std::string>& user_pixel_format) const {
  if (user_pixel_format.has_value()) {
    return validate_pixel_format(av_codec, user_pixel_format.value());
  }
  const AVPixelFormat* formats = get_supported_pixel_formats(av_codec);
  // Pick the codec's first supported format (often yuv420p), else fall back.
  return (formats && formats[0] != AV_PIX_FMT_NONE) ? formats[0]
                                                    : AV_PIX_FMT_YUV420P;
}

CpuDeviceInterface::CpuDeviceInterface(const StableDevice& device)
    : DeviceInterface(device) {
  STD_TORCH_CHECK(g_cpu, "CpuDeviceInterface was not registered!");
  STD_TORCH_CHECK(
      device_.type() == kStableCPU, "Unsupported device: must be CPU");
}

void CpuDeviceInterface::initialize(const SharedAVCodecContext& codec_context) {
  codec_context_ = codec_context;
}

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) {
  av_media_type_ = AVMEDIA_TYPE_VIDEO;
  video_stream_options_ = video_stream_options;
  resized_output_dims_ = resized_output_dims;
  output_pixel_format_ =
      get_output_pixel_format(video_stream_options_.output_dtype);

  // We can use swscale when we have a single resize transform.
  // With a single resize, we use swscale twice:
  // first for color conversion (YUV->RGB24), then for resize in RGB24 space.
  //
  // Note that this means swscale will not support the case of having several,
  // back-to-back resizes or other transforms.
  //
  // We calculate this value during initialization but we don't refer to it
  // until getColorConversionLibrary() is called. Calculating this value during
  // initialization saves us from having to save all of the transforms.
  are_transforms_sw_scale_compatible_ = transforms.empty() ||
      (transforms.size() == 1 && transforms[0]->is_resize());

  // Note that we do not expose this capability in the public API, only through
  // the core API.
  //
  // Same as above, we calculate this value during initialization and refer to
  // it in getColorConversionLibrary().
  user_requested_sw_scale_ = video_stream_options_.color_conversion_library ==
      ColorConversionLibrary::SWSCALE;

  // We can only use swscale when we have a single resize transform. Note that
  // we actually decide on whether or not to actually use swscale at the last
  // possible moment, when we actually convert the frame. This is because we
  // need to know the actual frame dimensions.
  if (transforms.size() == 1 && transforms[0]->is_resize()) {
    auto resize = dynamic_cast<ResizeTransform*>(transforms[0].get());
    STD_TORCH_CHECK(
        resize != nullptr, "ResizeTransform expected but not found!");
    sws_flags_ = resize->get_sws_flags();
  }

  // If we have any transforms, replace filters_ with the filter strings from
  // the transforms. As noted above, we decide between swscale and filtergraph
  // when we actually decode a frame.
  std::stringstream filters;
  bool first = true;
  for (const auto& transform : transforms) {
    if (!first) {
      filters << ",";
    }
    filters << transform->get_filter_graph_cpu();
    first = false;
  }
  if (!transforms.empty()) {
    // Note [Transform and Format Conversion Order]
    // We have to ensure that all user filters happen AFTER the explicit format
    // conversion. That is, we want the filters to be applied in the output RGB
    // color space, not the pixel format of the input frame.
    //
    // The output frame will always be in the output RGB format (RGB24 or
    // RGB48), as we specify the sink node with the appropriate format.
    // Filtergraph will automatically insert a format conversion to ensure the
    // output frame matches the pixel format specified in the sink. But by
    // default, it will insert it after the user filters. We need an explicit
    // format conversion to get the behavior we want.
    filters_ = get_format_filter_string(output_pixel_format_) + filters.str();
  }

  initialized_ = true;
}

void CpuDeviceInterface::initialize_audio(
    const AudioStreamOptions& audio_stream_options) {
  av_media_type_ = AVMEDIA_TYPE_AUDIO;
  audio_stream_options_ = audio_stream_options;
  initialized_ = true;
}

ColorConversionLibrary CpuDeviceInterface::get_color_conversion_library(
    const FrameDims& input_dims,
    const FrameDims& output_dims) const {
  // swscale requires widths to be multiples of 32:
  // https://stackoverflow.com/questions/74351955/turn-off-sw-scale-conversion-to-planar-yuv-32-byte-alignment-requirements
  bool are_widths_sw_scale_compatible =
      (input_dims.width % 32) == 0 && (output_dims.width % 32) == 0;

  // We want to use swscale for color conversion if possible because it is
  // faster than filtergraph. The following are the conditions we need to meet
  // to use it.
  //
  // Note that we treat the transform limitation differently from the width
  // limitation. That is, we consider the transforms being compatible with
  // swscale as a hard requirement. If the transforms are not compatiable,
  // then we will end up not applying the transforms, and that is wrong.
  //
  // The width requirement, however, is a soft requirement. Even if we don't
  // meet it, we let the user override it. We have tests that depend on this
  // behavior. Since we don't expose the ability to choose swscale or
  // filtergraph in our public API, this is probably okay. It's also the only
  // way that we can be certain we are testing one versus the other.
  if (are_transforms_sw_scale_compatible_ &&
      (user_requested_sw_scale_ || are_widths_sw_scale_compatible)) {
    return ColorConversionLibrary::SWSCALE;
  } else {
    return ColorConversionLibrary::FILTERGRAPH;
  }
}

void CpuDeviceInterface::convert_av_frame_to_frame_output(
    const AVFrame& av_frame,
    FrameOutput& frame_output,
    std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
  STD_TORCH_CHECK(initialized_, "CpuDeviceInterface was not initialized.");

  if (av_media_type_ == AVMEDIA_TYPE_AUDIO) {
    convert_audio_av_frame_to_frame_output(av_frame, frame_output);
  } else {
    convert_video_av_frame_to_frame_output(
        av_frame, frame_output, pre_allocated_output_tensor);
  }
}

// Note [preAllocatedOutputTensor with swscale and filtergraph]:
// Callers may pass a pre-allocated tensor, where the output.data tensor will
// be stored. This parameter is honored in any case, but it only leads to a
// speed-up when swscale is used. With swscale, we can tell ffmpeg to place the
// decoded frame directly into `preAllocatedtensor.mutable_data_ptr()`. We
// haven't yet found a way to do that with filtegraph.
// TODO: Figure out whether that's possible!
// Dimension order of the preAllocatedOutputTensor must be HWC, regardless of
// `dimension_order` parameter. It's up to callers to re-shape it if needed.
void CpuDeviceInterface::convert_video_av_frame_to_frame_output(
    const AVFrame& av_frame,
    FrameOutput& frame_output,
    std::optional<torch::stable::Tensor> pre_allocated_output_tensor) {
  // Note that we ignore the dimensions from the metadata; we don't even bother
  // storing them. The resized dimensions take priority. If we don't have any,
  // then we use the dimensions from the actual decoded frame. We use the actual
  // decoded frame and not the metadata for two reasons:
  //
  //   1. Metadata may be wrong. If we access to more accurate information, we
  //      should use it.
  //   2. Video streams can have variable resolution. This fact is not captured
  //      in the stream  metadata.
  //
  // Both cases cause problems for our batch APIs, as we allocate
  // FrameBatchOutputs based on the the stream metadata. But single-frame APIs
  // can still work in such situations, so they should.
  auto input_dims = FrameDims(av_frame.height, av_frame.width);
  auto output_dims = resized_output_dims_.value_or(input_dims);

  if (pre_allocated_output_tensor.has_value()) {
    auto shape = pre_allocated_output_tensor.value().sizes();
    STD_TORCH_CHECK(
        (shape.size() == 3) && (shape[0] == output_dims.height) &&
            (shape[1] == output_dims.width) && (shape[2] == 3),
        "Expected pre-allocated tensor of shape ",
        output_dims.height,
        "x",
        output_dims.width,
        "x3, got ",
        int_array_ref_to_string(shape));
  }

  auto color_conversion_library =
      get_color_conversion_library(input_dims, output_dims);
  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));

    auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);
    SwsConfig sws_config(
        av_frame.width,
        av_frame.height,
        av_frame_format,
        av_frame.colorspace,
        output_dims.width,
        output_dims.height,
        output_pixel_format_);

    if (!sw_scale_ || sw_scale_->get_config() != sws_config) {
      sw_scale_ = std::make_unique<SwScale>(sws_config, sws_flags_);
    }

    int result_height = sw_scale_->convert(av_frame, output_tensor);

    // If this check failed, it would mean that the frame wasn't reshaped to
    // the expected height.
    // TODO: Can we do the same check for width?
    STD_TORCH_CHECK(
        result_height == output_dims.height,
        "resultHeight != outputDims.height: ",
        result_height,
        " != ",
        output_dims.height);

    frame_output.data = output_tensor;
  } else if (color_conversion_library == ColorConversionLibrary::FILTERGRAPH) {
    output_tensor =
        convert_av_frame_to_tensor_using_filter_graph(av_frame, output_dims);

    // Similarly to above, if this check fails it means the frame wasn't
    // reshaped to its expected dimensions by filtergraph.
    auto shape = output_tensor.sizes();
    STD_TORCH_CHECK(
        (shape.size() == 3) && (shape[0] == output_dims.height) &&
            (shape[1] == output_dims.width) && (shape[2] == 3),
        "Expected output tensor of shape ",
        output_dims.height,
        "x",
        output_dims.width,
        "x3, got ",
        int_array_ref_to_string(shape));

    if (pre_allocated_output_tensor.has_value()) {
      // We have already validated that preAllocatedOutputTensor and
      // outputTensor have the same shape.
      torch::stable::copy_(pre_allocated_output_tensor.value(), output_tensor);
      frame_output.data = pre_allocated_output_tensor.value();
    } else {
      frame_output.data = output_tensor;
    }
  } else {
    STD_TORCH_CHECK(
        false,
        "Invalid color conversion library: ",
        static_cast<int>(color_conversion_library));
  }
}

torch::stable::Tensor
CpuDeviceInterface::convert_av_frame_to_tensor_using_filter_graph(
    const AVFrame& av_frame,
    const FrameDims& output_dims) {
  auto av_frame_format = static_cast<AVPixelFormat>(av_frame.format);

  FiltersConfig filters_config(
      av_frame.width,
      av_frame.height,
      av_frame_format,
      av_frame.sample_aspect_ratio,
      output_dims.width,
      output_dims.height,
      output_pixel_format_,
      filters_,
      time_base_);

  if (!filter_graph_ || prev_filters_config_ != filters_config) {
    filter_graph_ =
        std::make_unique<FilterGraph>(filters_config, video_stream_options_);
    prev_filters_config_ = std::move(filters_config);
  }
  return rgb_av_frame_to_tensor(*filter_graph_->convert(av_frame));
}

void CpuDeviceInterface::convert_audio_av_frame_to_frame_output(
    const AVFrame& src_av_frame,
    FrameOutput& frame_output) {
  AVSampleFormat src_sample_format =
      static_cast<AVSampleFormat>(src_av_frame.format);
  AVSampleFormat out_sample_format = AV_SAMPLE_FMT_FLTP;

  int src_sample_rate = src_av_frame.sample_rate;
  int out_sample_rate =
      audio_stream_options_.sample_rate.value_or(src_sample_rate);

  int src_num_channels = get_num_channels(codec_context_);
  STD_TORCH_CHECK(
      src_num_channels == get_num_channels(src_av_frame),
      "The frame has ",
      get_num_channels(src_av_frame),
      " channels, expected ",
      src_num_channels,
      ". If you are hitting this, it may be because you are using "
      "a buggy FFmpeg version. FFmpeg4 is known to fail here in some "
      "valid scenarios. Try to upgrade FFmpeg?");
  int out_num_channels =
      audio_stream_options_.num_channels.value_or(src_num_channels);

  bool must_convert =
      (src_sample_format != out_sample_format ||
       src_sample_rate != out_sample_rate ||
       src_num_channels != out_num_channels);

  UniqueAVFrame converted_av_frame;
  if (must_convert) {
    if (!swr_context_) {
      swr_context_.reset(create_swr_context(
          src_sample_format,
          out_sample_format,
          src_sample_rate,
          out_sample_rate,
          src_av_frame,
          out_num_channels));
    }

    converted_av_frame = convert_audio_av_frame_samples(
        swr_context_,
        src_av_frame,
        out_sample_format,
        out_sample_rate,
        out_num_channels);
  }
  const AVFrame& av_frame = must_convert ? *converted_av_frame : src_av_frame;

  AVSampleFormat format = static_cast<AVSampleFormat>(av_frame.format);
  STD_TORCH_CHECK(
      format == out_sample_format,
      "Something went wrong, the frame didn't get converted to the desired format. ",
      "Desired format = ",
      av_get_sample_fmt_name(out_sample_format),
      "source format = ",
      av_get_sample_fmt_name(format));

  int num_channels = get_num_channels(av_frame);
  STD_TORCH_CHECK(
      num_channels == out_num_channels,
      "Something went wrong, the frame didn't get converted to the desired ",
      "number of channels = ",
      out_num_channels,
      ". Got ",
      num_channels,
      " instead.");

  auto num_samples = av_frame.nb_samples;

  frame_output.data = torch::stable::empty({num_channels, num_samples});

  if (num_samples > 0) {
    uint8_t* output_channel_data =
        static_cast<uint8_t*>(frame_output.data.mutable_data_ptr());
    auto num_bytes_per_channel = num_samples * av_get_bytes_per_sample(format);
    for (auto channel = 0; channel < num_channels;
         ++channel, output_channel_data += num_bytes_per_channel) {
      std::memcpy(
          output_channel_data,
          av_frame.extended_data[channel],
          num_bytes_per_channel);
    }
  }
}

std::optional<torch::stable::Tensor>
CpuDeviceInterface::maybe_flush_audio_buffers() {
  // When sample rate conversion is involved, swresample buffers some of the
  // samples in-between calls to swr_convert (see the libswresample docs).
  // That's because the last few samples in a given frame require future
  // samples from the next frame to be properly converted. This function
  // flushes out the samples that are stored in swresample's buffers.
  if (!swr_context_) {
    return std::nullopt;
  }
  auto num_remaining_samples = // this is an upper bound
      swr_get_out_samples(swr_context_.get(), 0);

  if (num_remaining_samples == 0) {
    return std::nullopt;
  }

  int num_channels = audio_stream_options_.num_channels.value_or(
      get_num_channels(codec_context_));
  torch::stable::Tensor last_samples =
      torch::stable::empty({num_channels, num_remaining_samples});

  std::vector<uint8_t*> output_buffers(num_channels);
  for (auto i = 0; i < num_channels; i++) {
    output_buffers[i] = reinterpret_cast<uint8_t*>(
        select_row(last_samples, i).mutable_data_ptr<float>());
  }

  auto actual_num_remaining_samples = swr_convert(
      swr_context_.get(),
      output_buffers.data(),
      num_remaining_samples,
      nullptr,
      0);

  return torch::stable::narrow(
      last_samples,
      /*dim=*/1,
      /*start=*/0,
      /*length=*/actual_num_remaining_samples);
}

std::string CpuDeviceInterface::get_details() {
  return std::string("CPU Device Interface.");
}

UniqueAVFrame CpuDeviceInterface::convert_tensor_to_av_frame_for_encoding(
    const torch::stable::Tensor& frame,
    int frame_index,
    AVCodecContext* codec_context) {
  int in_height = static_cast<int>(frame.sizes()[1]);
  int in_width = static_cast<int>(frame.sizes()[2]);
  AVPixelFormat in_pixel_format = AV_PIX_FMT_GBRP;
  int out_width = codec_context->width;
  int out_height = codec_context->height;
  AVPixelFormat out_pixel_format = codec_context->pix_fmt;

  // Initialize and cache scaling context if it does not exist
  if (!encoding_sws_context_) {
    encoding_sws_context_.reset(sws_getContext(
        in_width,
        in_height,
        in_pixel_format,
        out_width,
        out_height,
        out_pixel_format,
        SWS_BICUBIC, // Used by FFmpeg CLI
        nullptr,
        nullptr,
        nullptr));
    STD_TORCH_CHECK(
        encoding_sws_context_ != nullptr, "Failed to create scaling context");
  }

  UniqueAVFrame av_frame(av_frame_alloc());
  STD_TORCH_CHECK(av_frame != nullptr, "Failed to allocate AVFrame");

  // Set output frame properties
  av_frame->format = out_pixel_format;
  av_frame->width = out_width;
  av_frame->height = out_height;
  av_frame->pts = frame_index;

  int status = av_frame_get_buffer(av_frame.get(), 0);
  STD_TORCH_CHECK(status >= 0, "Failed to allocate frame buffer");

  // Need to convert/scale the frame
  // Create temporary frame with input format
  UniqueAVFrame input_frame(av_frame_alloc());
  STD_TORCH_CHECK(input_frame != nullptr, "Failed to allocate input AVFrame");

  input_frame->format = in_pixel_format;
  input_frame->width = in_width;
  input_frame->height = in_height;

  uint8_t* tensor_data = static_cast<uint8_t*>(frame.mutable_data_ptr());

  int channel_size = in_height * in_width;
  // Since frames tensor is in NCHW, we must use a planar format.
  // FFmpeg only provides AV_PIX_FMT_GBRP for planar RGB,
  // so we reorder RGB -> GBR.
  input_frame->data[0] = tensor_data + channel_size;
  input_frame->data[1] = tensor_data + (2 * channel_size);
  input_frame->data[2] = tensor_data;

  input_frame->linesize[0] = in_width;
  input_frame->linesize[1] = in_width;
  input_frame->linesize[2] = in_width;

  status = sws_scale(
      encoding_sws_context_.get(),
      input_frame->data,
      input_frame->linesize,
      0,
      input_frame->height,
      av_frame->data,
      av_frame->linesize);
  STD_TORCH_CHECK(status == out_height, "sws_scale failed");
  return av_frame;
}

} // namespace facebook::torchcodec