20 #include "llvm/Support/Debug.h"
22 #define DEBUG_TYPE "tensor-swap-slices"
28 auto producerOp = dyn_cast<TilingInterface>(producer.
getOwner());
33 if (!llvm::all_of(sliceOp.getMixedStrides(),
isOneInteger))
36 FailureOr<TilingResult> tiledResult = producerOp.generateResultTileValue(
38 sliceOp.getMixedSizes());
39 if (failed(tiledResult))
44 llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
45 if (droppedDims.any()) {
46 assert(tiledResult->tiledValues.size() == 1 &&
47 "expected only a single tiled result value to replace the extract "
53 auto newSliceOp = tensor::ExtractSliceOp::create(
54 builder, sliceOp.getLoc(), sliceOp.
getType(),
55 tiledResult->tiledValues[0], offsets, sliceOp.getMixedSizes(), strides);
56 tiledResult->tiledValues[0] = newSliceOp;
65 if (sliceOps.empty()) {
67 { llvm::dbgs() <<
"expected candidate slices list to be non-empty"; });
70 if (sliceOps.size() != consumerOperands.size()) {
73 <<
"expected as many operands as the number of slices passed";
78 dyn_cast<TilingInterface>(consumerOperands.front()->getOwner());
81 for (
auto opOperand : consumerOperands.drop_front()) {
82 if (opOperand->getOwner() != consumerOp) {
85 <<
"expected all consumer operands to be from the same operation";
91 auto consumerOperandNums = llvm::map_to_vector(
92 consumerOperands, [](
OpOperand *opOperand) ->
unsigned {
97 for (
auto sliceOp : sliceOps) {
100 if (!llvm::all_of(sliceOp.getMixedStrides(),
isOneInteger))
105 allOffsets.emplace_back(std::move(offsets));
106 allSizes.emplace_back(std::move(sizes));
108 FailureOr<TilingResult> tiledResult =
109 consumerOp.getTiledImplementationFromOperandTiles(
110 builder, consumerOperandNums, allOffsets, allSizes);
111 if (failed(tiledResult))
IntegerAttr getIndexAttr(int64_t value)
Ty getType(Args &&...args)
Get or construct an instance of the type Ty with provided arguments.
This class helps build Operations.
This class represents an operand of an operation.
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
This is a value defined by a result of an operation.
Operation * getOwner() const
Returns the operation that owns this result.
unsigned getResultNumber() const
Returns the number of this result.
FailureOr< TilingResult > replaceExtractSliceWithTiledProducer(OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp)
Method to swap an tensor.extract_slice with its producer when the producer implements the TilingInter...
FailureOr< TilingResult > replaceInsertSlicesWithTiledConsumer(OpBuilder &builder, ArrayRef< tensor::InsertSliceOp > sliceOps, ArrayRef< OpOperand * > consumerOperands)
Method to swap tensor.insert_slices with their consumers when the consumer implements the TilingInter...
Include the generated interface declarations.
bool isOneInteger(OpFoldResult v)
Return true if v is an IntegerAttr with value 1.