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

Reimplement macro to make writing kernels more ergonomic #80

Merged
merged 67 commits into from
Sep 10, 2024

Conversation

wingertge
Copy link
Contributor

@wingertge wingertge commented Aug 22, 2024

This is a rewrite of the frontend macro to make writing kernels more ergonomic and "rust-like".

It replaces the Comptime struct with a #[comptime] attribute so values can be used like any other Rust constant.

Loops now use normal ranges and allow unrolling via #[unroll] rather than a verborse comptime parameter that is often set to a literal anyways. For conditional unrolling, the attribute supports inserting an expression: #[unroll(should_unroll)].

The custom vectorized types are replaced with Rust primitives since vectorization has always been stored in the IR anyways. The vectorization field on the custom types was almost entirely redundant.

Adds support for struct destructuring in locals.

Folding of comptime values into literals has also been improved.

Testing

All unit and integration tests (including frontend tests) are passing. Several backend codegen tests are also added.
All burn tests are passing.

@wingertge
Copy link
Contributor Author

All tests passing including frontend tests, actually ready for merge now

@wingertge wingertge marked this pull request as ready for review September 8, 2024 21:34
@wingertge wingertge changed the title [Draft] Reimplement macro and IR to make writing kernels more ergonomic Reimplement macro to make writing kernels more ergonomic Sep 8, 2024
@wingertge wingertge marked this pull request as draft September 9, 2024 11:55
@wingertge
Copy link
Contributor Author

Something in burn seems to be miscompiling, going back to draft until I can figure out what's going on

@wingertge wingertge marked this pull request as ready for review September 9, 2024 16:05
@wingertge
Copy link
Contributor Author

All burn tests are passing now.

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.

Great job, definitively a big improvement!

crates/cubecl-core/src/frontend/branch.rs Outdated Show resolved Hide resolved
crates/cubecl-core/src/frontend/branch.rs Outdated Show resolved Hide resolved
crates/cubecl-core/src/frontend/operation/base.rs Outdated Show resolved Hide resolved
crates/cubecl-core/src/runtime_tests/assign.rs Outdated Show resolved Hide resolved
crates/cubecl-linalg/src/matmul/cmma/block_loop.rs Outdated Show resolved Hide resolved
crates/cubecl-linalg/src/matmul/tiling2d/tile/loader.rs Outdated Show resolved Hide resolved
Comment on lines +36 to +39
cubecl-core = { path = "../cubecl-core", version = "0.2", default-features = false }
cubecl-cuda = { path = "../cubecl-cuda", version = "0.2", default-features = false }
cubecl-linalg = { path = "../cubecl-linalg", version = "0.2", default-features = false }
cubecl-wgpu = { path = "../cubecl-wgpu", version = "0.2", default-features = false }
Copy link
Member

Choose a reason for hiding this comment

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

Why are they needed?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

For the codegen comparison tests. Could move them but I put them here because they're testing the macro, not the backends in question. Or maybe just use the testgen macro thing.

@nathanielsimard nathanielsimard merged commit e83e8d9 into tracel-ai:main Sep 10, 2024
2 of 7 checks passed
@wingertge wingertge deleted the new-ir branch September 10, 2024 14:17
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