[mlir] Handle NativeCodeCallVoid in result patterns.

[jcai19]

Currently NativeCodeCallVoid is not supported in the result patterns. For example, below code will fail to build with an error message "referencing unbound symbol"

def Foo: NativeCodeCallVoid<"foo()">;

def AddToAddV2 : Pattern<
(TF_AddOp TF_NumberTensor:$arg0, TF_NumberTensor:$arg1),
[(TF_AddV2Op $arg0, $arg1), (Foo)]>;

MLIR tablegen-based pattern rewrites does not preserve attributes of the source op, with this change users could mannualy copy source attributes to the target op via NativeCodeCallVoid. This is a replacement of D157032.

[llvmbot]

@llvm/pr-subscribers-mlir

Changes

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 6bb79fb4b4cbe67..bc2731df1850838 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1184,17 +1184,22 @@ void PatternEmitter::emitRewriteLogic() {
       DagNode resultTree = pattern.getResultPattern(i);
       auto val = handleResultPattern(resultTree, offsets[i], 0);
       os << "\n";
-      // Resolve each symbol for all range use so that we can loop over them.
-      // We need an explicit cast to `SmallVector` to capture the cases where
-      // `{0}` resolves to an `Operation::result_range` as well as cases that
-      // are not iterable (e.g. vector that gets wrapped in additional braces by
-      // RewriterGen).
-      // TODO: Revisit the need for materializing a vector.
-      os << symbolInfoMap.getAllRangeUse(
-          val,
-          "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
-          "  tblgen_repl_values.push_back(v);\n}\n",
-          "\n");
+      if (resultTree.isNativeCodeCall() &&
+          resultTree.getNumReturnsOfNativeCode() == 0) {
+        os << val << ";\n";
+      } else {
+        // Resolve each symbol for all range use so that we can loop over them.
+        // We need an explicit cast to `SmallVector` to capture the cases where
+        // `{0}` resolves to an `Operation::result_range` as well as cases that
+        // are not iterable (e.g. vector that gets wrapped in additional braces by
+        // RewriterGen).
+        // TODO: Revisit the need for materializing a vector.
+        os << symbolInfoMap.getAllRangeUse(
+            val,
+            "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
+            "  tblgen_repl_values.push_back(v);\n}\n",
+            "\n");
+      }
     }
     os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }

[jpienaar]

The pattern looks wrong though, from PatternBase.td

  // ... In the case of more result patterns
  // than needed to replace the source op, only the last N results generated
  // by the last N result pattern is used to replace a N-result source op.

so in

def AddToAddV2 : Pattern<
(TF_AddOp TF_NumberTensor:$arg0, TF_NumberTensor:$arg1),
[(TF_AddV2Op $arg0, $arg1), (Foo)]>;

TF_AddOp should be replaced by what Foo produces (that's a requirement of DRR that these match), and so Foo should be returning a result here of the same type as TF_AddOp.

[jcai19]

Thanks for the clarification! My understanding is that since Foo is a void function (e.g. copying op attributes) and it does not return any values (i.e. return 0 result), "the last N results" will be the result of TF_AddV2Op pattern in this case. This is meant to replace reviews.llvm.org/D157032 because the code generated by the post-processing parameter I added in that patch is inserted after the source op is removed so I cannot copy the attributes before the original op is removed. If you think this change does not make sense, I can patch reviews.llvm.org/D157032 and move the post-processing code right before the original op is replaced (and removed). Would the latter make more sense to you?

[llvmbot]

@llvm/pr-subscribers-mlir-core

Changes

Currently NativeCodeCallVoid is not supported in the result patterns. For example, below code will fail to build with an error message "referencing unbound symbol"

def Foo: NativeCodeCallVoid<"foo()">;

def AddToAddV2 : Pattern<
(TF_AddOp TF_NumberTensor:$arg0, TF_NumberTensor:$arg1),
[(TF_AddV2Op $arg0, $arg1), (Foo)]>;

MLIR tablegen-based pattern rewrites does not preserve attributes of the source op, with this change users could mannualy copy source attributes to the target op via NativeCodeCallVoid. This is a replacement of D157032.

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

1 Files Affected:

• (modified) mlir/tools/mlir-tblgen/RewriterGen.cpp (+16-11)

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 6bb79fb4b4cbe67..bc2731df1850838 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -1184,17 +1184,22 @@ void PatternEmitter::emitRewriteLogic() {
       DagNode resultTree = pattern.getResultPattern(i);
       auto val = handleResultPattern(resultTree, offsets[i], 0);
       os << "\n";
-      // Resolve each symbol for all range use so that we can loop over them.
-      // We need an explicit cast to `SmallVector` to capture the cases where
-      // `{0}` resolves to an `Operation::result_range` as well as cases that
-      // are not iterable (e.g. vector that gets wrapped in additional braces by
-      // RewriterGen).
-      // TODO: Revisit the need for materializing a vector.
-      os << symbolInfoMap.getAllRangeUse(
-          val,
-          "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
-          "  tblgen_repl_values.push_back(v);\n}\n",
-          "\n");
+      if (resultTree.isNativeCodeCall() &&
+          resultTree.getNumReturnsOfNativeCode() == 0) {
+        os << val << ";\n";
+      } else {
+        // Resolve each symbol for all range use so that we can loop over them.
+        // We need an explicit cast to `SmallVector` to capture the cases where
+        // `{0}` resolves to an `Operation::result_range` as well as cases that
+        // are not iterable (e.g. vector that gets wrapped in additional braces by
+        // RewriterGen).
+        // TODO: Revisit the need for materializing a vector.
+        os << symbolInfoMap.getAllRangeUse(
+            val,
+            "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
+            "  tblgen_repl_values.push_back(v);\n}\n",
+            "\n");
+      }
     }
     os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }

[jcai19]

Abandon this change in favor of #66959.