[MLIR][Shape] Canonicalize subsequent `size_to_index` and `index_to_size`

Eliminate the subsequent applications of `size_to_index` and `index_to_size`.

Differential Revision: https://reviews.llvm.org/D82083
This commit is contained in:
Frederik Gossen 2020-06-25 11:59:19 +00:00
parent 32ea3397be
commit 66e0f66d8f
4 changed files with 22 additions and 0 deletions

View File

@ -249,6 +249,7 @@ def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
let assemblyFormat = "$arg attr-dict";
let hasFolder = 1;
let hasCanonicalizer = 1;
}
def Shape_JoinOp : Shape_Op<"join", []> {

View File

@ -397,6 +397,11 @@ OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
return {};
}
void IndexToSizeOp::getCanonicalizationPatterns(
OwningRewritePatternList &patterns, MLIRContext *context) {
patterns.insert<SizeToIndexToSizeCanonicalization>(context);
}
//===----------------------------------------------------------------------===//
// FromExtentsOp
//===----------------------------------------------------------------------===//

View File

@ -22,3 +22,7 @@ def IndexToSizeToIndexCanonicalization : Pat<
(Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
(replaceWithValue $arg)>;
def SizeToIndexToSizeCanonicalization : Pat<
(Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
(replaceWithValue $arg)>;

View File

@ -503,3 +503,15 @@ func @index_to_size_to_index(%index : index) -> index {
return %result : index
}
// -----
// Canonicalize redundant conversion from `size` to `index` and back.
// CHECK-LABEL: @size_to_index_to_size
// CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size
func @size_to_index_to_size(%size : !shape.size) -> !shape.size {
// CHECK: return %[[SIZE]] : !shape.size
%idx = shape.size_to_index %size
%result = shape.index_to_size %idx
return %result : !shape.size
}