37 template <
typename ConcreteType>
41 std::optional<APInt> apInt;
42 std::optional<APFloat> apFloat;
44 struct APIntOrFloatArray {
48 using RegionComputationFn =
49 std::function<APIntOrFloat(
const APIntOrFloatArray &)>;
54 controlFn(controlFn) {}
56 LogicalResult matchAndRewrite(LinalgOp linalgOp,
59 if (!linalgOp.hasPureTensorSemantics())
63 if (linalgOp.getNumDpsInits() != 1)
66 auto outputType = dyn_cast<ShapedType>(linalgOp->getResultTypes().front());
69 if (!outputType || !outputType.hasStaticShape())
72 if (!llvm::all_of(linalgOp.getDpsInputs(), [](
Value input) {
73 return isa<ShapedType>(input.getType());
78 auto getOperandElementType = [](
Value value) {
79 return cast<ShapedType>(value.getType()).getElementType();
82 llvm::map_range(linalgOp->getOperands(), getOperandElementType)))
86 auto elementType = outputType.getElementType();
87 if (!elementType.isIntOrFloat())
95 if (!llvm::all_of(linalgOp.getIndexingMapsArray(),
96 [](
AffineMap map) { return map.isPermutation(); }))
99 for (
OpOperand &operand : linalgOp.getDpsInitsMutable()) {
100 if (linalgOp.payloadUsesValueFromOperand(&operand))
105 if (!
static_cast<const ConcreteType *
>(
this)->matchIndexingMaps(linalgOp))
110 RegionComputationFn computeFn =
111 static_cast<const ConcreteType *
>(
this)->getRegionComputeFn(linalgOp);
116 int numInputs = linalgOp.getNumDpsInputs();
118 for (
const auto &en :
llvm::enumerate(linalgOp.getDpsInputOperands())) {
126 for (
OpOperand *operand : linalgOp.getDpsInputOperands()) {
127 if (!controlFn(operand))
132 int64_t numElements = outputType.getNumElements();
139 if (isa<FloatType>(elementType))
140 fpOutputValues.resize(numElements, APFloat(0.f));
142 intOutputValues.resize(numElements);
145 auto getDimPositions = [](
AffineMap map) {
149 dims.push_back(cast<AffineDimExpr>(result).getPosition());
155 for (
int i = 0; i < numInputs; ++i)
156 inputDims.push_back(getDimPositions(linalgOp.getIndexingMapsArray()[i]));
157 auto outputDims = getDimPositions(linalgOp.getIndexingMapsArray().back());
158 auto outputShape = outputType.getShape();
167 uint64_t dstLinearIndex = 0;
171 APIntOrFloatArray computeFnInputs;
173 auto inputShapes = llvm::to_vector<4>(
174 llvm::map_range(linalgOp.getDpsInputs(), [](
Value value) {
175 return cast<ShapedType>(value.getType()).getShape();
181 auto computeRemappedLinearIndex = [&](
int linearIndex) {
182 int totalCount = linearIndex;
183 for (
int dim = loopBounds.size() - 1; dim >= 0; --dim) {
184 indices[dim] = totalCount % loopBounds[dim];
185 totalCount /= loopBounds[dim];
188 for (
int dim = loopBounds.size() - 1; dim >= 0; --dim) {
189 for (
int i = 0; i < numInputs; ++i)
190 srcIndices[i][dim] = indices[inputDims[i][dim]];
191 dstIndices[dim] = indices[outputDims[dim]];
194 dstLinearIndex = dstIndices.front();
195 for (
int i = 0; i < numInputs; ++i)
196 srcLinearIndices[i] = srcIndices[i].front();
198 for (
int dim = 1; dim < outputType.getRank(); ++dim) {
199 dstLinearIndex = dstLinearIndex * outputShape[dim] + dstIndices[dim];
200 for (
int i = 0; i < numInputs; ++i)
201 srcLinearIndices[i] =
202 srcLinearIndices[i] * inputShapes[i][dim] + srcIndices[i][dim];
206 bool isFloat = isa<FloatType>(elementType);
209 for (
int i = 0; i < numInputs; ++i)
210 inFpRanges.push_back(inputValues[i].getValues<APFloat>());
212 computeFnInputs.apFloats.resize(numInputs, APFloat(0.f));
217 for (
int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
218 computeRemappedLinearIndex(linearIndex);
221 for (
int i = 0; i < numInputs; ++i)
222 computeFnInputs.apFloats[i] = inFpRanges[i][srcLinearIndices[i]];
226 fpOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apFloat;
230 for (
int i = 0; i < numInputs; ++i)
231 inIntRanges.push_back(inputValues[i].getValues<APInt>());
233 computeFnInputs.apInts.resize(numInputs);
238 for (
int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
239 computeRemappedLinearIndex(linearIndex);
242 for (
int i = 0; i < numInputs; ++i)
243 computeFnInputs.apInts[i] = inIntRanges[i][srcLinearIndices[i]];
247 intOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apInt;
265 struct FoldConstantTranspose :
public FoldConstantBase<FoldConstantTranspose> {
267 using FoldConstantBase::FoldConstantBase;
269 bool matchIndexingMaps(LinalgOp linalgOp)
const {
271 return linalgOp.getIndexingMapsArray().size() == 2;
274 RegionComputationFn getRegionComputeFn(LinalgOp linalgOp)
const {
277 if (!llvm::hasSingleElement(body))
279 auto yieldOp = dyn_cast<linalg::YieldOp>(body.
getTerminator());
284 for (
Value yieldVal : yieldOp.getValues()) {
285 auto yieldArg = dyn_cast<BlockArgument>(yieldVal);
286 if (!yieldArg || yieldArg.getOwner() != &body)
288 if (yieldArg.getArgNumber() != 0)
293 return [](
const APIntOrFloatArray &inputs) {
294 if (inputs.apFloats.empty())
295 return APIntOrFloat{inputs.apInts.front(), std::nullopt};
296 return APIntOrFloat{std::nullopt, inputs.apFloats.front()};
307 patterns.insert<FoldConstantTranspose>(context, controlFn);
Base type for affine expression.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
ArrayRef< AffineExpr > getResults() const
unsigned getNumResults() const
Block represents an ordered list of Operations.
Operation * getTerminator()
Get the terminator operation of this block.
An attribute that represents a reference to a dense vector or tensor object.
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
MLIRContext is the top-level object for a collection of MLIR operations.
This class represents an operand of an operation.
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...
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...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
std::function< bool(OpOperand *fusedOperand)> ControlFusionFn
Function type which is used to control when to stop fusion.
void populateConstantFoldLinalgOperations(RewritePatternSet &patterns, const ControlFusionFn &controlFn)
Patterns to constant fold Linalg operations.
Include the generated interface declarations.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
const FrozenRewritePatternSet & patterns
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
detail::constant_op_matcher m_Constant()
Matches a constant foldable operation.
OpInterfaceRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting a...