tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Relax][BlockBuilder] Use PrimValue to provide tir_vars

Open Lunderberg opened this issue 1 year ago • 2 comments

Prior to this commit, if a TIR variable was required to compute the output of BlockBuilder.call_te, but that TIR variable could not be inferred from the shape of any tensor arguments, it would be provided in an optional tir_vars argument to R.call_tir. In C++, this would be then be accessed as an optional call->args[2].as<ShapeExprNode>().

This extra argument can cause unexpected bugs. For example, the bug that was fixed in https://github.com/apache/tvm/pull/17086 was caused by RewriteDataflowReshape identifying the output buffer using prim_func->buffer_map.Get(prim_func->params.back()), which is only correct if tir_vars is empty. Rather than fixing these issues as they come up, it would be better to make the general Relax guarantees stronger by removing the tir_vars argument altogether.

Use of extra R.shape parameter to specify additional tir_vars predates the existence of relax::PrimValue, and is no longer required. This commit updates BlockBuilder.call_te to use additional relax.PrimValue arguments to handle symbolic values that cannot be inferred from tensor shapes, rather than tir_vars.

Lunderberg avatar Jun 12 '24 19:06 Lunderberg

cc @tqchen @Hzfengsy

masahi avatar Jun 12 '24 19:06 masahi

Rebased onto main to avoid stale CI results.

Lunderberg avatar Sep 11 '24 16:09 Lunderberg