enable decode_jpeg to support privateuseone device

Signed-off-by: taozhiwei <[email protected]>

Author: taozhiwei

test/test_image.py

@@ -160,6 +160,35 @@ def test_decode_bad_huffman_images():
     decode_jpeg(bad_huff)
 
 
+def test_decode_jpeg_privateuseone_custom_backend():
+    privateuseone_name = torch._C._get_privateuse1_backend_name()
+    device = torch.device(privateuseone_name)
+    data = torch.full((1, 2, 3), 233, dtype=torch.uint8)
+    # When the custom operator is not registered, an error should
+    # be reported and prompted to register decode_jpegs_privateuseone.
+    with pytest.raises(RuntimeError, match="decode_jpegs_privateuseone"):
+        decode_jpeg(data, device=device)
+
+    # Register a simple custom implementation to return the original data
+    called = {}
+    lib = torch.library.Library("image", "FRAGMENT")
+    try:
+        lib.define(
+            "decode_jpegs_privateuseone(Tensor[] input, int mode=0, bool apply_exif_orientation=False) -> Tensor[]"
+        )
+
+        @torch.library.impl(lib, "decode_jpegs_privateuseone", "CPU")
+        def _decode_jpegs_privateuseone(input, mode=0, apply_exif_orientation=False):
+            called["value"] = True
+            return input
+    except RuntimeError:
+        pass
+
+    output = decode_jpeg(data, device=device)
+    assert called.get("value") is True
+    torch.testing.assert_close(output, data)
+
+
 @pytest.mark.parametrize(
     "img_path",
     [

torchvision/io/image.py

@@ -210,6 +210,16 @@ def decode_jpeg(
             raise ValueError("All elements of the input list must be tensors.")
         if not all(t.device.type == "cpu" for t in input):
             raise ValueError("Input list must contain tensors on CPU.")
+        custom_privateuse1_name = torch._C._get_privateuse1_backend_name()
+        if device.type == custom_privateuse1_name or device.type == "privateuseone":
+            # When the target device is privateuseone, switch to calling the custom decode_jpegs_privateuseone.
+            # This operator needs to be pre-registered by the user through torch.library.define/impl.
+            decoder = getattr(torch.ops.image, "decode_jpegs_privateuseone", None)
+            if decoder is None:
+                raise RuntimeError(
+                    "decode_jpeg(device='privateuseone') need register torch.ops.image.decode_jpegs_privateuseone."
+                )
+            return decoder(input, mode.value, apply_exif_orientation)
         if device.type == "cuda":
             return torch.ops.image.decode_jpegs_cuda(input, mode.value, device)
         else:
@@ -218,6 +228,14 @@ def decode_jpeg(
     else:  # input is tensor
         if input.device.type != "cpu":
             raise ValueError("Input tensor must be a CPU tensor")
+        if device.type == custom_privateuse1_name or device.type == "privateuseone":
+            custom_privateuse1_name = torch._C._get_privateuse1_backend_name()
+            decoder = getattr(torch.ops.image, "decode_jpegs_privateuseone", None)
+            if decoder is None:
+                raise RuntimeError(
+                    "decode_jpeg(device='privateuseone') need register torch.ops.image.decode_jpegs_privateuseone."
+                )
+            return decoder([input], mode.value, apply_exif_orientation)[0]
         if device.type == "cuda":
             return torch.ops.image.decode_jpegs_cuda([input], mode.value, device)[0]
         else: