Add bounds checking to implicit GEMM to allow arbitrary input shapes

[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

Implicit GEMM was previously only padded on the input channel, now all 3 dimensions get implicitly padded to allow arbitrary sized inputs. The performance of the padded convolutions is slightly slower than the next matching size, but still faster than other convolution strategies.

Testing

All tests pass with implicit GEMM except the ones with groups.

[codecov]

Codecov Report

Attention: Patch coverage is 22.03390% with 92 lines in your changes missing coverage. Please review.

Project coverage is 85.05%. Comparing base (ce7d866) to head (ec8a356).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
...s/burn-jit/src/kernel/conv/conv2d/implicit_gemm.rs	22.03%	92 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2354      +/-   ##
==========================================
- Coverage   85.06%   85.05%   -0.01%     
==========================================
  Files         769      769              
  Lines       98702    98732      +30     
==========================================
+ Hits        83959    83976      +17     
- Misses      14743    14756      +13     

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

[louisfd]

LGTM