Panic due to indexing into an empty array during kernel launch

[wingertge]

I can't figure out the conditions that bring this about, but I'm currently getting strange crashes stemming from the cubecl_core::compute::launcher::TensorState<R>::push function in cubecl. Something causes the vector to be empty when it is being indexed in cubecl_core::compute::launcher::TensorState<R>::register_strides.

Relevant section in the crash log

thread 'tests::ad_deform_conv2d::tests::test_deform_conv2d_basic' panicked at <cargo_dir>\crates\cubecl-core\src\compute\launcher.rs:230:30:
index out of bounds: the len is 0 but the index is 0
stack backtrace:
   0: std::panicking::begin_panic_handler
             at /rustc/3f5fd8dd41153bc5fdca9427e9e05be2c767ba23/library\std\src\panicking.rs:652
   1: core::panicking::panic_fmt
             at /rustc/3f5fd8dd41153bc5fdca9427e9e05be2c767ba23/library\core\src\panicking.rs:72
   2: core::panicking::panic_bounds_check
             at /rustc/3f5fd8dd41153bc5fdca9427e9e05be2c767ba23/library\core\src\panicking.rs:274
   3: cubecl_core::compute::launcher::TensorState<R>::push
   4: burn_jit::kernel::conv::deform_conv2d::deform_col2img_kernel::launch
   5: burn_jit::kernel::conv::deform_conv2d::deform_conv2d_backward

Now that I have the error, no changes make it go away so I unfortunately can't manage to root cause the issue. The deform_col2img_kernel in this file: deform_conv2d.rs seems to be having this issue but I haven't yet had a chance to do an isolated repro.

[wingertge]

I've managed to root cause it. It seems to always occur if a lower ranked tensor parameter is used before a higher one, triggering the adjust_rank function. This function calls reigster_strides with an empty vec which will always panic.

[nathanielsimard]

Yeah I think adjust_rank is actually buggy, needs to improve it 😅