Define CpuDeviceInterface

Fixes: #619

* Commit defines `CpuDeviceInterface` and moves video `*OnCPU` methods
  from `SingleStreamDecoder` to it.
* Audio `*OnCPU` methods left in `SingleStreamDecoder`
* Constructor API of `DeviceInterface` was changed to allow passing
  `AVRational timeBase` required to initialize ffmpeg filter graph

Signed-off-by: Dmitry Rogozhkin <[email protected]>

Author: dvrogozh

src/torchcodec/_core/CMakeLists.txt

@@ -60,7 +60,8 @@ function(make_torchcodec_libraries
     set(decoder_sources
         AVIOContextHolder.cpp
         FFMPEGCommon.cpp
-	DeviceInterface.cpp
+        DeviceInterface.cpp
+        CpuDeviceInterface.cpp
         SingleStreamDecoder.cpp
         # TODO: lib name should probably not be "*_decoder*" now that it also
         # contains an encoder

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -0,0 +1,23 @@
+#include "src/torchcodec/_core/CpuDeviceInterface.h"
+
+namespace facebook::torchcodec {
+namespace {
+
+bool g_cpu = registerDeviceInterface(
+    torch::kCPU,
+    [](const torch::Device& device, const AVRational& timeBase) {
+      return new CpuDeviceInterface(device, timeBase);
+    });
+
+} // namespace
+
+CpuDeviceInterface::CpuDeviceInterface(
+    const torch::Device& device,
+    const AVRational& timeBase)
+    : DeviceInterface(device, timeBase) {
+  if (device_.type() != torch::kCPU) {
+    throw std::runtime_error("Unsupported device: " + device_.str());
+  }
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/CpuDeviceInterface.h

@@ -0,0 +1,77 @@
+// 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 "src/torchcodec/_core/DeviceInterface.h"
+
+namespace facebook::torchcodec {
+
+class CpuDeviceInterface : public DeviceInterface {
+ public:
+  CpuDeviceInterface(const torch::Device& device, const AVRational& timeBase);
+
+  virtual ~CpuDeviceInterface() {}
+
+  std::optional<const AVCodec*> findCodec(
+      [[maybe_unused]] const AVCodecID& codecId) override {
+    return std::nullopt;
+  }
+
+  void initializeContext(
+      [[maybe_unused]] AVCodecContext* codecContext) override {}
+
+  void convertAVFrameToFrameOutput(
+      const VideoStreamOptions& videoStreamOptions,
+      UniqueAVFrame& avFrame,
+      FrameOutput& frameOutput,
+      std::optional<torch::Tensor> preAllocatedOutputTensor =
+          std::nullopt) override;
+
+ private:
+  int convertAVFrameToTensorUsingSwsScale(
+      const UniqueAVFrame& avFrame,
+      torch::Tensor& outputTensor);
+
+  torch::Tensor convertAVFrameToTensorUsingFilterGraph(
+      const UniqueAVFrame& avFrame);
+
+  struct FilterGraphContext {
+    UniqueAVFilterGraph filterGraph;
+    AVFilterContext* sourceContext = nullptr;
+    AVFilterContext* sinkContext = nullptr;
+  };
+
+  struct DecodedFrameContext {
+    int decodedWidth;
+    int decodedHeight;
+    AVPixelFormat decodedFormat;
+    AVRational decodedAspectRatio;
+    int expectedWidth;
+    int expectedHeight;
+    bool operator==(const DecodedFrameContext&);
+    bool operator!=(const DecodedFrameContext&);
+  };
+
+  void createSwsContext(
+      const DecodedFrameContext& frameContext,
+      const enum AVColorSpace colorspace);
+
+  void createFilterGraph(
+      const DecodedFrameContext& frameContext,
+      const VideoStreamOptions& videoStreamOptions);
+
+  // color-conversion fields. Only one of FilterGraphContext and
+  // UniqueSwsContext should be non-null.
+  FilterGraphContext filterGraphContext_;
+  UniqueSwsContext swsContext_;
+
+  // Used to know whether a new FilterGraphContext or UniqueSwsContext should
+  // be created before decoding a new frame.
+  DecodedFrameContext prevFrameContext_;
+};
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -16,9 +16,10 @@ extern "C" {
 namespace facebook::torchcodec {
 namespace {
 
-bool g_cuda =
-    registerDeviceInterface(torch::kCUDA, [](const torch::Device& device) {
-      return new CudaDeviceInterface(device);
+bool g_cuda = registerDeviceInterface(
+    torch::kCUDA,
+    [](const torch::Device& device, const AVRational& timeBase) {
+      return new CudaDeviceInterface(device, timeBase);
     });
 
 // We reuse cuda contexts across VideoDeoder instances. This is because
@@ -164,8 +165,10 @@ AVBufferRef* getCudaContext(const torch::Device& device) {
 }
 } // namespace
 
-CudaDeviceInterface::CudaDeviceInterface(const torch::Device& device)
-    : DeviceInterface(device) {
+CudaDeviceInterface::CudaDeviceInterface(
+    const torch::Device& device,
+    const AVRational& timeBase)
+    : DeviceInterface(device, timeBase) {
   if (device_.type() != torch::kCUDA) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }

src/torchcodec/_core/CudaDeviceInterface.h

@@ -12,7 +12,7 @@ namespace facebook::torchcodec {
 
 class CudaDeviceInterface : public DeviceInterface {
  public:
-  CudaDeviceInterface(const torch::Device& device);
+  CudaDeviceInterface(const torch::Device& device, const AVRational& timeBase);
 
   virtual ~CudaDeviceInterface();
 

src/torchcodec/_core/DeviceInterface.cpp

@@ -45,11 +45,6 @@ bool registerDeviceInterface(
 }
 
 torch::Device createTorchDevice(const std::string device) {
-  // TODO: remove once DeviceInterface for CPU is implemented
-  if (device == "cpu") {
-    return torch::kCPU;
-  }
-
   std::scoped_lock lock(g_interface_mutex);
   std::string deviceType = getDeviceType(device);
   auto deviceInterface = std::find_if(
@@ -68,21 +63,17 @@ torch::Device createTorchDevice(const std::string device) {
 }
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device) {
+    const torch::Device& device,
+    const AVRational& timeBase) {
   auto deviceType = device.type();
-  // TODO: remove once DeviceInterface for CPU is implemented
-  if (deviceType == torch::kCPU) {
-    return nullptr;
-  }
-
   std::scoped_lock lock(g_interface_mutex);
   TORCH_CHECK(
       g_interface_map->find(deviceType) != g_interface_map->end(),
       "Unsupported device: ",
       device);
 
   return std::unique_ptr<DeviceInterface>(
-      (*g_interface_map)[deviceType](device));
+      (*g_interface_map)[deviceType](device, timeBase));
 }
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/DeviceInterface.h

@@ -27,7 +27,8 @@ namespace facebook::torchcodec {
 
 class DeviceInterface {
  public:
-  DeviceInterface(const torch::Device& device) : device_(device) {}
+  DeviceInterface(const torch::Device& device, const AVRational& timeBase)
+      : device_(device), timeBase_(timeBase) {}
 
   virtual ~DeviceInterface(){};
 
@@ -49,10 +50,11 @@ class DeviceInterface {
 
  protected:
   torch::Device device_;
+  AVRational timeBase_;
 };
 
-using CreateDeviceInterfaceFn =
-    std::function<DeviceInterface*(const torch::Device& device)>;
+using CreateDeviceInterfaceFn = std::function<
+    DeviceInterface*(const torch::Device& device, const AVRational& timeBase)>;
 
 bool registerDeviceInterface(
     torch::DeviceType deviceType,
@@ -61,6 +63,7 @@ bool registerDeviceInterface(
 torch::Device createTorchDevice(const std::string device);
 
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const torch::Device& device);
+    const torch::Device& device,
+    const AVRational& timeBase);
 
 } // namespace facebook::torchcodec