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

Unify kernels for basic and nested affine for loop #204

Merged
merged 17 commits into from
Oct 28, 2024

Conversation

EthanMeng324
Copy link
Contributor

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

Copy link
Member

@chhzh123 chhzh123 left a 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?

mlir/lib/Transforms/UnifyKernels.cpp Outdated Show resolved Hide resolved
@chhzh123
Copy link
Member

@EthanMeng324 Can you follow this guide to format your code and add necessary headers?

Copy link
Member

@chhzh123 chhzh123 left a 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

@@ -3,6 +3,7 @@
# pylint: disable=no-name-in-module, unexpected-keyword-arg, no-value-for-parameter

import numpy as np

Copy link
Member

Choose a reason for hiding this comment

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

Remove this line

from ..customize import customize


def unify_kernels(func1: Callable, func2: Callable, loop_num: int):
Copy link
Member

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.

Copy link
Member

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

Copy link
Member

Choose a reason for hiding this comment

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

(Move and rename)

Copy link
Member

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

Copy link
Member

@chhzh123 chhzh123 left a 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!

@chhzh123 chhzh123 merged commit 39c69a2 into cornell-zhang:main Oct 28, 2024
1 check passed
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