[Mlir-commits] [mlir] [MLIR][SCF] Add an API to fuse consumer to a producer within scf loop (PR #88712)

Oleksandr Alex Zinenko llvmlistbot at llvm.org
Thu Apr 18 04:58:19 PDT 2024


================
@@ -1100,6 +1101,451 @@ mlir::scf::tileConsumerAndFuseProducersUsingSCF(
                                    replacements};
 }
 
+//===----------------------------------------------------------------------===//
+// tileAndFuseConsumerUsingSCF implementation.
+//===----------------------------------------------------------------------===//
+
+/// We traverse through the use-def chain of the tensor.insert_slice op through
+/// the containing scf.for to fetch the first untiled consumer.  We also return
+/// the operand number of the consumer and the result number of the scf.for
+/// being consumed.
+static Operation *
+getUntiledConsumerFromSliceDestSCFFor(tensor::InsertSliceOp candidateSliceOp,
+                                      unsigned &operandNumber,
+                                      unsigned &resultNumber) {
+  // Step 1. Fetch the corresponding output.
+  Value sliceResult = candidateSliceOp.getResult();
+  Value::user_range users = sliceResult.getUsers();
+  auto yieldOp = cast<scf::YieldOp>(*users.begin());
+  for (Value operand : yieldOp->getOperands()) {
----------------
ftynse wrote:

Please document the assumption that the first (and only) user of the slice is the yield. I don't think we guarantee anything about the order of users in the user range.

https://github.com/llvm/llvm-project/pull/88712


More information about the Mlir-commits mailing list