[Perf] Interpolate optimizations

[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.

Related Issues/PRs

Also fixes #3076

Changes

Refactors interpolation to use NHWC layout to allow for potential vectorization when used with NHWC layout tensors (i.e. from convolution), as well as other perf optimizations. Also fixes a bug in is_contiguous and a clippy warning for slices.

Testing

All tests pass

[codecov]

Codecov Report

Attention: Patch coverage is 44.57364% with 143 lines in your changes missing coverage. Please review.

Project coverage is 81.15%. Comparing base (e096b0c) to head (4878d8f).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...ates/burn-cubecl/src/kernel/interpolate/bicubic.rs	17.24%	48 Missing ⚠️
...tes/burn-cubecl/src/kernel/interpolate/bilinear.rs	18.51%	44 Missing ⚠️
...-cubecl/src/kernel/interpolate/nearest_backward.rs	28.57%	25 Missing ⚠️
...ates/burn-cubecl/src/kernel/interpolate/nearest.rs	37.83%	23 Missing ⚠️
...ecl/src/kernel/conv/conv2d/implicit_gemm/launch.rs	66.66%	1 Missing ⚠️
...burn-cubecl/src/kernel/pool/avg_pool2d_backward.rs	66.66%	1 Missing ⚠️
crates/burn-cubecl/src/kernel/utils.rs	92.30%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3077      +/-   ##
==========================================
- Coverage   81.16%   81.15%   -0.01%     
==========================================
  Files         815      816       +1     
  Lines      117201   117265      +64     
==========================================
+ Hits        95121    95163      +42     
- Misses      22080    22102      +22     

☔ 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]

LGTM!

Tested locally and with the super-resolution model code (ESRGAN) which prompted this PR.

Just a minor comment on the contiguous checks

[laggui]

We should probably add more tests (e.g., your tensor of shape [1, 1, 1, 9] with strides [72, 1, 72, 8] which prompted this fix). Better illustrates what is expected to be contiguous or not, even in more complex cases.

Also @nathanielsimard I just realized we seem to have another is_contiguous fn in fusion, but the implementations are slightly diverging with the latest changes

burn/crates/burn-cubecl-fusion/src/base.rs

Lines 141 to 169 in e096b0c

	pub(crate) fn is_contiguous(shape: &[usize], strides: &[usize]) -> bool {
	if shape.is_empty() {
	return true;
	}
	if shape.len() == 1 {
	return strides[0] == 1;
	}
	let mut prev_stride = 1;
	let mut current_num_elems_shape = 1;
	for (i, (stride, shape)) in strides.iter().zip(shape).rev().enumerate() {
	if i > 0 {
	if current_num_elems_shape != *stride {
	return false;
	}
	if prev_stride >= *stride {
	return false;
	}
	}
	current_num_elems_shape *= shape;
	prev_stride = *stride;
	}
	true
	}

The definition of contiguous should probably match 😅

[wingertge]

I added a test, but leaving this unresolved because of the fusion thing

[laggui]

Unrelated to this PR, so I'll merge but will open a small PR to use the same is_contiguous fn across the crates.