23 #include "llvm/ADT/APFloat.h"
24 #include "llvm/ADT/FloatingPointMode.h"
25 #include "llvm/ADT/SmallVector.h"
39 template <
class SrcValType,
class TargetValType,
class TargetType>
42 const std::function<TargetValType(
const SrcValType &)> &toApply,
43 TargetType targetType) {
48 for (
auto val : toTransform.
getValues<SrcValType>()) {
49 auto transformedVal = toApply(val);
50 transformedValues.push_back(transformedVal);
54 auto inShape = toTransform.
getType();
55 auto outTy = inShape.cloneWith({}, targetType);
62 const std::function<APFloat(
const APFloat &)> &toApply,
68 if (isa<FloatType>(toCheck.getType().getElementType())) {
72 "Unexpected input tensor type: the "
73 "TOSA spec only allows floats");
87 "Non-const or non-dense input tensor");
92 if (isa<ConstOp>(toCheck.getDefiningOp())) {
97 "The reciprocal can only be folded if "
98 "it operates on a TOSA constant");
105 auto floatCheck = notifyIfNotFloat(toCheck, location, rewriter);
106 if (failed(floatCheck)) {
109 return notifyIfNoTosaDenseConstantTensor(toCheck, location, rewriter);
121 assert(unaryOp->getNumOperands() == 1);
122 auto inputOp = unaryOp->getOperand(0);
125 if (isa<SplatElementsAttr>(values)) {
131 return inputOp.hasOneUse();
134 template <
typename RangeType>
136 ShapedType outputType,
138 using ElementType = std::decay_t<decltype(*std::begin(data))>;
140 assert(inputType.getElementType() == outputType.getElementType());
142 if (inputType.getNumElements() == 0)
145 auto inputShape = inputType.getShape();
153 auto initialValue = *std::begin(data);
158 auto srcLinearIndex = it.index();
160 uint64_t dstLinearIndex = 0;
161 for (int64_t dim = inputShape.size() - 1; dim >= 0; --dim) {
163 auto sourceIndexForDim = srcLinearIndex % inputShape[dim];
164 srcLinearIndex /= inputShape[dim];
169 outputStrides[invertedPermValues[dim]] * sourceIndexForDim;
172 outputValues[dstLinearIndex] = it.value();
184 ShapedType outputType,
186 if (
auto data = attr.tryGetValues<
bool>())
187 return transposeType(*data, inputType, outputType, permValues);
189 if (
auto data = attr.tryGetValues<int8_t>())
190 return transposeType(*data, inputType, outputType, permValues);
192 if (
auto data = attr.tryGetValues<int16_t>())
193 return transposeType(*data, inputType, outputType, permValues);
195 if (
auto data = attr.tryGetValues<int32_t>())
196 return transposeType(*data, inputType, outputType, permValues);
198 if (
auto data = attr.tryGetValues<int64_t>())
199 return transposeType(*data, inputType, outputType, permValues);
201 if (
auto data = attr.tryGetValues<
float>())
202 return transposeType(*data, inputType, outputType, permValues);
204 if (
auto data = attr.tryGetValues<APFloat>())
205 return transposeType(*data, inputType, outputType, permValues);
210 struct TosaFoldConstantTranspose :
public OpRewritePattern<tosa::TransposeOp> {
213 LogicalResult matchAndRewrite(tosa::TransposeOp op,
215 auto outputType = cast<ShapedType>(op.getType());
217 if (!outputType.getElementType().isIntOrIndexOrFloat())
220 ElementsAttr inputValues;
224 if (!llvm::hasSingleElement(op.getInput1().getDefiningOp()->
getUsers()))
230 auto permValues = llvm::map_to_vector(
232 permAttr.getValues<APInt>(),
233 [](
const APInt &val) { return val.getSExtValue(); });
235 auto inputType = cast<ShapedType>(op.getInput1().getType());
237 auto resultAttr =
transpose(inputValues, inputType, outputType, permValues);
240 op,
"unsupported attribute or element type");
252 LogicalResult matchAndRewrite(ReciprocalOp recip,
254 auto inputTensor = recip.getInput1();
258 notifyIfNotConstantFloatTosaTensor(inputTensor, recip, rewriter);
259 if (failed(preCondCheck)) {
268 if (!constantUnaryOpShouldBeFolded(recip, inputValues)) {
270 recip,
"Currently, reciprocals will only be folded if the input "
271 "tensor has a single user");
275 auto newTensor = applyElementWise<APFloat, APFloat, FloatType>(
276 inputValues, &ReciprocalOp::calcOneElement,
290 int64_t remaining = index;
292 for (int64_t i = tensorShape.size() - 1; i >= 0; --i) {
293 position[i] = remaining % tensorShape[i];
294 remaining /= tensorShape[i];
305 int64_t multiplierTmp = 1;
306 for (int64_t i = position.size() - 1; i >= 0; --i) {
307 index += position[i] * multiplierTmp;
308 multiplierTmp *= tensorShape[i];
313 template <
typename OperationType>
314 llvm::APInt calculateReducedValue(
const mlir::ElementsAttr &oldTensorAttr,
316 int64_t reductionAxis,
317 int64_t reductionIndex) {
320 newShape[reductionAxis] = 1;
323 getPositionFromIndex(reductionIndex, newShape);
324 auto oldTensor = oldTensorAttr.getValues<llvm::APInt>();
326 position[reductionAxis] = 0;
327 int64_t indexAtOldTensor = getIndexFromPosition(position, oldShape);
328 llvm::APInt reducedValue = oldTensor[indexAtOldTensor];
330 for (int64_t reductionAxisVal = 1; reductionAxisVal < oldShape[reductionAxis];
331 ++reductionAxisVal) {
333 int64_t stride = std::accumulate(oldShape.begin() + reductionAxis + 1,
334 oldShape.end(), 1, std::multiplies<int>());
335 int64_t index = indexAtOldTensor + stride * reductionAxisVal;
337 OperationType::calcOneElement(reducedValue, oldTensor[index]);
342 template <
typename OperationType>
343 struct ReduceConstantOptimization :
public OpRewritePattern<OperationType> {
346 bool aggressiveReduceConstant)
348 aggressiveReduceConstant(aggressiveReduceConstant) {}
352 LogicalResult matchAndRewrite(OperationType op,
354 Value inputOp = op.getInput();
359 op,
"reduce input must be const operation");
361 if (!inputOp.
hasOneUse() && !this->aggressiveReduceConstant)
363 op,
"input operation has more than one user");
365 auto resultType = cast<ShapedType>(op.getOutput().getType());
367 if (!resultType.hasStaticShape())
370 auto reductionAxis = op.getAxis();
371 const auto denseElementsAttr = constOp.getValue();
372 const auto shapedOldElementsValues =
373 cast<ShapedType>(denseElementsAttr.getType());
375 if (!llvm::isa<IntegerType>(shapedOldElementsValues.getElementType()))
377 op,
"reduce input currently supported with integer type");
379 auto oldShape = shapedOldElementsValues.getShape();
380 auto newShape = resultType.getShape();
382 auto newNumOfElements = std::accumulate(newShape.begin(), newShape.end(), 1,
383 std::multiplies<int>());
386 for (int64_t reductionIndex = 0; reductionIndex < newNumOfElements;
390 newReducedTensor[reductionIndex] = calculateReducedValue<OperationType>(
391 denseElementsAttr, oldShape, reductionAxis, reductionIndex);
394 auto rankedTensorType = cast<RankedTensorType>(resultType);
400 const bool aggressiveReduceConstant;
407 bool aggressiveReduceConstant) {
408 patterns.
add<ReduceConstantOptimization<ReduceAllOp>>(
409 ctx, aggressiveReduceConstant);
410 patterns.
add<ReduceConstantOptimization<ReduceAnyOp>>(
411 ctx, aggressiveReduceConstant);
412 patterns.
add<ReduceConstantOptimization<ReduceMaxOp>>(
413 ctx, aggressiveReduceConstant);
414 patterns.
add<ReduceConstantOptimization<ReduceMinOp>>(
415 ctx, aggressiveReduceConstant);
416 patterns.
add<ReduceConstantOptimization<ReduceProdOp>>(
417 ctx, aggressiveReduceConstant);
418 patterns.
add<ReduceConstantOptimization<ReduceSumOp>>(
419 ctx, aggressiveReduceConstant);
424 patterns.
add<TosaFoldConstantTranspose>(ctx);
429 patterns.
add<TosaFoldConstantReciprocal>(ctx);
An attribute that represents a reference to a dense vector or tensor object.
auto getValues() const
Return the held element values as a range of the given type.
int64_t getNumElements() const
Returns the number of elements held by this attribute.
Type getElementType() const
Return the element type of this DenseElementsAttr.
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
ShapedType getType() const
Return the type of this ElementsAttr, guaranteed to be a vector or tensor with static shape.
An attribute that represents a reference to a dense integer vector or tensor object.
MLIRContext is the top-level object for a collection of MLIR operations.
user_range getUsers()
Returns a range of all users.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
bool hasOneUse() const
Returns true if this value has exactly one use.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
void populateTosaConstantReduction(MLIRContext *ctx, RewritePatternSet &patterns, bool aggressiveReduceConstant)
void populateTosaFoldConstantReciprocalPatterns(MLIRContext *ctx, RewritePatternSet &patterns)
void populateTosaFoldConstantTransposePatterns(MLIRContext *ctx, RewritePatternSet &patterns)
static void transpose(llvm::ArrayRef< int64_t > trans, SmallVector< int64_t > &shape)
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
std::conditional_t< std::is_same_v< Ty, mlir::Type >, mlir::Value, detail::TypedValue< Ty > > TypedValue
If Ty is mlir::Type this will select Value instead of having a wrapper around it.
SmallVector< int64_t > computeStrides(ArrayRef< int64_t > sizes)
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
SmallVector< int64_t > invertPermutationVector(ArrayRef< int64_t > permutation)
Helper method to apply to inverse a permutation.
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...