pytorch

090b3b95 - [PyTorch] Add Vulkan support and tests for at::select.int operator, 4 dim/rank tensor case (#96228)

[PyTorch] Add Vulkan support and tests for at::select.int operator, 4 dim/rank tensor case (#96228) Summary: Currently, selection along a dimension/rank is only supported for 3D/rank tensors in PyTorch Vulkan. This adds support for 4D/rank tensors at selection along batch, channel (depth), height, and width. Additionally: - The existing implementations have been name-refactored to reflect whether they operate on 3d or 4d tensors. - The params buffer for all select operations now use `ivec2` or `ivec4` only, for memory alignment safety. Test Plan: 1. `buck run --target-platforms ovr_config//platform/macos:arm64-fbsource //xplat/caffe2:pt_vulkan_api_test_binAppleMac\#macosx-arm64 -c pt.vulkan_full_precision=1` on Apple M1 MacBook 2. Confirm all tests pass with no regression, and the directly affected tests `select_4d_*`, refactored `select_3d_`, pass 3. Test output P636928908, in particular: ``` [...bunch of other tests...] [ RUN ] VulkanAPITest.select_3d_depth_small [ OK ] VulkanAPITest.select_3d_depth_small (1 ms) [ RUN ] VulkanAPITest.select_3d_depth_medium [ OK ] VulkanAPITest.select_3d_depth_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_depth_large [ OK ] VulkanAPITest.select_3d_depth_large (1 ms) [ RUN ] VulkanAPITest.select_3d_height_small [ OK ] VulkanAPITest.select_3d_height_small (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium [ OK ] VulkanAPITest.select_3d_height_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium1 [ OK ] VulkanAPITest.select_3d_height_medium1 (0 ms) [ RUN ] VulkanAPITest.select_3d_height_medium2 [ OK ] VulkanAPITest.select_3d_height_medium2 (0 ms) [ RUN ] VulkanAPITest.select_3d_height_large [ OK ] VulkanAPITest.select_3d_height_large (1 ms) [ RUN ] VulkanAPITest.select_3d_width_small [ OK ] VulkanAPITest.select_3d_width_small (0 ms) [ RUN ] VulkanAPITest.select_3d_width_medium [ OK ] VulkanAPITest.select_3d_width_medium (0 ms) [ RUN ] VulkanAPITest.select_3d_width_medium2 [ OK ] VulkanAPITest.select_3d_width_medium2 (0 ms) [ RUN ] VulkanAPITest.select_3d_width_large [ OK ] VulkanAPITest.select_3d_width_large (1 ms) [ RUN ] VulkanAPITest.select_4d_batch_small [ OK ] VulkanAPITest.select_4d_batch_small (0 ms) [ RUN ] VulkanAPITest.select_4d_batch_medium [ OK ] VulkanAPITest.select_4d_batch_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_batch_large [ OK ] VulkanAPITest.select_4d_batch_large (1 ms) [ RUN ] VulkanAPITest.select_4d_depth_small [ OK ] VulkanAPITest.select_4d_depth_small (1 ms) [ RUN ] VulkanAPITest.select_4d_depth_medium [ OK ] VulkanAPITest.select_4d_depth_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_depth_large [ OK ] VulkanAPITest.select_4d_depth_large (1 ms) [ RUN ] VulkanAPITest.select_4d_height_small [ OK ] VulkanAPITest.select_4d_height_small (0 ms) [ RUN ] VulkanAPITest.select_4d_height_medium [ OK ] VulkanAPITest.select_4d_height_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_height_large [ OK ] VulkanAPITest.select_4d_height_large (1 ms) [ RUN ] VulkanAPITest.select_4d_width_small [ OK ] VulkanAPITest.select_4d_width_small (0 ms) [ RUN ] VulkanAPITest.select_4d_width_medium [ OK ] VulkanAPITest.select_4d_width_medium (0 ms) [ RUN ] VulkanAPITest.select_4d_width_large [ OK ] VulkanAPITest.select_4d_width_large (1 ms) [...bunch of other tests...] [ FAILED ] 7 tests, listed below: [ FAILED ] VulkanAPITest.cat_dim1_singledepth_success [ FAILED ] VulkanAPITest.gru_success [ FAILED ] VulkanAPITest.gru_mclareninputs_success [ FAILED ] VulkanAPITest.gru_prepack_success [ FAILED ] VulkanAPITest.lstm_success [ FAILED ] VulkanAPITest.lstm_mclareninputs_success [ FAILED ] VulkanAPITest.lstm_prepack_success ``` Reviewed By: SS-JIA Differential Revision: D42623181 Pull Request resolved: https://github.com/pytorch/pytorch/pull/96228 Approved by: https://github.com/SS-JIA