Skip to content

Commit 3887fd1

Browse files
wingertgewingertge
authored
refactor: Move from CubeOption to Option (#4543)
* Remove R::supported_line_sizes * refactor: Metadata optimization * Revert temp fix * Rename `ShapeError` to `MetadataError` * Cleanup * Bump cubecl and cubek rev * Fix doc test * Refactor `CubeOption`
1 parent a50709b commit 3887fd1

14 files changed

Lines changed: 349 additions & 385 deletions

File tree

‎Cargo.lock‎

Lines changed: 254 additions & 273 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -179,10 +179,10 @@ portable-atomic = { version = "1.13.1" }
179179
portable-atomic-util = { version = "0.2.5", features = ["alloc"] }
180180

181181
### For the main burn branch. ###
182-
cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "074fb0473a9441565844cfdfe5a4d8f7d62aa299" }
183-
cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "074fb0473a9441565844cfdfe5a4d8f7d62aa299" }
184-
cubecl-zspace = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "074fb0473a9441565844cfdfe5a4d8f7d62aa299" }
185-
cubek = { git = "https://github.com/tracel-ai/cubek", default-features = false, rev = "312ee613666a1847f87a232b7b1e589ca2995015" }
182+
cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0400cc7d6c0758f54c9485558d78585ddd22da78" }
183+
cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0400cc7d6c0758f54c9485558d78585ddd22da78" }
184+
cubecl-zspace = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "0400cc7d6c0758f54c9485558d78585ddd22da78" }
185+
cubek = { git = "https://github.com/tracel-ai/cubek", default-features = false, rev = "738891e604bf3c1ec8f3ddbbb3a0924d3ed985f0" }
186186
### For local development. ###
187187
# cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
188188
# cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

