Consider adding this on an anonymous namespace as it looks an internal detail of the implementation. Also just checking whether this should be released at any point to avoid memory leaks.

Also just checking whether this should be released at any point to avoid memory leaks

Yeah this is a good point. Creating / allocating it at each call has some severe overhead so it makes sense to declare this as a global variable (related discussion: #2786 (comment)). But this means we never really know when to release it, and the memory will only be freed when the process is killed.

For reference: this is thread-safe

Good that it's thread-safe but it's still unclear to me whether we have to find a way to release it or if we can leave it be. We don't have such an idiom at TorchVision but I wonder if there are examples of resources on PyTorch core that are never released.

@ezyang how do you handle situations like this on core?

We leak a bunch of global statically initialized objects, as it is the easiest way to avoid destructor ordering problems. If nvjpeg is very simple library you might be able to write an RAII class for this object and have it destruct properly on shutdown.

Precedent in PyTorch is the cudnn convolution cache, look at that for some inspiration.

@ezyang Thanks for the insights.

@NicolasHug Given Ed's input, I think we can mark this as resolved if you want.

Not sure if the dispatcher would make sense here. Since this is the first IO method we add for GPU, it might be worth checking the naming conventions (_cuda) as this will be reproduced on the near future in other places. Thoughts @fmassa ?

I think using the dispatcher would be good, but I'm not sure how it handles constructor functions (like torch.empty / torch.rand).

Indeed, this function always takes CPU tensors, and it's up to a device argument to decide if we should dispatch to the CPU or the CUDA version.

@ezyang do you know if we can use the dispatcher to dispatch taking a torch.device into account, knowing that all tensors live in the CPU?

@fmassa How about reading the data on CPU since that's needed and then calling to() to move it on the right device. This can happen in the python side of things and remain hidden. Then after the binary data living in GPU, the dispatcher can be used as normal to decide if the decoding should happen on the GPU or CPU. Thoughts on this?

Still, nvjpeg requires the input data to live on CPU, so we would need to move it back to CPU again within the function, which would be inefficient. I would have preferred if we could pass the tensor directly as a CUDA tensor as well, but I'm not sure this is possible without further overheads

Thanks for the clarifications concerning nvjpeg. I think that we can investigate on future PRs how we could do this more elegantly. No need to block this PR.

Do we want to consider the mean different or the max difference here? What would be the minimum value so that max tests pass here?

The max error can be quite high unfortunately, the minimum threshold for all tests to pass seems to be 52, after which some tests start failing.
In test_decode_jpeg, we also test for the MAE (with the same threshold=2)

Hum, this looks suspicious that the decoding gives such large differences. Something to keep an eye on

nit: this is fine for now, but this adds extra overhead as we need to create a full Tensor just to extract the data pointer (and the Tensor construction is heavy). Given that we generally only have 3 channels that shouldn't be much of an issue, but still good to keep in mind.

Some alternatives would be to directly use the raw data_ptr adding the correct offsets, like

uint8_t * out_tensor_ptr = out_tensor.data_ptr<uint8_t>();
...
out_image.channel[c] = out_tensor_ptr + c * height * width;

Also, interesting that nvjpeg accept decoding images in both CHW and HWC formats -- I wonder if there is any performance implications by decoding it in CHW?

For reference: this is thread-safe

I'm wondering if we should clear this if the creation fails, otherwise we might not be able to run this function anymore as the handle will be invalid?

Since it's an opaque handle I think the use of free() may not be correct unless it's documented as being supported.

(I would hope that it simply leaves the handle as null if initialisation fails, although I don't see that in the docs - here it just reinits the handle without any freeing when hw backend fails though)

Hm that's a good point. What would you recommend instead of free?

Before pushing this I did a quick test by inserting

       nvjpegStatus_t create_status = nvjpegCreateSimple(&nvjpeg_handle);
       create_status = NVJPEG_STATUS_NOT_INITIALIZED;   // <- this

and all the tests were failing gracefully with E RuntimeError: nvjpegCreateSimple failed: 1. Since I was running the tests with pytest test/test_image.py -k cuda they were all in the same process and pytest was just catching the RuntimeErrors, so I assumed it was OK.

Perhaps nvjpegDestroy(nvjpeg_handle)? See: https://github.com/NVIDIA/CUDALibrarySamples/blob/master/nvJPEG/Image-Resize/imageResize.cpp#L426

I didn't use it because I was wondering whether nvjpegDestroy would properly work with a bad handle?

I guess if you assume anything can happen if initialisation fails then it might end up being an arbitrary value like 0xDEADBEEF and all you can do is reset it to null.