[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

258

Can you try VectorConvertToLLVMPattern<FPowIOp, PowIOp>?

[vzakhari]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

258

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

258

Do you need an exponent greater than i32?

[vzakhari]

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

258

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/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

 def LLVM_Prefetch : LLVM_ZeroResultIntrOp<"prefetch", [0]> {
   let arguments = (ins LLVM_Type:$addr, LLVM_Type:$rw, LLVM_Type:$hint,
                    LLVM_Type:$cache);
 }
 def LLVM_SinOp : LLVM_UnaryIntrinsicOp<"sin">;
 def LLVM_RoundOp : LLVM_UnaryIntrinsicOp<"round">;
 def LLVM_SqrtOp : LLVM_UnaryIntrinsicOp<"sqrt">;
 def LLVM_PowOp : LLVM_BinarySameArgsIntrinsicOp<"pow">;
+// TBD:
+// Should LLVM::PowIOp be restricted to support only
+// the cases supported by llvm.powi.* intrinsic (e.g.
+// exponent must be scalar 'i32',
+// see "POWI exponent does not match sizeof(int)" verification
+// in LegalizeDAG.cpp)?
+// Or, should LLVM::PowIOp translation to LLVM IR handle
+// all the cases allowed here?
 def LLVM_PowIOp : LLVM_BinaryIntrinsicOp<"powi">;
 def LLVM_BitReverseOp : LLVM_UnaryIntrinsicOp<"bitreverse">;
 def LLVM_CountLeadingZerosOp : LLVM_CountZerosIntrinsicOp<"ctlz">;
 def LLVM_CountTrailingZerosOp : LLVM_CountZerosIntrinsicOp<"cttz">;
 def LLVM_CtPopOp : LLVM_UnaryIntrinsicOp<"ctpop">;
 def LLVM_MaxNumOp : LLVM_BinarySameArgsIntrinsicOp<"maxnum">;
 def LLVM_MinNumOp : LLVM_BinarySameArgsIntrinsicOp<"minnum">;
 def LLVM_MaximumOp : LLVM_BinarySameArgsIntrinsicOp<"maximum">;

mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp

return LLVM::detail::handleMultidimensionalVectors(
           auto sqrt =
               rewriter.create<LLVM::SqrtOp>(loc, llvm1DVectorTy, operands[0]);
           return rewriter.create<LLVM::FDivOp>(loc, llvm1DVectorTy, one, sqrt);
         },
         rewriter);
   }
 };
 
