[Refactor] Buffer flatten

[junrushao]

No description provided.

[tqchen]

Thanks @junrushao1994 One thing that I noticed is the change of unordered map from strong ref object to weak ref(raw ptr). We need to be careful with such a change, especially in cases where an IR object can be mutated.

we need to be careful when using raw ptr in a mutator. Since the children can get de-allocated and re-allocated to another node (causing map info to be inaccurate). Prefer the strong ref version. See a related PR here apache/tvm#6004

Because of that reason, I think we could actually prefer to use strong ref version for unordered map (with ObjectPtrHash and Equal) for most of the maps of interest. The overhead of ref counting is not as large comparing to the gains. For a temporary data structure within a Visitor(that does not change the IR), we should be fine.

[junrushao]

@tqchen It certainly makes sense to avoid deallocation that causes hash map key to be invalid in StmtExprMutators. I wonder in which case it could happen? If we turn off the copy-on-write, I assume it would be valid during visiting. Is that correct?

[tqchen]

in some cases when we visit, say we want to replace a subtree A=> B, then after the replacement of the subtree A could no longer be referenced anywhere(thus de-allocated)

If the same address of A get re-allocated to another node during the visit, then the information in the mutator becomes stale. I cannot tell exactly which case it happens, but it indeed casued some bugs before, thus the PR from qualcomm

[junrushao]

If copy-on-write is disabled, then I assume the original IR will keep valid during the visit, which means deallocation of A will be deferred until the visit is finished. Is that correct?

[tqchen]

You are right. In our case there is possibility of COW, so we need to watch out for that

[junrushao]

Got it. I am going to change those data structure outlive the mutator to hold strong references

[junrushao]

A second thought leads me to changing every reference in the mutator to strong ones, in case helping cow

[junrushao]

@Hzfengsy This PR is basically done, please take a look when you have time :-)

There are two TODOs in the codebase:

1. Line 445. It checks specifically if the given loop is bound to threadIdx or vthread, as well as if the buffer is allocated in shared memory, which is a bit foggy to me because it couples storage scope and threads. What if we have a special storage scope like "shared.shuffled" used in cutlass?

          if (need_relax || (buffer->scope == "shared" && IsThreadBound(loop))) {
            // TODO
            dom_map[loop_var] = IntSetFromMinExtent(loop->min, loop->extent);
          }

2. Line 509: To pass this test (https://github.com/Hzfengsy/tvm-tensorir/blob/main/tests/python/tir/test_schedule_tensorize.py#L339), I changed ICHECK_EQ to ICHECK_GE, where the tensorize primitive seems to change the dimension of buffer access. I discussed with @spectrometerHBH a moment ago and he will take care of it.

Additionally, I handled double_buffer in a way different from the original codebase, please take a close look whether it makes sense. CC: @vinx13.

Thank a lot!

[vinx13]

use Map::Get

[vinx13]

should be ==?

[junrushao]

ooops