[mlir][math] Added math::FPowI conversion to LLVM dialect.

[vzakhari]

Summary

The operations are converted into LLVM::PowIOp.

Depends on D139805

[vzakhari]

Hi @wsmoses, can you please read through the TBD comments in this change and see if you have any comments? This is related to LLVM::PowIOp added in https://reviews.llvm.org/D116364. It is not quite aligned with the supported forms of llvm.powi.* LLVM intrinsic, so I wonder what the plan for this MLIR operation is.

[vzakhari]

@wsmoses, friendly ping regarding the questions above.

[vzakhari]

rebase

[Mogball]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

285

Can you try VectorConvertToLLVMPattern<FPowIOp, PowIOp>?

[vzakhari]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

285

LLVM::PowIOp is only guaranteed to support i32 exponent, so for other cases LLVM backend will not be able to handle the resulting LLVM IR. I described the two options here: https://discourse.llvm.org/t/rfc-adding-more-power-operations-into-math-dialect/63975/3?u=szakharin

I am leaning toward using LLVM::PowIOp for i32 exponent and use the outlined implementation otherwise.

[Mogball]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

285

Do you need an exponent greater than i32?

[vzakhari]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

285

I am going to use it for lowering Fortran exponentiation operator (**), and it does support i64 exponent.

[vzakhari]

This is a one-to-one conversion to LLVM::PowIOp. Flang will use MathToFuncs convertor (D139804) to convert unsupported exponent-width cases before engaging MathToLLVM.

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

 using FloorOpLowering =
     ConvertFMFMathToLLVMPattern<math::FloorOp, LLVM::FFloorOp>;
 using FmaOpLowering = ConvertFMFMathToLLVMPattern<math::FmaOp, LLVM::FMAOp>;
 using Log10OpLowering =
     ConvertFMFMathToLLVMPattern<math::Log10Op, LLVM::Log10Op>;
 using Log2OpLowering = ConvertFMFMathToLLVMPattern<math::Log2Op, LLVM::Log2Op>;
 using LogOpLowering = ConvertFMFMathToLLVMPattern<math::LogOp, LLVM::LogOp>;
 using PowFOpLowering = ConvertFMFMathToLLVMPattern<math::PowFOp, LLVM::PowOp>;
+using FPowIOpLowering =
+    ConvertFMFMathToLLVMPattern<math::FPowIOp, LLVM::PowIOp>;
 using RoundEvenOpLowering =
     ConvertFMFMathToLLVMPattern<math::RoundEvenOp, LLVM::RoundEvenOp>;
 using RoundOpLowering =
     ConvertFMFMathToLLVMPattern<math::RoundOp, LLVM::RoundOp>;
 using SinOpLowering = ConvertFMFMathToLLVMPattern<math::SinOp, LLVM::SinOp>;
 using SqrtOpLowering = ConvertFMFMathToLLVMPattern<math::SqrtOp, LLVM::SqrtOp>;
 using FTruncOpLowering =
     ConvertFMFMathToLLVMPattern<math::TruncOp, LLVM::FTruncOp>;
return LLVM::detail::handleMultidimensionalVectors(
               rewriter.create<LLVM::ConstantOp>(loc, llvm1DVectorTy, splatAttr);
           auto sqrt = rewriter.create<LLVM::SqrtOp>(
               loc, llvm1DVectorTy, operands[0], sqrtAttrs.getAttrs());
           return rewriter.create<LLVM::FDivOp>(
               loc, llvm1DVectorTy, ValueRange{one, sqrt}, divAttrs.getAttrs());
         },
         rewriter);
   }
 };
 
 struct ConvertMathToLLVMPass
     : public impl::ConvertMathToLLVMBase<ConvertMathToLLVMPass> {
   ConvertMathToLLVMPass() = default;
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     LLVMTypeConverter converter(&getContext());
patterns.add<
     CopySignOpLowering,
     CosOpLowering,
     CountLeadingZerosOpLowering,
     CountTrailingZerosOpLowering,
     CtPopFOpLowering,
     Exp2OpLowering,
     ExpM1OpLowering,
     ExpOpLowering,
+    FPowIOpLowering,
     FloorOpLowering,
     FmaOpLowering,
     Log10OpLowering,
     Log1pOpLowering,
     Log2OpLowering,
     LogOpLowering,
     PowFOpLowering,
     RoundEvenOpLowering,

mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir

func.func @rsqrt_multidim_vector(%arg0 : vector<4x3xf32>) {
   // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : vector<3xf32>
   // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[DIV]], %{{.*}}[0] : !llvm.array<4 x vector<3xf32>>
   %0 = math.rsqrt %arg0 : vector<4x3xf32>
   func.return
 }
 
 // -----
 
+// CHECK-LABEL: func @fpowi(
+// CHECK-SAME: f64
+func.func @fpowi(%arg0 : f64, %arg1 : i32) {
+  // CHECK: llvm.intr.powi(%arg0, %arg1) : (f64, i32) -> f64
+  %0 = math.fpowi %arg0, %arg1 : f64, i32
+  func.return
+}
+
+
+// -----
+
 // CHECK-LABEL: func @powf(
 // CHECK-SAME: f64
 func.func @powf(%arg0 : f64) {
   // CHECK: %[[POWF:.*]] = llvm.intr.pow(%arg0, %arg0) : (f64, f64) -> f64
   %0 = math.powf %arg0, %arg0 : f64
   func.return
 }