Implement constant folding for PowFOp

[jacksonfellows]

Summary

Add a constant folder for PowFOp. Analogous to existing folders for
floating point operators, but instead of using an APFloat method to
perform its operation it first converts the operands to doubles, calls
the built-in pow function, and then converts the result back into the
proper floating point type. This is necessary since APFloat lacks a
pow method. This behavior matches how constant folding is implemented
for the pow intrinsic in
LLVM (https://github.com/llvm/llvm-project/blob/689de5841c1c4c9b0fe711b61d26f7425cf99423/llvm/lib/Analysis/ConstantFolding.cpp#L2373).

[mehdi_amini]

Can you add a lit test for this?

[jacksonfellows]

In D95105#2511374, @mehdi_amini wrote:

Can you add a lit test for this?

Should I add it to this file: https://github.com/llvm/llvm-project/blob/main/mlir/test/Transforms/constant-fold.mlir?
Also should the cases mirror what already exists?

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2276

How does this work for floating point types with greater than 64 bits? e.g. FP80/FP128/etc.

[jacksonfellows]

As far as I can tell they would lose precision on the conversion. It seems like LLVM avoids this by only folding halfs, floats, and doubles. We could do the same by checking the operand types first, which is probably the right solution.

[jacksonfellows]

Since MLIR is also handling tensors and vectors, I'm not entirely sure how this check would be implemented (i.e. the operands can't always be cast to FloatAttr). The function constFoldBinaryOp currently handles these cases, so we'd have to re-implement some of that logic. Does someone with more experience with the existing folders have any good ideas about this?

[rriddle]

In D95105#2511444, @jacksonfellows wrote:

Since MLIR is also handling tensors and vectors, I'm not entirely sure how this check would be implemented (i.e. the operands can't always be cast to FloatAttr). The function constFoldBinaryOp currently handles these cases, so we'd have to re-implement some of that logic. Does someone with more experience with the existing folders have any good ideas about this?

Can you just add a check before calling constFoldBinaryOp that the result type of the powf when fed to mlir::getElementTypeOrSelf is one of the valid types?

[jacksonfellows]

Add check to only fold the pow of halfs, floats, and doubles.

This matches the behavior present in LLVM.

[jacksonfellows]

Add lit tests for powf constant folding.

[rriddle]

Thanks for adding the folding!

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

2276

I think we can include BF16 here as well.

2277

I wouldn't say "following LLVM" here. The limitation is due to how the folding is implemented, i.e. it goes to double which naturally can't support larger representations.

2281

This comment shouldn't be necessary given that the element types are already guaranteed to be the same.

[jacksonfellows]

Add BF16 to allowed FP types

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let description = [{
 
     // SIMD pointwise vector exponentiation
     %f = powf %g, %h : vector<4xf32>
 
     // Tensor pointwise exponentiation.
     %x = powf %y, %z : tensor<4x?xbf16>
     ```
   }];
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // PrefetchOp
 //===----------------------------------------------------------------------===//
 
 def PrefetchOp : Std_Op<"prefetch"> {
   let summary = "prefetch operation";

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

OpFoldResult OrOp::fold(ArrayRef<Attribute> operands) {
   if (lhs() == rhs())
     return rhs();
 
   return constFoldBinaryOp<IntegerAttr>(operands,
                                         [](APInt a, APInt b) { return a | b; });
 }
 
 //===----------------------------------------------------------------------===//
+// PowFOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult PowFOp::fold(ArrayRef<Attribute> operands) {
+  assert(operands.size() == 2 && "binary op takes two operands");
+  if (operands[0] && operands[1] &&
+      operands[0].getType() == operands[1].getType()) {
+    // make sure that types match (similar check performed in constFoldBinaryOp)
+    Type ty = getElementTypeOrSelf(operands[0]);
+    if (ty.isBF16() || ty.isF16() || ty.isF32() || ty.isF64()) {
+      // only fold the pow of floating point types that can be represented as a
+      // double without losing information
+      return constFoldBinaryOp<FloatAttr>(operands, [](APFloat a, APFloat b) {
+        bool unused;
+        // assume a and b are the same floating point type (i.e. share the same
+        // semantics)
+        APFloat res = APFloat(pow(FloatAttr::getValueAsDouble(a),
+                                  FloatAttr::getValueAsDouble(b)));
+        res.convert(a.getSemantics(), APFloat::rmNearestTiesToEven, &unused);
+        return res;
+      });
+    }
+  }
+  return {};
+}
+
+//===----------------------------------------------------------------------===//
 // PrefetchOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, PrefetchOp op) {
   p << PrefetchOp::getOperationName() << " " << op.memref() << '[';
   p.printOperands(op.indices());
   p << ']' << ", " << (op.isWrite() ? "write" : "read");
   p << ", locality<" << op.localityHint();

mlir/test/Transforms/constant-fold.mlir

 // -----
 
 // CHECK-LABEL: func @subview_scalar_fold
 func @subview_scalar_fold(%arg0: memref<f32>) -> memref<f32> {
   // CHECK-NOT: subview
   %c = subview %arg0[] [] [] : memref<f32> to memref<f32>
   return %c : memref<f32>
 }
+
+// -----
+
+// CHECK-LABEL: func @simple_powf
+func @simple_powf() -> f32 {
+  %0 = constant 4.5 : f32
+  %1 = constant 2.0 : f32
+
+  // CHECK-NEXT: [[C:%.+]] = constant 2.025{{0*}}e+01 : f32
+  %2 = powf %0, %1 : f32
+
+  // CHECK-NEXT: return [[C]]
+  return %2 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @powf_splat_tensor
+func @powf_splat_tensor() -> tensor<4xf32> {
+  %0 = constant dense<4.5> : tensor<4xf32>
+  %1 = constant dense<2.0> : tensor<4xf32>
+
+  // CHECK-NEXT: [[C:%.+]] = constant dense<2.025{{0*}}e+01> : tensor<4xf32>
+  %2 = powf %0, %1 : tensor<4xf32>
+
+  // CHECK-NEXT: return [[C]]
+  return %2 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @powf_dont_fold_f128
+func @powf_dont_fold_f128() -> f128 {
+  // CHECK-NEXT: [[A:%.+]] = constant 4.5{{0*}}e+00 : f128
+  %0 = constant 4.5 : f128
+  // CHECK-NEXT: [[B:%.+]] = constant 2.0{{0*}}e+00 : f128
+  %1 = constant 2.0 : f128
+
+  // CHECK-NEXT: %0 = powf [[A]], [[B]] : f128
+  %2 = powf %0, %1 : f128
+
+  // CHECK-NEXT: return %0
+  return %2 : f128
+}