[mlir] Add block arguments for input/output operands of 'linalg.tiled_loop`.

[ThomasRaoux]

Thanks @pifon2a!

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

2124–2135

Not directly related to this change but I believe we are not supposed to create op from another builder than rewriter in rewrite patterns. It will cause problem if the rewriter later decide this transformation should not happen. I'm not sure there is an easy way to reproduce it.

Could we use rewriter.inlineRegionBefore instead? I think something like that would work instead:

rewriter.eraseBlock(newTiledLoop.getBody());
rewriter.inlineRegionBefore(tiledLoop.region(), newTiledLoop.region(),
                            newTiledLoop.region().end());

Operation* terminator = newTiledLoop.getBody()->getTerminator();
rewriter.setInsertionPoint(terminator);
rewriter.create<linalg::YieldOp>(loc, newYieldArgs);
rewriter.eraseOp(terminator);

[pifon2a]

@ThomasRaoux Yes, this is dangerous, if the "inner" builder was created without rewriter.getListener() passed to it.

[ThomasRaoux]

In D101186#2712987, @pifon2a wrote:

@ThomasRaoux Yes, this is dangerous, if the "inner" builder was created without rewriter.getListener() passed to it.

Ha missed that, you're right!

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

let arguments = (ins Variadic<Index>:$lowerBound,
                        ArrayAttr:$iterator_types);
   let results = (outs Variadic<AnyRankedTensor>:$results);
   let regions = (region SizedRegion<1>:$region);
 
   let builders = [
     OpBuilder<(ins "ValueRange":$lowerBounds, "ValueRange":$upperBounds,
       "ValueRange":$steps, "ValueRange":$inputs, "ValueRange":$outputs,
       "ArrayAttr":$iteratorTypes,
-      CArg<"function_ref<void (OpBuilder &, Location, ValueRange)>",
-           "nullptr">:$bodyBuilderFn)>,
+      CArg<"function_ref<void (OpBuilder &, Location, /*ivs=*/ValueRange,"
+        "/*inputs=*/ValueRange, /*outputs=*/ValueRange)>",
   ];
 
   let extraClassDeclaration = [{
     /// Number of operands controlling the loop: lbs, ubs, steps
     unsigned getNumControlOperands() { return 3 * getNumLoops(); }
 
     ValueRange getInductionVars() {
-      return getBody()->getArguments();
+      return getBody()->getArguments().take_front(getNumLoops());
+    }
+    ValueRange getRegionInputArgs() {
+      return getBody()->getArguments().slice(getNumLoops(), inputs().size());
+    }
+    ValueRange getRegionOutputArgs() {
+      return getBody()->getArguments().take_back(outputs().size());
     }
 
 
     void setLowerBounds(ValueRange lowerBounds) {
       unsigned numLoops = getNumLoops();
       assert(lowerBounds.size() == numLoops &&
              "expected lower bounds for every loop dimension");
       for (unsigned i = 0; i < numLoops; ++i)

mlir/include/mlir/IR/OpImplementation.h

if (!result.hasValue())
       return emitError(getCurrentLocation(), "expected '('");
     return result.getValue();
   }
 
   virtual OptionalParseResult
   parseOptionalAssignmentList(SmallVectorImpl<OperandType> &lhs,
                               SmallVectorImpl<OperandType> &rhs) = 0;
 
+  /// Parse a list of assignments of the form
+  ///   (%x1 = %y1 : type1, %x2 = %y2 : type2, ...)
+  ParseResult parseAssignmentListWithTypes(SmallVectorImpl<OperandType> &lhs,
+                                           SmallVectorImpl<OperandType> &rhs,
+                                           SmallVectorImpl<Type> &types) {
+    OptionalParseResult result =
+        parseOptionalAssignmentListWithTypes(lhs, rhs, types);
+    if (!result.hasValue())
+      return emitError(getCurrentLocation(), "expected '('");
+    return result.getValue();
+  }
+
+  virtual OptionalParseResult
+  parseOptionalAssignmentListWithTypes(SmallVectorImpl<OperandType> &lhs,
+                                       SmallVectorImpl<OperandType> &rhs,
+                                       SmallVectorImpl<Type> &types) = 0;
   /// Parse a keyword followed by a type.
   ParseResult parseKeywordType(const char *keyword, Type &result) {
     return failure(parseKeyword(keyword) || parseType(result));
   }
 
   /// Add the specified type to the end of the specified type list and return
   /// success.  This is a helper designed to allow parse methods to be simple
   /// and chain through || operators.

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

static LogicalResult verify(linalg::YieldOp op) {
   }
   return op.emitOpError("expected parent op with LinalgOp interface");
 }
 
 //===----------------------------------------------------------------------===//
 // TiledLoopOp
 //===----------------------------------------------------------------------===//
 
-void TiledLoopOp::build(
-    OpBuilder &builder, OperationState &result, ValueRange lowerBounds,
-    ValueRange upperBounds, ValueRange steps, ValueRange inputs,
-    ValueRange outputs, ArrayAttr iteratorTypes,
-    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuilderFn) {
+void TiledLoopOp::build(OpBuilder &builder, OperationState &result,
+                        ValueRange lowerBounds, ValueRange upperBounds,
+                        ValueRange steps, ValueRange inputs, ValueRange outputs,
+                        ArrayAttr iteratorTypes,
+                        function_ref<void(OpBuilder &, Location, ValueRange,
+                                          ValueRange, ValueRange)>
+                            bodyBuilderFn) {
   result.addOperands(lowerBounds);
   result.addOperands(upperBounds);
   result.addOperands(steps);
   result.addOperands(inputs);
   result.addOperands(outputs);
   result.addAttribute(
       TiledLoopOp::getOperandSegmentSizeAttr(),
       builder.getI32VectorAttr({static_cast<int32_t>(lowerBounds.size()),
                                 static_cast<int32_t>(upperBounds.size()),
                                 static_cast<int32_t>(steps.size()),
                                 static_cast<int32_t>(inputs.size()),
                                 static_cast<int32_t>(outputs.size())}));
   result.addAttribute(getIteratorTypesAttrName(), iteratorTypes);
 
   // Add output types for `RankedTensorType` output arguments.
   for (Value output : outputs) {
     Type outputType = output.getType();
     if (outputType.isa<RankedTensorType>())
       result.addTypes(outputType);
   }
 
   OpBuilder::InsertionGuard guard(builder);
   unsigned numIVs = steps.size();
   SmallVector<Type, 8> argTypes(numIVs, builder.getIndexType());
+  for (Type type : TypeRange(inputs))
+    argTypes.push_back(type);
+  for (Type type : TypeRange(outputs))
+    argTypes.push_back(type);
   Region *bodyRegion = result.addRegion();
   Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes);
 
   if (bodyBuilderFn) {
     builder.setInsertionPointToStart(bodyBlock);
-    bodyBuilderFn(builder, result.location, bodyBlock->getArguments());
+    bodyBuilderFn(builder, result.location,
+                  bodyBlock->getArguments().take_front(numIVs),
+                  bodyBlock->getArguments().slice(numIVs, inputs.size()),
+                  bodyBlock->getArguments().take_back(outputs.size()));
     TiledLoopOp::ensureTerminator(*bodyRegion, builder, result.location);
   }
 }
 
 static void print(OpAsmPrinter &p, TiledLoopOp op) {
-  p << op.getOperationName() << " (" << op.getBody()->getArguments() << ") = ("
+  p << op.getOperationName() << " (" << op.getInductionVars() << ") = ("
     << op.lowerBound() << ") to (" << op.upperBound() << ") step (" << op.step()
     << ")";
 
-  if (!op.inputs().empty())
-    p << " ins (" << op.inputs() << ": " << TypeRange(op.inputs()) << ")";
-  if (!op.outputs().empty())
-    p << " outs (" << op.outputs() << ":" << TypeRange(op.outputs()) << ")";
+  if (!op.inputs().empty()) {
+    p << " ins (";
+    llvm::interleaveComma(llvm::zip(op.getRegionInputArgs(), op.inputs()), p,
+                          [&](auto it) {
+                            p << std::get<0>(it) << " = " << std::get<1>(it)
+                              << ": " << std::get<1>(it).getType();
+                          });
+    p << ")";
+  }
+  if (!op.outputs().empty()) {
+    p << " outs (";
+    llvm::interleaveComma(llvm::zip(op.getRegionOutputArgs(), op.outputs()), p,
+                          [&](auto it) {
+                            p << std::get<0>(it) << " = " << std::get<1>(it)
+                              << ": " << std::get<1>(it).getType();
+                          });
+    p << ")";
+  }
 
   if (llvm::any_of(op.iterator_types(), [](Attribute attr) {
         return attr.cast<StringAttr>().getValue() !=
                getParallelIteratorTypeName();
       })) {
     p << " iterators" << op.iterator_types() << "";
   }
 
static ParseResult parseTiledLoopOp(OpAsmParser &parser,
   SmallVector<OpAsmParser::OperandType, 4> steps;
   if (parser.parseKeyword("step") ||
       parser.parseOperandList(steps, ivs.size(),
                               OpAsmParser::Delimiter::Paren) ||
       parser.resolveOperands(steps, builder.getIndexType(), result.operands))
     return failure();
 
   // Parse input tensors.
-  SmallVector<OpAsmParser::OperandType, 4> inputs;
-  if (succeeded(parser.parseOptionalKeyword("ins"))) {
-    SmallVector<Type, 4> inputTypes;
+  SmallVector<OpAsmParser::OperandType, 4> inputs, input_region_args;
+  if (succeeded(parser.parseOptionalKeyword("ins"))) {
     llvm::SMLoc inputsOperandsLoc = parser.getCurrentLocation();
 
-    if (parser.parseLParen() || parser.parseOperandList(inputs) ||
-        parser.parseColonTypeList(inputTypes) || parser.parseRParen())
+    if (parser.parseAssignmentListWithTypes(input_region_args, inputs,
+                                            inputTypes))
       return failure();
 
     if (parser.resolveOperands(inputs, inputTypes, inputsOperandsLoc,
                                result.operands))
       return failure();
   }
 
   // Parse output tensors.
-  SmallVector<OpAsmParser::OperandType, 4> outputs;
-  if (succeeded(parser.parseOptionalKeyword("outs"))) {
-    SmallVector<Type, 4> outputTypes;
+  SmallVector<OpAsmParser::OperandType, 4> outputs, output_region_args;
+  if (succeeded(parser.parseOptionalKeyword("outs"))) {
     llvm::SMLoc outputsOperandsLoc = parser.getCurrentLocation();
 
-    if (parser.parseLParen() || parser.parseOperandList(outputs) ||
-        parser.parseColonTypeList(outputTypes) || parser.parseRParen())
+    if (parser.parseAssignmentListWithTypes(output_region_args, outputs,
+                                            outputTypes))
       return failure();
 
     if (parser.resolveOperands(outputs, outputTypes, outputsOperandsLoc,
                                result.operands))
       return failure();
     for (Type outputType : outputTypes)
       if (outputType.isa<RankedTensorType>())
         result.addTypes(outputType);
result.addAttribute(
       builder.getI32VectorAttr({static_cast<int32_t>(lower.size()),
                                 static_cast<int32_t>(upper.size()),
                                 static_cast<int32_t>(steps.size()),
                                 static_cast<int32_t>(inputs.size()),
                                 static_cast<int32_t>(outputs.size())}));
 
   // Parse the body.
   Region *body = result.addRegion();
-  SmallVector<Type, 4> types(ivs.size(), builder.getIndexType());
-  if (parser.parseRegion(*body, ivs, types))
+
+  SmallVector<Type, 4> region_types(ivs.size(), builder.getIndexType());
+  region_types.append(inputTypes);
+  region_types.append(outputTypes);
+
+  SmallVector<OpAsmParser::OperandType, 4> region_args(ivs);
+  region_args.append(input_region_args);
+  region_args.append(output_region_args);
+
+  if (parser.parseRegion(*body, region_args, region_types))
     return failure();
 
   // Parse optional attributes.
   parser.parseOptionalAttrDict(result.attributes);
 
   return success();
 }
 

 }
 
 static LogicalResult verify(TiledLoopOp op) {
   // Check if iterator types are provided for every loop dimension.
   if (op.iterator_types().size() != op.getNumLoops())
     return op.emitOpError("expected iterator types array attribute size = ")
            << op.iterator_types().size()
            << " to match the number of loops = " << op.getNumLoops();
+
+  // Check if types of input arguments match region args types.
+  for (auto &item :
+       llvm::enumerate(llvm::zip(op.inputs(), op.getRegionInputArgs()))) {
+    Value input, inputRegionArg;
+    unsigned index = item.index();
+    std::tie(input, inputRegionArg) = item.value();
+    if (input.getType() != inputRegionArg.getType())
+      return op.emitOpError("expected input arg ")
+             << index << " with type = " << input.getType()
+             << " to match region arg " << index + op.getNumLoops()
+             << " type = " << inputRegionArg.getType();
+  }
+
+  // Check if types of input arguments match region args types.
+  for (auto &item :
+       llvm::enumerate(llvm::zip(op.outputs(), op.getRegionOutputArgs()))) {
+    Value output, outputRegionArg;
+    unsigned index = item.index();
+    std::tie(output, outputRegionArg) = item.value();
+    if (output.getType() != outputRegionArg.getType())
+      return op.emitOpError("expected output arg ")
+             << index << " with type = " << output.getType()
+             << " to match region arg "
+             << index + op.getNumLoops() + op.inputs().size()
+             << " type = " << outputRegionArg.getType();
+  }
   return success();
 }
 
 namespace {
 
 // Folds away TiledLoopOp output tensors when the following conditions are met:
 // * result of `linalg.tiled_loop` has no uses
 // * output tensor is the argument of `linalg.yield`
 //
 // Example:
 //
-// %0 = linalg.tiled_loop ...  outs (%out, %out_buf:tensor<...>, memref<...>) {
+// %0 = linalg.tiled_loop ...  outs (%o_ = %out: tensor<...>,
+//                                   %obuf_ = %out_buf: memref<...>) {
 //   ...
-//   linalg.yield %out : tensor ...
+//   linalg.yield %o_ : tensor ...
 // }
 //
 // Becomes
 //
-// linalg.tiled_loop ...  outs (%out_buf:memref<...>) {
+// linalg.tiled_loop ...  outs (%obuf_ = %out_buf: memref<...>) {
 //   ...
 //   linalg.yield
 // }
 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());
 
     // 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, newYieldArgs,
+        regionOutputTensorArgs;
     int resultId = 0;
-    for (Value out : tiledLoop.outputs()) {
+    // Store ids of the corresponding old and new output operands.
+    SmallVector<std::pair<size_t, size_t>, 2> old_out_id_to_new;
+    for (auto item : llvm::enumerate(
+             llvm::zip(tiledLoop.outputs(), tiledLoop.getRegionOutputArgs()))) {
+      size_t index = item.index();
+      Value out = std::get<0>(item.value());
+      Value outRegionArg = std::get<1>(item.value());
+
       if (!out.getType().isa<RankedTensorType>()) {
+        old_out_id_to_new.push_back({index, newOutputOperands.size()});
         newOutputOperands.push_back(out);
+        regionOutputTensorArgs.push_back(outRegionArg);
         continue;
       }
       Value result = tiledLoop.getResult(resultId);
       Value yieldArg = yieldOp.getOperand(resultId);
-      if (yieldArg != out || !result.use_empty()) {
+      if (yieldArg != outRegionArg || !result.use_empty()) {
+        old_out_id_to_new.push_back({index, newOutputOperands.size()});
         newOutputOperands.push_back(out);
         newYieldArgs.push_back(yieldArg);
       }
       ++resultId;
     }
     if (newOutputOperands.size() == tiledLoop.outputs().size())
       return failure();
 
     Location loc = tiledLoop.getLoc();
     auto newTiledLoop = rewriter.create<TiledLoopOp>(
         loc, tiledLoop.lowerBound(), tiledLoop.upperBound(), tiledLoop.step(),
         tiledLoop.inputs(), newOutputOperands, tiledLoop.iterator_types());
 
     // Clone the region ignoring the def-chain for linalg.yield args:
     // unnecessary `subtensor_insert`, `tensor_load` and `cast` ops.
     BlockAndValueMapping bvm;
     bvm.map(tiledLoop.getInductionVars(), newTiledLoop.getInductionVars());
+    bvm.map(tiledLoop.getRegionInputArgs(), newTiledLoop.getRegionInputArgs());
+    for (const auto &item : old_out_id_to_new)
+      bvm.map(tiledLoop.getRegionOutputArgs()[item.first],
+              newTiledLoop.getRegionOutputArgs()[item.second]);
     OpBuilder innerBuilder =
         OpBuilder::atBlockEnd(newTiledLoop.getBody(), rewriter.getListener());
     for (auto &op : tiledLoop.getBody()->without_terminator())
       innerBuilder.clone(op, bvm);
     innerBuilder.create<linalg::YieldOp>(loc, newYieldArgs);
     rewriter.eraseOp(tiledLoop);
 
     return success();
   }
 };
 } // namespace
 
 void TiledLoopOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {

mlir/lib/Parser/Parser.cpp

auto parseElt = [&]() -> ParseResult {
         return failure();
       lhs.push_back(regionArg);
       rhs.push_back(operand);
       return success();
     };
     return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt);
   }
 
+  /// Parse a list of assignments of the form
+  ///   (%x1 = %y1 : type1, %x2 = %y2 : type2, ...).
+  OptionalParseResult
+  parseOptionalAssignmentListWithTypes(SmallVectorImpl<OperandType> &lhs,
+                                       SmallVectorImpl<OperandType> &rhs,
+                                       SmallVectorImpl<Type> &types) override {
+    if (failed(parseOptionalLParen()))
+      return llvm::None;
+
+    auto parseElt = [&]() -> ParseResult {
+      OperandType regionArg, operand;
+      Type type;
+      if (parseRegionArgument(regionArg) || parseEqual() ||
+          parseOperand(operand) || parseColon() || parseType(type))
+        return failure();
+      lhs.push_back(regionArg);
+      rhs.push_back(operand);
+      types.push_back(type);
+      return success();
+    };
+    return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt);
+  }
+
 private:
   /// The source location of the operation name.
   SMLoc nameLoc;
 
   /// Information about the result name specifiers.
   ArrayRef<OperationParser::ResultRecord> resultIDs;
 
   /// The abstract information of the operation.

mlir/test/Dialect/Linalg/canonicalize.mlir

func @fold_tiled_loop_results(%A: memref<192x192xf32>, %B: memref<192x192xf32>,
                               %C_tensor: tensor<192x192xf32>) {
   %cst = constant 0.000000e+00 : f32
   %c24 = constant 24 : index
   %c16 = constant 16 : index
   %c0 = constant 0 : index
   %c192 = constant 192 : index
   %useless = linalg.tiled_loop (%i, %j) = (%c0, %c0) to (%c192, %c192)
       step (%c24, %c16)
-      ins (%A, %B: memref<192x192xf32>, memref<192x192xf32>)
-      outs (%C_tensor, %C :tensor<192x192xf32>, memref<192x192xf32>) {
-        call @foo(%A, %B, %C) : (memref<192x192xf32>, memref<192x192xf32>, memref<192x192xf32>)-> ()
-    linalg.yield %C_tensor : tensor<192x192xf32>
+      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>)-> ()
+    linalg.yield %CT_ : tensor<192x192xf32>
   }
   return
 }
 
 // CHECK-LABEL: func @fold_tiled_loop_results(
 // CHECK-SAME:    %[[A:.*]]: [[TY:.*]], %[[B:.*]]: [[TY]], %[[C:.*]]: [[TY]],
 // CHECK-SAME:    %[[C_TENSOR:.*]]: tensor<{{.*}}>) {
 // CHECK:  %[[C24:.*]] = constant 24 : index
 // CHECK:  %[[C16:.*]] = constant 16 : index
 // CHECK:  %[[C0:.*]] = constant 0 : index
 // CHECK:  %[[C192:.*]] = constant 192 : index
 
 // CHECK-NOT: %{{.*}} = linalg.tiled_loop
 // CHECK:  linalg.tiled_loop (%{{.*}}, %{{.*}}) = (%[[C0]], %[[C0]])
 // CHECK-SAME: to (%[[C192]], %[[C192]]) step (%[[C24]], %[[C16]])
