[mlir-tblgen] Support either in Tablegen DRR.

Add a new directive `either` to specify the operands can be matched in either order

Reviewed By: jpienaar, Mogball

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

Author: ChiaHungDuan

mlir/docs/DeclarativeRewrites.md

@@ -774,6 +774,23 @@ Explicitly-specified return types will take precedence over return types
 inferred from op traits or user-defined builders. The return types of values
 replacing root op results cannot be overridden.
 
+### `either`
+
+The `either` directive is used to specify the operands may be matched in either
+order.
+
+```tablegen
+def : Pat<(TwoArgOp (either $firstArg, (AnOp $secondArg))),
+          (...)>;
+```
+
+The above pattern will accept either `"test.TwoArgOp"(%I32Arg, %AnOpArg)` and
+`"test.TwoArgOp"(%AnOpArg, %I32Arg)`.
+
+Only operand is supported with `either` and note that an operation with
+`Commutative` trait doesn't imply that it'll have the same behavior than
+`either` while pattern matching.
+
 ## Debugging Tips
 
 ### Run `mlir-tblgen` to see the generated content

mlir/include/mlir/IR/OpBase.td

@@ -2730,6 +2730,21 @@ def location;
 
 def returnType;
 
+// Directive used to specify the operands may be matched in either order. When
+// two adjacents are marked with `either`, it'll try to match the operands in
+// either ordering of constraints. Example:
+//
+// ```
+// (TwoArgOp (either $firstArg, (AnOp $secondArg)))
+// ```
+// The above pattern will accept either `"test.TwoArgOp"(%I32Arg, %AnOpArg)` and
+// `"test.TwoArgOp"(%AnOpArg, %I32Arg)`.
+//
+// Only operand is supported with `either` and note that an operation with
+// `Commutative` trait doesn't imply that it'll have the same behavior than
+// `either` while pattern matching.
+def either;
+
 //===----------------------------------------------------------------------===//
 // Attribute and Type generation
 //===----------------------------------------------------------------------===//

mlir/include/mlir/TableGen/Pattern.h

