Fix InternalError in StaticPlanBlockMemory when visiting DataflowBlockNode

[Thrsu]

This PR fixes an internal error #17488

This error happens because the visitor class StorageAllocatorBaseVisitor does not correctly handle DataflowBlockNode instances. Specifically, the VisitBindingBlock_ method is not overridden for DataflowBlockNode, leading to an empty block_stack_ when it is expected to contain the current block.

To fix this issue, we need to override the VisitBindingBlock_ method for const DataflowBlockNode* in the StorageAllocatorBaseVisitor class. By doing so, we ensure that the block_stack_ is correctly managed when visiting dataflow blocks, similar to how it is managed for regular binding blocks.

[Thrsu]

Could you please help review this pull request? @tqchen @yongwww

[tqchen]

cc @MasterJH5574 can you take a look

[MasterJH5574]

Hi @Thrsu, thanks for contributing! Actually we expect all dataflow blocks to have been removed in order to properly apply memory planning, given we no longer need the dataflow block information in subsequent lowering passes. Here is our lowering pipeline and you can see that we apply ToNonDataflow before other lowering passes, including the memory planning

tvm/python/tvm/relax/pipeline.py

Lines 89 to 92 in 3e386fd

	transform.ToNonDataflow(),
	transform.RemovePurityChecking(),
	transform.CallTIRRewrite(),
	transform.StaticPlanBlockMemory(),

So if possible, do you mind sharing the reason why you want the memory planning pass to work with dataflow blocks? A code example may be sufficient, thanks!

[tqchen]

@MasterJH5574 in that case, would be great to have a proper error message when assumption do not hold

[Thrsu]

Hi @Thrsu, thanks for contributing! Actually we expect all dataflow blocks to have been removed in order to properly apply memory planning, given we no longer need the dataflow block information in subsequent lowering passes. Here is our lowering pipeline and you can see that we apply ToNonDataflow before other lowering passes, including the memory planning

tvm/python/tvm/relax/pipeline.py

Lines 89 to 92 in 3e386fd

	transform.ToNonDataflow(),
	transform.RemovePurityChecking(),
	transform.CallTIRRewrite(),
	transform.StaticPlanBlockMemory(),

So if possible, do you mind sharing the reason why you want the memory planning pass to work with dataflow blocks? A code example may be sufficient, thanks!

Thanks for your response! @tqchen @MasterJH5574
I want to customize a pass pipeline for exploring the possibility of improving the optimization performance. However, it crashed unexpectedly due to the assumptions about the pass. This PR aims to enable the StaticPlanBlockMemory optimization to run correctly in such conditions.

[MasterJH5574]

@Thrsu I see. Thanks for elaborating! The memory planning pass also has a dependency on the transform.CallTIRRewrite pass, which you may also need to be careful about. Could you please add a unit test to this file, showing that the memory planning works for dataflow blocks? With the unit test we can merge this PR. Thanks!

https://github.com/apache/tvm/blob/main/tests/python/relax/test_transform_static_plan_block_memory.py

[Thrsu]

@Thrsu I see. Thanks for elaborating! The memory planning pass also has a dependency on the transform.CallTIRRewrite pass, which you may also need to be careful about. Could you please add a unit test to this file, showing that the memory planning works for dataflow blocks? With the unit test we can merge this PR. Thanks!

https://github.com/apache/tvm/blob/main/tests/python/relax/test_transform_static_plan_block_memory.py

@MasterJH5574 Thank you for the guidance! I've added a unit test to the specified file, demonstrating that the memory planning pass works for dataflow blocks as requested. Could you please review the latest changes? Thank you for your time and assistance!

[MasterJH5574]

@Thrsu Thank you for adding the unit test! It looks good to me :-)