refac

Author: NicolasHug

src/torchcodec/_core/ColorConverter.cpp

@@ -15,6 +15,17 @@
 
 namespace facebook::torchcodec {
 
+namespace {
+// Only ever used to name a device in the error message below.
+std::string printable(const StableDevice& device) {
+  std::string name = device_type_name(device.type());
+  if (device.type() != kStableCPU && device.index() >= 0) {
+    name += ":" + std::to_string(device.index());
+  }
+  return name;
+}
+} // namespace
+
 ColorConverter::ColorConverter(
     const StableDevice& device,
     OutputDtypeConfig output_dtype_config)
@@ -51,13 +62,13 @@ void ColorConverter::maybe_initialize_interface(OutputDtype output_dtype) {
 
 torch::stable::Tensor ColorConverter::convert(
     const AVFrame& av_frame,
-    const std::string& frame_device) {
+    const StableDevice& frame_device) {
   STD_TORCH_CHECK(
-      StableDevice(frame_device) == device_,
+      frame_device == device_,
       "This ColorConverter is on ",
-      device_to_string(device_),
+      printable(device_),
       " but the frame's samples are on ",
-      frame_device,
+      printable(frame_device),
       ". A ColorConverter only converts frames that are already on its own "
       "device: create one per device, or move the RGB output afterwards.");
 

src/torchcodec/_core/ColorConverter.h

@@ -28,7 +28,7 @@ class FORCE_PUBLIC_VISIBILITY ColorConverter {
   // converter's device: we refuse to move samples around behind your back.
   torch::stable::Tensor convert(
       const AVFrame& av_frame,
-      const std::string& frame_device);
+      const StableDevice& frame_device);
 
  private:
   void maybe_initialize_interface(OutputDtype output_dtype);

src/torchcodec/_core/DeviceInterface.cpp

@@ -47,16 +47,6 @@ StableDeviceType parse_device_type(const std::string& device_type) {
 
 } // namespace
 
-std::string device_to_string(const StableDevice& device) {
-  std::string name = device_type_name(device.type());
-  // A negative index means "unspecified" (e.g. device was just "cuda"); leave
-  // it off so the string round-trips and resolves to the current device.
-  if (device.type() != kStableCPU && device.index() >= 0) {
-    name += ":" + std::to_string(device.index());
-  }
-  return name;
-}
-
 bool register_device_interface(
     const DeviceInterfaceKey& key,
     CreateDeviceInterfaceFn create_interface) {

src/torchcodec/_core/DeviceInterface.h

@@ -229,10 +229,6 @@ FORCE_PUBLIC_VISIBILITY void validate_device_interface(
     const std::string& device,
     const std::string& variant = "default");
 
-// "cpu", "cuda", "cuda:1"...
-FORCE_PUBLIC_VISIBILITY std::string device_to_string(
-    const StableDevice& device);
-
 TORCHCODEC_THIRD_PARTY_API std::unique_ptr<DeviceInterface>
 create_device_interface(
     const StableDevice& device,

src/torchcodec/_core/custom_ops.cpp

@@ -88,15 +88,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) {
   m.def("_blocks_packet_decoder_send_eof(Tensor(a!) decoder) -> int");
   m.def("_blocks_packet_decoder_reset(Tensor(a!) decoder) -> ()");
   m.def(
-      "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, str, Tensor)");
+      "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, Device, Tensor)");
   m.def(
       "_blocks_create_color_converter(str device=\"cpu\", str output_dtype=\"uint8\") -> Tensor");
   m.def(
-      "_blocks_convert_frame(Tensor(a!) converter, Tensor frame, str device) -> Tensor");
+      "_blocks_convert_frame(Tensor(a!) converter, Tensor frame, Device device) -> Tensor");
   m.def(
       "_blocks_frame_metadata(Tensor frame) -> (str, str, str, int, int, int, float)");
   m.def(
-      "_blocks_frame_planes(Tensor frame, str device) -> (Tensor, Tensor, Tensor, Tensor)");
+      "_blocks_frame_planes(Tensor frame, Device device) -> (Tensor, Tensor, Tensor, Tensor)");
   m.def("_get_key_frame_indices(Tensor(a!) decoder) -> Tensor");
   m.def("get_json_metadata(Tensor(a!) decoder) -> str");
   m.def("get_container_json_metadata(Tensor(a!) decoder) -> str");
@@ -904,7 +904,7 @@ using OpsReceiveFrameOutput = std::tuple<
     int64_t,
     double,
     double,
-    std::string,
+    StableDevice,
     torch::stable::Tensor>;
 
 OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
@@ -919,7 +919,7 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
         static_cast<int64_t>(status),
         0.0,
         0.0,
-        std::string("cpu"),
+        StableDevice(kStableCPU),
         torch::stable::empty({int64_t(0)}, kStableUInt8));
   }
   AVRational time_base = decoder_ptr->time_base();
@@ -929,7 +929,7 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame(
   // ones a CUDA decoder had to decode on the CPU and upload.
   // TODO_API_BREAKDOWN DESIGN P1: Not sure we need to return the device at all,
   // the device should always match the device parameter now.
-  std::string device = device_to_string(decoder_ptr->device());
+  StableDevice device = decoder_ptr->device();
   torch::stable::Tensor storage =
       decoder_ptr->get_frame_storage(*av_frame).value_or(
           torch::stable::empty({int64_t(0)}, kStableUInt8));
@@ -956,7 +956,7 @@ torch::stable::Tensor _blocks_create_color_converter(
 torch::stable::Tensor _blocks_convert_frame(
     torch::stable::Tensor& converter,
     torch::stable::Tensor& frame,
-    std::string device) {
+    StableDevice device) {
   ColorConverter* converter_ptr =
       unwrap_tensor_to_pointer<ColorConverter>(converter);
   return converter_ptr->convert(
@@ -994,10 +994,10 @@ using OpsFramePlanesOutput = std::tuple<
 
 OpsFramePlanesOutput _blocks_frame_planes(
     torch::stable::Tensor& tensor_handle,
-    std::string device) {
+    StableDevice device) {
   AVFrame* av_frame = unwrap_tensor_to_pointer<AVFrame>(tensor_handle);
   std::vector<torch::stable::Tensor> planes =
-      frame_planes(*av_frame, StableDevice(device), tensor_handle);
+      frame_planes(*av_frame, device, tensor_handle);
 
   // Op schema wants a fixed number of planes, so we pad with empty tensors that
   // then get removed at the Python level.

src/torchcodec/decoders/_blocks/_frame.py

@@ -70,7 +70,7 @@ def __init__(
         handle: torch.Tensor,
         pts_seconds: float,
         duration_seconds: float,
-        device: str = "cpu",
+        device: torch.device = torch.device("cpu"),
         storage: torch.Tensor | None = None,
     ):
         self._handle = handle