[MLIR] Create folders for extsi/extui

[wsmoses]

Summary

Create folders/canonicalizers for extsi/extui. Specifically,

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

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

let description = [{
     %3 = arith.constant 2 : i3      // %3 is 0b010
     %4 = arith.extsi %3 : i3 to i6  // %4 is 0b000010
 
     %5 = arith.extsi %0 : vector<2 x i32> to vector<2 x i64>
     ```
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let verifier = [{ return verifyExtOp<IntegerType>(*this); }];
 }
 
 //===----------------------------------------------------------------------===//
 // ExtFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_ExtFOp : Arith_FToFCastOp<"extf"> {

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

 //===----------------------------------------------------------------------===//
 // BitcastOp
 //===----------------------------------------------------------------------===//
 
 // bitcast(bitcast(x)) -> x
 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

 // ExtUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ExtUIOp::fold(ArrayRef<Attribute> operands) {
   if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>())
     return IntegerAttr::get(
         getType(), lhs.getValue().zext(getType().getIntOrFloatBitWidth()));
 
+  if (auto lhs = getIn().getDefiningOp<ExtUIOp>()) {
+    getInMutable().assign(lhs.getIn());
+    return getResult();
+  }
+
   return {};
 }
 
 bool arith::ExtUIOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, IntegerType>(inputs, outputs);
 }
 
 //===----------------------------------------------------------------------===//
 // ExtSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ExtSIOp::fold(ArrayRef<Attribute> operands) {
   if (auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>())
     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
 //===----------------------------------------------------------------------===//
 
 bool arith::ExtFOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   return checkWidthChangeCast<std::greater, FloatType>(inputs, outputs);
 }
 

mlir/test/Dialect/Arithmetic/canonicalize.mlir

func @cmpOfExtUI(%arg0: i1) -> i1 {
   %ext = arith.extui %arg0 : i1 to i64
   %c0 = arith.constant 0 : i64
   %res = arith.cmpi ne, %ext, %c0 : i64
   return %res : 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]]
 func @indexCastOfSignExtend(%arg0: i8) -> index {
   %ext = arith.extsi %arg0 : i8 to i16
   %idx = arith.index_cast %ext : i16 to index
   return %idx : index
 }