Matmul 2d input config selection bugfix

[AshAnand34]

This pull request introduces enhancements to the matrix multiplication (matmul) logic to handle 2D inputs more robustly, including updates to stage size validation, error handling, and new test cases. The key changes focus on ensuring correctness for edge cases like 1xN and Mx1 matrices.

Enhancements to 2D input handling:

• Stage size adjustments for 2D inputs: Modified the matmul_selection logic to ensure stage sizes divide evenly when handling 1xN or Mx1 matrices. This includes iterative adjustments to stage_size_m and stage_size_n to meet divisibility constraints. (crates/cubecl-linalg/src/convolution/selection.rs, crates/cubecl-linalg/src/convolution/selection.rsL148-R185)

• Validation for 2D stage sizes: Added checks in the MatmulConfigFactory implementation to validate that stage sizes for 2D inputs (1xN or Mx1) divide the corresponding dimensions evenly. Errors are returned if the constraints are not met. (crates/cubecl-linalg/src/matmul/components/tile/accelerated.rs, crates/cubecl-linalg/src/matmul/components/tile/accelerated.rsR129-R151)

Testing improvements:

• New test cases for 2D inputs: Introduced a dedicated test, test_2d_input_matmul, to validate matrix multiplication for 1xN and Mx1 input scenarios. This ensures that the new logic handles these edge cases correctly. (crates/cubecl-linalg/tests/matmul_2d.rs, crates/cubecl-linalg/tests/matmul_2d.rsR1-R64)

Validate your PR with burn.

It is important that you make sure that you don't introduce any bugs in burn.

Instructions

• Create a new branch or fork of the burn repo
• Update the main Cargo.toml with this PR hash.
• Fix any broken tests or compilation errors in burn.
•  Submit a PR in burn with your fixes and link it here.

[louisfd]

It's true that we would need to optimize for degenerate vector cases 😄 However there are a few things to fix, hope my comments can guide you