+// `fpwosi` is converted into llvm.powi intrinsic operation.
+struct FPowSIOpLowering : public ConvertOpToLLVMPattern<math::FPowSIOp> {
+  using ConvertOpToLLVMPattern<math::FPowSIOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(math::FPowSIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto baseType = adaptor.getOperands()[0].getType();
+    auto powerType = adaptor.getOperands()[1].getType();
+
+    if (!baseType || !LLVM::isCompatibleType(baseType) ||
+        !powerType || !LLVM::isCompatibleType(powerType))
+      return failure();
+
+    auto loc = op.getLoc();
+    if (baseType.isa<FloatType>()) {
+      // TBD: while LLVM::PowIOp is defined quite loosely,
+      //      its support in LLVM IR is very narrow. In particular,
+      //      only scalar 'i32' exponent is supported.
+      //      Should we convert to an inline implementation in cases
+      //      that are not supported by llvm.powi intrinsic?
+      rewriter.replaceOpWithNewOp<LLVM::PowIOp>(op, baseType, adaptor.getOperands());
+      return success();
+    }
+
+    auto resultType = op.getResult().getType();
+    auto vectorType = resultType.dyn_cast<VectorType>();
+    if (!vectorType)
+      return failure();
+
+    auto elementType = vectorType.getElementType().dyn_cast<FloatType>();
+    if (!elementType)
+      return failure();
+
+    // Convert to a sequence of scalar LLVM::PowIOp's.
+    return LLVM::detail::handleMultidimensionalVectors(
+        op.getOperation(), adaptor.getOperands(), *getTypeConverter(),
+        [&](Type llvm1DVectorTy, ValueRange operands) {
+          Value result = rewriter.create<LLVM::UndefOp>(loc, llvm1DVectorTy);
+          auto vecType = llvm1DVectorTy.cast<VectorType>();
+          // This loop is not needed, if LLVM::PowIOp has normal elementwise
+          // semantics (which is not clear since it is defined as a MLIR
+          // equivalent of llvm.powi.* that applies a scalar exponent
+          // to a vector base).
+          for (unsigned idx = 0, numElements = vecType.getNumElements();
+               idx < numElements; ++idx) {
+            Value pos = rewriter.create<LLVM::ConstantOp>(
+                loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(idx));
+            Value baseOp =
+                rewriter.create<LLVM::ExtractElementOp>(loc, operands[0], pos);
+            Value powerOp =
+                rewriter.create<LLVM::ExtractElementOp>(loc, operands[1], pos);
+            Value eltResult =
+                rewriter
+                    .create<LLVM::PowIOp>(loc, baseOp.getType(),
+                                          ValueRange{baseOp, powerOp});
+            result = rewriter.create<LLVM::InsertElementOp>(
+                loc, llvm1DVectorTy, result, eltResult, pos);
+          }
+          return result;
+        },
+        rewriter);
+  }
+};
+
 struct IPowSIOpLowering : public ConvertOpToLLVMPattern<math::IPowSIOp> {
 private:
   /// Create linkonce_odr function to implement the power function with
   /// the given \p elementType scalar type.
   ///
   /// template <typename T>
   /// T __mlir_math_ipowsi_*(T b, T p) {
   ///   if (p == T(0))
patterns.add<
     ExpOpLowering,
     FloorOpLowering,
     FmaOpLowering,
     Log10OpLowering,
     Log1pOpLowering,
     Log2OpLowering,
     LogOpLowering,
     PowFOpLowering,
+    FPowSIOpLowering,
     IPowSIOpLowering,
     RoundOpLowering,
     RsqrtOpLowering,
     SinOpLowering,
     SqrtOpLowering
   >(converter);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertMathToLLVMPass() {
   return std::make_unique<ConvertMathToLLVMPass>();
 }

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

 // CHECK:           %[[VAL_23:.*]] = llvm.insertvalue %[[VAL_22]], %[[VAL_4]][0] : !llvm.array<2 x vector<3xi64>>
 
 // CHECK:           llvm.call @__mlir_math_ipowsi_i64(%{{.*}}, %{{.*}}) : (i64, i64) -> i64
 // CHECK:           llvm.call @__mlir_math_ipowsi_i64(%{{.*}}, %{{.*}}) : (i64, i64) -> i64
 // CHECK:           llvm.call @__mlir_math_ipowsi_i64(%{{.*}}, %{{.*}}) : (i64, i64) -> i64
   %0 = math.ipowsi %arg0, %arg1 : vector<2x3xi64>
   func.return
 }
