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

Migrate/jit/matmul tiling 2d #1472

Merged
merged 28 commits into from
Mar 22, 2024
Merged

Migrate/jit/matmul tiling 2d #1472

Changes from 1 commit
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
83fd099
refactor matmul files
louisfd Mar 11, 2024
d0b2ce4
wip refactor matmul
louisfd Mar 11, 2024
603df3f
everything is memco
louisfd Mar 12, 2024
423d23d
support local arrays
louisfd Mar 12, 2024
fa222a0
advancing tiling2d
louisfd Mar 12, 2024
4419f45
advancing tiling2d
louisfd Mar 13, 2024
9b6f2d1
advancing tiling2d
louisfd Mar 13, 2024
feddabc
tiling2d finished but buggy
louisfd Mar 13, 2024
7e07825
configurable unrolling
louisfd Mar 13, 2024
855a77c
not bugged
louisfd Mar 13, 2024
11f55fd
fails on unroll
louisfd Mar 13, 2024
d045fe6
stupid break
louisfd Mar 13, 2024
a0bb940
tiling2d no assumption works
louisfd Mar 14, 2024
5596e3b
Merge branch 'main' into migrate/jit/matmul_tiling_2d
louisfd Mar 14, 2024
4123842
clippy
louisfd Mar 14, 2024
ed18cfa
bounds check as bool
louisfd Mar 19, 2024
651010a
lhs rhs as enum
louisfd Mar 19, 2024
5e8786b
tiling 2d major refactor
louisfd Mar 19, 2024
ea82414
remove assign vec4
louisfd Mar 19, 2024
8b504f7
variable declarations above loops
louisfd Mar 19, 2024
947c473
fmt
louisfd Mar 19, 2024
4c97c75
clippy
louisfd Mar 19, 2024
66076db
Fix autotune + unroll
nathanielsimard Mar 20, 2024
7a6299a
move val
louisfd Mar 21, 2024
1e9c815
Merge branch 'migrate/jit/matmul_tiling_2d' of github.com:tracel-ai/b…
louisfd Mar 21, 2024
b3af965
clippy
louisfd Mar 21, 2024
6237478
fmt
louisfd Mar 21, 2024
94a81ee
Merge branch 'main' into migrate/jit/matmul_tiling_2d
louisfd Mar 21, 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
Next Next commit
lhs rhs as enum
  • Loading branch information
@louisfd
louisfd committed Mar 19, 2024
commit 651010ac297ac0b8fd5c47e7ada022f3ffadefcf
59 changes: 33 additions & 26 deletions crates/burn-jit/src/kernel/matmul/tiling2d_shader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,11 @@ pub(crate) struct MatmulTiling2dShader {
pub unroll: bool,
}

enum InputIdentifier {
Lhs,
Rhs,
}

impl MatmulTiling2dShader {
pub(crate) fn expand(self, scope: &mut Scope) {
// Inputs
Expand Down Expand Up @@ -161,8 +166,8 @@ impl MatmulTiling2dShader {
gpu!(scope, k = i * block_size_k);

if self.bounds_check_required {
// LHS
self.load_shared_memory_with_bound_check(
InputIdentifier::Lhs,
scope,
k,
dim_k,
Expand All @@ -175,11 +180,10 @@ impl MatmulTiling2dShader {
lhs,
offset_lhs,
shared_lhs,
true,
);

// RHS
self.load_shared_memory_with_bound_check(
InputIdentifier::Rhs,
scope,
k,
dim_k,
Expand All @@ -192,11 +196,10 @@ impl MatmulTiling2dShader {
rhs,
offset_rhs,
shared_rhs,
false,
);
} else {
// LHS
self.load_shared_memory(
InputIdentifier::Lhs,
scope,
k,
thread_col,
Expand All @@ -206,11 +209,10 @@ impl MatmulTiling2dShader {
lhs,
offset_lhs,
shared_lhs,
true,
);

// RHS
self.load_shared_memory(
InputIdentifier::Rhs,
scope,
k,
thread_row,
Expand All @@ -220,7 +222,6 @@ impl MatmulTiling2dShader {
rhs,
offset_rhs,
shared_rhs,
false,
);
}

Expand Down Expand Up @@ -253,6 +254,7 @@ impl MatmulTiling2dShader {
#[allow(clippy::too_many_arguments)]
fn load_shared_memory_with_bound_check(
&self,
input_identifier: InputIdentifier,
scope: &mut Scope,
k: Variable,
dim_k: Variable,
Expand All @@ -265,7 +267,6 @@ impl MatmulTiling2dShader {
input: Variable,
input_offset: Variable,
shared_memory: Variable,
is_lhs: bool,
) {
// How close is the thread to the end of the matrix.
// If < 4 then it is an edge case
Expand All @@ -291,14 +292,17 @@ impl MatmulTiling2dShader {

// Position in shared memory
let sm_position = scope.create_local(Elem::UInt);
if is_lhs {
gpu!(scope, sm_position = thread_idx_2 / 4u32);
gpu!(scope, sm_position *= block_size_k);
gpu!(scope, sm_position += current);
} else {
gpu!(scope, sm_position = current * block_size_n);
gpu!(scope, sm_position += thread_idx_2);
gpu!(scope, sm_position = sm_position / 4u32);
match input_identifier {
InputIdentifier::Lhs => {
gpu!(scope, sm_position = thread_idx_2 / 4u32);
gpu!(scope, sm_position *= block_size_k);
gpu!(scope, sm_position += current);
},
InputIdentifier::Rhs => {
gpu!(scope, sm_position = current * block_size_n);
gpu!(scope, sm_position += thread_idx_2);
gpu!(scope, sm_position = sm_position / 4u32);
}
}

// To pad with zeros if outside lhs
Expand Down Expand Up @@ -375,6 +379,7 @@ impl MatmulTiling2dShader {
#[allow(clippy::too_many_arguments)]
fn load_shared_memory(
&self,
input_identifier: InputIdentifier,
scope: &mut Scope,
k: Variable,
thread_idx_1: Variable,
Expand All @@ -384,7 +389,6 @@ impl MatmulTiling2dShader {
input: Variable,
input_offset: Variable,
shared_memory: Variable,
is_lhs: bool,
) {
let block_size_k: Variable = self.config.block_size_k.into();
let block_size_n: Variable = self.config.block_size_n.into();
Expand All @@ -403,14 +407,17 @@ impl MatmulTiling2dShader {
gpu!(scope, if(aligned_with_shared_memory).then(|scope|{

let sm_position = scope.create_local(Elem::UInt);
if is_lhs {
gpu!(scope, sm_position = thread_idx_2 / 4u32);
gpu!(scope, sm_position *= block_size_k);
gpu!(scope, sm_position += current);
} else {
gpu!(scope, sm_position = current * block_size_n);
gpu!(scope, sm_position += thread_idx_2);
gpu!(scope, sm_position = sm_position / 4u32);
match input_identifier {
InputIdentifier::Lhs => {
gpu!(scope, sm_position = thread_idx_2 / 4u32);
gpu!(scope, sm_position *= block_size_k);
gpu!(scope, sm_position += current);
},
InputIdentifier::Rhs => {
gpu!(scope, sm_position = current * block_size_n);
gpu!(scope, sm_position += thread_idx_2);
gpu!(scope, sm_position = sm_position / 4u32);
}
}

let position_0 = scope.create_local(Elem::UInt);
Expand Down