35 LogicalResult matchAndRewrite(math::PowFOp op,
41 PowFStrengthReduction::matchAndRewrite(math::PowFOp op,
44 Value x = op.getLhs();
46 FloatAttr scalarExponent;
53 auto isExponentValue = [&](
double value) ->
bool {
55 return scalarExponent.getValue().isExactlyValue(value);
57 if (isVector && vectorExponent.isSplat())
58 return vectorExponent.getSplatValue<FloatAttr>()
60 .isExactlyValue(value);
67 if (
auto vec = dyn_cast<VectorType>(op.getType()))
68 return vector::BroadcastOp::create(rewriter, op.getLoc(), vec, value);
73 if (isExponentValue(1.0)) {
79 if (isExponentValue(2.0)) {
85 if (isExponentValue(3.0)) {
87 arith::MulFOp::create(rewriter, op.getLoc(),
ValueRange({x, x}));
93 if (isExponentValue(-1.0)) {
94 Value one = arith::ConstantOp::create(
102 if (isExponentValue(0.5)) {
108 if (isExponentValue(-0.5)) {
114 if (isExponentValue(0.75)) {
115 Value powHalf = math::SqrtOp::create(rewriter, op.getLoc(), x);
116 Value powQuarter = math::SqrtOp::create(rewriter, op.getLoc(), powHalf);
130 template <
typename PowIOpTy,
typename DivOpTy,
typename MulOpTy>
133 unsigned exponentThreshold;
136 PowIStrengthReduction(
MLIRContext *context,
unsigned exponentThreshold = 3,
140 exponentThreshold(exponentThreshold) {}
142 LogicalResult matchAndRewrite(PowIOpTy op,
147 template <
typename PowIOpTy,
typename DivOpTy,
typename MulOpTy>
149 PowIStrengthReduction<PowIOpTy, DivOpTy, MulOpTy>::matchAndRewrite(
152 Value base = op.getLhs();
154 IntegerAttr scalarExponent;
161 int64_t exponentValue = 0;
163 exponentValue = scalarExponent.getInt();
164 else if (isVector && vectorExponent.isSplat())
165 exponentValue = vectorExponent.getSplatValue<IntegerAttr>().getInt();
170 auto bcast = [&loc, &op, &rewriter](
Value value) ->
Value {
171 if (
auto vec = dyn_cast<VectorType>(op.getType()))
172 return vector::BroadcastOp::create(rewriter, loc, vec, value);
178 if constexpr (std::is_same_v<PowIOpTy, math::FPowIOp>)
179 one = arith::ConstantOp::create(rewriter, loc,
182 one = arith::ConstantOp::create(rewriter, loc,
186 if (exponentValue == 0) {
191 bool exponentIsNegative =
false;
192 if (exponentValue < 0) {
193 exponentIsNegative =
true;
198 if (exponentValue > exponentThreshold)
211 for (
unsigned i = 1; i < exponentValue; ++i)
212 result = MulOpTy::create(rewriter, loc, result, base);
216 if (exponentIsNegative)
217 result = DivOpTy::create(rewriter, loc, bcast(one), result);
228 .add<PowFStrengthReduction,
229 PowIStrengthReduction<math::IPowIOp, arith::DivSIOp, arith::MulIOp>,
230 PowIStrengthReduction<math::FPowIOp, arith::DivFOp, arith::MulFOp>>(
IntegerAttr getIntegerAttr(Type type, int64_t value)
FloatAttr getFloatAttr(Type type, double value)
An attribute that represents a reference to a dense float vector or tensor object.
An attribute that represents a reference to a dense integer vector or tensor object.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
This class represents the benefit of a pattern match in a unitless scheme that ranges from 0 (very li...
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
virtual void replaceOp(Operation *op, ValueRange newValues)
Replace the results of the given (original) operation with the specified list of values (replacements...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns)
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
const FrozenRewritePatternSet & patterns
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
OpRewritePattern(MLIRContext *context, PatternBenefit benefit=1, ArrayRef< StringRef > generatedNames={})
Patterns must specify the root operation name they match against, and can also specify the benefit of...