Test for a inception module

[Ruinhuang]

Hi, I wan to use AOT to test concatenate op, but it has some error.

import numpy as np
import tvm
import tvm.relay as relay

from aot import compile
from tvm.relay.transform import gradient
from tvm.relay.testing import ctx_list, run_infer_type


def compute(data, axis):
    relay_op = relay.concatenate
    relay_x = []
    for x in data:
        relay_x.append(relay.var("input", relay.TensorType(x.shape, "float32")))
    y = relay_op(relay_x, axis)

    fwd_func = relay.Function(relay_x, y)
    fwd_func = run_infer_type(fwd_func)

    bwd_func = run_infer_type(gradient(fwd_func))

    print("data: ", data)
    tgt = tvm.target.create('llvm')
    ctx = tvm.context('llvm', 0)
    mod = relay.module.Module()
    intrp_wrapper = compile(bwd_func, mod, ctx=ctx, tgt=tgt)
    output = intrp_wrapper(*data)
    print("out: ", output)


def test_concatenate():
    """ Test for concatenate operator """
    print('\n----------------------Test start------------------------')

    def verify_concatenate(dshapes, dtype, axis):
        data = []
        for shape in dshapes:
            x = np.random.rand(*shape).astype(dtype)
            data.append(x)

        compute(data, axis)

    verify_concatenate([(2, 3), (3, 3)], 'float32', 0) # Success
    verify_concatenate([(2, 3), (3, 3), (4, 3), [5, 3]], 'float32', 0) # Success
    verify_concatenate([(2, 3), (3, 3), (4, 3)], 'float32', 0) # Failed
    verify_concatenate([(2, 3), (3, 3), (4, 3), (5, 3), (6, 3)], 'float32', 0) # Failed


if __name__ == '__main__':
    test_concatenate()
 

when i concatenate 2 or 4 array, it can be compiled. but when the parameter is 3 or 5, it can not be compiled, this is the error log:

