[Frontend][ONNX] Fix softmax converter when input shape is dynamic

[hope51607]

When I run the onnx version mobilenet_v1 from here without freeze_params, I meet an error:

tvm/python/tvm/relay/frontend/onnx.py:6006: in _convert_operator
    sym = convert_map[op_name](inputs, attrs, self._params)
tvm/python/tvm/relay/frontend/onnx.py:2426: in _impl_v1
    data_reshape = _op.reshape(inputs[0], newshape=reshape_shape)

...

        if isinstance(newshape, (tuple, list)):
            tempshape = []
            print("newshape: ", newshape)
            for shape in newshape:
                if isinstance(shape, _expr.IntImm):
                    tempshape.append(shape.value)
                else:
                    try:
                        print("shape: ", shape)
>                       tempshape.append(int(shape))
E                       TypeError: int() argument must be a string, a bytes-like object or a number, not 'Any'
tvm/python/tvm/relay/op/transform.py:311: TypeError

This is caused by Softmax encountering an input with a dynamic shape and using it to construct reshape.
I think simple solution is to directly create Softmax when axis==-1.
I also rewrite the Softmax test to cover dynamic case.

@blackkker @AndrewZhaoLuo @mbrookhart please give it a look.

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

[AndrewZhaoLuo]

I see, LGTM. Yes this dynamism support stuff can be confusing and inconsistent sometimes. There are failing tests with somewhat scary errors so please take a look at those.

[hope51607]

Hi @AndrewZhaoLuo, thanks for your reply.

Consider the testing relay graph

def @main(%in: Tensor[(1, 10), float32], %shape: Tensor[(2), int64]) {
  %0 = dyn.reshape(%in, %shape, newshape=[]);
  nn.softmax(%0, axis=1)
}

I found that when the target is cuda, the TIR lowering from softmax contains

allocate(T_softmax_exp: Pointer(warp float32), float32, [any_dim_3]), storage_scope = warp;

which seems to allocate warp memory with any.

I am not sure why the reason behind the problem is that cuda's softmax does not support dynamic shapes? or is it a softmax schedule issue?
I think this issue seems to be out of scope of this PR.
Maybe I should mark skip in these softmax tests with dynamic shapes when target is cuda?

Can anyone give me some hints?
Thanks.

[AndrewZhaoLuo]

I'm not aware of dynamism support for softmax in CUDA but my guess is it's a schedule issue. I'm fine with you marking the dynamic tests as skip with cuda target. Just add a TODO please

[hope51607]

@AndrewZhaoLuo got it. I've marked cuda as xfailed and added a TODO comment.