[proto] Performance improvements for equalize op

[vfdev-5]

Description:

• vectorized version of equalize_image_tensor. This is a great job done by @lezcano, thanks a lot for your help ! 🎉
• Speed up on CPU and CUDA

[---------------- equalize_image_tensor transforms measurements ---------------]
                                         |    main   |    new  
1 threads: ----------------------------------------------------
      <class 'torch.Tensor'> Image data  |  978.917  |  335.936

Times are in microseconds (us).

Timestamp: 20221012-160208
Torch version: 1.14.0.dev20221010+cu116
Torchvision version: 0.15.0a0
Num threads: 1

Time benchmark: RandomEqualize (1.0,) None
V2: RandomEqualize(p=1.0) torchvision.prototype.transforms._color
Stable: RandomEqualize(p=1.0) torchvision.transforms.transforms

[- RandomEqualize transforms measurements -]
                         |  stable  |    v2 
1 threads: ---------------------------------
      Tensor Image data  |  2.891   |  2.266

Times are in milliseconds (ms).

cc @datumbox @NicolasHug

[datumbox]

LGTM, thank you @vfdev-5 and @lezcano for making this fast! It's interesting that we can achieve this kind of performance optimization on the frontend without calling the native histogram methods. There might be some speed improvements to be achieved on the future on Core.

I see that all the tests pass, so this looks good to go.

[vfdev-5]

@datumbox speaking with Mario, we can improve the code a bit more, I'll send an update in a follow-up PR.

[lezcano]

It's interesting that we can achieve this kind of performance optimization on the frontend without calling the native histogram methods

For reference, note that the speed-ups do not come from the histogram method itself. I just used that way of implementing a histogram because it allows to have batches. The main speed-up comes from using batched operations all across vs having a for loop and running all these kernels once per channel.

[lezcano]

Probably generalising torch.bincount would be overkill, but it could be interesting to add a note on bincount such as "if you want to do this for a batch of indices, think of using index_add_ in the .. seealso::" section of the docs.