Skip to content

[Bug] Onnx frontend can import invalid squeeze axis #20422

Description

@debugger-zhang

In squeeze in python/tvm/relax/frontend/nn/op.py:
(1) The axis argument accepts int or list of int.
(2) The nominal default value of axis is -1 (this is just nominal since TVMScript is not Python code, so the default value is not really stored in module), which means removing all axis of dimensions 1, instead of removing the last dim.

In squeeze in python/tvm/relax/op/manipulate.py:
(1) If the axis is int, then it is converted to a list of one element.

In InferLayoutSqueeze in src/relax/op/tensor/manipulate.cc:
(1) The axis attribute is read as-is, and for each value of axis the corresponding element axis_str is set. This means all elements of axis can not be negative, otherwise this is an out-of-bound write.

In Squeeze in python/tvm/relax/frontend/onnx/onnx_frontend.py:
(1) The axis are stored as-is from onnx (see https://github.com/apache/tvm/blob/5398c27e76c7c858d7fae655496e68b297fc99fc/python/tvm/relax/frontend/onnx/onnx_frontend.py#L2424C9-L2424C36). This means if any value of axis is negative, then the result Relax op is invalid.

What is expected: axis must be normalized before op is created.

A simple reproduce:

import tvm, onnx
from tvm.relax.frontend.onnx import from_onnx
model = onnx.load_model_from_string(b'\x08\n\x12\x18onnx.utils.extract_model:\x8c\x02\n5\n\x0btranspose_1\n\x06val_23\x12\x07squeeze\x1a\x0cnode_squeeze"\x07Squeeze\x12\x1bExtracted from {main_graph}*\x16\x08\x01\x10\x07B\x06val_23J\x08\xfe\xff\xff\xff\xff\xff\xff\xffZ*\n\x0btranspose_1\x12\x1b\n\x19\x08\x01\x12\x15\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x02\x08\x01\n\x03\x08\x80\x08b"\n\x07squeeze\x12\x17\n\x15\x08\x01\x12\x11\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x03\x08\x80\x08j*\n\x0btranspose_1\x12\x1b\n\x19\x08\x01\x12\x15\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x02\x08\x01\n\x03\x08\x80\x08j"\n\x07squeeze\x12\x17\n\x15\x08\x01\x12\x11\n\x02\x08\x03\n\x02\x08 \n\x02\x08\x01\n\x03\x08\x80\x08B\x04\n\x00\x10\x14')
tvm.relax.transform.ConvertLayout({})(from_onnx(model))

In tvm 0.26, it will throw std::bad_alloc.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions