[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl #14951

Lunderberg · 2023-05-25T15:23:00Z

Prior to this commit, any trivial let binding of var1 = var2 is inlined. However, buffer definitions are not updated, so this can result in dangling tir::Var instances. This commit updates the tir.Simplify pass to keep trivial let bindings if they are used as part of a buffer definition.

Ideally, the trivial LetStmt variable would be inlined into the buffer definition as well as other expressions. However, because a buffer may be implicitly declared, the first usage may be within a constrained context. If that happens, the simplified shape/strides expression cannot be used to update the buffer definition, as that simplification is not valid at all possible usage points of the buffer.

for i in range(n):
    elem_offset = i
    view = T.Buffer(1, data=buf, elem_offset = elem_offset)
    if i == 0:
        # First occurrence in TIR is here, where elem_offset would
        # simplify to zero.
        view[0] = 1
    else:
        # But the same buffer is used here, where elem_offset doesn't
        # simplify to zero.
        view[0] = 2

This will be resolvable after #14778 lands, requiring all buffers to be declared with DeclBuffer prior to usage.

for i in range(n):
    elem_offset = i
    # All variables used by the DeclBuffer are valid across the entire
    # body of the DeclBuffer.
    view = T.decl_buffer(1, data=buf, elem_offset = elem_offset)
    if i == 0:
        view[0] = 1
    else:
        view[0] = 2

tvm-bot · 2023-05-25T15:23:03Z

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

cc @Hzfengsy, @junrushao, @quic-sanirudh, @shingjan _{See #10317 for details}

_{Generated by tvm-bot}

Prior to this commit, `ArgBinder` would always introduce a new variable to represent the input argument, even if the argument already a primitive type. This introduces trivial let bindings that are expected to be simplified out, but which can produce dangling `tir::Var` usage in some cases (see apache#14951). This commit updates `ArgBinder` to prefer using the original `tir::Var` when possible. That is, when a function takes `n: T.int32` as input, the packed function should produce a binding `n: T.int32 = T.tvm_struct_get(...)`, rather than producing a binding `arg_n = T.tvm_struct_get(...)` followed by `n = arg_n`.

Lunderberg · 2023-05-25T23:52:55Z

@tvm-bot rerun

Prior to this commit, `ArgBinder` would always introduce a new variable to represent the input argument, even if the argument already a primitive type. This introduces trivial let bindings that are expected to be simplified out, but which can produce dangling `tir::Var` usage in some cases (see #14951). This commit updates `ArgBinder` to prefer using the original `tir::Var` when possible. That is, when a function takes `n: T.int32` as input, the packed function should produce a binding `n: T.int32 = T.tvm_struct_get(...)`, rather than producing a binding `arg_n = T.tvm_struct_get(...)` followed by `n = arg_n`.

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

Prior to this commit, `ArgBinder` would always introduce a new variable to represent the input argument, even if the argument already a primitive type. This introduces trivial let bindings that are expected to be simplified out, but which can produce dangling `tir::Var` usage in some cases (see apache#14951). This commit updates `ArgBinder` to prefer using the original `tir::Var` when possible. That is, when a function takes `n: T.int32` as input, the packed function should produce a binding `n: T.int32 = T.tvm_struct_get(...)`, rather than producing a binding `arg_n = T.tvm_struct_get(...)` followed by `n = arg_n`.

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions.

masahi · 2023-06-02T21:53:27Z

tests/python/unittest/test_tir_transform_split_host_device.py

@@ -168,5 +168,21 @@ def main_kernel(n: T.int32):
        return mod


+class TestSplitHostDevice(BaseCompare):


Is this test relevant? Looks like a dup from #14988

Hmm, it was a dup from #14982, which I had based this PR on top of to make allow the CI to use that bugfix. I think the "squash and merge" should remove it from this PR as part of merging, but just to be on the safe side, rebasing onto main.

And after rebase, no changes remain in test_tir_transform_split_host_device.py.

Prior to this commit, any trivial let binding of `var1 = var2` is inlined. However, buffer definitions are not updated, so this can result in dangling `tir::Var` instances. This commit updates the `tir.Simplify` pass to keep trivial let bindings if they are used as part of a buffer definition. Ideally, the trivial `LetStmt` variable would be inlined into the buffer definition as well as other expressions. However, because a buffer may be implicitly declared, the first usage may be within a constrained context. If that happens, the simplified shape/strides expression cannot be used to update the buffer definition, as that simplification is not valid at all possible usage points of the buffer. ```python for i in range(n): elem_offset = i view = T.Buffer(1, data=buf, elem_offset = elem_offset) if i == 0: # First occurrence in TIR is here, where elem_offset would # simplify to zero. view[0] = 1 else: # But the same buffer is used here, where elem_offset doesn't # simplify to zero. view[0] = 2 ``` This will be resolvable after apache#14778 lands, requiring all buffers to be declared with `DeclBuffer` prior to usage. ```python for i in range(n): elem_offset = i # All variables used by the DeclBuffer are valid across the entire # body of the DeclBuffer. view = T.decl_buffer(1, data=buf, elem_offset = elem_offset) if i == 0: view[0] = 1 else: view[0] = 2 ```

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions.

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions. While they no longer cause issues with correct-ness, they are unnecessary in this case.

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

Prior to #14951, these can have erroneous simplifications when used in buffer definitions. While they no longer cause issues with correct-ness, they are unnecessary in this case.

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

…pache#14951) Prior to this commit, any trivial let binding of `var1 = var2` is inlined. However, buffer definitions are not updated, so this can result in dangling `tir::Var` instances. This commit updates the `tir.Simplify` pass to keep trivial let bindings if they are used as part of a buffer definition. Ideally, the trivial `LetStmt` variable would be inlined into the buffer definition as well as other expressions. However, because a buffer may be implicitly declared, the first usage may be within a constrained context. If that happens, the simplified shape/strides expression cannot be used to update the buffer definition, as that simplification is not valid at all possible usage points of the buffer. ```python for i in range(n): elem_offset = i view = T.Buffer(1, data=buf, elem_offset = elem_offset) if i == 0: # First occurrence in TIR is here, where elem_offset would # simplify to zero. view[0] = 1 else: # But the same buffer is used here, where elem_offset doesn't # simplify to zero. view[0] = 2 ``` This will be resolvable after apache#14778 lands, requiring all buffers to be declared with `DeclBuffer` prior to usage. ```python for i in range(n): elem_offset = i # All variables used by the DeclBuffer are valid across the entire # body of the DeclBuffer. view = T.decl_buffer(1, data=buf, elem_offset = elem_offset) if i == 0: view[0] = 1 else: view[0] = 2 ```

The functionality tested in this commit was added across several recent PRs, each of which tested their features in isolation. This PR adds unit tests to validate the end-to-end behavior of TIR subroutine calls. PRs building up to this point: - TVMScript - apache#14889 - apache#14915 - apache#14919 - apache#14941 - Functionality improvements of existing TIR passes - apache#14913 - apache#14914 - apache#14918 - apache#14951 - Changes to the TIR lowering flow - apache#14942 - apache#14985 - Codegen updates - apache#14958 - apache#14901 - Compatibility updates/fixes - apache#14892 - apache#14950 - apache#14943 - apache#14944 - apache#14945 - apache#14952 - apache#14982 - apache#14949

Lunderberg mentioned this pull request May 25, 2023

[TIR] Avoid re-defining var = arg_var in ArgBinder #14952

Merged

Lunderberg force-pushed the keep_trivial_letstmt_when_used_in_buffer branch 2 times, most recently from c62c8cf to 6dce15b Compare May 30, 2023 15:59

Lunderberg added a commit to Lunderberg/tvm that referenced this pull request Jun 2, 2023

Avoid Var-to-Var Let binding in AOTExecutorCodegen

af1628f

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions.

masahi approved these changes Jun 2, 2023

View reviewed changes

Lunderberg force-pushed the keep_trivial_letstmt_when_used_in_buffer branch from 6dce15b to d0c9df2 Compare June 3, 2023 19:38

Lunderberg mentioned this pull request Jun 3, 2023

[TIR] Handle callees on same target, different codegen #14988

Merged

Lunderberg added a commit to Lunderberg/tvm that referenced this pull request Jun 3, 2023

Avoid Var-to-Var Let binding in AOTExecutorCodegen

870ba06

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions.

masahi merged commit 5b9e9bd into apache:main Jun 4, 2023

Lunderberg deleted the keep_trivial_letstmt_when_used_in_buffer branch June 4, 2023 13:32

Lunderberg added a commit to Lunderberg/tvm that referenced this pull request Jun 5, 2023

Avoid Var-to-Var Let binding in AOTExecutorCodegen

5286508

Prior to apache#14951, these can have erroneous simplifications when used in buffer definitions.

Lunderberg mentioned this pull request Jun 6, 2023

[AOT] Avoid Var-to-Var Let binding in AOTExecutorCodegen #15033

Merged

ysh329 mentioned this pull request Jul 12, 2023

[Release] v0.13.0 Release Candidate Notes #15295

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl #14951

[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl #14951

Lunderberg commented May 25, 2023

tvm-bot commented May 25, 2023

Lunderberg commented May 25, 2023

masahi Jun 2, 2023

Lunderberg Jun 3, 2023

Lunderberg Jun 3, 2023

		@@ -168,5 +168,21 @@ def main_kernel(n: T.int32):
		return mod


		class TestSplitHostDevice(BaseCompare):

[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl #14951

[TIR] Keep trivial LetStmt in tir.Simplify when used in buffer decl #14951

Conversation

Lunderberg commented May 25, 2023

tvm-bot commented May 25, 2023

Lunderberg commented May 25, 2023

masahi Jun 2, 2023

Choose a reason for hiding this comment

Lunderberg Jun 3, 2023

Choose a reason for hiding this comment

Lunderberg Jun 3, 2023

Choose a reason for hiding this comment