[burn 0.21 / burn-fusion] Feature request: public API to construct a FusionTensor from a base-backend tensor (Client::adopt_base_tensor)
Environment
- burn 0.21.0 (burn-fusion 0.21.0, burn-cubecl 0.21.0, burn-cuda 0.21.0)
- Any
Fusion<B>backend; we hit this withburn_cuda::Cuda(which aliasesFusion<CubeBackend<CudaRuntime, …>>in default builds)
Summary
The fusion → base direction is already publicly solvable (thank you — this works well):
// burn_fusion::client::Client::resolve_tensor_float — pub
let base: B::FloatTensorPrimitive = fusion_tensor.client
.resolve_tensor_float::<B>(fusion_tensor); // device-side, drains the streamThe base → fusion direction has no public path:
FusionTensor::new(id, shape, dtype, client, stream)ispub(crate);- the
count: Arc<AtomicU32>field is private, so the struct cannot even be built via literal syntax from outside the crate; - we found no public API that registers an existing base tensor handle into the fusion pool.
What we need
One of:
Client::adopt_base_tensor(handle: B::FloatTensorPrimitive) -> FusionTensor<R>— the mirror ofresolve_tensor_float(attach the base primitive to the fusion client, fresh id, current stream), or- making
FusionTensor::newpub(plus a public constructor forcount), or - a
Tensor<Fusion<B>, D>::from_base(Tensor<B, D>)convenience on the burn side.
Use case (why the reverse direction matters)
We ship custom #[cube(launch_unchecked)] kernels that operate on raw
ComputeClient + Handle (obtained from the unfused CubeBackend's
CubeTensor, whose client/handle fields are public — thank you). For
models running on the default fusion-wrapped backend:
- inputs are now solved device-side via
resolve_tensor_float(zero host round-trip); - outputs of a bridge kernel live on the unfused backend. Returning them
to the fusion model currently requires a host round-trip
(
into_data()→from_data()), because there is no public way to wrap a base tensor back into aFusionTensor.
With adopt_base_tensor, a bridge kernel's result would re-enter the fusion
graph device-side and the whole pattern becomes allocation-free in both
directions.
Minimal demonstration of the asymmetry
use burn::tensor::{backend::Backend, Tensor};
type Unfused = CubeBackend<CudaRuntime, f32, i32, u8>;
type Fused = burn_cuda::Cuda; // Fusion<Unfused> in default builds
// fusion → base: public, device-side ✅
let base_t: Tensor<Unfused, 2> = fusion_to_cube(&fusion_t);
// base → fusion: no public path ❌
// let fusion_t: Tensor<Fused, 2> = Tensor::from_primitive(
// FusionTensor::new(...) // pub(crate) — cannot construct
// );Workaround we ship today
Host round-trip for outputs (Tensor::into_data() on the unfused tensor,
Tensor::from_data() on the fusion backend). Correct but costs a
device→host→device copy per bridge call.
Notes
- We deliberately do not ask to pierce the fusion pool internals — only to adopt an already-materialized base primitive under a fresh id.
- If a public path already exists that we missed, a pointer would fully resolve this request.
Source: tracel-ai/burn