[MLIR][Arith] Add constant folder for left shift

[wsmoses]

Summary

Add constant folder for left shift

[rriddle]

The test failures look real.

[wsmoses]

Fix test

[rriddle]

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

1859

We should likely check here that the shift is valid before trying to fold.

[wsmoses]

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

1859

Valid in what sense? Doesn't ShLIOp have the semantic of other left shifts that top bits shifted off are discarded?

[wsmoses]

Fix test

[rriddle]

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

1859

I assume (though it isn't documented) that the behavior is similar to LLVM in that shifting by amounts larger than the integer width is undefined: https://llvm.org/docs/LangRef.html#id150

There are also questions about wrapping behavior, but we don't model those things at this point.

[wsmoses]

Only handle shifts within width

mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td

let description = [{
     Example:
 
     ```mlir
     %1 = arith.constant 5 : i8                 // %1 is 0b00000101
     %2 = arith.constant 3 : i8
     %3 = arith.shli %1, %2 : (i8, i8) -> i8    // %3 is 0b00101000
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ShRUIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_ShRUIOp : Arith_IntBinaryOp<"shrui"> {
   let summary = "unsigned integer right-shift";

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

LogicalResult arith::SelectOp::verify() {
   if (conditionType != shapedConditionType) {
     return emitOpError() << "expected condition type to have the same shape "
                             "as the result type, expected "
                          << shapedConditionType << ", but got "
                          << conditionType;
   }
   return success();
 }
+//===----------------------------------------------------------------------===//
+// ShLIOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult arith::ShLIOp::fold(ArrayRef<Attribute> operands) {
+  // Don't fold if shifting more than the bit width.
+  bool bounded = false;
+  auto result =
+      constFoldBinaryOp<IntegerAttr>(operands, [&](APInt a, const APInt &b) {
+        bounded = b.ule(b.getBitWidth());
+        return std::move(a).shl(b);
+      });
+  return bounded ? result : Attribute();
+}
 
 //===----------------------------------------------------------------------===//
 // Atomic Enum
 //===----------------------------------------------------------------------===//
 
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
 Attribute mlir::arith::getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
                                             OpBuilder &builder, Location loc) {

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 func @test7(%arg0: i32) -> i1 {
   %cst = arith.constant 3.200000e+00 : f64
   %1 = arith.uitofp %arg0: i32 to f64
   %2 = arith.cmpf oge, %1, %cst : f64
   return %2 : i1
   // CHECK: %[[c3:.+]] = arith.constant 3 : i32
   // CHECK: arith.cmpi ugt, %[[arg0]], %[[c3]] : i32
 }
+
+// -----
+
+// CHECK-LABEL: @foldShl(
+// CHECK: %[[res:.+]] = arith.constant 4294967296 : i64
+// CHECK: return %[[res]]
+func @foldShl() -> i64 {
+  %c1 = arith.constant 1 : i64
+  %c32 = arith.constant 32 : i64
+  %r = arith.shli %c1, %c32 : i64
+  return %r : i64
+}
+
+// CHECK-LABEL: @nofoldShl(
+// CHECK: %[[res:.+]] = arith.shli
+// CHECK: return %[[res]]
+func @nofoldShl() -> i64 {
+  %c1 = arith.constant 1 : i64
+  %c132 = arith.constant 132 : i64
+  %r = arith.shli %c1, %c132 : i64
+  return %r : i64
+}
+
+// CHECK-LABEL: @nofoldShl2(
+// CHECK: %[[res:.+]] = arith.shli
+// CHECK: return %[[res]]
+func @nofoldShl2() -> i64 {
+  %c1 = arith.constant 1 : i64
+  %cm32 = arith.constant -32 : i64
+  %r = arith.shli %c1, %cm32 : i64
+  return %r : i64
+}