Decode audio in the Blocks API, as RawAudioSamples

PacketDecoder now follows the demuxer it was built from: RawFrames for a
VideoDemuxer, RawAudioSamples for an AudioDemuxer. Decoding is the same
avcodec_send_packet / avcodec_receive_frame pair either way, so this stays
one block rather than two.

RawAudioSamples carries the samples in the codec's own sample type (uint8
for u8, int16 for s16, float32 for flt...), always as a contiguous
[num_channels, num_samples] tensor. That is a copy, not a view: planar
formats put each channel in its own allocation and packed ones interleave
them, so neither is that shape as it stands, and an audio frame is only a
few kB. Normalizing to [-1, 1] is the converter's job, not the decoder's.

Audio is CPU-only, so a non-CPU device raises rather than being silently
ignored, and an unspecified one doesn't inherit a non-CPU default device.

Author: NicolasHug

src/torchcodec/_core/AudioCommon.cpp

@@ -0,0 +1,38 @@
+// 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 "AudioCommon.h"
+
+namespace facebook::torchcodec {
+
+torch::headeronly::ScalarType sample_format_dtype(
+    AVSampleFormat sample_format) {
+  switch (av_get_packed_sample_fmt(sample_format)) {
+    case AV_SAMPLE_FMT_U8:
+      return kStableUInt8;
+    case AV_SAMPLE_FMT_S16:
+      return kStableInt16;
+    case AV_SAMPLE_FMT_S32:
+      return kStableInt32;
+    case AV_SAMPLE_FMT_S64:
+      return kStableInt64;
+    case AV_SAMPLE_FMT_FLT:
+      return kStableFloat32;
+    case AV_SAMPLE_FMT_DBL:
+      return kStableFloat64;
+    default:
+      break;
+  }
+  const char* name = av_get_sample_fmt_name(sample_format);
+  STD_TORCH_CHECK(
+      false,
+      "Unsupported sample format '",
+      name == nullptr ? "unknown" : name,
+      "'.");
+  return kStableUInt8;
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/AudioCommon.h

@@ -0,0 +1,23 @@
+// 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.
+
+#pragma once
+
+#include "FFMPEGCommon.h"
+#include "StableABICompat.h"
+
+// Where FFmpeg's audio samples meet torch tensors. FFMPEGCommon deliberately
+// knows nothing about tensors, and these helpers are shared by the decode, the
+// conversion and the SingleStreamDecoder paths, so they live on their own.
+
+namespace facebook::torchcodec {
+
+// The dtype that holds `sample_format`'s samples exactly. Planar and packed
+// variants of a format share a sample type, which is why this doesn't care
+// which one it is given.
+torch::headeronly::ScalarType sample_format_dtype(AVSampleFormat sample_format);
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/PacketDecoder.cpp

@@ -6,7 +6,10 @@
 
 #include "PacketDecoder.h"
 
+#include "AudioCommon.h"
+
 #include <algorithm>
+#include <cstring>
 
 namespace facebook::torchcodec {
 
@@ -57,26 +60,49 @@ const AVCodec* find_decoder(
 PacketDecoder::PacketDecoder(
     const Demuxer& demuxer,
     const StableDevice& device,
-    std::optional<int> ffmpeg_thread_count) {
+    std::optional<int> ffmpeg_thread_count)
+    : media_type_(demuxer.media_type()) {
+  bool is_audio = media_type_ == AVMEDIA_TYPE_AUDIO;
+  STD_TORCH_CHECK(
+      !is_audio || device.type() == kStableCPU,
+      "Audio can only be decoded on the CPU.");
+
   device_interface_ = create_device_interface(device);
   STD_TORCH_CHECK(
       device_interface_ != nullptr,
       "Failed to create device interface. This should never happen, please report.");
 
   AVStream* stream = demuxer.active_stream();
   time_base_ = stream->time_base;
-  is_mpeg_ps_ =
-      std::string_view(demuxer.format_context()->iformat->name) == "mpeg";
-  if (const int32_t* matrix = get_display_matrix_from_stream(stream)) {
-    display_matrix_.emplace();
-    std::copy(
-        matrix, matrix + display_matrix_->size(), display_matrix_->begin());
+
+  if (is_audio) {
+    // Audio codecs are hardcoded to a single FFmpeg thread, see
+    // https://github.com/pytorch/torchcodec/issues/1253.
+    ffmpeg_thread_count = 1;
+  } else {
+    is_mpeg_ps_ =
+        std::string_view(demuxer.format_context()->iformat->name) == "mpeg";
+    if (const int32_t* matrix = get_display_matrix_from_stream(stream)) {
+      display_matrix_.emplace();
+      std::copy(
+          matrix, matrix + display_matrix_->size(), display_matrix_->begin());
+    }
   }
+
   const AVCodec* av_codec = find_decoder(stream, device_interface_.get());
   codec_context_ = create_and_open_codec_context(
       stream, av_codec, device_interface_.get(), ffmpeg_thread_count);
   device_interface_->initialize(codec_context_);
 
+  if (is_audio) {
+    // Nothing else to set up: unlike video, we hand out the samples in the
+    // codec's own format, so no conversion state is needed here. Note we
+    // deliberately do NOT set request_sample_fmt: what SingleStreamDecoder
+    // asks for (FLTP) is an optimization for its own conversion, and here it
+    // would hide what the codec natively produces.
+    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;
@@ -133,10 +159,12 @@ int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) {
   int status = device_interface_->receive_frame(av_frame);
   if (status == AVSUCCESS) {
     device_interface_->make_frame_standalone(av_frame);
-    // Attach a copy of the display matrix to the frame, so the ColorConverter
-    // can use it.
-    set_display_matrix_on_frame(
-        *av_frame, display_matrix_ ? display_matrix_->data() : nullptr);
+    if (media_type_ == AVMEDIA_TYPE_VIDEO) {
+      // Attach a copy of the display matrix to the frame, so the ColorConverter
+      // can use it.
+      set_display_matrix_on_frame(
+          *av_frame, display_matrix_ ? display_matrix_->data() : nullptr);
+    }
   }
   return status;
 }
@@ -240,4 +268,74 @@ std::vector<torch::stable::Tensor> frame_planes(
   return planes;
 }
 
+namespace {
+// Scatters `num_channels`-interleaved samples into one contiguous row per
+// channel. Templated on an integer of the right width rather than the actual
+// sample type: we're only moving bytes around, so all that matters is size.
+template <typename T>
+void deinterleave(
+    const uint8_t* src,
+    uint8_t* dst,
+    int num_channels,
+    int num_samples) {
+  const T* in = reinterpret_cast<const T*>(src);
+  T* out = reinterpret_cast<T*>(dst);
+  for (int channel = 0; channel < num_channels; ++channel) {
+    T* row = out + static_cast<int64_t>(channel) * num_samples;
+    for (int sample = 0; sample < num_samples; ++sample) {
+      row[sample] = in[static_cast<int64_t>(sample) * num_channels + channel];
+    }
+  }
+}
+} // namespace
+
+torch::stable::Tensor audio_samples(const AVFrame& av_frame) {
+  auto sample_format = static_cast<AVSampleFormat>(av_frame.format);
+  int num_channels = get_num_channels(av_frame);
+  int64_t num_samples = av_frame.nb_samples;
+
+  torch::stable::Tensor samples = torch::stable::empty(
+      {num_channels, num_samples}, sample_format_dtype(sample_format));
+  if (num_samples == 0) {
+    return samples;
+  }
+
+  int bytes_per_sample = av_get_bytes_per_sample(sample_format);
+  auto* dst = static_cast<uint8_t*>(samples.mutable_data_ptr());
+  int64_t bytes_per_channel = num_samples * bytes_per_sample;
+
+  if (av_sample_fmt_is_planar(sample_format)) {
+    for (int channel = 0; channel < num_channels; ++channel) {
+      // extended_data rather than data: the latter only holds
+      // AV_NUM_DATA_POINTERS (8) pointers, and we support more channels.
+      std::memcpy(
+          dst + channel * bytes_per_channel,
+          av_frame.extended_data[channel],
+          bytes_per_channel);
+    }
+    return samples;
+  }
+
+  const uint8_t* src = av_frame.extended_data[0];
+  int num_samples_int = static_cast<int>(num_samples);
+  switch (bytes_per_sample) {
+    case 1:
+      deinterleave<uint8_t>(src, dst, num_channels, num_samples_int);
+      break;
+    case 2:
+      deinterleave<uint16_t>(src, dst, num_channels, num_samples_int);
+      break;
+    case 4:
+      deinterleave<uint32_t>(src, dst, num_channels, num_samples_int);
+      break;
+    case 8:
+      deinterleave<uint64_t>(src, dst, num_channels, num_samples_int);
+      break;
+    default:
+      STD_TORCH_CHECK(
+          false, "Unexpected sample width: ", bytes_per_sample, " bytes.");
+  }
+  return samples;
+}
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/PacketDecoder.h

@@ -29,8 +29,9 @@ SharedAVCodecContext create_and_open_codec_context(
     DeviceInterface* device_interface,
     std::optional<int> thread_count);
 
-// Decode building block: turns compressed packets into decoded (YUV) frames.
-// Configured from a Demuxer's active stream; stateful. Not thread-safe.
+// Decode building block: turns compressed packets into decoded frames - (YUV)
+// pictures for a video stream, samples in the codec's own format for an audio
+// one. Configured from a Demuxer's active stream; stateful. Not thread-safe.
 class FORCE_PUBLIC_VISIBILITY PacketDecoder {
  public:
   explicit PacketDecoder(
@@ -67,10 +68,15 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder {
     return time_base_;
   }
 
+  AVMediaType media_type() const {
+    return media_type_;
+  }
+
  private:
   std::unique_ptr<DeviceInterface> device_interface_;
   SharedAVCodecContext codec_context_;
   AVRational time_base_ = {};
+  AVMediaType media_type_ = AVMEDIA_TYPE_VIDEO;
   // Stamped onto every frame we hand out, so downstream blocks can read the
   // rotation off the frame itself instead of knowing about the stream. Held by
   // value: we're only handed the Demuxer at construction and it may well be
@@ -111,4 +117,14 @@ FORCE_PUBLIC_VISIBILITY std::vector<torch::stable::Tensor> frame_planes(
     const StableDevice& device,
     const torch::stable::Tensor& tensor_handle);
 
+// A decoded audio frame's samples as a contiguous [num_channels, num_samples]
+// tensor whose dtype is the frame's own sample type: uint8 for u8, int16 for
+// s16, float32 for flt, and so on, planar or not. This is a copy rather than a
+// view: planar formats put each channel in its own allocation and packed ones
+// interleave them, so neither is a [C, N] tensor as it stands. An audio frame
+// is a few kB, so normalizing here buys a uniform layout for the price of a
+// memcpy - and it means a converter can treat the result as planar-of-dtype.
+FORCE_PUBLIC_VISIBILITY torch::stable::Tensor audio_samples(
+    const AVFrame& av_frame);
+
 } // namespace facebook::torchcodec

src/torchcodec/_core/StableABICompat.h

@@ -63,6 +63,7 @@ constexpr auto kStableXPU = torch::headeronly::DeviceType::XPU;
 // Scalar type constants
 constexpr auto kStableUInt8 = torch::headeronly::ScalarType::Byte;
 constexpr auto kStableUInt16 = torch::headeronly::ScalarType::UInt16;
+constexpr auto kStableInt16 = torch::headeronly::ScalarType::Short;
 constexpr auto kStableInt32 = torch::headeronly::ScalarType::Int;
 constexpr auto kStableInt64 = torch::headeronly::ScalarType::Long;
 constexpr auto kStableFloat32 = torch::headeronly::ScalarType::Float;

src/torchcodec/_core/_ffmpeg_op_names.py

@@ -41,6 +41,7 @@
         "_blocks_packet_decoder_send_eof",
         "_blocks_packet_decoder_reset",
         "_blocks_packet_decoder_receive_frame",
+        "_blocks_audio_packet_decoder_receive_frame",
         "_blocks_create_color_converter",
         "_blocks_convert_frame",
         "_blocks_frame_metadata",

src/torchcodec/_core/_ffmpeg_ops.py

@@ -121,6 +121,9 @@ def add_video_stream(
 _blocks_packet_decoder_receive_frame = (
     torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default
 )
+_blocks_audio_packet_decoder_receive_frame = (
+    torch.ops.torchcodec_ns._blocks_audio_packet_decoder_receive_frame.default
+)
 _blocks_create_color_converter = (
     torch.ops.torchcodec_ns._blocks_create_color_converter.default
 )

src/torchcodec/_core/custom_ops.cpp

@@ -91,6 +91,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) {
   m.def("_blocks_packet_decoder_reset(Tensor(a!) decoder) -> ()");
   m.def(
       "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, Device, Tensor)");
+  m.def(
+      "_blocks_audio_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, int, str)");
   m.def(
       "_blocks_create_color_converter(str device=\"cpu\", str output_dtype=\"uint8\") -> Tensor");
   m.def(
@@ -978,6 +980,46 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
       storage);
 }
 
+// (samples, status, pts_seconds, duration_seconds, sample_rate,
+// sample_format). `samples` is [num_channels, num_samples] in the frame's own
+// sample type; there is no frame handle because, unlike a video frame, nothing
+// downstream needs the AVFrame itself.
+using OpsReceiveAudioFrameOutput = std::
+    tuple<torch::stable::Tensor, int64_t, double, double, int64_t, std::string>;
+
+OpsReceiveAudioFrameOutput _blocks_audio_packet_decoder_receive_frame(
+    torch::stable::Tensor& decoder) {
+  PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer<PacketDecoder>(decoder);
+  STD_TORCH_CHECK(
+      decoder_ptr->media_type() == AVMEDIA_TYPE_AUDIO,
+      "This PacketDecoder decodes video, not audio.");
+
+  UniqueAVFrame av_frame(av_frame_alloc());
+  STD_TORCH_CHECK(av_frame != nullptr, "Failed to allocate AVFrame");
+  int status = decoder_ptr->receive_frame(av_frame);
+  if (status != AVSUCCESS) {
+    return std::make_tuple(
+        torch::stable::empty({int64_t(0)}, kStableUInt8),
+        static_cast<int64_t>(status),
+        0.0,
+        0.0,
+        static_cast<int64_t>(0),
+        std::string(""));
+  }
+
+  AVRational time_base = decoder_ptr->time_base();
+  const char* sample_format_name =
+      av_get_sample_fmt_name(static_cast<AVSampleFormat>(av_frame->format));
+  return std::make_tuple(
+      audio_samples(*av_frame),
+      static_cast<int64_t>(0),
+      pts_to_seconds(get_pts_or_dts(*av_frame), time_base),
+      pts_to_seconds(get_duration(*av_frame), time_base),
+      static_cast<int64_t>(av_frame->sample_rate),
+      std::string(
+          sample_format_name == nullptr ? "unknown" : sample_format_name));
+}
+
 torch::stable::Tensor _blocks_create_color_converter(
     std::string device,
     std::string output_dtype) {
@@ -1613,6 +1655,9 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) {
   m.impl(
       "_blocks_packet_decoder_receive_frame",
       TORCH_BOX(&_blocks_packet_decoder_receive_frame));
+  m.impl(
+      "_blocks_audio_packet_decoder_receive_frame",
+      TORCH_BOX(&_blocks_audio_packet_decoder_receive_frame));
   m.impl("_blocks_convert_frame", TORCH_BOX(&_blocks_convert_frame));
   m.impl("_blocks_frame_metadata", TORCH_BOX(&_blocks_frame_metadata));
   m.impl("_blocks_frame_planes", TORCH_BOX(&_blocks_frame_planes));

src/torchcodec/_core/sources.bzl

@@ -33,6 +33,7 @@ decoder_core_sources = [
     "CpuDeviceInterface.cpp",
     "Demuxer.cpp",
     "PacketDecoder.cpp",
+    "AudioCommon.cpp",
     "ColorConverter.cpp",
     "SingleStreamDecoder.cpp",
     "Encoder.cpp",

src/torchcodec/decoders/_blocks/__init__.py

@@ -21,7 +21,7 @@
 
 from ._color_converter import ColorConverter
 from ._demuxer import AudioDemuxer, StreamIndex, VideoDemuxer
-from ._frame import Packet, RawFrame
+from ._frame import Packet, RawAudioSamples, RawFrame
 from ._packet_decoder import PacketDecoder
 
 __all__ = [
@@ -31,6 +31,7 @@
     "ColorConverter",
     "Packet",
     "RawFrame",
+    "RawAudioSamples",
     "StreamIndex",
 ]