-// CHECK-SAME: ins (%[[A]], %[[B]]: memref<192x192xf32>, memref<192x192xf32>)
-// CHECK-SAME: outs (%[[C]]:memref<192x192xf32>) {
-// CHECK-NEXT:   call @foo(%[[A]], %[[B]], %[[C]])
+// CHECK-SAME: ins (%[[A_:.*]] = %[[A]]: memref<192x192xf32>, %[[B_:.*]] = %[[B]]: memref<192x192xf32>)
+// CHECK-SAME: outs (%[[C_:.*]] = %[[C]]: memref<192x192xf32>) {
+// CHECK-NEXT:   call @foo(%[[A_]], %[[B_]], %[[C_]])
 // CHECK-NEXT:   linalg.yield
 
 // -----
 
 func @dim_of_pad_op(%arg0 : tensor<2x?x?xf32>, %arg1 : index, %arg2 : index,
     %arg3: f32) -> (index, index, index)
 {
    %c0 = constant 0 : index

mlir/test/Dialect/Linalg/invalid.mlir

 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, %B: memref<192x192xf32>, memref<192x192xf32>)
-      outs (%C_tensor, %C :tensor<192x192xf32>, memref<192x192xf32>) {
-        call @foo(%A, %B, %C)
+      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
   }
   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, %B: memref<192x192xf32>, memref<192x192xf32>)
