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

[kuhar]

Summary

This pattern triggers when only the extension bits are truncated.

mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td

Constraint<And<[
       CPred<"getScalarOrElementWidth($1) > 0">]>>;
 
 def TruncationMatchesShiftAmount :
     Constraint<And<[
       CPred<"succeeded(getIntOrSplatIntValue($2))">,
       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
           (Arith_ShRSIOp $x, (ConstantLikeMatcher TypedAttrInterface:$c0))),
         (Arith_TruncIOp (Arith_ShRUIOp $x, (Arith_ConstantOp (cast<"TypedAttr"> $c0)))),
         [(TruncationMatchesShiftAmount $x, $tr, $c0)]>;
 
 // trunci(shrui(mul(sext(x), sext(y)), c)) -> mulsi_extended(x, y)

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

 }
 
 bool arith::TruncIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::less, IntegerType>(inputs, 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() {
   return verifyTruncateOp<IntegerType>(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // TruncFOp

mlir/test/Dialect/Arith/canonicalize.mlir

 //       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : i32 to i16
 //       CHECK:   return  %[[CST:.*]]
 func.func @truncExtui2(%arg0: i32) -> i16 {
   %extui = arith.extui %arg0 : i32 to i64
   %trunci = arith.trunci %extui : i64 to 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>
 //       CHECK:   return  %[[CST:.*]]
 func.func @truncExtuiVector(%arg0: vector<2xi32>) -> vector<2xi16> {
   %extsi = arith.extui %arg0 : vector<2xi32> to vector<2xi64>
   %trunci = arith.trunci %extsi : vector<2xi64> to vector<2xi16>
   return %trunci : vector<2xi16>

 //       CHECK:  %[[CST:.*]] = arith.trunci %[[ARG0:.+]] : i32 to i16
 //       CHECK:   return  %[[CST:.*]]
 func.func @truncExtsi2(%arg0: i32) -> i16 {
   %extsi = arith.extsi %arg0 : i32 to i64
   %trunci = arith.trunci %extsi : i64 to 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>
 //       CHECK:   return  %[[CST:.*]]
 func.func @truncExtsiVector(%arg0: vector<2xi32>) -> vector<2xi16> {
   %extsi = arith.extsi %arg0 : vector<2xi32> to vector<2xi64>
   %trunci = arith.trunci %extsi : vector<2xi64> to vector<2xi16>
   return %trunci : vector<2xi16>

mlir/test/Transforms/canonicalize.mlir

func.func @fold_trunci_vector(%arg0: vector<4xi1>) -> vector<4xi1> attributes {} {
   // CHECK-NEXT: return %[[ARG0]] : vector<4xi1>
   %0 = arith.extui %arg0 : vector<4xi1> to vector<4xi8>
   %1 = arith.trunci %0 : vector<4xi8> to vector<4xi1>
   return %1 : vector<4xi1>
 }
 
 // -----
 
-// 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
   return %1 : i2
 }
 
 // -----
 
-// 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>
   return %1 : vector<4xi2>
 }
 
 // -----