[Mlir-commits] [mlir] [mlir][TilingInterface] Make `tileAndFuseConsumerOfSlice` take surrounding loops as an argument. (PR #132082)

[Abhishek Varma]

================
@@ -1890,25 +1890,81 @@ getPerfectlyNestedLoopsOutsideOf(scf::ForOp loop) {
   return {nestLoops.rbegin(), nestLoops.rend()};
 }
 
+/// Check that the loop is perfectly nested.
+static bool
+isPerfectlyNestedForLoops(MutableArrayRef<LoopLikeOpInterface> loops) {
+  assert(!loops.empty() && "unexpected empty loop nest");
+  if (loops.size() == 1) {
+    return isa_and_nonnull<scf::ForOp>(loops.front().getOperation());
+  }
+  for (auto [outerLoop, innerLoop] :
+       llvm::zip_equal(loops.drop_back(), loops.drop_front())) {
+    auto outerFor = dyn_cast_or_null<scf::ForOp>(outerLoop.getOperation());
+    auto innerFor = dyn_cast_or_null<scf::ForOp>(innerLoop.getOperation());
+    if (!outerFor || !innerFor) {
+      return false;
+    }
+    auto outerBBArgs = outerFor.getRegionIterArgs();
+    auto innerIterArgs = innerFor.getInitArgs();
+    if (outerBBArgs.size() != innerIterArgs.size()) {
+      return false;
+    }
+
+    for (auto [outerBBArg, innerIterArg] :
+         llvm::zip(outerBBArgs, innerIterArgs)) {
+      if (!llvm::hasSingleElement(outerBBArg.getUses()) ||
+          innerIterArg != outerBBArg) {
+        return false;
+      }
+    }
+
+    auto outerYields =
+        cast<scf::YieldOp>(outerFor.getBody()->getTerminator())->getOperands();
+    auto innerResults = innerFor.getResults();
+    if (outerYields.size() != innerResults.size()) {
+      return false;
+    }
+    for (auto [outerYield, innerResult] :
+         llvm::zip(outerYields, innerResults)) {
+      if (!llvm::hasSingleElement(innerResult.getUses()) ||
+          outerYield != innerResult) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
 /// Fetch the untiled consumer of a scf.for's result which is yielded by a
 /// tensor.insert_slice. This function makes the following assumptions :
 /// 1.  tensor.insert_slice has scf.yield as its only user.
 /// 2.  scf.for's corresponding result has only one use.
 static FailureOr<OpOperand *>
 getUntiledConsumerFromSlice(RewriterBase &rewriter,
-                            tensor::InsertSliceOp candidateSliceOp) {
+                            tensor::InsertSliceOp candidateSliceOp,
+                            MutableArrayRef<LoopLikeOpInterface> loops) {
+  assert(!loops.empty() && "unexpected loops to be empty");
+  // 1. Expect slice to be part of the body of the inner most loop.
----------------
Abhishek-Varma wrote:

Extreme nit but for consistency : The numbering of step seems incoherent. Either this can be the 0-th step or the latter ones can be incremented by one starting at line 1963 "`// Step 1. Fetch the corresponding output.`".

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