[mlir] Introduce `linalg.tiled_yield` terminator for `linalg.tiled_loop`.

[pifon2a]

Summary

RFC

[herhut]

Just some comments as I looked at it anyway.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

754

This comment is off. There is no enclosing linalg generic op. Maybe explain that it updates the part of the enclosing tiled_loop ops result specified by the second operand with the values from the first operand.

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

574

If the assumption is that it always is canonicalized away to an in-place update, then the terminator does not have any of these effects. However, I don't know this part well enough to judge whether it models the naive state or the expected bufferized behavior.

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

956 ↗

(On Diff #358963)

Debuggin left-over.

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

307–308

Comment is now off

[pifon2a]

Address the comments.

[pifon2a]

Add a pattern to fold tensor.cast into tiled_yield.

[mehdi_amini]

(Marking as blocking while we're discussing on the RFC)

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
   }];
   let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
 }
 
 def Linalg_TiledLoopOp : Linalg_Op<"tiled_loop", [
      AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      RecursiveSideEffects,
-     SingleBlockImplicitTerminator<"linalg::YieldOp">
+     SingleBlockImplicitTerminator<"linalg::TiledYieldOp">
     ]> {
   let summary = "Linalg tiled loop operation";
   let description = [{
     This is a loop-like operation with additional properties. The arguments
     also include the input and the output tensors or memrefs and the attributes
     to specify the iterator types.
 
     Parsing TiledLoopOp will set all elements of the `iterator_types` attribute
     to "parallel" type, when it is absent from the custom format.
 
     Tensor-based version:
 
     The body region of the loop contains `extract_slice` operations applied to
     every tensor argument of TiledLoopOp.
 
     The body region must contain exactly one block that terminates with
-    `linalg.yield` with the operands resulting from `insert_slice` operations.
+    `linalg.tiled_yield`.
 
     Example:
 
     ```mlir
     %0 = linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
         ins(%lhs, %rhs : tensor<24x64xi8>, tensor<24x64xi8>)
         outs(%out : tensor<24x64xi8>)
         iterators("parallel")
         distribution("block_x") {
       %lhs_sub = tensor.extract_slice %lhs[%i, 0] [%c4, %c64] [1, 1]
           : tensor<24x64xi8> to tensor<?x?xi8>
       %rhs_sub = tensor.extract_slice %rhs[%i, 0] [%c4, %c64] [1, 1]
           : tensor<24x64xi8> to tensor<?x?xi8>
       %out_sub = tensor.extract_slice %out[%i, 0] [%c4, %c64] [1, 1]
           : tensor<24x64xi8> to tensor<?x?xi8>
 
       %result_sub = linalg.generic ...
 
-      %result = tensor.insert_slice %result_sub into %out[%i, 0][%c4, %c64][1, 1]
-        : tensor<?x?xi8> into tensor<24x64xi8>
-      linalg.yield %result : tensor<24x64xi8>
+      linalg.tiled_yield %result_sub to %out_sub : tensor<?x?xi8>
     }
     ```
 
     MemRef-based version:
 
     The body region of the loop contains `subview` operations applied to
     every memref argument of TiledLoopOp.
 
     The body region must contain exactly one block that terminates with
-    `linalg.yield` with no operands.
+    `linalg.tiled_yield` with no operands.
 
     Example:
 
     ```mlir
     linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
         ins(%lhs, %rhs : memref<24x64xi8>, memref<24x64xi8>)
         outs(%out : memref<24x64xi8>)
         iterators("parallel")
         distribution("block_x") {
       %lhs_sub = subview %lhs[%i, 0] [%c4, %c64] [1, 1]
           : memref<24x64xi8> to memref<?x?xi8>
       %rhs_sub = subview %rhs[%i, 0] [%c4, %c64] [1, 1]
           : memref<24x64xi8> to memref<?x?xi8>
       %out_sub = subview %out[%i, 0] [%c4, %c64] [1, 1]
           : memref<24x64xi8> to memref<?x?xi8>
 
       %result_sub = linalg.generic ...
-      linalg.yield
+      linalg.tiled_yield
     }
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$lowerBound,
                        Variadic<Index>:$upperBound,
                        Variadic<Index>:$step,
                        Variadic<AnyType>:$inputs,
OpOperand* findOutputOperand(Value value) {
       return it.getBase();
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
+def Linalg_TiledYieldOp : Linalg_Op<"tiled_yield",
+    [NoSideEffect, ReturnLike, Terminator, SameVariadicOperandSize]>,
+    Arguments<(ins Variadic<AnyType>:$tiles, Variadic<AnyType>:$outputs)> {
+  let summary = "Linalg tiled yield operation";
+  let description = [{
+    `linalg.tiled_yield` is a special terminator operation for the block inside
+    the region of `linalg.tiled_loop` op. It updates the part of the enclosing
+    `linalg.tiled_loop` result specifies by the `outputs` operand with the
+    values from the `tiles` operand.
+
+    Example:
+
+    ```mlir
+    linalg.tiled_loop ... outs(%out_ = %out : tensor<?x?xf32>) {
+      %output = tensor.extract_slice %out_... // or %output = %out_
+      %tile = "some_computation"
+      linalg.tiled_yield %tile in %output : tensor<?x?xf32>
+    ```
+  }];
+  let builders = [OpBuilder<(ins), [{ /* nothing to do */ }]>];
+  let hasCanonicalizer = 1;
+}
+
 def Linalg_IndexOp : Linalg_Op<"index", [NoSideEffect]>,
     Arguments<(ins Confined<I64Attr, [IntMinValue<0>]>:$dim)>,
     Results<(outs Index:$result)> {
   let summary = "linalg index operation";
   let description = [{
     The `linalg.index` operation returns the iteration index of the immediately
     enclosing linalg structured operation for the iteration dimension `dim`. The
     `dim` attribute specifies the position of the accessed dimension in the

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (op.getNumOperands() != 1)
       return op.emitOpError("expected single yield operand (got ")
              << op->getNumOperands() << ")";
     if (op.getOperand(0).getType() !=
         padTensorOp.getType().cast<ShapedType>().getElementType())
       return op.emitOpError("expected yield type to match shape element type");
     return success();
   }
 
-  if (auto tiledLoopOp = dyn_cast<linalg::TiledLoopOp>(parentOp)) {
-    // Check if output args with tensor types match results types.
-    SmallVector<Value, 2> tensorOuts;
-    llvm::copy_if(
-        tiledLoopOp.outputs(), std::back_inserter(tensorOuts),
-        [&](Value out) { return out.getType().isa<RankedTensorType>(); });
-    if (tensorOuts.size() != op.values().size())
-      return op.emitOpError("expected number of tensor output args = ")
-             << tensorOuts.size() << " to match the number of yield operands = "
-             << op.values().size();
-
-    TypeRange tensorTypes(llvm::makeArrayRef(tensorOuts));
-    for (auto &item :
-         llvm::enumerate(llvm::zip(tensorTypes, op.getOperandTypes()))) {
-      Type outType, resultType;
-      unsigned index = item.index();
-      std::tie(outType, resultType) = item.value();
-      if (outType != resultType)
-        return op.emitOpError("expected yield operand ")
-               << index << " with type = " << resultType
-               << " to match output arg type = " << outType;
-    }
-    return success();
-  }
   return op.emitOpError("expected parent op with LinalgOp interface");
 }
 
 //===----------------------------------------------------------------------===//
 // TiledLoopOp
 //===----------------------------------------------------------------------===//
 
 void TiledLoopOp::build(OpBuilder &builder, OperationState &result,
struct TiledLoopResultsFolder : public OpRewritePattern<linalg::TiledLoopOp> {
   using OpRewritePattern<linalg::TiledLoopOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(linalg::TiledLoopOp tiledLoop,
                                 PatternRewriter &rewriter) const final {
     if (tiledLoop.getNumResults() == 0)
       return failure();
 
     Block *block = tiledLoop.getBody();
-    auto yieldOp = cast<linalg::YieldOp>(block->getTerminator());
+    auto yieldOp = cast<linalg::TiledYieldOp>(block->getTerminator());
 
     // Match the pattern and collect output buffers that will replace the output
     // tensors and also the ops that will be ignored when cloning the body.
-    SmallVector<Value, 2> newOutputOperands, newYieldArgs;
+    SmallVector<Value, 2> newOutputOperands, newYieldTileArgs, newYieldOutArgs;
     int resultId = 0;
     // Store ids of the corresponding old and new output operands.
     SmallVector<int64_t, 2> oldOutputIdToNew(tiledLoop.outputs().size(),
                                              kNoMatch);
     // Store ids of the corresponding old and new results.
     SmallVector<int64_t, 2> oldResultIdToNew(tiledLoop.getNumResults(),
                                              kNoMatch);
     SmallVector<Value, 2> resultReplacement(tiledLoop.getNumResults());
     for (auto en : llvm::enumerate(
              llvm::zip(tiledLoop.outputs(), tiledLoop.getRegionOutputArgs()))) {
       size_t index = en.index();
       Value out = std::get<0>(en.value());
       Value outRegionArg = std::get<1>(en.value());
 
       if (!out.getType().isa<RankedTensorType>()) {
         oldOutputIdToNew[index] = newOutputOperands.size();
         newOutputOperands.push_back(out);
         continue;
       }
       Value result = tiledLoop.getResult(resultId);
-      Value yieldArg = yieldOp.getOperand(resultId);
-      if (yieldArg != outRegionArg || !result.use_empty()) {
+      Value yieldTileArg = yieldOp.tiles()[resultId];
+      Value yieldOutArg = yieldOp.outputs()[resultId];
+      if (yieldTileArg != outRegionArg || !result.use_empty()) {
         oldOutputIdToNew[index] = newOutputOperands.size();
-        oldResultIdToNew[resultId] = newYieldArgs.size();
+        oldResultIdToNew[resultId] = newYieldTileArgs.size();
         resultReplacement[resultId] = out;
         newOutputOperands.push_back(out);
-        newYieldArgs.push_back(yieldArg);
+        newYieldTileArgs.push_back(yieldTileArg);
+        newYieldOutArgs.push_back(yieldOutArg);
       }
       ++resultId;
     }
     if (newOutputOperands.size() == tiledLoop.outputs().size())
       return failure();
 
     Location loc = tiledLoop.getLoc();
     auto newTiledLoop = rewriter.create<TiledLoopOp>(
for (const auto &en : llvm::enumerate(oldOutputIdToNew)) {
       else
         bvm.map(tiledLoop.getRegionOutputArgs()[en.index()],
                 tiledLoop.outputs()[en.index()]);
     }
     OpBuilder innerBuilder =
         OpBuilder::atBlockEnd(newTiledLoop.getBody(), rewriter.getListener());
     for (auto &op : tiledLoop.getBody()->without_terminator())
       innerBuilder.clone(op, bvm);
-    innerBuilder.create<linalg::YieldOp>(
-        loc, llvm::to_vector<2>(llvm::map_range(
-                 newYieldArgs, [&](Value arg) { return bvm.lookup(arg); })));
+    innerBuilder.create<linalg::TiledYieldOp>(
+        loc,
+        llvm::to_vector<2>(llvm::map_range(
+            newYieldTileArgs, [&](Value arg) { return bvm.lookup(arg); })),
+        llvm::to_vector<2>(llvm::map_range(
+            newYieldOutArgs, [&](Value arg) { return bvm.lookup(arg); })));
 
     for (const auto &en : llvm::enumerate(oldResultIdToNew))
       if (en.value() != kNoMatch)
         resultReplacement[en.index()] = newTiledLoop.getResult(en.value());
     rewriter.replaceOp(tiledLoop, resultReplacement);
 
     return success();
   }
 };
 } // namespace
 
 void TiledLoopOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {
   results.insert<TiledLoopInputsFolder, TiledLoopResultsFolder>(context);
 }
 
 LogicalResult TiledLoopOp::fold(ArrayRef<Attribute>,
                                 SmallVectorImpl<OpFoldResult> &) {
   return foldMemRefCastInTiledLoopOp(*this);
 }
 
 //===----------------------------------------------------------------------===//
+// TiledYieldOp
+//===----------------------------------------------------------------------===//
+
+static void print(OpAsmPrinter &p, TiledYieldOp op) {
+  p << op.getOperationName();
+
+  if (!op.tiles().empty()) {
+    llvm::interleaveComma(llvm::zip(op.tiles(), op.outputs()), p, [&](auto it) {
+      p << ' ' << std::get<0>(it) << " in " << std::get<1>(it) << " : "
+        << std::get<1>(it).getType();
+    });
+  }
+  p.printOptionalAttrDict(op->getAttrs());
+}
+
+static ParseResult parseTiledYieldOp(OpAsmParser &parser,
+                                     OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 4> tiles, outputs;
+  SmallVector<Type, 4> types;
+
+  OpAsmParser::OperandType tile;
+  while (parser.parseOptionalOperand(tile).hasValue()) {
+    Type type;
+    OpAsmParser::OperandType output;
+    if (parser.parseKeyword("in") || parser.parseOperand(output) ||
+        parser.parseColon() || parser.parseType(type))
+      return failure();
+    tiles.push_back(tile);
+    outputs.push_back(output);
+    types.push_back(type);
+    parser.parseOptionalComma();
+  }
+  llvm::SMLoc loc = parser.getCurrentLocation();
+  if (parser.resolveOperands(tiles, types, loc, result.operands) ||
+      parser.resolveOperands(outputs, types, loc, result.operands))
+    return failure();
+
+  // Parse optional attributes.
+  parser.parseOptionalAttrDict(result.attributes);
+
+  return success();
+}
+
+static LogicalResult verify(TiledYieldOp op) {
+  // Check if output args with tensor types match results types.
+  auto loop = op->getParentOfType<TiledLoopOp>();
+  SmallVector<Value, 2> loopTensorOuts;
+  llvm::copy_if(
+      loop.outputs(), std::back_inserter(loopTensorOuts),
+      [&](Value out) { return out.getType().isa<RankedTensorType>(); });
+  if (loopTensorOuts.size() != op.tiles().size())
+    return op.emitOpError("expected number of tensor output args = ")
+           << loopTensorOuts.size()
+           << " to match the number of yield operands = " << op.tiles().size();
+
+  // Check if the `tiles` args types match the `outputs` args types.
+  SmallVector<Value, 2> loopTensorOutsBlockArgs;
+  llvm::copy_if(
+      loop.getRegionOutputArgs(), std::back_inserter(loopTensorOutsBlockArgs),
+      [&](Value out) { return out.getType().isa<RankedTensorType>(); });
+  for (auto en : llvm::enumerate(
+           llvm::zip(op.tiles(), op.outputs(), loopTensorOutsBlockArgs))) {
+    size_t index = en.index();
+    Type tileType = std::get<0>(en.value()).getType();
+    Value yieldOut = std::get<1>(en.value());
+    Type yieldOutType = yieldOut.getType();
+
+    if (tileType != yieldOutType)
+      return op.emitOpError("expected tile operand with type = ")
+             << tileType << " to match output type = " << yieldOutType;
+
+    // Check if yieldOut is either an output bbArg or a slice of it.
+    Value src = yieldOut;
+    if (auto extractSlice = llvm::dyn_cast_or_null<tensor::ExtractSliceOp>(
+            yieldOut.getDefiningOp()))
+      src = extractSlice.source();
+
+    Value loopBlockArg = std::get<2>(en.value());
+    if (src != loopBlockArg)
+      return op.emitOpError("expected output ")
+             << index << " to be a subset of the corresponding block argument";
+  }
+  return success();
+}
+
+namespace {
+/// Pattern to rewrite TiledYieldOp with tensor::CastOp arguments.
+///
+/// Example:
+/// ```
+///   %TILE_CAST = tensor.cast %TILE : tensor<16x16xf32> to tensor<?x?xf32>
+///   %OUT_CAST = tensor.cast %OUT : tensor<16x16xf32> to tensor<?x?xf32>
+///   linalg.tiled_yield %TILE_CAST in %OUT_CAST : tensor<?x?xf32>
+/// ```
+/// is rewritten into:
+/// ```
+///   linalg.tiled_yield %TILE in %OUT : tensor<16x16xf32>
+/// ```
+class TensorYieldOpCastFolder final : public OpRewritePattern<TiledYieldOp> {
+public:
+  using OpRewritePattern<TiledYieldOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TiledYieldOp tiledYieldOp,
+                                PatternRewriter &rewriter) const override {
+    bool foundTensorCasts = false;
+    SmallVector<Value, 4> tiles, outputs;
+    for (auto item : llvm::zip(tiledYieldOp.tiles(), tiledYieldOp.outputs())) {
+      Value tile, output;
+      std::tie(tile, output) = item;
+      auto tileCast = tile.getDefiningOp<tensor::CastOp>();
+      if (!tileCast)
+        continue;
+
+      auto outputCast = output.getDefiningOp<tensor::CastOp>();
+      if (!outputCast)
+        continue;
+
+      if (tileCast.source().getType() != outputCast.source().getType()) {
+        tiles.push_back(tile);
+        outputs.push_back(output);
+        continue;
+      }
+      foundTensorCasts = true;
+      tiles.push_back(tileCast.source());
+      outputs.push_back(outputCast.source());
+    }
+    if (!foundTensorCasts)
+      return failure();
+    rewriter.replaceOpWithNewOp<TiledYieldOp>(tiledYieldOp, tiles, outputs);
+    return success();
+  }
+};
+} // namespace
+
+void TiledYieldOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                               MLIRContext *context) {
+  results.add<TensorYieldOpCastFolder>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // IndexOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(IndexOp op) {
   auto linalgOp = dyn_cast<LinalgOp>(op->getParentOp());
   if (!linalgOp)
     return op.emitOpError("expected parent op with LinalgOp interface");
   if (linalgOp.getNumLoops() <= op.dim())

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

return
           ExtractSliceOp,
           scf::ForOp,
           InsertSliceOp,
           InitTensorOp,
           LinalgOp,
           ReturnOp,
           TiledLoopOp,
           VectorTransferOpInterface,
+          linalg::TiledYieldOp,
           linalg::YieldOp,
           scf::YieldOp>(op)
       // clang-format on
       || (none_of(op->getResultTypes(), isaTensor) &&
           none_of(op->getOperandTypes(), isaTensor));
 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`

 /// bufferization aliasing behavior, which indicates that the op must allocate
 /// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpResult.
 static Optional<OpResult> getAliasingOpResult(OpOperand &opOperand) {
   if (!hasKnownBufferizationAliasingBehavior(opOperand.getOwner()))
     return None;
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       // These terminators legitimately have no result.
-      .Case<ReturnOp, linalg::YieldOp, scf::YieldOp>(
+      .Case<ReturnOp, linalg::TiledYieldOp, linalg::YieldOp, scf::YieldOp>(
           [&](auto op) { return OpResult(); })
       // ConstantOp is never inplaceable.
       .Case([&](ConstantOp op) { return op->getResult(0); })
       // ExtractSliceOp is different: its result is not inplaceable on op.source
       // but when bufferized inplace, the result is an aliasing subregion of
       // op.source.
       .Case([&](ExtractSliceOp op) { return op->getResult(0); })
       // All other ops, including scf::ForOp, return the result of
static bool bufferizesToMemoryRead(OpOperand &opOperand) {
   // of the matching bbArg may. It is the responsibility of the caller to
   // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
   // conservative.
   if (auto callOp = dyn_cast<CallOpInterface>(opOperand.getOwner()))
     return true;
   if (auto linalgOp = dyn_cast<LinalgOp>(opOperand.getOwner()))
     return linalgOp.isInputTensor(&opOperand) ||
            linalgOp.isInitTensor(&opOperand);
+  // This is questionable. Should we consider TiledYieldOp as an op that
+  // bufferizes to "read" for the `tile` args and to "write" for the `output`
+  // args?
+  if (isa<TiledYieldOp>(opOperand.getOwner()))
+    return false;
   // All other cases are considered to bufferize to memory reads.
   // In particular, terminators are often the last use and need to be considered
   // as reads to return the proper value and avoid WAW clobbers.
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write.
 /// If inPlaceSpec is different from InPlaceSpec::None, additionally require the
 /// write to match the inplace specification.
 static bool
 bufferizesToMemoryWrite(OpOperand &opOperand,
                         InPlaceSpec inPlaceSpec = InPlaceSpec::None) {
   // These terminators are not writes.
-  if (isa<ReturnOp, linalg::YieldOp, scf::YieldOp>(opOperand.getOwner()))
+  if (isa<ReturnOp, linalg::TiledYieldOp, linalg::YieldOp, scf::YieldOp>(
+          opOperand.getOwner()))
     return false;
   // ExtractSliceOp alone doesn't bufferize to a memory write, one of its uses
   // may.
   if (isa<ExtractSliceOp>(opOperand.getOwner()))
     return false;
   // CallOpInterface alone doesn't bufferize to a memory write, one of the uses
   // of the matching bbArg may. It is the responsibility of the caller to
   // inspect bbArgs. In the absence of a BufferizationAliasInfo, we need to be
static LogicalResult bufferize(OpBuilder &b, linalg::YieldOp yieldOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(yieldOp);
   // No tensors -> success.
   if (!llvm::any_of(yieldOp.getOperandTypes(), isaTensor))
     return success();
-  // linalg::YieldOp nested under TiledLoop must just canonicalize.
-  if (yieldOp->getParentOfType<TiledLoopOp>())
-    return success();
   llvm_unreachable("unexpected yieldOp");
 }
 
 /// Bufferization for tensor::ExtractOp just translate to memref.load, it only
 /// reads the tensor.
 static LogicalResult bufferize(OpBuilder &b, tensor::ExtractOp extractOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(extractOp);
 
   Location loc = extractOp.getLoc();
   Value srcMemref = lookup(bvm, extractOp.tensor());
   Value l = b.create<memref::LoadOp>(loc, srcMemref, extractOp.indices());
   extractOp.replaceAllUsesWith(l);
   return success();
 }
+
+/// Bufferization for linalg::TiledYieldOp just results in later
+/// canonicalization.
+static LogicalResult bufferize(OpBuilder &b, linalg::TiledYieldOp yieldOp,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo) {
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 ///
 /// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace.
 /// ===========================================================
 ///
WalkResult result =
             InitTensorOp,
             InsertSliceOp,
             tensor::ExtractOp,
             LinalgOp,
             ReturnOp,
             TiledLoopOp,
             VectorTransferOpInterface,
             linalg::YieldOp,
+            linalg::TiledYieldOp,
             scf::YieldOp>([&](auto op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo);
       })
       .Case([&](CallOpInterface op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo, bufferizedFunctionTypes);
       })

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

static Value insertSliceIntoTensor(OpBuilder &b, Location loc,
                                    Value dest) {
   return b.create<tensor::InsertSliceOp>(
       loc, sliceOp.source().getType(), source, dest, sliceOp.offsets(),
       sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
       sliceOp.static_sizes(), sliceOp.static_strides());
 }
 
 template <typename LoopTy>
+static SmallVector<Value, 4>
+collectLoopYieldArgs(OpBuilder &b, LinalgOp clonedOp,
+                     ArrayRef<Value> tiledOperands,
+                     SmallVectorImpl<Value> &tensorResults) {
+
+  Location loc = clonedOp.getLoc();
+  SmallVector<Value, 4> yieldArgs;
+  unsigned resultIdx = 0;
+  for (OpOperand *opOperand : clonedOp.getOutputTensorOperands()) {
+    // TODO: use an interface/adaptor to avoid leaking position in
+    // `tiledOperands`.
+    Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
+    // Insert a insert_slice for each output tensor.
+    if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
+      yieldArgs.push_back(insertSliceIntoTensor(
+          b, loc, sliceOp, clonedOp->getResult(resultIdx), sliceOp.source()));
+    } else {
+      yieldArgs.push_back(clonedOp->getResult(resultIdx));
+    }
+    ++resultIdx;
+  }
+  tensorResults = yieldArgs;
+  return yieldArgs;
+}
+
+template <>
+SmallVector<Value, 4>
+collectLoopYieldArgs<TiledLoopOp>(OpBuilder &b, LinalgOp clonedOp,
+                                  ArrayRef<Value> tiledOperands,
+                                  SmallVectorImpl<Value> &tensorResults) {
+  auto outputTensorOperands = clonedOp.getOutputTensorOperands();
+  size_t numOutputTensors = outputTensorOperands.size();
+
+  SmallVector<Value, 4> yieldArgs(clonedOp->getResults());
+  auto tiledOutputOperands = tiledOperands.take_back(numOutputTensors);
+  yieldArgs.append(tiledOutputOperands.begin(), tiledOutputOperands.end());
+
+  return yieldArgs;
+}
+
+template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ValueRange tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
   if (llvm::all_of(tileSizes, isZero))
if (!options.interchangeVector.empty()) {
     invPermutationMap = inversePermutation(
         AffineMap::getPermutationMap(interchangeVector, b.getContext()));
     assert(invPermutationMap);
     applyPermutationToVector(loopRanges, interchangeVector);
     applyPermutationToVector(iteratorTypes, interchangeVector);
   }
 
   // 2. Create the tiled loops.
-  LinalgOp res = op;
+  LinalgOp clonedOp = op;
   SmallVector<Value, 4> ivs, tensorResults;
   auto tiledLoopBodyBuilder = [&](OpBuilder &b, Location loc,
                                   ValueRange localIvs,
                                   ValueRange iterArgs) -> scf::ValueVector {
     ivs.assign(localIvs.begin(), localIvs.end());
 
     // When an `interchangeVector` is present, it has been applied to the
     // loop ranges and the iterator types. Apply its inverse to the
auto tiledLoopBodyBuilder = [&](OpBuilder &b, Location loc,
 
     // TODO: use an interface/adaptor to avoid leaking position in
     // `tiledOperands`.
     SmallVector<Type, 4> resultTensorTypes;
     for (OpOperand *opOperand : op.getOutputTensorOperands())
       resultTensorTypes.push_back(
           tiledOperands[opOperand->getOperandNumber()].getType());
 
-    res = op.clone(b, loc, resultTensorTypes, tiledOperands);
+    clonedOp = op.clone(b, loc, resultTensorTypes, tiledOperands);
 
-    // Insert a insert_slice for each output tensor.
-    unsigned resultIdx = 0;
-    for (OpOperand *opOperand : op.getOutputTensorOperands()) {
-      // TODO: use an interface/adaptor to avoid leaking position in
-      // `tiledOperands`.
-      Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
-      if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
-        tensorResults.push_back(insertSliceIntoTensor(
-            b, loc, sliceOp, res->getResult(resultIdx), sliceOp.source()));
-      } else {
-        tensorResults.push_back(res->getResult(resultIdx));
-      }
-      ++resultIdx;
-    }
-    return scf::ValueVector(tensorResults.begin(), tensorResults.end());
+    auto yieldArgs =
+        collectLoopYieldArgs<LoopTy>(b, clonedOp, tiledOperands, tensorResults);
+    return {yieldArgs.begin(), yieldArgs.end()};
   };
   GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
                                  tiledLoopBodyBuilder, options.distribution,
                                  options.distributionTypes);
 
   // 3. Transform IndexOp results w.r.t. the tiling.
-  transformIndexOps(b, res, ivs, loopIndexToRangeIndex);
+  transformIndexOps(b, clonedOp, ivs, loopIndexToRangeIndex);
 
   // 4. Gather the newly created loops and return them with the new op.
   SmallVector<Operation *, 8> loops;
   loops.reserve(ivs.size());
   for (auto iv : ivs) {
     if (iv.isa<BlockArgument>()) {
       loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
       assert(loops.back() && "no owner found for induction variable!");
     } else {
       // TODO: Instead of doing this, try to recover the ops used instead of the
       // loop.
       loops.push_back(nullptr);
     }
   }
 
   // 5. Get the tensor results from the outermost loop if available. Otherwise
   // use the previously captured `tensorResults`.
   Operation *outermostLoop = nullptr;
   for (Operation *loop : loops)
     if ((outermostLoop = loop))
       break;
 
-  return TiledLinalgOp{
-      res, loops, outermostLoop ? outermostLoop->getResults() : tensorResults};
+  return TiledLinalgOp{clonedOp, loops,
+                       outermostLoop ? outermostLoop->getResults()
+                                     : tensorResults};
 }
 
 template <typename LoopTy>
 Optional<TiledLinalgOp> static tileLinalgOpImpl(
     OpBuilder &b, LinalgOp op, const LinalgTilingOptions &options) {
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
 
void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
   ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
   memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
   memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
   PadTensorOp::getCanonicalizationPatterns(patterns, ctx);
   ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
+  linalg::TiledYieldOp::getCanonicalizationPatterns(patterns, ctx);
   CanonicalizationPatternList<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::insert(patterns);
 }
 
 /// Populate the given list with patterns that apply Linalg tiling.
 static void insertTilingPatterns(RewritePatternSet &patterns,

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

void GenerateLoopNest<TiledLoopOp>::doit(
   SmallVector<ProcInfo, 2> procInfo;
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(loopRanges, lbs, ubs, steps);
 
   auto wrappedBuilderFn = [&](OpBuilder &nestedBuilder, Location nestedLoc,
                               ValueRange ivs, ValueRange inputs,
                               ValueRange outputs) {
     SmallVector<Value> outputTensors = linalgOp.getOutputTensorOperands();
-    scf::ValueVector results =
+    scf::ValueVector yieldArgs =
         bodyBuilderFn(nestedBuilder, nestedLoc, ivs, outputTensors);
-    nestedBuilder.create<linalg::YieldOp>(nestedLoc, results);
+    auto yieldArgsRef = llvm::makeArrayRef(yieldArgs);
+    nestedBuilder.create<linalg::TiledYieldOp>(
+        nestedLoc, yieldArgsRef.take_front(outputTensors.size()),
+        yieldArgsRef.drop_front(outputTensors.size()));
   };
 
   SmallVector<Value> inputOperands = linalgOp.getInputOperands();
   SmallVector<Value> outputOperands = linalgOp.getOutputOperands();
   auto tiledLoop =
       b.create<TiledLoopOp>(loc, lbs, ubs, steps, inputOperands, outputOperands,
                             b.getArrayAttr(iteratorTypes), wrappedBuilderFn);
   if (!distributionTypes.empty())

mlir/test/Dialect/Linalg/canonicalize.mlir

func @memref_cast_into_tiled_loop(%arg0: memref<192xf32>)  {
   // CHECK: linalg.tiled_loop
   // CHECK-SAME: outs (%{{.*}} = %{{.*}}: memref<192xf32>)
   linalg.tiled_loop (%arg3) = (%c0) to (%c192) step (%c24)
     outs (%out = %0: memref<192xf32, #map>) {
     %14 = affine.min affine_map<(d0) -> (-d0 + 192, 24)>(%arg3)
     %16 = memref.subview %out[%arg3] [%14] [1]
       : memref<192xf32, #map> to memref<?xf32, #map>
     linalg.fill(%cst, %16) : f32, memref<?xf32, #map>
-    linalg.yield
+    linalg.tiled_yield
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @tiled_yield_fold_tensor_cast
+func @tiled_yield_fold_tensor_cast(%init: tensor<200xf32>) -> tensor<200xf32> {
+  %c0 = constant 0 : index
+  %c10 = constant 10 : index
+  %c200 = constant 200 : index
+  %cst = constant 0.000000e+00 : f32
+  %0 = linalg.tiled_loop (%i) = (%c0) to (%c200) step (%c10)
+      outs (%init_ = %init: tensor<200xf32>) {
+    %sub_init = tensor.extract_slice %init_[%i] [%c10] [1]
+      : tensor<200xf32> to tensor<?xf32>
+    %sub_fill = linalg.fill(%cst, %sub_init)
+        : f32, tensor<?xf32> -> tensor<?xf32>
+    linalg.tiled_yield %sub_fill in %sub_init : tensor<?xf32>
+  }
+  return %0 : tensor<200xf32>
+}
+// CHECK: linalg.tiled_yield %{{.*}} in %{{.*}} : tensor<10xf32>
+
+// -----
+
 // CHECK-LABEL: zero_rank_reshape_multi
 func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
   // CHECK: return %arg0
   %0 = linalg.tensor_expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
   %1 = linalg.tensor_expand_shape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
   %2 = linalg.tensor_collapse_shape %1 [] : tensor<1x1xf32> into tensor<f32>
   return %2 : tensor<f32>
 }
func @fold_tiled_loop_results(%A: memref<48xf32>, %B: tensor<48xf32>,
   %c24 = constant 24 : index
   %c48 = constant 48 : index
   %useful, %useless = linalg.tiled_loop (%i) = (%c0) to (%c48) step (%c24)
       ins (%A_ = %A: memref<48xf32>)
       outs (%B_ = %B: tensor<48xf32>,
             %CT_ = %C_tensor: tensor<48xf32>,
             %C_ = %C: memref<48xf32>) {
         %result = call @foo(%A_, %B_, %C_)
-          : (memref<48xf32>, tensor<48xf32>, memref<48xf32>)-> (tensor<48xf32>)
-    linalg.yield %result, %CT_ : tensor<48xf32>, tensor<48xf32>
+          : (memref<48xf32>, tensor<48xf32>, memref<48xf32>) -> (tensor<48xf32>)
+        linalg.tiled_yield %result in %B_ : tensor<48xf32>,
+                           %CT_ in %CT_ :  tensor<48xf32>
   }
   return %useful : tensor<48xf32>
 }
 
 // CHECK-LABEL: func @fold_tiled_loop_results(
 // CHECK-SAME:   %[[A:.*]]: [[BUF_TY:memref<48xf32>]], %[[B:.*]]: [[TY:tensor<48xf32>]],
 // CHECK-SAME:   %[[C:.*]]: [[BUF_TY]],  %[[C_TENSOR:.*]]: [[TY]]) -> [[TY]] {
 
 // CHECK-DAG:  %[[C0:.*]] = constant 0 : index
 // CHECK-DAG:  %[[C24:.*]] = constant 24 : index
 // CHECK-DAG:  %[[C48:.*]] = constant 48 : index
 
 // CHECK-NOT: %{{.*}} = linalg.tiled_loop
 // CHECK:  %[[RESULT:.*]] = linalg.tiled_loop (%{{.*}}) = (%[[C0]])
 // CHECK-SAME: to (%[[C48]]) step (%[[C24]])
 // CHECK-SAME: ins (%[[A_:.*]] = %[[A]]: [[BUF_TY]])
 // CHECK-SAME: outs (%[[B_:.*]] = %[[B]]: [[TY]], %[[C_:.*]] = %[[C]]: [[BUF_TY]]) {
 // CHECK-NEXT:   %[[RES:.*]] = call @foo(%[[A_]], %[[B_]], %[[C_]])
-// CHECK-NEXT:   linalg.yield %[[RES]] :
+// CHECK-NEXT:   linalg.tiled_yield %[[RES]] in %[[B_]]
 
 // CHECK: return %[[RESULT]]
 
 // -----
 
 func private @foo(%A: memref<192xf32>, %B: tensor<192xf32>) -> tensor<192xf32>
 
 func @fold_tiled_loop_inputs(%A: memref<192xf32>, %A_tensor: tensor<192xf32>,
                              %B_tensor: tensor<192xf32>) -> tensor<192xf32> {
   %c0 = constant 0 : index
   %c24 = constant 24 : index
   %c192 = constant 192 : index
   %result = linalg.tiled_loop (%i) = (%c0) to (%c192) step (%c24)
       ins (%A_ = %A: memref<192xf32>, %AT_ = %A_tensor: tensor<192xf32>)
       outs (%BT_ = %B_tensor: tensor<192xf32>) {
     %0 = call @foo(%A_, %BT_) : (memref<192xf32>, tensor<192xf32>) -> tensor<192xf32>
-    linalg.yield %0 : tensor<192xf32>
+    linalg.tiled_yield %0 in %BT_ : tensor<192xf32>
   }
   return %result : tensor<192xf32>
 }
 
 // CHECK-LABEL: func @fold_tiled_loop_inputs
 // CHECK: %[[RESULT:.*]] = linalg.tiled_loop
 // CHECK-SAME: ins (%{{.*}} = %{{.*}}: memref<192xf32>)
 

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

%r1 = scf.for %i = %lb to %ub step %step iter_args(%t = %A) -> (tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
     scf.yield %t : tensor<?xf32>
   }
 
   // %r2 must be out of place because one use of %t in the subsequent production 
   // of %r3 is read.
   //      CHECK: linalg.tiled_loop
   // CHECK-NEXT: call
-  // CHECK-NEXT: linalg.yield
+  // CHECK-NEXT: linalg.tiled_yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["false"]}
   %r2 = linalg.tiled_loop (%i) = (%lb) to (%ub) step (%step)
         ins()
         outs(%t = %B: tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
-    linalg.yield %t : tensor<?xf32>
+    linalg.tiled_yield %t in %t : tensor<?xf32>
   }
 
   // %r3 bufferizes inplace fine.
   //      CHECK: linalg.tiled_loop
   // CHECK-NEXT: call
-  // CHECK-NEXT: linalg.yield
+  // CHECK-NEXT: linalg.tiled_yield
   // CHECK-NEXT: {__inplace_results_attr__ = ["true"]}
   %r3 = linalg.tiled_loop (%i) = (%lb) to (%ub) step (%step)
         ins()
         outs(%t = %B: tensor<?xf32>) {
     call @some_use(%t) : (tensor<?xf32>) -> ()
-    linalg.yield %t : tensor<?xf32>
+    linalg.tiled_yield %t in %t : tensor<?xf32>
   }
 
   return %r1, %r3: tensor<?xf32>, tensor<?xf32>
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @tiled_dot(%A: tensor<?xf32>, %B: tensor<?xf32>, %c: tensor<f32> {linalg.inplaceable = true},
   %c3 = constant 3 : index
   %c0 = constant 0 : index
 
   //     CHECK: %[[M:.*]] = memref.dim %[[A]], {{.*}} : memref<?xf32, #[[$DYN_1D_MAP:.*]]>
   %0 = tensor.dim %A, %c0 : tensor<?xf32>
 
   //     CHECK: linalg.tiled_loop {{.*}} to (%[[M]]) {{.*}} %[[A]]{{.*}}%[[B]]{{.*}}outs{{.*}}%[[c]]
   %1 = linalg.tiled_loop (%arg3) = (%c0) to (%0) step (%c3)
-       ins (%arg4 = %A: tensor<?xf32>, %use = %effecting : memref<?xf32>, %arg5 = %B: tensor<?xf32>)
+       ins (%arg4 = %A: tensor<?xf32>,
+            %use = %effecting : memref<?xf32>,
+            %arg5 = %B: tensor<?xf32>)
       outs (%arg6 = %c: tensor<f32>)
-      iterators["reduction"]
-  {
+      iterators["reduction"] {
     // CHECK-NOT:   alloc
 
     %2 = tensor.dim %arg4, %c0 : tensor<?xf32>
     %3 = affine.min #TILE_MAP(%arg3)[%2]
 
     //     CHECK:   %[[SV_A:.*]] = memref.subview {{.*}}
     %4 = tensor.extract_slice %arg4[%arg3] [%3] [1] : tensor<?xf32> to tensor<?xf32>
     %5 = tensor.dim %arg5, %c0 : tensor<?xf32>
     %6 = affine.min #TILE_MAP(%arg3)[%5]
 
     //     CHECK:   %[[SV_B:.*]] = memref.subview {{.*}}
     %7 = tensor.extract_slice %arg5[%arg3] [%6] [1] : tensor<?xf32> to tensor<?xf32>
 
     //     CHECK:   linalg.dot ins(%[[SV_A]], %[[SV_B]] : memref<?xf32, #[[$DYN_1D_MAP:.*]]>, memref<?xf32, #[[$DYN_1D_MAP:.*]]>) outs(%{{.*}} : memref<f32, #[[$DYN_0D_MAP]]>)
     %8 = linalg.dot ins(%4, %7 : tensor<?xf32>, tensor<?xf32>) outs(%arg6 : tensor<f32>) -> tensor<f32>
 
     //     CHECK:   call @some_use(%{{.*}}) : (memref<?xf32>) -> ()
     call @some_use(%use) : (memref<?xf32>) -> ()
 
-    linalg.yield %8 : tensor<f32>
-    //     CHECK:   linalg.yield
+    linalg.tiled_yield %8 in %arg6 : tensor<f32>
+    //     CHECK:   linalg.tiled_yield
     // CHECK-NOT:   tensor
   }
 
   //     CHECK: return
   // CHECK-NOT: tensor
   return %1 : tensor<f32>
 }
 

mlir/test/Dialect/Linalg/distribute-tiled-loop.mlir

 // RUN: mlir-opt -test-linalg-distribution %s | FileCheck %s
 
 func private @foo(%A: tensor<64x64xf32>,
                   %B: tensor<64x64xf32>) -> tensor<64x64xf32>
 
 func @distribute_for_gpu(%A: tensor<64x64xf32>,
                          %B: tensor<64x64xf32>) -> tensor<64x64xf32> {
   %c0 = constant 0 : index
   %c16 = constant 16 : index
   %c64 = constant 64 : index
   %c24 = constant 24 : index
   %0 = linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c64, %c64) step (%c24, %c16)
       ins (%A_ = %A: tensor<64x64xf32>) outs (%B_ = %B:tensor<64x64xf32>)
       distribution ["block_x", "block_y"] {
     %0 = call @foo(%A_, %B_)
       : (tensor<64x64xf32>, tensor<64x64xf32>) -> tensor<64x64xf32>
-    linalg.yield %0 : tensor<64x64xf32>
+    linalg.tiled_yield %0 in %B_ : tensor<64x64xf32>
   }
   return %0 : tensor<64x64xf32>
 }
 
 // CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 * 24)>
 // CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 * 16)>
 
 // CHECK-LABEL: func @distribute_for_gpu

mlir/test/Dialect/Linalg/fusion-tensor-pattern.mlir

 // TLOOP:    %[[AB_INIT_SUB:.*]] = tensor.extract_slice %[[AB_INIT_]][%[[IV0]], 0]
 
 // TLOOP:    %[[AB_SUB:.*]] = linalg.matmul
 // TLOOP-SAME:  ins(%[[A_SUB]], %[[B_]] : {{.*}}) outs(%[[AB_INIT_SUB]]
 
 // TLOOP:    %[[DIM_B_1:.*]] = tensor.dim %[[B_]], %[[C1]] : [[TY]]
 // TLOOP:    %[[DIM_C_1:.*]] = tensor.dim %[[C_]], %[[C1]] : [[TY]]
 
-// TLOOP:    %[[ABC_SUB_:.*]] = linalg.tiled_loop (%[[IV1:.*]], %[[IV2:.*]]) = 
+// TLOOP:    %[[ABC_SUB:.*]] = linalg.tiled_loop (%[[IV1:.*]], %[[IV2:.*]]) =
 // TLOOP-SAME: (%[[C0]], %[[C0]]) to (%[[DIM_C_1]], %[[DIM_B_1]])
 // TLOOP-SAME: step (%[[C64]], %[[C16]]) 
 // TLOOP-SAME: ins (%[[AB_SUB_:.*]] = %[[AB_SUB]]: [[TY]],
 // TLOOP-SAME:      %[[C__:.*]] = %[[C_]]: [[TY]])
 // TLOOP-SAME: outs (%[[ABC_INIT_SUB_:.*]] = %[[ABC_INIT_SUB]]: [[TY]])
 // TLOOP-SAME: iterators["parallel", "reduction"] {
 
 // TLOOP:      %[[AB_SUB_SUB:.*]] = tensor.extract_slice %[[AB_SUB_]][0, %[[IV2]]]
 // TLOOP:      %[[C__SUB:.*]] = tensor.extract_slice %[[C__]][%[[IV2]], %[[IV1]]]
-// TLOOP:      %[[ABS_INIT_SUB_SUB:.*]] = tensor.extract_slice %[[ABC_INIT_SUB_]][0, %[[IV1]]]
+// TLOOP:      %[[ABC_INIT_SUB_SUB:.*]] = tensor.extract_slice %[[ABC_INIT_SUB_]][0, %[[IV1]]]
 
 // TLOOP:      %[[ABC_SUB_SUB:.*]] = linalg.matmul
 // TLOOP-SAME:  ins(%[[AB_SUB_SUB]], %[[C__SUB]] : [[TY]], [[TY]])
-// TLOOP-SAME:  outs(%[[ABS_INIT_SUB_SUB]] : [[TY]]) -> [[TY]]
+// TLOOP-SAME:  outs(%[[ABC_INIT_SUB_SUB]] : [[TY]]) -> [[TY]]
 
-// TLOOP:      %[[RES0:.*]] = tensor.insert_slice %[[ABC_SUB_SUB]]
-// TLOOP-SAME:   into %[[ABC_INIT_SUB_]][0, %[[IV1]]]
-// TLOOP:      linalg.yield %[[RES0]] : [[TY]]
+// TLOOP:      linalg.tiled_yield %[[ABC_SUB_SUB]] in %[[ABC_INIT_SUB_SUB]] : [[TY]]
 // TLOOP:    }
-// TLOOP:    %[[RES1:.*]] = tensor.insert_slice %[[ABC_SUB_]] into %[[ABC_INIT_]][%[[IV0]], 0]
-// TLOOP:    linalg.yield %[[RES1]] : [[TY]]
+// TLOOP:    linalg.tiled_yield %[[ABC_SUB]] in %[[ABC_INIT_SUB]] : [[TY]]
 // TLOOP:  }
 // TLOOP:  return %[[ABC]] : [[TY]]
 
 // -----
 
 module {
   func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
                            %arg2: tensor<?x?xf32>) -> tensor<?x?xf32>{

 
 // TLOOP:    %[[AB_SUB:.*]] = linalg.matmul
 // TLOOP-SAME:  ins(%[[A_SUB]], %[[B_SUB]] : [[TY]], [[TY]])
 // TLOOP-SAME:  outs(%[[AB_SUB_INIT]] : [[TY]])
 
 // TLOOP:    %[[DOUBLE_AB:.*]] = linalg.generic
 // TLOOP-SAME:  ins(%[[AB_SUB]] : [[TY]]) outs(%[[INIT_SUB]] : [[TY]])
 
-// TLOOP:    %[[RESULT_SUB:.*]] = tensor.insert_slice
-// TLOOP-SAME:  %[[DOUBLE_AB:.*]] into %[[INIT_]][%[[IV0]], %[[IV1]]]
-
-// TLOOP:    linalg.yield %[[RESULT_SUB]] : [[TY]]
+// TLOOP:    linalg.tiled_yield %[[DOUBLE_AB]] in %[[INIT_SUB]] : [[TY]]
 // TLOOP:  }
 // TLOOP:  return %[[RESULT]] : [[TY]]
 
 // -----
 
 module {
   func @matmul_out_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
                       %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {

 // TLOOP-SAME:      %[[B_:.*]] = %[[B]]: [[TY]],
 // TLOOP-SAME:      %[[C0_F32_:.*]] = %[[C0_F32]]
 // TLOOP-SAME: outs (%[[OUT_:.*]] = %[[OUT]]: [[TY]]) {
 
 // TLOOP:    %[[DIM_A__1:.*]] = tensor.dim %[[A_]], %[[C1]] : [[TY]]
 // TLOOP:    %[[A_SUB:.*]] = tensor.extract_slice %[[A_]][%[[I]], 0]
 // TLOOP:    %[[B_SUB:.*]] = tensor.extract_slice %[[B_]][0, %[[J]]]
 // TLOOP:    %[[OUT_SUB:.*]] = tensor.extract_slice %[[OUT_]][%[[I]], %[[J]]]
-// TLOOP:    %[[INIT_SUB:.*]] = linalg.fill(%[[C0_F32_]], %[[OUT_SUB]])
+// TLOOP:    %[[OUT_SUB_2:.*]] = tensor.extract_slice %[[OUT_]][%[[I]], %[[J]]]
+// TLOOP:    %[[INIT_SUB:.*]] = linalg.fill(%[[C0_F32_]], %[[OUT_SUB_2]])
 
 // TLOOP:    %[[AB_SUB:.*]] = linalg.tiled_loop (%[[K:.*]]) = (%[[C0]])
 // TLOOP-SAME: to (%[[DIM_A__1]]) step (%[[C16]])
 // TLOOP-SAME: ins (%[[A_SUB_:.*]] = %[[A_SUB]]: [[TY]],
 // TLOOP-SAME:      %[[B_SUB_:.*]] = %[[B_SUB]]: [[TY]])
 // TLOOP-SAME: outs (%[[INIT_SUB_:.*]] = %[[INIT_SUB]]: [[TY]])
 // TLOOP-SAME: iterators["reduction"] {
 
 // TLOOP:      %[[A_SUB_SUB:.*]] = tensor.extract_slice %[[A_SUB_]][0, %[[K]]]
 // TLOOP:      %[[B_SUB_SUB:.*]] = tensor.extract_slice %[[B_SUB_]][%[[K]], 0]
 
 // TLOOP:      %[[AB_SUB_SUB:.*]] = linalg.matmul
 // TLOOP-SAME:   ins(%[[A_SUB_SUB]], %[[B_SUB_SUB]] : [[TY]], [[TY]])
 // TLOOP-SAME:   outs(%[[INIT_SUB_]] : [[TY]]) -> [[TY]]
-// TLOOP:      linalg.yield %[[AB_SUB_SUB]] : [[TY]]
+// TLOOP:      linalg.tiled_yield %[[AB_SUB_SUB]] in %[[INIT_SUB_]] : [[TY]]
 // TLOOP:    }
-// TLOOP:    %[[SUB_RESULT:.*]] = tensor.insert_slice %[[AB_SUB]]
-// TLOOP-SAME:  into %[[OUT_]][%[[I]], %[[J]]]
-// TLOOP:    linalg.yield %[[SUB_RESULT]] : [[TY]]
+// TLOOP:    linalg.tiled_yield %[[AB_SUB]] in %[[OUT_SUB]] : [[TY]]
 // TLOOP:  }
 // TLOOP:  return %[[AB]] : [[TY]]
 
 // -----
 
 module {
   func @generic_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
                       %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {

 // TLOOP-SAME:      %[[B_:.*]] = %[[B]]: [[TY]],
 // TLOOP-SAME:      %[[C0_F32_:.*]] = %[[C0_F32]]
 // TLOOP-SAME: outs (%[[OUT_:.*]] = %[[OUT]]: [[TY]]) {
 
 // TLOOP:    %[[DIM_A__1:.*]] = tensor.dim %[[A_]], %[[C1]] : [[TY]]
 // TLOOP:    %[[A_SUB:.*]] = tensor.extract_slice %[[A_]][%[[I]], 0]
 // TLOOP:    %[[B_SUB:.*]] = tensor.extract_slice %[[B_]][0, %[[J]]]
 // TLOOP:    %[[OUT_SUB:.*]] = tensor.extract_slice %[[OUT_]][%[[I]], %[[J]]]
+// TLOOP:    %[[OUT_SUB_2:.*]] = tensor.extract_slice %[[OUT_]][%[[I]], %[[J]]]
 // TLOOP:    %[[INIT_SUB:.*]] = linalg.generic
 // TLOOP-SAME: ins(%[[C0_F32_]]
-// TLOOP-SAME: outs(%[[OUT_SUB]]
+// TLOOP-SAME: outs(%[[OUT_SUB_2]]
 
 // TLOOP:    %[[AB_SUB:.*]] = linalg.tiled_loop (%[[K:.*]]) = (%[[C0]])
 // TLOOP-SAME: to (%[[DIM_A__1]]) step (%[[C16]])
 // TLOOP-SAME: ins (%[[A_SUB_:.*]] = %[[A_SUB]]: [[TY]],
 // TLOOP-SAME:      %[[B_SUB_:.*]] = %[[B_SUB]]: [[TY]])
 // TLOOP-SAME: outs (%[[INIT_SUB_:.*]] = %[[INIT_SUB]]: [[TY]])
 // TLOOP-SAME: iterators["reduction"] {
 
 // TLOOP:      %[[A_SUB_SUB:.*]] = tensor.extract_slice %[[A_SUB_]][0, %[[K]]]
 // TLOOP:      %[[B_SUB_SUB:.*]] = tensor.extract_slice %[[B_SUB_]][%[[K]], 0]
 
 // TLOOP:      %[[AB_SUB_SUB:.*]] = linalg.matmul
 // TLOOP-SAME:   ins(%[[A_SUB_SUB]], %[[B_SUB_SUB]] : [[TY]], [[TY]])
 // TLOOP-SAME:   outs(%[[INIT_SUB_]] : [[TY]]) -> [[TY]]
-// TLOOP:      linalg.yield %[[AB_SUB_SUB]] : [[TY]]
+// TLOOP:      linalg.tiled_yield %[[AB_SUB_SUB]] in %[[INIT_SUB_]] : [[TY]]
 // TLOOP:    }
-// TLOOP:    %[[SUB_RESULT:.*]] = tensor.insert_slice %[[AB_SUB]]
-// TLOOP-SAME:  into %[[OUT_]][%[[I]], %[[J]]]
-// TLOOP:    linalg.yield %[[SUB_RESULT]] : [[TY]]
+// TLOOP:    linalg.tiled_yield %[[AB_SUB]] in %[[OUT_SUB]] : [[TY]]
 // TLOOP:  }
 // TLOOP:  return %[[AB]] : [[TY]]
 

mlir/test/Dialect/Linalg/invalid.mlir

linalg.conv_2d_input_nhwc_filter_hwcf
     { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
     ins(%input, %filter : memref<1x3x4x2xf32>, memref<3x2x2x1xf32>)
     outs(%output : memref<1x2x3x1xf32>)
   return
 }
 
 // -----
 
-#map0 = affine_map<(d0) -> (24, -d0 + 192)>
-#map1 = affine_map<(d0, d1)[s0] -> (d0 * 192 + s0 + d1)>
-#map2 = affine_map<(d0) -> (16, -d0 + 192)>
-
 func private @foo(%A: memref<192x192xf32>, %B: memref<192x192xf32>,
                   %C: memref<192x192xf32>) -> ()
 
 func @tiled_loop_incorrent_num_yield_operands(%A: memref<192x192xf32>,
     %B: memref<192x192xf32>, %C: memref<192x192xf32>,
     %C_tensor: tensor<192x192xf32>) {
   %c24 = constant 24 : index
   %c0 = constant 0 : index
   %c192 = constant 192 : index
   %0 = linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192)
       step (%c24, %c24)
       ins (%A_ = %A: memref<192x192xf32>, %B_ = %B: memref<192x192xf32>)
       outs (%CT_ = %C_tensor: tensor<192x192xf32>,
             %C_ = %C: memref<192x192xf32>) {
         call @foo(%A_, %B_, %C_)
           : (memref<192x192xf32>, memref<192x192xf32>, memref<192x192xf32>)-> ()
     // expected-error @+1 {{expected number of tensor output args = 1 to match the number of yield operands = 0}}
-    linalg.yield
+    linalg.tiled_yield
   }
   return
 }
 
 // -----
 
+func @tiled_loop_incorrect_destination_for_tile(%A: tensor<4xf32>,
+                                                %B: tensor<4xf32>) {
+  %c2 = constant 2 : index
+  %c4 = constant 2 : index
+  %c0 = constant 0 : index
+  %0 = linalg.tiled_loop (%i) = (%c0) to (%c4) step (%c2)
+      ins (%A_ = %A: tensor<4xf32>)
+      outs (%B_ = %B: tensor<4xf32>) {
+    %A_sub = tensor.extract_slice %A_[%i][2][1]
+      : tensor<4xf32> to tensor<2xf32>
+    %B_sub = tensor.extract_slice %B_[%i][2][1]
+      : tensor<4xf32> to tensor<2xf32>
+    %c0_f32 = constant 0.0 : f32
+    %tile = linalg.fill(%c0_f32, %A_sub) : f32, tensor<2xf32> -> tensor<2xf32>
+    // expected-error @+1 {{expected output 0 to be a subset of the corresponding block argument}}
+    linalg.tiled_yield %tile in %A_sub : tensor<2xf32>
+  }
+  return
+}
+
+
+// -----
+
 #map0 = affine_map<(d0) -> (24, -d0 + 192)>
 #map1 = affine_map<(d0, d1)[s0] -> (d0 * 192 + s0 + d1)>
 #map2 = affine_map<(d0) -> (16, -d0 + 192)>
 
 func private @foo(%A: memref<192x192xf32>, %B: memref<192x192xf32>,
                   %C: memref<192x192xf32>) -> tensor<f32>
 
 func @tiled_loop_incorrent_yield_operand_type(%A: memref<192x192xf32>,
     %B: memref<192x192xf32>, %C: memref<192x192xf32>,
     %C_tensor: tensor<192x192xf32>) {
   %c24 = constant 24 : index
   %c0 = constant 0 : index
   %c192 = constant 192 : index
   %0 = linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192)
       step (%c24, %c24)
       ins (%A_ = %A: memref<192x192xf32>, %B_ = %B: memref<192x192xf32>)
       outs (%CT_ = %C_tensor: tensor<192x192xf32>,
             %C_ = %C: memref<192x192xf32>) {
         %1 = call @foo(%A_, %B_, %C_)
           : (memref<192x192xf32>, memref<192x192xf32>, memref<192x192xf32>)-> tensor<f32>
-    // expected-error @+1 {{expected yield operand 0 with type = 'tensor<f32>' to match output arg type = 'tensor<192x192xf32>}}
-    linalg.yield %1 : tensor<f32>
+    // expected-error @+1 {{expected tile operand with type = 'tensor<f32>' to match output type = 'tensor<192x192xf32>}}
+    "linalg.tiled_yield" (%1, %CT_) : (tensor<f32>, tensor<192x192xf32>) -> ()
   }
   return
 }
 
 // -----
 
 func private @foo(%A: memref<192x192xf32>, %B: memref<192x192xf32>,
-                  %C: memref<192x192xf32>) -> ()
+                  %C: memref<192x192xf32>) -> (tensor<192x192xf32>)
 
 func @tiled_loop_incorrent_iterator_types_count(%A: memref<192x192xf32>,
     %B: memref<192x192xf32>, %C: memref<192x192xf32>,
     %C_tensor: tensor<192x192xf32>) {
   %c24 = constant 24 : index
   %c0 = constant 0 : index
   %c192 = constant 192 : index
   // expected-error @+1 {{expected iterator types array attribute size = 1 to match the number of loops = 2}}
   %0 = "linalg.tiled_loop"(%c0, %c0, %c192, %c192, %c24, %c24, %A, %B, %C_tensor, %C) ( {
     ^bb0(%arg4: index, %arg5: index, %A_: memref<192x192xf32>,
          %B_: memref<192x192xf32>, %CT_: tensor<192x192xf32>,
          %C_: memref<192x192xf32>):
-      call @foo(%A_, %B_, %C_)
-          : (memref<192x192xf32>, memref<192x192xf32>, memref<192x192xf32>)-> ()
-      linalg.yield %CT_ : tensor<192x192xf32>
+      %tile = call @foo(%A_, %B_, %C_)
+          : (memref<192x192xf32>, memref<192x192xf32>, memref<192x192xf32>)
+          -> (tensor<192x192xf32>)
+      linalg.tiled_yield %tile in %CT_ : tensor<192x192xf32>
     }) {
       iterator_types = ["parallel"],
       operand_segment_sizes = dense<2> : vector<5xi32>
     } : (index, index, index, index, index, index, memref<192x192xf32>,
       memref<192x192xf32>, tensor<192x192xf32>, memref<192x192xf32>
     ) -> tensor<192x192xf32>
   return
 }
 
 // -----
 
 func private @foo(%A: memref<100xf32>) -> ()
 
 func @tiled_loop_incorrent_block_arg_type(%A: memref<192xf32>) {
   %c0 = constant 0 : index
   %c192 = constant 192 : index
   %c24 = constant 24 : index
   // expected-error @+1 {{expected output arg 0 with type = 'memref<192xf32>' to match region arg 1 type = 'memref<100xf32>'}}
   "linalg.tiled_loop"(%c0, %c192, %c24, %A) ( {
     ^bb0(%arg4: index, %A_: memref<100xf32>):
       call @foo(%A_) : (memref<100xf32>)-> ()
-      linalg.yield
+      linalg.tiled_yield
     }) {
       iterator_types = ["parallel"],
       operand_segment_sizes = dense<[1, 1, 1, 0, 1]> : vector<5xi32>
     } : (index, index, index, memref<192xf32>) -> ()
   return
 }
 
 // -----

mlir/test/Dialect/Linalg/roundtrip.mlir

%prod = linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
     %sum = linalg.generic #trait_4
         ins(%lhs_sub, %rhs_sub : tensor<?x?xi8>, tensor<?x?xi8>)
         outs(%out_sub : tensor<?x?xi8>) {
       ^bb(%l: i8, %r: i8, %o: i8) :
         %s = addi %l, %r : i8
         linalg.yield %s : i8
       } -> tensor<?x?xi8>
 
-    %sum_sub = tensor.insert_slice %sum into %out_[%i, 0][%c4, %c64][1, 1]
-      : tensor<?x?xi8> into tensor<24x64xi8>
-    linalg.yield %sum_sub : tensor<24x64xi8>
+    linalg.tiled_yield %sum in %out_sub : tensor<?x?xi8>
   }
   return %prod : tensor<24x64xi8>
 }
 // CHECK-LABEL: func @tiled_loop
 // CHECK-NOT: iterators[
 
 // -----
 
%acc = linalg.generic #trait_5
         : tensor<2x4x8xf32>, tensor<2x8xf32>, tensor<4xf32>)
       outs(%sub_out : tensor<4xf32>)  {
     ^bb0(%i3d: f32, %i2d: f32, %i1d: f32, %o: f32):
       %0 = addf %i3d, %i2d : f32
       %1 = addf %0, %i1d : f32
       linalg.yield %1 : f32
     } -> tensor<4xf32>
 
-    %sum_sub = tensor.insert_slice %acc into %o_[%j][%c4][1]
-      : tensor<4xf32> into tensor<24xf32>
-    linalg.yield %sum_sub : tensor<24xf32>
+    linalg.tiled_yield %acc in %sub_out : tensor<4xf32>
   }
   return %result : tensor<24xf32>
 }
 // CHECK-LABEL: func @tiled_loop_reduction
 // CHECK: iterators[
 
 // -----
 
linalg.generic #trait_6
           memref<2x8xf32, #map_2>,
           memref<4xf32, #map_3>)
       outs(%sub_out : memref<4xf32, #map_3>)  {
     ^bb0(%i3d: f32, %i2d: f32, %i1d: f32, %o: f32):
       %0 = addf %i3d, %i2d : f32
       %1 = addf %0, %i1d : f32
       linalg.yield %1 : f32
     }
-    linalg.yield
+    linalg.tiled_yield
   }
   return
 }
 // CHECK-LABEL: func @tiled_loop_on_buffers
 // CHECK: iterators[

mlir/test/Dialect/Linalg/tile-tensors.mlir

 
 // TLOOP: %[[SUB_ARG_0:.*]] = tensor.extract_slice %[[A0]][%[[I]], %[[K]]]
 // TLOOP: %[[SUB_ARG_1:.*]] = tensor.extract_slice %[[A1]][%[[K]], %[[J]]]
 // TLOOP: %[[SUB_ARG_2:.*]] = tensor.extract_slice %[[A2]][%[[I]], %[[J]]]
 
 // TLOOP: %[[PROD:.*]] = linalg.matmul ins(%[[SUB_ARG_0]], %[[SUB_ARG_1]]
 // TLOOP-SE: outs(%[[SUB_ARG_2]] : [[TY]]) -> [[TY]]
 
-// TLOOP: %[[O:.*]] = tensor.insert_slice %[[PROD]] into %[[A2]][%[[I]], %[[J]]]
-// TLOOP: linalg.yield %[[O]] : [[TY]]
+// TLOOP: linalg.tiled_yield %[[PROD]] in %[[SUB_ARG_2]] : [[TY]]
 
 // -----
 
 func @generic_op_tensors(
   %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index

mlir/test/Dialect/Linalg/tiled-loops.mlir

linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
     %3 = memref.subview %B_[0, %j] [192, %2] [1, 1]
       : memref<192x192xf32> to memref<192x?xf32, #map1>
     %4 = memref.subview %C_[%i, %j] [%0, %2] [1, 1]
       : memref<192x192xf32> to memref<?x?xf32, #map1>
     linalg.fill(%cst, %4) : f32, memref<?x?xf32, #map1>
     linalg.matmul ins(%1, %3 : memref<?x192xf32, #map1>,
                                memref<192x?xf32, #map1>)
                   outs(%4 : memref<?x?xf32, #map1>)
-    linalg.yield
+    linalg.tiled_yield
   }
   return
 }
 
 // CHECK-LABEL: @tiled_loop
 // CHECK-SAME:  %[[A:.*]]: memref<192x192xf32>, %[[B:.*]]: memref<192x192xf32>,
 // CHECK-SAME:  %[[C:.*]]: memref<192x192xf32>) {
 // CHECK:       %[[C24:.*]] = constant 24 : index
func @tiled_loop_reduction(%A: memref<192x192xf32>,
    %c192 = constant 192 : index
    %cst = constant 0.000000e+00 : f32
 
   linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192) step (%c24, %c16)
       ins (%A_ = %A: memref<192x192xf32>, %B_ = %B:  memref<192x192xf32>)
       outs (%C_ = %C: memref<f32>)
       iterators["reduction", "reduction"] {
     linalg.fill(%cst, %A_) : f32, memref<192x192xf32>
-    linalg.yield
+    linalg.tiled_yield
   }
   return
 }
 
 // CHECK-LABEL: @tiled_loop_reduction
 // CHECK:       %[[C24:.*]] = constant 24 : index
 // CHECK:       %[[C16:.*]] = constant 16 : index
 // CHECK:       %[[C0:.*]] = constant 0 : index
 // CHECK:       %[[C192:.*]] = constant 192 : index
 // CHECK:       scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C24]]
 // CHECK:         scf.for %{{.*}} = %[[C0]] to %[[C192]] step %[[C16]]
 // CHECK:           linalg.fill