[mlir][linalg] Use populateLinalgBufferizePatterns instead of copypasta.

[silvas]

Summary

The non-copypasta patterns set up proper type materializations, which
revealed the need for helper patterns that eliminate materializations,
which due to the complexity of the test cases (e.g. control flow),
tickles this particular issue.

The observation is that TestBufferPlacementPreparationPass is a
"finalizing" bufferization pass. That is, it does a "full" conversion
and needs all the materializations gone. So we expose a new helper
populateBufferizeEliminateMaterializationsPatternsAndLegality for
setting that up. This helper will soon be merged into
populateWithBufferizeOpConversionPatterns (which will be renamed
populateBufferizeFinalizationPatterns), but for now remains separate.

Note: LinalgBufferizePass is currently "finalizing", but I didn't apply
this there because I'm in the process of making that pass
non-finalizing (I'm making the finalizing process much nicer,
so that each dialect doesn't need its own finalizing bufferization).

[herhut]

Thanks for your work on cleaning up finalization. This all looks much nicer now.

I'd be happy to see the finalization part landing but would avoid introducing the dependency if it gets cleaned up soon enough.

mlir/lib/Transforms/Bufferize.cpp

104

Should this pattern fail if the operand is not a memref? Instead of creating wrong IR?

mlir/test/lib/Transforms/TestBufferPlacement.cpp

47

There is work underway by @dfki-albo and @dfki-heme to rephrase these tests so that they no longer depend on linalg operations. That will also remove the copypasta here.

[silvas]

mlir/test/lib/Transforms/TestBufferPlacement.cpp

47

Are there patches under-way? I'm going to get to this in the next couple days.

[silvas]

I'd be happy to see the finalization part landing but would avoid introducing the dependency if it gets cleaned up soon enough.

What do you mean "the dependency"? Or perhaps, what is the actionable request for this patch?

mlir/lib/Transforms/Bufferize.cpp

104

I don't think that is needed.

tensor (which is the only possibility other than memref), is illegal according to the TypeConverter that this pattern is meant to be used with, so operands cannot contain a tensor (the invariant is that operands are all legal types).

Most conversion patterns will create illegal IR / crash if invoked with an improper type converter. E.g. I can invoke mhlo to lmhlo patterns with a TypeConverter that converts tensor to i1 (with some bogus materialization) and then everything will crash.

[silvas]

In D89923#2348135, @silvas wrote:

I'd be happy to see the finalization part landing but would avoid introducing the dependency if it gets cleaned up soon enough.

What do you mean "the dependency"? Or perhaps, what is the actionable request for this patch?

Oh, I think I understand what you mean. I guess I'll just update this patch to not test the linalg case.

[silvas]

Ok, abandoning this patch. See the updated patch in https://reviews.llvm.org/D89979 which just generally refactors the pass to be a pure finalization test (and removes the dependence on linalg.generic bufferization).
I've added @dfki-albo @dfki-heme as reviewers there.

mlir/include/mlir/Transforms/Bufferize.h

private:
   SmallVector<DecomposeTypeConversionCallFn, 2> decomposeTypeConversions;
 };
 
 /// Marks ops used by bufferization for type conversion materializations as
 /// "legal" in the given ConversionTarget.
 ///
 /// This function should be called by all bufferization passes using
 /// BufferizeTypeConverter so that materializations work proprely. One exception
-/// is bufferization passes doing "full" conversions, where it can be desirable
-/// for even the materializations to remain illegal so that they are eliminated.
+/// is "finalizing" bufferization passes (such as those doing "full"
+/// conversions), where it can be desirable for even the materializations to
+/// remain illegal so that they are eliminated. For "finalizing" passes,
+/// populateBufferizeEliminateMaterializationsPatternsAndLegality is
+/// recommended.
 void populateBufferizeMaterializationLegality(ConversionTarget &target);
 
