[Transform] Operator legalizer V0 #96

MasterJH5574 · 2023-01-11T05:04:06Z

This PR is the very first version of operator legalizer, which leverages CallTE of Relax BlockBuilder and the existing TOPI functions or newly written TE functions to lower high-level operator calls down to CallTIRs with TIR PrimFuncs.

The legalizer can pass the existing unit tests, which are mostly written in October, 2022. Those unit tests guarantee the correctness on static shapes as much as possible. So for static shape cases, the legalizer is expected to work properly.

However, the test cases are far from enough, especially the support and robustness on symbolic shapes and other cases need to be further confirmed as a recent-future work. Another to-do is to well-document the pass.

MasterJH5574 · 2023-01-11T05:06:22Z

[Update.] Dependency cleared.

Depends on tlc-pack#352, tlc-pack#353, and tlc-pack#354 for them to be merged and cherry-pick to our fork. Those PRs fix a few bugs that were revealed during my work on the legalizer unit tests. A sync is preferably needed.

python/tvm/relax/transform/op_legalizer.py

tqchen · 2023-01-15T18:02:01Z

tests/python/relax/test_op_legalizer.py

+
+        @T.prim_func
+        def sum(
+            rxplaceholder: T.Buffer[(T.int64(1), T.int64(2), T.int64(3), T.int64(4)), "float32"],


Perhaps a good motivational pt for meta-programming later cc @MasterJH5574 @Hzfengsy

* [VM] Refactor and improve vm. - Have a separate function for RunInstCall. - Cache func_index lookup by table to avoid repeative lookup by str. - Move PackedFunc call arg stack to Frame to increase locality and avoid re-allocation in repeative calls. - Make frame stack of unique_ptr to avoid frame re-allocation and copy during frame.resize. - Pass curr_frame as arguments into sub-functions to make it explicit. * address review comments

python/tvm/relax/transform/legalize_ops.py

[Fix] Buffer slicing using index dtype as extent

…lc-pack#361)

This reverts commit 99d5afc.

This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

* [TIR][Fix] Buffer slicing using index dtype as extent (#13788) [Fix] Buffer slicing using index dtype as extent * [TIR][Fix] IndexDataTypeNormalizer not unwrapping float casting (#13789) * [Fix][TVMScript] Parse and print `tir_vars` of `call_tir` properly (tlc-pack#361) * [Transform] Operator legalizer * Documentation

This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

* [TIR][Fix] Buffer slicing using index dtype as extent (#13788) [Fix] Buffer slicing using index dtype as extent * [TIR][Fix] IndexDataTypeNormalizer not unwrapping float casting (#13789) * [Fix][TVMScript] Parse and print `tir_vars` of `call_tir` properly (tlc-pack#361) * [Transform] Operator legalizer * Documentation

This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

* [VM] Refactor and improve vm. - Have a separate function for RunInstCall. - Cache func_index lookup by table to avoid repeative lookup by str. - Move PackedFunc call arg stack to Frame to increase locality and avoid re-allocation in repeative calls. - Make frame stack of unique_ptr to avoid frame re-allocation and copy during frame.resize. - Pass curr_frame as arguments into sub-functions to make it explicit. * address review comments

* [TIR][Fix] Buffer slicing using index dtype as extent (#13788) [Fix] Buffer slicing using index dtype as extent * [TIR][Fix] IndexDataTypeNormalizer not unwrapping float casting (#13789) * [Fix][TVMScript] Parse and print `tir_vars` of `call_tir` properly (tlc-pack#361) * [Transform] Operator legalizer * Documentation

This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

* [VM] Refactor and improve vm. - Have a separate function for RunInstCall. - Cache func_index lookup by table to avoid repeative lookup by str. - Move PackedFunc call arg stack to Frame to increase locality and avoid re-allocation in repeative calls. - Make frame stack of unique_ptr to avoid frame re-allocation and copy during frame.resize. - Pass curr_frame as arguments into sub-functions to make it explicit. * address review comments

This PR migrates mlc-ai/relax#46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer mlc-ai/relax#96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch mlc-ai/relax#43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in mlc-ai/relax#96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once mlc-ai/relax#94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

This PR migrates #46 to new struct info infra, as part of our AD migration. Because we need do numerical testing for gradients, this PR depends on the operator legalizer #96. Also because the original version of legalizer did not handle the negative indexing case of `relax.mean`, this PR fixes it. To lower `collapse_sum_to`, `collapse_sum_like` properly, this PR migrates a previous patch #43 which introduces `collapse_sum` in topi. Now we can remove the skip marker in the legalizer test for `collapse_sum_to` and `collapse_sum_like`. The gradients of `cross_entropy` and `softmax_cross_entropy` are removed. And the former will be added back and adjust to new `cross_entropy` introduced in #96. Further plan in this PR: - [x] Add gradients for `log_softmax` and `nll_loss` once #94 is merged. - [x] Gradients for some tuple related operators such as `split` and `concat`. It can help us to test the correctness of AD when there are Tuple-I/O operators. - (Not in this PR) "Undefined Gradient" representation. As we know, the gradients of some operators w.r.t. specified inputs are undefined or meaningless, such as the partial gradient of `indices` in `take(x, indices)`. Relay directly uses `zeros_like` in this case as it won't affect gradient propagation. Another choice is to introduce a dummy Expr named `UndefinedGradient` to represent it. How do we handle this case in relax?

@MasterJH5574

@MasterJH5574 brings the very first version of legalizer for relax high-level operators in this PR: #96. This PR refactors the `LegalizeOps` pass according to the suggestions mentioned in the comments of the PR. In details, it includes: - Change the default legalization map to the attribute of operators. By this way we can make better use of op attribute and move the main body of the Pass to C++, making `LegalizeOps` a more formal Pass and bringing opportunities that we can register legalization without Python in the future. - Use decorator to register legalization instead of writing it in the Pass. It is more elegant and brings extensibility. Also by this way we can have a clear code structure since now we separate the implementation in different files. As we will have more and more legalization for operators, putting them in a single file is not promising. - Preserve the customize map part. It looks good for now.

MasterJH5574 force-pushed the mlc-dev/2023-01-11-legalizer-v0 branch 3 times, most recently from 4c72973 to d0966f0 Compare January 11, 2023 05:27

SiriusNEO mentioned this pull request Jan 11, 2023

[Tracking Issue] Relax training M0 migration and polishment #97

Closed

20 tasks

MasterJH5574 force-pushed the mlc-dev/2023-01-11-legalizer-v0 branch from d0966f0 to 27f79ba Compare January 11, 2023 19:09

MasterJH5574 mentioned this pull request Jan 11, 2023

[Tracking Issue] Relax high-level operator migration #62

Closed

21 tasks

SiriusNEO mentioned this pull request Jan 12, 2023

[Op] Migration: Gradients for some operators #98

Merged

2 tasks

MasterJH5574 force-pushed the relax branch from ca67f29 to e4ed2c3 Compare January 13, 2023 02:31

MasterJH5574 force-pushed the mlc-dev/2023-01-11-legalizer-v0 branch 2 times, most recently from 26fb3e8 to 86a609f Compare January 13, 2023 02:35

SiriusNEO reviewed Jan 13, 2023

View reviewed changes

python/tvm/relax/transform/op_legalizer.py Outdated Show resolved Hide resolved

python/tvm/relax/transform/op_legalizer.py Outdated Show resolved Hide resolved

Ubospica reviewed Jan 13, 2023

View reviewed changes

python/tvm/relax/transform/op_legalizer.py Outdated Show resolved Hide resolved

MasterJH5574 force-pushed the mlc-dev/2023-01-11-legalizer-v0 branch from 86a609f to 93310e0 Compare January 13, 2023 15:43