Add BackendRouter to handle multiple backends

[laggui]

Needs more tests :)

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

Closes #2276

Changes

Introduces a new BackendRouter responsible for forwarding tensor operations to the appropriate backend, given multiple backends.

This is achieved with the help of the intermediate representation defined for ReprBackend and the tensor/ops descriptions.

Testing

Modified the ag-news-train text classification example to run on cuda + wgpu. Also have a MWE:

use burn::{
    backend::{cuda_jit::CudaDevice, wgpu::WgpuDevice, CudaJit, Wgpu},
    tensor::Tensor,
};
use burn_router::{BackendRouter, ByteBridge, DirectChannel, MultiDevice2};

fn main() {
    type DirectByteChannel<Backends> = DirectChannel<Backends, ByteBridge<Backends>>;

    type DualBackend = BackendRouter<DirectByteChannel<(CudaJit, Wgpu)>>;

    let device2 = WgpuDevice::Cpu;
    let device1 = CudaDevice::new(0);

    // TODO: this is wack.. how to automatically implement From<B1::Device1> for MultiDevice2?
    let multi_device1 = MultiDevice2::Device1(device1);
    let multi_device2 = MultiDevice2::Device2(device2);
    let tensor1 = Tensor::<DualBackend, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &multi_device1);
    let tensor2 = Tensor::<DualBackend, 1>::from_floats([5.0, 6.0, 7.0, 8.0], &multi_device2);

    println!("Tensor 1:\n{tensor1}");
    println!("Tensor 2:\n{tensor2}");

    let tensor1 = tensor1.to_device(&multi_device2);

    let output = tensor1.add(tensor2);

    println!("Result:\n{output}");
}

Tensor 1:
Tensor {
  data:
[1.0, 2.0, 3.0, 4.0],
  shape:  [4],
  device:  Device1(CudaDevice { index: 0 }),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}
Tensor 2:
Tensor {
  data:
[5.0, 6.0, 7.0, 8.0],
  shape:  [4],
  device:  Device2(Cpu),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}
Result:
Tensor {
  data:
[6.0, 8.0, 10.0, 12.0],
  shape:  [4],
  device:  Device2(Cpu),
  backend:  "router<direct<(fusion<jit<cuda>>, fusion<jit<wgpu>>)>>",
  kind:  "Float",
  dtype:  "f32",
}

[codecov]

Codecov Report

Attention: Patch coverage is 90.61910% with 497 lines in your changes missing coverage. Please review.

Project coverage is 85.23%. Comparing base (604dbae) to head (67c7479).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-router/src/ops/op_int.rs	92.17%	84 Missing ⚠️
crates/burn-router/src/runner.rs	92.00%	84 Missing ⚠️
crates/burn-router/src/bridge/byte.rs	29.91%	82 Missing ⚠️
crates/burn-router/src/ops/op_qfloat.rs	0.00%	54 Missing ⚠️
crates/burn-router/src/ops/op_bool.rs	89.09%	29 Missing ⚠️
crates/burn-router/src/channel/direct.rs	76.10%	27 Missing ⚠️
crates/burn-ndarray/src/backend.rs	50.00%	21 Missing ⚠️
crates/burn-router/src/backend.rs	66.10%	20 Missing ⚠️
crates/burn-fusion/src/backend.rs	24.00%	19 Missing ⚠️
crates/burn-router/src/client/base.rs	76.66%	14 Missing ⚠️
... and 10 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2353      +/-   ##
==========================================
+ Coverage   84.95%   85.23%   +0.28%     
==========================================
  Files         771      785      +14     
  Lines       98678   103756    +5078     
==========================================
+ Hits        83828    88437    +4609     
- Misses      14850    15319     +469     

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

[laggui]

Stoopid windows CI. Wgpu as a test backend doesn't seem to work on windows (auto device doesn't detect anything)

[laggui]

I just disabled the ByteBridge tests for the windows CI. Seems like this is a known issue with wgpu because we are actively setting DISABLE_WGPU=1 for windows.

The BackendRouter implementation is pretty much complete. Just a couple things to do or improve:

• Quantization ops (qtensor and q_* ops not implemented in this PR)
• DirectChannel is only implemented with support for two backends at this time
  • Improvements: maybe a macro to implement it for up to 4 backends (with appropriate types generated)