[mlir][Linalg] Introduce a helper function for staged pattern application

[nicolasvasilache]

Summary

This revision introduces a helper function to allow applying rewrite patterns, interleaved with more global transformations, in a staged fashion:

1. the first stage consists of an OwningRewritePatternList. The RewritePattern in this list are applied once and in order.
2. the second stage consists of a single OwningRewritePattern that is applied greedily until convergence.
3. the third stage consists of applying a lambda, generally used for non-local transformation effects.

This allows creating custom fused transformations where patterns can be ordered and applied at a finer granularity than a sequence of traditional compiler passes.

A test that exercises these behaviors is added.

[nicolasvasilache]

Add missing comment.

[nicolasvasilache]

Rebase

[rriddle]

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

386

Only use std::function if you need to store it, otherwise use function_ref.

[ftynse]

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

375

it's a list of OwningRewritePatternLists

376

this does not match the code

377

OwningRewritePatternList

384

ArrayRef

385

const &

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

206

const &

208

May be you want to report which stage

208

Don't we want to bail out if we could not converge in one of the applications?

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

176

Could we just add a constructor template to OwningRewritePatternList that forwards its arguments to insert? This would remove the need to std::move and call insert here, and looks generally useful for lists that need only one call to insert.

[nicolasvasilache]

Address review.

[nicolasvasilache]

Address review.

[ftynse]

mlir/include/mlir/IR/PatternMatch.h

394

Doc doesn't match the code

397

This looks like it's calling the move constructor of T (pattern class), not constructing the pattern in-place from the arguments. I was thinking about literally forwarding to insert, something like

template <typename... Ts, typename ConstructorArg,
            typename... ConstructorArgs,
            typename = std::enable_if_t<sizeof...(Ts) != 0>>
OwningRewritePatternList(ConstructorArg &&arg, ConstructorArgs &&... args) {
  insert<Ts...>(std::forward<ConstructorArg>(arg), std::forward<ConstructorArgs>(args)...);
}

but I guess the pattern-move constructor works too, if you fix the documentation.

[nicolasvasilache]

mlir/include/mlir/IR/PatternMatch.h

397

Yeah, didn't manage to get the combination of emplace_back + templated ctor to work so I went with the move constructor. Feel free to give it a shot if you find the templatefu better (I don't necessarily).

[nicolasvasilache]

Disable test temporarily.

[nicolasvasilache]

Address review,

[rriddle]

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

386

nit: Drop the llvm::

mlir/include/mlir/IR/PatternMatch.h

397

You aren't properly forwarding t here, you need std::forward<T>(t)

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

205

nit: Drop the llvm::

208

Why aren't these llvm::dbgs calls wrapped by LLVM_DEBUG?

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

 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgMarker marker;
   /// Controls whether the pattern lowers to library calls, loop.for, affine.for
   /// or loop.parallel.
   LinalgLoweringType loweringType;
 };
 
+//===----------------------------------------------------------------------===//
+// 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 OwningRewritePatternList. Each
+///   OwningRewritePatternList in this list is applied once, in order.
+///   2. the second stage consists of a single OwningRewritePattern that is
+///   applied greedily until convergence.
+///   3. the third stage consists of applying a lambda, generally used for
+///   non-local transformation effects. This allows creating custom fused
+///   transformations where patterns can be ordered and applied at a finer
+///   granularity than a sequence of traditional compiler passes.
+LogicalResult applyStagedPatterns(
+    Operation *op, ArrayRef<OwningRewritePatternList> stage1Patterns,
+    const OwningRewritePatternList &stage2Patterns,
+    llvm::function_ref<LogicalResult(Operation *)> stage3Lambda = nullptr);
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_

