[mlir][linalg] Fix padding size calculation for Conv2d ops.
authorJerry Wu <cheyuw@google.com>
Mon, 9 May 2022 15:43:40 +0000 (08:43 -0700)
committerHanhan Wang <hanchung@google.com>
Mon, 9 May 2022 15:45:37 +0000 (08:45 -0700)
commitad7c49bef774191304b804e60821c0cdb6a390cb
treefa998a3b6ef914d16d44073ca5f28348a09c4916
parente287d647c61f5fbd054410d6ed9c37d5271f29ef
[mlir][linalg] Fix padding size calculation for Conv2d ops.

This patch fixed the padding size calculation for Conv2d ops when the stride > 1. It contains the changes below:

- Use addBound to add constraint for AffineApplyOp in getUpperBoundForIndex. So the result value can be mapped and retrieved later.

- Fixed the bound from AffineMinOp by adding as a closed bound. Originally the bound was added as an open upper bound, which results in the incorrect bounds when we multiply the values. For example:

```
%0 = affine.min affine_map<()[s0] -> (4, -s0 + 11)>()[iv0]
%1 = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%0]

If we add the affine.min as an open bound, addBound will internally transform it into the close bound "%0 <= 3". The following sliceBounds will derive the bound of %1 as "%1 <= 6" and return the open bound "%1 < 7", while the correct bound should be "%1 <= 8".
```

- In addition to addBound, I also changed sliceBounds to support returning closed upper bound, since for the size computation, we usually care about the closed bounds.

- Change the getUpperBoundForIndex to favor constant bounds when required. The sliceBounds will return a tighter but non-constant bounds, which can't be used for padding. The constantRequired option requires getUpperBoundForIndex to get the constant bounds when possible.

Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D124821
mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h
mlir/include/mlir/Dialect/Linalg/Utils/Utils.h
mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
mlir/lib/Dialect/Linalg/Utils/Utils.cpp
mlir/test/Dialect/Linalg/pad.mlir