[all-commits] [llvm/llvm-project] bb6d5c: [mlir][Transforms] `GreedyPatternRewriteDriver`: D...

[Matthias Springer via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: bb6d5c220004a5d7e466a669324001285a688918
      https://github.com/llvm/llvm-project/commit/bb6d5c220004a5d7e466a669324001285a688918
  Author: Matthias Springer <me at m-sp.org>
  Date:   2024-01-05 (Fri, 05 Jan 2024)

  Changed paths:
    M flang/test/Lower/array-temp.f90
    M mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
    M mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
    M mlir/test/Conversion/VectorToArmSME/vector-to-arm-sme.mlir
    M mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
    M mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
    M mlir/test/Dialect/ArmSME/arith-ops-to-sme.mlir
    M mlir/test/Dialect/LLVMIR/type-consistency.mlir
    M mlir/test/Dialect/Linalg/loops.mlir
    M mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir
    M mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir
    M mlir/test/Dialect/Math/algebraic-simplification.mlir
    M mlir/test/Dialect/Math/expand-math.mlir
    M mlir/test/Dialect/Math/polynomial-approximation.mlir
    M mlir/test/Dialect/Mesh/resharding-spmdization.mlir
    M mlir/test/Dialect/NVGPU/transform-create-async-groups.mlir
    M mlir/test/Dialect/SCF/loop-pipelining.mlir
    M mlir/test/Dialect/SparseTensor/sparse_1d.mlir
    M mlir/test/Dialect/SparseTensor/sparse_affine.mlir
    M mlir/test/Dialect/SparseTensor/sparse_concat.mlir
    M mlir/test/Dialect/SparseTensor/sparse_storage.mlir
    M mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir
    M mlir/test/Dialect/Tosa/constant-op-fold.mlir
    M mlir/test/Dialect/Tosa/tosa-decompose-depthwise.mlir
    M mlir/test/Dialect/Tosa/tosa-decompose-transpose-conv.mlir
    M mlir/test/Dialect/Vector/vector-broadcast-lowering-transforms.mlir
    M mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
    M mlir/test/Dialect/Vector/vector-mask-lowering-transforms.mlir
    M mlir/test/Dialect/Vector/vector-multi-reduction-lowering.mlir
    M mlir/test/Dialect/Vector/vector-scalable-create-mask-lowering.mlir

  Log Message:
  -----------
  [mlir][Transforms] `GreedyPatternRewriteDriver`: Do not CSE constants during iterations (#75897)

The `GreedyPatternRewriteDriver` tries to iteratively fold ops and apply
rewrite patterns to ops. It has special handling for constants: they are
CSE'd and sometimes moved to parent regions to allow for additional
CSE'ing. This happens in `OperationFolder`.

To allow for efficient CSE'ing, `OperationFolder` maintains an internal
lookup data structure to find the existing constant ops with the same
value for each `IsolatedFromAbove` region:
```c++
/// A mapping between an insertion region and the constants that have been
/// created within it.
DenseMap<Region *, ConstantMap> foldScopes;
```

Rewrite patterns are allowed to modify operations. In particular, they
may move operations (including constants) from one region to another
one. Such an IR rewrite can make the above lookup data structure
inconsistent.

We encountered such a bug in a downstream project. This bug materialized
in the form of an op that uses the result of a constant op from a
different `IsolatedFromAbove` region (that is not accessible).

This commit changes the behavior of the `GreedyPatternRewriteDriver`
such that `OperationFolder` is used to CSE constants at the beginning of
each iteration (as the worklist is populated), but no longer during an
iteration. `OperationFolder` is no longer used after populating the
worklist, so we do not have to care about inconsistent state in the
`OperationFolder` due to IR rewrites. The `GreedyPatternRewriteDriver`
now performs the op folding by itself instead of calling
`OperationFolder::tryToFold`.

This change changes the order of constant ops in test cases, but not the
region in which they appear. All broken test cases were fixed by turning
`CHECK` into `CHECK-DAG`.

Alternatives considered: The state of `OperationFolder` could be
partially invalidated with every `notifyOperationModified` notification.
That is more fragile than the solution in this commit because incorrect
rewriter API usage can lead to missing notifications and hard-to-debug
`IsolatedFromAbove` violations. (It did not fix the above mention bug in
a downstream project, which could be due to incorrect rewriter API usage
or due to another conceptual problem that I missed.) Moreover, ops are
frequently getting modified during a greedy pattern rewrite, so we would
likely keep invalidating large parts of the state of `OperationFolder`
over and over.

Migration guide: Turn `CHECK` into `CHECK-DAG` in test cases. Constant
ops are no longer folded during a greedy pattern rewrite. If you rely on
folding (and rematerialization) of constant ops during a greedy pattern
rewrite, turn the folder into a pattern.