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

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

Change LoopVector and ValueVector to SmallVector.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D136926

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCF.h b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
index 84b0ad3..c1e7bc3 100644 (file)
--- a/mlir/include/mlir/Dialect/SCF/IR/SCF.h
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCF.h
@@ -61,11 +61,11 @@ ForeachThreadOp getForeachThreadOpThreadIndexOwner(Value val);
 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

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index 4db555c..e87d4f9 100644 (file)
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
@@ -821,7 +821,7 @@ struct PadOpTilingPattern : public OpRewritePattern<tensor::PadOp> {
     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
index 6924107..d8e22cf 100644 (file)
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -523,7 +523,7 @@ LoopNest mlir::scf::buildLoopNest(
     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 @@ LoopNest mlir::scf::buildLoopNest(
   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(

diff --git a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
index 82ea071..fa3cdb7 100644 (file)
--- a/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
+++ b/mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp
@@ -179,12 +179,6 @@ struct RewriteExtractSliceFromCollapseShapeUsingScfFor
     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,
@@ -193,15 +187,10 @@ struct RewriteExtractSliceFromCollapseShapeUsingScfFor
               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();
   }