[MLIR][Arith] Canonicalize cmpi of extui/extsi

Canonicalize cmpi(eq, ext a, ext b) and cmpi(ne, ext a, ext b)

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D120620
This commit is contained in:
William S. Moses 2022-02-26 14:40:51 -05:00
parent 17d71347b2
commit 2af81c6978
4 changed files with 70 additions and 0 deletions

View File

@ -1104,6 +1104,7 @@ def Arith_CmpIOp : Arith_CompareOpOfAnyRank<"cmpi"> {
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
}
//===----------------------------------------------------------------------===//

View File

@ -106,6 +106,28 @@ def XOrINotCmpI :
(Arith_ConstantOp ConstantAttr<I1Attr, "1">)),
(Arith_CmpIOp (InvertPredicate $pred), $a, $b)>;
//===----------------------------------------------------------------------===//
// CmpIOp
//===----------------------------------------------------------------------===//
// cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
def CmpIExtSI :
Pat<(Arith_CmpIOp $pred,
(Arith_ExtSIOp $a),
(Arith_ExtSIOp $b)),
(Arith_CmpIOp $pred, $a, $b),
[(Constraint<CPred<"$0.getType() == $1.getType()">> $a, $b),
(Constraint<CPred<"$0.getValue() == arith::CmpIPredicate::eq || $0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
// cmpi(== or !=, a ext iNN, b ext iNN) == cmpi(== or !=, a, b)
def CmpIExtUI :
Pat<(Arith_CmpIOp $pred,
(Arith_ExtUIOp $a),
(Arith_ExtUIOp $b)),
(Arith_CmpIOp $pred, $a, $b),
[(Constraint<CPred<"$0.getType() == $1.getType()">> $a, $b),
(Constraint<CPred<"$0.getValue() == arith::CmpIPredicate::eq || $0.getValue() == arith::CmpIPredicate::ne">> $pred)]>;
//===----------------------------------------------------------------------===//
// IndexCastOp
//===----------------------------------------------------------------------===//

View File

@ -1322,6 +1322,11 @@ OpFoldResult arith::CmpIOp::fold(ArrayRef<Attribute> operands) {
return BoolAttr::get(getContext(), val);
}
void arith::CmpIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
MLIRContext *context) {
patterns.insert<CmpIExtSI, CmpIExtUI>(context);
}
//===----------------------------------------------------------------------===//
// CmpFOp
//===----------------------------------------------------------------------===//

View File

@ -176,6 +176,48 @@ func @extSIOfExtSI(%arg0: i1) -> i64 {
// -----
// CHECK-LABEL: @cmpIExtSINE
// CHECK: %[[comb:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
// CHECK: return %[[comb]]
func @cmpIExtSINE(%arg0: i8, %arg1: i8) -> i1 {
%ext0 = arith.extsi %arg0 : i8 to i64
%ext1 = arith.extsi %arg1 : i8 to i64
%res = arith.cmpi ne, %ext0, %ext1 : i64
return %res : i1
}
// CHECK-LABEL: @cmpIExtSIEQ
// CHECK: %[[comb:.+]] = arith.cmpi eq, %arg0, %arg1 : i8
// CHECK: return %[[comb]]
func @cmpIExtSIEQ(%arg0: i8, %arg1: i8) -> i1 {
%ext0 = arith.extsi %arg0 : i8 to i64
%ext1 = arith.extsi %arg1 : i8 to i64
%res = arith.cmpi eq, %ext0, %ext1 : i64
return %res : i1
}
// CHECK-LABEL: @cmpIExtUINE
// CHECK: %[[comb:.+]] = arith.cmpi ne, %arg0, %arg1 : i8
// CHECK: return %[[comb]]
func @cmpIExtUINE(%arg0: i8, %arg1: i8) -> i1 {
%ext0 = arith.extui %arg0 : i8 to i64
%ext1 = arith.extui %arg1 : i8 to i64
%res = arith.cmpi ne, %ext0, %ext1 : i64
return %res : i1
}
// CHECK-LABEL: @cmpIExtUIEQ
// CHECK: %[[comb:.+]] = arith.cmpi eq, %arg0, %arg1 : i8
// CHECK: return %[[comb]]
func @cmpIExtUIEQ(%arg0: i8, %arg1: i8) -> i1 {
%ext0 = arith.extui %arg0 : i8 to i64
%ext1 = arith.extui %arg1 : i8 to i64
%res = arith.cmpi eq, %ext0, %ext1 : i64
return %res : i1
}
// -----
// CHECK-LABEL: @andOfExtSI
// CHECK: %[[comb:.+]] = arith.andi %arg0, %arg1 : i8
// CHECK: %[[ext:.+]] = arith.extsi %[[comb]] : i8 to i64