[Mlir-commits] [mlir] db31da2 - [MLIR][Arith] Add constant folder for left shift

[William S. Moses]

Author: William S. Moses
Date: 2022-03-02T12:00:23-05:00
New Revision: db31da279fb39234efbe8a01fa39e8b678e6798b

URL: https://github.com/llvm/llvm-project/commit/db31da279fb39234efbe8a01fa39e8b678e6798b
DIFF: https://github.com/llvm/llvm-project/commit/db31da279fb39234efbe8a01fa39e8b678e6798b.diff

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

Add constant folder for left shift

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D120661

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td
    mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
    mlir/test/Dialect/Arithmetic/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td b/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td
index ea60a83998ad7..ac8348c890262 100644
--- a/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td
+++ b/mlir/include/mlir/Dialect/Arithmetic/IR/ArithmeticOps.td
@@ -510,6 +510,7 @@ def Arith_ShLIOp : Arith_IntBinaryOp<"shli"> {
     %3 = arith.shli %1, %2 : (i8, i8) -> i8    // %3 is 0b00101000
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp b/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
index c381247e4148a..f7f3da1e49486 100644
--- a/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
+++ b/mlir/lib/Dialect/Arithmetic/IR/ArithmeticOps.cpp
@@ -1850,6 +1850,20 @@ LogicalResult arith::SelectOp::verify() {
   }
   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

diff  --git a/mlir/test/Dialect/Arithmetic/canonicalize.mlir b/mlir/test/Dialect/Arithmetic/canonicalize.mlir
index 4d1a415b82a5c..ac50c4367a380 100644
--- a/mlir/test/Dialect/Arithmetic/canonicalize.mlir
+++ b/mlir/test/Dialect/Arithmetic/canonicalize.mlir
@@ -948,3 +948,35 @@ func @test7(%arg0: i32) -> 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
+}