Relax the restriction that parsed arguments must match op arguments in ODS.

[Mokosha]

Summary

In particular, when an op defines one or more optional attributes, we should
be able to write pattern rewrites that omit those attributes.

[jpienaar]

Could you please update the documentation and description of the change?

Seems like you are assuming some C++ like default call behavior from the patterns, or how do you plan to generalize this?

And some tests for error cases would be good.

[Mokosha]

Fix clang-format to use camel case instead of snake case.

[Mokosha]

Fix cases where we may have location directives

[Mokosha]

Fix some tests to use actual required args and add failure tests

[Mokosha]

Add some documentation and refine the tests a bit further.

[Mokosha]

I've added some documentation that hopefully highlights the intent behind the change. The way that we generate rewrite implementations for ops (i.e. the build() call that we use) means that we can omit some of the attributes if they're marked optional and everything should "just work". The tests that were added will hopefully catch this if anything fails in the future. PTAL (and let me know if I'm missing something egregious here!)

mlir/docs/DeclarativeRewrites.md

void AOp::build(OpBuilder &builder, OperationState &state,
   state.addTypes({type});
 }
 ```
 
 Failing to define such a builder will result in an error at C++ compilation time
 saying the call to `AOp::build()` cannot be resolved because of the number of
 parameters mismatch.
 
+#### Optional Op Arguments
+
+In the case of ODS-generated rewrite rules, we use the special `build()` method
+that was described in the [previous section](#building-operations). When
+specifying optional attributes as arguments, they may be omitted from the
+patterns in certain cases, namely when omitting the result from
+the parameterized lists passed to the `build()` function won't cause the op
+creation to fail. For example, suppose we have the following `OptionalOp`:
+
+```tablegen
+def OptionalOp : Op<"optional_op"> {
+    let arguments = (ins OptionalAttr<BoolAttr>:$opt);
+}
+```
+
+In this case, the following would be the valid rewrite rules for `OptionalOp`:
+
+```tablegen
+def : Pat<(AOp $input, $attr), (OptionalOp)>;
+def : Pat<(AOp $input, $attr), (OptionalOp ConstBoolAttrTrue)>;
+def : Pat<(AOp $input, $attr), (OptionalOp ConstBoolAttrFalse)>;
+```
+
+The number of arguments required to construct an op is therefore determined by
+the last non-optional (required) argument to the Op. Take for example the
+following definition of `OptioanlOp`:
+
+```tablegen
+def OptionalOp : Op<"optional_op"> {
+    let arguments = (ins
+        OptionalAttr<BoolAttr>:$opt1,
+        BoolAttr:$req,
+        OptionalAttr<BoolAttr>:$opt2
+    );
+}
+```
+
+In this case, `$opt1` is always required, since there is no way to construct the
+op without specifying `$opt1` in order to also specify `$req`. This means that
+`OptionalOp` requires at least two arguments. However, `$opt2` is "truly"
+optional, and may be omitted from result rewrite patterns:
+
+```tablegen
+// NOT OK -- $req is not defined
+def : Pat<(AOp $i, $a), (OptionalOp)>;
+
+// NOT OK -- $req is not defined
+def : Pat<(AOp $i, $a), (OptionalOp ConstBoolAttrTrue)>;
+
+// OK
+def : Pat<(AOp $i, $a), (OptionalOp ConstBoolAttrTrue ConstBoolAttrTrue)>;
+
+// OK
+def : Pat<(AOp $i, $a),
+          (OptionalOp ConstBoolAttrTrue ConstBoolAttrTrue ConstBoolAttrTrue)>;
+```
+
+Optional attributes can also be omitted in the source pattern, but this is
+identical to using `$_` to denote an unused argument:
+
+```tablegen
+// The following are identical:
+def : Pat<(OptionalOp $opt, $req),
+          (OptionalOp $opt, $req, ConstBoolAttrTrue)>;
+
+def : Pat<(OptionalOp $opt, $req, $_),
+          (OptionalOp $opt, $req, ConstBoolAttrTrue)>;
+```
+
 #### Generating DAG of operations
 
 `dag` objects can be nested to generate a DAG of operations:
 
 ```tablegen
 def : Pat<(AOp $input, $attr), (COp (BOp), $attr)>;
 ```
 

mlir/include/mlir/TableGen/Operator.h

public:
   }
 
   // Returns the number of variadic operands in this operation.
   unsigned getNumVariableLengthOperands() const;
 
   // Returns the total number of arguments.
   int getNumArgs() const { return arguments.size(); }
 
+  // Returns the total number of required arguments.
+  int getNumRequiredArgs() const;
+
   // Returns true of the operation has a single variadic arg.
   bool hasSingleVariadicArg() const;
 
   // Returns true if the operation has a single variadic result.
   bool hasSingleVariadicResult() const {
     return getNumResults() == 1 && getResult(0).isVariadic();
   }
 

mlir/lib/TableGen/Operator.cpp

 }
 
 unsigned Operator::getNumVariableLengthOperands() const {
   return llvm::count_if(operands, [](const NamedTypeConstraint &c) {
     return c.constraint.isVariableLength();
   });
 }
 
+int Operator::getNumRequiredArgs() const {
+  int numRequiredArgs = getNumArgs();
+  while (numRequiredArgs > 0) {
+    auto arg = getArg(numRequiredArgs - 1);
+    if (arg.is<NamedTypeConstraint *>()) {
+      if (!arg.get<NamedTypeConstraint *>()->isOptional()) {
+        break;
+      }
+    } else if (!arg.get<NamedAttribute *>()->attr.isOptional()) {
+      break;
+    }
+    --numRequiredArgs;
+  }
+  return numRequiredArgs;
+}
+
 bool Operator::hasSingleVariadicArg() const {
   return getNumArgs() == 1 && getArg(0).is<NamedTypeConstraint *>() &&
          getOperand(0).isVariadic();
 }
 
 Operator::arg_iterator Operator::arg_begin() const { return arguments.begin(); }
 
 Operator::arg_iterator Operator::arg_end() const { return arguments.end(); }

mlir/lib/TableGen/Pattern.cpp

for (int i = 0; i != numTreeArgs; ++i) {
       }
     }
 
     return;
   }
 
   if (tree.isOperation()) {
     auto &op = getDialectOp(tree);
-    auto numOpArgs = op.getNumArgs();
 
     // The pattern might have the last argument specifying the location.
     bool hasLocDirective = false;
     if (numTreeArgs != 0) {
-      if (auto lastArg = tree.getArgAsNestedDag(numTreeArgs - 1))
+      if (auto lastArg = tree.getArgAsNestedDag(numTreeArgs - 1)) {
         hasLocDirective = lastArg.isLocationDirective();
       }
+    }
 
-    if (numOpArgs != numTreeArgs - hasLocDirective) {
-      auto err = formatv("op '{0}' argument number mismatch: "
-                         "{1} in pattern vs. {2} in definition",
-                         op.getOperationName(), numTreeArgs, numOpArgs);
+    const int numRequiredArgs = op.getNumRequiredArgs();
+    const int opArgs = numTreeArgs - hasLocDirective;
+    if (opArgs < numRequiredArgs) {
+      auto err = formatv("too few arguments for op '{0}': "
+                         "{1} in pattern vs. {2} required in definition",
+                         op.getOperationName(), opArgs, numRequiredArgs);
+      PrintFatalError(&def, err);
+    }
+
+    if (opArgs > op.getNumArgs()) {
+      auto err = formatv("too many arguments for op '{0}': "
+                         "{1} in pattern vs. {2} allowed in definition",
+                         op.getOperationName(), opArgs, op.getNumArgs());
       PrintFatalError(&def, err);
     }
 
     // The name attached to the DAG node's operator is for representing the
     // results generated from this op. It should be remembered as bound results.
     if (!treeName.empty()) {
       LLVM_DEBUG(llvm::dbgs()
                  << "found symbol bound to op result: " << treeName << '\n');
       verifyBind(infoMap.bindOpResult(treeName, op), treeName);
     }
 
-    for (int i = 0; i != numTreeArgs; ++i) {
+    for (int i = 0; i != opArgs; ++i) {
       if (auto treeArg = tree.getArgAsNestedDag(i)) {
         // This DAG node argument is a DAG node itself. Go inside recursively.
         collectBoundSymbols(treeArg, infoMap, isSrcPattern);
         continue;
       }
 
       if (isSrcPattern) {
         // We can only bind symbols to op arguments in source pattern. Those

mlir/test/lib/Dialect/Test/TestOps.td

 def : Pat<(OpC $input), (OpB $input, ConstantAttr<I32Attr, "17">:$attr)>;
 
 // Test string enum attribute in rewrites.
 def : Pat<(StrEnumAttrOp StrCaseA), (StrEnumAttrOp StrCaseB)>;
 // Test integer enum attribute in rewrites.
 def : Pat<(I32EnumAttrOp I32Case5), (I32EnumAttrOp I32Case10)>;
 def : Pat<(I64EnumAttrOp I64Case5), (I64EnumAttrOp I64Case10)>;
 
+// Test partial matches for optional attributes
+def OpAttrMatch5 : TEST_Op<"match_optional_op_attribute1"> {
+  let arguments = (ins
+    BoolAttr:$required_attr,
+    OptionalAttr<I32Attr>:$optional_attr1,
+    OptionalAttr<I32Attr>:$optional_attr2
+  );
+}
+def OpAttrMatch6 : TEST_Op<"match_optional_op_attribute2"> {
+  let arguments = (ins
+    BoolAttr:$required_attr1,
+    OptionalAttr<I32Attr>:$optional_attr1,
+    OptionalAttr<I32Attr>:$optional_attr2,
+    BoolAttr:$required_attr2
+  );
+}
+def : Pat<(OpAttrMatch5 ConstBoolAttrTrue, $optional),
+          (OpAttrMatch5 ConstBoolAttrFalse, $optional)>;
+def : Pat<(OpAttrMatch5 ConstBoolAttrFalse, $optional, $_),
+          (OpAttrMatch5 ConstBoolAttrTrue, $optional, $optional)>;
+def : Pat<(OpAttrMatch6 ConstBoolAttrTrue, $optional, $optional2, ConstBoolAttrFalse),
+          (OpAttrMatch5 ConstBoolAttrTrue, $optional)>;
+def : Pat<(OpAttrMatch6 ConstBoolAttrFalse, $optional, $optional2, $required2),
+          (OpAttrMatch5 $required2, $optional, $optional2)>;
+def : Pat<(OpAttrMatch6 $required, $optional, $optional2, $required2),
+          (OpAttrMatch5 $required)>;
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Multi-result Ops)
 
 def MultiResultOpKind1: I64EnumAttrCase<"kind1", 1>;
 def MultiResultOpKind2: I64EnumAttrCase<"kind2", 2>;
 def MultiResultOpKind3: I64EnumAttrCase<"kind3", 3>;
 def MultiResultOpKind4: I64EnumAttrCase<"kind4", 4>;
 def MultiResultOpKind5: I64EnumAttrCase<"kind5", 5>;

mlir/test/mlir-tblgen/op-error.td

 // RUN: not mlir-tblgen -gen-op-decls -I %S/../../include -DERROR1 %s 2>&1 | FileCheck --check-prefix=ERROR1 %s
 // RUN: not mlir-tblgen -gen-op-decls -I %S/../../include -DERROR2 %s 2>&1 | FileCheck --check-prefix=ERROR2 %s
 // RUN: not mlir-tblgen -gen-op-decls -I %S/../../include -DERROR3 %s 2>&1 | FileCheck --check-prefix=ERROR3 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR4 %s 2>&1 | FileCheck --check-prefix=ERROR4 %s
+// 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
 
 include "mlir/IR/OpBase.td"
 
 def Test_Dialect : Dialect {
   let name = "test_dialect";
 }
 
 #ifdef ERROR1

 #ifdef ERROR3
 // ERROR3: error: expected an argument with default value after other arguments with default values
 def OpDefaultValueNotTrailing : Op<Test_Dialect, "default_value"> {
   let builders = [
     OpBuilderDAG<(ins CArg<"int", "42">, CArg<"int">)>
   ];
 }
 #endif
+
+def OpAttrMatch : Op<Test_Dialect, "optional_attr_match"> {
+  let arguments = (ins
+    BoolAttr:$required_attr1,
+    OptionalAttr<I32Attr>:$optional_attr1,
+    OptionalAttr<I32Attr>:$optional_attr2,
+    BoolAttr:$required_attr2
+  );
+  let results = (outs I32);
+}
+
+#ifdef ERROR4
+// ERROR4: error: too few arguments for op 'test_dialect.optional_attr_match'
+def : Pat<(OpAttrMatch ConstBoolAttrTrue, $optional, $optional2, ConstBoolAttrFalse),
+          (OpAttrMatch ConstBoolAttrFalse)>;
+#endif
+
+#ifdef ERROR5
+// ERROR5: error: too few arguments for op 'test_dialect.optional_attr_match'
+def : Pat<(OpAttrMatch:$op1 ConstBoolAttrTrue, $optional, $optional2, ConstBoolAttrFalse),
+          (OpAttrMatch ConstBoolAttrFalse, $optional, $optional, (location $op1))>;
+#endif
+
+#ifdef ERROR6
+// ERROR6: error: too many arguments for op 'test_dialect.optional_attr_match'
+def : Pat<(OpAttrMatch ConstBoolAttrTrue, $o, $o2, ConstBoolAttrFalse),
+          (OpAttrMatch ConstBoolAttrFalse, $o, $o, $o, $o)>;
+#endif

mlir/tools/mlir-tblgen/RewriterGen.cpp

void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) {
   os << formatv("auto {0} = ::llvm::dyn_cast_or_null<{2}>({1}); "
                 "(void){0};\n",
                 castedName, opName, op.getQualCppClassName());
   // Skip the operand matching at depth 0 as the pattern rewriter already does.
   if (depth != 0) {
     // Skip if there is no defining operation (e.g., arguments to function).
     os << formatv("if (!{0}) return failure();\n", castedName);
   }
-  if (tree.getNumArgs() != op.getNumArgs()) {
+  if (tree.getNumArgs() < op.getNumRequiredArgs() ||
+      tree.getNumArgs() > op.getNumArgs()) {
     PrintFatalError(loc, formatv("op '{0}' argument number mismatch: {1} in "
-                                 "pattern vs. {2} in definition",
+                                 "pattern vs. [{2},{3}] required in definition",
                                  op.getOperationName(), tree.getNumArgs(),
-                                 op.getNumArgs()));
+                                 op.getNumRequiredArgs(), op.getNumArgs()));
   }
 
   // If the operand's name is set, set to that variable.
   auto name = tree.getSymbol();
   if (!name.empty())
     os << formatv("{0} = {1};\n", name, castedName);
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {

 
 std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
                                              int depth) {
   LLVM_DEBUG(llvm::dbgs() << "create op for pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   Operator &resultOp = tree.getDialectOp(opMap);
-  auto numOpArgs = resultOp.getNumArgs();
   auto numPatArgs = tree.getNumArgs();
 
   bool hasLocationDirective;
   std::string locToUse;
   std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
   auto inPattern = numPatArgs - hasLocationDirective;
-  if (numOpArgs != inPattern) {
+  if (resultOp.getNumRequiredArgs() > inPattern) {
     PrintFatalError(loc,
-                    formatv("resultant op '{0}' argument number mismatch: "
-                            "{1} in pattern vs. {2} in definition",
-                            resultOp.getOperationName(), inPattern, numOpArgs));
+                    formatv("resultant op '{0}' fas too few arguments: "
+                            "{1} in pattern vs. {2} required in definition",
+                            resultOp.getOperationName(), inPattern,
+                            resultOp.getNumRequiredArgs()));
+  }
+
+  if (inPattern > resultOp.getNumArgs()) {
+    PrintFatalError(loc, formatv("resultant op '{0}' has too many arguments: "
+                                 "{1} in pattern vs. {2} allowed in definition",
+                                 resultOp.getOperationName(), inPattern,
+                                 resultOp.getNumArgs()));
   }
 
   // A map to collect all nested DAG child nodes' names, with operand index as
   // the key. This includes both bound and unbound child nodes.
   ChildNodeIndexNameMap childNodeNames;
 
   // 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
void PatternEmitter::createSeparateLocalVarsForOpArgs(
     DagNode node, ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);
 
   // Now prepare operands used for building this op:
   // * If the operand is non-variadic, we create a `Value` local variable.
   // * If the operand is variadic, we create a `SmallVector<Value>` local
   //   variable.
 
+  int numOpArgs = node.getNumArgs();
+  if (auto lastArg = node.getArgAsNestedDag(numOpArgs - 1)) {
+    numOpArgs -= lastArg.isLocationDirective();
+  }
+
   int valueIndex = 0; // An index for uniquing local variable names.
-  for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
+  for (int argIndex = 0; argIndex < numOpArgs; ++argIndex) {
     const auto *operand =
         resultOp.getArg(argIndex).dyn_cast<NamedTypeConstraint *>();
     // We do not need special handling for attributes.
     if (!operand)
       continue;
 
     raw_indented_ostream::DelimitedScope scope(os);
     std::string varName;
for (int argIndex = 0; argIndex < numOpArgs; ++argIndex) {
     // Update to use the newly created local variable for building the op later.
     childNodeNames[argIndex] = varName;
   }
 }
 
 void PatternEmitter::supplyValuesForOpArgs(
     DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) {
   Operator &resultOp = node.getDialectOp(opMap);
-  for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();
-       argIndex != numOpArgs; ++argIndex) {
+
+  int numOpArgs = node.getNumArgs();
+  if (auto lastArg = node.getArgAsNestedDag(numOpArgs - 1)) {
+    numOpArgs -= lastArg.isLocationDirective();
+  }
+
+  for (int argIndex = 0; argIndex != numOpArgs; ++argIndex) {
     // Start each argument on its own line.
     os << ",\n    ";
 
     Argument opArg = resultOp.getArg(argIndex);
     // Handle the case of operand first.
     if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
       if (!operand->name.empty())
         os << "/*" << operand->name << "=*/";
       os << childNodeNames.lookup(argIndex);
       continue;
     }
 
     // The argument in the op definition.
     auto opArgName = resultOp.getArgName(argIndex);
     if (auto subTree = node.getArgAsNestedDag(argIndex)) {
+      if (subTree.isLocationDirective())
+        continue;
       if (!subTree.isNativeCodeCall())
         PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                              "for creating attribute");
       os << formatv("/*{0}=*/{1}", opArgName,
                     handleReplaceWithNativeCodeCall(subTree, depth));
     } else {
       auto leaf = node.getArgAsLeaf(argIndex);
       // The argument in the result DAG pattern.
os << formatv("::mlir::SmallVector<::mlir::Value, 4> "
                 "tblgen_values; (void)tblgen_values;\n");
   os << formatv("::mlir::SmallVector<::mlir::NamedAttribute, 4> "
                 "tblgen_attrs; (void)tblgen_attrs;\n");
 
   const char *addAttrCmd =
       "if (auto tmpAttr = {1}) {\n"
       "  tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), "
       "tmpAttr);\n}\n";
-  for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
+  for (int argIndex = 0, e = node.getNumArgs(); argIndex < e; ++argIndex) {
+    auto subTree = node.getArgAsNestedDag(argIndex);
+    if (subTree && subTree.isLocationDirective())
+      continue;
+
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
       // The argument in the op definition.
       auto opArgName = resultOp.getArgName(argIndex);
-      if (auto subTree = node.getArgAsNestedDag(argIndex)) {
+      if (subTree) {
         if (!subTree.isNativeCodeCall())
           PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                                "for creating attribute");
         os << formatv(addAttrCmd, opArgName,
                       handleReplaceWithNativeCodeCall(subTree, depth + 1));
       } else {
         auto leaf = node.getArgAsLeaf(argIndex);
         // The argument in the result DAG pattern.