24 #include "llvm/ADT/APFloat.h"
25 #include "llvm/ADT/FloatingPointMode.h"
26 #include "llvm/ADT/SmallVector.h"
40 template <
class SrcValType,
class TargetValType,
class TargetType>
43 const std::function<TargetValType(
const SrcValType &)> &toApply,
44 TargetType targetType) {
49 for (
auto val : toTransform.
getValues<SrcValType>()) {
50 auto transformedVal = toApply(val);
51 transformedValues.push_back(transformedVal);
55 auto inShape = toTransform.
getType();
56 auto outTy = inShape.cloneWith({}, targetType);
63 const std::function<APFloat(
const APFloat &)> &toApply,
64 FloatType targetType);
69 if (isa<FloatType>(toCheck.getType().getElementType())) {
73 "Unexpected input tensor type: the "
74 "TOSA spec only allows floats");
88 "Non-const or non-dense input tensor");
93 if (isa<ConstOp>(toCheck.getDefiningOp())) {
98 "The reciprocal can only be folded if "
99 "it operates on a TOSA constant");
106 auto floatCheck = notifyIfNotFloat(toCheck, location, rewriter);
107 if (failed(floatCheck)) {
110 return notifyIfNoTosaDenseConstantTensor(toCheck, location, rewriter);
122 assert(unaryOp->getNumOperands() == 1);
123 auto inputOp = unaryOp->getOperand(0);
126 if (isa<SplatElementsAttr>(values)) {
132 return inputOp.hasOneUse();
135 template <
typename RangeType>
137 ShapedType outputType,
139 using ElementType = std::decay_t<decltype(*std::begin(data))>;
141 assert(inputType.getElementType() == outputType.getElementType());
143 if (inputType.getNumElements() == 0)
146 auto inputShape = inputType.getShape();
154 auto initialValue = *std::begin(data);
159 auto srcLinearIndex = it.index();
161 uint64_t dstLinearIndex = 0;
162 for (int64_t dim = inputShape.size() - 1; dim >= 0; --dim) {
164 auto sourceIndexForDim = srcLinearIndex % inputShape[dim];
165 srcLinearIndex /= inputShape[dim];
170 outputStrides[invertedPermValues[dim]] * sourceIndexForDim;
173 outputValues[dstLinearIndex] = it.value();
181 template <
typename T>
182 std::optional<ArrayRef<T>> tryGetDenseResourceValues(ElementsAttr attr) {
183 if (
auto denseResource = dyn_cast<DenseResourceElementsAttr>(attr)) {
190 bool isSplat =
false;
196 return blob->template getDataAs<T>();
207 ShapedType outputType,
210 if (
auto data = attr.tryGetValues<
bool>())
211 return transposeType(*data, inputType, outputType, permValues);
213 if (
auto data = attr.tryGetValues<int8_t>())
214 return transposeType(*data, inputType, outputType, permValues);
216 if (
auto data = attr.tryGetValues<int16_t>())
217 return transposeType(*data, inputType, outputType, permValues);
219 if (
auto data = attr.tryGetValues<int32_t>())
220 return transposeType(*data, inputType, outputType, permValues);
222 if (
auto data = attr.tryGetValues<int64_t>())
223 return transposeType(*data, inputType, outputType, permValues);
225 if (
auto data = attr.tryGetValues<
float>())
226 return transposeType(*data, inputType, outputType, permValues);
228 if (
auto data = attr.tryGetValues<APFloat>())
229 return transposeType(*data, inputType, outputType, permValues);
232 if (isa<DenseResourceElementsAttr>(attr)) {
233 auto elementTy = attr.getElementType();
235 if (
auto data = tryGetDenseResourceValues<bool>(attr);
236 data && elementTy.isInteger(1))
237 return transposeType(*data, inputType, outputType, permValues);
239 if (
auto data = tryGetDenseResourceValues<int8_t>(attr);
240 data && elementTy.isInteger(8))
241 return transposeType(*data, inputType, outputType, permValues);
243 if (
auto data = tryGetDenseResourceValues<int16_t>(attr);
244 data && elementTy.isInteger(16))
245 return transposeType(*data, inputType, outputType, permValues);
247 if (
auto data = tryGetDenseResourceValues<int32_t>(attr);
248 data && elementTy.isInteger(32))
249 return transposeType(*data, inputType, outputType, permValues);
251 if (
auto data = tryGetDenseResourceValues<int64_t>(attr);
252 data && elementTy.isInteger(64))
253 return transposeType(*data, inputType, outputType, permValues);
255 if (
auto data = tryGetDenseResourceValues<float>(attr);
256 data && elementTy.isF32())
257 return transposeType(*data, inputType, outputType, permValues);
263 struct TosaFoldConstantTranspose :
public OpRewritePattern<tosa::TransposeOp> {
266 LogicalResult matchAndRewrite(tosa::TransposeOp op,
268 auto outputType = cast<ShapedType>(op.getType());
270 if (!outputType.getElementType().isIntOrIndexOrFloat())
273 ElementsAttr inputValues;
277 if (!llvm::hasSingleElement(op.getInput1().getDefiningOp()->getUsers()))
280 auto permValues = llvm::map_to_vector(
281 op.getPerms(), [](
const int32_t v) { return static_cast<int64_t>(v); });
283 auto inputType = cast<ShapedType>(op.getInput1().getType());
285 auto resultAttr =
transpose(inputValues, inputType, outputType, permValues);
288 op,
"unsupported attribute or element type");
300 LogicalResult matchAndRewrite(ReciprocalOp recip,
302 auto inputTensor = recip.getInput1();
306 notifyIfNotConstantFloatTosaTensor(inputTensor, recip, rewriter);
307 if (failed(preCondCheck)) {
316 if (!constantUnaryOpShouldBeFolded(recip, inputValues)) {
318 recip,
"Currently, reciprocals will only be folded if the input "
319 "tensor has a single user");
323 auto newTensor = applyElementWise<APFloat, APFloat, FloatType>(
324 inputValues, &ReciprocalOp::calcOneElement,
338 int64_t remaining = index;
340 for (int64_t i = tensorShape.size() - 1; i >= 0; --i) {
341 position[i] = remaining % tensorShape[i];
342 remaining /= tensorShape[i];
353 int64_t multiplierTmp = 1;
354 for (int64_t i = position.size() - 1; i >= 0; --i) {
355 index += position[i] * multiplierTmp;
356 multiplierTmp *= tensorShape[i];
361 template <
typename OperationType>
362 llvm::APInt calculateReducedValue(
const mlir::ElementsAttr &oldTensorAttr,
364 int64_t reductionAxis,
365 int64_t reductionIndex) {
368 newShape[reductionAxis] = 1;
371 getPositionFromIndex(reductionIndex, newShape);
372 auto oldTensor = oldTensorAttr.getValues<llvm::APInt>();
374 position[reductionAxis] = 0;
375 int64_t indexAtOldTensor = getIndexFromPosition(position, oldShape);
376 llvm::APInt reducedValue = oldTensor[indexAtOldTensor];
378 for (int64_t reductionAxisVal = 1; reductionAxisVal < oldShape[reductionAxis];
379 ++reductionAxisVal) {
381 int64_t stride = std::accumulate(oldShape.begin() + reductionAxis + 1,
382 oldShape.end(), 1, std::multiplies<int>());
383 int64_t index = indexAtOldTensor + stride * reductionAxisVal;
385 OperationType::calcOneElement(reducedValue, oldTensor[index]);
390 template <
typename OperationType>
391 struct ReduceConstantOptimization :
public OpRewritePattern<OperationType> {
394 bool aggressiveReduceConstant)
396 aggressiveReduceConstant(aggressiveReduceConstant) {}
400 LogicalResult matchAndRewrite(OperationType op,
402 Value inputOp = op.getInput();
407 op,
"reduce input must be const operation");
409 if (!inputOp.
hasOneUse() && !this->aggressiveReduceConstant)
411 op,
"input operation has more than one user");
413 auto resultType = cast<ShapedType>(op.getOutput().getType());
415 if (!resultType.hasStaticShape())
418 auto reductionAxis = op.getAxis();
419 const auto denseElementsAttr = constOp.getValues();
420 const auto shapedOldElementsValues =
421 cast<ShapedType>(denseElementsAttr.getType());
423 if (!llvm::isa<IntegerType>(shapedOldElementsValues.getElementType()))
425 op,
"reduce input currently supported with integer type");
427 auto oldShape = shapedOldElementsValues.getShape();
428 auto newShape = resultType.getShape();
430 auto newNumOfElements = std::accumulate(newShape.begin(), newShape.end(), 1,
431 std::multiplies<int>());
434 for (int64_t reductionIndex = 0; reductionIndex < newNumOfElements;
438 newReducedTensor[reductionIndex] = calculateReducedValue<OperationType>(
439 denseElementsAttr, oldShape, reductionAxis, reductionIndex);
442 auto rankedTensorType = cast<RankedTensorType>(resultType);
448 const bool aggressiveReduceConstant;
455 bool aggressiveReduceConstant) {
456 patterns.add<ReduceConstantOptimization<ReduceAllOp>>(
457 ctx, aggressiveReduceConstant);
458 patterns.add<ReduceConstantOptimization<ReduceAnyOp>>(
459 ctx, aggressiveReduceConstant);
460 patterns.add<ReduceConstantOptimization<ReduceMaxOp>>(
461 ctx, aggressiveReduceConstant);
462 patterns.add<ReduceConstantOptimization<ReduceMinOp>>(
463 ctx, aggressiveReduceConstant);
464 patterns.add<ReduceConstantOptimization<ReduceProductOp>>(
465 ctx, aggressiveReduceConstant);
466 patterns.add<ReduceConstantOptimization<ReduceSumOp>>(
467 ctx, aggressiveReduceConstant);
472 patterns.add<TosaFoldConstantTranspose>(ctx);
477 patterns.add<TosaFoldConstantReciprocal>(ctx);
This class represents a processed binary blob of data.
ArrayRef< char > getData() const
Return the raw underlying data of this blob.
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.
static bool isValidRawBuffer(ShapedType type, ArrayRef< char > rawBuffer, bool &detectedSplat)
Returns true if the given buffer is a valid raw buffer for the given type.
MLIRContext is the top-level object for a collection of MLIR operations.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
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)
const FrozenRewritePatternSet & patterns
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...