[mlir][tblgen] Fix emitting wrong index for `either` directive.

[jpienaar]

Good catch!

mlir/lib/TableGen/Pattern.cpp

if (!treeName.empty()) {
       LLVM_DEBUG(llvm::dbgs()
                  << "found symbol bound to op result: " << treeName << '\n');
       verifyBind(infoMap.bindOpResult(treeName, op), treeName);
     }
 
     // The operand in `either` DAG should be bound to the operation in the
     // parent DagNode.
     auto collectSymbolInEither = [&](DagNode parent, DagNode tree,
-                                     int &opArgIdx) {
+                                     int opArgIdx) {
       for (int i = 0; i < tree.getNumArgs(); ++i, ++opArgIdx) {
         if (DagNode subTree = tree.getArgAsNestedDag(i)) {
           collectBoundSymbols(subTree, infoMap, isSrcPattern);
         } else {
           auto argName = tree.getArgName(i);
-          if (!argName.empty() && argName != "_")
+          if (!argName.empty() && argName != "_") {
             verifyBind(infoMap.bindOpArgument(parent, argName, op, opArgIdx),
                        argName);
           }
         }
+      }
     };
 
     for (int i = 0, opArgIdx = 0; i != numTreeArgs; ++i, ++opArgIdx) {
       if (auto treeArg = tree.getArgAsNestedDag(i)) {
         if (treeArg.isEither()) {
           collectSymbolInEither(tree, treeArg, opArgIdx);
+          // `either` DAG is *flattened*. For example,
+          //
+          //  (FooOp (either arg0, arg1), arg2)
+          //
+          //  can be viewed as:
+          //
+          //  (FooOp arg0, arg1, arg2)
+          ++opArgIdx;
         } else {
           // This DAG node argument is a DAG node itself. Go inside recursively.
           collectBoundSymbols(treeArg, infoMap, isSrcPattern);
         }
         continue;
       }
 
       if (isSrcPattern) {

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

 // Test Patterns (either)
 
 def TestEitherOpA : TEST_Op<"either_op_a"> {
   let arguments = (ins AnyInteger:$arg0, AnyInteger:$arg1, AnyInteger:$arg2);
   let results = (outs I32:$output);
 }
 
 def TestEitherOpB : TEST_Op<"either_op_b"> {
-  let arguments = (ins AnyInteger:$arg0);
+  let arguments = (ins AnyInteger:$arg0, AnyInteger:$arg1);
   let results = (outs I32:$output);
 }
 
-def : Pat<(TestEitherOpA (either I32:$arg1, I16:$arg2), $_),
-          (TestEitherOpB $arg2)>;
+def : Pat<(TestEitherOpA (either I32:$arg1, I16:$arg2), $x),
+          (TestEitherOpB $arg2, $x)>;
 
-def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1), I16:$arg2), $_),
-          (TestEitherOpB $arg2)>;
+def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1, $_), I16:$arg2), $x),
+          (TestEitherOpB $arg2, $x)>;
 
-def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1),
-                                 (TestEitherOpB I16:$arg2)),
-                          $_),
-          (TestEitherOpB $arg2)>;
+def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1, $_),
+                                 (TestEitherOpB I16:$arg2, $_)),
+                          $x),
+          (TestEitherOpB $arg2, $x)>;
 
 //===----------------------------------------------------------------------===//
 // Test Patterns (Location)
 
 // Test that we can specify locations for generated ops.
 def : Pat<(TestLocationSrcOp:$res1
            (TestLocationSrcOp:$res2
             (TestLocationSrcOp:$res3 $input))),