+/// Populate patterns and set up legality for eliminating materialization ops.
+///
+/// Some bufferization passes are "finalizing". They expect that all tensors are
+/// converted to memrefs, including function arguments, basic block arguments,
+/// etc. In these cases, all materializations that may have been inserted for
+/// structural legality of the IR need to be eliminated.
+void populateBufferizeEliminateMaterializationsPatternsAndLegality(
+    MLIRContext *context, BufferizeTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target);
+
 /// Helper conversion pattern that encapsulates a BufferizeTypeConverter
 /// instance.
 template <typename SourceOp>
 class BufferizeOpConversionPattern : public OpConversionPattern<SourceOp> {
 public:
   explicit BufferizeOpConversionPattern(MLIRContext *context,
                                         BufferizeTypeConverter &converter,
                                         PatternBenefit benefit = 1)

mlir/lib/Transforms/Bufferize.cpp

BufferizeTypeConverter::getResultConversionKind(Type origin, Type converted) {
   return KeepAsFunctionResult;
 }
 
 void mlir::populateBufferizeMaterializationLegality(ConversionTarget &target) {
   target.addLegalOp<TensorLoadOp, TensorToMemrefOp>();
 };
 
 //===----------------------------------------------------------------------===//
+// populateBufferizeEliminateMaterializationsPatternsAndLegality
+//===----------------------------------------------------------------------===//
+
+namespace {
+class BufferizeTensorLoadOp : public OpConversionPattern<TensorLoadOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorLoadOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    TensorLoadOp::Adaptor adaptor(operands);
+    rewriter.replaceOp(op, adaptor.memref());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+class BufferizeTensorToMemrefOp : public OpConversionPattern<TensorToMemrefOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(TensorToMemrefOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    TensorToMemrefOp::Adaptor adaptor(operands);
+    rewriter.replaceOp(op, adaptor.tensor());
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateBufferizeEliminateMaterializationsPatternsAndLegality(
+    MLIRContext *context, BufferizeTypeConverter &typeConverter,
+    OwningRewritePatternList &patterns, ConversionTarget &target) {
+  patterns.insert<BufferizeTensorLoadOp, BufferizeTensorToMemrefOp>(
+      typeConverter, context);
+  target.addIllegalOp<TensorLoadOp, TensorToMemrefOp>();
+}
+
+//===----------------------------------------------------------------------===//
 // BufferizeFuncOpConverter
 //===----------------------------------------------------------------------===//
 
 /// Performs the actual function signature rewriting step.
 LogicalResult BufferizeFuncOpConverter::matchAndRewrite(
     mlir::FuncOp funcOp, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   auto funcType = funcOp.getType();

mlir/test/lib/Transforms/TestBufferPlacement.cpp

 //===- TestBufferPlacement.cpp - Test for buffer placement ------*- C++ -*-===//
 //
 // 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 logic for testing buffer placement including its
 // utility converters.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Bufferize.h"
 
 using namespace mlir;
 
 namespace {
 /// This pass tests the computeAllocPosition helper method and bufferize
 /// operation converters. Furthermore, this pass converts linalg operations on
 /// tensors to linalg operations on buffers to prepare them for the
 /// BufferPlacement pass that can be applied afterwards.
 /// `allowMemrefFunctionResults` informs the buffer placement to allow functions
 /// that have memref typed results. Buffer assignment operation converters will
 /// be adapted respectively. It will also allow memref typed results to escape
 /// from the deallocation.
 template <bool allowMemrefFunctionResults>
 struct TestBufferPlacementPreparationPass
     : mlir::PassWrapper<
           TestBufferPlacementPreparationPass<allowMemrefFunctionResults>,
           OperationPass<ModuleOp>> {
 
-  /// Converts tensor-type generic linalg operations to memref ones using
-  /// bufferize.
-  /// TODO: Avoid the copy-pasta by exposing the pattern from BufferPlacement.h
-  /// This is limited by not wanting BufferPlacement to depend on Linalg. Fixing
-  /// this probably requires an OpConversionPattern over generic Operation*. For
-  /// now only RewritePattern but not ConversionPattern allow this.
-
-  class GenericOpConverter
-      : public BufferizeOpConversionPattern<linalg::GenericOp> {
-  public:
-    using BufferizeOpConversionPattern<
-        linalg::GenericOp>::BufferizeOpConversionPattern;
-
-    LogicalResult
-    matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
-                    ConversionPatternRewriter &rewriter) const final {
-      linalg::GenericOpAdaptor adaptor(operands,
-                                       op.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(),
-                              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);
-        }
-      }
-
-      // Generate a new linalg operation that works on buffers.
-      auto linalgOp = rewriter.create<linalg::GenericOp>(
-          loc,
-          /*resultTensorTypes=*/ArrayRef<Type>{},
-          /*inputs=*/adaptor.inputs(),
-          /*outputBuffers=*/newOutputBuffers,
-          /*initTensors=*/ValueRange{}, op.indexing_maps(), op.iterator_types(),
-          op.docAttr(), op.library_callAttr(), op.symbol_sourceAttr());
-
-      // Create a new block in the region of the new Generic Op.
-      Block &oldBlock = op.getRegion().front();
-      Region &newRegion = linalgOp.region();
-      Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
-                                             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()))
-        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));
-
-      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(op, newOutputBuffers);
-      return success();
-    }
-  };
-
-  void populateTensorLinalgToBufferLinalgConversionPattern(
-      MLIRContext *context, BufferizeTypeConverter &converter,
-      OwningRewritePatternList &patterns) {
-    populateWithBufferizeOpConversionPatterns<mlir::ReturnOp, mlir::ReturnOp,
-                                              linalg::CopyOp>(
-        context, converter, patterns);
-    patterns.insert<GenericOpConverter>(context, converter);
-  }
-
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<TestDialect>();
     registry.insert<linalg::LinalgDialect>();
   }
 
   void runOnOperation() override {
     MLIRContext &context = this->getContext();
     ConversionTarget target(context);
converter.addDecomposeValueConversion([](OpBuilder &builder, Location loc,
         Value res = builder.create<GetTupleElementOp>(
             loc, resultType.getType(i), value, builder.getI32IntegerAttr(i));
         values.push_back(res);
       }
       return success();
     });
 
     OwningRewritePatternList patterns;
-    populateTensorLinalgToBufferLinalgConversionPattern(&context, converter,
-                                                        patterns);
+    populateWithBufferizeOpConversionPatterns<mlir::ReturnOp, mlir::ReturnOp,
+                                              linalg::CopyOp>(
+        &context, converter, patterns);
+    linalg::populateLinalgBufferizePatterns(&context, converter, patterns);
+    populateBufferizeEliminateMaterializationsPatternsAndLegality(
+        &context, converter, patterns, target);
     if (failed(applyFullConversion(this->getOperation(), target, patterns)))
       this->signalPassFailure();
   };
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestBufferPlacementPreparationPass() {