burn-cubecl 0.22.0

Generic backend that can be compiled just-in-time to any shader language target
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

use crate::{CubeBackend, CubeDevice, kernel, tensor::CubeTensor};
use burn_backend::tensor::{BoolTensor, FloatTensor, IntTensor, QuantizedTensor};
use burn_backend::{DType, Shape};
pub use burn_cubecl_fusion::{CubeFusionHandle, FallbackOperation, Tiles};
use burn_fusion::{
    FusionBackend, FusionRuntime,
    stream::{FallbackOp, OrderedExecution},
};
use burn_ir::{BackendIr, TensorHandle};

mod registry;
pub use burn_cubecl_fusion::optim::{CubeOptimization, CubeOptimizationState, FusedOperation};
pub use registry::{
    BUILTIN_NAMES, CubeFuser, OptimizationProvider, RegistryError, register, remove,
};

impl burn_fusion::Optimization<FusionCubeRuntime> for CubeOptimization {
    fn execute(
        &mut self,
        context: &mut burn_fusion::stream::Context<
            <FusionCubeRuntime as FusionRuntime>::FusionHandle,
        >,
        execution: &OrderedExecution<FusionCubeRuntime>,
    ) {
        self.run(context, &|index| {
            let operation = execution.operation_within_optimization(index);
            Box::new(FallbackOperationWrapper::new(operation))
        })
    }

    fn to_state(&self) -> CubeOptimizationState {
        Self::to_state(self)
    }

    fn from_state(device: &CubeDevice, state: CubeOptimizationState) -> Self {
        registry::restore(device, state)
    }
}

struct FallbackOperationWrapper<O: Clone> {
    operation: O,
}

impl<O: Clone> FallbackOperationWrapper<O> {
    fn new(op: O) -> Self {
        Self { operation: op }
    }
}

impl FallbackOperation for FallbackOperationWrapper<FallbackOp<FusionCubeRuntime>> {
    fn run(&self, context: &mut burn_fusion::stream::Context<CubeFusionHandle>) {
        // Through `FallbackOp`, so unfused work inside a fused block obeys the
        // same rule as unfused work outside one: an operation whose input a
        // failure claims does not run, and its outputs take that failure.
        self.operation.execute(&mut context.handles);
    }
}

impl BackendIr for CubeBackend {
    type Handle = CubeFusionHandle;

    fn float_tensor(handle: TensorHandle<Self::Handle>) -> FloatTensor<Self> {
        into_tensor(handle.handle, handle.shape)
    }

    fn int_tensor(handle: TensorHandle<Self::Handle>) -> IntTensor<Self> {
        into_tensor(handle.handle, handle.shape)
    }

    fn bool_tensor(handle: TensorHandle<Self::Handle>) -> BoolTensor<Self> {
        into_tensor(handle.handle, handle.shape)
    }

    fn quantized_tensor(handle: TensorHandle<Self::Handle>) -> QuantizedTensor<Self> {
        into_tensor(handle.handle, handle.shape)
    }

    fn float_tensor_handle(tensor: FloatTensor<Self>) -> Self::Handle {
        tensor.into()
    }

    fn int_tensor_handle(tensor: IntTensor<Self>) -> Self::Handle {
        tensor.into()
    }

    fn bool_tensor_handle(tensor: BoolTensor<Self>) -> Self::Handle {
        tensor.into()
    }

    fn quantized_tensor_handle(tensor: QuantizedTensor<Self>) -> Self::Handle {
        tensor.into()
    }
}

impl FusionRuntime for FusionCubeRuntime {
    type OptimizationState = CubeOptimizationState;
    type Optimization = CubeOptimization;
    type FusionHandle = CubeFusionHandle;
    type FusionDevice = CubeDevice;

    fn fusers(device: CubeDevice) -> Vec<Box<dyn burn_fusion::OperationFuser<Self::Optimization>>> {
        registry::fusers(&device)
    }
}

/// Fusion runtime for JIT runtimes.
#[derive(Debug)]
pub struct FusionCubeRuntime;

impl FusionBackend for CubeBackend {
    type FusionRuntime = FusionCubeRuntime;

    type FullPrecisionBackend = CubeBackend;

    fn cast_float(tensor: FloatTensor<Self>, dtype: DType) -> Self::Handle {
        kernel::cast(tensor, dtype).into()
    }

    fn memory_persistent(device: &Self::Device, enabled: bool) {
        use cubecl::MemoryAllocationMode;

        let client = device.client();
        let mode = match enabled {
            true => MemoryAllocationMode::Persistent,
            false => MemoryAllocationMode::Auto,
        };
        // Safety: called from the fusion execution thread, whose stream is the
        // one every fused operation allocates on.
        unsafe { client.allocation_mode(mode) };
    }
}

fn into_tensor(handle: CubeFusionHandle, shape: Shape) -> CubeTensor {
    CubeTensor {
        client: handle.client.clone(),
        handle: handle.handle.clone(),
        device: handle.device.clone(),
        meta: Box::new(handle.metadata(shape)),
        dtype: handle.dtype,
        qparams: handle.qparams.clone(),
    }
}

impl From<CubeTensor> for CubeFusionHandle {
    fn from(value: CubeTensor) -> Self {
        Self {
            client: value.client.clone(),
            handle: value.handle.clone(),
            device: value.device.clone(),
            strides: value.meta.strides.clone(),
            tiles: value.meta.is_tiled().then(|| Tiles {
                shape: value.meta.shape.clone(),
                tiling: value.meta.tiling,
            }),
            dtype: value.dtype,
            qparams: value.qparams.clone(),
        }
    }
}