[linalg] Add basic support for dynamic tensor results in TensorToBuffers.cpp.

[pifon2a]

Summary

The simplest case is when the indexing maps are DimIds in every component. This covers cwise ops.

Also:

• Expose populateConvertLinalgOnTensorsToBuffersPatterns in Transforms.h
• Expose emitLoopRanges in Transforms.h

[nicolasvasilache]

Very cool, thanks @pifon2a !
Next step, Gaussian-elimination like range-inference ? ;)

[pifon2a]

Rebase

[nicolasvasilache]

Nice, thanks for this!

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 namespace mlir {
+
+class BufferAssignmentTypeConverter;
+
 namespace linalg {
 
 struct LinalgFusionOptions;
 struct LinalgTilingOptions;
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
struct TiledAndFusedLinalgOps {
   SmallVector<Operation *, 4> unfusedLoops;
 };
 
 /// Populates patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
     MLIRContext *context, SmallVectorImpl<OwningRewritePatternList> &patterns,
     ArrayRef<int64_t> tileSizes);
 
+/// Populates the given list with patterns to convert Linalg operations on
+/// tensors to buffers.
+void populateConvertLinalgOnTensorsToBuffersPatterns(
+    MLIRContext *context, BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns);
+
 /// Performs standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `interchangeVector`
 /// The permutation is expressed as a list of integers that specify
 /// the new ordering of the loop nest. The length of `interchangeVector`
 /// must be equal to the length of `tileSizes`.
 /// An empty vector is interpreted as the identity permutation and the
 /// transformation returns early.
 ///

 /// true.
 ///
 /// Returns the modified linalg op (the modification happens in place) as well
 /// as all the copy ops created.
 Optional<LinalgOp> promoteSubViews(OpBuilder &b, LinalgOp op,
                                    LinalgPromotionOptions options,
                                    OperationFolder *folder = nullptr);
 
+/// Creates a number of ranges equal to the number of dimensions in the `map`.
+/// The returned ranges correspond to the loop ranges, in the proper order, for
+/// which new loops will be created.
+/// The function supports only maps that are invertible and have results of type
+/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
+/// It expects a non-inverted, concatenated map and last values in
+/// allViewSizes will be applied to the symbols in the map if it contains any.
+SmallVector<Range, 4> emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
+                                     ValueRange viewSizes);
+
 /// Emit a suitable vector form for a Linalg op with fully static shape.
 void vectorizeLinalgOp(OpBuilder &builder, Operation *op);
 
 /// Emits a loop nest of `LoopTy` with the proper body for `op`.
 template <typename LoopTy>
 Optional<LinalgLoops> linalgLowerOpToLoops(OpBuilder &builder, Operation *op);
 
 /// Emits a loop nest of `scf.for` with the proper body for `op`.

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

 static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
                                         Optional<AffineMap> permutation) {
   return permutation ? applyMapToValues(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(),
                                         permutation.getValue(), ivs)
                      : SmallVector<Value, 4>(ivs.begin(), ivs.end());
 }
 
-/// Creates a number of ranges equal to the number of dimensions in the `map`.
-/// The returned ranges correspond to the loop ranges, in the proper order, for
-/// which new loops will be created.
-/// The function supports only maps that are invertible and have results of type
-/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
-/// It expects a non-inverted, concatenated map and last values in
-/// allViewSizes will be applied to the symbols in the map if it contains any.
-static SmallVector<Range, 4> emitLoopRanges(OpBuilder &b, Location loc,
-                                            AffineMap map,
-                                            ValueRange viewSizes) {
-  unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
-  unsigned numSym = map.getNumSymbols();
-  assert(viewSizes.size() == numRes + numSym &&
-         "viewSizes must contain sizes of all views and values for symbols");
-  SmallVector<Range, 4> res(numDims);
-  for (unsigned idx = 0; idx < numRes; ++idx) {
-    auto result = map.getResult(idx);
-    if (auto d = result.dyn_cast<AffineDimExpr>()) {
-      if (res[d.getPosition()].offset)
-        continue;
-      res[d.getPosition()] =
-          Range{std_constant_index(0), viewSizes[idx], std_constant_index(1)};
-    }
-
-    // If the access pattern is of form (m, n)[s] -> (m + n - s floordiv 2),
-    // then the bounds are:
-    //   (s floordiv 2) <= m <= (size(m) + s floordiv 2 - s + 1).
-    // where size(n) is applied to the symbol s.
-    // This is done statically now.
-    if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
-      auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
-      auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
-      if (!lhs || !rhs || binOp.getKind() != AffineExprKind::Add ||
-          lhs.getKind() != AffineExprKind::Add ||
-          rhs.getKind() != mlir::AffineExprKind::Mul)
-        continue;
-
-      auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
-      auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
-      auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
-      auto minusOne = rhs.getRHS().dyn_cast<AffineConstantExpr>();
-      if (!m || !n || !fDiv || !minusOne ||
-          fDiv.getKind() != AffineExprKind::FloorDiv ||
-          fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
-          fDiv.getRHS().getKind() != AffineExprKind::Constant)
-        continue;
-
-      auto s = fDiv.getLHS().dyn_cast<AffineSymbolExpr>();
-      if (minusOne.getValue() != -1)
-        continue;
-
-      int mPos = m.getPosition();
-      AffineExpr one = getAffineConstantExpr(1, s.getContext());
-      AffineExpr sizeOfM = getAffineSymbolExpr(numSym, s.getContext());
-      // Construction of upper bound (size(m) + s floordiv 2 - s + 1).
-      AffineExpr upperOffsetExpr = sizeOfM + fDiv + one - s;
-      AffineMap fromMap = AffineMap::get(numDims, numSym + 1, fDiv);
-      AffineMap toMap = AffineMap::get(numDims, numSym + 1, upperOffsetExpr);
-      SmallVector<Value, 8> values(viewSizes.begin(),
-                                   viewSizes.begin() + numDims);
-      values.insert(values.end(), viewSizes.begin() + numRes, viewSizes.end());
-      values.push_back(viewSizes[mPos]);
-      // Construction of the lower bound (s floordiv 2).
-      Value from = applyMapToValues(b, loc, fromMap, values).front();
-      Value to = applyMapToValues(b, loc, toMap, values).front();
-      res[mPos] = Range{from, to, std_constant_index(1)};
-    }
-  }
-  return res;
-}
-
 template <typename IndexedValueType, typename OpType>
 static void inlineRegionAndEmitStore(OpType op, ArrayRef<Value> indexedValues,
                                      ArrayRef<SmallVector<Value, 8>> indexing,
                                      ArrayRef<Value> outputBuffers) {
   assert(op.getOperation()->getNumRegions() == 1 &&
          "Expected single region op");
   auto &b = ScopedContext::getBuilderRef();
   auto &block = op.region().front();
if (isa<ConvDHWOp>(op))
     return linalgOpToLoopsImpl<LoopTy, ConvDHWOp>(op, builder);
   if (isa<ConvNDHWCOp>(op))
     return linalgOpToLoopsImpl<LoopTy, ConvNDHWCOp>(op, builder);
   if (isa<ConvNCDHWOp>(op))
     return linalgOpToLoopsImpl<LoopTy, ConvNCDHWOp>(op, builder);
   llvm_unreachable("Unexpected op in linalgOpToLoopsImpl");
 }
 
