[Mlir-commits] [mlir] [mlir][MathToLLVM] Fix math.fpowi incorrectly being mapped to llvm.powi (PR #226089)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mattéo Rizza Murgier (matteo-rm)

<details>
<summary>Changes</summary>

`math.fpowi` allows vector exponents, while `llvm.intr.powi` only takes scalar exponents. This led to vector `math.fpowi` ops being lowered to invalid IR.

This PR adds support for constant splat exponents, and leaves other vector cases unconverted.

---
Full diff: https://github.com/llvm/llvm-project/pull/226089.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp (+54-2) 
- (modified) mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir (+48) 


``````````diff
diff --git a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
index 29e1678bdd38a..ca7888a9aec76 100644
--- a/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
+++ b/mlir/lib/Conversion/MathToLLVM/MathToLLVM.cpp
@@ -15,6 +15,7 @@
 #include "mlir/Conversion/LLVMCommon/VectorPattern.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 
@@ -95,8 +96,6 @@ using Log10OpLowering =
 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 =
@@ -406,6 +405,59 @@ struct RsqrtOpLowering
   }
 };
 
+struct FPowIOpLowering
+    : public ConvertOpToLLVMPattern<math::FPowIOp,
+                                    /*FailOnUnsupportedFP=*/true> {
+  using ConvertOpToLLVMPattern<
+      math::FPowIOp, /*FailOnUnsupportedFP=*/true>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(math::FPowIOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    const auto &typeConverter = *this->getTypeConverter();
+    auto llvmOperandType = typeConverter.convertType(op.getLhs().getType());
+    if (!llvmOperandType)
+      return failure();
+
+    auto loc = op.getLoc();
+    Value exponent = adaptor.getRhs();
+    if (isa<VectorType>(op.getRhs().getType())) {
+      SplatElementsAttr splatAttr;
+      if (!matchPattern(op.getRhs(), m_Constant(&splatAttr)))
+        return rewriter.notifyMatchFailure(op, "expected a splat exponent");
+
+      auto exponentType = typeConverter.convertType(
+          getElementTypeOrSelf(op.getRhs().getType()));
+      if (!exponentType)
+        return failure();
+      exponent = LLVM::ConstantOp::create(
+          rewriter, loc, exponentType,
+          rewriter.getIntegerAttr(exponentType,
+                                  splatAttr.getSplatValue<APInt>()));
+    }
+
+    ConvertFastMath<math::FPowIOp, LLVM::PowIOp> attrs(op);
+
+    if (!isa<LLVM::LLVMArrayType>(llvmOperandType)) {
+      rewriter.replaceOpWithNewOp<LLVM::PowIOp>(
+          op, TypeRange{llvmOperandType},
+          ValueRange{adaptor.getLhs(), exponent}, attrs.getProperties(),
+          attrs.getDiscardableAttrs());
+      return success();
+    }
+
+    return LLVM::detail::handleMultidimensionalVectors(
+        op.getOperation(), ValueRange{adaptor.getLhs()}, typeConverter,
+        [&](Type llvm1DVectorTy, ValueRange operands) {
+          return LLVM::PowIOp::create(rewriter, loc, TypeRange{llvm1DVectorTy},
+                                      ValueRange{operands[0], exponent},
+                                      attrs.getProperties(),
+                                      attrs.getDiscardableAttrs());
+        },
+        rewriter);
+  }
+};
+
 struct IsNaNOpLowering
     : public ConvertOpToLLVMPattern<math::IsNaNOp,
                                     /*FailOnUnsupportedFP=*/true> {
diff --git a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
index 3cb902a64f18b..3dd9d74ba7043 100644
--- a/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
+++ b/mlir/test/Conversion/MathToLLVM/math-to-llvm.mlir
@@ -583,6 +583,54 @@ func.func @fpowi(%arg0 : f64, %arg1 : i32) {
   func.return
 }
 
+// -----
+
+// CHECK-LABEL: func @fpowi_splat_vector(
+// CHECK-SAME: vector<4xf32>
+func.func @fpowi_splat_vector(%arg0 : vector<4xf32>) {
+  // CHECK: %[[EXP:.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: llvm.intr.powi(%arg0, %[[EXP]]) : (vector<4xf32>, i32) -> vector<4xf32>
+  %cst = arith.constant dense<2> : vector<4xi32>
+  %0 = math.fpowi %arg0, %cst : vector<4xf32>, vector<4xi32>
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @fpowi_splat_scalable_vector(
+// CHECK-SAME: vector<[4]xf32>
+func.func @fpowi_splat_scalable_vector(%arg0 : vector<[4]xf32>) {
+  // CHECK: %[[EXP:.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: llvm.intr.powi(%arg0, %[[EXP]]) : (vector<[4]xf32>, i32) -> vector<[4]xf32>
+  %cst = arith.constant dense<2> : vector<[4]xi32>
+  %0 = math.fpowi %arg0, %cst : vector<[4]xf32>, vector<[4]xi32>
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @fpowi_multidim_vector(
+// CHECK-SAME: vector<4x4xf32>
+func.func @fpowi_multidim_vector(%arg0 : vector<4x4xf32>) {
+  // CHECK: %[[EXP:.*]] = llvm.mlir.constant(2 : i32) : i32
+  // CHECK: %[[EXTRACT:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.array<4 x vector<4xf32>>
+  // CHECK: %[[POWI:.*]] = llvm.intr.powi(%[[EXTRACT]], %[[EXP]]) : (vector<4xf32>, i32) -> vector<4xf32>
+  // CHECK: %[[INSERT:.*]] = llvm.insertvalue %[[POWI]], %{{.*}}[0] : !llvm.array<4 x vector<4xf32>>
+  %cst = arith.constant dense<2> : vector<4x4xi32>
+  %0 = math.fpowi %arg0, %cst : vector<4x4xf32>, vector<4x4xi32>
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @fpowi_non_splat_vector(
+// CHECK-SAME: vector<4xf32>
+func.func @fpowi_non_splat_vector(%arg0 : vector<4xf32>, %arg1 : vector<4xi32>) {
+  // CHECK: math.fpowi
+  // CHECK-NOT: llvm.intr.powi
+  %0 = math.fpowi %arg0, %arg1 : vector<4xf32>, vector<4xi32>
+  func.return
+}
 
 // -----
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/226089