[Mlir-commits] [mlir] [MLIR][NVVM] Pass an i32 exponent to libdevice when lowering math.fpowi (PR #228852)
Dhairyashil R G
llvmlistbot at llvm.org
Sun Oct 4 00:49:11 PDT 2026
https://github.com/dhairyashilRG created https://github.com/llvm/llvm-project/pull/228852
libdevice's `__nv_powif` and `__nv_powi` take an `int` exponent, but `convert-math-to-nvvm` passed the exponent of `math.fpowi` through with whatever type it had. An i64 exponent gave a call to `@__nv_powif(f32, i64)`.
Flang hits this on NVPTX with `real ** integer(8)`. The call has the wrong type, and a file that also uses an `integer(4)` exponent fails to compile with "Lowering to LLVM IR failed".
This patch sign-extends an exponent narrower than 32 bits, and leaves `math.fpowi` alone when the exponent is wider or `index`, so a later pass can lower it (MathToFuncs in Flang). i32 exponents are lowered as before.
Assisted-by: Claude
>From 88a61c1fd416739b9b8e305b17c7a161b3604110 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Sun, 4 Oct 2026 12:30:50 +0530
Subject: [PATCH] [MLIR][NVVM] Pass an i32 exponent to libdevice when lowering
math.fpowi
libdevice declares __nv_powif(float, int) and __nv_powi(double, int).
convert-math-to-nvvm built the callee from the exponent type as it was,
so an i64 exponent declared and called __nv_powif(f32, i64), and a
module using two exponent widths failed to verify.
Sign-extend an exponent narrower than 32 bits, and do not lower
math.fpowi with a wider or index exponent.
Assisted-by: Claude
---
mlir/lib/Conversion/MathToNVVM/MathToNVVM.cpp | 59 +++++++++++++++----
.../Conversion/GPUToNVVM/gpu-to-nvvm.mlir | 16 +++++
.../Conversion/MathToNVVM/math-to-nvvm.mlir | 17 ++++++
3 files changed, 82 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Conversion/MathToNVVM/MathToNVVM.cpp b/mlir/lib/Conversion/MathToNVVM/MathToNVVM.cpp
index e6343b77cea71..83425a60c5126 100644
--- a/mlir/lib/Conversion/MathToNVVM/MathToNVVM.cpp
+++ b/mlir/lib/Conversion/MathToNVVM/MathToNVVM.cpp
@@ -49,16 +49,54 @@ static void populateIntOpPatterns(const LLVMTypeConverter &converter,
benefit);
}
-template <typename OpTy>
-static void populateFloatIntOpPatterns(const LLVMTypeConverter &converter,
- RewritePatternSet &patterns,
- PatternBenefit benefit,
- StringRef f32Func, StringRef f64Func) {
- patterns.add<ScalarizeVectorOpLowering<OpTy>>(converter, benefit);
- patterns.add<OpToFuncCallLowering<OpTy>>(converter, f32Func, f64Func, "", "",
- /*i32Func=*/"", benefit);
+/// libdevice's __nv_powi and __nv_powif take an i32 exponent, so math.fpowi
+/// with a wider or index exponent is not lowered.
+static bool exponentFitsInI32(math::FPowIOp op) {
+ auto type =
+ dyn_cast<IntegerType>(getElementTypeOrSelf(op.getRhs().getType()));
+ return type && type.getWidth() <= 32;
}
+namespace {
+struct FPowIScalarizeLowering
+ : public ScalarizeVectorOpLowering<math::FPowIOp> {
+ using ScalarizeVectorOpLowering::ScalarizeVectorOpLowering;
+
+ LogicalResult
+ matchAndRewrite(math::FPowIOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (!exponentFitsInI32(op))
+ return rewriter.notifyMatchFailure(op, "exponent does not fit in i32");
+ return ScalarizeVectorOpLowering::matchAndRewrite(op, adaptor, rewriter);
+ }
+};
+
+struct FPowIOpLowering : public OpToFuncCallLowering<math::FPowIOp> {
+ using OpToFuncCallLowering::OpToFuncCallLowering;
+
+ LogicalResult
+ matchAndRewrite(math::FPowIOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (!exponentFitsInI32(op))
+ return rewriter.notifyMatchFailure(op, "exponent does not fit in i32");
+ // Return before creating anything in the cases OpToFuncCallLowering
+ // does not handle.
+ if (!op->getParentOfType<FunctionOpInterface>() ||
+ !isa<Float16Type, BFloat16Type, Float32Type, Float64Type>(
+ adaptor.getLhs().getType()))
+ return rewriter.notifyMatchFailure(op, "cannot be lowered to libdevice");
+ // Sign-extend a narrower exponent.
+ Value exponent = adaptor.getRhs();
+ if (!exponent.getType().isInteger(32))
+ exponent = LLVM::SExtOp::create(rewriter, op.getLoc(),
+ rewriter.getI32Type(), exponent);
+ Value operands[] = {adaptor.getLhs(), exponent};
+ return OpToFuncCallLowering::matchAndRewrite(
+ op, OpAdaptor(ValueRange(operands), op), rewriter);
+ }
+};
+} // namespace
+
// Custom pattern for sincos since it returns two values
struct SincosOpLowering : public ConvertOpToLLVMPattern<math::SincosOp> {
using ConvertOpToLLVMPattern<math::SincosOp>::ConvertOpToLLVMPattern;
@@ -223,8 +261,9 @@ void mlir::populateLibDeviceConversionPatterns(
"__nv_log2", "__nv_fast_log2f");
populateOpPatterns<math::PowFOp>(converter, patterns, benefit, "__nv_powf",
"__nv_pow", "__nv_fast_powf");
- populateFloatIntOpPatterns<math::FPowIOp>(converter, patterns, benefit,
- "__nv_powif", "__nv_powi");
+ patterns.add<FPowIScalarizeLowering>(converter, benefit);
+ patterns.add<FPowIOpLowering>(converter, "__nv_powif", "__nv_powi", "", "",
+ /*i32Func=*/"", benefit);
populateOpPatterns<math::RoundOp>(converter, patterns, benefit, "__nv_roundf",
"__nv_round");
populateOpPatterns<math::RoundEvenOp>(converter, patterns, benefit,
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
index 29aaa1f26bce4..4f0bba637eecc 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir
@@ -1109,6 +1109,22 @@ gpu.module @test_module_53 {
// CHECK: llvm.call @__nv_powi(%{{.*}}, %{{.*}}) : (f64, i32) -> f64
func.return %result32, %result64 : f32, f64
}
+
+ // A narrower exponent is sign-extended, also in a vector.
+ // CHECK-LABEL: func @gpu_powi_narrow
+ func.func @gpu_powi_narrow(%arg_f64 : vector<2xf64>, %arg_i16 : vector<2xi16>) -> vector<2xf64> {
+ // CHECK-COUNT-2: llvm.call @__nv_powi(%{{.*}}, %{{.*}}) : (f64, i32) -> f64
+ %result = math.fpowi %arg_f64, %arg_i16 : vector<2xf64>, vector<2xi16>
+ func.return %result : vector<2xf64>
+ }
+
+ // A vector with an exponent wider than i32 is not scalarized.
+ // CHECK-LABEL: func @gpu_powi_wide
+ func.func @gpu_powi_wide(%arg_f32 : vector<2xf32>, %arg_i64 : vector<2xi64>) -> vector<2xf32> {
+ // CHECK: math.fpowi %{{.*}}, %{{.*}} : vector<2xf32>, vector<2xi64>
+ %result = math.fpowi %arg_f32, %arg_i64 : vector<2xf32>, vector<2xi64>
+ func.return %result : vector<2xf32>
+ }
}
gpu.module @test_module_54 {
diff --git a/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir b/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
index ac14ab15fd80c..aecf4448a900a 100644
--- a/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
+++ b/mlir/test/Conversion/MathToNVVM/math-to-nvvm.mlir
@@ -27,3 +27,20 @@ func.func @fpclass_scalar(%arg: f32) -> (i1, i1, i1) {
%nan = math.isnan %arg : f32
return %inf, %finite, %nan : i1, i1, i1
}
+
+// libdevice's __nv_powif takes an i32 exponent. A narrower exponent is
+// sign-extended; a wider or index exponent is not lowered.
+
+// CHECK-LABEL: func.func @fpowi(
+func.func @fpowi(%x: f32, %e16: i16, %e32: i32, %e64: i64, %idx: index) -> (f32, f32, f32, f32) {
+ // CHECK: %[[EXT:.*]] = llvm.sext %{{.*}} : i16 to i32
+ // CHECK: llvm.call @__nv_powif(%{{.*}}, %[[EXT]]) : (f32, i32) -> f32
+ %a = math.fpowi %x, %e16 : f32, i16
+ // CHECK: llvm.call @__nv_powif(%{{.*}}, %{{.*}}) : (f32, i32) -> f32
+ %b = math.fpowi %x, %e32 : f32, i32
+ // CHECK: math.fpowi %{{.*}}, %{{.*}} : f32, i64
+ %c = math.fpowi %x, %e64 : f32, i64
+ // CHECK: math.fpowi %{{.*}}, %{{.*}} : f32, index
+ %d = math.fpowi %x, %idx : f32, index
+ return %a, %b, %c, %d : f32, f32, f32, f32
+}
More information about the Mlir-commits
mailing list