[MLIR][arith] More float op folders

[csigg]

Summary

Fold arith.fadd %x, -0.0 -> %x and similarly for fsub, fmul, fdiv.

Fold arith.fmin %x, %x -> %x, arith.fmin %x, +inf -> %x and similarly for fmax.

[bondhugula]

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

580–585

It'll be useful to as well add an m_ZeroFloat along the lines of the int m_Zero.

581–582

I think this won't work for non f32 Values. You need to create an APFloat zero with the right float semantics. All your test cases are with f32 -- could you add/try one with f16 or f64?

[csigg]

Adding float value matchers and using the same pattern for int value matchers for consistency.

[csigg]

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

581–582

It did work for non f32 values. The initialization value didn't matter, but I admit initializing with 0.0f was confusing. I'm using a different c'tor now that produces a bogus value. Plus I changed some tests to use f16 and f64.

[mehdi_amini]

@jpienaar mentioned to me that `add %x

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

589

Actually I think that %x + 0.0 can't be folded to %x, but %x + -0.0 can!

That's because -0. + 0. == +0.0.

Seems very relevant to https://llvm.discourse.group/t/rfc-fastmath-flags-support-in-mlir-arith-dialect/6049/2

[mehdi_amini]

Can you add context and motivation in the description also? "Remove 'Commutative' interface from fmin/fmax " deserve to be explained, I would even say that it deserve a separate revision from adding folders on simple arithmetic.

[csigg]

Addressing reviewer comments:

• don't fold addf(x, +0) and subf(x, -0)
• split out Commutative tag change

[csigg]

In D118244#3273453, @mehdi_amini wrote:

I would even say that it deserve a separate revision from adding folders on simple arithmetic.

Split out into https://reviews.llvm.org/D118318.

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

589

Thanks a lot Mehdi for catching this! I changed it to only fold x + -0, -0 + x and x - +0.

[csigg]

Rebase

[csigg]

Rebase.

[csigg]

Rebase.

[csigg]

Also fold minf(%x, +inf) -> %x, fix test.

[csigg]

Simplify 1.0 matcher, add test for minf(inf, %x) -> %x.

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

def Arith_SubFOp : Arith_FloatBinaryOp<"subf"> {
   let summary = "floating point subtraction operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MaxFOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MaxFOp : Arith_FloatBinaryOp<"maxf", [Commutative]> {
+def Arith_MaxFOp : Arith_FloatBinaryOp<"maxf"> {
   let summary = "floating-point maximum operation";
   let description = [{
     Syntax:
 
     ```
     operation ::= ssa-id `=` `arith.maxf` ssa-use `,` ssa-use `:` type
     ```
 
     Returns the maximum of the two arguments, treating -0.0 as less than +0.0.
     If one of the arguments is NaN, then the result is also NaN.
 
     Example:
 
     ```mlir
     // Scalar floating-point maximum.
     %a = arith.maxf %b, %c : f64
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MaxSIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_MaxSIOp : Arith_IntBinaryOp<"maxsi"> {
   let summary = "signed integer maximum operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MaxUIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_MaxUIOp : Arith_IntBinaryOp<"maxui"> {
   let summary = "unsigned integer maximum operation";
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MinFOp
 //===----------------------------------------------------------------------===//
 
-def Arith_MinFOp : Arith_FloatBinaryOp<"minf", [Commutative]> {
+def Arith_MinFOp : Arith_FloatBinaryOp<"minf"> {
   let summary = "floating-point minimum operation";
   let description = [{
     Syntax:
 
     ```
     operation ::= ssa-id `=` `arith.minf` ssa-use `,` ssa-use `:` type
     ```
 
     Returns the minimum of the two arguments, treating -0.0 as less than +0.0.
     If one of the arguments is NaN, then the result is also NaN.
 
     Example:
 
     ```mlir
     // Scalar floating-point minimum.
     %a = arith.minf %b, %c : f64
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MinSIOp
 //===----------------------------------------------------------------------===//
 
 def Arith_MinSIOp : Arith_IntBinaryOp<"minsi"> {
   let summary = "signed integer minimum operation";

mlir/include/mlir/IR/Matchers.h

if (auto attr = foldedOp.front().get<Attribute>().dyn_cast<AttrT>()) {
         *bind_value = attr;
       return true;
     }
     return false;
   }
 };
 
 /// The matcher that matches a constant scalar / vector splat / tensor splat
+/// float operation and binds the constant float value.
+struct constant_float_op_binder {
+  FloatAttr::ValueType *bind_value;
+
+  /// Creates a matcher instance that binds the value to bv if match succeeds.
+  constant_float_op_binder(FloatAttr::ValueType *bv) : bind_value(bv) {}
+
+  bool match(Operation *op) {
+    Attribute attr;
+    if (!constant_op_binder<Attribute>(&attr).match(op))
+      return false;
+    auto type = op->getResult(0).getType();
+
+    if (type.isa<FloatType>())
+      return attr_value_binder<FloatAttr>(bind_value).match(attr);
+    if (type.isa<VectorType, RankedTensorType>()) {
+      if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
+        return attr_value_binder<FloatAttr>(bind_value)
+            .match(splatAttr.getSplatValue<Attribute>());
+      }
+    }
+    return false;
+  }
+};
+
+/// The matcher that matches a constant scalar / vector splat / tensor splat
 /// integer operation and binds the constant integer value.
 struct constant_int_op_binder {
   IntegerAttr::ValueType *bind_value;
 
   /// Creates a matcher instance that binds the value to bv if match succeeds.
   constant_int_op_binder(IntegerAttr::ValueType *bv) : bind_value(bv) {}
 
   bool match(Operation *op) {

 }
 
 /// Entry point for matching a pattern over an Operation.
 template <typename Pattern>
 inline bool matchPattern(Operation *op, const Pattern &pattern) {
   return const_cast<Pattern &>(pattern).match(op);
 }
 
+/// Matches a constant holding a scalar/vector/tensor float (splat) and
+/// writes the float value to bind_value.
+inline detail::constant_float_op_binder
+m_ConstantFloat(FloatAttr::ValueType *bind_value) {
+  return detail::constant_float_op_binder(bind_value);
+}
+
 /// Matches a constant holding a scalar/vector/tensor integer (splat) and
 /// writes the integer value to bind_value.
 inline detail::constant_int_op_binder
 m_ConstantInt(IntegerAttr::ValueType *bind_value) {
   return detail::constant_int_op_binder(bind_value);
 }
 
 template <typename OpType, typename... Matchers>

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

 // AddIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::AddIOp::fold(ArrayRef<Attribute> operands) {
   // addi(x, 0) -> x
   if (matchPattern(getRhs(), m_Zero()))
     return getLhs();
 
-  // add(sub(a, b), b) -> a
+  // addi(sub(a, b), b) -> a
   if (auto sub = getLhs().getDefiningOp<SubIOp>())
     if (getRhs() == sub.getRhs())
       return sub.getLhs();
 
-  // add(b, sub(a, b)) -> a
+  // addi(b, sub(a, b)) -> a
   if (auto sub = getRhs().getDefiningOp<SubIOp>())
     if (getLhs() == sub.getRhs())
       return sub.getLhs();
 
   return constFoldBinaryOp<IntegerAttr>(
       operands, [](APInt a, const APInt &b) { return std::move(a) + b; });
 }
 
void arith::XOrIOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<XOrINotCmpI>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
+// Returns whether 'value' is a scalar/vector/tensor with zero (negative or
+// positive) float values.
+static bool isAnyZeroFloat(Value value) {
+  APFloat floatValue(0.0f);
+  return matchPattern(value, m_ConstantFloat(&floatValue)) &&
+         floatValue.isZero();
+}
+
 OpFoldResult arith::AddFOp::fold(ArrayRef<Attribute> operands) {
+  APFloat floatValue(0.0f);
+  // addf(x, 0) -> x
+  if (isAnyZeroFloat(getRhs()))
+    return getLhs();
+
+  // addf(0, x) -> x
+  if (isAnyZeroFloat(getLhs()))
+    return getRhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a + b; });
 }
 
 //===----------------------------------------------------------------------===//
 // SubFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::SubFOp::fold(ArrayRef<Attribute> operands) {
+  APFloat floatValue(0.0f);
+  // subf(x, 0) -> x
+  if (isAnyZeroFloat(getRhs()))
+    return getLhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a - b; });
 }
 
 //===----------------------------------------------------------------------===//
+// MaxFOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult arith::MaxFOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "maxf takes two operands");
+
+  // maxf(x,x) -> x
+  if (getLhs() == getRhs())
+    return getRhs();
+
+  // maxf(c,x) -> maxf(x,c)
+  if (operands.front() && !operands.back()) {
+    std::swap(getOperation()->getOpOperand(0), getOperation()->getOpOperand(1));
+    return getResult();
+  }
+
+  return constFoldBinaryOp<FloatAttr>(
+      operands,
+      [](const APFloat &a, const APFloat &b) { return llvm::maximum(a, b); });
+}
+
+//===----------------------------------------------------------------------===//
 // MaxSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MaxSIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // maxsi(x,x) -> x
   if (getLhs() == getRhs())
OpFoldResult MaxUIOp::fold(ArrayRef<Attribute> operands) {
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](const APInt &a, const APInt &b) {
                                           return llvm::APIntOps::umax(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
+// MinFOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult arith::MinFOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "minf takes two operands");
+
+  // minf(x,x) -> x
+  if (getLhs() == getRhs())
+    return getRhs();
+
+  // minf(c,x) -> minf(x,c)
+  if (operands.front() && !operands.back()) {
+    std::swap(getOperation()->getOpOperand(0), getOperation()->getOpOperand(1));
+    return getResult();
+  }
+
+  return constFoldBinaryOp<FloatAttr>(
+      operands,
+      [](const APFloat &a, const APFloat &b) { return llvm::minimum(a, b); });
+}
+
+//===----------------------------------------------------------------------===//
 // MinSIOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult MinSIOp::fold(ArrayRef<Attribute> operands) {
   assert(operands.size() == 2 && "binary operation takes two operands");
 
   // minsi(x,x) -> x
   if (getLhs() == getRhs())
return constFoldBinaryOp<IntegerAttr>(operands,
                                           return llvm::APIntOps::umin(a, b);
                                         });
 }
 
 //===----------------------------------------------------------------------===//
 // MulFOp
 //===----------------------------------------------------------------------===//
 
+// Returns whether 'value' is a scalar/vector/tensor with one float values.
+static bool isOneFloat(Value value) {
+  APFloat floatValue(0.0f), inverseValue(0.0f);
+  return matchPattern(value, m_ConstantFloat(&floatValue)) &&
+         floatValue.getExactInverse(&inverseValue) &&
+         floatValue == inverseValue;
+}
+
 OpFoldResult arith::MulFOp::fold(ArrayRef<Attribute> operands) {
+  APFloat floatValue(0.0f), inverseValue(0.0f);
+  // mulf(x, 1) -> x
+  if (isOneFloat(getRhs()))
+    return getLhs();
+
+  // mulf(1, x) -> x
+  if (isOneFloat(getLhs()))
+    return getRhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a * b; });
 }
 
 //===----------------------------------------------------------------------===//
 // DivFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivFOp::fold(ArrayRef<Attribute> operands) {
+  APFloat floatValue(0.0f), inverseValue(0.0f);
+  // divf(x, 1) -> x
+  if (isOneFloat(getRhs()))
+    return getLhs();
+
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a / b; });
 }
 
 //===----------------------------------------------------------------------===//
 // Utility functions for verifying cast ops
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Arithmetic/canonicalize.mlir

func @test_minui(%arg0 : i8) -> (i8, i8, i8, i8) {
   %1 = arith.minui %arg0, %maxIntCst : i8
   %2 = arith.minui %arg0, %minIntCst : i8
   %3 = arith.minui %arg0, %c0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // -----
 
+// CHECK-LABEL: @test_minf(
+func @test_minf(%arg0 : f32) -> (f32, f32) {
+  // CHECK-NEXT:  %[[C0:.+]] = arith.constant
+  // CHECK-NEXT:  %[[X:.+]] = arith.minf %arg0, %[[C0]]
+  // CHECK-NEXT:  return %[[X]], %arg0
+  %c0 = arith.constant 0.0 : f32
+  %0 = arith.minf %c0, %arg0 : f32
+  %1 = arith.minf %arg0, %arg0 : f32
+  return %0, %1 : f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_maxf(
+func @test_maxf(%arg0 : f32) -> (f32, f32) {
+  // CHECK-NEXT:  %[[C0:.+]] = arith.constant
+  // CHECK-NEXT:  %[[X:.+]] = arith.maxf %arg0, %[[C0]]
+  // CHECK-NEXT:   return %[[X]], %arg0
+  %c0 = arith.constant 0.0 : f32
+  %0 = arith.maxf %c0, %arg0 : f32
+  %1 = arith.maxf %arg0, %arg0 : f32
+  return %0, %1 : f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_addf(
+func @test_addf(%arg0 : f32) -> (f32, f32, f32) {
+  // CHECK-NEXT:  %[[C2:.+]] = arith.constant 2.0
+  // CHECK-NEXT:   return %arg0, %arg0, %[[C2]]
+  %c0 = arith.constant 0.0 : f32
+  %c1 = arith.constant 1.0 : f32
+  %0 = arith.addf %arg0, %c0 : f32
+  %1 = arith.addf %c0, %arg0 : f32
+  %2 = arith.addf %c1, %c1 : f32
+  return %0, %1, %2 : f32, f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_subf(
+func @test_subf(%arg0 : f32) -> (f32, f32) {
+  // CHECK-NEXT:  %[[C1:.+]] = arith.constant -1.0
+  // CHECK-NEXT:   return %arg0, %[[C1]]
+  %c0 = arith.constant 0.0 : f32
+  %c1 = arith.constant 1.0 : f32
+  %0 = arith.subf %arg0, %c0 : f32
+  %1 = arith.subf %c0, %c1 : f32
+  return %0, %1 : f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_mulf(
+func @test_mulf(%arg0 : f32) -> (f32, f32, f32) {
+  // CHECK-NEXT:  %[[C4:.+]] = arith.constant 4.0
+  // CHECK-NEXT:   return %arg0, %arg0, %[[C4]]
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %0 = arith.mulf %arg0, %c1 : f32
+  %1 = arith.mulf %c1, %arg0 : f32
+  %2 = arith.mulf %c2, %c2 : f32
+  return %0, %1, %2 : f32, f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_divf(
+func @test_divf(%arg0 : f32) -> (f32, f32) {
+  // CHECK-NEXT:  %[[C5:.+]] = arith.constant 5.000000e-01
+  // CHECK-NEXT:   return %arg0, %[[C5]]
+  %c1 = arith.constant 1.0 : f32
+  %c2 = arith.constant 2.0 : f32
+  %0 = arith.divf %arg0, %c1 : f32
+  %1 = arith.divf %c1, %c2 : f32
+  return %0, %1 : f32, f32
+}
+
+// -----
+
+// CHECK-LABEL: @test_cmpf(
+func @test_cmpf(%arg0 : f32) -> (i1, i1, i1, i1) {
+//   CHECK-DAG:   %[[T:.*]] = arith.constant true
+//   CHECK-DAG:   %[[F:.*]] = arith.constant false
+//       CHECK:   return %[[F]], %[[F]], %[[T]], %[[T]]
+  %nan = arith.constant 0x7fffffff : f32
+  %0 = arith.cmpf olt, %nan, %arg0 : f32
+  %1 = arith.cmpf olt, %arg0, %nan : f32
+  %2 = arith.cmpf ugt, %nan, %arg0 : f32
+  %3 = arith.cmpf ugt, %arg0, %nan : f32
+  return %0, %1, %2, %3 : i1, i1, i1, i1
+}
+
+// -----
+
 // CHECK-LABEL: @constant_FPtoUI(
 func @constant_FPtoUI() -> i32 {
   // CHECK: %[[C0:.+]] = arith.constant 2 : i32
   // CHECK: return %[[C0]]
   %c0 = arith.constant 2.0 : f32
   %res = arith.fptoui %c0 : f32 to i32
   return %res : i32
 }

 // CHECK-LABEL: @constant_UItoFP(
 func @constant_UItoFP() -> f32 {
   // CHECK: %[[C0:.+]] = arith.constant 2.000000e+00 : f32
   // CHECK: return %[[C0]]
   %c0 = arith.constant 2 : i32
   %res = arith.sitofp %c0 : i32 to f32
   return %res : f32
 }
-
-// -----
-// CHECK-LABEL: @constant_MinMax(
-func @constant_MinMax(%arg0 : f32) -> f32 {
-  // CHECK:  %[[const:.+]] = arith.constant
-  // CHECK:  %[[min:.+]] = arith.minf %arg0, %[[const]] : f32
-  // CHECK:  %[[res:.+]] = arith.maxf %[[min]], %[[const]] : f32
-  // CHECK:   return %[[res]]
-  %const = arith.constant 0.0 : f32
-  %min = arith.minf %const, %arg0 : f32
-  %res = arith.maxf %const, %min : f32
-  return %res : f32
-}
-
-// -----
-// CHECK-LABEL: @cmpf_nan(
-func @cmpf_nan(%arg0 : f32) -> (i1, i1, i1, i1) {
-//   CHECK-DAG:   %[[T:.*]] = arith.constant true
-//   CHECK-DAG:   %[[F:.*]] = arith.constant false
-//       CHECK:   return %[[F]], %[[F]], %[[T]], %[[T]]
-  %nan = arith.constant 0x7fffffff : f32
-  %0 = arith.cmpf olt, %nan, %arg0 : f32
-  %1 = arith.cmpf olt, %arg0, %nan : f32
-  %2 = arith.cmpf ugt, %nan, %arg0 : f32
-  %3 = arith.cmpf ugt, %arg0, %nan : f32
-  return %0, %1, %2, %3 : i1, i1, i1, i1
-}