[mlir] Update DRR patterns to notify the rewriter why a pattern fails to match.

[rriddle]

Summary

This adds support in RewriterGen for calling into the new PatternRewriter::notifyMatchFailure hook. This lets derived pattern rewriters display this information to users, an example from DialectConversion is shown below:

Legalizing operation : 'std.and'(0x60e0000066a0) {
  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'std.and -> (spv.BitwiseAnd)' {
    ** Failure : operand 0 of op 'std.and' failed to satisfy constraint: '8/16/32/64-bit integer or vector of 8/16/32/64-bit integer values of length 2/3/4'
  } -> FAILURE : pattern failed to match

  * Pattern : 'std.and -> (spv.LogicalAnd)' {
    ** Failure : operand 0 of op 'std.and' failed to satisfy constraint: 'bool or vector of bool values of length 2/3/4'
  } -> FAILURE : pattern failed to match
} -> FAILURE : no matched legalization pattern

Depends On D76313

[rriddle]

Fix typo

[rriddle]

Remove unnecessary .data()

[jpienaar]

Conciser than expected, thanks!

mlir/tools/mlir-tblgen/RewriterGen.cpp

86

match failure?

395

Missing space between h and (

422

Could you extract depth=0 as variable and use it here and in line 409?

[rriddle]

Resolve comments

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an operand.
   void emitOperandMatch(DagNode tree, int argIndex, int depth, int indent);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an attribute.
   void emitAttributeMatch(DagNode tree, int argIndex, int depth, int indent);
 
+  // Emits C++ for checking a match with a corresponding match failure
+  // diagnostic.
+  void emitMatchCheck(int depth, const FmtObjectBase &matchFmt,
+                      const llvm::formatv_object_base &failureFmt);
+
   //===--------------------------------------------------------------------===//
   // Rewrite utilities
   //===--------------------------------------------------------------------===//
 
   // The entry point for handling a result pattern rooted at `resultTree`. This
   // method dispatches to concrete handlers according to `resultTree`'s kind and
   // returns a symbol representing the whole value pack. Callers are expected to
   // further resolve the symbol according to the specific use case.
if (!matcher.isUnspecified()) {
     if (!matcher.isOperandMatcher()) {
       PrintFatalError(
           loc, formatv("the {1}-th argument of op '{0}' should be an operand",
                        op.getOperationName(), argIndex + 1));
     }
 
     // Only need to verify if the matcher's type is different from the one
     // of op definition.
-    if (operand->constraint != matcher.getAsConstraint()) {
+    Constraint constraint = matcher.getAsConstraint();
+    if (operand->constraint != constraint) {
       if (operand->isVariadic()) {
         auto error = formatv(
             "further constrain op {0}'s variadic operand #{1} unsupported now",
             op.getOperationName(), argIndex);
         PrintFatalError(loc, error);
       }
       auto self =
           formatv("(*castedOp{0}.getODSOperands({1}).begin()).getType()", depth,
                   argIndex);
-      os.indent(indent) << "if (!("
-                        << std::string(tgfmt(matcher.getConditionTemplate(),
-                                             &fmtCtx.withSelf(self)))
-                        << ")) return failure();\n";
+      emitMatchCheck(
+          depth,
+          tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),
+          formatv("\"operand {0} of op '{1}' failed to satisfy constraint: "
+                  "'{2}'\"",
+                  operand - op.operand_begin(), op.getOperationName(),
+                  constraint.getDescription()));
     }
   }
 
   // Capture the value
   auto name = tree.getArgName(argIndex);
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     // We need to subtract the number of attributes before this operand to get
     // the index in the operand list.
     auto numPrevAttrs = std::count_if(
         op.arg_begin(), op.arg_begin() + argIndex,
         [](const Argument &arg) { return arg.is<NamedAttribute *>(); });
 
     os.indent(indent) << formatv("{0} = castedOp{1}.getODSOperands({2});\n",
                                  name, depth, argIndex - numPrevAttrs);
   }
 }
 
 void PatternEmitter::emitAttributeMatch(DagNode tree, int argIndex, int depth,
                                         int indent) {
-
   Operator &op = tree.getDialectOp(opMap);
   auto *namedAttr = op.getArg(argIndex).get<NamedAttribute *>();
   const auto &attr = namedAttr->attr;
 
   os.indent(indent) << "{\n";
   indent += 2;
   os.indent(indent) << formatv(
       "auto tblgen_attr = op{0}->getAttrOfType<{1}>(\"{2}\");"
       "(void)tblgen_attr;\n",
       depth, attr.getStorageType(), namedAttr->name);
 
   // TODO(antiagainst): This should use getter method to avoid duplication.
   if (attr.hasDefaultValue()) {
     os.indent(indent) << "if (!tblgen_attr) tblgen_attr = "
                       << std::string(tgfmt(attr.getConstBuilderTemplate(),
                                            &fmtCtx, attr.getDefaultValue()))
                       << ";\n";
   } else if (attr.isOptional()) {
     // For a missing attribute that is optional according to definition, we
     // should just capture a mlir::Attribute() to signal the missing state.
     // That is precisely what getAttr() returns on missing attributes.
   } else {
-    os.indent(indent) << "if (!tblgen_attr) return failure();\n";
+    emitMatchCheck(depth, tgfmt("tblgen_attr", &fmtCtx),
+                   formatv("\"expected op '{0}' to have attribute '{1}' "
+                           "of type '{2}'\"",
+                           op.getOperationName(), namedAttr->name,
+                           attr.getStorageType()));
   }
 
   auto matcher = tree.getArgAsLeaf(argIndex);
   if (!matcher.isUnspecified()) {
     if (!matcher.isAttrMatcher()) {
       PrintFatalError(
           loc, formatv("the {1}-th argument of op '{0}' should be an attribute",
                        op.getOperationName(), argIndex + 1));
     }
 
     // If a constraint is specified, we need to generate C++ statements to
     // check the constraint.
-    os.indent(indent) << "if (!("
-                      << std::string(tgfmt(matcher.getConditionTemplate(),
-                                           &fmtCtx.withSelf("tblgen_attr")))
-                      << ")) return failure();\n";
+    emitMatchCheck(
+        depth,
+        tgfmt(matcher.getConditionTemplate(), &fmtCtx.withSelf("tblgen_attr")),
+        formatv("\"op '{0}' attribute '{1}' failed to satisfy constraint: "
+                "{2}\"",
+                op.getOperationName(), namedAttr->name,
+                matcher.getAsConstraint().getDescription()));
   }
 
   // Capture the value
   auto name = tree.getArgName(argIndex);
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     os.indent(indent) << formatv("{0} = tblgen_attr;\n", name);
   }
 
   indent -= 2;
   os.indent(indent) << "}\n";
 }
 
+void PatternEmitter::emitMatchCheck(
+    int depth, const FmtObjectBase &matchFmt,
+    const llvm::formatv_object_base &failureFmt) {
+  // {0} The match depth (used to get the operation that failed to match).
+  // {1} The format for the match string.
+  // {2} The format for the failure string.
+  const char *matchStr = R"(
+    if (!({1})) {
+      return rewriter.notifyMatchFailure(op{0}, [&](::mlir::Diagnostic &diag) {
+        diag << {2};
+      });
+    })";
+  os << llvm::formatv(matchStr, depth, matchFmt.str(), failureFmt.str())
+     << "\n";
+}
+
 void PatternEmitter::emitMatchLogic(DagNode tree) {
   LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n");
-  emitOpMatch(tree, 0);
+  int depth = 0;
+  emitOpMatch(tree, depth);
 
   for (auto &appliedConstraint : pattern.getConstraints()) {
     auto &constraint = appliedConstraint.constraint;
     auto &entities = appliedConstraint.entities;
 
     auto condition = constraint.getConditionTemplate();
-    auto cmd = "if (!({0})) return failure();\n";
-
     if (isa<TypeConstraint>(constraint)) {
       auto self = formatv("({0}.getType())",
                           symbolInfoMap.getValueAndRangeUse(entities.front()));
-      os.indent(4) << formatv(cmd,
-                              tgfmt(condition, &fmtCtx.withSelf(self.str())));
+      emitMatchCheck(
+          depth, tgfmt(condition, &fmtCtx.withSelf(self.str())),
+          formatv("\"value entity '{0}' failed to satisfy constraint: {1}\"",
+                  entities.front(), constraint.getDescription()));
+
     } else if (isa<AttrConstraint>(constraint)) {
       PrintFatalError(
           loc, "cannot use AttrConstraint in Pattern multi-entity constraints");
     } else {
       // TODO(b/138794486): replace formatv arguments with the exact specified
       // args.
       if (entities.size() > 4) {
         PrintFatalError(loc, "only support up to 4-entity constraints now");
       }
       SmallVector<std::string, 4> names;
       int i = 0;
       for (int e = entities.size(); i < e; ++i)
         names.push_back(symbolInfoMap.getValueAndRangeUse(entities[i]));
       std::string self = appliedConstraint.self;
       if (!self.empty())
         self = symbolInfoMap.getValueAndRangeUse(self);
       for (; i < 4; ++i)
         names.push_back("<unused>");
-      os.indent(4) << formatv(cmd,
-                              tgfmt(condition, &fmtCtx.withSelf(self), names[0],
-                                    names[1], names[2], names[3]));
+      emitMatchCheck(depth,
+                           names[1], names[2], names[3]),
+                     formatv("\"entities '{0}' failed to satisfy constraint: "
+                             "{1}\"",
+                             llvm::join(entities, ", "),
+                             constraint.getDescription()));
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "--- done emitting match logic ---\n");
 }
 
 void PatternEmitter::collectOps(DagNode tree,
                                 llvm::SmallPtrSetImpl<const Operator *> &ops) {
   // Check if this tree is an operation.