[Mlir-commits] [mlir] 42092c0 - [mlir][index] Add `index.mul` identity folders

[Jeff Niu]

Author: Jeff Niu
Date: 2023-06-26T11:49:13-07:00
New Revision: 42092c071a848ea1864314ae0a740ed7c83d521d

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

LOG: [mlir][index] Add `index.mul` identity folders

Fold `mul(x, 1)` and `mul(x, 0)`.

Depends on D153736

Reviewed By: rriddle

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Index/IR/IndexOps.cpp
    mlir/test/Dialect/Index/index-canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Index/IR/IndexOps.cpp b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
index fb6f891300225..8762ce81d43ca 100644
--- a/mlir/lib/Dialect/Index/IR/IndexOps.cpp
+++ b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
@@ -142,9 +142,22 @@ OpFoldResult SubOp::fold(FoldAdaptor adaptor) {
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MulOp::fold(FoldAdaptor adaptor) {
-  return foldBinaryOpUnchecked(
-      adaptor.getOperands(),
-      [](const APInt &lhs, const APInt &rhs) { return lhs * rhs; });
+  if (OpFoldResult result = foldBinaryOpUnchecked(
+          adaptor.getOperands(),
+          [](const APInt &lhs, const APInt &rhs) { return lhs * rhs; });
+      !result.isNull())
+    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())
+      return rhs;
+  }
+
+  return {};
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Index/index-canonicalize.mlir b/mlir/test/Dialect/Index/index-canonicalize.mlir
index 56a3cb4c6031b..940e6b6812e5e 100644
--- a/mlir/test/Dialect/Index/index-canonicalize.mlir
+++ b/mlir/test/Dialect/Index/index-canonicalize.mlir
@@ -521,3 +521,13 @@ func.func @cmp_maxs(%arg0: index) -> (i1, i1) {
   // CHECK: return %true, %false
   return %1, %2 : i1, i1
 }
+
+// CHECK-LABEL: @mul_identity
+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
+}