[mlir][math] powf(a, b) drop support when a < 0
#126338
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -17,8 +17,13 @@ | |
| #include "mlir/Dialect/Vector/IR/VectorOps.h" | ||
| #include "mlir/IR/Builders.h" | ||
| #include "mlir/IR/ImplicitLocOpBuilder.h" | ||
| #include "mlir/IR/Matchers.h" | ||
| #include "mlir/IR/PatternMatch.h" | ||
| #include "mlir/IR/TypeUtilities.h" | ||
| #include "mlir/Transforms/DialectConversion.h" | ||
| #include "llvm/ADT/APFloat.h" | ||
| #include "llvm/Support/LogicalResult.h" | ||
| #include <cmath> | ||
|
|
||
| using namespace mlir; | ||
|
|
||
|
|
@@ -311,40 +316,113 @@ static LogicalResult convertFPowIOp(math::FPowIOp op, | |
| return success(); | ||
| } | ||
|
|
||
| // Convert Powf(float a, float b) for some special cases | ||
| // where b == 1.0, b == 0.0, b == 0.5, b == -0.5, b == -1.0, and b % 2 == 0 | ||
| static LogicalResult convertSpecialPowfOp(math::PowFOp op, | ||
| PatternRewriter &rewriter) { | ||
| ImplicitLocOpBuilder b(op->getLoc(), rewriter); | ||
| Value operandA = op.getOperand(0); | ||
| Value operandB = op.getOperand(1); | ||
| auto baseType = operandB.getType(); | ||
|
|
||
| auto &sem = dyn_cast<mlir::FloatType>(getElementTypeOrSelf(baseType)) | ||
| .getFloatSemantics(); | ||
|
|
||
| auto valueB = APFloat(sem); | ||
| if (!matchPattern(operandB, m_ConstantFloat(&valueB))) { | ||
| // Not a constant, return failure | ||
| return failure(); | ||
| } | ||
| float floatValueB = valueB.convertToFloat(); | ||
|
||
|
|
||
| if (floatValueB == 1.0f) { | ||
| // a^1 -> a | ||
| rewriter.replaceOp(op, operandA); | ||
| return success(); | ||
| } | ||
|
|
||
| if (floatValueB == 0.0) { | ||
| // a^0 -> 1 | ||
| Value one = | ||
| createFloatConst(op->getLoc(), operandA.getType(), 1.0, rewriter); | ||
| rewriter.replaceOp(op, one); | ||
| return success(); | ||
| } | ||
|
|
||
| if (floatValueB == 0.5f) { | ||
| // a^(1/2) -> sqrt(a) | ||
| Value sqrt = b.create<math::SqrtOp>(operandA); | ||
| rewriter.replaceOp(op, sqrt); | ||
| return success(); | ||
| } | ||
|
|
||
| if (floatValueB == -0.5f) { | ||
| // a^(-1/2) -> 1 / sqrt(a) | ||
| Value rsqrt = b.create<math::RsqrtOp>(operandA); | ||
| rewriter.replaceOp(op, rsqrt); | ||
| return success(); | ||
| } | ||
|
|
||
| if (floatValueB == -1.0f) { | ||
| // a^(-1) -> 1 / a | ||
| Value one = | ||
| createFloatConst(op->getLoc(), operandA.getType(), 1.0, rewriter); | ||
| Value div = b.create<arith::DivFOp>(one, operandA); | ||
| rewriter.replaceOp(op, div); | ||
| return success(); | ||
| } | ||
|
|
||
| // Check if the power is an integer | ||
| if (floatValueB != std::floor(floatValueB)) { | ||
|
||
| // We don't handle non-integer powers here, return failure | ||
| return failure(); | ||
| } | ||
|
|
||
| auto sign = std::signbit(floatValueB) ? -1 : 1; | ||
| auto absIntValueB = std::abs(static_cast<int>(floatValueB)); | ||
|
|
||
| auto cstOne = | ||
| createFloatConst(op->getLoc(), operandA.getType(), 1.0, rewriter); | ||
| auto base = operandA; | ||
| if (sign == -1) { | ||
| base = b.create<arith::DivFOp>(cstOne, base); | ||
| } | ||
| auto current = base; | ||
| auto result = cstOne; | ||
| while (absIntValueB > 0) { | ||
| if (absIntValueB & 1) { | ||
| result = b.create<arith::MulFOp>(result, current); | ||
| } | ||
| current = b.create<arith::MulFOp>(current, current); | ||
| absIntValueB >>= 1; | ||
| } | ||
| rewriter.replaceOp(op, result); | ||
| return success(); | ||
| } | ||
|
|
||
| // Converts Powf(float a, float b) (meaning a^b) to exp^(b * ln(a)) | ||
| // Restricting a >= 0 | ||
|
||
| static LogicalResult convertPowfOp(math::PowFOp op, PatternRewriter &rewriter) { | ||
| ImplicitLocOpBuilder b(op->getLoc(), rewriter); | ||
| Value operandA = op.getOperand(0); | ||
| Value operandB = op.getOperand(1); | ||
| Type opType = operandA.getType(); | ||
| Value zero = createFloatConst(op->getLoc(), opType, 0.00, rewriter); | ||
| Value one = createFloatConst(op->getLoc(), opType, 1.00, rewriter); | ||
|
|
||
| Value logA = b.create<math::LogOp>(opType, operandA); | ||
| Value mult = b.create<arith::MulFOp>(opType, operandB, logA); | ||
| Value expResult = b.create<math::ExpOp>(opType, mult); | ||
|
|
||
| // First, we select between the exp value and the adjusted value for odd | ||
| // powers of negatives. Then, we ensure that one is produced if `b` is zero. | ||
| // This corresponds to `libm` behavior, even for `0^0`. Without this check, | ||
| // `exp(0 * ln(0)) = exp(0 *-inf) = exp(-nan) = -nan`. | ||
| Value zeroCheck = | ||
| b.create<arith::CmpFOp>(arith::CmpFPredicate::OEQ, operandB, zero); | ||
| Value finalResult = | ||
| b.create<arith::SelectOp>(op->getLoc(), zeroCheck, one, expResult); | ||
| rewriter.replaceOp(op, finalResult); | ||
| return success(); | ||
| } | ||
|
|
||
|
|
@@ -660,6 +738,7 @@ void mlir::populateExpandExp2FPattern(RewritePatternSet &patterns) { | |
| } | ||
|
|
||
| void mlir::populateExpandPowFPattern(RewritePatternSet &patterns) { | ||
| patterns.add(convertSpecialPowfOp); | ||
|
||
| patterns.add(convertPowfOp); | ||
| } | ||
|
|
||
|
|
||
|
Contributor
There was a problem hiding this comment. At least we need all the special cases are tested in this file.
Contributor
Author
There was a problem hiding this comment. my apology for that tihs PR got verbose, I made appropriate tests and runs well!! |
ita9naiwa marked this conversation as resolved.
Show resolved
Hide resolved
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Here, since
math.powfrequires float arguments (I just checked MathOps.td), the cast really shouldn't ever fail, so I think you can simply usecastinstead ofdyn_cast. You weren't checking for a null return value fromdyn_castanyway.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done!