‎crates/burn-cubecl-fusion/src/optim/matmul/args.rs‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use cubecl::{
99
prelude::*,
1010
quant::scheme::{QuantLevel, QuantScheme},
1111
std::{
12-
CubeOption, CubeOptionExpand, FastDivmod,
12+
FastDivmod,
1313
quant::{
1414
RunWithQuantType,
1515
view::{QuantizedView, run_with_quant_type},
@@ -91,13 +91,13 @@ impl MatmulArgs for FusedMatmulArgs {
9191
comptime![inputs.config.clone()],
9292
);
9393
let batch_acc = match comptime![inputs.c.clone()] {
94-
Option::Some(c) => CubeOption::new_Some(input_batch_layout(
94+
Some(c) => Option::Some(input_batch_layout(
9595
&inputs.global,
9696
&batch_shape,
9797
comptime![c],
9898
comptime![inputs.config.clone()],
9999
)),
100-
Option::None => CubeOption::new_None(),
100+
None => Option::new_None(),
101101
};
102102
let batch_out = BatchLayout::new(batch_strides_out, batch_shape.clone());
103103

@@ -159,9 +159,9 @@ impl MatmulArgs for FusedMatmulArgs {
159159

160160
fn view_acc<Lhs: Numeric, Rhs: Numeric, EO: Numeric>(
161161
state: &Self::State<Lhs, Rhs, EO>,
162-
) -> CubeOption<View<Line<EO>, BatchedCoords>> {
162+
) -> Option<View<Line<EO>, BatchedCoords>> {
163163
match comptime![state.c.clone()] {
164-
Option::Some(c) => {
164+
Some(c) => {
165165
let view = global_view(
166166
&state.inputs,
167167
&state.locals,
@@ -170,9 +170,9 @@ impl MatmulArgs for FusedMatmulArgs {
170170
comptime![state.config.clone()],
171171
comptime![state.out_layout_config],
172172
);
173-
CubeOption::new_Some(view)
173+
Option::Some(view)
174174
}
175-
Option::None => CubeOption::new_None(),
175+
None => Option::new_None(),
176176
}
177177
}
178178

@@ -181,8 +181,8 @@ impl MatmulArgs for FusedMatmulArgs {
181181
batch: usize,
182182
) -> usize {
183183
match state.c_batch {
184-
CubeOption::Some(c_batch) => c_batch.to_source_pos(batch),
185-
CubeOption::None => batch,
184+
Some(c_batch) => c_batch.to_source_pos(batch),
185+
None => batch,
186186
}
187187
}
188188

@@ -453,7 +453,7 @@ pub struct FusedMatmulState {
453453
locals: LocalArgs,
454454
a_batch: VirtualLayout<Coords1d, Coords1d>,
455455
b_batch: VirtualLayout<Coords1d, Coords1d>,
456-
c_batch: CubeOption<VirtualLayout<Coords1d, Coords1d>>,
456+
c_batch: Option<VirtualLayout<Coords1d, Coords1d>>,
457457
out_batch: VirtualLayout<Coords1d, Coords1d>,
458458
#[cube(comptime)]
459459
config: FuseBlockConfig,
@@ -483,7 +483,7 @@ impl FusedMatmulState {
483483
locals: &mut LocalArgs,
484484
a_batch: VirtualLayout<usize, usize>,
485485
b_batch: VirtualLayout<usize, usize>,
486-
c_batch: CubeOption<VirtualLayout<usize, usize>>,
486+
c_batch: Option<VirtualLayout<usize, usize>>,
487487
out_batch: VirtualLayout<usize, usize>,
488488
batch_shape: Sequence<FastDivmod<u32>>,
489489
#[comptime] config: &FuseBlockConfig,

‎crates/burn-cubecl-fusion/src/optim/matmul/optimization.rs‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@ impl FusedMatmulLaunch<'_> {
478478
config.clone(),
479479
self.matmul.lhs.clone(),
480480
self.matmul.rhs.clone(),
481-
Option::None,
481+
None,
482482
self.matmul.out.clone(),
483483
),
484484
outputs,
@@ -502,7 +502,7 @@ impl FusedMatmulLaunch<'_> {
502502
config.clone(),
503503
self.matmul.lhs.clone(),
504504
self.matmul.rhs.clone(),
505-
Option::None,
505+
None,
506506
self.matmul.out.clone(),
507507
),
508508
outputs,
@@ -529,7 +529,7 @@ impl FusedMatmulLaunch<'_> {
529529
config.clone(),
530530
self.matmul.lhs.clone(),
531531
self.matmul.rhs.clone(),
532-
Option::None,
532+
None,
533533
self.matmul.out.clone(),
534534
),
535535
outputs,
@@ -553,7 +553,7 @@ impl FusedMatmulLaunch<'_> {
553553
config.clone(),
554554
self.matmul.lhs.clone(),
555555
self.matmul.rhs.clone(),
556-
Option::None,
556+
None,
557557
self.matmul.out.clone(),
558558
),
559559
outputs,
@@ -573,7 +573,7 @@ impl FusedMatmulLaunch<'_> {
573573
config.clone(),
574574
self.matmul.lhs.clone(),
575575
self.matmul.rhs.clone(),
576-
Option::None,
576+
None,
577577
self.matmul.out.clone(),
578578
),
579579
outputs,
@@ -593,7 +593,7 @@ impl FusedMatmulLaunch<'_> {
593593
config.clone(),
594594
self.matmul.lhs.clone(),
595595
self.matmul.rhs.clone(),
596-
Option::None,
596+
None,
597597
self.matmul.out.clone(),
598598
),
599599
outputs,
@@ -613,7 +613,7 @@ impl FusedMatmulLaunch<'_> {
613613
config.clone(),
614614
self.matmul.lhs.clone(),
615615
self.matmul.rhs.clone(),
616-
Option::None,
616+
None,
617617
self.matmul.out.clone(),
618618
),
619619
outputs,

‎crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/launch.rs‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,8 @@ use crate::{
1010
use cubecl::{
1111
Runtime,
1212
ir::{ElemType, FloatKind, StorageType},
13-
prelude::{ComputeClient, ScalarArg, SequenceArg},
13+
prelude::*,
1414
server::LaunchError,
15-
std::CubeOptionArgs,
1615
};
1716
use cubek::reduce::{
1817
LineMode, ReduceDtypes,
@@ -113,10 +112,10 @@ impl<R: Runtime> TraceRunner<R> for FusedReduceBroadcastedLaunch<'_> {
113112
}
114113

115114
let block_end = match configs.len() > index {
116-
true => CubeOptionArgs::Some(ElemwiseFuseBlockLaunch::new(
115+
true => OptionArgs::Some(ElemwiseFuseBlockLaunch::new(
117116
configs.last().cloned().unwrap(),
118117
)),
119-
false => CubeOptionArgs::None,
118+
false => OptionArgs::None,
120119
};
121120

122121
// TODO: Ensure parallel is selected.

‎crates/burn-cubecl-fusion/src/optim/reduce_broadcasted/unit.rs‎

Lines changed: 24 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,7 @@ use crate::{
55
},
66
optim::reduce::args::{FusedReduceArgs, FusedReduceInput, FusedReduceOutput},
77
};
8-
use cubecl::{
9-
Runtime,
10-
prelude::*,
11-
std::{CubeOption, CubeOptionExpand, tensor::r#virtual::VirtualTensor},
12-
};
8+
use cubecl::{Runtime, prelude::*, std::tensor::r#virtual::VirtualTensor};
139
use cubek::reduce::{
1410
LineMode, ReduceInstruction, ReducePrecision,
1511
components::{
@@ -66,7 +62,7 @@ pub fn reduce_kernel_broadcasted(
6662
outputs: &mut GlobalArgs,
6763
reduce_axis: usize,
6864
blocks: Sequence<ReduceFuseBlock>,
69-
block_end: CubeOption<ElemwiseFuseBlock>,
65+
block_end: Option<ElemwiseFuseBlock>,
7066
) {
7167
#[unroll]
7268
for i in 0..blocks.len() {
@@ -122,7 +118,7 @@ fn reduce_many(
122118
outputs: &mut GlobalArgs,
123119
reduce_axis: usize,
124120
blocks: Sequence<ReduceFuseBlock>,
125-
block_end: CubeOption<ElemwiseFuseBlock>,
121+
block_end: Option<ElemwiseFuseBlock>,
126122
) {
127123
let mut axis_size = 0;
128124

@@ -155,31 +151,28 @@ fn reduce_many(
155151
);
156152
}
157153

158-
match block_end {
159-
CubeOption::Some(block) => {
160-
let global_index = ABSOLUTE_POS;
161-
let width = comptime!(block.config.width as u32);
162-
let num_iter = axis_size / usize::cast_from(width);
163-
164-
for i in 0..num_iter {
165-
// Register block local inputs.
166-
let values = Registry::<FuseArg, Line<f32>>::new();
167-
let args = comptime![Vec::<FuseArg>::new()];
168-
let index = global_index * num_iter + i;
169-
let mut locals = init_locals(inputs, outputs, &block.config);
170-
171-
fuse_on_write::<f32>(
172-
inputs,
173-
outputs,
174-
&mut locals,
175-
index,
176-
values,
177-
args,
178-
&block.config.clone(),
179-
)
180-
}
154+
if let Some(block) = block_end {
155+
let global_index = ABSOLUTE_POS;
156+
let width = comptime!(block.config.width as u32);
157+
let num_iter = axis_size / usize::cast_from(width);
158+
159+
for i in 0..num_iter {
160+
// Register block local inputs.
161+
let values = Registry::<FuseArg, Line<f32>>::new();
162+
let args = comptime![Vec::<FuseArg>::new()];
163+
let index = global_index * num_iter + i;
164+
let mut locals = init_locals(inputs, outputs, &block.config);
165+
166+
fuse_on_write::<f32>(
167+
inputs,
168+
outputs,
169+
&mut locals,
170+
index,
171+
values,
172+
args,
173+
&block.config.clone(),
174+
)
181175
}
182-
CubeOption::None => {}
183176
}
184177
}
185178

‎crates/burn-cubecl/src/kernel/conv/conv_transpose2d/col2im.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use burn_backend::{
1717
use cubecl::{
1818
calculate_cube_count_elemwise,
1919
prelude::*,
20-
std::{CubeOption, CubeOptionExpand, FastDivmod, tensor::layout::linear::LinearView},
20+
std::{FastDivmod, tensor::layout::linear::LinearView},
2121
};
2222
use cubek::convolution::components::ConvSetupError;
2323

@@ -240,7 +240,7 @@ struct Col2ImArgs {
240240
#[cube(launch_unchecked, address_type = "dynamic")]
241241
fn col2im_kernel<E: Numeric>(
242242
columns: &Tensor<E>,
243-
bias: &CubeOption<Tensor<E>>,
243+
bias: &Option<Tensor<E>>,
244244
image: &mut LinearView<E, ReadWrite>,
245245
image_shape: Sequence<FastDivmod<usize>>,
246246
args: &Col2ImArgs,
@@ -296,7 +296,7 @@ fn col2im_kernel<E: Numeric>(
296296
}
297297

298298
match bias {
299-
CubeOption::Some(bias) => image[ABSOLUTE_POS] = val + bias[ch_im],
300-
CubeOption::None => image[ABSOLUTE_POS] = val,
299+
Some(bias) => image[ABSOLUTE_POS] = val + bias[ch_im],
300+
None => image[ABSOLUTE_POS] = val,
301301
}
302302
}

‎crates/burn-cubecl/src/kernel/conv/conv_transpose2d/transpose_direct.rs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use burn_backend::{Shape, ops::ConvTransposeOptions};
88
use cubecl::{
99
calculate_cube_count_elemwise,
1010
prelude::*,
11-
std::{CubeOption, CubeOptionExpand, FastDivmod, tensor::layout::linear::LinearView},
11+
std::{FastDivmod, tensor::layout::linear::LinearView},
1212
};
1313
use cubek::convolution::components::ConvSetupError;
1414

@@ -27,7 +27,7 @@ struct ConvArgs {
2727
fn conv_transpose2d_direct_kernel<E: Numeric>(
2828
input: &Tensor<E>,
2929
weight: &Tensor<E>,
30-
bias: &CubeOption<Tensor<E>>,
30+
bias: &Option<Tensor<E>>,
3131
output: &mut LinearView<E, ReadWrite>,
3232
out_shape: Sequence<FastDivmod<usize>>,
3333
args: ConvArgs,
@@ -71,10 +71,8 @@ fn conv_transpose2d_direct_kernel<E: Numeric>(
7171
let idx_input_batch = batch * input.stride(0);
7272
let idx_weight_oc = out_c * weight.stride(1);
7373

74-
let mut sum = match bias {
75-
CubeOption::Some(bias) => bias[oc_out],
76-
CubeOption::None => E::from_int(0),
77-
};
74+
let bias: Option<E> = bias.map(|bias| bias[oc_out]);
75+
let mut sum = bias.unwrap_or_default();
7876

7977
let numerator_h_base = out_y + args.padding_0;
8078
let numerator_w_base = out_x + args.padding_1;

‎crates/burn-cubecl/src/kernel/conv/conv_transpose3d.rs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use cubecl::{
22
calculate_cube_count_elemwise,
33
prelude::*,
4-
std::{CubeOption, CubeOptionExpand, FastDivmod, tensor::layout::linear::LinearView},
4+
std::{FastDivmod, tensor::layout::linear::LinearView},
55
};
66

77
use crate::{
@@ -30,7 +30,7 @@ struct ConvArgs {
3030
fn conv_transpose3d_kernel<E: Numeric>(
3131
input: &Tensor<E>,
3232
weight: &Tensor<E>,
33-
bias: &CubeOption<Tensor<E>>,
33+
bias: &Option<Tensor<E>>,
3434
output: &mut LinearView<E, ReadWrite>,
3535
out_shape: Sequence<FastDivmod<usize>>,
3636
args: ConvArgs,
@@ -80,10 +80,8 @@ fn conv_transpose3d_kernel<E: Numeric>(
8080
let index_input_batch = batch * input.stride(0);
8181
let index_weight_out_c = out_channel * weight.stride(1);
8282

83-
let mut sum = match bias {
84-
CubeOption::Some(bias) => bias[out_c_out],
85-
CubeOption::None => E::from_int(0),
86-
};
83+
let bias: Option<E> = bias.map(|bias| bias[out_c_out]);
84+
let mut sum = bias.unwrap_or_default();
8785

8886
let numerator_d_base = out_z + args.padding_0;
8987
let numerator_h_base = out_y + args.padding_1;

0 commit comments

Comments
 (0)