mlir/test/mlir-tblgen/pattern.mlir

 }
 
 //===----------------------------------------------------------------------===//
 // Test either directive
 //===----------------------------------------------------------------------===//
 
 // CHECK: @either_dag_leaf_only
 func.func @either_dag_leaf_only_1(%arg0 : i32, %arg1 : i16, %arg2 : i8) -> () {
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %0 = "test.either_op_a"(%arg0, %arg1, %arg2) : (i32, i16, i8) -> i32
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %1 = "test.either_op_a"(%arg1, %arg0, %arg2) : (i16, i32, i8) -> i32
   return
 }
 
 // CHECK: @either_dag_leaf_dag_node
 func.func @either_dag_leaf_dag_node(%arg0 : i32, %arg1 : i16, %arg2 : i8) -> () {
-  %0 = "test.either_op_b"(%arg0) : (i32) -> i32
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  %0 = "test.either_op_b"(%arg0, %arg0) : (i32, i32) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %1 = "test.either_op_a"(%0, %arg1, %arg2) : (i32, i16, i8) -> i32
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %2 = "test.either_op_a"(%arg1, %0, %arg2) : (i16, i32, i8) -> i32
   return
 }
 
 // CHECK: @either_dag_node_dag_node
 func.func @either_dag_node_dag_node(%arg0 : i32, %arg1 : i16, %arg2 : i8) -> () {
-  %0 = "test.either_op_b"(%arg0) : (i32) -> i32
-  %1 = "test.either_op_b"(%arg1) : (i16) -> i32
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  %0 = "test.either_op_b"(%arg0, %arg0) : (i32, i32) -> i32
+  %1 = "test.either_op_b"(%arg1, %arg1) : (i16, i16) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %2 = "test.either_op_a"(%0, %1, %arg2) : (i32, i32, i8) -> i32
-  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  // CHECK: "test.either_op_b"(%arg1, %arg2) : (i16, i8) -> i32
   %3 = "test.either_op_a"(%1, %0, %arg2) : (i32, i32, i8) -> i32
   return
 }
 
 //===----------------------------------------------------------------------===//
 // Test that ops without type deduction can be created with type builders.
 //===----------------------------------------------------------------------===//
 

mlir/tools/mlir-tblgen/RewriterGen.cpp

emitMatchCheck(opName, /*matchStr=*/castedName,
                    formatv("\"{0} is not {1} type\"", castedName,
                            op.getQualCppClassName()));
 
   // 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(), nextOperand = 0; i != e; ++i) {
-    auto opArg = op.getArg(i);
+  for (int i = 0, opArgIdx = 0, e = tree.getNumArgs(), nextOperand = 0; i != e;
+       ++i, ++opArgIdx) {
+    auto opArg = op.getArg(opArgIdx);
     std::string argName = formatv("op{0}", depth + 1);
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (argTree.isEither()) {
-        emitEitherOperandMatch(tree, argTree, castedName, i, nextOperand,
+        emitEitherOperandMatch(tree, argTree, castedName, opArgIdx, nextOperand,
                                depth);
+        ++opArgIdx;
         continue;
       }
       if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
         if (operand->isVariableLength()) {
           auto error = formatv("use nested DAG construct to match op {0}'s "
                                "variadic operand #{1} unsupported now",
-                               op.getOperationName(), i);
+                               op.getOperationName(), opArgIdx);
           PrintFatalError(loc, error);
         }
       }
 
       os << "{\n";
 
       // Attributes don't count for getODSOperands.
       // TODO: Operand is a Value, check if we should remove `getDefiningOp()`.
for (int i = 0, opArgIdx = 0, e = tree.getNumArgs(), nextOperand = 0; i != e;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
       auto operandName =
           formatv("{0}.getODSOperands({1})", castedName, nextOperand);
       emitOperandMatch(tree, castedName, operandName.str(),
                        /*operandMatcher=*/tree.getArgAsLeaf(i),
-                       /*argName=*/tree.getArgName(i),
-                       /*argIndex=*/i);
+                       /*argName=*/tree.getArgName(i), opArgIdx);
       ++nextOperand;
     } else if (opArg.is<NamedAttribute *>()) {
-      emitAttributeMatch(tree, opName, i, depth);
+      emitAttributeMatch(tree, opName, opArgIdx, depth);
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 }