Skip to content
Commit d0aeb74e authored by Quentin Colombet's avatar Quentin Colombet
Browse files

[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
parent 9dc0b167
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment