[mlir][index] Add identity folders for add and sub

[Mogball]

Summary

Depends on D153736

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

static OpFoldResult foldBinaryOpChecked(
   return IntegerAttr::get(IndexType::get(lhs.getContext()), *result64);
 }
 
 //===----------------------------------------------------------------------===//
 // AddOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
-  return foldBinaryOpUnchecked(
+  if (OpFoldResult result = foldBinaryOpUnchecked(
           adaptor.getOperands(),
-      [](const APInt &lhs, const APInt &rhs) { return lhs + rhs; });
+          [](const APInt &lhs, const APInt &rhs) { return lhs + rhs; }))
+    return result;
+
+  if (auto rhs = dyn_cast_or_null<IntegerAttr>(adaptor.getRhs())) {
+    // Fold `add(x, 0) -> x`.
+    if (rhs.getValue().isZero())
+      return getLhs();
+  }
+
+  return {};
 }
 
 //===----------------------------------------------------------------------===//
 // SubOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
-  return foldBinaryOpUnchecked(
+  if (OpFoldResult result = foldBinaryOpUnchecked(
           adaptor.getOperands(),
-      [](const APInt &lhs, const APInt &rhs) { return lhs - rhs; });
+          [](const APInt &lhs, const APInt &rhs) { return lhs - rhs; }))
+    return result;
+
+  if (auto rhs = dyn_cast_or_null<IntegerAttr>(adaptor.getRhs())) {
+    // Fold `sub(x, 0) -> x`.
+    if (rhs.getValue().isZero())
+      return getLhs();
+  }
+
+  return {};
 }
 
 //===----------------------------------------------------------------------===//
 // MulOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
   if (OpFoldResult result = foldBinaryOpUnchecked(
           adaptor.getOperands(),
-          [](const APInt &lhs, const APInt &rhs) { return lhs * rhs; });
-      !result.isNull())
+          [](const APInt &lhs, const APInt &rhs) { return lhs * rhs; }))
     return result;
 
   if (auto rhs = dyn_cast_or_null<IntegerAttr>(adaptor.getRhs())) {
     // Fold `mul(x, 1) -> x`.
     if (rhs.getValue().isOne())
       return getLhs();
     // Fold `mul(x, 0) -> 0`.
     if (rhs.getValue().isZero())

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

 func.func @mul_identity(%arg0: index) -> (index, index) {
   %idx0 = index.constant 0
   %idx1 = index.constant 1
   %0 = index.mul %arg0, %idx0
   %1 = index.mul %arg0, %idx1
   // CHECK: return %idx0, %arg0
   return %0, %1 : index, index
 }
+
+// CHECK-LABEL: @add_identity
+func.func @add_identity(%arg0: index) -> index {
+  %idx0 = index.constant 0
+  %0 = index.add %arg0, %idx0
+  // CHECK-NEXT: return %arg0
+  return %0 : index
+}
+
+// CHECK-LABEL: @sub_identity
+func.func @sub_identity(%arg0: index) -> index {
+  %idx0 = index.constant 0
+  %0 = index.sub %arg0, %idx0
+  // CHECK-NEXT: return %arg0
+  return %0 : index
+}