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

[jacquesguan]

Summary

This patch adds a conditional binary constant folder which allow to exit when the constants not meet the fold condition. And use it for PowFOp to make it able to fold the constant dense.

[Mogball]

mlir/include/mlir/Dialect/CommonFolders.h

77

This duplicates a lot of code with constFoldBinaryOp. Can you change the latter to call into this function?

[jacquesguan]

Address comment.

[jacquesguan]

mlir/include/mlir/Dialect/CommonFolders.h

77

Done, thanks.

mlir/include/mlir/Dialect/CommonFolders.h

if (operands[0].isa<SplatElementsAttr>() &&
     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 {};
 }
 
+template <
+    class AttrElementT, class ElementValueT = typename AttrElementT::ValueType,
+    class CalculationT =
+        function_ref<Optional<ElementValueT>(ElementValueT, ElementValueT)>>
+Attribute constFoldBinaryOpConditional(ArrayRef<Attribute> operands,
+                                       const CalculationT &calculate) {
+  assert(operands.size() == 2 && "binary op takes two operands");
+  if (!operands[0] || !operands[1])
+    return {};
+  if (operands[0].getType() != operands[1].getType())
+    return {};
+
+  if (operands[0].isa<AttrElementT>() && operands[1].isa<AttrElementT>()) {
+    auto lhs = operands[0].cast<AttrElementT>();
+    auto rhs = operands[1].cast<AttrElementT>();
+
+    auto calRes = calculate(lhs.getValue(), rhs.getValue());
+
+    if (!calRes)
+      return {};
+
+    return AttrElementT::get(lhs.getType(), *calRes);
+  }
+
+  if (operands[0].isa<SplatElementsAttr>() &&
+      operands[1].isa<SplatElementsAttr>()) {
+    // Both operands are splats so we can avoid expanding the values out and
+    // just fold based on the splat value.
+    auto lhs = operands[0].cast<SplatElementsAttr>();
+    auto rhs = operands[1].cast<SplatElementsAttr>();
+
+    auto elementResult = calculate(lhs.getSplatValue<ElementValueT>(),
+                                   rhs.getSplatValue<ElementValueT>());
+    if (!elementResult)
+      return {};
+
+    return DenseElementsAttr::get(lhs.getType(), *elementResult);
+  } else if (operands[0].isa<ElementsAttr>() &&
+             operands[1].isa<ElementsAttr>()) {
+    // Operands are ElementsAttr-derived; perform an element-wise fold by
+    // expanding the values.
+    auto lhs = operands[0].cast<ElementsAttr>();
+    auto rhs = operands[1].cast<ElementsAttr>();
+
+    auto lhsIt = lhs.value_begin<ElementValueT>();
+    auto rhsIt = rhs.value_begin<ElementValueT>();
+    SmallVector<ElementValueT, 4> elementResults;
+    elementResults.reserve(lhs.getNumElements());
+    for (size_t i = 0, e = lhs.getNumElements(); i < e; ++i, ++lhsIt, ++rhsIt) {
+      auto elementResult = calculate(*lhsIt, *rhsIt);
+      if (!elementResult)
+        return {};
+      elementResults.push_back(*elementResult);
+    }
+
+    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");

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

OpFoldResult math::Log2Op::fold(ArrayRef<Attribute> operands) {
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // PowFOp folder
 //===----------------------------------------------------------------------===//
 
 OpFoldResult math::PowFOp::fold(ArrayRef<Attribute> operands) {
-  auto ft = getType().dyn_cast<FloatType>();
-  if (!ft)
-    return {};
-
-  APFloat vals[2]{APFloat(ft.getFloatSemantics()),
-                  APFloat(ft.getFloatSemantics())};
-  for (int i = 0; i < 2; ++i) {
-    if (!operands[i])
-      return {};
-
-    auto attr = operands[i].dyn_cast<FloatAttr>();
-    if (!attr)
-      return {};
-
-    vals[i] = attr.getValue();
-  }
-
-  if (ft.getWidth() == 64)
-    return FloatAttr::get(
-        getType(), pow(vals[0].convertToDouble(), vals[1].convertToDouble()));
-
-  if (ft.getWidth() == 32)
-    return FloatAttr::get(
-        getType(), powf(vals[0].convertToFloat(), vals[1].convertToFloat()));
+  return constFoldBinaryOpConditional<FloatAttr>(
+      operands, [](const APFloat &a, const APFloat &b) -> Optional<APFloat> {
+        if (a.getSizeInBits(a.getSemantics()) == 64 &&
+            b.getSizeInBits(b.getSemantics()) == 64)
+          return APFloat(pow(a.convertToDouble(), b.convertToDouble()));
+
+        if (a.getSizeInBits(a.getSemantics()) == 32 &&
+            b.getSizeInBits(b.getSemantics()) == 32)
+          return APFloat(powf(a.convertToFloat(), b.convertToFloat()));
 
         return {};
+      });
 }
 
 OpFoldResult math::SqrtOp::fold(ArrayRef<Attribute> operands) {
   auto constOperand = operands.front();
   if (!constOperand)
     return {};
 
   auto attr = constOperand.dyn_cast<FloatAttr>();

mlir/test/Dialect/Math/canonicalize.mlir

 // 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
 }
 
+// CHECK-LABEL: @powf_fold_vec
+// CHECK: %[[cst:.+]] = arith.constant dense<[1.000000e+00, 4.000000e+00, 9.000000e+00, 1.600000e+01]> : vector<4xf32>
+// CHECK: return %[[cst]]
+func.func @powf_fold_vec() -> (vector<4xf32>) {
+  %v1 = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : vector<4xf32>
+  %v2 = arith.constant dense<[2.0, 2.0, 2.0, 2.0]> : vector<4xf32>
+  %0 = math.powf %v1, %v2 : vector<4xf32>
+  return %0 : vector<4xf32>
+}
+
 // CHECK-LABEL: @sqrt_fold
 // CHECK: %[[cst:.+]] = arith.constant 2.000000e+00 : f32
 // CHECK: return %[[cst]]
 func.func @sqrt_fold() -> f32 {
   %c = arith.constant 4.0 : f32
   %r = math.sqrt %c : f32
   return %r : f32
 }