[MLIR] Replace a not of a comparison with appropriate comparison

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3021

Please use m_Constant to match constants.

3030

Are all of these mlir:: necessary? If not, please drop them.

mlir/test/Dialect/Standard/canonicalize.mlir

405

Can you add tests for the other cases?

[wsmoses]

Test all integer comparisons

[chelini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3015

I would make this top-level comment more generic, something like Replace a not of a comparison operation, for example: not(cmp eq A, B) => cmp ne A, B.

mlir/test/Dialect/Standard/canonicalize.mlir

651

No newline.

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3015

And also mention that "not" is in practice "xor 1".

3028

The i / 1 - i trick deserves a comment. (Personally, I would have written a function and just called it twice).

3030

I suppose you want to continue here rather than error out. Otherwise you seem to be missing on a valid an optimizable case of xor %value, 1 : i1.

[wsmoses]

Address comments

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3023

Please avoid i and other single letter variable name and find something meaningful instead. Even operand_id would be better.

Also, in general we don't need to have such loop. Instead we should canonicalize commutative operations to have constant always on the same side.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3023

(operandId given we use camelCase variables)

[wsmoses]

Remove for loop

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3016

Nit: trailing dot please.

3024

Nit: let's mention that we expect constants to have been canonicalized to be the second operand.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
     // SIMD vector element-wise bitwise integer xor.
     %f = xor %g, %h : vector<4xi32>
 
     // Tensor element-wise bitwise integer xor.
     %x = xor %y, %z : tensor<4x?xi8>
     ```
   }];
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ZeroExtendIOp
 //===----------------------------------------------------------------------===//
 
 def ZeroExtendIOp : Std_Op<"zexti", [NoSideEffect,
     DeclareOpInterfaceMethods<VectorUnrollOpInterface>] #

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult XOrOp::fold(ArrayRef<Attribute> operands) {
   /// xor(x,x) -> 0
   if (lhs() == rhs())
     return Builder(getContext()).getZeroAttr(getType());
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a ^ b; });
 }
 
+namespace {
+/// Replace a not of a comparison operation, for example: not(cmp eq A, B) =>
+/// cmp ne A, B. Note that a logical not is implemented as xor 1, val
+struct NotICmp : public OpRewritePattern<XOrOp> {
+  using OpRewritePattern<XOrOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(XOrOp op,
+                                PatternRewriter &rewriter) const override {
+
+    APInt constValue;
+    if (!matchPattern(op.getOperand(1), m_ConstantInt(&constValue)))
+      return failure();
+
+    if (constValue != 1)
+      return failure();
+
+    auto prev = op.getOperand(0).getDefiningOp<CmpIOp>();
+    if (!prev)
+      return failure();
+
+    switch (prev.predicate()) {
+    case CmpIPredicate::eq:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::ne, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::ne:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::eq, prev.lhs(),
+                                          prev.rhs());
+      return success();
+
+    case CmpIPredicate::slt:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::sge, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::sle:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::sgt, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::sgt:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::sle, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::sge:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::slt, prev.lhs(),
+                                          prev.rhs());
+      return success();
+
+    case CmpIPredicate::ult:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::uge, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::ule:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::ugt, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::ugt:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::ule, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    case CmpIPredicate::uge:
+      rewriter.replaceOpWithNewOp<CmpIOp>(op, CmpIPredicate::ult, prev.lhs(),
+                                          prev.rhs());
+      return success();
+    }
+    return failure();
+  }
+};
+} // namespace
+
+void XOrOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                        MLIRContext *context) {
+  results.insert<NotICmp>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ZeroExtendIOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ZeroExtendIOp op) {
   auto srcType = getElementTypeOrSelf(op.getOperand().getType());
   auto dstType = getElementTypeOrSelf(op.getType());
 

mlir/test/Dialect/Standard/canonicalize.mlir

 func @select_cmp_ne_select(%arg0: i64, %arg1: i64) -> i64 {
   %0 = cmpi ne, %arg0, %arg1 : i64
   %1 = select %0, %arg0, %arg1 : i64
   return %1 : i64
 }
 
 // -----
 
 // CHECK-LABEL: @indexCastOfSignExtend
 //       CHECK:   %[[res:.+]] = index_cast %arg0 : i8 to index
 //       CHECK:   return %[[res]]
 func @indexCastOfSignExtend(%arg0: i8) -> index {
   %ext = sexti %arg0 : i8 to i16
   %idx = index_cast %ext : i16 to index
   return %idx : index
 }
 

 //       CHECK:   return %[[add]]
 func @tripleSubSub3(%arg0: index) -> index {
   %c17 = constant 17 : index
   %c42 = constant 42 : index
   %add1 = subi %arg0, %c17 : index
   %add2 = subi %add1, %c42 : index
   return %add2 : index
 }
+
+// CHECK-LABEL: @notCmpEQ
+//       CHECK:   %[[cres:.+]] = cmpi ne, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpEQ(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "eq", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpEQ2
+//       CHECK:   %[[cres:.+]] = cmpi ne, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpEQ2(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "eq", %arg0, %arg1 : i8
+  %ncmp = xor %true, %cmp : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpNE
+//       CHECK:   %[[cres:.+]] = cmpi eq, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpNE(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "ne", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpSLT
+//       CHECK:   %[[cres:.+]] = cmpi sge, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpSLT(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "slt", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpSLE
+//       CHECK:   %[[cres:.+]] = cmpi sgt, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpSLE(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "sle", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpSGT
+//       CHECK:   %[[cres:.+]] = cmpi sle, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpSGT(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "sgt", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpSGE
+//       CHECK:   %[[cres:.+]] = cmpi slt, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpSGE(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "sge", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpULT
+//       CHECK:   %[[cres:.+]] = cmpi uge, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpULT(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "ult", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpULE
+//       CHECK:   %[[cres:.+]] = cmpi ugt, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpULE(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "ule", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpUGT
+//       CHECK:   %[[cres:.+]] = cmpi ule, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpUGT(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "ugt", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}
+
+// CHECK-LABEL: @notCmpUGE
+//       CHECK:   %[[cres:.+]] = cmpi ult, %arg0, %arg1 : i8
+//       CHECK:   return %[[cres]]
+func @notCmpUGE(%arg0: i8, %arg1: i8) -> i1 {
+  %true = constant true
+  %cmp = cmpi "uge", %arg0, %arg1 : i8
+  %ncmp = xor %cmp, %true : i1
+  return %ncmp : i1
+}