@@ -186,6 +186,9 @@ class DagNode {
   // Returns true if this DAG node is wrapping native code call.
   bool isNativeCodeCall() const;
 
+  // Returns whether this DAG is an `either` specifier.
+  bool isEither() const;
+
   // Returns true if this DAG node is an operation.
   bool isOperation() const;
 

mlir/lib/TableGen/Pattern.cpp

@@ -113,7 +113,7 @@ bool DagNode::isNativeCodeCall() const {
 
 bool DagNode::isOperation() const {
   return !isNativeCodeCall() && !isReplaceWithValue() &&
-         !isLocationDirective() && !isReturnTypeDirective();
+         !isLocationDirective() && !isReturnTypeDirective() && !isEither();
 }
 
 llvm::StringRef DagNode::getNativeCodeTemplate() const {
@@ -142,7 +142,9 @@ Operator &DagNode::getDialectOp(RecordOperatorMap *mapper) const {
 }
 
 int DagNode::getNumOps() const {
-  int count = isReplaceWithValue() ? 0 : 1;
+  // We want to get number of operations recursively involved in the DAG tree.
+  // All other directives should be excluded.
+  int count = isOperation() ? 1 : 0;
   for (int i = 0, e = getNumArgs(); i != e; ++i) {
     if (auto child = getArgAsNestedDag(i))
       count += child.getNumOps();
@@ -184,6 +186,11 @@ bool DagNode::isReturnTypeDirective() const {
   return dagOpDef->getName() == "returnType";
 }
 
+bool DagNode::isEither() const {
+  auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
+  return dagOpDef->getName() == "either";
+}
+
 void DagNode::print(raw_ostream &os) const {
   if (node)
     node->print(os);
@@ -764,22 +771,25 @@ void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
   if (tree.isOperation()) {
     auto &op = getDialectOp(tree);
     auto numOpArgs = op.getNumArgs();
+    int numEither = 0;
 
-    // The pattern might have trailing directives.
+    // We need to exclude the trailing directives and `either` directive groups
+    // two operands of the operation.
     int numDirectives = 0;
     for (int i = numTreeArgs - 1; i >= 0; --i) {
       if (auto dagArg = tree.getArgAsNestedDag(i)) {
         if (dagArg.isLocationDirective() || dagArg.isReturnTypeDirective())
           ++numDirectives;
-        else
-          break;
+        else if (dagArg.isEither())
+          ++numEither;
       }
     }
 
-    if (numOpArgs != numTreeArgs - numDirectives) {
-      auto err = formatv("op '{0}' argument number mismatch: "
-                         "{1} in pattern vs. {2} in definition",
-                         op.getOperationName(), numTreeArgs, numOpArgs);
+    if (numOpArgs != numTreeArgs - numDirectives + numEither) {
+      auto err =
+          formatv("op '{0}' argument number mismatch: "
+                  "{1} in pattern vs. {2} in definition",
+                  op.getOperationName(), numTreeArgs + numEither, numOpArgs);
       PrintFatalError(&def, err);
     }
 
@@ -791,10 +801,30 @@ void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
       verifyBind(infoMap.bindOpResult(treeName, op), treeName);
     }
 
-    for (int i = 0; i != numTreeArgs; ++i) {
+    // The operand in `either` DAG should be bound to the operation in the
+    // parent DagNode.
+    auto collectSymbolInEither = [&](DagNode parent, DagNode tree,
+                                     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 != "_")
+            verifyBind(infoMap.bindOpArgument(parent, argName, op, opArgIdx),
+                       argName);
+        }
+      }
+    };
+
+    for (int i = 0, opArgIdx = 0; i != numTreeArgs; ++i, ++opArgIdx) {
       if (auto treeArg = tree.getArgAsNestedDag(i)) {
-        // This DAG node argument is a DAG node itself. Go inside recursively.
-        collectBoundSymbols(treeArg, infoMap, isSrcPattern);
+        if (treeArg.isEither()) {
+          collectSymbolInEither(tree, treeArg, opArgIdx);
+        } else {
+          // This DAG node argument is a DAG node itself. Go inside recursively.
+          collectBoundSymbols(treeArg, infoMap, isSrcPattern);
+        }
         continue;
       }
 
@@ -806,7 +836,7 @@ void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
         if (!treeArgName.empty() && treeArgName != "_") {
           LLVM_DEBUG(llvm::dbgs() << "found symbol bound to op argument: "
                                   << treeArgName << '\n');
-          verifyBind(infoMap.bindOpArgument(tree, treeArgName, op, i),
+          verifyBind(infoMap.bindOpArgument(tree, treeArgName, op, opArgIdx),
                      treeArgName);
         }
       }

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

@@ -1328,6 +1328,30 @@ def : Pat<(OneI32ResultOp),
               (replaceWithValue $results__2),
               ConstantAttr<I32Attr, "2">)>;
 
+//===----------------------------------------------------------------------===//
+// 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 results = (outs I32:$output);
+}
+
+def : Pat<(TestEitherOpA (either I32:$arg1, I16:$arg2), $_),
+          (TestEitherOpB $arg2)>;
+
+def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1), I16:$arg2), $_),
+          (TestEitherOpB $arg2)>;
+
+def : Pat<(TestEitherOpA (either (TestEitherOpB I32:$arg1),
+                                 (TestEitherOpB I16:$arg2)),
+                          $_),
+          (TestEitherOpB $arg2)>;
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Location)
 

mlir/test/mlir-tblgen/pattern.mlir

@@ -531,6 +531,40 @@ func @redundantTest(%arg0: i32) -> i32 {
   return %0 : i32
 }
 
+//===----------------------------------------------------------------------===//
+// Test either directive
+//===----------------------------------------------------------------------===//
+
+// CHECK: @either_dag_leaf_only
+func @either_dag_leaf_only_1(%arg0 : i32, %arg1 : i16, %arg2 : i8) -> () {
+  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  %0 = "test.either_op_a"(%arg0, %arg1, %arg2) : (i32, i16, i8) -> i32
+  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  %1 = "test.either_op_a"(%arg1, %arg0, %arg2) : (i16, i32, i8) -> i32
+  return
+}
+
+// CHECK: @either_dag_leaf_dag_node
+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
+  %1 = "test.either_op_a"(%0, %arg1, %arg2) : (i32, i16, i8) -> i32
+  // CHECK: "test.either_op_b"(%arg1) : (i16) -> i32
+  %2 = "test.either_op_a"(%arg1, %0, %arg2) : (i16, i32, i8) -> i32
+  return
+}
+
+// CHECK: @either_dag_node_dag_node
+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
+  %2 = "test.either_op_a"(%0, %1, %arg2) : (i32, i32, i8) -> i32
+  // CHECK: "test.either_op_b"(%arg1) : (i16) -> 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.
 //===----------------------------------------------------------------------===//