+SmallVector<Range, 4> mlir::linalg::emitLoopRanges(OpBuilder &b, Location loc,
+                                                   AffineMap map,
+                                                   ValueRange viewSizes) {
+  unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
+  unsigned numSym = map.getNumSymbols();
+  assert(viewSizes.size() == numRes + numSym &&
+         "viewSizes must contain sizes of all views and values for symbols");
+  SmallVector<Range, 4> res(numDims);
+  for (unsigned idx = 0; idx < numRes; ++idx) {
+    auto result = map.getResult(idx);
+    if (auto d = result.dyn_cast<AffineDimExpr>()) {
+      if (res[d.getPosition()].offset)
+        continue;
+      res[d.getPosition()] =
+          Range{std_constant_index(0), viewSizes[idx], std_constant_index(1)};
+    }
+
+    // If the access pattern is of form (m, n)[s] -> (m + n - s floordiv 2),
+    // then the bounds are:
+    //   (s floordiv 2) <= m <= (size(m) + s floordiv 2 - s + 1).
+    // where size(n) is applied to the symbol s.
+    // This is done statically now.
+    if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
+      auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
+      auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
+      if (!lhs || !rhs || binOp.getKind() != AffineExprKind::Add ||
+          lhs.getKind() != AffineExprKind::Add ||
+          rhs.getKind() != mlir::AffineExprKind::Mul)
+        continue;
+
+      auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
+      auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
+      auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
+      auto minusOne = rhs.getRHS().dyn_cast<AffineConstantExpr>();
+      if (!m || !n || !fDiv || !minusOne ||
+          fDiv.getKind() != AffineExprKind::FloorDiv ||
+          fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
+          fDiv.getRHS().getKind() != AffineExprKind::Constant)
+        continue;
+
+      auto s = fDiv.getLHS().dyn_cast<AffineSymbolExpr>();
+      if (minusOne.getValue() != -1)
+        continue;
+
+      int mPos = m.getPosition();
+      AffineExpr one = getAffineConstantExpr(1, s.getContext());
+      AffineExpr sizeOfM = getAffineSymbolExpr(numSym, s.getContext());
+      // Construction of upper bound (size(m) + s floordiv 2 - s + 1).
+      AffineExpr upperOffsetExpr = sizeOfM + fDiv + one - s;
+      AffineMap fromMap = AffineMap::get(numDims, numSym + 1, fDiv);
+      AffineMap toMap = AffineMap::get(numDims, numSym + 1, upperOffsetExpr);
+      SmallVector<Value, 8> values(viewSizes.begin(),
+                                   viewSizes.begin() + numDims);
+      values.insert(values.end(), viewSizes.begin() + numRes, viewSizes.end());
+      values.push_back(viewSizes[mPos]);
+      // Construction of the lower bound (s floordiv 2).
+      Value from = applyMapToValues(b, loc, fromMap, values).front();
+      Value to = applyMapToValues(b, loc, toMap, values).front();
+      res[mPos] = Range{from, to, std_constant_index(1)};
+    }
+  }
+  return res;
+}
+
 /// Emits a loop nest with the proper body for `op`.
 template <typename LoopTy>
 Optional<LinalgLoops> mlir::linalg::linalgLowerOpToLoops(OpBuilder &builder,
                                                          Operation *op) {
   return linalgOpToLoopsImplSwitch<LoopTy>(op, builder);
 }
 
 template Optional<LinalgLoops>

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

 //===- TensorsToBuffers.cpp - Transformation from tensors to buffers ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the conversion from tensors to buffers on Linalg
 // operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferPlacement.h"
 
