[all-commits] [llvm/llvm-project] d0aeb7: [mlir][MemRef] Simplify extract_strided_metadata(e...

[qcolombet via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: d0aeb74e8869e5db23c079b98c5e1f325aeeeefe
      https://github.com/llvm/llvm-project/commit/d0aeb74e8869e5db23c079b98c5e1f325aeeeefe
  Author: Quentin Colombet <quentin.colombet at gmail.com>
  Date:   2022-09-22 (Thu, 22 Sep 2022)

  Changed paths:
    M mlir/lib/Dialect/MemRef/Transforms/SimplifyExtractStridedMetadata.cpp
    M mlir/test/Dialect/MemRef/simplify-extract-strided-metadata.mlir

  Log Message:
  -----------
  [mlir][MemRef] Simplify extract_strided_metadata(expand_shape)

Add a pattern to the pass that simplifies
extract_strided_metadata(other_op(memref)).

The new pattern gets rid of the expand_shape operation while
materializing its effects on the sizes, and the strides of
the base object.

In other words, this simplification replaces:
```
baseBuffer, offset, sizes, strides =
             extract_strided_metadata(expand_shape(memref))
```

With

```
baseBuffer, offset, baseSizes, baseStrides =
    extract_strided_metadata(memref)
sizes#reassIdx =
    baseSizes#reassDim / product(expandShapeSizes#j,
                                 for j in group excluding
                                   reassIdx)
strides#reassIdx =
    baseStrides#reassDim * product(expandShapeSizes#j,
                                   for j in
                                     reassIdx+1..
                                       reassIdx+group.size-1)
```

Where `reassIdx` is a reassociation index for the group at
`reassDim` and `expandShapeSizes#j` is either:
- The constant size at dimension j, derived directly from the
  result type of the expand_shape op, or
- An affine expression: baseSizes#reassDim / product of all
  constant sizes in expandShapeSizes.

Note: baseBuffer and offset are unaffected by the expand_shape
operation.

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