[mlir][Arithmetic] Fold NegF in MulF and DivF.

[jacquesguan]

Summary

This patch adds the following combination:

mulf(negf(x), negf(y)) -> mulf(x, y)
divf(negf(x), negf(y)) -> divf(x, y)

[jacquesguan]

ping.

[jacquesguan]

ping.

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

let description = [{
     // Tensor pointwise multiplication.
     %x = arith.mulf %y, %z : tensor<4x?xbf16>
     ```
 
     TODO: In the distant future, this will accept optional attributes for fast
     math, contraction, rounding mode, and other controls.
   }];
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DivFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_DivFOp : Arith_FloatBinaryOp<"divf"> {
   let summary = "floating point division operation";
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RemFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_RemFOp : Arith_FloatBinaryOp<"remf"> {
   let summary = "floating point division remainder operation";

mlir/lib/Dialect/Arithmetic/IR/ArithmeticCanonicalization.td

def OrOfExtUI :
     Pat<(Arith_OrIOp (Arith_ExtUIOp $x), (Arith_ExtUIOp $y)), (Arith_ExtUIOp (Arith_OrIOp $x, $y)),
       [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
 // or extsi(x), extsi(y) -> extsi(or(x,y))
 def OrOfExtSI :
     Pat<(Arith_OrIOp (Arith_ExtSIOp $x), (Arith_ExtSIOp $y)), (Arith_ExtSIOp (Arith_OrIOp $x, $y)),
       [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
 
+//===----------------------------------------------------------------------===//
+// MulFOp
+//===----------------------------------------------------------------------===//
+
+// mulf(negf(x), negf(y)) -> mulf(x,y)
+def MulFOfNegF :
+    Pat<(Arith_MulFOp (Arith_NegFOp $x), (Arith_NegFOp $y)), (Arith_MulFOp $x, $y),
+      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+
+//===----------------------------------------------------------------------===//
+// DivFOp
+//===----------------------------------------------------------------------===//
+
+// divf(negf(x), negf(y)) -> divf(x,y)
+def DivFOfNegF :
+    Pat<(Arith_DivFOp (Arith_NegFOp $x), (Arith_NegFOp $y)), (Arith_DivFOp $x, $y),
+      [(Constraint<CPred<"$0.getType() == $1.getType()">> $x, $y)]>;
+
 #endif // ARITHMETIC_PATTERNS

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

OpFoldResult arith::MulFOp::fold(ArrayRef<Attribute> operands) {
   // mulf(x, 1) -> x
   if (matchPattern(getRhs(), m_OneFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a * b; });
 }
 
+void arith::MulFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                MLIRContext *context) {
+  patterns.add<MulFOfNegF>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // DivFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::DivFOp::fold(ArrayRef<Attribute> operands) {
   // divf(x, 1) -> x
   if (matchPattern(getRhs(), m_OneFloat()))
     return getLhs();
 
   return constFoldBinaryOp<FloatAttr>(
       operands, [](const APFloat &a, const APFloat &b) { return a / b; });
 }
 
+void arith::DivFOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                MLIRContext *context) {
+  patterns.add<DivFOfNegF>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // RemFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::RemFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(operands,
                                       [](const APFloat &a, const APFloat &b) {
                                         APFloat Result(a);

mlir/test/Dialect/Arithmetic/canonicalize.mlir

func.func @test_mulf(%arg0 : f32) -> (f32, f32, f32, f32) {
   %c2 = arith.constant 2.0 : f32
   %0 = arith.mulf %c2, %arg0 : f32
   %1 = arith.mulf %arg0, %c1 : f32
   %2 = arith.mulf %c1, %arg0 : f32
   %3 = arith.mulf %c2, %c2 : f32
   return %0, %1, %2, %3 : f32, f32, f32, f32
 }
 
+// CHECK-LABEL: @test_mulf1(
+func.func @test_mulf1(%arg0 : f32, %arg1 : f32) -> (f32) {
+  // CHECK-NEXT:  %[[X:.+]] = arith.mulf %arg0, %arg1 : f32
+  // CHECK-NEXT:  return %[[X]]
+  %0 = arith.negf %arg0 : f32
+  %1 = arith.negf %arg1 : f32
+  %2 = arith.mulf %0, %1 : f32
+  return %2 : f32
+}
+
 // -----
 
 // CHECK-LABEL: @test_divf(
 func.func @test_divf(%arg0 : f64) -> (f64, f64) {
   // CHECK-NEXT:  %[[C5:.+]] = arith.constant 5.000000e-01
   // CHECK-NEXT:   return %arg0, %[[C5]]
   %c1 = arith.constant 1.0 : f64
   %c2 = arith.constant 2.0 : f64
   %0 = arith.divf %arg0, %c1 : f64
   %1 = arith.divf %c1, %c2 : f64
   return %0, %1 : f64, f64
 }
 
+// CHECK-LABEL: @test_divf1(
+func.func @test_divf1(%arg0 : f32, %arg1 : f32) -> (f32) {
+  // CHECK-NEXT:  %[[X:.+]] = arith.divf %arg0, %arg1 : f32
+  // CHECK-NEXT:  return %[[X]]
+  %0 = arith.negf %arg0 : f32
+  %1 = arith.negf %arg1 : f32
+  %2 = arith.divf %0, %1 : f32
+  return %2 : f32
+}
+
 // -----
 
 // CHECK-LABEL: @test_cmpf(
 func.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