commit b05eb92a167f98dd716ffbc51f487857839ef150
author Stephan Herhut <herhut@google.com> Mon Nov 11 02:33:49 2019 -0800
committer A. Unique TensorFlower <gardener@tensorflow.org> Mon Nov 11 02:34:21 2019 -0800
tree d79edecc3b8fd39af46caf33a533e173f3c335a4
parent c902682b0ee21c7b91228df2bfe252ee14bc26ae

Also consider index constants when folding integer arithmetics with constants.

PiperOrigin-RevId: 279698088

include/mlir/IR/Matchers.h

diff --git a/include/mlir/IR/Matchers.h b/include/mlir/IR/Matchers.h
index aba63d6e..1a1869b 100644
--- a/include/mlir/IR/Matchers.h
+++ b/include/mlir/IR/Matchers.h
@@ -98,7 +98,7 @@
       return false;
     auto type = op->getResult(0)->getType();
 
-    if (type.isa<IntegerType>()) {
+    if (type.isIntOrIndex()) {
       return attr_value_binder<IntegerAttr>(bind_value).match(attr);
     }
     if (type.isa<VectorType>() || type.isa<RankedTensorType>()) {

test/Transforms/canonicalize.mlir

diff --git a/test/Transforms/canonicalize.mlir b/test/Transforms/canonicalize.mlir
index 134b967..bd2d7de 100644
--- a/test/Transforms/canonicalize.mlir
+++ b/test/Transforms/canonicalize.mlir
@@ -74,6 +74,15 @@
   return %y: i32
 }
 
+// CHECK-LABEL: func @addi_zero_index
+func @addi_zero_index(%arg0: index) -> index {
+  // CHECK-NEXT: return %arg0
+  %c0_index = constant 0 : index
+  %y = addi %c0_index, %arg0 : index
+  return %y: index
+}
+
+
 // CHECK-LABEL: func @addi_zero_vector
 func @addi_zero_vector(%arg0: vector<4 x i32>) -> vector<4 x i32> {
   // CHECK-NEXT: return %arg0
@@ -101,6 +110,17 @@
   return %y: i32
 }
 
+// CHECK-LABEL: func @muli_zero_index
+func @muli_zero_index(%arg0: index) -> index {
+  // CHECK-NEXT: %[[CST:.*]] = constant 0 : index
+  %c0_index = constant 0 : index
+
+  %y = muli %c0_index, %arg0 : index
+
+  // CHECK-NEXT: return %[[CST]]
+  return %y: index
+}
+
 // CHECK-LABEL: func @muli_zero_vector
 func @muli_zero_vector(%arg0: vector<4 x i32>) -> vector<4 x i32> {
   // CHECK-NEXT: %cst = constant dense<0> : vector<4xi32>
@@ -131,6 +151,14 @@
   return %y: i32
 }
 
+// CHECK-LABEL: func @muli_one_index
+func @muli_one_index(%arg0: index) -> index {
+  // CHECK-NEXT: return %arg0
+  %c0_index = constant 1 : index
+  %y = muli %c0_index, %arg0 : index
+  return %y: index
+}
+
 // CHECK-LABEL: func @muli_one_vector
 func @muli_one_vector(%arg0: vector<4 x i32>) -> vector<4 x i32> {
   // CHECK-NEXT: return %arg0
@@ -177,6 +205,15 @@
   return %1 : i32
 }
 
+//CHECK-LABEL: func @and_zero_index
+func @and_zero_index(%arg0: index) -> index {
+  // CHECK-NEXT: %[[CST:.*]] = constant 0 : index
+  %c0_index = constant 0 : index
+  // CHECK-NEXT: return %[[CST]]
+  %1 = and %arg0, %c0_index : index
+  return %1 : index
+}
+
 //CHECK-LABEL: func @and_zero_vector
 func @and_zero_vector(%arg0: vector<4xi32>) -> vector<4xi32> {
   // CHECK-NEXT: %cst = constant dense<0> : vector<4xi32>
@@ -224,6 +261,14 @@
   return %1 : i32
 }
 
+//CHECK-LABEL: func @or_zero_index
+func @or_zero_index(%arg0: index) -> index {
+  %c0_index = constant 0 : index
+  // CHECK-NEXT: return %arg0
+  %1 = or %arg0, %c0_index : index
+  return %1 : index
+}
+
 //CHECK-LABEL: func @or_zero_vector
 func @or_zero_vector(%arg0: vector<4xi32>) -> vector<4xi32> {
   // CHECK-NEXT: return %arg0