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] Add preserve-unit-iters #11585

Merged

Conversation

junrushao
Copy link
Member

@junrushao junrushao commented Jun 5, 2022

Follow-up of #11578, which enforces structural stability of TIR by
avoiding over-simplification in affine analysis. On the other hand, it
is possible that over-simplification could be desirable behavior.
Therefore, following the precedent of preserve-unit-loops in
Compute-At, this PR introduces preserve-unit-iters in block binding
for cases where users don't need structural stability (which is
admittedly rare).

This PR does not affect any existing functionalities.

Example:

# The original TIR
for i in T.serial(2):
    with T.block("C"):
        k = T.axis.reduce(2, i)

# Our operation is:
Split(i, [1, 2], preserve-unit-iters=True/False)

# preserve-unit-iters = True
for i_0, i_1 in T.grid(1, 2):
    with T.block("C"):
        k = T.axis.reduce(2, i_0 * 2 + i_1)

# preserve-unit-iters = False
for i_0, i_1 in T.grid(1, 2):
    with T.block("C"):
        k = T.axis.reduce(2, i_1)

@junrushao
Copy link
Member Author

Per discussion with @spectrometerHBH. CC: @wrongtest @Hzfengsy @jinhongyii would you guys like to review this PR? Thanks a lot!

@Hzfengsy
Copy link
Member

Hzfengsy commented Jun 6, 2022

Can you add unittests for it? Then we can merge it quickly

@junrushao
Copy link
Member Author

Per discussion with @Hzfengsy: we are going to use True as default given it's not gonna to affect existing functionalities. I'm going to update the PR later this week :-)

@junrushao junrushao force-pushed the feature/2022-06-05/preserve-unit-iter branch from 1ebeab2 to 0529a35 Compare June 16, 2022 03:53
Follow-up of apache#11578, which enforces structural stability of TIR by
avoiding over-simplification in affine analysis. On the other hand, it
is possible that over-simplification could be desirable behavior.
Therefore, following the precedent of `preserve-unit-loops` in
`Compute-At`, this PR introduces `preserve-unit-iters` in block binding
for cases where users don't need structural stability (which is
admittedly rare).

This PR does not affect any existing functionalities.

Example:

```python
for i in T.serial(2):
    with T.block("C"):
        k = T.axis.reduce(2, i)

Split(i, [1, 2], preserve-unit-iters=True/False)

for i_0, i_1 in T.grid(1, 2):
    with T.block("C"):
        k = T.axis.reduce(2, i_0 * 2 + i_1)

for i_0, i_1 in T.grid(1, 2):
    with T.block("C"):
        k = T.axis.reduce(2, i_1)
```
@junrushao junrushao force-pushed the feature/2022-06-05/preserve-unit-iter branch from 0529a35 to 9b85836 Compare June 16, 2022 03:55
@Hzfengsy Hzfengsy merged commit 89e1a6c into apache:main Jun 16, 2022
@junrushao
Copy link
Member Author

Script to upgrade tuning record database:

import argparse
import json


def _parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--json-tuning-record",
        type=str,
        help="path to the json tuning record file",
        required=True,
    )
    parser.add_argument(
        "--updated-json-tuning-record",
        type=str,
        help="path to the new json tuning record file after upgrade",
        required=True,
    )
    return parser.parse_args()


ARGS = _parse_args()


def main():
    with open(ARGS.json_tuning_record, "r", encoding="utf-8") as i_f:
        lines = [json.loads(line) for line in i_f]
    for line in lines:
        for inst in line[1][0][0]:
            if inst[0] in ["Split", "Fuse"]:
                inst[2] = [1]
    with open(ARGS.updated_json_tuning_record, "w", encoding="utf-8") as o_f:
        for line in lines:
            o_f.write(json.dumps(line) + "\n")


if __name__ == "__main__":
    main()

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.

3 participants