29 assert(((resType.getNumDynamicDims() == dynOutDims.size()) ||
30 dynOutDims.empty()) &&
31 "Either none or all output dynamic dims must be specified!");
40 for (
const auto [idx, val] :
enumerate(resType.getShape())) {
41 bool isDimDynamic = ShapedType::isDynamic(val);
42 bool updatePadHigh = !isDimDynamic || !dynOutDims.empty();
59 return PadOp::create(b, loc, resType, source, low, high, pad, nofold);
65 auto tensorTy = cast<RankedTensorType>(rankedTensor.
getType());
68 if (en.value() == ShapedType::kDynamic)
69 dynamicDims.push_back(
70 tensor::DimOp::create(b, loc, rankedTensor, en.index()));
75 FailureOr<RankedTensorType>
78 if (transposeVector.empty())
79 return rankedTensorType;
82 transposeVector.size() !=
static_cast<size_t>(rankedTensorType.getRank()))
89 RankedTensorType transposedTensorType =
90 RTTBuilder(rankedTensorType).
setShape(transposedShape);
91 return transposedTensorType;
96 const llvm::SmallBitVector &
dropDims) {
97 auto srcType = cast<ShapedType>(src.
getType());
98 int64_t rank = srcType.getRank();
99 assert(rank ==
static_cast<int64_t
>(
dropDims.size()) &&
100 "dropDims dimension does not match src tensor rank");
103 [&](
unsigned dim) { return srcType.getShape()[dim] == 1; }) &&
104 "Dropping non unit dimension");
109 int64_t nextDimToGroup = 0;
110 llvm::SmallBitVector keptDims(
dropDims);
112 int64_t lastSetBit = keptDims.find_last();
113 for (int64_t
setBit : keptDims.set_bits()) {
117 int64_t upTo =
setBit == lastSetBit ? rank - 1 :
setBit;
118 auto seq = llvm::seq_inclusive(nextDimToGroup, upTo);
119 reassocMaps.emplace_back(llvm::make_range(seq.begin(), seq.end()));
120 nextDimToGroup =
setBit + 1;
122 return tensor::CollapseShapeOp::create(b, loc, src, reassocMaps);
126 llvm::SmallBitVector droppedDims = op.getDroppedDims();
128 RankedTensorType resultType = op.getDestType();
131 for (int64_t resultDim = 0; resultDim < resultType.getRank(); ++resultDim) {
132 if (droppedDims.test(resultDim)) {
135 if (resultType.getDimSize(resultDim) != 1)
140 {op.getSource(), srcDim}, {op.getResult(), resultDim});
141 if (failed(equalDimSize) || !*equalDimSize)
150 llvm::SmallBitVector droppedDims = op.getDroppedDims();
151 int64_t resultDim = 0;
154 RankedTensorType sourceType = op.getSourceType();
155 for (int64_t dim = 0, e = sourceType.getRank(); dim < e; ++dim) {
156 if (droppedDims.test(dim)) {
159 if (sourceType.getDimSize(dim) != 1)
164 {op.getSource(), dim}, {op.getResult(), resultDim});
165 if (failed(equalDimSize) || !*equalDimSize)
static void setBit(char *rawData, size_t bitPos, bool value)
Set a bit to a specific value.
Base type for affine expression.
IntegerAttr getIndexAttr(int64_t value)
MLIRContext * getContext() const
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
This class helps build Operations.
This class represents a single result from folding an operation.
This is a builder type that keeps local references to arguments.
Builder & setShape(ArrayRef< int64_t > newShape)
static FailureOr< bool > areEqual(const Variable &var1, const Variable &var2)
Compute whether the given variables are equal.
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...
Type getType() const
Return the type of this value.
OpFoldResult makeComposedFoldedAffineApply(OpBuilder &b, Location loc, AffineMap map, ArrayRef< OpFoldResult > operands, bool composeAffineMin=false)
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)
PadOp createPadHighOp(RankedTensorType resType, Value source, Value pad, bool nofold, Location loc, OpBuilder &builder, ValueRange dynOutDims={})
SmallVector< Value > createDynamicDimValues(OpBuilder &b, Location loc, Value rankedTensor)
bool isCastLikeInsertSliceOp(InsertSliceOp op)
A tensor.insert_slice is a cast-like operation if it merely rank-extends the source tensor or inserts...
CollapseShapeOp dropGivenUnitDims(OpBuilder &b, Location loc, Value src, const llvm::SmallBitVector &dropDims)
Create tensor.collapse_shape to drop unit dimensions in dropDims in tensor src.
bool isCastLikeExtractSliceOp(ExtractSliceOp op)
A tensor.extract_slice is a cast-like operation if it merely rank-reduces unit dimensions of the sour...
OpFoldResult getMixedSize(OpBuilder &builder, Location loc, Value value, int64_t dim)
Return the dimension of the given tensor value.
FailureOr< RankedTensorType > computeTransposedType(RankedTensorType rankedTensorType, ArrayRef< int64_t > transposeVector)
Returns the transposed rankedTensorType if transposeVector is non-empty.
Include the generated interface declarations.
void bindDims(MLIRContext *ctx, AffineExprTy &...exprs)
Bind a list of AffineExpr references to DimExpr at positions: [0 .
void applyPermutationToVector(SmallVector< T, N > &inVec, ArrayRef< int64_t > permutation)
Apply the permutation defined by permutation to inVec.
SmallVector< int64_t > dropDims(ArrayRef< int64_t > inputPerm, ArrayRef< int64_t > dropPositions)
Returns a permutation vector that drop the input dims in dropPositions from inputPerm.
bool isPermutationVector(ArrayRef< int64_t > interchange)
Method to check if an interchange vector is a permutation.