-      outs (%C_tensor, %C :tensor<192x192xf32>, memref<192x192xf32>) {
-        %1 = call @foo(%A, %B, %C)
+      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>
   }
   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_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):  // no predecessors
-      call @foo(%A, %B, %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 %C_tensor : tensor<192x192xf32>
+      linalg.yield %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
+    }) {
+      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

 func @tiled_loop(%lhs: tensor<24x64xi8>, %rhs: tensor<24x64xi8>,
                  %out: tensor<24x64xi8>) -> tensor<24x64xi8> {
  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %c4 = constant 4 : index
  %c24 = constant 24 : index
  %c64 = constant 64 : index
  %prod = linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
-      ins(%lhs, %rhs : tensor<24x64xi8>, tensor<24x64xi8>)
-      outs(%out : tensor<24x64xi8>) {
-    %lhs_sub = subtensor %lhs[%i, 0] [%c4, %c64] [1, 1]
+      ins(%lhs_ = %lhs: tensor<24x64xi8>, %rhs_ = %rhs: tensor<24x64xi8>)
+      outs(%out_ = %out: tensor<24x64xi8>) {
+    %lhs_sub = subtensor %lhs_[%i, 0] [%c4, %c64] [1, 1]
         : tensor<24x64xi8> to tensor<?x?xi8>
-    %rhs_sub = subtensor %rhs[%i, 0] [%c4, %c64] [1, 1]
+    %rhs_sub = subtensor %rhs_[%i, 0] [%c4, %c64] [1, 1]
         : tensor<24x64xi8> to tensor<?x?xi8>
-    %out_sub = subtensor %out[%i, 0] [%c4, %c64] [1, 1]
+    %out_sub = subtensor %out_[%i, 0] [%c4, %c64] [1, 1]
         : tensor<24x64xi8> to tensor<?x?xi8>
 
     %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 = subtensor_insert %sum into %out[%i, 0][%c4, %c64][1, 1]
+    %sum_sub = subtensor_insert %sum into %out_[%i, 0][%c4, %c64][1, 1]
       : tensor<?x?xi8> into tensor<24x64xi8>
     linalg.yield %sum_sub : tensor<24x64xi8>
   }
   return %prod : tensor<24x64xi8>
 }
 // CHECK-LABEL: func @tiled_loop
 // CHECK-NOT: iterators[
 
func @tiled_loop_reduction(%input_3d: tensor<16x24x32xf32>,
   %c2 = constant 2 : index
   %c4 = constant 4 : index
   %c8 = constant 8 : index
   %X = memref.dim %input_3d, %c0 : tensor<16x24x32xf32>
   %Y = memref.dim %input_3d, %c1 : tensor<16x24x32xf32>
   %Z = memref.dim %input_3d, %c2 : tensor<16x24x32xf32>
   %result = linalg.tiled_loop (%i, %j, %k)
       = (%c0, %c0, %c0) to (%X, %Y, %Z) step (%c2, %c4, %c8)
-      ins(%input_3d, %input_2d: tensor<16x24x32xf32>, tensor<16x32xf32>)
-      outs( %output: tensor<24xf32>)
+      ins(%i3d_ = %input_3d: tensor<16x24x32xf32>,
+          %i2d_ = %input_2d: tensor<16x32xf32>,
+          %i1d_ = %input_1d: tensor<24xf32>)
+      outs(%o_ =  %output: tensor<24xf32>)
       iterators["reduction", "parallel", "reduction"] {
-    %sub_3d = subtensor %input_3d[%i, %j, %k][2, 4, 8][1, 1, 1]
+    %sub_3d = subtensor %i3d_[%i, %j, %k][2, 4, 8][1, 1, 1]
       : tensor<16x24x32xf32> to tensor<2x4x8xf32>
-    %sub_2d = subtensor %input_2d[%i, %k][2, 8][1, 1]
+    %sub_2d = subtensor %i2d_[%i, %k][2, 8][1, 1]
       : tensor<16x32xf32> to tensor<2x8xf32>
-    %sub_1d = subtensor %input_1d[%j] [4] [1]
+    %sub_1d = subtensor %i1d_[%j] [4] [1]
       : tensor<24xf32> to tensor<4xf32>
-    %sub_out = subtensor %output[%j] [4] [1]
+    %sub_out = subtensor %o_[%j] [4] [1]
       : tensor<24xf32> to tensor<4xf32>
     %acc = linalg.generic #trait_5
       ins(%sub_3d, %sub_2d, %sub_1d
         : 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 = subtensor_insert %acc into %output[%j][%c4][1]
+    %sum_sub = subtensor_insert %acc into %o_[%j][%c4][1]
       : tensor<4xf32> into tensor<24xf32>
     linalg.yield %sum_sub : tensor<24xf32>
   }
   return %result : tensor<24xf32>
 }
 // CHECK-LABEL: func @tiled_loop_reduction
 // CHECK: iterators[
 
func @tiled_loop_on_buffers(%input_3d: memref<16x24x32xf32>,
   %c2 = constant 2 : index
   %c4 = constant 4 : index
   %c8 = constant 8 : index
   %X = memref.dim %input_3d, %c0 : memref<16x24x32xf32>
   %Y = memref.dim %input_3d, %c1 : memref<16x24x32xf32>
   %Z = memref.dim %input_3d, %c2 : memref<16x24x32xf32>
   linalg.tiled_loop (%i, %j, %k) = (%c0, %c0, %c0)
       to (%X, %Y, %Z) step (%c2, %c4, %c8)
-      ins(%input_3d, %input_2d: memref<16x24x32xf32>, memref<16x32xf32>)
-      outs( %output: memref<24xf32>)
+      ins(%i3d_ = %input_3d: memref<16x24x32xf32>,
+          %i2d_ = %input_2d: memref<16x32xf32>,
+          %i1d_ = %input_1d: memref<24xf32>)
+      outs(%o_ =  %output: memref<24xf32>)
       iterators["reduction", "parallel", "reduction"] {
-    %sub_3d = memref.subview %input_3d[%i, %j, %k][2, 4, 8][1, 1, 1]
+    %sub_3d = memref.subview %i3d_[%i, %j, %k][2, 4, 8][1, 1, 1]
       : memref<16x24x32xf32> to memref<2x4x8xf32, #map_1>
-    %sub_2d = memref.subview %input_2d[%i, %k][2, 8][1, 1]
+    %sub_2d = memref.subview %i2d_[%i, %k][2, 8][1, 1]
       : memref<16x32xf32> to memref<2x8xf32, #map_2>
-    %sub_1d = memref.subview %input_1d[%j] [4] [1]
+    %sub_1d = memref.subview %i1d_[%j] [4] [1]
       : memref<24xf32> to memref<4xf32, #map_3>
-    %sub_out = memref.subview %output[%j] [4] [1]
+    %sub_out = memref.subview %o_[%j] [4] [1]
       : memref<24xf32> to memref<4xf32, #map_3>
     linalg.generic #trait_6
       ins(%sub_3d, %sub_2d, %sub_1d
         : memref<2x4x8xf32, #map_1>,
           memref<2x8xf32, #map_2>,
           memref<4xf32, #map_3>)
       outs(%sub_out : memref<4xf32, #map_3>)  {
     ^bb0(%i3d: f32, %i2d: f32, %i1d: f32, %o: f32):