[mlir][Arith] Fix folder of CmpIOp to not fail when element type is not integer.

[mravishankar]

Summary

The folder used cast<IntegerType> which would segfault if the type were
a vector type. Handle this case appropriately and avoid failure.

[hanchung]

mlir/test/Dialect/Arith/canonicalize.mlir

1683–1684

delete?

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

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::arith;
 
 //===----------------------------------------------------------------------===//
 // Pattern helpers
 //===----------------------------------------------------------------------===//
 

 static Attribute getBoolAttribute(Type type, MLIRContext *ctx, bool value) {
   auto boolAttr = BoolAttr::get(ctx, value);
   ShapedType shapedType = type.dyn_cast_or_null<ShapedType>();
   if (!shapedType)
     return boolAttr;
   return DenseElementsAttr::get(shapedType, boolAttr);
 }
 
+static Optional<int64_t> getIntegerWidth(Type t) {
+  if (auto intType = t.dyn_cast<IntegerType>()) {
+    return intType.getWidth();
+  }
+  if (auto vectorIntType = t.dyn_cast<VectorType>()) {
+    return vectorIntType.getElementType().cast<IntegerType>().getWidth();
+  }
+  return llvm::None;
+}
+
 OpFoldResult arith::CmpIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "cmpi takes two operands");
 
   // cmpi(pred, x, x)
   if (getLhs() == getRhs()) {
     auto val = applyCmpPredicateToEqualOperands(getPredicate());
     return getBoolAttribute(getType(), getContext(), val);
   }
 
   if (matchPattern(getRhs(), m_Zero())) {
     if (auto extOp = getLhs().getDefiningOp<ExtSIOp>()) {
       // extsi(%x : i1 -> iN) != 0  ->  %x
-      if (extOp.getOperand().getType().cast<IntegerType>().getWidth() == 1 &&
+      Optional<int64_t> integerWidth =
+          getIntegerWidth(extOp.getOperand().getType());
+      if (integerWidth && integerWidth.value() == 1 &&
           getPredicate() == arith::CmpIPredicate::ne)
         return extOp.getOperand();
     }
     if (auto extOp = getLhs().getDefiningOp<ExtUIOp>()) {
       // extui(%x : i1 -> iN) != 0  ->  %x
-      if (extOp.getOperand().getType().cast<IntegerType>().getWidth() == 1 &&
+      Optional<int64_t> integerWidth =
+          getIntegerWidth(extOp.getOperand().getType());
+      if (integerWidth && integerWidth.value() == 1 &&
           getPredicate() == arith::CmpIPredicate::ne)
         return extOp.getOperand();
     }
   }
 
   // Move constant to the right side.
   if (operands[0] && !operands[1]) {
     // Do not use invertPredicate, as it will change eq to ne and vice versa.

mlir/test/Dialect/Arith/canonicalize.mlir

func.func @cmpi_const_right(%arg0: i64)
   %8 = arith.cmpi ult, %c1, %arg0 : i64
   %9 = arith.cmpi ugt, %c1, %arg0 : i64
   return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9
       : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
 
 // -----
 
-// CHECK-LABEL: @cmpOfExtSI
+// CHECK-LABEL: @cmpOfExtSI(
 //  CHECK-NEXT:   return %arg0
 func.func @cmpOfExtSI(%arg0: i1) -> i1 {
   %ext = arith.extsi %arg0 : i1 to i64
   %c0 = arith.constant 0 : i64
   %res = arith.cmpi ne, %ext, %c0 : i64
   return %res : i1
 }
 
-// CHECK-LABEL: @cmpOfExtUI
+// CHECK-LABEL: @cmpOfExtUI(
 //  CHECK-NEXT:   return %arg0
 func.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: @cmpOfExtSIVector(
+//  CHECK-NEXT:   return %arg0
+func.func @cmpOfExtSIVector(%arg0: vector<4xi1>) -> vector<4xi1> {
+  %ext = arith.extsi %arg0 : vector<4xi1> to vector<4xi64>
+  %c0 = arith.constant dense<0> : vector<4xi64>
+  %res = arith.cmpi ne, %ext, %c0 : vector<4xi64>
+  return %res : vector<4xi1>
+}
+
+// CHECK-LABEL: @cmpOfExtUIVector(
+//  CHECK-NEXT:   return %arg0
+func.func @cmpOfExtUIVector(%arg0: vector<4xi1>) -> vector<4xi1> {
+  %ext = arith.extui %arg0 : vector<4xi1> to vector<4xi64>
+  %c0 = arith.constant dense<0> : vector<4xi64>
+  %res = arith.cmpi ne, %ext, %c0 : vector<4xi64>
+  return %res : vector<4xi1>
+}
+
+// -----
+
 // CHECK-LABEL: @extSIOfExtUI
 //       CHECK:   %[[res:.+]] = arith.extui %arg0 : i1 to i64
 //       CHECK:   return %[[res]]
 func.func @extSIOfExtUI(%arg0: i1) -> i64 {
   %ext1 = arith.extui %arg0 : i1 to i8
   %ext2 = arith.extsi %ext1 : i8 to i64
   return %ext2 : i64
 }

 // CHECK-LABEL: @xorxor3(
 //       CHECK-NOT: xori
 //       CHECK:   return %arg0
 func.func @xorxor3(%a : i32, %b : i32) -> i32 {
   %c = arith.xori %b, %a : i32
   %res = arith.xori %b, %c : i32
   return %res : i32
 }
+
+// -----