mlir/include/mlir/IR/PatternMatch.h

 //===----------------------------------------------------------------------===//
 // Pattern-driven rewriters
 //===----------------------------------------------------------------------===//
 
 class OwningRewritePatternList {
   using PatternListT = std::vector<std::unique_ptr<RewritePattern>>;
 
 public:
+  OwningRewritePatternList() = default;
+
+  /// Construct a OwningRewritePatternList populated with one pattern of type
+  /// 'T' constructed from the given arguments
+  template <typename T>
+  OwningRewritePatternList(T &&t) {
+    patterns.emplace_back(std::make_unique<T>(t));
+  }
+
   PatternListT::iterator begin() { return patterns.begin(); }
   PatternListT::iterator end() { return patterns.end(); }
   PatternListT::const_iterator begin() const { return patterns.begin(); }
   PatternListT::const_iterator end() const { return patterns.end(); }
   void clear() { patterns.clear(); }
 
   //===--------------------------------------------------------------------===//
   // Pattern Insertion
   //===--------------------------------------------------------------------===//
 
   /// Add an instance of each of the pattern types 'Ts' to the pattern list with
-  /// the given arguments.
+  /// the given arguments. Return a reference to `this` for chaining insertions.
   /// Note: ConstructorArg is necessary here to separate the two variadic lists.
   template <typename... Ts, typename ConstructorArg,
             typename... ConstructorArgs,
             typename = std::enable_if_t<sizeof...(Ts) != 0>>
-  void insert(ConstructorArg &&arg, ConstructorArgs &&... args) {
+  OwningRewritePatternList &insert(ConstructorArg &&arg,
+                                   ConstructorArgs &&... args) {
     // The following expands a call to emplace_back for each of the pattern
     // types 'Ts'. This magic is necessary due to a limitation in the places
     // that a parameter pack can be expanded in c++11.
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
     using dummy = int[];
     (void)dummy{
         0, (patterns.emplace_back(std::make_unique<Ts>(arg, args...)), 0)...};
+    return *this;
   }
 
 private:
   PatternListT patterns;
 };
 
 /// This class manages optimization and execution of a group of rewrite
 /// patterns, providing an API for finding and applying, the best match against

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

LogicalResult mlir::linalg::LinalgBaseVectorizationPattern::matchAndRewrite(
   if (failed(marker.checkAndNotify(rewriter, linalgOp)))
     return failure();
   if (failed(vectorizeLinalgOpPrecondition(op)))
     return failure();
   vectorizeLinalgOp(rewriter, op);
   rewriter.eraseOp(op);
   return success();
 }
+
+LogicalResult mlir::linalg::applyStagedPatterns(
+    Operation *op, ArrayRef<OwningRewritePatternList> stage1Patterns,
+    const OwningRewritePatternList &stage2Patterns,
+    llvm::function_ref<LogicalResult(Operation *)> stage3Lambda) {
+  for (const auto &patterns : stage1Patterns) {
+    if (!applyPatternsAndFoldGreedily(op, patterns)) {
+      llvm::dbgs() << "Underlying first stage rewrite did not converge";
+      return failure();
+    }
+    if (!applyPatternsAndFoldGreedily(op, stage2Patterns)) {
+      llvm::dbgs() << "Underlying second stage rewrite did not converge";
+      return failure();
+    }
+    if (stage3Lambda) {
+      if (failed(stage3Lambda(op)))
+        return failure();
+    }
+  }
+  return success();
+}

mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
+
+func @matmul(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+                  %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+                  %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
+  linalg.matmul(%A, %B, %C) {__internal_linalg_transform__ = "START"} :
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
+    memref<1584x1584xf32, offset: 0, strides: [1584, 1]>
+  return
+}
+
+// CHECK-LABEL:func @matmul
+//      CHECK: vector.broadcast {{.*}} : f32 to vector<8x16xf32>
+//      CHECK: store {{.*}}[] : memref<vector<8x16xf32>>
+//
+//      CHECK: vector.broadcast {{.*}} : f32 to vector<16x12xf32>
+//      CHECK: store {{.*}}[] : memref<vector<16x12xf32>>
+//
+//      CHECK: vector.broadcast {{.*}} : f32 to vector<8x12xf32>
+//      CHECK: store {{.*}}[] : memref<vector<8x12xf32>>
+//
+//      CHECK: linalg.copy
+//      CHECK: linalg.copy
+//      CHECK: linalg.copy
+//
+//      CHECK: vector.contract
+// CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME:   : vector<8x16xf32>, vector<16x12xf32> into vector<8x12xf32>
+//
+//      CHECK: linalg.copy

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

struct TestLinalgTransforms
   TestLinalgTransforms() = default;
   TestLinalgTransforms(const TestLinalgTransforms &pass) {}
 
   void runOnFunction() override;
 
   Option<bool> testPatterns{*this, "test-patterns",
                             llvm::cl::desc("Test a mixed set of patterns"),
                             llvm::cl::init(false)};
+  Option<bool> testMatmulToVectorPatterns1dTiling{
+      *this, "test-matmul-to-vector-patterns-tile-1d",
+      llvm::cl::desc(
+          "Test a fused pass that applies patterns from matmul to vectors via "
+          "1-d tiling"),
+      llvm::cl::init(false)};
+  Option<bool> testMatmulToVectorPatterns2dTiling{
+      *this, "test-matmul-to-vector-patterns-tile-2d",
+      llvm::cl::desc(
+          "Test a fused pass that applies patterns from matmul to vectors via "
+          "2-d tiling"),
+      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//
static void applyPatterns(FuncOp funcOp) {
   applyPatternsAndFoldGreedily(funcOp, patterns);
 
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }
 
+OwningRewritePatternList
+getMatmulToVectorCanonicalizationPatterns(MLIRContext *context) {
+  OwningRewritePatternList patterns;
+  AffineApplyOp::getCanonicalizationPatterns(patterns, context);
+  AffineMinOp::getCanonicalizationPatterns(patterns, context);
+  AffineMaxOp::getCanonicalizationPatterns(patterns, context);
+  AllocOp::getCanonicalizationPatterns(patterns, context);
+  SubViewOp::getCanonicalizationPatterns(patterns, context);
+  ViewOp::getCanonicalizationPatterns(patterns, context);
+  MatmulOp::getCanonicalizationPatterns(patterns, context);
+  return patterns;
+}
+
+void fillL1TilingAndMatmulToVectorPatterns(
+    MLIRContext *context, StringRef startMarker,
+    SmallVectorImpl<OwningRewritePatternList> &patternsVector) {
+  patternsVector.emplace_back(LinalgTilingPattern<MatmulOp>(
+      context,
+      LinalgTilingOptions().setTileSizes({8, 12, 16}).setInterchange({1, 0, 2}),
+      LinalgMarker({startMarker}, "L1")));
+
+  patternsVector.emplace_back(LinalgPromotionPattern<MatmulOp>(
+      context, LinalgPromotionOptions(), LinalgMarker({"L1"}, "VEC")));
+
+  patternsVector.emplace_back(
+      LinalgVectorizationPattern<MatmulOp>(context, LinalgMarker({"VEC"})));
+  patternsVector.back()
+      .insert<LinalgVectorizationPattern<FillOp>,
+              LinalgVectorizationPattern<CopyOp>>(context);
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
-  if (testPatterns)
-    return applyPatterns(getFunction());
+  if (testPatterns) {
+    applyPatterns(getFunction());
+  } else {
+    SmallVector<OwningRewritePatternList, 4> stage1Patterns;
+    if (testMatmulToVectorPatterns1dTiling) {
+      fillL1TilingAndMatmulToVectorPatterns(&getContext(), "START",
+                                            stage1Patterns);
+    } else if (testMatmulToVectorPatterns2dTiling) {
+      stage1Patterns.emplace_back(
+          LinalgTilingPattern<MatmulOp>(&getContext(),
+                                        LinalgTilingOptions()
+                                            .setTileSizes({768, 264, 768})
+                                            .setInterchange({1, 2, 0}),
+                                        LinalgMarker({"START"}, "L2")));
+      fillL1TilingAndMatmulToVectorPatterns(&getContext(), "L2",
+                                            stage1Patterns);
+    }
+    OwningRewritePatternList stage2Patterns =
+        getMatmulToVectorCanonicalizationPatterns(&getContext());
+    applyStagedPatterns(getFunction(), stage1Patterns, stage2Patterns);
+  }
+
+  // Drop the marker.
+  getFunction().walk([](LinalgOp op) {
+    op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
+  });
 }
 
 namespace mlir {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace mlir