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

[TVMScript] Improve printer for TIR syntax sugar #9680

Merged
merged 8 commits into from
Dec 10, 2021

Conversation

shingjan
Copy link

@shingjan shingjan commented Dec 8, 2021

This PR intends to improve the printing of three kinds of syntax sugar introduced in TVMScript from #9492, #9620 and #9634. Examples below are supposedly the output of print(some_tir_prim_func.script()).

  1. For reads & writes:

Before this PR

@T.prim_func
def func(a: T.handle, b: T.handle, c: T.handle) -> None:
            ...
            T.reads([C[vi, vj], A[vi, vk], B[vj, vk]])
            T.writes([C[vi, vj], A[vi, vk]])
            ...

After this PR, note that brackets are removed

@T.prim_func
def func(a: T.handle, b: T.handle, c: T.handle) -> None:
            ...
            T.reads(C[vi, vj], A[vi, vk], B[vj, vk])
            T.writes(C[vi, vj], A[vi, vk])
            ...
  1. For loops:

Before this PR

@T.prim_func
def func(a: T.handle) -> None:
    ...
    for i in T.serial(0, 128):
        for j in T.parallel(0, 128):
            for k in T.vectorized(0, 128):
                for x in T.unroll(0, 128):
                    for y in T.thread_binding(0, 128, thread="threadIdx.x"):
                        ...

After this PR, note that all loops with 0 as min will be printed out as syntax sugar

    ...
    for i in T.serial(128):
        for j in T.parallel(128):
            for k in T.vectorized(128):
                for x in T.unroll(128):
                    for y in T.thread_binding(128, thread="threadIdx.x"):
                        ...
  1. For T.match_buffer():

Before this PR

@T.prim_func
def elementwise(
    a: T.handle,
    b: T.handle,
) -> None:
    A = T.match_buffer(a, (128, 128, 128, 128))
    B = T.match_buffer(b, (128, 128, 128, 128))
    # some computation

After this PR, note that we assume T.match_buffer with two or less arguments specified are syntax-sugarred.

@T.prim_func
def elementwise(
    a: T.Buffer[[128, 128, 128, 128], "float32"],
    b:T.Buffer[[128, 128, 128, 128], "float32"],
) -> None:
    # some computation

cc: @vinx13 @junrushao1994 @Hzfengsy

@shingjan shingjan force-pushed the tir-syntax-sugar-printer branch from e47a560 to 0e0202e Compare December 8, 2021 20:42
@shingjan shingjan marked this pull request as ready for review December 9, 2021 01:18
@shingjan shingjan changed the title [WIP][TVMScript] Improve printer for TIR syntax sugar [TVMScript] Improve printer for TIR syntax sugar Dec 9, 2021
@shingjan shingjan requested a review from vinx13 December 9, 2021 01:25
Copy link
Member

@vinx13 vinx13 left a comment

Choose a reason for hiding this comment

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

Please fix CI error and double check if works locally

src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
shingjan added 2 commits December 9, 2021 00:00
Copy link
Member

@vinx13 vinx13 left a comment

Choose a reason for hiding this comment

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

from the CI error looks like there are some issues with reads/writes in previous PR

src/printer/tvmscript_printer.cc Show resolved Hide resolved
tests/python/unittest/test_tvmscript_roundtrip.py Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
shingjan added 2 commits December 9, 2021 17:37
@shingjan shingjan requested review from Hzfengsy and vinx13 December 10, 2021 02:25
Copy link
Member

@vinx13 vinx13 left a comment

Choose a reason for hiding this comment

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

a few minor issues

src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
src/printer/tvmscript_printer.cc Outdated Show resolved Hide resolved
@shingjan shingjan requested a review from vinx13 December 10, 2021 03:37
@vinx13 vinx13 merged commit 40ecfec into apache:main Dec 10, 2021
@shingjan shingjan deleted the tir-syntax-sugar-printer branch December 10, 2021 18:05
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 7, 2022
yangulei pushed a commit to yangulei/tvm that referenced this pull request Jan 11, 2022
yangulei pushed a commit to yangulei/tvm that referenced this pull request Jan 12, 2022
ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
qsqqsqqsq-intellif pushed a commit to qsqqsqqsq-intellif/tvm that referenced this pull request Apr 29, 2022
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