[mlir][arith] Fix CmpIOP folding for vector types.

[akuegel]

Summary

It failed for vector types before, but now it folds them correctly.

[herhut]

Thanks!

[herhut]

Can you change the subject to have an additional tag, e.g., [mlir][arith]. That provides a little more context. You could also state that it did fail for vectors before but now it folds them correctly.

[dcaballe]

LGTM. Just a comment about the test.

mlir/test/Dialect/Arithmetic/canonicalize.mlir

22

Is the issue only with ugt? Would it make sense to keep this test as is and then copy this test and use a vector type for all the predicates instead?

[akuegel]

Address review comment (improve test).

[akuegel]

Simplify test.

[akuegel]

In D116151#3206517, @herhut wrote:

Can you change the subject to have an additional tag, e.g., [mlir][arith]. That provides a little more context. You could also state that it did fail for vectors before but now it folds them correctly.

Done.

mlir/test/Dialect/Arithmetic/canonicalize.mlir

22

You are right, it is better to copy the whole test. Done.

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

static bool applyCmpPredicateToEqualOperands(arith::CmpIPredicate predicate) {
   case arith::CmpIPredicate::sgt:
   case arith::CmpIPredicate::ult:
   case arith::CmpIPredicate::ugt:
     return false;
   }
   llvm_unreachable("unknown cmpi predicate kind");
 }
 
+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);
+}
+
 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 BoolAttr::get(getContext(), val);
+    return getBoolAttribute(getType(), getContext(), val);
   }
 
   auto lhs = operands.front().dyn_cast_or_null<IntegerAttr>();
   auto rhs = operands.back().dyn_cast_or_null<IntegerAttr>();
   if (!lhs || !rhs)
     return {};
 
   auto val = applyCmpPredicate(getPredicate(), lhs.getValue(), rhs.getValue());

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 // RUN: mlir-opt %s -canonicalize --split-input-file | FileCheck %s
 
 // Test case: Folding of comparisons with equal operands.
 // CHECK-LABEL: @cmpi_equal_operands
 //   CHECK-DAG:   %[[T:.*]] = arith.constant true
 //   CHECK-DAG:   %[[F:.*]] = arith.constant false
 //       CHECK:   return %[[T]], %[[T]], %[[T]], %[[T]], %[[T]],
 //  CHECK-SAME:          %[[F]], %[[F]], %[[F]], %[[F]], %[[F]]
-func @cmpi_equal_operands(%arg0: i64)
+func @cmpi_equal_operands(%arg0: i64, %arg1: i64)
     -> (i1, i1, i1, i1, i1, i1, i1, i1, i1, i1) {
   %0 = arith.cmpi eq, %arg0, %arg0 : i64
   %1 = arith.cmpi sle, %arg0, %arg0 : i64
   %2 = arith.cmpi sge, %arg0, %arg0 : i64
   %3 = arith.cmpi ule, %arg0, %arg0 : i64
   %4 = arith.cmpi uge, %arg0, %arg0 : i64
   %5 = arith.cmpi ne, %arg0, %arg0 : i64
   %6 = arith.cmpi slt, %arg0, %arg0 : i64
   %7 = arith.cmpi sgt, %arg0, %arg0 : i64
   %8 = arith.cmpi ult, %arg0, %arg0 : i64
-  %9 = arith.cmpi ugt, %arg0, %arg0 : i64
+  %9 = arith.cmpi ugt, %arg1, %arg1 : i64
   return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9
       : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
 
+// Test case: Folding of comparisons with equal vector operands.
+// CHECK-LABEL: @cmpi_equal_vector_operands
+//   CHECK-DAG:   %[[T:.*]] = arith.constant dense<true>
+//   CHECK-DAG:   %[[F:.*]] = arith.constant dense<false>
+//       CHECK:   return %[[T]], %[[T]], %[[T]], %[[T]], %[[T]],
+//  CHECK-SAME:          %[[F]], %[[F]], %[[F]], %[[F]], %[[F]]
+func @cmpi_equal_vector_operands(%arg0: vector<1x8xi64>, %arg1: vector<1x8xi64>)
+    -> (vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>,
+        vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>,
+  vector<1x8xi1>, vector<1x8xi1>) {
+  %0 = arith.cmpi eq, %arg0, %arg0 : vector<1x8xi64>
+  %1 = arith.cmpi sle, %arg0, %arg0 : vector<1x8xi64>
+  %2 = arith.cmpi sge, %arg0, %arg0 : vector<1x8xi64>
+  %3 = arith.cmpi ule, %arg0, %arg0 : vector<1x8xi64>
+  %4 = arith.cmpi uge, %arg0, %arg0 : vector<1x8xi64>
+  %5 = arith.cmpi ne, %arg0, %arg0 : vector<1x8xi64>
+  %6 = arith.cmpi slt, %arg0, %arg0 : vector<1x8xi64>
+  %7 = arith.cmpi sgt, %arg0, %arg0 : vector<1x8xi64>
+  %8 = arith.cmpi ult, %arg0, %arg0 : vector<1x8xi64>
+  %9 = arith.cmpi ugt, %arg1, %arg1 : vector<1x8xi64>
+  return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9
+      : vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>,
+        vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>, vector<1x8xi1>,
+  vector<1x8xi1>, vector<1x8xi1>
+}
+
 // -----
 
 // 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