Add index::CmpOp canonicalization.

[weiweichen]

Summary

Add canonicalization pattern for index::CmpOp

[bzcheeseman]

mlir/lib/Dialect/Index/IR/IndexOps.cpp

569

plz remove else after return

582

Could you make this failure reason a little clearer? I'm not sure what this means.

Also - is there a way to check this up at the top and exit early? It'd reduce the indentation of most of the meat of this function.

[Mogball]

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

17

Please move this include to the source file

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

544–546

mlir/lib/Dialect/Index/IR/IndexOps.cpp

549–584

559

Please drop all the mlir:: prefixes

569

582

[weiweichen]

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

17

Running into error that's from the declaration of the canonicalize function in generated header file, need to add the include into all source files including IndexOps.h and there seem to be ordering related with ordering. Would just include in IndexOps.h a bit easier here?

In file included from /Users/weiwei.chen/research/modularml/modular/third-party/llvm-project/mlir/lib/Dialect/Index/IR/InferIntRangeInterfaceImpls.cpp:9:
In file included from /Users/weiwei.chen/research/modularml/modular/third-party/llvm-project/mlir/include/mlir/Dialect/Index/IR/IndexOps.h:36:
/Users/weiwei.chen/research/modularml/modular/.derived/third-party/cmake-debug-build-llvm/tools/mlir/include/mlir/Dialect/Index/IR/IndexOps.h.inc:1115:63: error: no type named 'PatternRewriter' in namespace 'mlir'
  static ::mlir::LogicalResult canonicalize(CmpOp op, ::mlir::PatternRewriter &rewriter);

[Mogball]

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

17

Please add a forward declaration for it

[weiweichen]

• Address some review comments.

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

544–546

done

mlir/lib/Dialect/Index/IR/IndexOps.cpp

549–584

done

569

done

582

yes, did a bunch of code reorganizing to reflect the suggestions here

[weiweichen]

Add forward declaration.

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

17

ah, yes!!

[Mogball]

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

26–33

[weiweichen]

• remove empty lines.

mlir/include/mlir/Dialect/Index/IR/IndexOps.h

 //===- IndexOps.h - Index operation declarations ------------------*- C++-*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_INDEX_IR_INDEXOPS_H
 #define MLIR_DIALECT_INDEX_IR_INDEXOPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Dialect/Index/IR/IndexAttrs.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // Forward Declarations
 //===----------------------------------------------------------------------===//
 
-namespace mlir::index {
+namespace mlir {
+class PatternRewriter;
+namespace index {
 enum class IndexCmpPredicate : uint32_t;
 class IndexCmpPredicateAttr;
-} // namespace mlir::index
+} // namespace index
+} // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // ODS-Generated Declarations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Index/IR/IndexOps.h.inc"
 
 #endif // MLIR_DIALECT_INDEX_IR_INDEXOPS_H

mlir/include/mlir/Dialect/Index/IR/IndexOps.td

