[MediaTek] mtk_neuron compile error: 3D permutation not support

[ruitard]

🐛 Describe the bug

I want to adapt the OLMoE model on the MTK platform, refer olmoe

I got the RuntimeError:

Traceback (most recent call last):
  File "/home/zhangrui5/projects/example/model.py", line 129, in <module>
    delegated_program = edge_program.to_backend(partitioner)
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/program/_program.py", line 93, in wrapper
    return func(self, *args, **kwargs)
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/program/_program.py", line 1332, in to_backend
    new_edge_programs[name] = to_backend(program, partitioner)
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 396, in _
    tagged_graph_module = _partition_and_lower(
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 319, in _partition_and_lower
    partitioned_module = _partition_and_lower_one_graph_module(
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 249, in _partition_and_lower_one_graph_module
    lowered_submodule = to_backend(
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/functools.py", line 889, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/exir/backend/backend_api.py", line 113, in _
    preprocess_result: PreprocessResult = cls.preprocess(
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/executorch/backends/mediatek/preprocess.py", line 88, in preprocess
    model_bytes = mtk_neuron.compile(mlir_str, " ".join(compile_options))
  File "/home/zhangrui5/.conda/envs/mtk/lib/python3.10/site-packages/mtk_neuron/mtk_neuron.py", line 127, in compile
    raise RuntimeError(f'Compile error:\n{status["log"]}')
RuntimeError: Compile error:
NIR[0]: TransposeLayer
 ├ MDLA: Cannot support Int16 input
 ├ MDLA: Cannot support Int16 output
 ├ EDPA: false Transpose is not supported
 ├ EDPA: permSupport == perm 3D permutation not support
NIR[2]: ReshapeLayer
 ├ MDLA: Cannot support Int16 input
 ├ MDLA: Cannot support Int16 output
 ├ EDPA: unsupported operation
NIR[5]: ReshapeLayer
 ├ MDLA: Cannot support Int16 input
 ├ MDLA: Cannot support Int16 output
 ├ EDPA: unsupported operation
NIR[8]: ReshapeLayer
 ├ MDLA: Cannot support Int16 input
 ├ MDLA: Cannot support Int16 output
 ├ EDPA: unsupported operation
NIR[11]: ReshapeLayer
 ├ MDLA: Cannot support Int16 input
 ├ MDLA: Cannot support Int16 output
 ├ EDPA: unsupported operation

It might be caused by https://github.com/huggingface/transformers/blob/main/src/transformers/models/olmoe/modeling_olmoe.py#L632

Versions

Could you please provide some suggestions? @neuropilot-captain

cc @cccclai @neuropilot-captain @cbilgin

[JacobSzwejbka]

@cccclai