caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][math] powi with negative exponent should invert at the end (#135735)

Browse Source 
Previously, an FPowI operation would invert the base *before* performing
a sequence of multiplications, but this led to discrepancies between
LLVM pow intrinsic folding and that coming from the math dialect.

See compiler-rt's version, which does the inverse at the end of the
calculation: compiler-rt/lib/builtins/powidf2.c
This commit is contained in:
Asher Mancinelli
2025-04-15 08:02:11 -07:00
committed by GitHub
parent f215a61891
commit 9ab2dea61e
2 changed files with 29 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
mlir/lib/Dialect/Math/Transforms/AlgebraicSimplification.cpp
View File

@@ -197,11 +197,6 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
if (exponentValue > exponentThreshold)
return failure();
// Inverse the base for negative exponent, i.e. for
// `[fi]powi(x, negative_exponent)` set `x` to `1 / x`.
if (exponentIsNegative)
base = rewriter.create<DivOpTy>(loc, bcast(one), base);
Value result = base;
// Transform to naive sequence of multiplications:
// * For positive exponent case replace:
@@ -215,6 +210,11 @@ PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
for (unsigned i = 1; i < exponentValue; ++i)
result = rewriter.create<MulOpTy>(loc, result, base);
// Inverse the base for negative exponent, i.e. for
// `[fi]powi(x, negative_exponent)` set `x` to `1 / x`.
if (exponentIsNegative)
result = rewriter.create<DivOpTy>(loc, bcast(one), result);
rewriter.replaceOp(op, result);
return success();
}

48
mlir/test/Dialect/Math/algebraic-simplification.mlir
View File

@@ -135,11 +135,11 @@ func.func @ipowi_exp_two(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi32
// CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1> : vector<4xi32>
// CHECK: %[[SCALAR0:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
// CHECK: %[[VECTOR0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
// CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
// CHECK: %[[SMUL:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
// CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
// CHECK: %[[VMUL:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL]], %[[VMUL]]
// CHECK: %[[SMUL:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
// CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[SMUL]]
// CHECK: %[[VMUL:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
// CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[VMUL]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 2 : i32
%v1 = arith.constant dense <2> : vector<4xi32>
%0 = math.ipowi %arg0, %c1 : i32
@@ -162,13 +162,13 @@ func.func @ipowi_exp_three(%arg0: i32, %arg1: vector<4xi32>) -> (i32, vector<4xi
// CHECK: %[[SCALAR0:.*]] = arith.muli %[[SMUL0]], %[[ARG0]]
// CHECK: %[[VMUL0:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
// CHECK: %[[VECTOR0:.*]] = arith.muli %[[VMUL0]], %[[ARG1]]
// CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[ARG0]]
// CHECK: %[[SMUL1:.*]] = arith.muli %[[SCALAR1]], %[[SCALAR1]]
// CHECK: %[[SMUL2:.*]] = arith.muli %[[SMUL1]], %[[SCALAR1]]
// CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[ARG1]]
// CHECK: %[[VMUL1:.*]] = arith.muli %[[VECTOR1]], %[[VECTOR1]]
// CHECK: %[[VMUL2:.*]] = arith.muli %[[VMUL1]], %[[VECTOR1]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL2]], %[[VMUL2]]
// CHECK: %[[SMUL1:.*]] = arith.muli %[[ARG0]], %[[ARG0]]
// CHECK: %[[SMUL2:.*]] = arith.muli %[[SMUL1]], %[[ARG0]]
// CHECK: %[[SCALAR1:.*]] = arith.divsi %[[CST_S]], %[[SMUL2]]
// CHECK: %[[VMUL1:.*]] = arith.muli %[[ARG1]], %[[ARG1]]
// CHECK: %[[VMUL2:.*]] = arith.muli %[[VMUL1]], %[[ARG1]]
// CHECK: %[[VECTOR1:.*]] = arith.divsi %[[CST_V]], %[[VMUL2]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 3 : i32
%v1 = arith.constant dense <3> : vector<4xi32>
%0 = math.ipowi %arg0, %c1 : i32
@@ -225,11 +225,11 @@ func.func @fpowi_exp_two(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf32
// CHECK-DAG: %[[CST_V:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
// CHECK: %[[SCALAR0:.*]] = arith.mulf %[[ARG0]], %[[ARG0]]
// CHECK: %[[VECTOR0:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
// CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[ARG0]]
// CHECK: %[[SMUL:.*]] = arith.mulf %[[SCALAR1]], %[[SCALAR1]]
// CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[ARG1]]
// CHECK: %[[VMUL:.*]] = arith.mulf %[[VECTOR1]], %[[VECTOR1]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL]], %[[VMUL]]
// CHECK: %[[SMUL:.*]] = arith.mulf %[[ARG0]], %[[ARG0]]
// CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[SMUL]]
// CHECK: %[[VMUL:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
// CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[VMUL]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 2 : i32
%v1 = arith.constant dense <2> : vector<4xi32>
%0 = math.fpowi %arg0, %c1 : f32, i32
@@ -252,13 +252,13 @@ func.func @fpowi_exp_three(%arg0: f32, %arg1: vector<4xf32>) -> (f32, vector<4xf
// CHECK: %[[SCALAR0:.*]] = arith.mulf %[[SMUL0]], %[[ARG0]]
// CHECK: %[[VMUL0:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
// CHECK: %[[VECTOR0:.*]] = arith.mulf %[[VMUL0]], %[[ARG1]]
// CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[ARG0]]
// CHECK: %[[SMUL1:.*]] = arith.mulf %[[SCALAR1]], %[[SCALAR1]]
// CHECK: %[[SMUL2:.*]] = arith.mulf %[[SMUL1]], %[[SCALAR1]]
// CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[ARG1]]
// CHECK: %[[VMUL1:.*]] = arith.mulf %[[VECTOR1]], %[[VECTOR1]]
// CHECK: %[[VMUL2:.*]] = arith.mulf %[[VMUL1]], %[[VECTOR1]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SMUL2]], %[[VMUL2]]
// CHECK: %[[SMUL1:.*]] = arith.mulf %[[ARG0]], %[[ARG0]]
// CHECK: %[[SMUL2:.*]] = arith.mulf %[[SMUL1]], %[[ARG0]]
// CHECK: %[[SCALAR1:.*]] = arith.divf %[[CST_S]], %[[SMUL2]]
// CHECK: %[[VMUL1:.*]] = arith.mulf %[[ARG1]], %[[ARG1]]
// CHECK: %[[VMUL2:.*]] = arith.mulf %[[VMUL1]], %[[ARG1]]
// CHECK: %[[VECTOR1:.*]] = arith.divf %[[CST_V]], %[[VMUL2]]
// CHECK: return %[[SCALAR0]], %[[VECTOR0]], %[[SCALAR1]], %[[VECTOR1]]
%c1 = arith.constant 3 : i32
%v1 = arith.constant dense <3> : vector<4xi32>
%0 = math.fpowi %arg0, %c1 : f32, i32