[TIR] StmtFunctor RenewDefs

[Hzfengsy]

In this PR, I introduce a StmtFunctor ReGenerateDef for deep copy all definition nodes in PrimFunc (including Var, Buffer, and IterVar). This functor can create a new PrimFunc with the same behavior as the old one but contains different Nodes.

This Functor may help TIR fusion or inline multiple PrimFuncs

cc @tqchen @jinhongyii

[tqchen]

cc @junrushao1994 @vinx13 @jinhongyii @MasterJH5574 please help to review

[MasterJH5574]

“Regenerate” itself is already a word. So I suggest we using “RegenerateDef” instead of “ReGenerateDef” 🤔

[Hzfengsy]

Please take another look @tqchen @jinhongyii @junrushao1994 @MasterJH5574

[tqchen]

Some additional ideas for possible regressions

Symbolic shape case that needs remap

Symbolic vars needs to be mapped to the same one

def fn(a: T.handle, b: T.handle, n: T.int32)
   m = T.var("int32")
   A =  match_buffer((n, m), a)
   B = match_bufer((n, m*2), b) 

[tqchen]

Some high level comments, no correctness issue but would be good to add test coverage.

[junrushao]

Some additional ideas for possible regressions

Symbolic shape case that needs remap

Symbolic vars needs to be mapped to the same one

def fn(a: T.handle, b: T.handle, n: T.int32)
   m = T.var("int32")
   A =  match_buffer((n, m), a)
   B = match_bufer((n, m*2), b) 

Thanks for this interesting case!

As a further step, there might be usecases where buffers are declared as divisible by power-of-2, for example:

def fn(a: T.handle):
    m = T.var("int32")
    A = match_buffer(a, (m * 8, ), "float32")

@tqchen do you think this is a case we want to support?

[tqchen]

@junrushao1994 what you are describing is not a case we support right now, but can be something we want to support in the future and indeed an useful case

[junrushao]

Any updates?

[Hzfengsy]

@junrushao1994 @tqchen please take another look at it.
also cc @spectrometerHBH @wrongtest @jinhongyii @MasterJH5574