coffezhou opened a new issue, #18002:
URL: https://github.com/apache/tvm/issues/18002
### 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:
```c
[array([0., 0., 0., 0., 0.], dtype=float32)]
```
However, when legalize the LayerNormalization, TVM crashes as follows:
```c
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=0
```
### Environment
OS: Ubuntu 20.04
TVM: 0.21.dev0 (3db71bb3a)
### 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.
```python
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()
```
[testcase.zip](https://github.com/user-attachments/files/20359014/testcase.zip)
### Triage
* needs-triage
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
