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

[TIR][Compute-at] Enable complex floordiv/floormod expressions in compute_at #14854

Merged
merged 6 commits into from
Jun 14, 2023

Conversation

abhikran-quic
Copy link
Contributor

@abhikran-quic abhikran-quic commented May 15, 2023

Right now, compute_at is not able to handle movement of blocks if the indices of tensors contain complex floordiv/floormod operations. For ex. if a layout of tensor contains (width % 2) // 2, compute_at fails with an error.

Check failed: (dom_var.defined()) is false: ValueError: BufferRegion pattern match failed: v2 % T.int64(2) // T.int64(2)

This change is to handle complex expressions in compute_at.

@tvm-bot
Copy link
Collaborator

tvm-bot commented May 15, 2023

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

Generated by tvm-bot

…pute_at

Right now, compute_at is not able to handle movement of blocks if the indices
of tensors contain complex floordiv/floormod operations.
For ex. if a layout of tensor contains (width % 2) // 2,
compute_at fails with an error.
This change is to handle complex expressions.
@abhikran-quic
Copy link
Contributor Author

abhikran-quic commented May 15, 2023

cc @Hzfengsy , @junrushao , @quic-sanirudh , @shingjan

Copy link
Contributor

@quic-sanirudh quic-sanirudh left a comment

Choose a reason for hiding this comment

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

LGTM, but I'll prefer if someone else also reviews this as I'm not too familiar with arith pattern matching.

@Hzfengsy
Copy link
Member

cc @wrongtest-intellif

tests/python/unittest/test_tir_schedule_compute_at.py Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
@abhikran-quic
Copy link
Contributor Author

Thank you @wrongtest-intellif for the review. Currently, I'm working on fixing your comments.

Copy link
Contributor Author

@abhikran-quic abhikran-quic left a comment

Choose a reason for hiding this comment

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

Thank you @wrongtest-intellif for the review comments. Apologies for replying late. I have fixed the review comments.

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
Copy link
Contributor Author

@abhikran-quic abhikran-quic left a comment

Choose a reason for hiding this comment

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

Thank you @wrongtest-intellif for sharing more details. Sorry for coming back late but I was trying to fix your review comments and finding an appropriate test case to verify the LOC changed. I have updated the PR with new changes. Request you to please review the PR for any additional feedback.

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
tests/python/unittest/test_tir_schedule_compute_at.py Outdated Show resolved Hide resolved
@abhikran-quic abhikran-quic force-pushed the compute_at_floordivmod branch 3 times, most recently from 17090e9 to c2e7c5e Compare June 7, 2023 08:11
@abhikran-quic
Copy link
Contributor Author

Hi @wrongtest-intellif , Could you please review this PR ?

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
Copy link
Contributor Author

@abhikran-quic abhikran-quic left a comment

Choose a reason for hiding this comment

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

Thank you @wrongtest-intellif for the comments. I have fixed them. Please review the PR for any more feedback.

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved
@abhikran-quic abhikran-quic force-pushed the compute_at_floordivmod branch 2 times, most recently from b2db447 to 5b4f145 Compare June 12, 2023 14:31
- Merge conditions for floordiv and floormods respectively to handle recursion
- Update test case to handle non-bijective index map
- Update new_required to be in between [required_min, +inf]
Copy link
Contributor

@wrongtest-intellif wrongtest-intellif left a comment

Choose a reason for hiding this comment

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

Thanks, now it looks good to me.

@abhikran-quic
Copy link
Contributor Author

@tvm-bot rerun

@quic-sanirudh quic-sanirudh merged commit bd24133 into apache:main Jun 14, 2023
18 checks passed
@abhikran-quic abhikran-quic deleted the compute_at_floordivmod branch June 14, 2023 06:45
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.

None yet

5 participants