/// Fuses producer into consumer if the producer is structurally feasible and
/// the fusion would not violate dependencies.
-/// Implements the fusion part of the "tileAndFuse on buffers"
-/// transformation and thus requires the `consumerdIdx`^th operand of `consumer`
-/// to be a `subview` op (generally obtained by applying the tiling
-/// transformation).
-Optional<FusionInfo> fuseProducerOfBuffer(OpBuilder &b, LinalgOp consumer,
- unsigned consumerIdx,
+/// Implements the fusion part of the "tileAndFuse on buffers" transformation
+/// and thus requires the `consumerOpOperand` to be a `subview` op (generally
+/// obtained by applying the tiling transformation).
+Optional<FusionInfo> fuseProducerOfBuffer(OpBuilder &b,
+ OpOperand &consumerOpOperand,
const LinalgDependenceGraph &graph);
/// Tensor counterpart of `fuseProducerOfBuffer`.
/// This implements the fusion part of the "tileAndFuse on tensors"
-/// transformation and thus requires the `consumerdIdx`^th operand of `consumer`
-/// to be the result of a `subtensor` op (generally obtained by applying the
-/// tiling transformation).
-Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b, LinalgOp consumer,
- unsigned consumerIdx);
+/// transformation and thus requires the `consumerOpOperand` to be a `subtensor`
+/// op (generally obtained by applying the tiling transformation).
+Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
+ OpOperand &consumerOpOperand);
+/// Tensor counterpart of `fuseProducerOfBuffer`.
+/// This implements the fusion part of the "tileAndFuse on tensors"
+/// transformation and thus requires the `consumerOpOperand` to be a `subtensor`
+/// op (generally obtained by applying the tiling transformation).
+/// Assumes `producerOfTensor` is a Linalg op that produces `consumerOpOperand`.
+Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
+ OpResult producerOpResult,
+ OpOperand &consumerOpOperand);
/// Fuse linalg operation on tensors, with the producer of the operand at
/// position `consumerIdx` of the consumer.
Optional<SmallVector<Value, 1>> fuseTensorOps(PatternRewriter &rewriter,
- Operation *consumer,
- unsigned consumerIdx);
+ OpOperand &consumerOpOperand);
/// Like `getShape`, but only returns statically-known information, without
/// generating any new IR. For each shape dimension, returns >=0 if that
/// `producer.getOutputBuffers()`.
/// 2. Tensor case: `producerIdx` is the index of the tensor in
/// `producer.getResults()`.
-static LinalgOp fuse(OpBuilder &b, LinalgOp producer, unsigned producerIdx,
- LinalgOp consumer, unsigned consumerIdx) {
- AffineMap producerMap = producer.getOutputIndexingMap(producerIdx);
- LLVM_DEBUG(llvm::dbgs() << "Producer Idx: " << producerIdx
+static LinalgOp fuse(OpBuilder &b, LinalgOp producerOp,
+ unsigned producerOutNumber, OpOperand &consumerOpOperand) {
+ AffineMap producerMap = producerOp.getOutputIndexingMap(producerOutNumber);
+ LLVM_DEBUG(llvm::dbgs() << "Producer Idx: " << producerOutNumber
<< ", producer map: " << producerMap << "\n");
DenseMap<unsigned, Range> fusedLoopsAndRanges;
- Location loc = consumer.getLoc();
- Value shapedOperand = consumer.getShapedOperand(consumerIdx);
+ Value shapedOperand = consumerOpOperand.get();
for (auto en : llvm::enumerate(producerMap.getResults())) {
unsigned posInProducerLoop = en.value().cast<AffineDimExpr>().getPosition();
- fusedLoopsAndRanges[posInProducerLoop] =
- getRangeFromOperandShape(b, loc, shapedOperand, en.index());
+ fusedLoopsAndRanges[posInProducerLoop] = getRangeFromOperandShape(
+ b, consumerOpOperand.getOwner()->getLoc(), shapedOperand, en.index());
}
- return fuse(b, producer, fusedLoopsAndRanges);
+ return fuse(b, producerOp, fusedLoopsAndRanges);
}
// Encode structural fusion safety preconditions.
}
static Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
-findFusableProducer(LinalgOp consumer, unsigned consumerIdx,
+findFusableProducer(OpOperand &consumerOpOperand,
const LinalgDependenceGraph &dependenceGraph) {
- assert(consumer.hasBufferSemantics() && "revisit usage of shaped operand");
+ LinalgOp consumerOp = cast<LinalgOp>(consumerOpOperand.getOwner());
+ assert(consumerOp.hasBufferSemantics() && "revisit usage of shaped operand");
// Only consider RAW and WAW atm.
for (auto depType : {
LinalgDependenceGraph::DependenceType::WAW,
}) {
for (auto dependence : llvm::make_filter_range(
- dependenceGraph.getDependencesInto(consumer, depType),
- [consumerIdx](
- LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
- return elem.indexingOpView->getOperandNumber() == consumerIdx;
+ dependenceGraph.getDependencesInto(consumerOp, depType),
+ [&](LinalgDependenceGraph::LinalgDependenceGraphElem elem) {
+ return elem.indexingOpView->get() == consumerOpOperand.get() &&
+ elem.indexingOpView->getOperandNumber() ==
+ consumerOpOperand.getOperandNumber();
})) {
- // Check that the dependence is indeed on the input `consumerIdx` view.
- Value consumedView = dependence.indexingOpView->get();
- if (!isSameSubView(consumer.getShapedOperand(consumerIdx), consumedView))
- continue;
-
// Consumer consumes this view, `isStructurallyFusableProducer` also
// checks whether it is a strict subview of the producer view.
auto producer = cast<LinalgOp>(dependence.dependentOpView->getOwner());
- Value producedView = dependence.dependentOpView->get();
LLVM_DEBUG(llvm::dbgs()
<< "\n"
<< LinalgDependenceGraph::getDependenceTypeStr(depType)
<< dependence.dependentOpView->getOperandNumber() -
producer.getNumInputs()
<< "\n");
- (void)producedView;
// Simple fusability checks.
- if (!isFusableInto(dependenceGraph, consumer, consumedView, producer))
+ if (!isFusableInto(dependenceGraph, consumerOp, consumerOpOperand.get(),
+ producer))
continue;
return dependence;
}
Optional<FusionInfo>
-mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, LinalgOp consumer,
- unsigned consumerIdx,
+mlir::linalg::fuseProducerOfBuffer(OpBuilder &b, OpOperand &consumerOpOperand,
const LinalgDependenceGraph &graph) {
Optional<LinalgDependenceGraph::LinalgDependenceGraphElem> fusableDependence =
- findFusableProducer(consumer, consumerIdx, graph);
+ findFusableProducer(consumerOpOperand, graph);
if (!fusableDependence)
return {};
LinalgOp producerOp =
cast<LinalgOp>(fusableDependence->dependentOpView->getOwner());
// If producer is already in the same block as consumer, we are done.
- if (consumer->getBlock() == producerOp->getBlock())
+ if (consumerOpOperand.get().getParentBlock() ==
+ fusableDependence->dependentOpView->get().getParentBlock())
return {};
unsigned producerIdx =
fusableDependence->dependentOpView->getOperandNumber() -
producerOp.getNumInputs();
- Value consumerView = consumer.getShapedOperand(consumerIdx);
// Must be a subview or a slice to guarantee there are loops we can fuse
// into.
- auto subView = consumerView.getDefiningOp<SubViewOp>();
- auto slice = consumerView.getDefiningOp<SliceOp>();
+ auto subView = consumerOpOperand.get().getDefiningOp<SubViewOp>();
+ auto slice = consumerOpOperand.get().getDefiningOp<SliceOp>();
if (!subView && !slice) {
LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subview or slice)");
return {};
// Fuse `producer` just before `consumer`.
OpBuilder::InsertionGuard g(b);
- b.setInsertionPoint(consumer.getOperation());
- ScopedContext scope(b, consumer.getLoc());
- LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumer << "\n");
+ b.setInsertionPoint(consumerOpOperand.getOwner());
+ ScopedContext scope(b, consumerOpOperand.getOwner()->getLoc());
+ LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: "
+ << *consumerOpOperand.getOwner() << "\n");
- auto fusedProducer = fuse(b, producerOp, producerIdx, consumer, consumerIdx);
+ auto fusedProducer = fuse(b, producerOp, producerIdx, consumerOpOperand);
return FusionInfo{producerOp, fusedProducer};
}
/// Walk back use-def chain through scf::For yields.
/// Sets `producer` and `outputIndex` if it finds a producer LinalgOp
-static void getProducerOfTensor(Value tensor, LinalgOp &producer,
- unsigned &outputIndex) {
+static void getProducerOfTensor(Value tensor, OpResult &opResult) {
if (!tensor.getType().isa<RankedTensorType>())
return;
while (true) {
+ LLVM_DEBUG(llvm::dbgs() << "\ngetProducerOfTensor: " << tensor);
if (auto linalgOp = tensor.getDefiningOp<LinalgOp>()) {
- producer = linalgOp;
- outputIndex = tensor.cast<OpResult>().getResultNumber();
+ opResult = tensor.cast<OpResult>();
return;
}
if (auto subTensorOp = tensor.getDefiningOp<SubTensorOp>()) {
}
if (auto blockArg = tensor.dyn_cast<BlockArgument>()) {
if (auto forOp = blockArg.getDefiningOp<scf::ForOp>()) {
- tensor = forOp.getResult(blockArg.getArgNumber());
+ tensor = *(forOp.getIterOperands().begin() + blockArg.getArgNumber());
continue;
}
}
}
}
-Optional<FusionInfo> mlir::linalg::fuseProducerOfTensor(OpBuilder &b,
- LinalgOp consumer,
- unsigned consumerIdx) {
- Value inputTensor = consumer.getInput(consumerIdx);
- LinalgOp producerOp;
- unsigned producerIdx;
- getProducerOfTensor(inputTensor, producerOp, producerIdx);
+Optional<FusionInfo>
+mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpOperand &consumerOpOperand) {
+ Value inputTensor = consumerOpOperand.get();
+ OpResult producerOpResult;
+ getProducerOfTensor(inputTensor, producerOpResult);
+ if (!producerOpResult) {
+ LLVM_DEBUG(llvm::dbgs() << "\nUnable to find producer");
+ return {};
+ }
+ return fuseProducerOfTensor(b, producerOpResult, consumerOpOperand);
+}
+
+Optional<FusionInfo>
+mlir::linalg::fuseProducerOfTensor(OpBuilder &b, OpResult producerOpResult,
+ OpOperand &consumerOpOperand) {
+ auto producerOp = dyn_cast<LinalgOp>(producerOpResult.getOwner());
+ assert(producerOp && "expected Linalg producer");
+ LinalgOp consumerOp = cast<LinalgOp>(consumerOpOperand.getOwner());
+ Value inputTensor = consumerOpOperand.get();
// Must be a subtensor to guarantee there are loops we can fuse into.
auto subTensor = inputTensor.getDefiningOp<SubTensorOp>();
- if (!subTensor || !producerOp) {
- LLVM_DEBUG(llvm::dbgs() << "\nNot fusable (not a subtensor)");
+ if (!subTensor) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "\nNot fusable, not a subtensor: " << inputTensor);
return {};
}
// If producer is already in the same block as consumer, we are done.
- if (consumer->getBlock() == producerOp->getBlock())
+ if (consumerOpOperand.get().getParentBlock() ==
+ producerOpResult.getParentBlock())
return {};
// Insert fused `producer` just before `consumer`.
OpBuilder::InsertionGuard g(b);
- b.setInsertionPoint(consumer.getOperation());
- ScopedContext scope(b, consumer.getLoc());
- LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumer << "\n");
- LinalgOp fusedProducer =
- fuse(b, producerOp, producerIdx, consumer, consumerIdx);
+ b.setInsertionPoint(consumerOp);
+ ScopedContext scope(b, consumerOp->getLoc());
+ LLVM_DEBUG(llvm::dbgs() << "Fuse into consumer: " << *consumerOp << "\n");
+ LinalgOp fusedProducer = fuse(
+ b, producerOp, producerOpResult.getResultNumber(), consumerOpOperand);
// Replace use.
// Canonicalizations are not guaranteed to have happened before constructing
// `fusedProducer`. In the tensor case this can result in temporary type
// mismatches. Insert a `tensor.cast` op to propagate the transformation
// invariant that types are compatible.
- Value def = fusedProducer->getResult(producerIdx);
- OpOperand &use = consumer->getOpOperand(consumerIdx);
- Type consumerType = use.get().getType();
+ Value def = fusedProducer->getResult(producerOpResult.getResultNumber());
+ Type consumerType = consumerOpOperand.get().getType();
if (consumerType != def.getType())
def = b.create<tensor::CastOp>(fusedProducer.getLoc(), consumerType, def);
- use.set(def);
- return FusionInfo{producerOp, fusedProducer};
+ consumerOpOperand.set(def);
+ return FusionInfo{cast<LinalgOp>(producerOpResult.getOwner()), fusedProducer};
}
/// Prune all dimensions that are of reduction iterator type from `map`.
// in the meanwhile disallow such a fusion.
DenseMap<Operation *, AffineMap> fusedProducerIndexingMap;
for (LinalgOp op : reverse(ops)) {
- for (auto operandIndex :
- llvm::seq<unsigned>(0, op.getNumShapedOperands())) {
+ for (OpOperand &opOperand : op.getShapedOpOperands()) {
Optional<LinalgDependenceGraph::LinalgDependenceGraphElem>
- fusableDependence =
- findFusableProducer(op, operandIndex, dependenceGraph);
+ fusableDependence = findFusableProducer(opOperand, dependenceGraph);
if (!fusableDependence)
continue;
LinalgOp producerOp =
op.emitRemark(
"unhandled non permutation indexing map for fused view in "
"producer for operand at index ")
- << operandIndex;
+ << opOperand.getOperandNumber();
return FusableOpDependencesTy{};
}
op.emitRemark(
"unhandled case where indexing map for fused view in the consumer "
"is not a projected permutation while fusing at index ")
- << operandIndex;
+ << opOperand.getOperandNumber();
return FusableOpDependencesTy{};
}
}
static Optional<SmallVector<Value, 1>>
-fuseTensorOpsImpl(LinalgOp producer, LinalgOp consumer, unsigned consumerIdx,
+fuseTensorOpsImpl(LinalgOp producer, OpOperand &consumerOpOperand,
PatternRewriter &rewriter) {
+ LinalgOp consumer = cast<LinalgOp>(consumerOpOperand.getOwner());
+ unsigned consumerIdx = consumerOpOperand.getOperandNumber();
if (!areTensorOpsFusable(producer, consumer, consumerIdx))
return llvm::None;
} // namespace
Optional<SmallVector<Value, 1>>
-mlir::linalg::fuseTensorOps(PatternRewriter &rewriter, Operation *consumer,
- unsigned consumerIdx) {
- if (consumerIdx >= consumer->getNumOperands())
- return llvm::None;
- Operation *producer = consumer->getOperand(consumerIdx).getDefiningOp();
+mlir::linalg::fuseTensorOps(PatternRewriter &rewriter,
+ OpOperand &consumerOpOperand) {
+ Operation *producer = consumerOpOperand.get().getDefiningOp();
if (!producer || producer->getNumResults() != 1)
return llvm::None;
// Fuse when consumer is GenericOp or IndexedGenericOp.
- if (!isa<GenericOp, IndexedGenericOp>(consumer) ||
+ if (!isa<GenericOp, IndexedGenericOp>(consumerOpOperand.getOwner()) ||
!isa<GenericOp, IndexedGenericOp>(producer))
return llvm::None;
- return fuseTensorOpsImpl(cast<LinalgOp>(producer), cast<LinalgOp>(consumer),
- consumerIdx, rewriter);
+ return fuseTensorOpsImpl(cast<LinalgOp>(producer), consumerOpOperand,
+ rewriter);
}
namespace {
LogicalResult matchAndRewrite(LinalgOpTy op,
PatternRewriter &rewriter) const override {
// Find the first operand that is defined by another generic op on tensors.
- for (auto operandNum : llvm::seq<unsigned>(0, op->getNumOperands())) {
- Operation *producer = op->getOperand(operandNum).getDefiningOp();
+ for (OpOperand &opOperand : op.getShapedOpOperands()) {
+ Operation *producer = opOperand.get().getDefiningOp();
if (!producer)
continue;
Optional<SmallVector<Value, 1>> fusedOpResults =
- fuseTensorOps(rewriter, op, operandNum);
+ fuseTensorOps(rewriter, opOperand);
if (fusedOpResults) {
rewriter.replaceOp(op, *fusedOpResults);
if (producer->use_empty())