[mlir][Arithmetic] Add constant folder for negf.

[ftynse]

Do we expect folders for other unary operations or is the function template premature generalization?

mlir/include/mlir/Dialect/CommonFolders.h

79

Nit: I'd consider Calculation &&calculate here.

102

Nit: don't specify the explicit number of stack arguments in SmallVector unless you have a good reason to pick a specific number.

[jacquesguan]

address comment.

[jacquesguan]

In D123293#3435823, @ftynse wrote:

Do we expect folders for other unary operations or is the function template premature generalization?

I will try to create more cosntant folers for unary operatrions by this function template.

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

let description = [{
 
     // SIMD vector element-wise negation value.
     %f = arith.negf %g : vector<4xf32>
 
     // Tensor element-wise negation value.
     %x = arith.negf %y : tensor<4x?xf8>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
 def Arith_AddFOp : Arith_FloatBinaryOp<"addf", [Commutative]> {
   let summary = "floating point addition operation";

mlir/include/mlir/Dialect/CommonFolders.h

if (operands[0].isa<SplatElementsAttr>() &&
     SmallVector<ElementValueT, 4> elementResults;
     elementResults.reserve(lhs.getNumElements());
     for (size_t i = 0, e = lhs.getNumElements(); i < e; ++i, ++lhsIt, ++rhsIt)
       elementResults.push_back(calculate(*lhsIt, *rhsIt));
     return DenseElementsAttr::get(lhs.getType(), elementResults);
   }
   return {};
 }
+
+/// Performs constant folding `calculate` with element-wise behavior on the one
+/// attributes in `operands` and returns the result if possible.
+template <class AttrElementT,
+          class ElementValueT = typename AttrElementT::ValueType,
+          class CalculationT = function_ref<ElementValueT(ElementValueT)>>
+Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,
+                           const CalculationT &calculate) {
+  assert(operands.size() == 1 && "unary op takes one operands");
+  if (!operands[0])
+    return {};
+
+  if (operands[0].isa<AttrElementT>()) {
+    auto op = operands[0].cast<AttrElementT>();
+
+    return AttrElementT::get(op.getType(), calculate(op.getValue()));
+  }
+  if (operands[0].isa<SplatElementsAttr>()) {
+    // Both operands are splats so we can avoid expanding the values out and
+    // just fold based on the splat value.
+    auto op = operands[0].cast<SplatElementsAttr>();
+
+    auto elementResult = calculate(op.getSplatValue<ElementValueT>());
+    return DenseElementsAttr::get(op.getType(), elementResult);
+  } else if (operands[0].isa<ElementsAttr>()) {
+    // Operands are ElementsAttr-derived; perform an element-wise fold by
+    // expanding the values.
+    auto op = operands[0].cast<ElementsAttr>();
+
+    auto opIt = op.value_begin<ElementValueT>();
+    SmallVector<ElementValueT, 4> elementResults;
+    elementResults.reserve(op.getNumElements());
+    for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt)
+      elementResults.push_back(calculate(*opIt));
+    return DenseElementsAttr::get(op.getType(), elementResults);
+  }
+  return {};
+}
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_COMMONFOLDERS_H

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

 }
 
 void arith::XOrIOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
   patterns.add<XOrINotCmpI>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// NegFOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult arith::NegFOp::fold(ArrayRef<Attribute> operands) {
+  return constFoldUnaryOp<FloatAttr>(operands,
+                                     [](const APFloat &a) { return -a; });
+}
+
+//===----------------------------------------------------------------------===//
 // AddFOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult arith::AddFOp::fold(ArrayRef<Attribute> operands) {
   // addf(x, -0) -> x
   if (matchPattern(getRhs(), m_NegZeroFloat()))
     return getLhs();
 

mlir/test/Dialect/Arithmetic/canonicalize.mlir

 // CHECK: %[[res:.+]] = arith.shrsi
 // CHECK: return %[[res]]
 func @nofoldShrs2() -> i64 {
   %c1 = arith.constant 8 : i64
   %cm32 = arith.constant -32 : i64
   %r = arith.shrsi %c1, %cm32 : i64
   return %r : i64
 }
+
+// -----
+
+// CHECK-LABEL: @test_negf(
+// CHECK: %[[res:.+]] = arith.constant -2.0
+// CHECK: return %[[res]]
+func @test_negf() -> (f32) {
+  %c = arith.constant 2.0 : f32
+  %0 = arith.negf %c : f32
+  return %0: f32
+}