[Bug] [TVMScript] Progress hang on lowerTE during build dynamic Relay IR caused by endless ReprPrintTIR.

[cccxinli]

Progress hang on lowerTE during build dynamic Relay IR.

Expected behavior

Build success.

Actual behavior

Hang on.

Environment

OS: Linux
tvm version: v0.12.0

Steps to reproduce

import tvm
from tvm import relay
import numpy as np

inp0_shape = (10, 10, 10, 10)
inp1_shape = (10, )

inp0 = relay.var("inp0", relay.TensorType([relay.Any() for i in range(len(inp0_shape))], "float32"))
inp1 = relay.var("inp1", relay.TensorType([relay.Any() for i in range(len(inp1_shape))], "float32"))
add0 = relay.op.add(inp0, inp1)
add1 = relay.op.add(add0, inp1)
for i in range(15):
    add2 = relay.op.add(add1, inp1)
    cur = add2
    add1 = cur
div0 = relay.op.divide(cur, inp0)
func = relay.Function([inp0, inp1], div0)

data0_np = np.random.uniform(size=inp0_shape).astype("float32")
data1_np = np.random.uniform(size=inp1_shape).astype("float32")

op_res = relay.create_executor(
    kind="vm", device=tvm.cpu(), target="llvm"
).evaluate(func)(data0_np, data1_np)

Triage

• needs-triage

cc @Hzfengsy @junrushao @shingjan

[cccxinli]

It will work normal if delete the string replacement in python/tvm/relay/op/_tensor.py#L254.

 assert y[ndim2 - i] == 1, "Incompatible broadcast type %s and %s" % (
                    x[ndim1 - i],
                    y[ndim2 - i],
                )

to

assert y[ndim2 - i] == 1, "Incompatible broadcast type."

[Hzfengsy]

I'm not sure what happens, but it is not related to TVMScript I guess

[cccxinli]

I tried to debug it, and I found:

1. Why call ReprPrintTIR function when lower Relay IR？
   It will step into function LowerShapeFunc (src/relay/backend/te_compiler.cc:770) because node is dynamic, and will call shape function for node. The shape function of add op has been registered so that related to the broadcast_shape_func (python/tvm/relay/op/_tensor.py).
   The _broadcast_shape_func is decorated by script. During parse_python in source_to_op (python/tvm/te/hybrid/parser.py), the string replacement in code:

 assert y[ndim2 - i] == 1, "Incompatible broadcast type %s and %s" % (
                    x[ndim1 - i],
                    y[ndim2 - i],
                )

will call x[ndim1 - i].repr() which globally registered in src/node/repr_printer.cc:90. Finally step into ReprPrinTIR in src/script/printer/tir/utils.h:170, and will endless in SetCommonPrefix in this function.

2. Why SetCommonPrefix endless?
   This case will call broadcast_shape_func several times because it include several op.add. I count the number of times of step into Visit object function (src/script/printer/ir_docsifier.cc:93) in SetCommonPrefix. Result as follows, The next one is about eight times bigger than the last one.
   

The output of _broadcast_shape_func like:

It include several shape_p0 which is the previous output. I think that's probably the reason that visit the same shape_p0 several times.

3. How to resolve?
   The Visitor in SetCommonPrefix need to avoid visiting repetition Tensor I guess.