[Mlir-commits] [mlir] [mlir][drr] Add warning for simple case of mismatched variadic. (PR #84040)

[Jacques Pienaar]

https://github.com/jpienaar created https://github.com/llvm/llvm-project/pull/84040

When a variadic argument is expected but not provided the compilation fails later with a difficult to follow compilation error. Add a simple check to catch one such case.

This is not yet general as it doesn't yet check leaf nodes.

>From 490712968757b6b8aba473d29002d6e37169dcc0 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Tue, 5 Mar 2024 08:50:08 -0800
Subject: [PATCH] [mlir][drr] Add warning for simple case of mismatched
 variadic.

When a variadic argument is expected but not provided the compilation
fails later with a difficult to follow compilation error. Add a simple
check to catch one such case.

This is not yet general as it doesn't yet check leaf nodes.
---
 mlir/test/mlir-tblgen/rewriter-errors.td | 15 +++++++++++++
 mlir/tools/mlir-tblgen/RewriterGen.cpp   | 28 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+)

diff --git a/mlir/test/mlir-tblgen/rewriter-errors.td b/mlir/test/mlir-tblgen/rewriter-errors.td
index f74fb7a6981026..0b6d8c3fec6134 100644
--- a/mlir/test/mlir-tblgen/rewriter-errors.td
+++ b/mlir/test/mlir-tblgen/rewriter-errors.td
@@ -5,6 +5,7 @@
 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR5 %s 2>&1 | FileCheck --check-prefix=ERROR5 %s
 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR6 %s 2>&1 | FileCheck --check-prefix=ERROR6 %s
 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR7 %s 2>&1 | FileCheck --check-prefix=ERROR7 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR8 %s 2>&1 | FileCheck --check-prefix=ERROR8 %s
 
 include "mlir/IR/OpBase.td"
 include "mlir/IR/PatternBase.td"
@@ -64,3 +65,17 @@ def : Pat<(OpB:$result $val, $attr), (OpA $val, $val), [(AnyInteger:$result)]>;
 // ERROR7: [[@LINE+1]]:1: error: type constraint requires exactly one argument
 def : Pat<(OpB:$opB $val, $attr), (OpA $val, $val), [(AnyInteger $opB, $val)]>;
 #endif
+
+def OpC : A_Op<"op_c">, Results<(outs AnyInteger)>;
+def OpD : A_Op<"op_d">, Arguments<(ins Variadic<AnyInteger>:$vargs)>, Results<(outs AnyInteger)>;
+
+#ifdef ERROR8
+// Check that op with variadic operand gets variadic operand in target,
+//
+// FIXME: this should be an error.
+def : Pat<(OpB:$opB $val, $attr), (OpD $val)>;
+
+// ERROR8: [[@LINE+2]]
+// ERROR8-SAME: op expects variadic operand `vargs`, while provided is non-variadic
+def : Pat<(OpB:$opB $val, $attr), (OpD (OpC))>;
+#endif
diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 426a3482960be4..e63a065a070847 100644
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Support/IndentedOstream.h"
+#include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/CodeGenHelpers.h"
 #include "mlir/TableGen/Format.h"
@@ -18,6 +19,7 @@
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/Pattern.h"
 #include "mlir/TableGen/Predicate.h"
+#include "mlir/TableGen/Property.h"
 #include "mlir/TableGen/Type.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/SetVector.h"
@@ -1518,10 +1520,36 @@ std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
   // the key. This includes both bound and unbound child nodes.
   ChildNodeIndexNameMap childNodeNames;
 
+  // If the argument is a type constraint, then its an operand. Check if the
+  // op's argument is variadic that the argument in the pattern is too.
+  auto checkIfMatchedVariadic = [&](int i) {
+    // FIXME: This does not yet check for variable/leaf case.
+    // FIXME: Change so that native code call can be handled.
+    const auto *operand =
+        llvm::dyn_cast_if_present<NamedTypeConstraint *>(resultOp.getArg(i));
+    if (!operand || !operand->isVariadic())
+      return;
+
+    auto child = tree.getArgAsNestedDag(i);
+    if (!child)
+      return;
+
+    // Skip over replaceWithValues.
+    while (child.isReplaceWithValue()) {
+      if (!(child = child.getArgAsNestedDag(0)))
+        return;
+    }
+    if (!child.isNativeCodeCall() && !child.isVariadic())
+      PrintFatalError(loc, formatv("op expects variadic operand `{0}`, while "
+                                   "provided is non-variadic",
+                                   resultOp.getArgName(i)));
+  };
+
   // First go through all the child nodes who are nested DAG constructs to
   // create ops for them and remember the symbol names for them, so that we can
   // use the results in the current node. This happens in a recursive manner.
   for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {
+    checkIfMatchedVariadic(i);
     if (auto child = tree.getArgAsNestedDag(i))
       childNodeNames[i] = handleResultPattern(child, i, depth + 1);
   }