[CUDA] Support memory reuse for dynamic shared memory

[jinhongyii]

This PR applies part of the algorithm of storage rewrite to detect memory reuse possibility for dynamic shared memory, and does rewrite to the body. This functionality is important in matmul because we may want to allocate the shared memory for data fetch and the shared memory for write back in the same location.

[jinhongyii]

cc: @junrushao1994 @vinx13 @masahi

[masahi]

Hmm, for reuse on the constant-size dynamic shmem, I thought the existing storage_write pass would do the job, #8571 (comment). I wonder why it is not sufficient?

[vinx13]

I think this is to support reusing heterogenous data type allocs

[jinhongyii]

@masahi

1. storage_rewrite can't handle buffer with different dtypes
2. storage_rewrite can't allocate a buffer in the place of another 2 buffers. For example, you can take a look at https://github.com/jinhongyii/tvm/blob/dyn-shared-mem/tests/python/unittest/test_tir_transform_merge_dynamic_shared_memory_allocations.py#L80-L104
   C_sh's allocation reuses the memory of A_sh and B_sh. However, if we use the original algorithm of storage rewrite, it can only allocate C_sh in the place of A_sh or B_sh. The result allocation size would be block * block * 6

[masahi]

Thanks @jinhongyii, that you are able to reduce the alloc size from block * block * 3 * 4 to block * block * 4 in the test is very cool!

[junrushao]

Let's double check and get it in if there is no further issue :-)

[jinhongyii]

@vinx13 @Hzfengsy comments are all addressed. Please take another look.

[vinx13]

Otherwise Lgtm

[Hzfengsy]

Thanks @jinhongyii