Skip to content

[Perf] Convolution migration to NHWC #3090

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 18 commits into from
Apr 29, 2025
Merged

Conversation

wingertge
Copy link
Contributor

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.

Related Issues/PRs

Requires tracel-ai/cubecl#646 to get merged first

Changes

Reworks im2col and direct convolution to use NHWC layout, making the 2D forward pass completely NHWC. This saves unnecessary transpositions and is a prerequisite for eventually moving weights to NHWC overall. It also improves performance.

Testing

All tests pass with both the specific algorithms and autotune enabled.

@wingertge wingertge marked this pull request as draft April 27, 2025 15:39
@wingertge wingertge marked this pull request as ready for review April 28, 2025 14:32
Copy link

codecov bot commented Apr 28, 2025

Codecov Report

Attention: Patch coverage is 43.67816% with 196 lines in your changes missing coverage. Please review.

Project coverage is 81.06%. Comparing base (4360f36) to head (45a4393).
Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
...rates/burn-cubecl/src/kernel/conv/conv2d/im2col.rs 38.55% 102 Missing ⚠️
...rates/burn-cubecl/src/kernel/conv/conv2d/direct.rs 46.08% 62 Missing ⚠️
crates/burn-cubecl/src/kernel/contiguous.rs 26.31% 14 Missing ⚠️
crates/burn-cubecl/src/kernel/utils.rs 60.00% 10 Missing ⚠️
crates/burn-cubecl-fusion/src/shared/kernel.rs 0.00% 4 Missing ⚠️
crates/burn-cubecl-fusion/src/shared/io.rs 0.00% 2 Missing ⚠️
crates/burn-cubecl/src/kernel/conv/conv2d/base.rs 80.00% 1 Missing ⚠️
...ecl/src/kernel/conv/conv2d/implicit_gemm/launch.rs 66.66% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3090      +/-   ##
==========================================
- Coverage   81.08%   81.06%   -0.03%     
==========================================
  Files         817      817              
  Lines      117326   117408      +82     
==========================================
+ Hits        95131    95173      +42     
- Misses      22195    22235      +40     

☔ 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 laggui self-requested a review April 29, 2025 11:40
Copy link
Member

@laggui laggui left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ran the tests locally for cuda and wgpu/vulkan. Got a couple of failures on vulkan with f16 for convolutions, but it looks like they are also failing on main. So something broke since the 0.17 release.

With the failing tests I found a print statement for into data but that will be fixed in #3114.

Also linux-std runner keeps running out of disk space with all the updates we're doing recently. I think I'm just gonna disable the caching on this runner, more of an annoyance than anything else!

Sorry for the rant 😅 just happened to stumble upon multiple issues while testing the changes..

TL;DR: this PR looks good to me!

@laggui laggui merged commit ceab6d4 into tracel-ai:main Apr 29, 2025
26 of 29 checks passed
@wingertge wingertge deleted the perf/conv branch April 29, 2025 17:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants