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.

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

Migrate/jit/matmul tiling 2d #1472

Merged
merged 28 commits into from
Mar 22, 2024
Merged

Migrate/jit/matmul tiling 2d #1472

merged 28 commits into from
Mar 22, 2024

Conversation

louisfd
Copy link
Member

@louisfd louisfd commented Mar 14, 2024

Matmul tiling 2d kernels for #1422

Bonus stuff:

  • macro for configurable unrolling of loops
  • tiling 2d (unpadded) will automatically use the kernel with no bound checks when shapes are divisible by block sizes

Copy link

codecov bot commented Mar 14, 2024

Codecov Report

Attention: Patch coverage is 93.61702% with 63 lines in your changes are missing coverage. Please review.

Project coverage is 85.87%. Comparing base (0a8a3cc) to head (94a81ee).

Files Patch % Lines
crates/burn-jit/src/kernel/matmul/base.rs 62.19% 31 Missing ⚠️
.../burn-jit/src/codegen/dialect/gpu/vectorization.rs 0.00% 7 Missing ⚠️
crates/burn-jit/src/fusion/tracing/builder.rs 0.00% 5 Missing ⚠️
crates/burn-jit/src/kernel/matmul/tiling2d.rs 97.16% 4 Missing ⚠️
.../src/kernel/matmul/tiling2d_shader/write_output.rs 95.95% 4 Missing ⚠️
...t/src/kernel/matmul/tiling2d_shader/computation.rs 95.52% 3 Missing ⚠️
crates/burn-wgpu/src/compiler/wgsl/compiler.rs 81.25% 3 Missing ⚠️
...ernel/matmul/tiling2d_shader/load_shared_memory.rs 99.20% 2 Missing ⚠️
crates/burn-jit/src/kernel/matmul/tune/base.rs 87.50% 2 Missing ⚠️
...rates/burn-jit/src/codegen/dialect/gpu/variable.rs 50.00% 1 Missing ⚠️
... and 1 more
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1472      +/-   ##
==========================================
+ Coverage   85.79%   85.87%   +0.07%     
==========================================
  Files         657      660       +3     
  Lines       74022    74807     +785     
==========================================
+ Hits        63505    64238     +733     
- Misses      10517    10569      +52     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Copy link
Member

@nathanielsimard nathanielsimard left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only one comment, otherwise looks good

Comment on lines 117 to 127
let tmp = scope.create_local(Elem::UInt);
let position_0 = scope.create_local(Elem::UInt);
let position_1 = scope.create_local(Elem::UInt);
let position_2 = scope.create_local(Elem::UInt);
let position_3 = scope.create_local(Elem::UInt);
let remain_n = scope.create_local(Elem::Bool);

let val_0 = scope.zero(elem);
let val_1 = scope.zero(elem);
let val_2 = scope.zero(elem);
let val_3 = scope.zero(elem);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we extract those else where, the declaration is probably done in a for loop. Also, not sure we need to initialize val_0, val_1, val_2 and val_3 to zero.

@louisfd louisfd merged commit dd699a9 into main Mar 22, 2024
15 checks passed
@louisfd louisfd deleted the migrate/jit/matmul_tiling_2d branch March 22, 2024 12:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants