[mlir] Move some linalg patterns around.

[silvas]

Summary

The bufferization patterns are moved to the .cpp file, which is
preferred in the codebase when it makes sense.

The LinalgToStandard patterns are kept a header because they are
expected to be used individually. However, they are moved to
LinalgToStandard.h which is the file corresponding to where they are
defined.

This also removes TensorCastOpConverter, which is handled by
populateStdBufferizePatterns now. Eventually, the constant op lowering
will be handled as well, but it there are currently holdups on moving
it (see https://reviews.llvm.org/D89916).

[nicolasvasilache]

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

765

I wouldn't hide these, the finer granularity of control that individual patterns bring is useful for this type of rewrite.
Pass granularity is quite inconvenient for this.
I do not plan concrete uses of this in the short term but I expect I'd have to revert this in the future ..

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

197

These ones make sense to hide at the pass granularity.

[silvas]

Address comments.

[silvas]

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

765

Makes sense! I moved them to LinalgToStandard.h instead (which matches with the .cpp file that they seem to be defined in anyway).

[nicolasvasilache]

Thanks @silvas !

mlir/include/mlir/Conversion/LinalgToStandard/LinalgToStandard.h

 //===- LinalgToStandard.h - Utils to convert from the linalg dialect ------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_LINALGTOSTANDARD_LINALGTOSTANDARD_H_
 #define MLIR_CONVERSION_LINALGTOSTANDARD_LINALGTOSTANDARD_H_
 
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 class ModuleOp;
 template <typename T>
 class OperationPass;
 
+namespace linalg {
+
+//===----------------------------------------------------------------------===//
+// Patterns to convert a LinalgOp to std.call @external library implementation.
+//===----------------------------------------------------------------------===//
+// These patterns are exposed individually because they are expected to be
+// typically used individually.
+
+// Create a new call to the type-canonicalized `LinalgOp::getLibraryCallName()`
+// function. The implementation of the function can be either in the same module
+// or in an externally linked library.
+// This is a generic entry point for all LinalgOp, except for CopyOp and
+// IndexedGenericOp, for which omre specialized patterns are provided.
+class LinalgOpToLibraryCallRewrite : public RewritePattern {
+public:
+  LinalgOpToLibraryCallRewrite()
+      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Rewrite pattern specialization for CopyOp, kicks in when both input and
+/// output permutations are left unspecified or are the identity.
+class CopyOpToLibraryCallRewrite : public OpRewritePattern<CopyOp> {
+public:
+  using OpRewritePattern<CopyOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(CopyOp op,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Rewrite CopyOp with permutations into a sequence of TransposeOp and
+/// permutation-free CopyOp. This interplays with TransposeOpConversion and
+/// LinalgConversion<CopyOp> to create a path to the LLVM dialect.
+class CopyTransposeRewrite : public OpRewritePattern<CopyOp> {
+public:
+  using OpRewritePattern<CopyOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(CopyOp op,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Conversion pattern specialization for IndexedGenericOp, has special handling
+/// for the extra index operands.
+class IndexedGenericOpToLibraryCallRewrite
+    : public OpRewritePattern<IndexedGenericOp> {
+public:
+  using OpRewritePattern<IndexedGenericOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(IndexedGenericOp op,
+                                PatternRewriter &rewriter) const override;
+};
+
+/// Populate the given list with patterns that convert from Linalg to Standard.
+void populateLinalgToStandardConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
+} // namespace linalg
+
 /// Create a pass to convert Linalg operations to the Standard dialect.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertLinalgToStandardPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_LINALGTOSTANDARD_LINALGTOSTANDARD_H_

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

ConvOpVectorization(MLIRContext *context, SmallVector<bool, 4> msk)
     this->mask = msk;
   }
 
   LogicalResult matchAndRewrite(ConvOp minOp,
                                 PatternRewriter &rewriter) const override;
 };
 
 //===----------------------------------------------------------------------===//
-// Patterns to convert a LinalgOp to std.call @external library implementation.
-//===----------------------------------------------------------------------===//
-// Create a new call to the type-canonicalized `LinalgOp::getLibraryCallName()`
-// function. The implementation of the function can be either in the same module
-// or in an externally linked library.
-// This is a generic entry point for all LinalgOp, except for CopyOp and
-// IndexedGenericOp, for which omre specialized patterns are provided.
-class LinalgOpToLibraryCallRewrite : public RewritePattern {
-public:
-  LinalgOpToLibraryCallRewrite()
-      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
-
-  LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override;
-};
-
-/// Rewrite pattern specialization for CopyOp, kicks in when both input and
-/// output permutations are left unspecified or are the identity.
-class CopyOpToLibraryCallRewrite : public OpRewritePattern<CopyOp> {
-public:
-  using OpRewritePattern<CopyOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(CopyOp op,
-                                PatternRewriter &rewriter) const override;
-};
-
-/// Rewrite CopyOp with permutations into a sequence of TransposeOp and
-/// permutation-free CopyOp. This interplays with TransposeOpConversion and
-/// LinalgConversion<CopyOp> to create a path to the LLVM dialect.
-class CopyTransposeRewrite : public OpRewritePattern<CopyOp> {
-public:
-  using OpRewritePattern<CopyOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(CopyOp op,
-                                PatternRewriter &rewriter) const override;
-};
-
-/// Conversion pattern specialization for IndexedGenericOp, has special handling
-/// for the extra index operands.
-class IndexedGenericOpToLibraryCallRewrite
-    : public OpRewritePattern<IndexedGenericOp> {
-public:
-  using OpRewritePattern<IndexedGenericOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(IndexedGenericOp op,
-                                PatternRewriter &rewriter) const override;
-};
-
-/// Populate the given list with patterns that convert from Linalg to Standard.
-void populateLinalgToStandardConversionPatterns(
-    OwningRewritePatternList &patterns, MLIRContext *ctx);
-
-//===----------------------------------------------------------------------===//
-// Buffer allocation patterns.
-//===----------------------------------------------------------------------===//
-
-/// Generic BufferizeConversionPattern that matches any Operation* and
-/// dispatches internally. This avoids template instantiating one pattern for
-/// each LinalgOp op.
-class LinalgOpConverter : public BufferizeConversionPattern {
-public:
-  LinalgOpConverter(MLIRContext *context, BufferizeTypeConverter &converter)
-      : BufferizeConversionPattern(context, converter) {}
-
-  LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final;
-};
-
-/// TensorConstantOp conversion inserts a linearized 1-D vector constant that is
-/// stored in memory. A linalg.reshape is introduced to convert to the desired
-/// n-D buffer form.
-class TensorConstantOpConverter
-    : public BufferizeOpConversionPattern<ConstantOp> {
-public:
-  using BufferizeOpConversionPattern<ConstantOp>::BufferizeOpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final;
-};
-
-/// TensorCastOp converts 1-1 to MemRefCastOp.
-class TensorCastOpConverter
-    : public BufferizeOpConversionPattern<TensorCastOp> {
-public:
-  using BufferizeOpConversionPattern<
-      TensorCastOp>::BufferizeOpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(TensorCastOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const final;
-};
-
-//===----------------------------------------------------------------------===//
 // Support for staged pattern application.
 //===----------------------------------------------------------------------===//
 /// Helper function to allow applying rewrite patterns, interleaved with more
 /// global transformations, in a staged fashion:
 ///   1. the first stage consists of a list of FrozenRewritePatternList. Each
 ///   FrozenRewritePatternList in this list is applied once, in order.
 ///   2. the second stage consists of a single OwningRewritePattern that is
 ///   applied greedily until convergence.

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

 //===- Bufferize.cpp - Bufferization of linalg ops ------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/Bufferize.h"
 #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/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace ::mlir;
 using namespace ::mlir::linalg;
 
LinalgOp res = cast<LinalgOp>(linalgOp.clone(rewriter, linalgOp.getLoc(),
                                                newOperands));
   // Need to mutate the operands_segment_sizes in the resulting op.
   res.setNumOutputBuffers(outputs.size());
   res.setNumInitTensors(0);
   // Replace the results of the old op with the new output buffers.
   rewriter.replaceOp(linalgOp, outputs);
 }
 
-LogicalResult mlir::linalg::LinalgOpConverter::matchAndRewrite(
-    Operation *op, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
+//===----------------------------------------------------------------------===//
+// Buffer allocation patterns.
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Generic BufferizeConversionPattern that matches any Operation* and
+/// dispatches internally. This avoids template instantiating one pattern for
+/// each LinalgOp op.
+class LinalgOpConverter : public BufferizeConversionPattern {
+public:
+  LinalgOpConverter(MLIRContext *context, BufferizeTypeConverter &converter)
+      : BufferizeConversionPattern(context, converter) {}
+
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+
     LinalgOp linalgOp = dyn_cast<linalg::LinalgOp>(op);
     if (!linalgOp)
       return failure();
 
     // We abuse the GenericOpAdaptor here.
     // TODO: Manually create an Adaptor that captures inputs, output_buffers and
     // init_tensors for all linalg::LinalgOp interface ops.
     linalg::GenericOpAdaptor adaptor(operands, op->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 = linalgOp.getLoc();
     SmallVector<Value, 2> newOutputBuffers(adaptor.output_buffers().begin(),
                                            adaptor.output_buffers().end());
 
-  if (failed(allocateBuffersForResults(loc, linalgOp, adaptor, newOutputBuffers,
-                                       rewriter))) {
-    linalgOp.emitOpError() << "Failed to allocate buffers for tensor results.";
+    if (failed(allocateBuffersForResults(loc, linalgOp, adaptor,
+                                         newOutputBuffers, rewriter))) {
+      linalgOp.emitOpError()
+          << "Failed to allocate buffers for tensor results.";
       return failure();
     }
 
     // Delegate to the linalg generic pattern.
     if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
       finalizeBufferAllocation(rewriter, genericOp, adaptor.inputs(),
                                newOutputBuffers);
       return success();
     }
 
     finalizeBufferAllocation(rewriter, linalgOp, adaptor.inputs(),
                              newOutputBuffers);
     return success();
   }
+};
+} // namespace
 
-LogicalResult mlir::linalg::TensorConstantOpConverter::matchAndRewrite(
-    ConstantOp op, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  RankedTensorType rankedTensorType = op.getType().dyn_cast<RankedTensorType>();
+namespace {
+/// TensorConstantOp conversion inserts a linearized 1-D vector constant that is
+/// stored in memory. A linalg.reshape is introduced to convert to the desired
+/// n-D buffer form.
+class TensorConstantOpConverter
+    : public BufferizeOpConversionPattern<ConstantOp> {
+public:
+  using BufferizeOpConversionPattern<ConstantOp>::BufferizeOpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ConstantOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+
+    RankedTensorType rankedTensorType =
+        op.getType().dyn_cast<RankedTensorType>();
     if (!rankedTensorType)
       return failure();
     if (llvm::any_of(rankedTensorType.getShape(), [](int64_t s) {
           return s == 0 || ShapedType::isDynamic(s);
         }))
       return failure();
 
     int64_t nElements = 1;
     for (int64_t s : rankedTensorType.getShape())
       nElements *= s;
     Type elementType = rankedTensorType.getElementType();
     MemRefType memrefType =
         converter.convertType(op.getType()).cast<MemRefType>();
     VectorType flatVectorType = VectorType::get({nElements}, elementType);
     MemRefType memrefOfFlatVectorType = MemRefType::get({}, flatVectorType);
     MemRefType flatMemrefType = MemRefType::get({nElements}, elementType);
 
     Location loc = op.getLoc();
     auto attr = op.getValue().cast<DenseElementsAttr>();
     Value alloc =
         rewriter.create<AllocOp>(loc, memrefOfFlatVectorType, ValueRange{});
     Value cstVec = rewriter.create<ConstantOp>(loc, flatVectorType,
                                                attr.reshape(flatVectorType));
     rewriter.create<StoreOp>(loc, cstVec, alloc);
 
     Value memref =
         rewriter.create<vector::TypeCastOp>(loc, flatMemrefType, alloc);
     if (rankedTensorType.getRank() > 1) {
       // Introduce a linalg.reshape to flatten the memref.
       AffineMap collapseAllDims = AffineMap::getMultiDimIdentityMap(
           /*numDims=*/rankedTensorType.getRank(), op.getContext());
       memref = rewriter.create<linalg::ReshapeOp>(
           loc, memrefType, memref,
           rewriter.getAffineMapArrayAttr(collapseAllDims));
     }
     rewriter.replaceOp(op, memref);
 
     return success();
   }
-
-LogicalResult mlir::linalg::TensorCastOpConverter::matchAndRewrite(
-    TensorCastOp op, ArrayRef<Value> operands,
-    ConversionPatternRewriter &rewriter) const {
-  if (op.getType().hasRank())
-    return failure();
-  Type t = UnrankedMemRefType::get(op.getType().getElementType(),
-                                   /*memorySpace=*/0);
-  rewriter.replaceOpWithNewOp<MemRefCastOp>(op, t, operands.front());
-  return success();
-}
+};
+} // namespace
 
 namespace {
 
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
              converter.isLegal(&funcOp.getBody());
     });
 
     converter.setResultConversionKind<RankedTensorType, MemRefType>(
         BufferizeTypeConverter::AppendToArgumentsList);
 
     OwningRewritePatternList patterns;
     populateLinalgBufferizePatterns(&context, converter, patterns);
+    populateStdBufferizePatterns(&context, converter, patterns);
     populateWithBufferizeOpConversionPatterns<mlir::ReturnOp, mlir::ReturnOp,
                                               linalg::CopyOp>(
         &context, converter, patterns);
     if (failed(applyFullConversion(this->getOperation(), target,
                                    std::move(patterns))))
       this->signalPassFailure();
   }
 };
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createLinalgBufferizePass() {
   return std::make_unique<LinalgBufferizePass>();
 }
 void mlir::linalg::populateLinalgBufferizePatterns(
     MLIRContext *context, BufferizeTypeConverter &converter,
     OwningRewritePatternList &patterns) {
   patterns.insert<
       // clang-format off
       LinalgOpConverter,
-      TensorCastOpConverter,
       TensorConstantOpConverter
       // clang-format on
       >(context, converter);
 }

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRLinalgAnalysis
   MLIRLinalgEDSC
   MLIRLinalg
   MLIRLinalgUtils
   MLIRSCF
   MLIRSCFTransforms
   MLIRPass
   MLIRStandard
+  MLIRStandardOpsTransforms
   MLIRStandardToLLVM
   MLIRTransforms
   MLIRTransformUtils
   MLIRVector
   MLIRVectorToSCF
 )