[mlir][arith] Add missing canon pattern trunci(ext*i(x)) -> ext*i(x)

This pattern triggers when only the extension bits are truncated.

Reviewed By: dcaballe

Differential Revision: https://reviews.llvm.org/D149286
This commit is contained in:
Jakub Kuderski 2023-04-27 11:13:46 -04:00
parent 05d0caef60
commit ab85aec1af
4 changed files with 43 additions and 13 deletions

View File

@ -319,6 +319,20 @@ def TruncationMatchesShiftAmount :
CPred<"(getScalarOrElementWidth($0) - getScalarOrElementWidth($1)) == "
"*getIntOrSplatIntValue($2)">]>>;
// trunci(extsi(x)) -> extsi(x), when only the sign-extension bits are truncated
def TruncIExtSIToExtSI :
Pat<(Arith_TruncIOp:$tr (Arith_ExtSIOp:$ext $x)),
(Arith_ExtSIOp $x),
[(ValueWiderThan $ext, $tr),
(ValueWiderThan $tr, $x)]>;
// trunci(extui(x)) -> extui(x), when only the zero-extension bits are truncated
def TruncIExtUIToExtUI :
Pat<(Arith_TruncIOp:$tr (Arith_ExtUIOp:$ext $x)),
(Arith_ExtUIOp $x),
[(ValueWiderThan $ext, $tr),
(ValueWiderThan $tr, $x)]>;
// trunci(shrsi(x, c)) -> trunci(shrui(x, c))
def TruncIShrSIToTrunciShrUI :
Pat<(Arith_TruncIOp:$tr

View File

@ -1290,8 +1290,9 @@ bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
void arith::TruncIOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
MLIRContext *context) {
patterns.add<TruncIShrSIToTrunciShrUI, TruncIShrUIMulIToMulSIExtended,
TruncIShrUIMulIToMulUIExtended>(context);
patterns.add<TruncIExtSIToExtSI, TruncIExtUIToExtUI, TruncIShrSIToTrunciShrUI,
TruncIShrUIMulIToMulSIExtended, TruncIShrUIMulIToMulUIExtended>(
context);
}
LogicalResult arith::TruncIOp::verify() {

View File

@ -629,6 +629,16 @@ func.func @truncExtui2(%arg0: i32) -> i16 {
return %trunci : i16
}
// CHECK-LABEL: @truncExtui3
// CHECK: %[[ARG0:.+]]: i8
// CHECK: %[[CST:.*]] = arith.extui %[[ARG0:.+]] : i8 to i16
// CHECK: return %[[CST:.*]] : i16
func.func @truncExtui3(%arg0: i8) -> i16 {
%extui = arith.extui %arg0 : i8 to i32
%trunci = arith.trunci %extui : i32 to i16
return %trunci : i16
}
// CHECK-LABEL: @truncExtuiVector
// CHECK: %[[ARG0:.+]]: vector<2xi32>
// CHECK: %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : vector<2xi32> to vector<2xi16>
@ -658,6 +668,16 @@ func.func @truncExtsi2(%arg0: i32) -> i16 {
return %trunci : i16
}
// CHECK-LABEL: @truncExtsi3
// CHECK: %[[ARG0:.+]]: i8
// CHECK: %[[CST:.*]] = arith.extsi %[[ARG0:.+]] : i8 to i16
// CHECK: return %[[CST:.*]] : i16
func.func @truncExtsi3(%arg0: i8) -> i16 {
%extsi = arith.extsi %arg0 : i8 to i32
%trunci = arith.trunci %extsi : i32 to i16
return %trunci : i16
}
// CHECK-LABEL: @truncExtsiVector
// CHECK: %[[ARG0:.+]]: vector<2xi32>
// CHECK: %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : vector<2xi32> to vector<2xi16>

View File

@ -1107,14 +1107,10 @@ func.func @fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi1> attributes {}
// -----
// TODO Canonicalize this into:
// arith.extui %arg0 : i1 to i2
// CHECK-LABEL: func @do_not_fold_trunci
// CHECK-LABEL: func @fold_trunci
// CHECK-SAME: (%[[ARG0:[0-9a-z]*]]: i1)
func.func @do_not_fold_trunci(%arg0: i1) -> i2 attributes {} {
// CHECK-NEXT: arith.extui %[[ARG0]] : i1 to i8
// CHECK-NEXT: %[[RES:[0-9a-z]*]] = arith.trunci %{{.*}} : i8 to i2
func.func @fold_trunci(%arg0: i1) -> i2 attributes {} {
// CHECK-NEXT: %[[RES:[0-9a-z]*]] = arith.extui %[[ARG0]] : i1 to i2
// CHECK-NEXT: return %[[RES]] : i2
%0 = arith.extui %arg0 : i1 to i8
%1 = arith.trunci %0 : i8 to i2
@ -1123,11 +1119,10 @@ func.func @do_not_fold_trunci(%arg0: i1) -> i2 attributes {} {
// -----
// CHECK-LABEL: func @do_not_fold_trunci_vector
// CHECK-LABEL: func @fold_trunci_vector
// CHECK-SAME: (%[[ARG0:[0-9a-z]*]]: vector<4xi1>)
func.func @do_not_fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi2> attributes {} {
// CHECK-NEXT: arith.extui %[[ARG0]] : vector<4xi1> to vector<4xi8>
// CHECK-NEXT: %[[RES:[0-9a-z]*]] = arith.trunci %{{.*}} : vector<4xi8> to vector<4xi2>
func.func @fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi2> attributes {} {
// CHECK-NEXT: %[[RES:[0-9a-z]*]] = arith.extui %[[ARG0]] : vector<4xi1> to vector<4xi2>
// CHECK-NEXT: return %[[RES]] : vector<4xi2>
%0 = arith.extui %arg0 : vector<4xi1> to vector<4xi8>
%1 = arith.trunci %0 : vector<4xi8> to vector<4xi2>