llvm-project
be06c79c - [mlir][linalg] Enable Vectorization of 0-D tensor.extract (#119079)

Commit
1 year ago
[mlir][linalg] Enable Vectorization of 0-D tensor.extract (#119079) This patch removes an assert in `vectorizeTensorExtract` that was blocking the vectorization of 0-D tensor.extract operations, e.g.: ```mlir %1 = tensor.extract %src[] : tensor<f32> ``` As demonstrated by the included tests, this case is already effectively supported. **Context** The removed assert was introduced in #109580 as a guard, pending proper support and testing for 0-D tensors. This PR addresses that previously undocumented TODO. Apologies for the oversight! **Updates and Tests** * Revised the existing test `@negative_no_loop` to ensure the `vectorize_nd_extract` attribute is included, allowing the vectorizer to process it. The test was renamed and variables updated for clarity. * Added a new test `@extract_scalar_from_0d_into_1d` to cover "mixed" 0-D/1-D tensor extraction, e.g.: ```mlir %res = linalg.generic { indexing_maps = [#map], iterator_types = ["parallel"] } outs(%init : tensor<1xf32>) { ^bb0(%in: f32): %1 = tensor.extract %src[] : tensor<f32> linalg.yield %1 : f32 } -> tensor<1xf32> return %res : tensor<1xf32> ``` **Additional updates** I also took the liberty and improved test coverage for 0-D tensor in the vectorizer tests: * Added a specific test for "0D linalg.generic" in "vectorization-with-patterns.mlir". * Renamed several tests in "vectorization-with-patterns.mlir" to clarify that the 0-D case is now covered.
Author
Parents
Loading