Skip to content

[TVMScript] Improve printer for TIR syntax sugar#9680

Merged
vinx13 merged 8 commits into
apache:mainfrom
shingjan:tir-syntax-sugar-printer
Dec 10, 2021
Merged

[TVMScript] Improve printer for TIR syntax sugar#9680
vinx13 merged 8 commits into
apache:mainfrom
shingjan:tir-syntax-sugar-printer

Conversation

@shingjan

@shingjan shingjan commented Dec 8, 2021

Copy link
Copy Markdown

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
Comment thread src/printer/tvmscript_printer.cc Outdated
@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

@vinx13 vinx13 left a comment

Copy link
Copy Markdown
Member

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

Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
shingjan added 2 commits December 9, 2021 00:00

@vinx13 vinx13 left a comment

Copy link
Copy Markdown
Member

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

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

@vinx13 vinx13 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

a few minor issues

Comment thread src/printer/tvmscript_printer.cc
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
Comment thread src/printer/tvmscript_printer.cc Outdated
@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