-using namespace mlir;
-
 namespace {
+
+using namespace ::mlir;
+using namespace ::mlir::linalg;
+
+SmallVector<Range, 4>
+computeLoopRanges(Location loc, linalg::GenericOp linalgOp, OpBuilder *b) {
+  auto indexingMaps = llvm::to_vector<4>(
+      linalgOp.indexing_maps().getAsValueRange<AffineMapAttr>());
+  auto inputIndexingMaps =
+      llvm::makeArrayRef(indexingMaps).take_front(linalgOp.getNumInputs());
+
+  mlir::edsc::ScopedContext scope(*b, loc);
+  return emitLoopRanges(scope.getBuilderRef(), loc,
+                        concatAffineMaps(inputIndexingMaps),
+                        getShape(*b, linalgOp));
+}
+
+Value maybeConvertToIndex(Location loc, Value val, OpBuilder *b) {
+  if (val.getType().isIndex())
+    return val;
+  return b->create<IndexCastOp>(loc, val, b->getIndexType());
+}
+
+LogicalResult allocateBuffersForResults(Location loc,
+                                        linalg::GenericOp linalgOp,
+                                        linalg::GenericOpAdaptor &adaptor,
+                                        SmallVectorImpl<Value> *resultBuffers,
+                                        OpBuilder *b) {
+  // Lazily compute loopRanges.
+  SmallVector<Range, 4> loopRanges;
+
+  // Allocate a buffer for every tensor result.
+  for (auto en : llvm::enumerate(linalgOp.getResultTypes())) {
+    size_t resultIndex = en.index();
+    Type resultType = en.value();
+
+    auto tensorType = resultType.dyn_cast<RankedTensorType>();
+    if (tensorType == nullptr) {
+      linalgOp.emitOpError()
+          << "tensor to buffer conversion expects ranked tensor results";
+      return failure();
+    }
+    auto tensorShape = tensorType.getShape();
+    auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
+
+    // Allocate buffers for init tensors that are assumed to fold onto the first
+    // results.
+    // TODO: update this assumption because the reality is more complex
+    // under linalg on tensor based transformations.
+    bool foldedInitTensor = resultIndex < linalgOp.getNumInitTensors();
+    if (foldedInitTensor) {
+      // Dealing with an init tensor requires distinguishing between 1-use
+      // and many-use cases which would create aliasing and WAR hazards.
+      Value initTensor = linalgOp.getInitTensor(resultIndex);
+      Value initBuffer = adaptor.init_tensors()[resultIndex];
+      if (initTensor.hasOneUse()) {
+        resultBuffers->push_back(initBuffer);
+        continue;
+      }
+      SmallVector<Value, 4> dynOperands;
+      for (auto dim : llvm::enumerate(tensorShape)) {
+        if (dim.value() == TensorType::kDynamicSize) {
+          dynOperands.push_back(b->create<DimOp>(loc, initTensor, dim.index()));
+        }
+      }
+      auto alloc = b->create<AllocOp>(loc, memrefType, dynOperands);
+      b->create<linalg::CopyOp>(loc, initBuffer, alloc);
+      resultBuffers->push_back(alloc);
+      continue;
+    }
+
+    // Allocate buffers for statically-shaped results.
+    if (memrefType.hasStaticShape()) {
+      resultBuffers->push_back(b->create<AllocOp>(loc, memrefType));
+      continue;
+    }
+
+    // Perform a naive shape inference for the dynamically-shaped results.
+    // Extract the required element out of the vector.
+    SmallVector<Value, 4> dynOperands;
+    auto resultIndexingMap = linalgOp.getOutputIndexingMap(resultIndex);
+    for (auto shapeElement : llvm::enumerate(tensorType.getShape())) {
+      if (loopRanges.empty())
+        loopRanges = computeLoopRanges(loc, linalgOp, b);
+
+      if (shapeElement.value() != ShapedType::kDynamicSize)
+        continue;
+
+      AffineExpr expr = resultIndexingMap.getResult(shapeElement.index());
+      switch (expr.getKind()) {
+      case AffineExprKind::DimId: {
+        int64_t loopIndex = expr.cast<AffineDimExpr>().getPosition();
+        Value size = maybeConvertToIndex(loc, loopRanges[loopIndex].size, b);
+        dynOperands.push_back(size);
+        break;
+      }
+      default:
+        return failure();
+      }
+    }
+    resultBuffers->push_back(b->create<AllocOp>(loc, memrefType, dynOperands));
+  }
+  return success();
+}
+
 /// A pattern to convert Generic Linalg operations which work on tensors to
 /// use buffers. A buffer is allocated using BufferAssignmentPlacer for
 /// each operation result. BufferPlacement pass should be later used to move
 /// Alloc operations to the correct positions and insert the missing Dealloc
 /// operations in the correct places.
 class GenericOpConverter
     : public BufferAssignmentOpConversionPattern<linalg::GenericOp> {
 public:
   using BufferAssignmentOpConversionPattern<
       linalg::GenericOp>::BufferAssignmentOpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
+  matchAndRewrite(linalg::GenericOp linalgOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
-    linalg::GenericOpAdaptor adaptor(operands,
-                                     op.getOperation()->getAttrDictionary());
+    linalg::GenericOpAdaptor adaptor(
+        operands, linalgOp.getOperation()->getAttrDictionary());
 
     // All inputs need to be turned into buffers first. Until then, bail out.
     if (llvm::any_of(adaptor.inputs(),
                      [](Value in) { return !in.getType().isa<MemRefType>(); }))
       return failure();
 
     // All init_tensors need to be turned into buffers first. Until then, bail
     // out.
     if (llvm::any_of(adaptor.init_tensors(),
                      [](Value in) { return !in.getType().isa<MemRefType>(); }))
       return failure();
 
-    Location loc = op.getLoc();
-    SmallVector<Value, 2> newOutputBuffers;
-    newOutputBuffers.reserve(op.getNumOutputs());
-    newOutputBuffers.append(adaptor.output_buffers().begin(),
+    Location loc = linalgOp.getLoc();
+    SmallVector<Value, 2> newOutputBuffers(adaptor.output_buffers().begin(),
                                            adaptor.output_buffers().end());
 
-    // Update all types to memref types.
-    // Assume the init tensors fold onto the first results.
-    // TODO: update this assumption because the reality is more complex under
-    // linalg on tensor based transformations.
-    for (auto en : llvm::enumerate(op.getResultTypes())) {
-      auto type = en.value().cast<ShapedType>();
-      if (!type.hasStaticShape())
-        return rewriter.notifyMatchFailure(
-            op, "dynamic shapes not currently supported");
-      auto memrefType = MemRefType::get(type.getShape(), type.getElementType());
-      bool foldedInitTensor = en.index() < op.getNumInitTensors();
-      if (foldedInitTensor) {
-        // Dealing with an init tensor requires distinguishing between 1-use
-        // and many-use cases which would create aliasing and WAR hazards.
-        Value initTensor = op.getInitTensor(en.index());
-        Value initBuffer = adaptor.init_tensors()[en.index()];
-        if (initTensor.hasOneUse()) {
-          newOutputBuffers.push_back(initBuffer);
-          continue;
-        }
-        auto alloc = rewriter.create<AllocOp>(loc, memrefType);
-        rewriter.create<linalg::CopyOp>(loc, initBuffer, alloc);
-        newOutputBuffers.push_back(alloc);
-      } else {
-        auto alloc = rewriter.create<AllocOp>(loc, memrefType);
-        newOutputBuffers.push_back(alloc);
-      }
+    if (failed(allocateBuffersForResults(loc, linalgOp, adaptor,
+                                         &newOutputBuffers, &rewriter))) {
+      linalgOp.emitOpError()
+          << "Failed to allocate buffers for tensor results.";
+      return failure();
     }
 
     // Generate a new linalg operation that works on buffers.
-    auto linalgOp = rewriter.create<linalg::GenericOp>(
+    auto newLinalgOp = rewriter.create<linalg::GenericOp>(
         loc,
-        /*resultTensorTypes=*/ArrayRef<Type>{},
+        /*resultTensorTypes=*/llvm::None,
         /*inputs=*/adaptor.inputs(),
         /*outputBuffers=*/newOutputBuffers,
-        /*initTensors=*/ValueRange{}, op.indexing_maps(), op.iterator_types(),
-        op.docAttr(), op.library_callAttr(), op.symbol_sourceAttr());
+        /*initTensors=*/llvm::None, linalgOp.indexing_maps(),
+        linalgOp.iterator_types(), linalgOp.docAttr(),
+        linalgOp.library_callAttr(), linalgOp.symbol_sourceAttr());
 
     // Create a new block in the region of the new Generic Op.
-    Block &oldBlock = op.getRegion().front();
-    Region &newRegion = linalgOp.region();
+    Block *oldBlock = linalgOp.getBody();
+    Region &newRegion = newLinalgOp.region();
     Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
-                                           oldBlock.getArgumentTypes());
+                                           oldBlock->getArgumentTypes());
 
-    // Add the result arguments that do not come from init_tensors to the new
-    // block.
-    // TODO: update this assumption because the reality is more complex under
-    // linalg on tensor based transformations.
-    for (Value v :
-         ValueRange(newOutputBuffers).drop_front(adaptor.init_tensors().size()))
+    // Add the result arguments to the new block.
+    for (Value v : newOutputBuffers)
       newBlock->addArgument(v.getType().cast<MemRefType>().getElementType());
 
     // Clone the body of the old block to the new block.
     BlockAndValueMapping mapping;
-    for (unsigned i = 0; i < oldBlock.getNumArguments(); i++)
-      mapping.map(oldBlock.getArgument(i), newBlock->getArgument(i));
+    mapping.map(oldBlock->getArguments(), newBlock->getArguments());
 
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPointToEnd(newBlock);
-    for (auto &op : oldBlock.getOperations()) {
+    for (auto &op : oldBlock->getOperations()) {
       Operation *clonedOp = rewriter.clone(op, mapping);
       mapping.map(op.getResults(), clonedOp->getResults());
     }
 
     // Replace the results of the old op with the new output buffers.
-    rewriter.replaceOp(op, newOutputBuffers);
+    rewriter.replaceOp(linalgOp, newOutputBuffers);
     return success();
   }
 };
 
-/// Populate the given list with patterns to convert Linalg operations on
-/// tensors to buffers.
-static void populateConvertLinalgOnTensorsToBuffersPattern(
-    MLIRContext *context, BufferAssignmentTypeConverter *converter,
-    OwningRewritePatternList *patterns) {
-  populateWithBufferAssignmentOpConversionPatterns<
-      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, converter,
-                                                      patterns);
-  patterns->insert<GenericOpConverter>(context, converter);
-}
-
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct ConvertLinalgOnTensorsToBuffers
     : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
                            [&](Type type) { return type.isa<MemRefType>(); }) &&
              converter.isLegal(&funcOp.getBody());
     });
 
     converter.setResultConversionKind<RankedTensorType, MemRefType>(
         BufferAssignmentTypeConverter::AppendToArgumentsList);
 
     OwningRewritePatternList patterns;
