Open
Description
Expected behavior
LegalizeOps should process the LayerNormalization of valid onnx model correctly.
Actual behavior
The onnx model can be run by onnxruntime and outputs the following results:
[array([0., 0., 0., 0., 0.], dtype=float32)]However, when legalize the LayerNormalization, TVM crashes as follows:
Traceback (most recent call last):
File "/home/carla/Documents/test/test.py", line 41, in <module>
main()
File "/home/carla/Documents/test/test.py", line 37, in main
tvm_model = relax.transform.LegalizeOps()(tvm_model)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/carla/Documents/tvm/python/tvm/ir/transform.py", line 238, in __call__
return _ffi_transform_api.RunPass(self, mod)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "tvm/ffi/cython/./function.pxi", line 228, in tvm.ffi.core.Function.__call__
File "tvm/ffi/cython/./function.pxi", line 281, in tvm.ffi.core.tvm_ffi_callback
File "/home/carla/Documents/tvm/python/tvm/relax/transform/legalize_ops/nn.py", line 613, in _nn_layer_norm
return bb.call_te(
File "/home/carla/Documents/tvm/python/tvm/relax/block_builder.py", line 356, in call_te
tir_func, call_args, output_sinfo, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
File "/home/carla/Documents/tvm/python/tvm/relax/utils.py", line 354, in gen_call_tir_inputs
te_out = func(*te_args, **te_kwargs)
File "/home/carla/Documents/tvm/python/tvm/topi/nn/layer_norm.py", line 48, in layer_norm
return cpp.nn.layer_norm(data, gamma, beta, axis, epsilon)
File "tvm/ffi/cython/./function.pxi", line 228, in tvm.ffi.core.Function.__call__
tvm.error.InternalError: Check failed: shape.size() == indices.size() (1 vs. 0) : Tensor dimension mismatch in read ndim = 1, indices.size=0Environment
OS: Ubuntu 20.04
TVM: 0.21.dev0 (3db71bb)
Steps to reproduce
This bug can be reproduced by the following code with the model in the attachment. As shown in the code, the model can be executed by onnxruntime.
import sys
import numpy as np
import onnx
import onnxruntime
import tvm
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
import pickle
def main():
onnx_model = onnx.load("a552.onnx")
shape_onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
onnx.save(shape_onnx_model, '1111.onnx')
with open("inputs.pkl", "rb") as fp:
inputs = pickle.load(fp)
try:
ort_session = onnxruntime.InferenceSession(
onnx_model.SerializeToString(), providers=["CPUExecutionProvider"]
)
ort_output = ort_session.run([], inputs)
except Exception as e:
print(e)
sys.exit(1)
print("ONNXRuntime:\n", ort_output)
# Convert the onnx model into relax through the onnx importer.
tvm_model = from_onnx(onnx_model, keep_params_in_input=True)
# Convert operators for inference mode.
tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)
# Legalize any relax ops into tensorir.
tvm_model = relax.transform.LegalizeOps()(tvm_model)
if __name__ == "__main__":
main()Triage
- needs-triage