diff --git a/mlir/test/mlir-tblgen/rewriter-indexing.td b/mlir/test/mlir-tblgen/rewriter-indexing.td --- a/mlir/test/mlir-tblgen/rewriter-indexing.td +++ b/mlir/test/mlir-tblgen/rewriter-indexing.td @@ -51,6 +51,9 @@ // 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((*a.getODSResults(0).begin()).getLoc() def test3 : Pat<(BOp $attr, (AOp:$a $input)), (BOp $attr, (AOp $input), (location $a))>; diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp --- a/mlir/tools/mlir-tblgen/RewriterGen.cpp +++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp @@ -83,9 +83,10 @@ 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. @@ -379,7 +380,7 @@ // Next handle DAG leaf: operand or attribute if (opArg.is()) { // emitOperandMatch's argument indexing counts attributes. - emitOperandMatch(tree, castedName, i, depth); + emitOperandMatch(tree, castedName, i, nextOperand, depth); ++nextOperand; } else if (opArg.is()) { emitAttributeMatch(tree, opName, i, depth); @@ -393,7 +394,8 @@ } 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(); auto matcher = tree.getArgAsLeaf(argIndex); @@ -418,7 +420,7 @@ PrintFatalError(loc, error); } auto self = formatv("(*{0}.getODSOperands({1}).begin()).getType()", - opName, argIndex); + opName, operandIndex); emitMatchCheck( opName, tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),