Add option to flush queue instead of waiting for completion.

[ArthurBrussee]

Checklist

• Confirmed that run-checks all script has been executed.

Related Issues/PRs

This adds the last bit of infrastructure needed for custom burn<->wgpu interop, after initializing from an existing wgpu server (#1788), and getting WGPU resources (#1861).

Changes

After this PR, the setup for custom interop roughly looks like:

init_existing_device(adapter, device, wgpu_queue);

// Do some burn operations.
let x = Tensor::from_floats();
let y = x * 2 + 5.0; // Some calculations we really need in our wgpu code.

// Or however you have access to the client.
let client = y.into_primitive().client;

// Flush burn operations - they now live on the wgpu queue.
client.submit();

// Get a reference to y binding.
let y_wgpu_binding= client.get_resource(y);

// Run custom commands on the queue. wgpu makes guarantees that any resources have been
// calculated up to this point.
run_my_custom_wgpu_code(wgpu_queue, y_wgpu_binding);

This does require a seperate encoder for your custom wgpu code. I think that's ok - a custom wgpu app might have it's own already anyway - but in the future we could expose some notion of the servers encoder to share resources.

[nathanielsimard]

Do we want to have a submit function here or pass an option to the sync method? I'm unsure which one is the most intuitive. It's a bit weird to have an exposed submit function along execute and sync.

fn sync(&self, options: SyncOptions);

#[derive(Clone, Copy)]
enum SyncOptions {
    /// Sync the execution queue.
    Queue,
    //// Sync the execution queue and wait for all tasks to be completed.
    Full,
}

[ArthurBrussee]

Ah alright! Yeah that might be more consistent.

It does slightly complicate things as for better or worse it kind of forces this to be exposed on the backend trait.

It's also not quite clear where to keep something like SyncOptions as it's used in a bunch of crates that can't really depend on each other.

Lastly, changed the names a bit ("type" instead of options as there might be other options - flush&wait vs queue&full - the former being the more common terminlogy in the graphics world anyway).

Lmk how this looks!

[codecov]

Codecov Report

Attention: Patch coverage is 42.30769% with 45 lines in your changes missing coverage. Please review.

Project coverage is 86.10%. Comparing base (4b174a8) to head (85c6cb1).
Report is 4 commits behind head on main.

Files	Patch %	Lines
crates/burn-candle/src/backend.rs	0.00%	15 Missing ⚠️
crates/burn-tch/src/backend.rs	0.00%	14 Missing ⚠️
crates/burn-compute/src/channel/mpsc.rs	0.00%	7 Missing ⚠️
crates/burn-autodiff/src/backend.rs	0.00%	2 Missing ⚠️
crates/burn-compute/src/channel/cell.rs	0.00%	2 Missing ⚠️
crates/burn-fusion/src/backend.rs	0.00%	2 Missing ⚠️
crates/burn-jit/src/backend.rs	0.00%	2 Missing ⚠️
crates/burn-tensor/src/tensor/backend/base.rs	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1864      +/-   ##
==========================================
- Coverage   86.24%   86.10%   -0.14%     
==========================================
  Files         774      776       +2     
  Lines       90159    90410     +251     
==========================================
+ Hits        77759    77851      +92     
- Misses      12400    12559     +159     

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

[nathanielsimard]

Awesome! 👏