[mlir][linalg] Move generalization pattern to Transforms (NFC).

[gysit]

Summary

Move the generalization pattern to the other Linalg transforms to make it available to the codegen strategy.

[gysit]

Fix comment.

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

 ///
 /// E.g. the permutation `(i,j,k) -> (j,k,i)` is expressed with
 /// `interchangeVector = [1,2,0]`. All values in `interchangeVector` must be
 /// integers, in the range 0..`op.rank` without duplications
 /// (i.e. `[1,1,2]` is an invalid permutation).
 void interchangeGenericOp(PatternRewriter &rewriter, GenericOp genericOp,
                           ArrayRef<unsigned> interchangeVector);
 
+/// Creates a GenericOp from the given named operation `namedOp`. Assumes
+/// `namedOp` is not a GenericOp and has a region builder.
+GenericOp generalizeNamedOp(PatternRewriter &rewriter, LinalgOp namedOp);
+
 /// Callback function type used to perform the allocation for the promoted
 /// `subView`. In `boundingSubViewsize` a best attempt is made to find the
 /// smallest constant value for the size of the buffer needed for each
 /// dimension. If that is not possible, contains the dynamic size of the
 /// subview. The call back should return the buffer to use.
 using AllocBufferCallbackFn = std::function<Optional<Value>(
     OpBuilder &b, memref::SubViewOp subView,
     ArrayRef<Value> boundingSubViewSize, DataLayout &layout)>;

 // be applied as a rewrite pattern.
 //===----------------------------------------------------------------------===//
 /// Emits a `generic` operation with the `indexing_maps` and `iterator_types`
 /// permutated according to `permutation`.
 LogicalResult
 interchangeGenericOpPrecondition(GenericOp genericOp,
                                  ArrayRef<unsigned> interchangeVector);
 
+/// Generalize named operations to generic operations.
+LogicalResult generalizeNamedOpPrecondition(Operation *op);
+
 /// Promote std.subviews feeding linalg operations.
 LogicalResult promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options);
 
 /// Rewrite a linalg.generic into a suitable vector.contraction op.
 LogicalResult vectorizeLinalgOpPrecondition(Operation *op);
 
 //===----------------------------------------------------------------------===//

 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// The interchange vector to reorder the iterators and indexing_maps dims.
   SmallVector<unsigned, 8> interchangeVector;
 };
 
 ///
+/// Linalg generalization pattern.
+///
+/// Apply the `generalization` transformation as a pattern.
+/// `filter` controls LinalgTransformMarker matching and update when specified.
+/// See `generalization` for more details.
+struct LinalgGeneralizationPattern : public RewritePattern {
+  // Entry point to match any LinalgOp OpInterface.
+  LinalgGeneralizationPattern(
+      MLIRContext *context,
+      LinalgTransformationFilter filter = LinalgTransformationFilter(),
+      PatternBenefit benefit = 1);
+  // Entry point to match a specific Linalg op.
+  LinalgGeneralizationPattern(
+      StringRef opName, MLIRContext *context,
+      LinalgTransformationFilter filter = LinalgTransformationFilter(),
+      PatternBenefit benefit = 1);
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override;
+
+private:
+  /// LinalgTransformMarker handles special attribute manipulations.
+  LinalgTransformationFilter filter;
+};
+
+///
 /// Linalg promotion patterns.
 ///
 /// Apply the `promoteSubViews` transformation as a pattern.
 /// `filter` controls LinalgTransformMarker matching and update when specified.
 /// See `promoteSubViews` for more details.
 struct LinalgBasePromotionPattern : public RewritePattern {
   /// Entry point to match any LinalgOp OpInterface.
   /// MatchAnyOpTag-based constructor with a mandatory `filter`.

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

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "linalg-generalization"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-// Creates a linalg.generic op from the given `namedOp`. Returns a null op if
-// the given `namedOp` does not have a region builder.
-static GenericOp createGenericOpFromNamedOp(LinalgOp namedOp,
-                                            PatternRewriter &rewriter) {
+LogicalResult mlir::linalg::generalizeNamedOpPrecondition(Operation *op) {
+  LinalgOp namedOp = dyn_cast<LinalgOp>(op);
+  // Check if the operation is a LinalgOp but not a GenericOp.
+  if (!namedOp || isa<GenericOp>(op))
+    return failure();
+  // Check if the operation has a region builder.
+  if (!namedOp.getRegionBuilder())
+    return failure();
+  return success();
+}
+
+GenericOp mlir::linalg::generalizeNamedOp(PatternRewriter &rewriter,
+                                          LinalgOp namedOp) {
   SmallVector<Value> inputOperands = namedOp.getInputOperands();
   SmallVector<Value> outputOperands = namedOp.getOutputOperands();
   SmallVector<AffineMap> indexingMaps = namedOp.getIndexingMaps();
   SmallVector<StringRef> iterators = llvm::to_vector<4>(
       namedOp.iterator_types().getAsValueRange<StringAttr>());
   SmallVector<RankedTensorType> resultTypes = namedOp.getOutputTensorTypes();
   SmallVector<Type> types(resultTypes.begin(), resultTypes.end());
 
   // Inline the existing region if the named operation has a region attached.
   if (namedOp->getNumRegions() == 1) {
     GenericOp genericOp =
         rewriter.create<GenericOp>(namedOp.getLoc(), types, inputOperands,
                                    outputOperands, indexingMaps, iterators);
     rewriter.inlineRegionBefore(namedOp->getRegion(0), genericOp.region(),
                                 genericOp.region().begin());
     return genericOp;
   }
 
   // Otherwise use the region builder to generate a new region.
   // TODO: Remove this path once all linag operations have a region attached.
   auto regionBuilder = namedOp.getRegionBuilder();
-  if (!regionBuilder) {
-    LLVM_DEBUG(llvm::dbgs() << "no region builder for op: " << namedOp << "\n");
-    return nullptr;
-  }
+  assert(regionBuilder && "expect the operation to have region builder");
   return rewriter.create<GenericOp>(
       namedOp.getLoc(), types, inputOperands, outputOperands, indexingMaps,
       iterators,
       [&regionBuilder](OpBuilder &bodyBuilder, Location loc, ValueRange) {
         ImplicitLocOpBuilder b(loc, bodyBuilder);
         regionBuilder(b, *bodyBuilder.getBlock());
       });
 }

 
 struct GeneralizeConvOp
     : public LinalgGeneralizationPattern<GeneralizeConvOp, ConvOp> {
   using LinalgGeneralizationPattern::LinalgGeneralizationPattern;
 
   GenericOp createGenericOp(ConvOp convOp, OpBuilder &builder) const;
 };
 
-/// Catch-all pattern for converting all named ops with a region builder into
-/// linalg.generic.
-struct LinalgNamedOpGeneralizationPattern : RewritePattern {
-  LinalgNamedOpGeneralizationPattern(MLIRContext *context,
-                                     LinalgTransformationFilter marker,
-                                     PatternBenefit benefit = 1)
-      : RewritePattern(MatchAnyOpTypeTag(), benefit, context),
-        marker(std::move(marker)) {}
-
-  LogicalResult matchAndRewrite(Operation *rootOp,
-                                PatternRewriter &rewriter) const override {
-    auto linalgOp = dyn_cast<LinalgOp>(rootOp);
-    if (!linalgOp)
-      return failure();
-    if (failed(marker.checkAndNotify(rewriter, linalgOp)))
-      return failure();
-
-    // No nothing to do for linalg.generic.
-    if (isa<GenericOp>(rootOp))
-      return failure();
-
-    GenericOp genericOp = createGenericOpFromNamedOp(linalgOp, rewriter);
-    if (!genericOp)
-      return failure();
-
-    rewriter.replaceOp(rootOp, genericOp.getResults());
-    marker.replaceLinalgTransformationFilter(rewriter,
-                                             genericOp.getOperation());
-    return success();
-  }
-
-private:
-  LinalgTransformationFilter marker;
-};
-
 struct LinalgGeneralizationPass
     : public LinalgGeneralizationBase<LinalgGeneralizationPass> {
   void runOnFunction() override;
 };
 
 } // namespace
 
 void LinalgGeneralizationPass::runOnFunction() {

 
 void mlir::linalg::populateLinalgConvGeneralizationPatterns(
     RewritePatternSet &patterns, LinalgTransformationFilter marker) {
   patterns.add<GeneralizeConvOp>(patterns.getContext(), marker);
 }
 
 void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(
     RewritePatternSet &patterns, LinalgTransformationFilter marker) {
-  patterns.add<LinalgNamedOpGeneralizationPattern>(patterns.getContext(),
-                                                   marker);
+  patterns.add<LinalgGeneralizationPattern>(patterns.getContext(), marker);
 }
 
 std::unique_ptr<OperationPass<FuncOp>> mlir::createLinalgGeneralizationPass() {
   return std::make_unique<LinalgGeneralizationPass>();
 }

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

LogicalResult mlir::linalg::GenericOpInterchangePattern::matchAndRewrite(
   rewriter.updateRootInPlace(genericOp, [&]() {
     interchangeGenericOp(rewriter, genericOp, interchangeVector);
     // New filter if specified.
     filter.replaceLinalgTransformationFilter(rewriter, genericOp);
   });
   return success();
 }
 
+/// Linalg generalization pattern.
+mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
+    MLIRContext *context, LinalgTransformationFilter filter,
+    PatternBenefit benefit)
+    : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter) {}
+
+mlir::linalg::LinalgGeneralizationPattern::LinalgGeneralizationPattern(
+    StringRef opName, MLIRContext *context, LinalgTransformationFilter filter,
+    PatternBenefit benefit)
+    : RewritePattern(opName, benefit, context, {}), filter(filter) {}
+
+LogicalResult mlir::linalg::LinalgGeneralizationPattern::matchAndRewrite(
+    Operation *op, PatternRewriter &rewriter) const {
+  if (failed(filter.checkAndNotify(rewriter, op)))
+    return failure();
+  if (failed(generalizeNamedOpPrecondition(op)))
+    return failure();
+
+  GenericOp genericOp = generalizeNamedOp(rewriter, op);
+  rewriter.replaceOp(op, genericOp.getResults());
+  filter.replaceLinalgTransformationFilter(rewriter, genericOp);
+  return success();
+}
+
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     MLIRContext *context, LinalgTransformationFilter filter,
     LinalgPromotionOptions options, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
 mlir::linalg::LinalgBasePromotionPattern::LinalgBasePromotionPattern(
     StringRef opName, MLIRContext *context, LinalgPromotionOptions options,