Create kernel on-device for transforms.functional.gaussian_blur

[drhead]

The current implementation of gaussian blur creates the kernels on CPU and then moves them to the device and dtype of the image. This goes against best practices for pytorch optimization, since it induces a forced device sync.

To fix it, I have passed the device and dtype parameters over to the linspace operation that creates the initial tensor. This should remove the forced device sync. It is worth noting that this may result in it being calculated in a lower precision than normal. This would of course involve an extra call for the cast to get it to the correct precision.

There's probably room for more optimization in this, especially when GaussianBlur is used as a module -- really, I think it would be more appropriate to create the kernels one time and reuse them instead of constructing new ones. This PR solves the more critical issue as is, though.

Closes #8401

cc @vfdev-5

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/8426

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 8 New Failures, 14 Unrelated Failures

As of commit ad00447 with merge base a5f531a ():

NEW FAILURES - The following jobs have failed:

• CMake / windows (windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-macos (3.11, macos-m1-stable) / macos-job (gh)
  test/test_transforms_v2.py::TestLinearTransform::test_transform[cpu-dtype1-make_image_tensor]
• Tests / unittests-macos (3.9, macos-m1-stable) / macos-job (gh)
  test/test_transforms_v2.py::TestLinearTransform::test_transform[cpu-dtype1-make_image_tensor]
• Tests / unittests-windows (3.10, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.11, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.12, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.8, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Tests / unittests-windows (3.9, windows.4xlarge, cpu) / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• CMake / linux (linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• CMake / linux (linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.10, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.11, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.12, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.8, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Prototype tests on Linux / unittests-prototype (3.9, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.10, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.11, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.12, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.8, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.8, linux.g5.4xlarge.nvidia.gpu, cuda, 11.8) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Tests / unittests-linux (3.9, linux.12xlarge, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[NicolasHug]

Thank you for the PR @drhead , I'll merge after the CI passes.

Note that the V2 version in torchvision.transforms.v2.functional.gaussian_blur would already create the kernel on-device, so perhaps you can migrate to the V2 version already.