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

[Jian Cai]

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

>From 781ab12010a045984460c0cf3323c6f06ba906d4 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 for supplemental patterns to be able to access the source op.
---
 mlir/tools/mlir-tblgen/RewriterGen.cpp | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 6bb79fb4b4cbe67..545e08de85a481a 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,20 +1209,10 @@ void PatternEmitter::emitRewriteLogic() {
           "  tblgen_repl_values.push_back(v);\n}\n",
           "\n");
     }
+    processSupplementalPatterns();
     os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }
 
-  // Process supplemtal patterns.
-  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";
-  }
-
   LLVM_DEBUG(llvm::dbgs() << "--- done emitting rewrite logic ---\n");
 }