let description = [{
     // Not equal comparison.
     %2 = index.cmp ne(%a, %b)
     ```
   }];
 
   let arguments = (ins IndexCmpPredicateAttr:$pred, Index:$lhs, Index:$rhs);
   let results = (outs I1:$result);
   let assemblyFormat = "`` $pred `(` $lhs `,` $rhs `)` attr-dict";
   let hasFolder = 1;
+  let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SizeOfOp
 //===----------------------------------------------------------------------===//
 
 def Index_SizeOfOp : IndexOp<"sizeof"> {
   let summary = "size in bits of the index type";
   let description = [{

mlir/lib/Dialect/Index/IR/IndexOps.cpp

 //===- IndexOps.cpp - Index operation definitions --------------------------==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Index/IR/IndexOps.h"
 #include "mlir/Dialect/Index/IR/IndexAttrs.h"
 #include "mlir/Dialect/Index/IR/IndexDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::index;
 
 //===----------------------------------------------------------------------===//
std::optional<bool> result64 = foldCmpOfMaxOrMin(
         lhsOp, cstA.getValue(), rhs.getValue(), 64, getPred());
     std::optional<bool> result32 =
         foldCmpOfMaxOrMin(lhsOp, cstA.getValue().trunc(32),
                           rhs.getValue().trunc(32), 32, getPred());
     // Fold if the 32-bit and 64-bit results are the same.
     if (result64 && result32 && *result64 == *result32)
       return BoolAttr::get(getContext(), *result64);
   }
 
   return {};
 }
 
+/// Canonicalize
+/// `x - y cmp 0` to `x cmp y`. or `x - y cmp 0` to `x cmp y`.
+/// `0 cmp x - y` to `y cmp x`. or `0 cmp x - y` to `y cmp x`.
+LogicalResult CmpOp::canonicalize(CmpOp op, PatternRewriter &rewriter) {
+  IntegerAttr cmpRhs;
+  IntegerAttr cmpLhs;
+
+  bool rhsIsZero = matchPattern(op.getRhs(), m_Constant(&cmpRhs)) &&
+                   cmpRhs.getValue().isZero();
+  bool lhsIsZero = matchPattern(op.getLhs(), m_Constant(&cmpLhs)) &&
+                   cmpLhs.getValue().isZero();
+  if (!rhsIsZero && !lhsIsZero)
+    return rewriter.notifyMatchFailure(op.getLoc(),
+                                       "cmp is not comparing something with 0");
+  SubOp subOp = rhsIsZero ? op.getLhs().getDefiningOp<index::SubOp>()
+                          : op.getRhs().getDefiningOp<index::SubOp>();
+  if (!subOp)
+    return rewriter.notifyMatchFailure(
+        op.getLoc(), "non-zero operand is not a result of subtraction");
+
+  index::CmpOp newCmp;
+  if (rhsIsZero)
+    newCmp = rewriter.create<index::CmpOp>(op.getLoc(), op.getPred(),
+                                           subOp.getLhs(), subOp.getRhs());
+  else
+    newCmp = rewriter.create<index::CmpOp>(op.getLoc(), op.getPred(),
+                                           subOp.getRhs(), subOp.getLhs());
+  rewriter.replaceOp(op, newCmp);
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
 void ConstantOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   SmallString<32> specialNameBuffer;
   llvm::raw_svector_ostream specialName(specialNameBuffer);
   specialName << "idx" << getValueAttr().getValue();

mlir/test/Dialect/Index/index-canonicalize.mlir

func.func @xor() -> index {
   %rhs = index.constant 1
   // CHECK: %[[A:.*]] = index.constant 4
   %0 = index.xor %lhs, %rhs
   // CHECK: return %[[A]]
   return %0 : index
 }
 
 // CHECK-LABEL: @cmp
-func.func @cmp() -> (i1, i1, i1, i1) {
+func.func @cmp(%arg0: index) -> (i1, i1, i1, i1, i1, i1) {
   %a = index.constant 0
   %b = index.constant -1
   %c = index.constant -2
   %d = index.constant 4
 
   %0 = index.cmp slt(%a, %b)
   %1 = index.cmp ugt(%b, %a)
   %2 = index.cmp ne(%d, %a)
   %3 = index.cmp sgt(%b, %a)
 
+  %4 = index.sub %a, %arg0
+  %5 = index.cmp sgt(%4, %a)
+
+  %6 = index.sub %a, %arg0
+  %7 = index.cmp sgt(%a, %6)
+
   // CHECK-DAG: %[[TRUE:.*]] = index.bool.constant true
   // CHECK-DAG: %[[FALSE:.*]] = index.bool.constant false
+  // CHECK-DAG: [[IDX0:%.*]] = index.constant 0
+  // CHECK-DAG: [[V4:%.*]] = index.cmp sgt([[IDX0]], %arg0)
+  // CHECK-DAG: [[V5:%.*]] = index.cmp sgt(%arg0, [[IDX0]])
   // CHECK: return %[[FALSE]], %[[TRUE]], %[[TRUE]], %[[FALSE]]
-  return %0, %1, %2, %3 : i1, i1, i1, i1
+  return %0, %1, %2, %3, %5, %7 : i1, i1, i1, i1, i1, i1
 }
 
 // CHECK-LABEL: @cmp_nofold
 func.func @cmp_nofold() -> i1 {
   %lhs = index.constant 1
   %rhs = index.constant 0x100000000
   // 32-bit result differs from 64-bit.
   // CHECK: index.cmp slt