[mlir] Change scf::LoopNest to store 'results'.

[csigg]

Summary

This fixes the case where scf::LoopNest::loops is empty.

Change LoopVector and ValueVector to SmallVector.

[nicolasvasilache]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

66–67

hmm .. who are the clients of this?

I am thinking that whoever is the client needs to either check !loops.empty or ValueRange != null, so whatever the case some check is needed at the client.
From an API perspective, it seems confusing to say that the you can take the result of an empty loop nest, I'd expect an assert(!loops.empty) here.

Is there another tradeoff I am not yet seeing that led you to this API change?

[csigg]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

66–67

I expect that I can default-construct a LoopNest instance and call getResults() without crashing.

My particular use-case is that we call buildLoopNest() with empty lbs and empty iterArgs, which returns such a default-constructed LoopNest. One could argue that buildLoopNest().getResults() should return the ValueVector (which imho should be a SmallVector instead of a std::vector like everywhere else in MLIR) of the builder in that case, but that's an orthogonal issue.

[csigg]

Add 'results' member to LoopNest.

[csigg]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

66–67

I changed this revision to include a 'results' vector in LoopNest to handle the empty lbs case properly.

[csigg]

Nicolas, could you take another look please?

[nicolasvasilache]

This is missing a test.

mlir/include/mlir/Dialect/SCF/IR/SCF.h

64–65

Thanks for cleaning this up.

66–67

I am sorry but I don't think that is a good change, you are now introducing opportunities to have a decoupling between loops and results that does not make sense to me..

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

824

It is hard to see from this PR without tests whether this the one place where you see failures.

If so, why not just:

if (nest.loops.empty())
  return failure();
rewriter.replaceOp(op, nest.getResults());

?

The absence of a LoopNest is not the same as a loop nest without results ..

[nicolasvasilache]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

66–67

My particular use-case is that we call buildLoopNest() with empty lbs and empty iterArgs, which returns such a default-constructed LoopNest. One could argue that buildLoopNest().getResults() should return the ValueVector (which imho should be a SmallVector instead of a std::vector like everywhere else in MLIR) of the builder in that case, but that's an orthogonal issue.

I could see the APIs evolving to returning a FailureOr<LoopNest> and dropping the default constructor / enforcing non-empty loop nest at construction.

mlir/lib/Dialect/SCF/IR/SCF.cpp

526

Just make the API return FailureOr<LoopNest> / llvm::None ?

[csigg]

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

824

This is not where I saw the failure, this change is just necessary because I changed LoopNest's result from a member function to a member.

mlir/lib/Dialect/SCF/IR/SCF.cpp

526

That would call the bodyBuilder, create IR, and then return Failure? That sounds unintuitive.

[nicolasvasilache]

mlir/lib/Dialect/SCF/IR/SCF.cpp

526

I wouldn't call it here but just return failure.
Maybe this corner case has been used as a "way to inline" the computation.

This code is quite old, seems like an opportunity to refactor to force users to call inline when appropriate and have the LoopNest building part either build a loop nest or fail as one would expect.

Atm the impl, seem to be able to either:

1. build a LoopNest with a body (with or without results / iter args)
2. build a LoopNest without a body (with or without results / iter args)
3. inline the body without creating a LoopNest.

This API is doing too much and splitting 3. out would be the much cleaner solution IMO.

[csigg]

Clean up test TODO.

[csigg]

This API is doing too much and splitting 3. out would be the much cleaner solution IMO.

That is your opinion, but the API is currently used for all three cases and avoiding the foot-gun for 3. is an orthogonal improvement.

I noticed only now that this change allows cleaning up a TODO item in the test.

[ftynse]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

64–65

The point of using std::vector instead of SmallVector was to avoid excessive copying of stack-allocated data since LoopNest instances are returned by value.

66–67

How about changing getResults to getResultsOr(ValueRange)? The latter would work similarly to Optional::getValueOr and return the argument when no loops were created. This seems like it could satisfy the use case without changing how the builder works.

[nicolasvasilache]

Sync'd offline with Alex who authored this builder.

My objection is I view buildX that 1. does not create X and 2. does not fail but 3. does something else as problematic.
We discussed an analogy: 0-D load is a corner case but still creates a load; 0-D loop nest OTOH will not create a loop nest so we need to put a conditional somewhere.
I am advocating to put the conditional at the client rather than hide it inside the builder because of my point just above.

OTOH Alex justly remarked that scf::buildLoopNest is not the same as scf::buildScfForOps and that it builds some abstract loop nest that we are ok defining as "empty loop nest with the builder called once".

Leaving @ftynse with the approval decision once this LGTH.

[csigg]

Thanks for reviewing and sorry for the delay. Just to be clear, although I think it's a minor improvement, I'm not attached to this change. If you prefer to keep things as is, I will just abandon it.

mlir/include/mlir/Dialect/SCF/IR/SCF.h

64–65

I don't think the extra copies matter. Both Value and scf::ForOp are just a pointer.

66–67

That doesn't really help in e.g. TestTensorTransforms.cpp or my use case because there is no convenient way to access the body builder return value.

[ftynse]

mlir/include/mlir/Dialect/SCF/IR/SCF.h

66–67

Good point.

