Skip to content

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

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

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

@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