diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h --- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h +++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h @@ -61,11 +61,11 @@ bool insideMutuallyExclusiveBranches(Operation *a, Operation *b); /// An owning vector of values, handy to return from functions. -using ValueVector = std::vector; -using LoopVector = std::vector; +using ValueVector = SmallVector; +using LoopVector = SmallVector; 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 diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp @@ -821,7 +821,7 @@ 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(); } diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -523,7 +523,7 @@ 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 @@ -573,9 +573,9 @@ builder.create(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( diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp --- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp +++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp @@ -179,12 +179,6 @@ SmallVector lbs(numTiledDims, zero); SmallVector 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, @@ -193,15 +187,10 @@ helper.emitLoopNestBody(nestedBuilder, loc, outputIvs); // Insert the slice into the destination. - insertResult = nestedBuilder.create( - loc, tile, iterArgs[0], insertParams); - return {insertResult}; + return {nestedBuilder.create( + 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(); }