[mlir] Enable perfect forwarding in RewritePatternSet::add

[laszlokindrat]

Summary

This patch modifies the implementation of RewritePatternSet::add to perfectly forward its arguments to pattern constructors. Without this, code like the following compiles but, due to the limited lifetime of the temporary TypeConverter, can produce unexpected behavior:

RewritePatternSet patterns(context);
patterns.add<SomeOpConversion, OtherOpConversion>(TypeConverter(), context);

if (failed(applyPartialConversion(getOperation(), target, std::move(patterns))))
  return signalPassFailure();

The patch also changes the linalg fusion pattern implementation to correctly fill the test pattern set given the new behavior.

Author: Laszlo Kindrat <[e-mail]>

[laszlokindrat]

Could someone please land this? (I can't commit to the repo.) Thanks!

mlir/include/mlir/IR/PatternMatch.h

template <typename... Ts, typename ConstructorArg,
             typename... ConstructorArgs,
             typename = std::enable_if_t<sizeof...(Ts) != 0>>
   RewritePatternSet &add(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'.
     (void)std::initializer_list<int>{
-        0, (addImpl<Ts>(/*debugLabels=*/llvm::None, arg, args...), 0)...};
+        0, (addImpl<Ts>(/*debugLabels=*/llvm::None,
+                        std::forward<ConstructorArg>(arg),
+                        std::forward<ConstructorArgs>(args)...),
+            0)...};
     return *this;
   }
   /// An overload of the above `add` method that allows for attaching a set
   /// of debug labels to the attached patterns. This is useful for labeling
   /// groups of patterns that may be shared between multiple different
   /// passes/users.
   template <typename... Ts, typename ConstructorArg,
             typename... ConstructorArgs,

mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp

 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
-template <LinalgTilingLoopType LoopType>
-static void fillFusionPatterns(MLIRContext *context,
-                               const LinalgDependenceGraph &dependenceGraph,
-                               RewritePatternSet &patterns) {
-  patterns.add<LinalgTileAndFusePattern<MatmulOp>,
-               LinalgTileAndFusePattern<Conv2DOp>>(
-      context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({2}),
-      LinalgTransformationFilter(
-          StringAttr::get(context, "basic_fusion"),
-          StringAttr::get(context, "after_basic_fusion")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_basic_fusion_producer")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_basic_fusion_original")));
-
-  patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
-      context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({0}),
-      LinalgTransformationFilter(StringAttr::get(context, "lhs_fusion"),
-                                 StringAttr::get(context, "after_lhs_fusion")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_lhs_fusion_producer")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_lhs_fusion_original")));
-
-  patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
-      context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({2}),
-      LinalgTransformationFilter(StringAttr::get(context, "out_fusion"),
-                                 StringAttr::get(context, "after_out_fusion")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_out_fusion_producer")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_out_fusion_original")));
-
-  patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
-      context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({1}),
-      LinalgTransformationFilter(StringAttr::get(context, "rhs_fusion"),
-                                 StringAttr::get(context, "after_rhs_fusion")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_rhs_fusion_producer")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_rhs_fusion_original")));
-
-  patterns.add<LinalgTileAndFusePattern<MatmulOp>>(
+/// Use this to safely fill patterns for this test, since RewritePatternSet::add
+/// forwards Rvalues only to the first pattern.
+template <typename OpTy, LinalgTilingLoopType LoopType>
+static void fillFusionPattern(MLIRContext *context,
+                              RewritePatternSet &patterns,
+                              const Twine &testCase,
+                              ArrayRef<int64_t> tileSizes,
+                              ArrayRef<int64_t> indicesToFuse) {
+  patterns.add<LinalgTileAndFusePattern<OpTy>>(
       context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64, 16}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({0, 2}),
+      LinalgTilingOptions().setTileSizes(tileSizes).setLoopType(LoopType),
+      LinalgFusionOptions().setIndicesToFuse(indicesToFuse),
       LinalgTransformationFilter(
-          StringAttr::get(context, "two_operand_fusion"),
-          StringAttr::get(context, "after_two_operand_fusion")),
+          StringAttr::get(context, testCase + "_fusion"),
+          StringAttr::get(context, "after_" + testCase + "_fusion")),
       LinalgTransformationFilter(
           ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_two_operand_fusion_producer")),
+          StringAttr::get(context, "after_" + testCase + "_fusion_producer")),
       LinalgTransformationFilter(
           ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_two_operand_fusion_original")));
+          StringAttr::get(context, "after_" + testCase + "_fusion_original")));
+}
 
-  patterns.add<LinalgTileAndFusePattern<GenericOp>>(
-      context, dependenceGraph,
-      LinalgTilingOptions().setTileSizes({32, 64}).setLoopType(LoopType),
-      LinalgFusionOptions().setIndicesToFuse({0, 1}),
-      LinalgTransformationFilter(
-          StringAttr::get(context, "transpose_fusion"),
-          StringAttr::get(context, "after_transpose_fusion")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_transpose_fusion_producer")),
-      LinalgTransformationFilter(
-          ArrayRef<StringAttr>(),
-          StringAttr::get(context, "after_transpose_fusion_original")));
+template <LinalgTilingLoopType LoopType>
+static void fillFusionPatterns(MLIRContext *context,
+                               const LinalgDependenceGraph &dependenceGraph,
+                               RewritePatternSet &patterns) {
+  fillFusionPattern<Conv2DOp, LoopType>(context, dependenceGraph, patterns,
+                                        /*testCase=*/"basic",
+                                        /*tileSizes=*/{32, 64, 16},
+                                        /*indicesToFuse=*/{2});
+
+  auto fillMatmulPattern = [&](const Twine &testCase,
+                               ArrayRef<int64_t> indicesToFuse) {
+    fillFusionPattern<MatmulOp, LoopType>(context, dependenceGraph, patterns,
+                                          testCase, /*tileSizes=*/{32, 64, 16},
+                                          indicesToFuse);
+  };
+  fillMatmulPattern(/*testCase=*/"basic",
+                    /*indicesToFuse=*/{2});
+  fillMatmulPattern(/*testCase=*/"lhs",
+                    /*indicesToFuse=*/{0});
+  fillMatmulPattern(/*testCase=*/"out",
+                    /*indicesToFuse=*/{2});
+  fillMatmulPattern(/*testCase=*/"rhs",
+                    /*indicesToFuse=*/{1});
+  fillMatmulPattern(/*testCase=*/"two_operand",
+                    /*indicesToFuse=*/{0, 2});
+
+  fillFusionPattern<GenericOp, LoopType>(context, dependenceGraph, patterns,
+                                         /*testCase=*/"transpose",
+                                         /*tileSizes=*/{32, 64},
+                                         /*indicesToFuse=*/{0, 1});
 }
 
 namespace {
 template <LinalgTilingLoopType LoopType>
 struct TestLinalgFusionTransforms
     : public PassWrapper<TestLinalgFusionTransforms<LoopType>,
                          OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgFusionTransforms)