[mlir][StandardDialect] Add some folding for operations in standard dialect.

[mravishankar]

Summary

Add the following canonicalization

• and(x, 1) -> x
• subi(x, 0) -> x

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

430

This doesn't make sense to me, did you intend to check for all ones?

mlir/test/Transforms/constant-fold.mlir

98

Can you add a non i1 test case?

[mravishankar]

Fixing and folding bug

[mravishankar]

Minor fixes to comments

[mravishankar]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

430

Yes indeed. Sorry about that, and thanks for catching this. Fixed it and generalized to test for splat tensor/vector constants as well.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

I don't think most of this is necessary, can you just use the existing m_ConstantInt and then check that result(if matched) isAllOnesValue?

[mravishankar]

Addressing comments.

[mravishankar]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

Thanks for the suggestion. Done.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

Why do you need the check for IntegerType here? I wouldn't expect that to be necessary.

445

intValue.isAllOnesValue()

[mravishankar]

Address comments

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

I think this wouldnt be valid to do for IndexTypes since there is no defined bitwidth for index types.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

I don't see the problem though, this seems conservative even in the case of IndexType. We have a maximum representation for an index constant, even if we don't know the exact bitwidth it will end up being. If all of the bits are one in the maximum representation, it is guaranteed that all bits will be one in any other representation. Also, given that we already fold index values in other operations we would have to create a new policy where we don't fold them(which seems a bit excessive IMO).

[mravishankar]

Address comments

[mravishankar]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

442

Ah good point. With use of isAllOnesValue here I dont need the bitwidth explicitly. So I can remove the check for integer. Added test for index type as well to make sure it behaves as expected.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 //===----------------------------------------------------------------------===//
 // AndOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult AndOp::fold(ArrayRef<Attribute> operands) {
   /// and(x, 0) -> 0
   if (matchPattern(rhs(), m_Zero()))
     return rhs();
+  /// and(x, allOnes) -> x
+  auto getIntegerType = [](Value v) -> IntegerType {
+    Type t = v.getType();
+    if (t.isa<IntegerType>())
+      return t.cast<IntegerType>();
+    if (VectorType vType = t.dyn_cast<VectorType>()) {
+      return vType.getElementType().dyn_cast<IntegerType>();
+    }
+    if (TensorType tType = t.dyn_cast<TensorType>()) {
+      return tType.getElementType().dyn_cast<IntegerType>();
+    }
+    return nullptr;
+  };
+  if (IntegerType intType = getIntegerType(this->getResult())) {
+    APInt intValue;
+    if (matchPattern(rhs(), m_ConstantInt(&intValue))) {
+      if (intValue == APInt::getAllOnesValue(intType.getWidth()))
+        return lhs();
+    }
+  }
   /// and(x,x) -> x
   if (lhs() == rhs())
     return rhs();
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a & b; });
 }
 

 //===----------------------------------------------------------------------===//
 // SubIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SubIOp::fold(ArrayRef<Attribute> operands) {
   // subi(x,x) -> 0
   if (getOperand(0) == getOperand(1))
     return Builder(getContext()).getZeroAttr(getType());
+  // subi(x,0) -> x
+  if (matchPattern(rhs(), m_Zero()))
+    return lhs();
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a - b; });
 }
 
 //===----------------------------------------------------------------------===//
 // SubViewOp
 //===----------------------------------------------------------------------===//

mlir/test/Transforms/constant-fold.mlir

func @simple_addi() -> i32 {
   %2 = addi %0, %1 : i32
 
   // CHECK-NEXT: return [[C]]
   return %2 : i32
 }
 
 // -----
 
