[Mlir-commits] [mlir] [mlir] Move supplemental patterns before op replacement (PR #66959)

[Jian Cai]

https://github.com/jcai19 created https://github.com/llvm/llvm-project/pull/66959

This moves the C++ code generated from supplemental patterns before op replacement. It is necessary if the supllemental patterns need to access the source op.

>From 396ae03e4ae4acb41eb28b8bde91410ff36d74be Mon Sep 17 00:00:00 2001
From: Jian Cai <caij2003 at gmail.com>
Date: Wed, 20 Sep 2023 15:36:06 -0700
Subject: [PATCH] [mlir] Move supplemental patterns before op replacement

This moves the C++ code generated from supplemental patterns before op replacement. It is necessary if the supllemental patterns need to access the source op.
---
 mlir/tools/mlir-tblgen/RewriterGen.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 6bb79fb4b4cbe67..bc8d16b8478f662 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1172,9 +1172,22 @@ void PatternEmitter::emitRewriteLogic() {
       os << val << ";\n";
   }
 
+  auto processSupplementalPatterns = [&]() {
+    int numSupplementalPatterns = pattern.getNumSupplementalPatterns();
+    for (int i = 0, offset = -numSupplementalPatterns;
+         i < numSupplementalPatterns; ++i) {
+      DagNode resultTree = pattern.getSupplementalPattern(i);
+      auto val = handleResultPattern(resultTree, offset++, 0);
+      if (resultTree.isNativeCodeCall() &&
+          resultTree.getNumReturnsOfNativeCode() == 0)
+        os << val << ";\n";
+    }
+  };
+
   if (numExpectedResults == 0) {
     assert(replStartIndex >= numResultPatterns &&
            "invalid auxiliary vs. replacement pattern division!");
+    processSupplementalPatterns();
     // No result to replace. Just erase the op.
     os << "rewriter.eraseOp(op0);\n";
   } else {
@@ -1196,6 +1209,7 @@ void PatternEmitter::emitRewriteLogic() {
           "  tblgen_repl_values.push_back(v);\n}\n",
           "\n");
     }
+    processSupplementalPatterns();
     os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }