[Mlir-commits] [mlir] [mlir] Fix loop pipelining when the operand of `yield` is not defined in the loop body (PR #75423)

[Thomas Raoux]

================
@@ -551,21 +560,24 @@ LogicalResult LoopPipelinerInternal::createKernel(
   for (const auto &retVal :
        llvm::enumerate(forOp.getBody()->getTerminator()->getOperands())) {
     Operation *def = retVal.value().getDefiningOp();
-    assert(def && "Only support loop carried dependencies of distance 1");
-    unsigned defStage = stages[def];
-    if (defStage > 0) {
+    assert(def && "Only support loop carried dependencies of distance of 1 or "
+                  "defined outside the loop");
+    auto defStage = stages.find(def);
+    if (defStage != stages.end() && defStage->second > 0)
       setValueMapping(forOp.getRegionIterArgs()[retVal.index()],
                       newForOp->getResult(retVal.index()),
-                      maxStage - defStage + 1);
-    }
+                      maxStage - defStage->second + 1);
+    else
----------------
ThomasRaoux wrote:

nit: please add  braces for multi-line if/else

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