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

Add expand-copy pass that expands memref.copy as a nested structure o… #3452

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

gabrielrodcanal
Copy link
Contributor

@gabrielrodcanal gabrielrodcanal commented Nov 14, 2024

Pass that expands memref.copy as a nested structure of affine for loops with load/store. This is a functional port of the expand-copy pass in Google's HEIR: https://heir.dev/docs/passes/#-expand-copy.

Not sure if this is interesting for the community, but is something I needed for my own work. If it is, I will provide tests soon.

@gabrielrodcanal gabrielrodcanal added the transformations Changes or adds a transformatio label Nov 14, 2024
@gabrielrodcanal gabrielrodcanal self-assigned this Nov 14, 2024

def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
expand_copy_pass = PatternRewriteWalker(
GreedyRewritePatternApplier([ExpandCopy()]),
Copy link
Member

Choose a reason for hiding this comment

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

I don't think there's a need for the greedy applier if it's just one rewrite

)
for_indices.append(args[0])

symbols = list(map(lambda dim: AffineExpr.dimension(dim), range(ndims)))
Copy link
Member

Choose a reason for hiding this comment

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

I'd recommend list/tuple comprehensions over map these days, they're noticeably faster

Copy link
Member

@superlopuh superlopuh left a comment

Choose a reason for hiding this comment

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

Sure, why not :) But would definitely be great for it to be well-tested. Out of curiosity, what will you be using it for?

@gabrielrodcanal
Copy link
Contributor Author

Sure, why not :) But would definitely be great for it to be well-tested. Out of curiosity, what will you be using it for?

Just latency estimation in terms of clock cycles. Strictly speaking I wouldn't need this pass to do that, but the tool I'm using works on the affine dialect, so it was easier writing this pass than changing the tool.

@gabrielrodcanal
Copy link
Contributor Author

gabrielrodcanal commented Nov 14, 2024

Then I will apply your comments and include tests in the following days :) @superlopuh .

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
transformations Changes or adds a transformatio
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants