[mlir][tosa]Tosa to Linalg indexingmap bug fix

https://github.com/llvm/llvm-project/compare/main...AmosLewis:llvm-project:tosa-to-linalg

I meet this bug when I tried to add aten.where.self ops to tosa in torch-mlir project. I write a wip patch to fix this.


> (mlir_venv) nod% python -m e2e_testing.main -f 'ElementwiseAtenWhereSelf' --config=tosa -v 
> Compiling ElementwiseAtenWhereSelfModule_basic...
> FAIL - "ElementwiseAtenWhereSelfModule_basic"
> 
> Unexpected outcome summary:
> 
> ****** Failed tests - 1 tests
>     FAIL - "ElementwiseAtenWhereSelfModule_basic"
>         Compilation error: Traceback (most recent call last):
>           File "/home/chi/src/ubuntu20/shark/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir_e2e_test/framework.py", line 290, in compile_and_run_test
>             compiled = config.compile(test.program_factory())
>           File "/home/chi/src/ubuntu20/shark/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir_e2e_test/configs/tosa_backend.py", line 35, in compile
>             return self.backend.compile(module)
>           File "/home/chi/src/ubuntu20/shark/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir_e2e_test/tosa_backends/linalg_on_tensors.py", line 57, in compile
>             run_pipeline_with_repro_report(
>           File "/home/chi/src/ubuntu20/shark/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/compiler_utils.py", line 73, in run_pipeline_with_repro_report
>             raise TorchMlirCompilerError(trimmed_message) from None
>         torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering TOSA to Linalg-on-Tensors failed with the following diagnostics:
>         error: 'linalg.generic' op expected indexing_map #2 to have 4 dim(s) to match the number of loops
>         note: see current operation: 
>         %2 = "linalg.generic"(%1, %arg1, %arg2, %0) ({
>         ^bb0(%arg3: i1, %arg4: f32, %arg5: f32, %arg6: f32):
>         %3 = "arith.select"(%arg3, %arg4, %arg5) : (i1, f32, f32) -> f32
>         "linalg.yield"(%3) : (f32) -> ()
>         }) {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<() -> ()>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], iterator_types = ["parallel", "parallel", "parallel", "parallel"], operand_segment_sizes = array<i32: 3, 1>} : (tensor<1x5x5xi1>, tensor<1x12x5x5xf32>, tensor<f32>, tensor<1x12x5x5xf32>) -> tensor<1x12x5x5xf32>

llvm / llvm-project

[mlir][tosa]Tosa to Linalg indexingmap bug fix #58154