[all-commits] [llvm/llvm-project] a5c802: [mlir] fold more eagerly in structured op splitting

[ftynse via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: a5c802a429e2746c3d5190b2f3ed781911c62ed8
      https://github.com/llvm/llvm-project/commit/a5c802a429e2746c3d5190b2f3ed781911c62ed8
  Author: Alex Zinenko <zinenko at google.com>
  Date:   2022-07-12 (Tue, 12 Jul 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
    M mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    M mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
    M mlir/lib/Dialect/Linalg/Transforms/Split.cpp
    M mlir/test/Dialect/Linalg/multisize-tiling-full.mlir
    M mlir/test/Dialect/Linalg/transform-op-split.mlir

  Log Message:
  -----------
  [mlir] fold more eagerly in structured op splitting

Existing implementation of structured op splitting creates several
affine.apply and affine.min operations in its subshape computation.
As these shapes are further used in data slice extraction, this may lead
to slice shapes being dynamic even when the original shapes and the
splitting point are static. This is particularly visible when splitting
is combined with further subsetting transformations such as tiling. Use
composition and folding more aggressively in splitting to avoid this.

In particular, introduce a `createComposedAffineMin` function that the
affine map used in "min" with the maps used by any `affine.apply` that
may be feeding the operands to the "min". This enables production of
more static shapes. Also introduce a `createComposedFoldedAffineApply`
function that combines the existing `createComposedAffineApply` with
in-place folding to propagate constants produced by zero-input affine
maps. Using these when splitting allows the subsequent canonicalizer
pass to recover static shapes for structured ops.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D129379