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

Implement Const Folding for Mul, Add and Sub #42

Merged
merged 25 commits into from
Mar 5, 2024
Merged

Conversation

BradenEverson
Copy link
Member

Nodes:

  • Implement is_zero() method, I have it working for all of the normal primative types, but I'll have to check with y'all how bf16 and the complex numbers should be tested.
  • Implement is_const() method that returns a Literal, used in const folding to iterate through the literal and check that each element is 1 for folding multiplication.

Const Folding:

  • Create replace_index() that takes in two NodeIdentifiers and replaces all references to the first one with references to the second. This effectively writes the first node out of the compute graph and sets the correct dependents all at the same time, so it's very nice. If we ever add more operations though we'll have to add it to the huge match tree, there's probably a better way with macros or something to do this but it works for now.
  • fold_consts(): Builds a queue of Nodes connected to the node we start at and traverses through them while making sure we don't pass the mod limit. For every node traveled, check the operation. If the operation is mul, sub or add, we can do const folding. For sub and add, if one node is a constant zero we use replace_index() and remove all references to the current node and replace it with the non-zero node. For multiplication, if one node is zeroes, we replace the current node with the zeroed node. If one node is ones, replace the current node with the non-one node. This will work as long as any dimension altering multiplication is handled by a different operation, mat_mul will be its own operation, correct?

Implemented a few unit tests to check that fold_consts() returns true when it should and doesn't when it shouldn't. Also compiled several compute graphs in tests for all of the use cases I've talked about. Feel free to let me know what you all think!

@BradenEverson BradenEverson linked an issue Mar 4, 2024 that may be closed by this pull request
@BradenEverson BradenEverson merged commit 23a9a76 into master Mar 5, 2024
1 of 2 checks passed
@BradenEverson BradenEverson deleted the const_fold branch March 5, 2024 15:05
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.

Constant folding
1 participant