burn_cubecl/
fusion.rs

1use crate::{CubeBackend, CubeDevice, kernel, tensor::CubeTensor};
2use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor};
3use burn_backend::{DType, Shape};
4pub use burn_cubecl_fusion::{CubeFusionHandle, FallbackOperation, Tiles};
5use burn_fusion::{
6    FusionBackend, FusionRuntime,
7    stream::{FallbackOp, OrderedExecution},
8};
9use burn_ir::{BackendIr, TensorHandle};
10
11mod registry;
12pub use burn_cubecl_fusion::optim::{CubeOptimization, CubeOptimizationState, FusedOperation};
13pub use registry::{
14    BUILTIN_NAMES, CubeFuser, OptimizationProvider, RegistryError, register, remove,
15};
16
17impl burn_fusion::Optimization<FusionCubeRuntime> for CubeOptimization {
18    fn execute(
19        &mut self,
20        context: &mut burn_fusion::stream::Context<
21            <FusionCubeRuntime as FusionRuntime>::FusionHandle,
22        >,
23        execution: &OrderedExecution<FusionCubeRuntime>,
24    ) {
25        self.run(context, &|index| {
26            let operation = execution.operation_within_optimization(index);
27            Box::new(FallbackOperationWrapper::new(operation))
28        })
29    }
30
31    fn to_state(&self) -> CubeOptimizationState {
32        Self::to_state(self)
33    }
34
35    fn from_state(device: &CubeDevice, state: CubeOptimizationState) -> Self {
36        registry::restore(device, state)
37    }
38}
39
40struct FallbackOperationWrapper<O: Clone> {
41    operation: O,
42}
43
44impl<O: Clone> FallbackOperationWrapper<O> {
45    fn new(op: O) -> Self {
46        Self { operation: op }
47    }
48}
49
50impl FallbackOperation for FallbackOperationWrapper<FallbackOp<FusionCubeRuntime>> {
51    fn run(&self, context: &mut burn_fusion::stream::Context<CubeFusionHandle>) {
52        // Through `FallbackOp`, so unfused work inside a fused block obeys the
53        // same rule as unfused work outside one: an operation whose input a
54        // failure claims does not run, and its outputs take that failure.
55        self.operation.execute(&mut context.handles);
56    }
57}
58
59impl BackendIr for CubeBackend {
60    type Handle = CubeFusionHandle;
61
62    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FloatTensor<Self> {
63        into_tensor(handle.handle, handle.shape)
64    }
65
66    fn int_tensor(handle: TensorHandle<Self::Handle>) -> IntTensor<Self> {
67        into_tensor(handle.handle, handle.shape)
68    }
69
70    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> BoolTensor<Self> {
71        into_tensor(handle.handle, handle.shape)
72    }
73
74    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> QuantizedTensor<Self> {
75        into_tensor(handle.handle, handle.shape)
76    }
77
78    fn float_tensor_handle(tensor: FloatTensor<Self>) -> Self::Handle {
79        tensor.into()
80    }
81
82    fn int_tensor_handle(tensor: IntTensor<Self>) -> Self::Handle {
83        tensor.into()
84    }
85
86    fn bool_tensor_handle(tensor: BoolTensor<Self>) -> Self::Handle {
87        tensor.into()
88    }
89
90    fn quantized_tensor_handle(tensor: QuantizedTensor<Self>) -> Self::Handle {
91        tensor.into()
92    }
93}
94
95impl FusionRuntime for FusionCubeRuntime {
96    type OptimizationState = CubeOptimizationState;
97    type Optimization = CubeOptimization;
98    type FusionHandle = CubeFusionHandle;
99    type FusionDevice = CubeDevice;
100
101    fn fusers(device: CubeDevice) -> Vec<Box<dyn burn_fusion::OperationFuser<Self::Optimization>>> {
102        registry::fusers(&device)
103    }
104}
105
106/// Fusion runtime for JIT runtimes.
107#[derive(Debug)]
108pub struct FusionCubeRuntime;
109
110impl FusionBackend for CubeBackend {
111    type FusionRuntime = FusionCubeRuntime;
112
113    type FullPrecisionBackend = CubeBackend;
114
115    fn cast_float(tensor: FloatTensor<Self>, dtype: DType) -> Self::Handle {
116        kernel::cast(tensor, dtype).into()
117    }
118
119    fn memory_persistent(device: &Self::Device, enabled: bool) {
120        use cubecl::MemoryAllocationMode;
121
122        let client = device.client();
123        let mode = match enabled {
124            true => MemoryAllocationMode::Persistent,
125            false => MemoryAllocationMode::Auto,
126        };
127        // Safety: called from the fusion execution thread, whose stream is the
128        // one every fused operation allocates on.
129        unsafe { client.allocation_mode(mode) };
130    }
131}
132
133fn into_tensor(handle: CubeFusionHandle, shape: Shape) -> CubeTensor {
134    CubeTensor {
135        client: handle.client.clone(),
136        handle: handle.handle.clone(),
137        device: handle.device.clone(),
138        meta: Box::new(handle.metadata(shape)),
139        dtype: handle.dtype,
140        qparams: handle.qparams.clone(),
141    }
142}
143
144impl From<CubeTensor> for CubeFusionHandle {
145    fn from(value: CubeTensor) -> Self {
146        Self {
147            client: value.client.clone(),
148            handle: value.handle.clone(),
149            device: value.device.clone(),
150            strides: value.meta.strides.clone(),
151            tiles: value.meta.is_tiled().then(|| Tiles {
152                shape: value.meta.shape.clone(),
153                tiling: value.meta.tiling,
154            }),
155            dtype: value.dtype,
156            qparams: value.qparams.clone(),
157        }
158    }
159}
burn_cubecl/fusion.rs

burn_cubecl/
fusion.rs
