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

Jump to bottom

[Refactor - JIT] Gather Scatter new implementations #1356

Merged
merged 9 commits into from
Feb 26, 2024
Merged

[Refactor - JIT] Gather Scatter new implementations #1356

Changes from 1 commit
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
a691c05
Add Gather
nathanielsimard Feb 22, 2024
be563ec
WIP;2u
nathanielsimard Feb 22, 2024
858208a
WIP
nathanielsimard Feb 23, 2024
e82f352
WIP
nathanielsimard Feb 23, 2024
f845e2c
FIx bugs
nathanielsimard Feb 23, 2024
e807d99
Cleanup
nathanielsimard Feb 23, 2024
601c429
Fix bugs
nathanielsimard Feb 23, 2024
8104a11
Cleanup
nathanielsimard Feb 23, 2024
3059589
Cleanup
nathanielsimard Feb 26, 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
Prev Previous commit
Cleanup
  • Loading branch information
@nathanielsimard
nathanielsimard committed Feb 26, 2024
commit 30595892a23e9462063076ae144508dd54dc1522
5 changes: 2 additions & 3 deletions crates/burn-wgpu/src/kernel/index/scatter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,11 +210,10 @@ pub(crate) fn scatter<R: Runtime, E: JitElement, I: JitElement, const D: usize>(
current *= val;
num_elems_per_workgroup *= tensor.shape.dims[index];
});

let workgroup = elemwise_workgroup(num_elems_per_workgroup, WORKGROUP_DEFAULT);
// Fake strides.
// Fake strides of the virtual output where the strides of dim is hardcoded to one.
indices.strides = strides;

let workgroup = elemwise_workgroup(num_elems_per_workgroup, WORKGROUP_DEFAULT);
execute_dynamic::<R, ScatterEagerKernel<R, E>, E>(
&[
EagerHandle::new(&tensor.handle, &tensor.strides, &tensor.shape.dims),
Expand Down
Loading