-
Notifications
You must be signed in to change notification settings - Fork 29
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
Unify kernels for basic and nested affine for loop #204
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
As now the dialect and the frontend have been merged, can you also add the Python tests under the tests
folder?
@EthanMeng324 Can you follow this guide to format your code and add necessary headers? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just some minor issues. I'd suggest changing all the unify_kernels
to unify
in order to provide a simple API for users. Only change the Python-side name, but no need to change the C++ name. Otherwise LGTM
allo/ir/transform.py
Outdated
@@ -3,6 +3,7 @@ | |||
# pylint: disable=no-name-in-module, unexpected-keyword-arg, no-value-for-parameter | |||
|
|||
import numpy as np | |||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove this line
allo/ir/unify_kernels.py
Outdated
from ..customize import customize | ||
|
||
|
||
def unify_kernels(func1: Callable, func2: Callable, loop_num: int): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Add docstring for this function to explain its functionality. It is important to explain what loop_num
is.
allo/ir/unify_kernels.py
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Move this file to allo/primitive/unify.py
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
(Move and rename)
tests/test_unify_kernels.py
Outdated
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Change name to test_unify
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great! Thanks for updating the code!
Create a MLIR level unify functionality for two kernels. Merge the same affine for loop and use conditional branch to support different logic.
Migrate from an unmerged PR in hcl-dialect: cornell-zhang/hcl-dialect#217