-    populateConvertLinalgOnTensorsToBuffersPattern(&context, &converter,
+    populateConvertLinalgOnTensorsToBuffersPatterns(&context, &converter,
+                                                    &patterns);
+    populateWithBufferAssignmentOpConversionPatterns<
+        mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(&context, &converter,
                                                         &patterns);
     if (failed(applyFullConversion(this->getOperation(), target, patterns)))
       this->signalPassFailure();
   }
 };
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertLinalgOnTensorsToBuffersPass() {
   return std::make_unique<ConvertLinalgOnTensorsToBuffers>();
 }
+
+void mlir::linalg::populateConvertLinalgOnTensorsToBuffersPatterns(
+    MLIRContext *context, BufferAssignmentTypeConverter *converter,
+    OwningRewritePatternList *patterns) {
+  patterns->insert<GenericOpConverter>(context, converter);
+}

mlir/test/Dialect/Linalg/tensors-to-buffers.mlir

 // RUN: mlir-opt -convert-linalg-on-tensors-to-buffers -buffer-placement -split-input-file %s | FileCheck %s
 
 #map0 = affine_map<(d0) -> (d0)>
 
-// CHECK-LABEL: func @multiple_results_generic_op
-func @multiple_results_generic_op(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
-    %0, %1 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel"]}
-      ins(%arg0 : tensor<4xf32>) {
+// CHECK-LABEL: func @multiple_results
+func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+    %0, %1 = linalg.generic {
+      indexing_maps = [#map0, #map0, #map0],
+      iterator_types = ["parallel"]
+    } ins(%arg0 : tensor<4xf32>) {
       ^bb0(%gen_arg1: f32):
         %tmp1 = exp %gen_arg1 : f32
         linalg.yield %tmp1, %tmp1 : f32, f32
     } -> tensor<4xf32>, tensor<4xf32>
     return %0, %1 : tensor<4xf32>, tensor<4xf32>
 }
 //      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]], %[[ARG2_RESULT:.*]]: [[TYPE]])
 //      CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]]
 //      CHECK: %[[SECOND_ALLOC:.*]] = alloc() : [[TYPE]]

 //      CHECK: return
 
 // -----
 
 #map0 = affine_map<(d0) -> (d0)>
 
 // CHECK-LABEL: func @chained_operations
 func @chained_operations(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-    %0 = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
-      ins(%arg0 : tensor<4xf32>) {
+    %0 = linalg.generic {
+      indexing_maps = [#map0, #map0],
+      iterator_types = ["parallel"]
+    } ins(%arg0 : tensor<4xf32>) {
       ^bb0(%gen_arg1: f32):
         %tmp1 = exp %gen_arg1 : f32
         linalg.yield %tmp1 : f32
     } -> tensor<4xf32>
-    %1 = linalg.generic {args_in = 1 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0], iterator_types = ["parallel"]}
-      ins(%0 : tensor<4xf32>) {
+
+    %1 = linalg.generic {
+      indexing_maps = [#map0, #map0],
+      iterator_types = ["parallel"]
+    } ins(%0 : tensor<4xf32>) {
       ^bb0(%gen_arg2: f32):
         %tmp2 = exp %gen_arg2 : f32
         linalg.yield %tmp2 : f32
     } -> tensor<4xf32>
     return %1 : tensor<4xf32>
 }
 //      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]], %[[ARG1_RESULT:.*]]: [[TYPE]])
 //      CHECK: %[[FIRST_ALLOC:.*]] = alloc() : [[TYPE]]
 //      CHECK: linalg.generic

 
 // -----
 
 // CHECK-LABEL: func @no_linalg_op
 func @no_linalg_op(%arg0: f32) -> (f32, f32) {
   %0 = mulf %arg0, %arg0 : f32
   return %0, %0 : f32, f32
 }
-//      CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]]) -> ([[TYPE]], [[TYPE]])
-//      CHECK: %[[RESULT:.*]] = mulf %[[NEW_ARG0]], %[[NEW_ARG0]] : [[TYPE]]
-//      CHECK: return %[[RESULT]], %[[RESULT]] : [[TYPE]], [[TYPE]]
+// CHECK: (%[[NEW_ARG0:.*]]: [[TYPE:.*]]) -> ([[TYPE]], [[TYPE]])
+// CHECK: %[[RESULT:.*]] = mulf %[[NEW_ARG0]], %[[NEW_ARG0]] : [[TYPE]]
+// CHECK: return %[[RESULT]], %[[RESULT]] : [[TYPE]], [[TYPE]]
+
+// -----
+
+#map_2d = affine_map<(d0, d1) -> (d0, d1)>
+#map_2d_inv = affine_map<(d0, d1) -> (d1, d0)>
+
+func @dynamic_results(%arg0: tensor<?x?xf32>)
+         -> (tensor<?x?xf32>, tensor<?x?xf32>) {
+    %0, %1 = linalg.generic {
+      indexing_maps = [#map_2d, #map_2d, #map_2d_inv],
+      iterator_types = ["parallel", "parallel"]
+    } ins(%arg0 : tensor<?x?xf32>) {
+      ^bb0(%gen_arg1: f32):
+        %tmp1 = exp %gen_arg1 : f32
+        linalg.yield %tmp1, %tmp1 : f32, f32
+    } -> tensor<?x?xf32>, tensor<?x?xf32>
+    return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
+}
+
+// CHECK: #map0 = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK: #map1 = affine_map<(d0, d1) -> (d1, d0)>
+
+// CHECK-LABEL: func @dynamic_results
+// CHECK-SAME: (%[[INPUT:.*]]: [[TYPE:.*]], %[[OUT_1:.*]]: [[TYPE]], %[[OUT_2:.*]]: [[TYPE]]) {
+// CHECK: %[[C0:.*]] = constant 0 : index
+// CHECK: %[[DIM_0:.*]] = dim %[[INPUT]], %[[C0]] : [[TYPE]]
+// CHECK: %[[C1:.*]] = constant 1 : index
+// CHECK: %[[DIM_1:.*]] = dim %[[INPUT]], %[[C1]] : [[TYPE]]
+// CHECK: %[[OUT_BUF_1:.*]] = alloc(%[[DIM_0]], %[[DIM_1]]) : [[TYPE]]
+// CHECK: %[[OUT_BUF_2:.*]] = alloc(%[[DIM_1]], %[[DIM_0]]) : [[TYPE]]
+
+// CHECK: linalg.generic {indexing_maps = [#map0, #map0, #map1], {{.*}}}
+// CHECK-SAME: ins(%[[INPUT]] : [[TYPE]])
+// CHECK-SAME: outs(%[[OUT_BUF_1]], %[[OUT_BUF_2]] : [[TYPE]], [[TYPE]]) {
+
+// CHECK: linalg.copy(%[[OUT_BUF_1]], %[[OUT_1]]) : [[TYPE]], [[TYPE]]
+// CHECK: dealloc %[[OUT_BUF_1]] : [[TYPE]]
+// CHECK: linalg.copy(%[[OUT_BUF_2]], %[[OUT_2]]) : [[TYPE]], [[TYPE]]
+// CHECK: dealloc %[[OUT_BUF_2]] : [[TYPE]]
+// CHECK: return