Traceback (most recent call last):

  File "/home/rui.huang/tvm-0813/tests/python/relay/train/models/test_aot.py", line 50, in <module>
    test_concatenate_grad()

  File "/home/rui.huang/tvm-0813/tests/python/relay/train/models/test_aot.py", line 45, in test_concatenate_grad
    verify_concatenate([(2, 3), (3, 3), (4, 3)], 'float32', 0)

  File "/home/rui.huang/tvm-0813/tests/python/relay/train/models/test_aot.py", line 41, in verify_concatenate
    compute(data, axis)

  File "/home/rui.huang/tvm-0813/tests/python/relay/train/models/test_aot.py", line 26, in compute
    intrp_wrapper = compile(bwd_func, mod, ctx=ctx, tgt=tgt)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 264, in compile
    func = compiler.visit(func)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 43, in visit
    res = self.visit_function(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 162, in visit_function
    return CPPFunction(func.params, self.visit(func.body), func.checked_type.ret_type)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 43, in visit
    res = self.visit_function(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 162, in visit_function
    return CPPFunction(func.params, self.visit(func.body), func.checked_type.ret_type)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 43, in visit
    res = self.visit_function(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 162, in visit_function
    return CPPFunction(func.params, self.visit(func.body), func.checked_type.ret_type)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 47, in visit
    res = self.visit_let(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 139, in visit_let
    cpp_value = self.visit(let.value)

  File "/home/rui.huang/tvm-0813/python/tvm/relay/expr_functor.py", line 45, in visit
    res = self.visit_call(expr)

  File "/home/rui.huang/tvm-0813/relay-aot/python/aot/aot.py", line 130, in visit_call
    assert (call.attrs == None)

AssertionError

I make a break point at aot.py(def visit_call(self, call: Expr) -> Expr:), the log shows:

call = {Call} v0.0.3\nfree_var %p0: Tensor[(9, 3), float32]\nsplit(%p0, indices_or_sections=[2, 5]) /* ty=(Tensor[(2, 3), float32], Tensor[(3, 3), float32], Tensor[(4, 3), float32]) */
 _checked_type_ = {TupleType} v0.0.3\n(Tensor[(2, 3), float32], Tensor[(3, 3), float32], Tensor[(4, 3), float32])
 args = {Array} [Var(p0, ty=TensorType([9, 3], float32))]
 attrs = {SplitAttrs} relay.attrs.SplitAttrs(0x1b9f7d0)
 op = {Op} v0.0.3\nsplit
 span = {NoneType} None
 type_args = {Array} [TensorType([9, 3], float32)]
self = {AoTCompiler} <aot.aot.AoTCompiler object at 0x7f1fa7eb6fd0>

As i know when it can be compiled, the call.attrs should be None, but when the parameter is odd number, the call.attrs is SplitAttrs type.
By the way, the concatenate's backward op is split, I implement it by this way:

@register_gradient("concatenate")
def concatenate_grad(orig, grad):
    """
    Return concatenate gradient
    :param orig: attrs(axis)
    :param grad: initial gradient computed for execution result of concatenate
    :return:
    """

    axis = orig.attrs.axis

    # Compute the data shape after split
    data_type = orig.type_args[0]
    indices = []
    split_indices = 0

    for field in data_type.fields[:-1]:
        split_indices += field.shape[axis]
        indices.append(split_indices)

    return [split(grad, indices, axis)]

it can be compiled by using JIT

Can anyone give me some suggestion? thanks very much

[Ruinhuang]

@MarisaKirisame
could you give me some suggestion? thanks

[slyubomirsky]

I'll study this example. We've been a little pressed for time, but it seems like something we should be able to debug.

Thank you for the amount of detail you've provided! It will be helpful

[slyubomirsky]

Marisa asks: Have you been able to run this without AoT? He thinks that there may be a bug with fusion for the tuple output case and if that's the case, then this will fail with the default interpreter, too. If not, the bug is likely within AoT

[slyubomirsky]

Fwiw I got a type-checking error in Relay when I tried the code you provided so there may be reasons unrelated to AoT for why it failed

[shiyf369]

@slyubomirsky Hi, I modified the code into the JIT executor, it can be executed normally. So we think this is supposed to be a bug of AOT executor. Could you please take a look at it and give us some suggestion? thanks a lot

Modified JIT executor:

def compute(data, axis):
    relay_op = relay.concatenate
    relay_x = []
    for x in data:
        relay_x.append(relay.var("input", relay.TensorType(x.shape, "float32")))
    y = relay_op(relay_x, axis)

    fwd_func = relay.Function(relay_x, y)
    fwd_func = run_infer_type(fwd_func)

    bwd_func = run_infer_type(gradient(fwd_func))

    print("data: ", data)
    tgt = tvm.target.create('llvm')
    ctx = tvm.context('llvm', 0)
    # mod = relay.module.Module()
    # intrp_wrapper = compile(bwd_func, mod, ctx=ctx, tgt=tgt)

    intrp_wrapper = relay.create_executor(ctx=ctx, target=tgt).evaluate(bwd_func)

    output = intrp_wrapper(*data)
    print("out: ", output)

Running result of JIT executor:

----------------------Test start------------------------
data:  [array([[0.9766186 , 0.75620013, 0.82304156],
       [0.2569066 , 0.18649143, 0.4070408 ]], dtype=float32), array([[0.32892945, 0.6248476 , 0.5253376 ],
       [0.8726906 , 0.2102789 , 0.9473281 ],
       [0.31386733, 0.43369105, 0.10988309]], dtype=float32), array([[0.80398893, 0.5437144 , 0.2712833 ],
       [0.7965053 , 0.02169295, 0.14893158],
       [0.7103855 , 0.0458868 , 0.8380036 ],
       [0.66297364, 0.60816246, 0.11053047]], dtype=float32)]
out:  ([[0.9766186  0.75620013 0.82304156]
 [0.2569066  0.18649143 0.4070408 ]
 [0.32892945 0.6248476  0.5253376 ]
 [0.8726906  0.2102789  0.9473281 ]
 [0.31386733 0.43369105 0.10988309]
 [0.80398893 0.5437144  0.2712833 ]
 [0.7965053  0.02169295 0.14893158]
 [0.7103855  0.0458868  0.8380036 ]
 [0.66297364 0.60816246 0.11053047]],([[1. 1. 1.]
 [1. 1. 1.]],[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]],[[1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]
 [1. 1. 1.]]))

Process finished with exit code 0

[shiyf369]

@slyubomirsky We use the tvm master version of 2019.08.10. Is it possible that the version you use is older?

[MarisaKirisame]

@shiyf369 I will look into it, this error is mysterious.

[shiyf369]

@MarisaKirisame Thanks a lot

[MarisaKirisame]

I am 90% sure this is a bug in ops fusion.

@tqchen can you take a look? that two line should be in %27.

[MarisaKirisame]

@Ruinhuang can you rebase onto apache/tvm#3871?
I tried running on 3871, and it get through.

[Ruinhuang]

@MarisaKirisame okay, i will try it. This problem is too odd, thanks for your great help~

[MarisaKirisame]

@Ruinhuang if it fix it please reply so I can close.

[Ruinhuang]

@MarisaKirisame i have test the case, it can be executed, you can close it. thanks again~

[MarisaKirisame]

@Ruinhuang it is fixed. can you upstream concat_grad?

[MarisaKirisame]

@Ruinhuang @shiyf369 any update?

[Ruinhuang]

@Ruinhuang it is fixed. can you upstream concat_grad?

I will add concat_grad to the tvm code recently.
No new issues so far