[Mlir-commits] [mlir] [mlir] Avoid repeated moves in RewritePatternSet (PR #220871)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

Multi-pattern add APIs forwarded the same constructor arguments once for each pattern type. An rvalue was therefore moved from repeatedly, and every pattern after the first could receive moved-from state.

Pass constructor arguments as lvalues when constructing multiple patterns, while preserving perfect forwarding when constructing a single pattern. Cover add, addWithLabel, and insert in the existing PatternMatch unit test.

Found by Coverity.

Assisted-by: Codex

---
Full diff: https://github.com/llvm/llvm-project/pull/220871.diff


2 Files Affected:

- (modified) mlir/include/mlir/IR/PatternMatch.h (+22-15) 
- (modified) mlir/unittests/IR/PatternMatchTest.cpp (+46) 


``````````diff
diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
index 83477c79ff582..531051e42b824 100644
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -859,11 +859,8 @@ class RewritePatternSet {
             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'.
-    (addImpl<Ts>(/*debugLabels=*/{}, std::forward<ConstructorArg>(arg),
-                 std::forward<ConstructorArgs>(args)...),
-     ...);
+    addImplForEach<Ts...>(/*debugLabels=*/{}, std::forward<ConstructorArg>(arg),
+                          std::forward<ConstructorArgs>(args)...);
     return *this;
   }
   /// An overload of the above `add` method that allows for attaching a set
@@ -876,11 +873,8 @@ class RewritePatternSet {
   RewritePatternSet &addWithLabel(ArrayRef<StringRef> debugLabels,
                                   ConstructorArg &&arg,
                                   ConstructorArgs &&...args) {
-    // The following expands a call to emplace_back for each of the pattern
-    // types 'Ts'.
-    (addImpl<Ts>(debugLabels, std::forward<ConstructorArg>(arg),
-                 std::forward<ConstructorArgs>(args)...),
-     ...);
+    addImplForEach<Ts...>(debugLabels, std::forward<ConstructorArg>(arg),
+                          std::forward<ConstructorArgs>(args)...);
     return *this;
   }
 
@@ -944,11 +938,8 @@ class RewritePatternSet {
             typename... ConstructorArgs,
             typename = std::enable_if_t<sizeof...(Ts) != 0>>
   RewritePatternSet &insert(ConstructorArg &&arg, ConstructorArgs &&...args) {
-    // The following expands a call to emplace_back for each of the pattern
-    // types 'Ts'.
-    (addImpl<Ts>(/*debugLabels=*/{}, std::forward<ConstructorArg>(arg),
-                 std::forward<ConstructorArgs>(args)...),
-     ...);
+    addImplForEach<Ts...>(/*debugLabels=*/{}, std::forward<ConstructorArg>(arg),
+                          std::forward<ConstructorArgs>(args)...);
     return *this;
   }
 
@@ -998,6 +989,22 @@ class RewritePatternSet {
   }
 
 private:
+  /// Add an instance of each pattern type. A single pattern receives perfectly
+  /// forwarded arguments. Multiple patterns receive lvalues so that a shared
+  /// rvalue argument is copied instead of moved from repeatedly.
+  template <typename... Ts, typename ConstructorArg,
+            typename... ConstructorArgs>
+  void addImplForEach(ArrayRef<StringRef> debugLabels, ConstructorArg &&arg,
+                      ConstructorArgs &&...args) {
+    if constexpr (sizeof...(Ts) == 1) {
+      (addImpl<Ts>(debugLabels, std::forward<ConstructorArg>(arg),
+                   std::forward<ConstructorArgs>(args)...),
+       ...);
+    } else {
+      (addImpl<Ts>(debugLabels, arg, args...), ...);
+    }
+  }
+
   /// Add an instance of the pattern type 'T'. Return a reference to `this` for
   /// chaining insertions.
   template <typename T, typename... Args>
diff --git a/mlir/unittests/IR/PatternMatchTest.cpp b/mlir/unittests/IR/PatternMatchTest.cpp
index 1c67bfc284d32..f4ca5c4c49454 100644
--- a/mlir/unittests/IR/PatternMatchTest.cpp
+++ b/mlir/unittests/IR/PatternMatchTest.cpp
@@ -9,6 +9,9 @@
 #include "mlir/IR/PatternMatch.h"
 #include "gtest/gtest.h"
 
+#include <string>
+#include <vector>
+
 #include "../../test/lib/Dialect/Test/TestDialect.h"
 #include "../../test/lib/Dialect/Test/TestOps.h"
 
@@ -35,6 +38,49 @@ TEST(OpRewritePatternTest, GetGeneratedNames) {
 }
 } // end anonymous namespace
 
+namespace {
+template <int ID>
+struct ConstructorArgumentPattern : RewritePattern {
+  ConstructorArgumentPattern(MLIRContext *context, std::string argument,
+                             std::vector<std::string> *arguments)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {
+    arguments->push_back(std::move(argument));
+  }
+
+  LogicalResult matchAndRewrite(Operation *, PatternRewriter &) const override {
+    return failure();
+  }
+};
+
+using ConstructorArgumentPattern0 = ConstructorArgumentPattern<0>;
+using ConstructorArgumentPattern1 = ConstructorArgumentPattern<1>;
+
+TEST(RewritePatternSetTest, CopyArgumentsWhenAddingMultiplePatterns) {
+  MLIRContext context;
+  std::vector<std::string> arguments;
+
+  RewritePatternSet addPatterns(&context);
+  addPatterns.add<ConstructorArgumentPattern0, ConstructorArgumentPattern1>(
+      &context, std::string("add"), &arguments);
+  EXPECT_EQ(arguments, (std::vector<std::string>{"add", "add"}));
+
+  arguments.clear();
+  RewritePatternSet labeledPatterns(&context);
+  labeledPatterns
+      .addWithLabel<ConstructorArgumentPattern0, ConstructorArgumentPattern1>(
+          {"test"}, &context, std::string("addWithLabel"), &arguments);
+  EXPECT_EQ(arguments,
+            (std::vector<std::string>{"addWithLabel", "addWithLabel"}));
+
+  arguments.clear();
+  RewritePatternSet insertPatterns(&context);
+  insertPatterns
+      .insert<ConstructorArgumentPattern0, ConstructorArgumentPattern1>(
+          &context, std::string("insert"), &arguments);
+  EXPECT_EQ(arguments, (std::vector<std::string>{"insert", "insert"}));
+}
+} // end anonymous namespace
+
 namespace {
 LogicalResult anOpRewritePatternFunc(test::OpA op, PatternRewriter &rewriter) {
   return failure();

``````````

</details>


https://github.com/llvm/llvm-project/pull/220871