Introduce affine terminator

Due to legacy reasons (ML/CFG function separation), regions in affine control
flow operations require contained blocks not to have terminators.  This is
inconsistent with the notion of the block and may complicate code motion
between regions of affine control operations and other regions.

Introduce `affine.terminator`, a special terminator operation that must be used
to terminate blocks inside affine operations and transfers the control back to
he region enclosing the affine operation.  For brevity and readability reasons,
allow `affine.for` and `affine.if` to omit the `affine.terminator` in their
regions when using custom printing and parsing format.  The custom parser
injects the `affine.terminator` if it is missing so as to always have it
present in constructed operations.

Update transformations to account for the presence of terminator.  In
particular, most code motion transformation between loops should leave the
terminator in place, and code motion between loops and non-affine blocks should
drop the terminator.

PiperOrigin-RevId: 240536998

1 files changed, 5 insertions, 5 deletions

diff --git a/mlir/lib/Transforms/LoopTiling.cpp b/mlir/lib/Transforms/LoopTiling.cpp
index d9f74808ad8..c235190b4b7 100644
--- a/mlir/lib/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Transforms/LoopTiling.cpp
@@ -89,10 +89,12 @@ FunctionPassBase *mlir::createLoopTilingPass(uint64_t cacheSizeBytes) {
 }
 
 // Move the loop body of AffineForOp 'src' from 'src' into the specified
-// location in destination's body.
+// location in destination's body, ignoring the terminator.
 static inline void moveLoopBody(AffineForOp src, AffineForOp dest,
                                 Block::iterator loc) {
-  dest.getBody()->getOperations().splice(loc, src.getBody()->getOperations());
+  auto &insts = src.getBody()->getOperations();
+  dest.getBody()->getOperations().splice(loc, insts, insts.begin(),
+                                         std::prev(insts.end()));
 }
 
 // Move the loop body of AffineForOp 'src' from 'src' to the start of dest's
@@ -202,7 +204,6 @@ LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band,
     FuncBuilder b(topLoop);
     // Loop bounds will be set later.
     auto pointLoop = b.create<AffineForOp>(loc, 0, 0);
-    pointLoop.createBody();
     pointLoop.getBody()->getOperations().splice(
         pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
         topLoop);
@@ -217,7 +218,6 @@ LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band,
     FuncBuilder b(topLoop);
     // Loop bounds will be set later.
     auto tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0);
-    tileSpaceLoop.createBody();
     tileSpaceLoop.getBody()->getOperations().splice(
         tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(),
         topLoop);
@@ -264,7 +264,7 @@ static void getTileableBands(Function &f,
     AffineForOp currInst = root;
     do {
       band.push_back(currInst);
-    } while (currInst.getBody()->getOperations().size() == 1 &&
+    } while (currInst.getBody()->getOperations().size() == 2 &&
              (currInst = currInst.getBody()->front().dyn_cast<AffineForOp>()));
     bands->push_back(band);
   };