Fusion 装饰器

impl<B: FusionBackend> BackendTypes for Fusion<B> {
    type Device = B::Device;

    type FloatTensorPrimitive = FusionTensor<B::FusionRuntime>;
    type IntTensorPrimitive = FusionTensor<B::FusionRuntime>;
    type BoolTensorPrimitive = FusionTensor<B::FusionRuntime>;
    type QuantizedTensorPrimitive = FusionTensor<B::FusionRuntime>;
    // ...
}

dispatch<fusion<cubecl<wgpu<msl>>>>
dispatch<cubecl<wgpu<msl>>>

        impl<B: FusionBackend> Operation<B::FusionRuntime> for $name<B> {
            fn execute(
                &self,
                handles: &mut HandleContainer<B::Handle>,
            ) -> Result<(), $crate::ExecutionError> {
                let lhs = handles.get_float_tensor::<B>(&self.desc.lhs);
                let rhs = handles.get_float_tensor::<B>(&self.desc.rhs);
                let output = $ops(lhs, rhs);

                handles.register_float_tensor::<B>(&self.desc.out.id, output);

                Ok(())
            }
        }

        // Create output tensors returned by this operation
        let outputs = repr
            .outputs()
            .map(|output| {
                FusionTensor::new(
                    output.id,
                    output.shape.clone(),
                    output.dtype,
                    self.clone(),
                    stream,
                )
            })
            .collect();

        // By doing this comparison, we reduce the number of bytes transferred in the device handle
        // queue.
        if size_of::<O>() < size_of::<UnfusedOp<R>>() {
            // Here the [`O`] type is smaller than the [`UnfusedOp`] type, so it's better to
            // transfer it directly.
            self.server.submit(move |server| {
                let operation = UnfusedOp::new(operation, stream);
                server.register(stream, repr, operation);
            });
        } else {
            // ...
        }

        outputs

marker("from_floats x3");
let a = small(&device, 1.0);
let b = small(&device, 2.0);
let c = small(&device, 3.0);
marker("let d = (a * b + c).exp()");
let d = (a * b + c).exp();
marker("d.into_data()");
let v = d.into_data().to_vec::<f32>().unwrap();
marker(&format!("result {v:?}"));

--- from_floats x3
--- let d = (a * b + c).exp()
--- d.into_data()
  [registered] Init in=[] out=[t1:NotInit]
  [block] 1 op(s): Init
  [registered] Init in=[] out=[t2:NotInit]
  [block] 1 op(s): Init
  [registered] Init in=[] out=[t3:NotInit]
  [block] 1 op(s): Init
  [registered] NumericFloat(F32, Mul) in=[t1:ReadWrite, t2:ReadWrite] out=[t4:NotInit]
  [registered] NumericFloat(F32, Add) in=[t4:ReadWrite, t3:ReadWrite] out=[t5:NotInit]
  [registered] Float(F32, Exp) in=[t5:ReadWrite] out=[t6:NotInit]
  [block] 3 op(s): NumericFloat(F32, Mul), NumericFloat(F32, Add), Float(F32, Exp)
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
--- result [148.41316, 22026.467, 24154958.0, 195729570000.0]

let x = (a * 2.0).exp();
let h = std::thread::spawn(move || {
  println!("spawned thread stream: {}", StreamId::current());
  let y = x.clone() + 1.0;
  settle();
  println!("--- spawned: y.into_data()");
  let v = y.into_data().to_vec::<f32>().unwrap();
  // ...
});

main thread stream: StreamId(0)
--- main: let x = (a * 2.0).exp()
--- spawn: let y = x.clone() + 1.0 on another thread
spawned thread stream: StreamId(1)
  [registered] Init in=[] out=[t1:NotInit]
  [block] 1 op(s): Init
  [registered] NumericFloat(F32, MulScalar) in=[t1:ReadWrite] out=[t2:NotInit]
  [registered] Float(F32, Exp) in=[t2:ReadWrite] out=[t3:NotInit]
  [block] 2 op(s): NumericFloat(F32, MulScalar), Float(F32, Exp)
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
--- spawned: y.into_data()
  [registered] NumericFloat(F32, AddScalar) in=[t4:ReadWrite] out=[t5:NotInit]
  [block] 1 op(s): NumericFloat(F32, AddScalar)
  [launch] [DSU-100-16384] burn_cubecl::kernel::binary::kernel_scalar_binop::KernelScalarBinop<burn_cubecl::kernel::binary::AddOp>
--- spawned: result [8.389055, 55.59815, 404.42868, 2981.9578]

--- let y = chain(x, 100)
  [block] 1 op(s), first Init
  [block] 65 op(s), first NumericFloat(F32, MulScalar)
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
--- y.into_data()
  [block] 35 op(s), first NumericFloat(F32, AddScalar)
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
--- done

