Fix MLIR DRR matching when attributes are interleaved with operands

The ODSOperand indexing should ignore the attributes.

Differential Revision: https://reviews.llvm.org/D94281

diff --git a/mlir/test/mlir-tblgen/rewriter-indexing.td b/mlir/test/mlir-tblgen/rewriter-indexing.td
index c21b04f..a6b4032 100644 (file)
--- a/mlir/test/mlir-tblgen/rewriter-indexing.td
+++ b/mlir/test/mlir-tblgen/rewriter-indexing.td
@@ -51,6 +51,9 @@ def test2 : Pat<(COp $attr1, $op1, $attr2, (AOp $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))>;

diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
index 9fca15b..1a66517 100644 (file)
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -83,9 +83,10 @@ private:
   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 @@ void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) {
     // 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);
@@ -393,7 +394,8 @@ void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) {
 }
 
 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);
@@ -418,7 +420,7 @@ void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,
         PrintFatalError(loc, error);
       }
       auto self = formatv("(*{0}.getODSOperands({1}).begin()).getType()",
-                          opName, argIndex);
+                          opName, operandIndex);
       emitMatchCheck(
           opName,
           tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),