Create DeviceInterface in addStream

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

Author: dvrogozh

src/torchcodec/_core/CudaDevice.cpp

@@ -16,9 +16,9 @@ extern "C" {
 namespace facebook::torchcodec {
 namespace {
 
-bool g_cuda = registerDeviceInterface("cuda", [](const std::string& device) {
-  return new CudaDevice(device);
-});
+bool g_cuda = registerDeviceInterface(
+    torch::kCUDA,
+    [](const torch::Device& device) { return new CudaDevice(device); });
 
 // We reuse cuda contexts across VideoDeoder instances. This is because
 // creating a cuda context is expensive. The cache mechanism is as follows:
@@ -164,7 +164,7 @@ AVBufferRef* getCudaContext(const torch::Device& device) {
 }
 } // namespace
 
-CudaDevice::CudaDevice(const std::string& device) : DeviceInterface(device) {
+CudaDevice::CudaDevice(const torch::Device& device) : DeviceInterface(device) {
   if (device_.type() != torch::kCUDA) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }

src/torchcodec/_core/CudaDevice.h

@@ -12,7 +12,7 @@ namespace facebook::torchcodec {
 
 class CudaDevice : public DeviceInterface {
  public:
-  CudaDevice(const std::string& device);
+  CudaDevice(const torch::Device& device);
 
   virtual ~CudaDevice(){};
 

src/torchcodec/_core/DeviceInterface.cpp

@@ -12,7 +12,7 @@ namespace facebook::torchcodec {
 
 namespace {
 std::mutex g_interface_mutex;
-std::map<std::string, CreateDeviceInterfaceFn> g_interface_map;
+std::map<torch::DeviceType, CreateDeviceInterfaceFn> g_interface_map;
 
 std::string getDeviceType(const std::string& device) {
   size_t pos = device.find(':');
@@ -25,7 +25,7 @@ std::string getDeviceType(const std::string& device) {
 } // namespace
 
 bool registerDeviceInterface(
-    const std::string deviceType,
+    torch::DeviceType deviceType,
     CreateDeviceInterfaceFn createInterface) {
   std::scoped_lock lock(g_interface_mutex);
   TORCH_CHECK(
@@ -36,15 +36,36 @@ bool registerDeviceInterface(
   return true;
 }
 
-std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const std::string device) {
+torch::Device createTorchDevice(const std::string device) {
   // TODO: remove once DeviceInterface for CPU is implemented
   if (device == "cpu") {
-    return nullptr;
+    return torch::kCPU;
   }
 
   std::scoped_lock lock(g_interface_mutex);
   std::string deviceType = getDeviceType(device);
+  auto deviceInterface = std::find_if(
+      g_interface_map.begin(),
+      g_interface_map.end(),
+      [&](const std::pair<torch::DeviceType, CreateDeviceInterfaceFn>& arg) {
+        return device.rfind(
+                   torch::DeviceTypeName(arg.first, /*lcase*/ true), 0) == 0;
+      });
+  TORCH_CHECK(
+      deviceInterface != g_interface_map.end(), "Unsupported device: ", device);
+
+  return torch::Device(device);
+}
+
+std::unique_ptr<DeviceInterface> createDeviceInterface(
+    const torch::Device& device) {
+  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: ",

src/torchcodec/_core/DeviceInterface.h

@@ -26,7 +26,7 @@ namespace facebook::torchcodec {
 
 class DeviceInterface {
  public:
-  DeviceInterface(const std::string& device) : device_(device) {}
+  DeviceInterface(const torch::Device& device) : device_(device) {}
 
   virtual ~DeviceInterface(){};
 
@@ -53,13 +53,15 @@ class DeviceInterface {
 };
 
 using CreateDeviceInterfaceFn =
-    std::function<DeviceInterface*(const std::string& device)>;
+    std::function<DeviceInterface*(const torch::Device& device)>;
 
 bool registerDeviceInterface(
-    const std::string deviceType,
+    torch::DeviceType deviceType,
     const CreateDeviceInterfaceFn createInterface);
 
+torch::Device createTorchDevice(const std::string device);
+
 std::unique_ptr<DeviceInterface> createDeviceInterface(
-    const std::string device);
+    const torch::Device& device);
 
 } // namespace facebook::torchcodec

src/torchcodec/_core/SingleStreamDecoder.cpp

@@ -94,9 +94,9 @@ SingleStreamDecoder::SingleStreamDecoder(
 
 SingleStreamDecoder::~SingleStreamDecoder() {
   for (auto& [streamIndex, streamInfo] : streamInfos_) {
-    auto& device = streamInfo.videoStreamOptions.device;
-    if (device) {
-      device->releaseContext(streamInfo.codecContext.get());
+    auto& deviceInterface = streamInfo.deviceInterface;
+    if (deviceInterface) {
+      deviceInterface->releaseContext(streamInfo.codecContext.get());
     }
   }
 }
@@ -386,7 +386,7 @@ torch::Tensor SingleStreamDecoder::getKeyFrameIndices() {
 void SingleStreamDecoder::addStream(
     int streamIndex,
     AVMediaType mediaType,
-    DeviceInterface* device,
+    const torch::Device& device,
     std::optional<int> ffmpegThreadCount) {
   TORCH_CHECK(
       activeStreamIndex_ == NO_ACTIVE_STREAM,
@@ -414,6 +414,7 @@ void SingleStreamDecoder::addStream(
   streamInfo.timeBase = formatContext_->streams[activeStreamIndex_]->time_base;
   streamInfo.stream = formatContext_->streams[activeStreamIndex_];
   streamInfo.avMediaType = mediaType;
+  streamInfo.deviceInterface = createDeviceInterface(device);
 
   // This should never happen, checking just to be safe.
   TORCH_CHECK(
@@ -425,9 +426,10 @@ void SingleStreamDecoder::addStream(
   // TODO_CODE_QUALITY it's pretty meh to have a video-specific logic within
   // addStream() which is supposed to be generic
   if (mediaType == AVMEDIA_TYPE_VIDEO) {
-    if (device) {
+    if (streamInfo.deviceInterface) {
       avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
-          device->findCodec(streamInfo.stream->codecpar->codec_id)
+          streamInfo.deviceInterface
+              ->findCodec(streamInfo.stream->codecpar->codec_id)
               .value_or(avCodec));
     }
   }
@@ -445,8 +447,8 @@ void SingleStreamDecoder::addStream(
 
   // TODO_CODE_QUALITY same as above.
   if (mediaType == AVMEDIA_TYPE_VIDEO) {
-    if (device) {
-      device->initializeContext(codecContext);
+    if (streamInfo.deviceInterface) {
+      streamInfo.deviceInterface->initializeContext(codecContext);
     }
   }
 
@@ -476,7 +478,7 @@ void SingleStreamDecoder::addVideoStream(
   addStream(
       streamIndex,
       AVMEDIA_TYPE_VIDEO,
-      videoStreamOptions.device.get(),
+      videoStreamOptions.device,
       videoStreamOptions.ffmpegThreadCount);
 
   auto& streamMetadata =
@@ -1217,11 +1219,11 @@ SingleStreamDecoder::convertAVFrameToFrameOutput(
       formatContext_->streams[activeStreamIndex_]->time_base);
   if (streamInfo.avMediaType == AVMEDIA_TYPE_AUDIO) {
     convertAudioAVFrameToFrameOutputOnCPU(avFrame, frameOutput);
-  } else if (!streamInfo.videoStreamOptions.device) {
+  } else if (!streamInfo.deviceInterface) {
     convertAVFrameToFrameOutputOnCPU(
         avFrame, frameOutput, preAllocatedOutputTensor);
-  } else if (streamInfo.videoStreamOptions.device) {
-    streamInfo.videoStreamOptions.device->convertAVFrameToFrameOutput(
+  } else if (streamInfo.deviceInterface) {
+    streamInfo.deviceInterface->convertAVFrameToFrameOutput(
         streamInfo.videoStreamOptions,
         avFrame,
         frameOutput,
@@ -1564,10 +1566,8 @@ SingleStreamDecoder::FrameBatchOutput::FrameBatchOutput(
       videoStreamOptions, streamMetadata);
   int height = frameDims.height;
   int width = frameDims.width;
-  torch::Device device = (videoStreamOptions.device)
-      ? videoStreamOptions.device->device()
-      : torch::kCPU;
-  data = allocateEmptyHWCTensor(height, width, device, numFrames);
+  data = allocateEmptyHWCTensor(
+      height, width, videoStreamOptions.device, numFrames);
 }
 
 torch::Tensor allocateEmptyHWCTensor(

src/torchcodec/_core/SingleStreamDecoder.h

@@ -139,7 +139,7 @@ class SingleStreamDecoder {
     std::optional<int> height;
     std::optional<ColorConversionLibrary> colorConversionLibrary;
     // By default we use CPU for decoding for both C++ and python users.
-    std::shared_ptr<DeviceInterface> device;
+    torch::Device device = torch::kCPU;
   };
 
   struct AudioStreamOptions {
@@ -358,6 +358,8 @@ class SingleStreamDecoder {
     // Used to know whether a new FilterGraphContext or UniqueSwsContext should
     // be created before decoding a new frame.
     DecodedFrameContext prevFrameContext;
+
+    std::unique_ptr<DeviceInterface> deviceInterface;
   };
 
   // --------------------------------------------------------------------------
@@ -460,7 +462,7 @@ class SingleStreamDecoder {
   void addStream(
       int streamIndex,
       AVMediaType mediaType,
-      DeviceInterface* device = nullptr,
+      const torch::Device& device = torch::kCPU,
       std::optional<int> ffmpegThreadCount = std::nullopt);
 
   // Returns the "best" stream index for a given media type. The "best" is

src/torchcodec/_core/custom_ops.cpp

@@ -243,8 +243,7 @@ void _add_video_stream(
     }
   }
   if (device.has_value()) {
-    videoStreamOptions.device =
-        createDeviceInterface(std::string(device.value()));
+    videoStreamOptions.device = createTorchDevice(std::string(device.value()));
   }
 
   auto videoDecoder = unwrapTensorToGetDecoder(decoder);