commit ca1a96364a568be6a2e65e6f12328b23dd35bafc
author Jacques Pienaar <jpienaar@google.com> Mon Jun 17 15:06:15 2024 -0700
committer GitHub <noreply@github.com> Mon Jun 17 15:06:15 2024 -0700
tree c4936073ae9b5603f7960565c0cb2267ca8bbcdb
parent ca22469a101e73fd2b2057d98d3f880f064764e5

[mlir][drr] Fix variadic destination emission (#95855)

Its possible for handleResultPattern to emit helpers, these helpers
cannot be interleaved with pushing into the array. Emit into a separate
string to enable helpers to be emitted before the population of vector.

Signed-off-by: Jacques Pienaar <jpienaar@google.com>

mlir/test/mlir-tblgen/rewriter-static-matcher.td

diff --git a/mlir/test/mlir-tblgen/rewriter-static-matcher.td b/mlir/test/mlir-tblgen/rewriter-static-matcher.td
index 2907923..7a84dfd 100644
--- a/mlir/test/mlir-tblgen/rewriter-static-matcher.td
+++ b/mlir/test/mlir-tblgen/rewriter-static-matcher.td
@@ -35,6 +35,16 @@
   let results = (outs AnyInteger);
 }
 
+def DOp : NS_Op<"d_op", []> {
+  let arguments = (ins
+    Variadic<AnyInteger>:$any_integer
+  );
+
+  let results = (outs AnyInteger);
+}
+
+def Foo : NativeCodeCall<"foo($_builder, $0)">;
+
 // Test static matcher for duplicate DagNode
 // ---
 
@@ -53,3 +63,8 @@
 // CHECK: if(::mlir::failed([[$DAG_MATCHER]](rewriter, op1, tblgen_ops
 def : Pat<(COp $_, (BOp I32Attr:$attr, I32:$int)),
           (COp $attr, $int)>;
+
+// CHECK: auto [[$VAR:.*]] = foo(
+// CHECK: ::llvm::SmallVector<::mlir::Value, 4> [[$ARR:tblgen_variadic_values_.*]];
+// CHECK: [[$ARR]].push_back([[$VAR]]);
+def : Pat<(AOp $x), (DOp (variadic (Foo $x)))>;

mlir/tools/mlir-tblgen/RewriterGen.cpp

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index d8e16d9..b8a3d5b 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1261,20 +1261,23 @@
 std::string PatternEmitter::handleVariadic(DagNode tree, int depth) {
   assert(tree.isVariadic());
 
+  std::string output;
+  llvm::raw_string_ostream oss(output);
   auto name = std::string(formatv("tblgen_variadic_values_{0}", nextValueId++));
   symbolInfoMap.bindValue(name);
-  os << "::llvm::SmallVector<::mlir::Value, 4> " << name << ";\n";
+  oss << "::llvm::SmallVector<::mlir::Value, 4> " << name << ";\n";
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     if (auto child = tree.getArgAsNestedDag(i)) {
-      os << name << ".push_back(" << handleResultPattern(child, i, depth + 1)
-         << ");\n";
+      oss << name << ".push_back(" << handleResultPattern(child, i, depth + 1)
+          << ");\n";
     } else {
-      os << name << ".push_back("
-         << handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i))
-         << ");\n";
+      oss << name << ".push_back("
+          << handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i))
+          << ");\n";
     }
   }
 
+  os << oss.str();
   return name;
 }