+
+// -----
+
+// CHECK-LABEL: func @fpowsi(
+// CHECK-SAME: %[[ARG0:.+]]: f64,
+// CHECK-SAME: %[[ARG1:.+]]: i64,
+// CHECK-SAME: %[[ARG2:.+]]: f16,
+// CHECK-SAME: %[[ARG3:.+]]: i32)
+func.func @fpowsi(%arg0: f64, %arg1: i64, %arg2: f16, %arg3: i32) {
+// CHECK: "llvm.intr.powi"(%[[ARG0]], %[[ARG1]]) : (f64, i64) -> f64
+  %0 = math.fpowsi %arg0, %arg1 : (f64, i64) -> f64
+// CHECK: "llvm.intr.powi"(%[[ARG2]], %[[ARG3]]) : (f16, i32) -> f16
+  %1 = math.fpowsi %arg2, %arg3 : (f16, i32) -> f16
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL:   func.func @fpowsi_vec(
+// CHECK-SAME:                          %[[VAL_0:.*]]: vector<2x3xf64>,
+// CHECK-SAME:                          %[[VAL_1:.*]]: vector<2x3xi32>) {
+func.func @fpowsi_vec(%arg0: vector<2x3xf64>, %arg1: vector<2x3xi32>) {
+// CHECK:           %[[VAL_4:.*]] = llvm.mlir.undef : !llvm.array<2 x vector<3xf64>>
+// CHECK:           %[[VAL_5:.*]] = llvm.extractvalue %[[VAL_2:.*]][0] : !llvm.array<2 x vector<3xf64>>
+// CHECK:           %[[VAL_6:.*]] = llvm.extractvalue %[[VAL_3:.*]][0] : !llvm.array<2 x vector<3xi32>>
+// CHECK:           %[[VAL_7:.*]] = llvm.mlir.undef : vector<3xf64>
+// CHECK:           %[[VAL_8:.*]] = llvm.mlir.constant(0 : i64) : i64
+// CHECK:           %[[VAL_9:.*]] = llvm.extractelement %[[VAL_5]]{{\[}}%[[VAL_8]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_10:.*]] = llvm.extractelement %[[VAL_6]]{{\[}}%[[VAL_8]] : i64] : vector<3xi32>
+// CHECK:           %[[VAL_11:.*]] = "llvm.intr.powi"(%[[VAL_9]], %[[VAL_10]]) : (f64, i32) -> f64
+// CHECK:           %[[VAL_12:.*]] = llvm.insertelement %[[VAL_11]], %[[VAL_7]]{{\[}}%[[VAL_8]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_13:.*]] = llvm.mlir.constant(1 : i64) : i64
+// CHECK:           %[[VAL_14:.*]] = llvm.extractelement %[[VAL_5]]{{\[}}%[[VAL_13]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_15:.*]] = llvm.extractelement %[[VAL_6]]{{\[}}%[[VAL_13]] : i64] : vector<3xi32>
+// CHECK:           %[[VAL_16:.*]] = "llvm.intr.powi"(%[[VAL_14]], %[[VAL_15]]) : (f64, i32) -> f64
+// CHECK:           %[[VAL_17:.*]] = llvm.insertelement %[[VAL_16]], %[[VAL_12]]{{\[}}%[[VAL_13]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_18:.*]] = llvm.mlir.constant(2 : i64) : i64
+// CHECK:           %[[VAL_19:.*]] = llvm.extractelement %[[VAL_5]]{{\[}}%[[VAL_18]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_20:.*]] = llvm.extractelement %[[VAL_6]]{{\[}}%[[VAL_18]] : i64] : vector<3xi32>
+// CHECK:           %[[VAL_21:.*]] = "llvm.intr.powi"(%[[VAL_19]], %[[VAL_20]]) : (f64, i32) -> f64
+// CHECK:           %[[VAL_22:.*]] = llvm.insertelement %[[VAL_21]], %[[VAL_17]]{{\[}}%[[VAL_18]] : i64] : vector<3xf64>
+// CHECK:           %[[VAL_23:.*]] = llvm.insertvalue %[[VAL_22]], %[[VAL_4]][0] : !llvm.array<2 x vector<3xf64>>
+
+// CHECK:           "llvm.intr.powi"({{.*}}, {{.*}}) : (f64, i32) -> f64
+// CHECK:           "llvm.intr.powi"({{.*}}, {{.*}}) : (f64, i32) -> f64
+// CHECK:           "llvm.intr.powi"({{.*}}, {{.*}}) : (f64, i32) -> f64
+  %0 = math.fpowsi %arg0, %arg1 : (vector<2x3xf64>, vector<2x3xi32>) -> vector<2x3xf64>
+  func.return
+}