+// CHECK: func @simple_and
+// CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]]: i1
+// CHECK-SAME: [[ARG1:%[a-zA-Z0-9]+]]: i32)
+func @simple_and(%arg0 : i1, %arg1 : i32) -> (i1, i32) {
+  %c1 = constant 1 : i1
+  %cAllOnes_32 = constant 4294967295 : i32
+
+  // CHECK: [[C31:%.*]] = constant 31 : i32
+  %c31 = constant 31 : i32
+  %1 = and %arg0, %c1 : i1
+  %2 = and %arg1, %cAllOnes_32 : i32
+
+  // CHECK: [[VAL:%.*]] = and [[ARG1]], [[C31]]
+  %3 = and %2, %c31 : i32
+
+  // CHECK: return [[ARG0]], [[VAL]]
+  return %1, %3 : i1, i32
+}
+
+// -----
+
+// CHECK: func @tensor_and
+// CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]]: tensor<2xi32>
+func @tensor_and(%arg0 : tensor<2xi32>) -> tensor<2xi32> {
+  %cAllOnes_32 = constant dense<4294967295> : tensor<2xi32>
+
+  // CHECK: [[C31:%.*]] = constant dense<31> : tensor<2xi32>
+  %c31 = constant dense<31> : tensor<2xi32>
+
+  // CHECK: [[CMIXED:%.*]] = constant dense<[31, -1]> : tensor<2xi32>
+  %c_mixed = constant dense<[31, 4294967295]> : tensor<2xi32>
+
+  %0 = and %arg0, %cAllOnes_32 : tensor<2xi32>
+
+  // CHECK: [[T1:%.*]] = and [[ARG0]], [[C31]]
+  %1 = and %0, %c31 : tensor<2xi32>
+
+  // CHECK: [[T2:%.*]] = and [[T1]], [[CMIXED]]
+  %2 = and %1, %c_mixed : tensor<2xi32>
+
+  // CHECK: return [[T2]]
+  return %2 : tensor<2xi32>
+}
+
+// -----
+
+// CHECK: func @vector_and
+// CHECK-SAME: [[ARG0:%[a-zA-Z0-9]+]]: vector<2xi32>
+func @vector_and(%arg0 : vector<2xi32>) -> vector<2xi32> {
+  %cAllOnes_32 = constant dense<4294967295> : vector<2xi32>
+
+  // CHECK: [[C31:%.*]] = constant dense<31> : vector<2xi32>
+  %c31 = constant dense<31> : vector<2xi32>
+
+  // CHECK: [[CMIXED:%.*]] = constant dense<[31, -1]> : vector<2xi32>
+  %c_mixed = constant dense<[31, 4294967295]> : vector<2xi32>
+
+  %0 = and %arg0, %cAllOnes_32 : vector<2xi32>
+
+  // CHECK: [[T1:%.*]] = and [[ARG0]], [[C31]]
+  %1 = and %0, %c31 : vector<2xi32>
+
+  // CHECK: [[T2:%.*]] = and [[T1]], [[CMIXED]]
+  %2 = and %1, %c_mixed : vector<2xi32>
+
+  // CHECK: return [[T2]]
+  return %2 : vector<2xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func @addi_splat_vector
 func @addi_splat_vector() -> vector<8xi32> {
   %0 = constant dense<1> : vector<8xi32>
   %1 = constant dense<5> : vector<8xi32>
 
   // CHECK-NEXT: [[C:%.+]] = constant dense<6> : vector<8xi32>
   %2 = addi %0, %1 : vector<8xi32>
 
func @subf_splat_vector() -> vector<4xf32> {
   %2 = subf %0, %1 : vector<4xf32>
 
   // CHECK-NEXT: return [[C]]
   return %2 : vector<4xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @simple_subi
-func @simple_subi() -> i32 {
+//      CHECK: func @simple_subi
+// CHECK-SAME:   [[ARG0:%[a-zA-Z0-9]+]]
+func @simple_subi(%arg0 : i32) -> (i32, i32) {
   %0 = constant 4 : i32
   %1 = constant 1 : i32
+  %2 = constant 0 : i32
 
   // CHECK-NEXT:[[C3:%.+]] = constant 3 : i32
-  %2 = subi %0, %1 : i32
+  %3 = subi %0, %1 : i32
+  %4 = subi %arg0, %2 : i32
 
-  // CHECK-NEXT: return [[C3]]
-  return %2 : i32
+  // CHECK-NEXT: return [[C3]], [[ARG0]]
+  return %3, %4 : i32, i32
 }
 
 // -----
 
 // CHECK-LABEL: func @subi_splat_tensor
 func @subi_splat_tensor() -> tensor<4xi32> {
   %0 = constant dense<4> : tensor<4xi32>
   %1 = constant dense<1> : tensor<4xi32>