Vectorize pooling for optimization

[wingertge]

Pull Request Template

Checklist

• Confirmed that run-checks all script has been executed.
• Made sure the book is up to date with changes in this PR.

Changes

Uses the custom transpose kernel from implicit GEMM convolution to switch all pooling operations to NHWC, and vectorizes along the channel dimension. This improves performance by approx. 20-25% when processing contiguous NCHW tensors, more when processing contiguous NHWC (i.e. output from implicit GEMM convolution). While migrating the kernels, I also renamed all variables that were just numbered to something more semantically meaningful.

Testing

All tests pass, and changes to the kernels are kept as minimal as possible.

[codecov]

Codecov Report

Attention: Patch coverage is 75.40323% with 61 lines in your changes missing coverage. Please review.

Project coverage is 82.32%. Comparing base (5a74fc1) to head (f8c2776).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-cubecl/src/kernel/pool/pool2d.rs	46.15%	14 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/max_pool2d.rs	71.73%	13 Missing ⚠️
...burn-cubecl/src/kernel/pool/adaptive_avg_pool2d.rs	69.44%	11 Missing ⚠️
...cl/src/kernel/pool/adaptive_avg_pool2d_backward.rs	77.27%	10 Missing ⚠️
crates/burn-cubecl/src/kernel/pool/avg_pool2d.rs	68.42%	6 Missing ⚠️
...burn-cubecl/src/kernel/pool/max_pool2d_backward.rs	90.47%	4 Missing ⚠️
...burn-cubecl/src/kernel/pool/avg_pool2d_backward.rs	91.42%	3 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2905   +/-   ##
=======================================
  Coverage   82.32%   82.32%           
=======================================
  Files         867      867           
  Lines      118525   118582   +57     
=======================================
+ Hits        97570    97628   +58     
+ Misses      20955    20954    -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[laggui]

Awesome 🔥