[mlir][Math] Support fold Log2Op with constant dense.

[jacquesguan]

Summary

This patch is similar to D129108, it adds a conditional unary constant folder which allow to exit when the constants not meet the fold condition. And use it for Log2Op to make it able to fold the constant dense.

mlir/include/mlir/Dialect/CommonFolders.h

return constFoldBinaryOpConditional<AttrElementT>(
       operands,
       [&](ElementValueT a, ElementValueT b) -> Optional<ElementValueT> {
         return calculate(a, b);
       });
 }
 
 /// 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,
+template <
+    class AttrElementT, class ElementValueT = typename AttrElementT::ValueType,
+    class CalculationT = function_ref<Optional<ElementValueT>(ElementValueT)>>
+Attribute constFoldUnaryOpConditional(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()));
+    auto res = calculate(op.getValue());
+    if (!res)
+      return {};
+    return AttrElementT::get(op.getType(), *res);
   }
   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);
+    if (!elementResult)
+      return {};
+    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> elementResults;
     elementResults.reserve(op.getNumElements());
-    for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt)
-      elementResults.push_back(calculate(*opIt));
+    for (size_t i = 0, e = op.getNumElements(); i < e; ++i, ++opIt) {
+      auto elementResult = calculate(*opIt);
+      if (!elementResult)
+        return {};
+      elementResults.push_back(*elementResult);
+    }
     return DenseElementsAttr::get(op.getType(), elementResults);
   }
   return {};
 }
 
+template <class AttrElementT,
+          class ElementValueT = typename AttrElementT::ValueType,
+          class CalculationT = function_ref<ElementValueT(ElementValueT)>>
+Attribute constFoldUnaryOp(ArrayRef<Attribute> operands,
+                           const CalculationT &&calculate) {
+  return constFoldUnaryOpConditional<AttrElementT>(
+      operands,
+      [&](ElementValueT a) -> Optional<ElementValueT> { return calculate(a); });
+}
+
 template <
     class AttrElementT, class TargetAttrElementT,
     class ElementValueT = typename AttrElementT::ValueType,
     class TargetElementValueT = typename TargetAttrElementT::ValueType,
     class CalculationT = function_ref<TargetElementValueT(ElementValueT, bool)>>
 Attribute constFoldCastOp(ArrayRef<Attribute> operands, Type resType,
                           const CalculationT &calculate) {
   assert(operands.size() == 1 && "Cast op takes one operand");

mlir/lib/Dialect/Math/IR/MathOps.cpp

OpFoldResult math::CtPopOp::fold(ArrayRef<Attribute> operands) {
   });
 }
 
 //===----------------------------------------------------------------------===//
 // Log2Op folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::Log2Op::fold(ArrayRef<Attribute> operands) {
-  auto constOperand = operands.front();
-  if (!constOperand)
-    return {};
-
-  auto attr = constOperand.dyn_cast<FloatAttr>();
-  if (!attr)
-    return {};
-
-  auto ft = getType().cast<FloatType>();
-
-  APFloat apf = attr.getValue();
-
-  if (apf.isNegative())
+  return constFoldUnaryOpConditional<FloatAttr>(
+      operands, [](const APFloat &a) -> Optional<APFloat> {
+        if (a.isNegative())
           return {};
 
-  if (ft.getWidth() == 64)
-    return FloatAttr::get(getType(), log2(apf.convertToDouble()));
+        if (a.getSizeInBits(a.getSemantics()) == 64)
+          return APFloat(log2(a.convertToDouble()));
 
-  if (ft.getWidth() == 32)
-    return FloatAttr::get(getType(), log2f(apf.convertToFloat()));
+        if (a.getSizeInBits(a.getSemantics()) == 32)
+          return APFloat(log2f(a.convertToFloat()));
 
         return {};
+      });
 }
 
 //===----------------------------------------------------------------------===//
 // PowFOp folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::PowFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOpConditional<FloatAttr>(

mlir/test/Dialect/Math/canonicalize.mlir

 // CHECK:  %[[res:.+]] = math.log2 %[[cst]] : f64
   // CHECK: return %[[res]]
 func.func @log2_nofold2_64() -> f64 {
   %c = arith.constant -1.0 : f64
   %r = math.log2 %c : f64
   return %r : f64
 }
 
+// CHECK-LABEL: @log2_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<[0.000000e+00, 1.000000e+00, 1.58496249, 2.000000e+00]> : vector<4xf32>
+// CHECK: return %[[cst]]
+func.func @log2_fold_vec() -> (vector<4xf32>) {
+  %v1 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : vector<4xf32>
+  %0 = math.log2 %v1 : vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
 // CHECK-LABEL: @powf_fold
 // CHECK: %[[cst:.+]] = arith.constant 4.000000e+00 : f32
 // CHECK: return %[[cst]]
 func.func @powf_fold() -> f32 {
   %c = arith.constant 2.0 : f32
   %r = math.powf %c, %c : f32
   return %r : f32
 }

mlir/test/mlir-cpu-runner/math-polynomial-approx.mlir

 }
 
 func.func @log2() {
   // CHECK: 3.81887
   %0 = arith.constant 14.112233 : f32
   %1 = math.log2 %0 : f32
   vector.print %1 : f32
 
-  // CHECK: -2, -0.415037, 0, 0.321928
+  // CHECK: -2, -0.415038, 0, 0.321928
   %2 = arith.constant dense<[0.25, 0.75, 1.0, 1.25]> : vector<4xf32>
   %3 = math.log2 %2 : vector<4xf32>
   vector.print %3 : vector<4xf32>
 
   // CHECK: -3.32193, -2.32193, -1.73697, -1.32193, -1, -0.736966, -0.514573, -0.321928
   %4 = arith.constant dense<[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8]> : vector<8xf32>
   %5 = math.log2 %4 : vector<8xf32>
   vector.print %5 : vector<8xf32>