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

[AmosLewis]

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>

[AmosLewis]

@rsuderman here is my commit for review https://reviews.llvm.org/D135343

[lhutton1]

Closing this issue since it seems the referenced fix was merged. Please reopen if that's not the case and the issue still persists.

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Chi_Liu (AmosLewis)

[https://github.com/llvm/llvm-project/compare/main...AmosLewis:llvm-project:tosa-to-linalg](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:
&gt;         error: 'linalg.generic' op expected indexing_map #<!-- -->2 to have 4 dim(s) to match the number of loops
&gt;         note: see current operation: 
&gt;         %2 = "linalg.generic"(%1, %arg1, %arg2, %0) ({
&gt;         ^bb0(%arg3: i1, %arg4: f32, %arg5: f32, %arg6: f32):
&gt;         %3 = "arith.select"(%arg3, %arg4, %arg5) : (i1, f32, f32) -&gt; f32
&gt;         "linalg.yield"(%3) : (f32) -&gt; ()
&gt;         }) {indexing_maps = [affine_map&lt;(d0, d1, d2, d3) -&gt; (d0, d2, d3)&gt;, affine_map&lt;(d0, d1, d2, d3) -&gt; (d0, d1, d2, d3)&gt;, affine_map&lt;() -&gt; ()&gt;, affine_map&lt;(d0, d1, d2, d3) -&gt; (d0, d1, d2, d3)&gt;], iterator_types = ["parallel", "parallel", "parallel", "parallel"], operand_segment_sizes = array&lt;i32: 3, 1&gt;} : (tensor&lt;1x5x5xi1&gt;, tensor&lt;1x12x5x5xf32&gt;, tensor&lt;f32&gt;, tensor&lt;1x12x5x5xf32&gt;) -&gt; tensor&lt;1x12x5x5xf32&gt;