Andrzej Warzyński d68a4b93d3
[mlir][linalg] Add support for masked vectorization of tensor.insert_slice (1/N) (#122927)
For context, `tensor.insert_slice` is vectorized using a
`vector.transfer_read` + `vector.transfer_write` pair.
An unmasked example is shown below:

```mlir
// BEFORE VECTORIZATION
%res = tensor.insert_slice
  %slice into %dest[0, %c2]
  [5, 1] [1, 1] : tensor<5x1xi32> into tensor<5x3xi32>

// AFTER VECTORIZATION
%read = vector.transfer_read %source[%c0, %c0], %pad 
  : tensor<5x1xi32>, vector<8x1xi32>
%res = vector.transfer_write %read, %dest[%c0, %c2] 
  : vector<8x1xi32>, tensor<5x3xi32>
```

This PR refactors `InsertSliceVectorizePattern` (which is used to
vectorize `tensor.extract_slice`) to enable masked vectorization. ATM,
only `vector.transfer_read` is masked. If `vector.transfer_write` also
requires masking, the vectorizer will bail out. This will be addressed
in a sub-sequent PR.

Summary of changes:
  * Added an argument to specify vector sizes (behavior remains
    unchanged if vector sizes are not specified).
  * Renamed `InsertSliceVectorizePattern` to `vectorizeAsInsertSliceOp`
    and integrated into (alongside other hooks for vectorization) in
    `linalg::vectorize`.
  * Removed `populateInsertSliceVectorizationPatterns`, as
    `InsertSliceVectorizePattern` was its only pattern.
  * Updated `vectorizeAsInsertSliceOp` to support masking for the
    "read" operation.
  * Updated `@pad_and_insert_slice_dest` in
    "vectorization-pad-patterns.mlir" to reflect the removal of
    `populateInsertSliceVectorizationPatterns` from
    `ApplyPadVectorizationPatternsOps`.
2025-02-02 14:51:25 +00:00
..