[MLIR][Math] Add constant folder for powf

[wsmoses]

Summary

Constant fold powf, given two constant operands and a compatible type

[ftynse]

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

93

Should we rather check for ft.getFloatSemantics() being Double/Float here?

[wsmoses]

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

93

I'm not aware of a good check on fltSemantics for that (https://llvm.org/doxygen/structllvm_1_1fltSemantics.html). If you know a better check, would be happy to use (and presumably also check the log2 one above). If I recall from the log2 case this is what LLVM itself does presently.

[ftynse]

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

93

&APFloatBase::getIEEESingle() == &ft.getFloatSemantics() ? AFAIK, these are singleton objects return by-reference so you can just compare their addresses.

[bondhugula]

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

76

The op definition guarantees that the result is float-like. Things enforced by the verifier need not be re-checked.

82

unsigned

82–84

Use llvm::all_of on operands?

mlir/test/Dialect/Math/canonicalize.mlir

79

This and the one above can be a CHECK-NEXT.

[ftynse]

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

76

Float-like means it's a float scalar, a vector of floats or a tensor of floats. The logic below only applies to float scalars, so the check is required.

[bondhugula]

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

76

Oh, missed that, thanks.

mlir/include/mlir/Dialect/Math/IR/MathOps.td

let description = [{
 
     // SIMD pointwise vector exponentiation
     %f = math.powf %g, %h : vector<4xf32>
 
     // Tensor pointwise exponentiation.
     %x = math.powf %y, %z : tensor<4x?xbf16>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RsqrtOp
 //===----------------------------------------------------------------------===//
 
 def Math_RsqrtOp : Math_FloatUnaryOp<"rsqrt"> {
   let summary = "reciprocal of sqrt (1 / sqrt of the specified value)";

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

OpFoldResult math::Log2Op::fold(ArrayRef<Attribute> operands) {
 
   if (apf.isNegative())
     return {};
 
   if (ft.getWidth() == 64)
     return FloatAttr::get(getType(), log2(apf.convertToDouble()));
 
   if (ft.getWidth() == 32)
-    return FloatAttr::get(getType(), log2f(apf.convertToDouble()));
+    return FloatAttr::get(getType(), log2f(apf.convertToFloat()));
+
+  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 {};
 }
 
 /// Materialize an integer or floating point constant.
 Operation *math::MathDialect::materializeConstant(OpBuilder &builder,
                                                   Attribute value, Type type,
                                                   Location loc) {
   return builder.create<arith::ConstantOp>(loc, value, type);
 }

mlir/test/Dialect/Math/canonicalize.mlir

 // CHECK: %[[cst:.+]] = arith.constant -1.000000e+00 : f64
 // CHECK:  %[[res:.+]] = math.log2 %[[cst]] : f64
   // CHECK: return %[[res]]
 func @log2_nofold2_64() -> f64 {
   %c = arith.constant -1.0 : f64
   %r = math.log2 %c : f64
   return %r : f64
 }
+
+// CHECK-LABEL: @powf_fold
+// CHECK: %[[cst:.+]] = arith.constant 4.000000e+00 : f32
+// CHECK: return %[[cst]]
+func @powf_fold() -> f32 {
+  %c = arith.constant 2.0 : f32
+  %r = math.powf %c, %c : f32
+  return %r : f32
+}