[Unity] Support simple dynamic-shape-aware fusion

[Hzfengsy]

This PR adds support for simple dynamic-shape-aware fusion, which is the first step towards supporting dynamic shapes. The main changes are as follows:

• Fix FuncStructInfo in well-formed checks
• Renew symbolic var defs in fuse_ops to prevent malformed functions

cc @tqchen @MasterJH5574 @jinhongyii

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[masahi]

Hi @Hzfengsy, it seems this change broke one of CUTLASS BYOC tests: pytest test_codegen_cutlass.py::test_matmul_offload[float16-x_shape17-y_shape17-False-none-none] -s.

If I remove SymbolicVarRenewMutator::Renew, the test works. I don't know why it creates an undefined sym var for a BYOC function. Can you take a look?

cc @junrushao @vinx13 @yelite

[Hzfengsy]

Hi @masahi, thanks for reporting the bug.

SymbolicVarRenewMutator::Renew works to re-define all symbolic vars to prevent the same symbolic var in different functions.

[Hzfengsy]

@masahi, After a deep dive, it's may not the issue of FuseOps, as the output is well-formed and pass the check. The error is reported at the pass RunCodegen

[masahi]

Yes, it might be a bug in RunCodegen. And IRs produced by with or without SymbolicVarRenewMutator::Renew look identical at least textually. So I'm confused what makes SymbolicVarRenewMutator::Renew break RunCodegen.

[masahi]

However, it does not update the out_sinfo of the body

Why not? Does this mean out_sinfo continue referring to an old sym var?

[Hzfengsy]

Why not? Does this mean out_sinfo refer to the old sym var?

Sorry, it updated. My bad

[masahi]

https://github.com/apache/tvm/blob/unity/src/relax/transform/run_codegen.cc#L105

Maybe we shouldn't use the original function's ret_struct_info there. After RunCodegen, it indeed references a dubious sym var:

@R.function                                                                                                                                          
def main(x: R.Tensor((1, 1, 16, 15), dtype="float16"), y: R.Tensor((3, 2, "a", 15, 2), dtype="float16")) -> R.Tensor((3, 2, "a", 16, 2), dtype="float
16"):                                                                     
    a = T.int64()                                                         
    a_1 = T.int64()   <- this shouldn't exisit                                                    
    with R.dataflow():                                                                                                                               
        gv = R.call_dps_packed("fused_relax_matmul_cutlass", (x, y), out_sinfo=R.Tensor((3, 2, a_1, 16, 2), dtype="float16"))                     
      R.output(gv)                                                      
    return   

[masahi]

@Hzfengsy Is this line in ExprMutator still correct?

https://github.com/apache/tvm/blob/unity/src/relax/ir/expr_functor.cc#L592

I think the mutated function would continue referring to the old sym var if the sym var is used in a return sinfo.

[Hzfengsy]

I think the mutated function would continue referring to the old sym var if the sym var is used in a return sinfo.

We need to update the ret_struct_info. However, I'm suprised that the undefined var in ret_sinfo is not reported by well-formed checker

[masahi]

I think this is a problem unique to BYOC. RunCodegen replaces a call to a function with an opaque call_dps_packed call. After RunCodegen runs, the original function is gone but we still need its ret_struct_info.

[masahi]

The fix has been added in the PR #14291