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

[TensorIR][M1c] LowerInitBlock #7806

Merged
merged 1 commit into from
Apr 8, 2021
Merged

Conversation

Hzfengsy
Copy link
Member

@Hzfengsy Hzfengsy commented Apr 7, 2021

This PR is the first part of the stage M1c, TensorIR upstreaming plan (#7527) on lowering init block.

We just use if branch to replace the init block. e.g.

Before

with tir.block([16, 16, tir.reduce_axis(16)]) as [i, j ,k]:
    with tir.init():
        C[i, j] = tir.float32(0)
    C[i, j] += A[i, k] * B[j, k]

After

with tir.block([16, 16, tir.reduce_axis(16)]) as [i, j ,k]:
    if k == 0:
        C[i, j] = tir.float32(0)
    C[i, j] += A[i, k] * B[j, k]

Please note that we only provide the default behavior for lowering the init part, which is not performent. We will provide primitives for users to manually "decompose" init part.

cc @tqchen @junrushao1994 @jroesch @comaniac @jcf94

@tqchen tqchen changed the title [TensorIR] introduce pass lower_init_block [TensorIR][PASS] LowerInitBlock Apr 7, 2021
@tqchen tqchen self-assigned this Apr 7, 2021
Copy link
Member

@tqchen tqchen left a comment

Choose a reason for hiding this comment

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

minor nits, overall LGTM

src/tir/transforms/lower_init_block.cc Outdated Show resolved Hide resolved
@Hzfengsy Hzfengsy force-pushed the LowerInitBlock branch 2 times, most recently from 3ff3602 to 87d5a58 Compare April 7, 2021 15:04
src/tir/transforms/lower_init_block.cc Outdated Show resolved Hide resolved
Copy link
Contributor

@comaniac comaniac left a comment

Choose a reason for hiding this comment

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

Overall LGTM.
One suggestion: Since we have defined a clear stage (i.e., M1c) for this series of PRs, it would be better for all PR submitters to use the unified PR titles (for exampe: [TensorIR][M1c] ...) so that they can be easily queried.

return Block(n);
}

static Stmt RealizeInitBlock(const Stmt& init, const Array<IterVar>& iter_vars) {
Copy link
Contributor

Choose a reason for hiding this comment

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

Could you clarify the meaning of the word "realize" here?

Copy link
Member

Choose a reason for hiding this comment

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

What about "concretize"? My understanding is that the init-block is something "conceptual" for initializing "zero point" of a reduction, and this pass is used to make it concrete, lowering to "if-then-else" statement

Copy link
Contributor

Choose a reason for hiding this comment

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

I got your point, and I feel "LowerInitBlock" itself is the best name to describe what it does, so maybe the simplest way is manually inlining this function so that we won't be bothered by its name, or just give it a general helper name to let people know this is just the implementation of lowering init block.

Copy link
Member

Choose a reason for hiding this comment

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

This helper is something a bit standalone (making the init block to If-Then-Else), and inlining this function might make the logic in the function above a bit vague, so what about we just call it "DoLowering"?

Copy link
Contributor

Choose a reason for hiding this comment

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

DoLowering sounds good to me.

@tqchen tqchen changed the title [TensorIR][PASS] LowerInitBlock [TensorIR][M1c] LowerInitBlock Apr 7, 2021


def test_lower_reduction():
origin_mod = WithInit()
Copy link
Member

Choose a reason for hiding this comment

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

Just a note here: per discussion with @Hzfengsy, we need to remove "()" here in future PRs, so that the script parser API is consistent across PrimFunc and IRModule

Copy link
Member

@tqchen tqchen left a comment

Choose a reason for hiding this comment

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

@Hzfengsy please address last comment by @comaniac (DoLowering) and we can merge this in

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
@tqchen tqchen merged commit 24e62ca into apache:main Apr 8, 2021
@tqchen
Copy link
Member

tqchen commented Apr 8, 2021

Thanks @comaniac @junrushao1994 for reviewing. Thanks @Hzfengsy

@Hzfengsy Hzfengsy deleted the LowerInitBlock branch April 9, 2021 14:33
tmoreau89 pushed a commit to tmoreau89/tvm that referenced this pull request Apr 11, 2021
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
@junrushao junrushao linked an issue Apr 12, 2021 that may be closed by this pull request
29 tasks
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request May 6, 2021
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request May 6, 2021
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request May 6, 2021
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
trevor-m pushed a commit to neo-ai/tvm that referenced this pull request May 11, 2021
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
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.

[RFC][Tracking Issue] TensorIR Scheduling
4 participants