[MLIR] Create folders for extsi/extui

Create folders/canonicalizers for extsi/extui. Specifically,

extui(extui(x)) -> extui(x)
extsi(extsi(x)) -> extsi(x)
extsi(extui(x)) -> extui(x)

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D116515

Author: wsmoses

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

@@ -818,6 +818,7 @@ def Arith_ExtSIOp : Arith_IToICastOp<"extsi"> {
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let verifier = [{ return verifyExtOp<IntegerType>(*this); }];
 }
 

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

@@ -128,4 +128,12 @@ def IndexCastOfExtSI :
 def BitcastOfBitcast :
     Pat<(Arith_BitcastOp (Arith_BitcastOp $x)), (replaceWithValue $x)>;
 
+//===----------------------------------------------------------------------===//
+// ExtSIOp
+//===----------------------------------------------------------------------===//
+
+// extsi(extui(x iN : iM) : iL) -> extui(x : iL)
+def ExtSIOfExtUI :
+    Pat<(Arith_ExtSIOp (Arith_ExtUIOp $x)), (Arith_ExtUIOp $x)>;
+
 #endif // ARITHMETIC_PATTERNS

mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp

@@ -788,6 +788,11 @@ OpFoldResult arith::ExtUIOp::fold(ArrayRef<Attribute> operands) {
     return IntegerAttr::get(
         getType(), lhs.getValue().zext(getType().getIntOrFloatBitWidth()));
 
+  if (auto lhs = getIn().getDefiningOp<ExtUIOp>()) {
+    getInMutable().assign(lhs.getIn());
+    return getResult();
+  }
+
   return {};
 }
 
@@ -804,13 +809,23 @@ OpFoldResult arith::ExtSIOp::fold(ArrayRef<Attribute> operands) {
     return IntegerAttr::get(
         getType(), lhs.getValue().sext(getType().getIntOrFloatBitWidth()));
 
+  if (auto lhs = getIn().getDefiningOp<ExtSIOp>()) {
+    getInMutable().assign(lhs.getIn());
+    return getResult();
+  }
+
   return {};
 }
 
 bool arith::ExtSIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, IntegerType>(inputs, outputs);
 }
 
+void arith::ExtSIOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<ExtSIOfExtUI>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ExtFOp
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Arithmetic/canonicalize.mlir

@@ -70,6 +70,35 @@ func @cmpOfExtUI(%arg0: i1) -> i1 {
 
 // -----
 
+// CHECK-LABEL: @extSIOfExtUI
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 : i1 to i64
+//       CHECK:   return %[[res]]
+func @extSIOfExtUI(%arg0: i1) -> i64 {
+  %ext1 = arith.extui %arg0 : i1 to i8
+  %ext2 = arith.extsi %ext1 : i8 to i64
+  return %ext2 : i64
+}
+
+// CHECK-LABEL: @extUIOfExtUI
+//       CHECK:   %[[res:.+]] = arith.extui %arg0 : i1 to i64
+//       CHECK:   return %[[res]]
+func @extUIOfExtUI(%arg0: i1) -> i64 {
+  %ext1 = arith.extui %arg0 : i1 to i8
+  %ext2 = arith.extui %ext1 : i8 to i64
+  return %ext2 : i64
+}
+
+// CHECK-LABEL: @extSIOfExtSI
+//       CHECK:   %[[res:.+]] = arith.extsi %arg0 : i1 to i64
+//       CHECK:   return %[[res]]
+func @extSIOfExtSI(%arg0: i1) -> i64 {
+  %ext1 = arith.extsi %arg0 : i1 to i8
+  %ext2 = arith.extsi %ext1 : i8 to i64
+  return %ext2 : i64
+}
+
+// -----
+
 // CHECK-LABEL: @indexCastOfSignExtend
 //       CHECK:   %[[res:.+]] = arith.index_cast %arg0 : i8 to index
 //       CHECK:   return %[[res]]