Custom fusion

[ArthurBrussee]

Pull Request Template

Related Issues/PRs

#1970

Changes

Expose enough information to allow custom operations to use fusing. Also expose a function to resolve a tensor from a fusion backend.

Draft until tracel-ai/cubecl#256 lands & some other bits are cleaned up.

This also fixes #1970 by bailing on fusing when going over max bindings.

[nathanielsimard]

cloning every time is a bit wasteful. I would only do it if builder_current_state.estimate_bindings() + new_op.nodes() > self.max_bindings

[ArthurBrussee]

That doesn't really work - we don't know how many nodes someone is adding in the build function. And we can't easily keep track of it as adding them mutates the builder.

estimate_bindings() also seems more expensive than a clone in the first place?

I've at least added a fast path for the first operation which should be common. If that's not enough might need to rethink how these operations are added so we can count them more easily.

[ArthurBrussee]

I've also added a fix for a deadlock in fusion on wasm in this PR now. It might or might not make things run a bit smoothly on other platforms too, as it was incorrectly holding the server lock while waiting for a read.

[codecov]

Codecov Report

Attention: Patch coverage is 56.29921% with 111 lines in your changes missing coverage. Please review.

Project coverage is 82.91%. Comparing base (6e71aaf) to head (fc63b22).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/burn-tensor/src/repr/operation.rs	0.00%	28 Missing ⚠️
crates/burn-fusion/src/client/mutex.rs	35.13%	24 Missing ⚠️
crates/burn-fusion/src/stream/context.rs	0.00%	19 Missing ⚠️
crates/burn-fusion/src/server.rs	40.00%	18 Missing ⚠️
...s/burn-jit/src/kernel/conv/conv2d/implicit_gemm.rs	0.00%	9 Missing ⚠️
...urn-jit/src/kernel/conv/deform_conv_transpose2d.rs	0.00%	2 Missing ⚠️
...rates/burn-jit/src/kernel/reduce/subcube/kernel.rs	33.33%	2 Missing ⚠️
...tes/burn-jit/src/kernel/reduce/subcube/mean_dim.rs	0.00%	2 Missing ⚠️
...tes/burn-jit/src/kernel/reduce/subcube/prod_dim.rs	0.00%	2 Missing ⚠️
...ates/burn-jit/src/kernel/reduce/subcube/sum_dim.rs	0.00%	2 Missing ⚠️
... and 3 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2486      +/-   ##
==========================================
- Coverage   82.91%   82.91%   -0.01%     
==========================================
  Files         810      812       +2     
  Lines      104904   105334     +430     
==========================================
+ Hits        86984    87336     +352     
- Misses      17920    17998      +78     

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