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

IR Refactor #199

Merged
merged 13 commits into from
Oct 29, 2024
Merged

IR Refactor #199

merged 13 commits into from
Oct 29, 2024

Conversation

wingertge
Copy link
Contributor

Refactors IR to make it more ergonomic by separating the output of an operation from the operation itself, and separating the type of any given variable from the kind/identifier.
Moves builtin variables into a separate enum so they no longer fill the VariableKind enum with a bunch of semantically identical duplicates. This is more in line with other variable kinds that use IDs to indentify the specific variable instance.
Moves atomic operations to their own enum.
Moves Select from Branch to Operation.
Adds explicit Cast operator to avoid confusing implicit casting on Assign.
Makes Assign a first class Operation since now that Cast is explicit it's essentially a register copy and thus has different semantics compared to other operations (for example, it could be safely inlined and removed without affecting execution).

There's a lot of churn but it's overall cleaner, and actually lowers LOC by almost 400. It also simplifies any future optimization passes that may be added to cubecl-opt. Things that only care about the output can now access it directly, and vice versa for the operation and inputs.

Testing

All tests pass in both burn and cubecl. I've already prepared a burn version where all legacy kernels are migrated to the refactored variables.

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[allow(dead_code, missing_docs, clippy::large_enum_variant)] // Some variants might not be used with different flags
pub enum Operation {
Assign(Variable),
Copy link
Member

Choose a reason for hiding this comment

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

Maybe Identity is better here? Assign without an output is a bit strange.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Normally it would be called copy but that was already used. Maybe we should just rename Copy to CopyMemory (which is in line with the SPIR-V name and memcpy), and call this one Copy. Because that's also how it's still used: to make a copy of a mutable variable.

@nathanielsimard nathanielsimard merged commit a5c9dee into tracel-ai:main Oct 29, 2024
5 checks passed
@wingertge wingertge deleted the feat/ir-refactor branch October 30, 2024 12:11
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