25 for (
const auto &map :
enumerate(reassociation)) {
27 map.value().getResults().front().cast<
AffineDimExpr>().getPosition();
29 map.value().getResults().back().cast<
AffineDimExpr>().getPosition();
30 for (
auto dim : llvm::seq_inclusive(startPos, endPos)) {
31 expandedDimToCollapsedDim[dim] = map.index();
34 return expandedDimToCollapsedDim;
44 if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
48 AffineMap map = reassociationMap[dimIndex];
54 for (
auto dim : llvm::seq_inclusive(startPos, endPos)) {
55 dynamicDims.push_back(builder.
createOrFold<tensor::DimOp>(loc, src, dim));
57 expr = (expr ? expr * currExpr : currExpr);
72 return llvm::to_vector<4>(llvm::map_range(
73 llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
74 return getCollapsedOutputDimFromInputShape(
75 builder, loc, dim, src, dstStaticShape, reassociation);
85 if (!ShapedType::isDynamic(dstStaticShape[dimIndex])) {
89 unsigned sourceDimPos = expandedDimToCollapsedDim[dimIndex];
90 unsigned startPos = reassociation[sourceDimPos]
95 unsigned endPos = reassociation[sourceDimPos]
100 int64_t linearizedStaticDim = 1;
102 llvm::enumerate(dstStaticShape.slice(startPos, endPos - startPos + 1))) {
103 if (d.index() + startPos ==
static_cast<unsigned>(dimIndex))
105 assert(!ShapedType::isDynamic(d.value()) &&
106 "single dimension cannot be expanded into multiple dynamic "
108 linearizedStaticDim *= d.value();
111 builder.
create<tensor::DimOp>(loc, src, sourceDimPos).getResult();
130 return llvm::to_vector<4>(llvm::map_range(
131 llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
132 return getExpandedOutputDimFromInputShape(builder, loc, dim, src,
133 dstStaticShape, reassociation,
134 expandedDimToCollapsedDim);
142 return dstStaticShape.size() >
144 llvm::cast<ShapedType>(src.
getType()).getRank())
146 builder, loc, src, dstStaticShape, reassocation)
148 builder, loc, src, dstStaticShape, reassocation);
151 template <
typename OpTy>
153 :
public ReifyRankedShapedTypeOpInterface::ExternalModel<
154 ReifyExpandOrCollapseShapeOp<OpTy>, OpTy> {
159 auto reshapeOp = cast<OpTy>(op);
161 b, loc, reshapeOp.getSrc(), reshapeOp.getResultType().getShape(),
162 reshapeOp.getReassociationMaps()));
170 :
public ReifyRankedShapedTypeOpInterface::ExternalModel<ReifyPadOp,
175 auto padOp = cast<PadOp>(op);
177 auto lowPad = padOp.getMixedLowPad();
178 auto highPad = padOp.getMixedHighPad();
180 for (
auto dim : llvm::seq<int64_t>(0, padOp.getSourceType().getRank())) {
181 if (!padOp.getResultType().isDynamicDim(dim)) {
182 shapes.push_back(b.
getIndexAttr(padOp.getResultType().getDimSize(dim)));
188 mapOperands.push_back(
189 b.
createOrFold<tensor::DimOp>(loc, padOp.getSource(), dim));
190 mapOperands.push_back(lowPad[dim]);
191 mapOperands.push_back(highPad[dim]);
199 reifiedReturnShapes.emplace_back(std::move(shapes));
209 ExpandShapeOp::attachInterface<
211 CollapseShapeOp::attachInterface<
213 PadOp::attachInterface<ReifyPadOp>(*ctx);
static OpFoldResult getCollapsedOutputDimFromInputShape(OpBuilder &builder, Location loc, int64_t dimIndex, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociationMap)
For reshape op compute the shape at dimension dimIndex of the output in terms of shape of the src,...
static llvm::DenseMap< int64_t, int64_t > getExpandedDimToCollapsedDimMap(ArrayRef< AffineMap > reassociation)
Compute a map that for a given dimension of the expanded type gives the dimension in the collapsed ty...
static OpFoldResult getExpandedOutputDimFromInputShape(OpBuilder &builder, Location loc, int64_t dimIndex, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation, llvm::DenseMap< int64_t, int64_t > &expandedDimToCollapsedDim)
For an expanding reshape op, compute the value for a dimension of the output from the shape of the in...
static SmallVector< OpFoldResult, 4 > getReshapeOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassocation)
static SmallVector< OpFoldResult, 4 > getExpandedOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation)
Given the src of an expanding reshape op, the reassociation maps and the result type,...
static SmallVector< OpFoldResult, 4 > getCollapsedOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src, ArrayRef< int64_t > dstStaticShape, ArrayRef< AffineMap > reassociation)
Given the src of a collapsing reshape op and its reassociation maps, compute the shape of the result ...
A dimensional identifier appearing in an affine expression.
Base type for affine expression.
AffineExpr floorDiv(uint64_t v) const
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap get(MLIRContext *context)
Returns a zero result affine map with no dimensions or symbols: () -> ().
ArrayRef< AffineExpr > getResults() const
IntegerAttr getIndexAttr(int64_t value)
AffineExpr getAffineSymbolExpr(unsigned position)
AffineExpr getAffineDimExpr(unsigned position)
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
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 helps build Operations.
void createOrFold(SmallVectorImpl< Value > &results, Location location, Args &&...args)
Create an operation of specific op type at the current insertion point, and immediately try to fold i...
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents a single result from folding an operation.
Operation is the basic unit of execution within MLIR.
Location getLoc()
The source location the operation was defined or derived from.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
AffineApplyOp makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Returns a composed AffineApplyOp by composing map and operands with other AffineApplyOps supplying th...
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands)
Constructs an AffineApplyOp that applies map to operands after composing the map with the maps of any...
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
void registerInferTypeOpInterfaceExternalModels(mlir::DialectRegistry ®istry)
Registers external models for Infer Type interfaces for tensor ops.
This header declares functions that assist transformations in the MemRef dialect.
LogicalResult reifyResultShapes(OpBuilder &b, Operation *op, ReifiedRankedShapedTypeDims &reifiedReturnShapes)
Reify the shape of the result of an operation (typically in terms of the shape of its operands).
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
LogicalResult reifyResultShapes(Operation *op, OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) const
This class represents an efficient way to signal success or failure.