[mlir][linalg] Expose function to create op on buffers during bufferization.

[tpopp]

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

79–109

I think this might need a clang format run.

212

This looks like "replace new with new"

[pifon2a]

Thank you!

[pifon2a]

Addressed the comments.

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

 /// collapsing (introducing) unit-extent dimensions.
 void populateFoldUnitDimsReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns);
 
 /// Populates the given list with patterns to bufferize linalg ops.
 void populateLinalgBufferizePatterns(BufferizeTypeConverter &converter,
                                      RewritePatternSet &patterns);
 
+/// Create linalg op on buffers given the original tensor-based operation and
+/// the buffers for the outputs.
+LinalgOp createLinalgOpOnBuffers(ConversionPatternRewriter &rewriter,
+                                 LinalgOp linalgOp, ValueRange inputs,
+                                 ValueRange outputs);
+
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to inline constant operands into linalg generic ops.
 void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
 
 /// Pattern to convert TiledLoopOp to SCF loops.

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

for (auto en : llvm::enumerate(linalgOp->getResultTypes())) {
     }
 
     resultBuffers.push_back(b.create<memref::AllocOp>(
         loc, memrefType, getDynOperands(loc, resultTensor, b)));
   }
   return success();
 }
 
-/// Specialization for `linalg::GenericOp`.
-/// A pattern to convert Generic Linalg operations which work on tensors to
-/// use buffers. BufferPlacement pass should be later used to move
-/// Alloc operations to the correct positions and insert the missing Dealloc
-/// operations in the correct places.
-static void
-finalizeBufferAllocationForGenericOp(ConversionPatternRewriter &rewriter,
-                                     GenericOp genericOp, ValueRange inputs,
+/// Create linalg op on buffers given the original tensor-based operation and
+/// the buffers for the outputs.
+LinalgOp
+mlir::linalg::createLinalgOpOnBuffers(ConversionPatternRewriter &rewriter,
+                                      LinalgOp linalgOp, ValueRange inputs,
                                       ValueRange outputs) {
+  if (auto genericOp = mlir::dyn_cast<GenericOp>(*linalgOp)) {
     // Generate a new linalg operation that works on buffers.
     auto newGenericOp = rewriter.create<GenericOp>(
         genericOp.getLoc(),
         /*resultTensorTypes=*/llvm::None,
         /*inputs=*/inputs,
         /*outputs=*/outputs, genericOp.indexing_maps(),
         genericOp.iterator_types(), genericOp.docAttr(),
         genericOp.library_callAttr());
 
     // Create a new block in the region of the new Generic Op.
     Block *oldBlock = genericOp.getBody();
     Region &newRegion = newGenericOp.region();
     Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
                                            oldBlock->getArgumentTypes());
 
     // Clone the body of the old block to the new block.
     BlockAndValueMapping mapping;
     mapping.map(oldBlock->getArguments(), newBlock->getArguments());
 
     OpBuilder::InsertionGuard guard(rewriter);
     rewriter.setInsertionPointToEnd(newBlock);
     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(genericOp, outputs);
+    return newGenericOp;
   }
-
-/// Specialization for all other `linalg::LinalgOp`.
-static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter,
-                                     linalg::LinalgOp linalgOp,
-                                     ValueRange inputs, ValueRange outputs) {
-  assert(!isa<linalg::GenericOp>(linalgOp.getOperation()));
   SmallVector<Value, 8> newOperands = inputs;
   newOperands.append(outputs.begin(), outputs.end());
-  linalgOp.clone(rewriter, linalgOp.getLoc(),
+  return linalgOp.clone(rewriter, linalgOp.getLoc(),
                         /*resultTypes=*/ArrayRef<Type>{}, newOperands);
-  // Replace the results of the old op with the new output buffers.
-  rewriter.replaceOp(linalgOp, outputs);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization patterns.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
matchAndRewrite(LinalgOp op, ArrayRef<Value> operands,
     Location loc = op.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
     if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
                                          newOutputBuffers, rewriter))) {
       return op.emitOpError()
              << "Failed to allocate buffers for tensor results.";
     }
-
-    // Delegate to the linalg generic pattern.
-    if (auto genericOp = dyn_cast<linalg::GenericOp>(*op)) {
-      finalizeBufferAllocationForGenericOp(rewriter, genericOp,
-                                           adaptor.inputs(), newOutputBuffers);
-      return success();
-    }
-
-    finalizeBufferAllocation(rewriter, op, adaptor.inputs(), newOutputBuffers);
+    LinalgOp newLinalgOp = createLinalgOpOnBuffers(
+        rewriter, op, adaptor.inputs(), newOutputBuffers);
+    // Replace the results of the old op with the new output buffers.
+    rewriter.replaceOp(newLinalgOp, newOutputBuffers);
     return success();
   }
 };
 
 /// Convert `extract_slice %t [offsets][sizes][strides] -> %st` to an
 /// alloc + copy pattern.
 /// ```
 ///   %a = alloc(sizes)