mlir/include/mlir/Dialect/SCF/IR/SCF.h

 ForeachThreadOp getForeachThreadOpThreadIndexOwner(Value val);
 
 /// Return true if ops a and b (or their ancestors) are in mutually exclusive
 /// regions/blocks of an IfOp.
 // TODO: Consider moving this functionality to RegionBranchOpInterface.
 bool insideMutuallyExclusiveBranches(Operation *a, Operation *b);
 
 /// An owning vector of values, handy to return from functions.
-using ValueVector = std::vector<Value>;
-using LoopVector = std::vector<scf::ForOp>;
+using ValueVector = SmallVector<Value>;
+using LoopVector = SmallVector<scf::ForOp>;
 struct LoopNest {
-  ResultRange getResults() { return loops.front().getResults(); }
   LoopVector loops;
+  ValueVector results;
 };
 
 /// Creates a perfect nest of "for" loops, i.e. all loops but the innermost
 /// contain only another loop and a terminator. The lower, upper bounds and
 /// steps are provided as `lbs`, `ubs` and `steps`, which are expected to be of
 /// the same size. `iterArgs` points to the initial values of the loop iteration
 /// arguments, which will be forwarded through the nest to the innermost loop.
 /// The body of the loop is populated using `bodyBuilder`, which accepts an

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

struct PadOpTilingPattern : public OpRewritePattern<tensor::PadOp> {
 
   LogicalResult matchAndRewrite(tensor::PadOp op,
                                 PatternRewriter &rewriter) const override {
     tensor::PadOp newPadOp;
     LoopNest loopNest;
     if (failed(tilePadOp(rewriter, op, newPadOp, loopNest, options)))
       return failure();
     // Replace all uses of the original tensor::PadOp.
-    rewriter.replaceOp(op, loopNest.getResults()[0]);
+    rewriter.replaceOp(op, loopNest.results.front());
     return success();
   }
 
   LinalgTilingOptions options;
 };
 } // namespace
 
 namespace {

mlir/lib/Dialect/SCF/IR/SCF.cpp

LoopNest mlir::scf::buildLoopNest(
   // If there are no bounds, call the body-building function and return early.
   if (lbs.empty()) {
     ValueVector results =
         bodyBuilder ? bodyBuilder(builder, loc, ValueRange(), iterArgs)
                     : ValueVector();
     assert(results.size() == iterArgs.size() &&
            "loop nest body must return as many values as loop has iteration "
            "arguments");
-    return LoopNest();
+    return LoopNest{{}, std::move(results)};
   }
 
   // First, create the loop structure iteratively using the body-builder
   // callback of `ForOp::build`. Do not create `YieldOp`s yet.
   OpBuilder::InsertionGuard guard(builder);
   SmallVector<scf::ForOp, 4> loops;
   SmallVector<Value, 4> ivs;
   loops.reserve(lbs.size());
ValueVector results = bodyBuilder
                             : ValueVector();
   assert(results.size() == iterArgs.size() &&
          "loop nest body must return as many values as loop has iteration "
          "arguments");
   builder.setInsertionPointToEnd(loops.back().getBody());
   builder.create<scf::YieldOp>(loc, results);
 
   // Return the loops.
-  LoopNest res;
-  res.loops.assign(loops.begin(), loops.end());
-  return res;
+  ValueVector nestResults;
+  llvm::copy(loops.front().getResults(), std::back_inserter(nestResults));
+  return LoopNest{std::move(loops), std::move(nestResults)};
 }
 
 LoopNest mlir::scf::buildLoopNest(
     OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs,
     ValueRange steps,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilder) {
   // Delegate to the main function by wrapping the body builder.
   return buildLoopNest(builder, loc, lbs, ubs, steps, llvm::None,

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

LogicalResult emitReplacement(tensor::ExtractSliceOp op, Value dest,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     const unsigned numTiledDims = helper.getIterationSpaceSizes().size();
     auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     auto one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
     SmallVector<Value> lbs(numTiledDims, zero);
     SmallVector<Value> steps(numTiledDims, one);
 
-    // Below, we pass out the result of the loop body builder lambda via the
-    // `insertResult` variable. In certain cases, no loops will be created, but
-    // the body builder will still execute. In this case, the results will not
-    // be passed to the LoopNest object.
-    // TODO: remove this workaround if `scf::buildLoopNest` behavior is updated.
-    Value insertResult = nullptr;
     scf::LoopNest nest = scf::buildLoopNest(
         rewriter, loc, lbs, helper.getIterationSpaceSizes(), steps, dest,
         [&](OpBuilder &nestedBuilder, Location loc, ValueRange outputIvs,
             ValueRange iterArgs) -> scf::ValueVector {
           auto [tile, insertParams] =
               helper.emitLoopNestBody(nestedBuilder, loc, outputIvs);
 
           // Insert the slice into the destination.
-          insertResult = nestedBuilder.create<tensor::InsertSliceOp>(
-              loc, tile, iterArgs[0], insertParams);
-          return {insertResult};
+          return {nestedBuilder.create<tensor::InsertSliceOp>(
+              loc, tile, iterArgs[0], insertParams)};
         });
-
-    if (!nest.loops.empty())
-      rewriter.replaceOp(op, nest.getResults());
-    else
-      rewriter.replaceOp(op, insertResult);
+    rewriter.replaceOp(op, nest.results);
 
     return success();
   }
 };
 
 struct RewriteExtractSliceFromCollapseShapeUsingScfForeach
     : public RewriteExtractSliceFromCollapseShapeBase {
   RewriteExtractSliceFromCollapseShapeUsingScfForeach(MLIRContext *context)