[mlir][scf] Fix bug in peelForLoop

[springerm]

Summary

Insertion point should be set before creating new operations.

[nicolasvasilache]

can you also add a simple test?

[springerm]

In D107326#2922290, @nicolasvasilache wrote:

can you also add a simple test?

This bug is only triggered when calling peelForLoop with a builder who's insertion point is somewhere else. That cannot happen with the peel-for-loop pass that I am using for unit tests.

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

LogicalResult mlir::scf::peelForLoop(RewriterBase &b, ForOp forOp,
   if (stepInt == static_cast<int64_t>(1))
     return failure();
 
   auto loc = forOp.getLoc();
   AffineExpr dim0, dim1, dim2;
   bindDims(b.getContext(), dim0, dim1, dim2);
   // New upper bound: %ub - (%ub - %lb) mod %step
   auto modMap = AffineMap::get(3, 0, {dim1 - ((dim1 - dim0) % dim2)});
+  b.setInsertionPoint(forOp);
   Value splitBound = b.createOrFold<AffineApplyOp>(
       loc, modMap,
       ValueRange{forOp.lowerBound(), forOp.upperBound(), forOp.step()});
 
   // Set new upper loop bound.
   Value previousUb = forOp.upperBound();
   b.updateRootInPlace(forOp,
                       [&]() { forOp.upperBoundMutable().assign(splitBound); });