pub struct FusionTensor<R: FusionRuntime> {
    /// Tensor id.
    pub id: TensorId,
    /// The shape of the tensor.
    pub shape: Shape,
    /// The fusion client.
    pub client: Client<R>,
    /// The datatype of the tensor.
    pub dtype: DType,
    /// The current stream id this tensor is on.
    pub stream: StreamId,
    pub(crate) count: Arc<AtomicU32>,
}

    pub fn into_ir(mut self) -> TensorIr {
        let current = StreamId::current();
        if self.stream != current {
            self = self.shared_view(current);
        }

        let count = self.count.load(Ordering::Acquire);
        let status = self.status(count);

        let mut shape_out = Shape::from(Vec::<usize>::new());
        core::mem::swap(&mut self.shape, &mut shape_out);

        if let TensorStatus::ReadWrite = status {
            // Avoids an unwanted drop on the same thread.
            //
            // Since `drop` is called after `into_ir`, we must not register a drop if the tensor
            // was consumed with a `ReadWrite` status.
            self.count.fetch_add(1, Ordering::Acquire);
        }

        TensorIr {
            status,
            shape: shape_out,
            id: self.id,
            dtype: self.dtype,
        }
    }

        let register = move || {
            if StreamId::current() == stream {
                client.register(stream, OperationIr::Drop(ir), DropOp { id });
            } else {
                client.foreign_drop(stream, ir);
            }
        };

        // ...
        if std::thread::panicking() {
            deferred::defer(register);
        } else {
            register();
        }

let d = (a.clone() * b.clone() + c.clone()).exp();
drop(d);
marker("device.sync()");
device.sync().unwrap();
marker(&format!("launches since start: {}", take_launches().0));

let a = Tensor::<1>::from_floats([1.0, 2.0], &metal);
let b = Tensor::<1>::from_floats([3.0, 4.0], &other); // Device::default(), another Cube device
let bad = Tensor::cat(vec![a.clone(), b], 0);
let worse = bad.clone() + 1.0;
let good = a * 10.0;
println!("{:?}", good.try_to_data().map(|d| d.to_vec::<f32>().unwrap()));
// then bad.try_to_data(), worse.try_to_data()

--- register: bad = cat([a, b]); worse = bad + 1; good = a * 10
--- good.try_to_data()

thread 'DSU-100-16384' (3347454) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-ir-0.22.1/src/handle.rs:333:32:
Should have handle for tensor TensorId { value: 3 }
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Ok([10.0, 20.0])
--- bad.try_to_data()
Err: An error happened during execution
Caused by:
  Should have handle for tensor TensorId { value: 3 }
--- worse.try_to_data()
Err: An error happened during execution
Caused by:
  Should have handle for tensor TensorId { value: 3 }

dispatch<cubecl<wgpu<msl>>>
run 0: 8 ops -> 8 launch(es)
       2 burn_cubecl::kernel::binary::kernel_scalar_binop::KernelScalarBinop<burn_cubecl::kernel::binary::Mul...
       2 burn_cubecl::kernel::binary::kernel_scalar_binop::KernelScalarBinop<burn_cubecl::kernel::binary::Add...
       4 burn_cubecl::kernel::unary_float::unary_float::UnaryFloat<burn_cubecl::kernel::unary_float::unary_ba...
run 1: 8 ops -> 8 launch(es)
run 2: 8 ops -> 8 launch(es)
...
dispatch<fusion<cubecl<wgpu<msl>>>>
run 0: 8 ops -> 1 launch(es)
       1 burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
run 1: 8 ops -> 1 launch(es)
run 2: 8 ops -> 1 launch(es)
...

--- let d = (a * b + c).exp(); drop(d)
--- device.sync()
  [registered] Init in=[] out=[t1:NotInit]
  [block] 1 op(s): Init
  [registered] Init in=[] out=[t2:NotInit]
  [block] 1 op(s): Init
  [registered] Init in=[] out=[t3:NotInit]
  [block] 1 op(s): Init
  [registered] NumericFloat(F32, Mul) in=[t1:ReadOnly, t2:ReadOnly] out=[t4:NotInit]
  [registered] NumericFloat(F32, Add) in=[t4:ReadWrite, t3:ReadOnly] out=[t5:NotInit]
  [registered] Float(F32, Exp) in=[t5:ReadWrite] out=[t6:NotInit]
  [registered] Drop in=[t6:ReadWrite] out=[]
  [block] 4 op(s): NumericFloat(F32, Mul), NumericFloat(F32, Add), Float(F32, Exp), Drop
--- launches since start: 0