[all-commits] [llvm/llvm-project] 82cc2f: [mlir][linalg] Refactor vectorization hooks to imp...

[Andrzej Warzyński via All-commits]

  Branch: refs/heads/users/banach-space/vector/update_create_write
  Home:   https://github.com/llvm/llvm-project
  Commit: 82cc2fe047642aa50b8e6780060be8abc6181dc9
      https://github.com/llvm/llvm-project/commit/82cc2fe047642aa50b8e6780060be8abc6181dc9
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-05-29 (Thu, 29 May 2025)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    M mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp
    M mlir/test/Dialect/LLVM/transform-e2e.mlir
    M mlir/test/Dialect/Linalg/vectorization.mlir
    M mlir/test/Dialect/Linalg/vectorization/insert-slice-with-patterns.mlir
    M mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
    M mlir/test/Dialect/Linalg/vectorization/pad-with-patterns.mlir
    M mlir/test/Dialect/Vector/transform-vector.mlir

  Log Message:
  -----------
  [mlir][linalg] Refactor vectorization hooks to improve code reuse

This patch refactors two vectorization hooks in Vectorization.cpp:
 * `createWriteOrMaskedWrite` gains a new parameter for write indices,
   aligning it with its counterpart `createReadOrMaskedRead`.
 * `vectorizeAsInsertSliceOp` is updated to reuse both of the above
   hooks, rather than re-implementing similar logic.

CONTEXT
-------
This is effectively a refactoring of the logic for vectorizing
`tensor.insert_slice`. Recent updates added masking support:
  * https://github.com/llvm/llvm-project/pull/122927
  * https://github.com/llvm/llvm-project/pull/123031

At the time, reuse of the shared `create*` hooks wasn't feasible due to
missing parameters and overly rigid assumptions. This patch resolves
that and moves us closer to a more maintainable structure.

CHANGES IN `vectorizeAsInsertSliceOp`
-------------------------------------
* Introduces a clear distinction between the destination tensor and the
  vector to store, via named variables like `destType`/`vecToStoreType`,
  `destShape`/`vecToStoreShape`, etc.
* Ensures the correct rank and shape are used for attributes like
  in_bounds. For example, the size of the in_bounds array now matches
  the source vector rank, not the tensor rank.
* Drops the assumption that `vecToStoreRank == destRank` — this doesn't
  hold in many real examples.
*  Deduces mask dimensions from `vecToStoreShape` (vector) instead of
   `destShape` (tensor). (Eventually we should not require
   `inputVecSizesForLeadingDims` at all — mask shape should be inferred.)

NEW HELPER: `isMaskTriviallyFoldable`
-------------------------------------
Adds a utility to detect when masking is unnecessary. This avoids
inserting redundant masks and reduces the burden on canonicalization to
clean them up later.

Example where masking is provably unnecessary:
```mlir
%2 = vector.mask %1 {
  vector.transfer_write %0, %arg1[%c0, %c0, %c0, %c0, %c0, %c0]
    {in_bounds = [true, true, true]}
    : vector<1x2x3xf32>, tensor<9x8x7x1x2x3xf32>
} : vector<1x2x3xi1> -> tensor<9x8x7x1x2x3xf32>
```

Also, without this hook, tests are more complicated and require more
matching.

TEST CHANGES
-----------
This patch primarily affects vectorization of:
  * `tensor.insert_slice`, now refactored to use shared hooks.

`tensor.pad` vectorization patterns, which internally use
`tensor.insert_slice`, are also _effectively_ updated. Note, only
pad-with-patterns.mlir is affected.

Most test updates involve the insertion of masks that were previously
missing — this reflects a correctness fix, not a regression. In all
cases, the added masks are indeed required.

You’ll also notice more repeated constants (`arith.constant 0 : index`),
due to increased use of helper hooks. This will be cleaned up separately
via a constant cache (see #138265 for discussion).

NOTE FOR REVIEWERS
------------------
This is a fairly substantial rewrite. You may find it easier to review
`createWriteOrMaskedWrite` as a new method rather than diffing
line-by-line.

TODOs (future PRs)
------------------
Further alignment of `createWriteOrMaskedWrite` and
`createReadOrMaskedRead`:
  * Move `createWriteOrMaskedWrite` next to `createReadOrMaskedRead` (in
    VectorUtils.cpp)
  * Make `createReadOrMaskedRead` leverage `isMaskTriviallyFoldable`.
  * Extend `isMaskTriviallyFoldable` with value-bounds-analysis. See the
    updated test in transform-vector.mlir for an example that would
    benefit from this.

(* This method will eventually be moved out of Vectorization.cpp, which isn't the right long-term home for it.)


  Commit: ca24a26136e450e71cd57e3a9951bf10147e23f3
      https://github.com/llvm/llvm-project/commit/ca24a26136e450e71cd57e3a9951bf10147e23f3
  Author: Andrzej Warzynski <andrzej.warzynski at arm.com>
  Date:   2025-05-29 (Thu, 29 May 2025)

  Changed paths:
    M mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

  Log Message:
  -----------
  [mlir][linalg] Simplify `createWriteOrMaskedWrite` (NFC)

This patch removes `inputVecSizesForLeadingDims` from the parameter list
of `createWriteOrMaskedWrite`. That argument is unnecessary — vector sizes
can be obtained from the `vecToStore` parameter. Since this doesn't change
behavior or test results, it's marked as NFC.

Additional cleanups:
  * Renamed `vectorToStore` to `vecToStore` for consistency and brevity.
  * Rewrote a conditional at the end of the function to use early exit,
    improving readability:

```cpp
  // BEFORE:
  if (maskingRequried) {
    Value maskForWrite = ...;
    write = maskOperation(write, maskForWrite);
  }
  return write;

  // AFTER
  if (!maskingRequried)
    return write;

  Value maskFroWrite = ...;
  return vector::maskOperation(builder, write, maskForWrite);
```

This change addresses a TODO from #141244.


Compare: https://github.com/llvm/llvm-project/compare/53d435df92a4...ca24a26136e4

To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications