Fix MLIR DRR matching when attributes are interleaved with operands

[mehdi_amini]

Summary

The ODSOperand indexing should ignore the attributes.

mlir/test/mlir-tblgen/rewriter-indexing.td

 // CHECK: castedOp0.getODSOperands(0);
 // CHECK: castedOp0.getODSOperands(1).begin()).getDefiningOp()
 def test2 : Pat<(COp $attr1, $op1, $attr2, (AOp $op2)),
                 (BOp $attr1, $op2)>;
 
 
 // Check rewriting with a DAG subtree in the result and remapping a location.
 // CHECK: struct test3 : public ::mlir::RewritePattern {
+// We expect ODSOperand 0 here, the attribute before the operand in BOp
+// definition shouldn't shift the counter.
+// CHECK: op1 = (*castedOp0.getODSOperands(0).begin()).getDefiningOp();
 // CHECK: rewriter.create<test::BOp>((*a.getODSResults(0).begin()).getLoc()
 def test3 : Pat<(BOp $attr, (AOp:$a $input)),
                 (BOp $attr, (AOp $input), (location $a))>;
 

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
   // Emits C++ statements for matching using a native code call.
   void emitNativeCodeMatch(DagNode tree, StringRef name, int depth);
 
   // Emits C++ statements for matching the op constrained by the given DAG
   // `tree` returning the op's variable name.
   void emitOpMatch(DagNode tree, StringRef opName, int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
-  // DAG `tree` as an operand.
+  // DAG `tree` as an operand. operandIndex is the index in the DAG excluding
+  // the preceding attributes.
   void emitOperandMatch(DagNode tree, StringRef opName, int argIndex,
-                        int depth);
+                        int operandIndex, int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an attribute.
   void emitAttributeMatch(DagNode tree, StringRef opName, int argIndex,
                           int depth);
 
   // Emits C++ for checking a match with a corresponding match failure
   // diagnostic.
if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       os << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName);
       os.unindent() << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
       // emitOperandMatch's argument indexing counts attributes.
-      emitOperandMatch(tree, castedName, i, depth);
+      emitOperandMatch(tree, castedName, i, nextOperand, depth);
       ++nextOperand;
     } else if (opArg.is<NamedAttribute *>()) {
       emitAttributeMatch(tree, opName, i, depth);
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 }
 
 void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,
-                                      int argIndex, int depth) {
+                                      int argIndex, int operandIndex,
+                                      int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *operand = op.getArg(argIndex).get<NamedTypeConstraint *>();
   auto matcher = tree.getArgAsLeaf(argIndex);
 
   // If a constraint is specified, we need to generate C++ statements to
   // check the constraint.
   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.
     Constraint constraint = matcher.getAsConstraint();
     if (operand->constraint != constraint) {
       if (operand->isVariableLength()) {
         auto error = formatv(
             "further constrain op {0}'s variadic operand #{1} unsupported now",
             op.getOperationName(), argIndex);
         PrintFatalError(loc, error);
       }
       auto self = formatv("(*{0}.getODSOperands({1}).begin()).getType()",
-                          opName, argIndex);
+                          opName, operandIndex);
       emitMatchCheck(
           opName,
           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()));
     }