From 79a291a7f27995fd9ea6d35d6773b67e758edad9 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Mon, 17 Jun 2024 15:00:28 +0000
Subject: [PATCH] [mlir][drr] Fix variadic destination emission
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 | 15 +++++++++++++++
mlir/tools/mlir-tblgen/RewriterGen.cpp | 15 +++++++++------
2 files changed, 24 insertions(+), 6 deletions(-)
diff --git a/mlir/test/mlir-tblgen/rewriter-static-matcher.td b/mlir/test/mlir-tblgen/rewriter-static-matcher.td
index 2907923cb28e4..7a84dfd5706b1 100644
--- a/mlir/test/mlir-tblgen/rewriter-static-matcher.td
+++ b/mlir/test/mlir-tblgen/rewriter-static-matcher.td
@@ -35,6 +35,16 @@ def COp : NS_Op<"c_op", []> {
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 @@ def : Pat<(AOp (BOp I32Attr:$attr, I32:$int)),
// 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)))>;
diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index d8e16d98fd756..b8a3d5b49938e 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1261,20 +1261,23 @@ std::string PatternEmitter::handleResultPattern(DagNode resultTree,
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;
}