masahi commented on pull request #7425: URL: https://github.com/apache/tvm/pull/7425#issuecomment-788690275
@tqchen @junrushao1994 @vinx13 @ZihengJiang @zxybazh I came to a conclusion that While node doesn't need a special handling in `storage_rewrite`. The first observation is that even if I remove all `ForNode` handling from `StoragePlanRewriter`, all tests in `test_tir_transform_storage_rewrite.py` except [test_parallel_alloc()](https://github.com/apache/tvm/blob/7340c02d0efe0f5eb5692fb9f4cc7573c5d056cb/tests/python/unittest/test_tir_transform_storage_rewrite.py#L269) pass. If we look at the visitor for `ForNode`, https://github.com/apache/tvm/blob/7340c02d0efe0f5eb5692fb9f4cc7573c5d056cb/src/tir/transforms/storage_rewrite.cc#L440-L452 it only does something special when `attach_map_` has an entry for this node. Here comes the second observation: the only case where`attach_map_` can have an entry for `ForNode` is if this `ForNode` is a parallel for loop, due to these lines: https://github.com/apache/tvm/blob/7340c02d0efe0f5eb5692fb9f4cc7573c5d056cb/src/tir/transforms/storage_rewrite.cc#L766-L772 Together, these two handler for `ForNode` lift allocation inside an inner loop and attach merged allocation under the parallel loop scope (via `MakeAttach` function at https://github.com/apache/tvm/blob/7340c02d0efe0f5eb5692fb9f4cc7573c5d056cb/src/tir/transforms/storage_rewrite.cc#L447). This is tested in `parallel_alloc`. For other kinds of `For` loop, a merged allocation is placed at the global scope, see https://github.com/apache/tvm/blob/7340c02d0efe0f5eb5692fb9f4cc7573c5d056cb/src/tir/transforms/storage_rewrite.cc#L457-L461. Since `While` node doesn't involve threading, I think we can always lift allocation done inside `While` loop into the global scope. That means `WhileNode` should be handled in the same way non-parallel `ForNode` are handled, i.e. we don't need a special handling logic for `WhileNode`. I think I nailed it, thoughts? ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
