Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Migrate to cubecl IR refactor #2418

Merged
merged 6 commits into from
Oct 30, 2024
Merged

[WIP] Migrate to cubecl IR refactor #2418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3ca9e06
Update to IR refator
wingertge Oct 23, 2024
dee64a2
Fix argmin/argmax and update cubecl
wingertge Oct 24, 2024
29200b9
Merge branch 'main' into feat/cubecl-ir-refactor
wingertge Oct 26, 2024
0026950
Merge branch 'main' into feat/cubecl-ir-refactor
wingertge Oct 30, 2024
8d6febb
Bump cubecl
wingertge Oct 30, 2024
d24048c
Update to cubecl changes
wingertge Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 78 additions & 50 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 4 additions & 4 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,14 +152,14 @@ ahash = { version = "0.8.11", default-features = false }
portable-atomic-util = { version = "0.2.2", features = ["alloc"] }

### For the main burn branch. ###
# cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "63da837b5ae78ff1a3b7363fe3af2b02c2bc864f" }
# cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "63da837b5ae78ff1a3b7363fe3af2b02c2bc864f" }
cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "99404b1e29946832a42b72a5c26d4cf42c67692e" }
cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "99404b1e29946832a42b72a5c26d4cf42c67692e" }
### For local development. ###
# cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
# cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }
### For the release. ###
cubecl = { version="0.3.0", default-features = false }
cubecl-common = { version="0.3.0", default-features = false }
# cubecl = { version = "0.3.0", default-features = false }
# cubecl-common = { version = "0.3.0", default-features = false }

### For xtask crate ###
tracel-xtask = { version = "~1.1" }
Expand Down
55 changes: 18 additions & 37 deletions crates/burn-jit/src/kernel/conv/conv2d/transpose_direct.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use cubecl::{
cpa,
ir::{Elem, IntKind, KernelDefinition, Scope, Variable, Visibility},
ir::{
Builtin, Elem, IntKind, Item, KernelDefinition, Scope, Variable, VariableKind, Visibility,
},
CubeCountSettings, Execution, InputInfo, KernelExpansion, KernelIntegrator, KernelSettings,
OutputInfo,
};
Expand Down Expand Up @@ -39,7 +41,7 @@ impl<E: JitElement> Conv2dTransposeComputeShader<E> {
let weight = self.weight;
let bias = self.bias;
let output = self.output;
let id = Variable::AbsolutePos;
let id = Variable::builtin(Builtin::AbsolutePos);

let input_stride_0 = scope.create_local(Elem::UInt);
let input_stride_1 = scope.create_local(Elem::UInt);
Expand Down Expand Up @@ -92,34 +94,13 @@ impl<E: JitElement> Conv2dTransposeComputeShader<E> {
cpa!(scope, kernel_size_0 = shape(weight, 2u32));
cpa!(scope, kernel_size_1 = shape(weight, 3u32));

let conv_stride_0 = Variable::GlobalScalar {
id: 0,
elem: Elem::UInt,
};
let conv_stride_1 = Variable::GlobalScalar {
id: 1,
elem: Elem::UInt,
};
let dilation_0 = Variable::GlobalScalar {
id: 2,
elem: Elem::UInt,
};
let dilation_1 = Variable::GlobalScalar {
id: 3,
elem: Elem::UInt,
};
let padding_0 = Variable::GlobalScalar {
id: 4,
elem: Elem::UInt,
};
let padding_1 = Variable::GlobalScalar {
id: 5,
elem: Elem::UInt,
};
let groups = Variable::GlobalScalar {
id: 6,
elem: Elem::UInt,
};
let conv_stride_0 = Variable::new(VariableKind::GlobalScalar(0), Item::new(Elem::UInt));
let conv_stride_1 = Variable::new(VariableKind::GlobalScalar(1), Item::new(Elem::UInt));
let dilation_0 = Variable::new(VariableKind::GlobalScalar(2), Item::new(Elem::UInt));
let dilation_1 = Variable::new(VariableKind::GlobalScalar(3), Item::new(Elem::UInt));
let padding_0 = Variable::new(VariableKind::GlobalScalar(4), Item::new(Elem::UInt));
let padding_1 = Variable::new(VariableKind::GlobalScalar(5), Item::new(Elem::UInt));
let groups = Variable::new(VariableKind::GlobalScalar(6), Item::new(Elem::UInt));

let stride_0_i = scope.create_local(Elem::Int(IntKind::I32));
let stride_1_i = scope.create_local(Elem::Int(IntKind::I32));
Expand Down Expand Up @@ -222,9 +203,9 @@ impl<E: JitElement> Conv2dTransposeComputeShader<E> {
cpa!(scope, index_input_b = b * input_stride_0);
cpa!(scope, index_weight_oc = oc * weight_stride_1);

let prod = scope.create_local(output.item());
let prod_tmp = scope.create_local(output.item());
let sum = scope.create_local(output.item());
let prod = scope.create_local(output.item);
let prod_tmp = scope.create_local(output.item);
let sum = scope.create_local(output.item);
cpa!(scope, sum = bias[oc_out]);

let kh = scope.create_local(Elem::UInt);
Expand Down Expand Up @@ -314,10 +295,10 @@ impl<R: JitRuntime, E: JitElement> Kernel for Conv2dTransposeEagerKernel<R, E> {
let mut scope = Scope::root();
let item = E::cube_elem().into();

let input = Variable::GlobalInputArray { id: 0, item };
let weight = Variable::GlobalInputArray { id: 1, item };
let bias = Variable::GlobalInputArray { id: 2, item };
let output = Variable::GlobalOutputArray { id: 0, item };
let input = Variable::new(VariableKind::GlobalInputArray(0), item);
let weight = Variable::new(VariableKind::GlobalInputArray(1), item);
let bias = Variable::new(VariableKind::GlobalInputArray(2), item);
let output = Variable::new(VariableKind::GlobalOutputArray(0), item);

scope.write_global_custom(output);

Expand Down
Loading