[mlir][Arith] Fold integer shift op with zero.

[jacquesguan]

Summary

This revision folds arith.shrui, arith.shrsi and arith.shli with zero
rhs to lhs.

[jacquesguan]

Add missing check.

[ftynse]

mlir/test/Dialect/Arith/canonicalize.mlir

2112

Nit: don't pattern-match SSA value names, they are transient. Capture the name at its definition point instead with %[[NAME:.+]]

[aartbik]

mlir/test/Dialect/Arith/canonicalize.mlir

2112

+1

capturing the arguments is simple as well, e.g.

CHECK-LABEL: func.func @foldShli0(
CHECK-SAME: %[[A:.*]]: i64)
// CHECK: return %[[A]] : i64

[jacquesguan]

Address comment.

mlir/test/Dialect/Arith/canonicalize.mlir

2112

Done, thanks.

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

LogicalResult arith::SelectOp::verify() {
   }
   return success();
 }
 //===----------------------------------------------------------------------===//
 // ShLIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShLIOp::fold(ArrayRef<Attribute> operands) {
+  // shli(x, 0) -> x
+  if (matchPattern(getRhs(), m_Zero()))
+    return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       operands, [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.shl(b);
       });
   return bounded ? result : Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // ShRUIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShRUIOp::fold(ArrayRef<Attribute> operands) {
+  // shrui(x, 0) -> x
+  if (matchPattern(getRhs(), m_Zero()))
+    return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       operands, [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.lshr(b);
       });
   return bounded ? result : Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // ShRSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::ShRSIOp::fold(ArrayRef<Attribute> operands) {
+  // shrsi(x, 0) -> x
+  if (matchPattern(getRhs(), m_Zero()))
+    return getLhs();
   // Don't fold if shifting more than the bit width.
   bool bounded = false;
   auto result = constFoldBinaryOp<IntegerAttr>(
       operands, [&](const APInt &a, const APInt &b) {
         bounded = b.ule(b.getBitWidth());
         return a.ashr(b);
       });
   return bounded ? result : Attribute();

mlir/test/Dialect/Arith/canonicalize.mlir

func.func @wideMulToMulSIExtendedBadShift2(%a: i32, %b: i32) -> i32 {
   %x = arith.extsi %a: i32 to i64
   %y = arith.extsi %b: i32 to i64
   %m = arith.muli %x, %y: i64
   %c31 = arith.constant 31: i64
   %sh = arith.shrui %m, %c31 : i64
   %hi = arith.trunci %sh: i64 to i32
   return %hi : i32
 }
+
+// CHECK-LABEL: @foldShli0
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+//       CHECK:   return %[[ARG]] : i64
+func.func @foldShli0(%x : i64) -> i64 {
+  %c0 = arith.constant 0 : i64
+  %r = arith.shli %x, %c0 : i64
+  return %r : i64
+}
+
+// CHECK-LABEL: @foldShrui0
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+//       CHECK:   return %[[ARG]] : i64
+func.func @foldShrui0(%x : i64) -> i64 {
+  %c0 = arith.constant 0 : i64
+  %r = arith.shrui %x, %c0 : i64
+  return %r : i64
+}
+
+// CHECK-LABEL: @foldShrsi0
+// CHECK-SAME: (%[[ARG:.*]]: i64)
+//       CHECK:   return %[[ARG]] : i64
+func.func @foldShrsi0(%x : i64) -> i64 {
+  %c0 = arith.constant 0 : i64
+  %r = arith.shrsi %x, %c0 : i64
+  return %r : i64
+}