[mlir] Enable DRR variadic operand matching

[logan]

Summary

This commit enables DRR rewriter to match a fixed number of sub-operands as a variadic operand:

def : Pat<(ConcatenateOp (variadic $input0, $input1)),
          (SomeOtherOp $input0, $input1)>;

[Chia-hungDuan]

LGTM, given that I'm not currently working on MLIR. Will leave @jpienaar to make the decision

mlir/include/mlir/IR/PatternBase.td

219

space?

mlir/lib/TableGen/Pattern.cpp

304–305

Could you update the comment a little bit?

305–308

I'm thinking if we can hide the -1 here and there. Like if we do

dagAndConstant->hasRangeIndex()

For the function's default argument

bindOpArgument(..., int rangeIndex = dagAndConstant::kInvalidRangeIndex) { ...}

It may look better.

mlir/tools/mlir-tblgen/RewriterGen.cpp

714

Nit: single line statement

[jpienaar]

Nice, so this is a more refined variant of https://reviews.llvm.org/D121046 that enables more general support?

[logan]

In D157359#4580594, @jpienaar wrote:

Nice, so this is a more refined variant of https://reviews.llvm.org/D121046 that enables more general support?

Kind of. This enables the sub-dag pattern matching into a variadic operand and checks HasExactValues<n> automatically. It also avoids the need of Front and DropFront sequences.

However, the limitation of this approach is that I only match variadic operand with exactly N sub-trees. I don't support "N sub-tree + tails". Technically doable by defining a (variantFrontNAndTails $sub0, $sub1, ..., $sub_n_m1, $tails) where $tails matches the remaining range. But to keep this simple, I didn't implement variantFrontNAndTails.

[logan]

Changes in Revision 2:

1. Allow binding a symbol to variadic DAG node, which means the full operand_range
2. Fix staticMatcher problem rooted at variadic DAG node. Static matcher must not start from variadic node
3. Introduces kInvalidRangeIndex and replace relevant -1 with kInvalidRangeIndex.
4. Add comments to describe the behavior of DagAndConstant.

[logan]

Changes in Revision 3:

1. Update some comments to clarify how we map SymbolInfo to the variable name (depending on whether SymbolInfo has a range index).

[logan]

mlir/include/mlir/IR/PatternBase.td

219

No. This character is intended to be colon :. Like the symbol binding to the dag result, this symbol binding refers to the full operand_range.

Revision 1 forgot to implement this feature. I added them in Revision 2 and added more description on Line 224-227.

[Chia-hungDuan]

Still LGTM! The DagAndConstant looks better now!

[jpienaar]

Does variadic also work in the target dag?

mlir/include/mlir/IR/PatternBase.td

226

Just to avoid folks thinking there is a convention lets call it input_a and input_b (given we have the __n convention elsewhere that is load bearing).

mlir/include/mlir/TableGen/Pattern.h

257

Why is this one public? Looking at usage, it feels more like optional<int64_t>. Range is also confusing as I'd expect 2 integers for a range. (looking below it almost feels like this is as if index is an array with at most 2 elements)

313

Does this mean only one level of nesting is allowed here? E.g., if there was a variadic of variadic what would this be?

332

Why 3?

[logan]

Does variadic also work in the target dag?

No, this commit doesn't support that. I couldn't figure it out at this moment.

mlir/include/mlir/TableGen/Pattern.h

257

I am not a fan of std::optional because it takes an extra bool. But I don't have strong feeling on this, I updated the change anyway.

313

Correct. I don't support nested variadic DAG. mlir/tools/mlir-tblgen/RewriterGen.cpp Line 612 emit an error if we detect a variadic DAG.

def : Pat<(SomeOp (variadic $a, $b, (variadic $c, $d))), ...>;  // error

332

I didn't write the op definition, but I tried to say OpWithThreeOutput has 3 outputs. Since it is a multi-result value, DagAndConstant keep tracks of the number of outputs.

[logan]

Address comments.

[logan]

Thanks. Landed as rG08d7377b67358496a409080fac22f3f7c077fb63.

mlir/docs/DeclarativeRewrites.md

 We say the above op has 3 declared operands and 3 declared results. But at
 runtime, an instance can have 3 values corresponding to `$input2` and 2 values
 correspond to `$output2`; we say it has 5 actual operands and 4 actual results.
 A variadic operand/result is a considered as a declared value that can
 correspond to multiple actual values.
 
 [TODO]
 
+#### Match variadic operand
+
+Use the `variadic` DAG node to match a variadic operand with a fixed number of
+actual sub-operands.
+
+For example, assume that `ConcatenateOp` is an operation with a variadic
+operand:
+
+```tablegen
+def ConcatenateOp : TEST_Op<"concatenate"> {
+  let arguments = (ins
+    Variadic<AnyTensor>:$inputs,
+    I32Attr:$axis
+  );
+
+  let results = (outs
+    AnyTensor$output
+  );
+}
+```
+
+We can match `ConcatenateOp` with exactly 2 actual operands with:
+
+```tablegen
+def : Pat<(ConcatenateOp (variadic $input0, $input1), $axis),
+          ...>;
+```
+
+The variadic sub-operands can be sub-DAGs to be matched:
+
+```tablegen
+def : Pat<(ConcatenateOp (variadic (SomeOp $a), (AnotherOp $b, $c)), $axis),
+          (OtherOp $a, $b, $c)>;
+```
+
+The variadic DAG can be bound to a symbol, which refers to the full
+`operand_range`:
+
+```tablegen
+def : Pat<(ConcatenateOp (variadic:$inputs $input0, $input1),
+                         ConstantAttr<I32Attr, "0">),
+          (VStackOp $inputs)>;
+```
+
 ### Supplying additional constraints
 
 Constraints can be placed on op arguments when matching. But sometimes we need
 to also place constraints on the matched op's results or sometimes need to limit
 the matching with some constraints that cover both the arguments and the
 results. The third parameter to `Pattern` (and `Pat`) is for this purpose.
 
 For example, we can write

mlir/include/mlir/IR/PatternBase.td

 // 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;
 
+// Directive used to match variadic operands. This directive only matches if
+// the variadic operand has the same length as the specified formal
+// sub-dags.
+//
+// ```
+// (VariadicOp (variadic:$input1 $input1a, $input1b),
+//             (variadic:$input2 $input2a, $input2b, $input2c),
+//             $attr1, $attr2)
+// ```
+//
+// The pattern above only matches if the `$input1` operand is of length 2,
+// `$input2` is of length 3, and all sub-dags match respectively. The `$input1`
+// symbol denotes the full variadic operand range. The `$input1a` symbol
+// denotes the first operand in the variadic sub-operands.
+def variadic;
+
 //===----------------------------------------------------------------------===//
 // Common value constraints
 //===----------------------------------------------------------------------===//
 
 def HasNoUseOf: Constraint<
     CPred<"$_self.use_empty()">, "has no use">;
 
 #endif // PATTERNBASE_TD

mlir/include/mlir/TableGen/Pattern.h

 #include "mlir/Support/LLVM.h"
 #include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Operator.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSet.h"
 
+#include <optional>
 #include <unordered_map>
 
 namespace llvm {
 class DagInit;
 class Init;
 class Record;
 } // namespace llvm
 
public:
   bool isReturnTypeDirective() const;
 
   // 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 whether this DAG is an `variadic` specifier.
+  bool isVariadic() const;
+
   // Returns true if this DAG node is an operation.
   bool isOperation() const;
 
   // Returns the native code call template inside this DAG node.
   // Precondition: isNativeCodeCall()
   StringRef getNativeCodeTemplate() const;
 
   // Returns the number of values will be returned by the native helper

 // values in a suitable way.
 class SymbolInfoMap {
 public:
   explicit SymbolInfoMap(ArrayRef<SMLoc> loc) : loc(loc) {}
 
   // Class for information regarding a symbol.
   class SymbolInfo {
   public:
     // Returns a type string of a variable.
     std::string getVarTypeStr(StringRef name) const;
 
     // Returns a string for defining a variable named as `name` to store the
     // value bound by this symbol.
     std::string getVarDecl(StringRef name) const;
 
     // Returns a string for defining an argument which passes the reference of
     // the variable.
     std::string getArgDecl(StringRef name) const;
 
     // Returns a variable name for the symbol named as `name`.
     std::string getVarName(StringRef name) const;
 
   private:
     // Allow SymbolInfoMap to access private methods.
     friend class SymbolInfoMap;
 
-    // DagNode and DagLeaf are accessed by value which means it can't be used as
-    // identifier here. Use an opaque pointer type instead.
-    using DagAndConstant = std::pair<const void *, int>;
+    // Structure to uniquely distinguish different locations of the symbols.
+    //
+    // * If a symbol is defined as an operand of an operation, `dag` specifies
+    //   the DAG of the operation, `operandIndexOrNumValues` specifies the
+    //   operand index, and `variadicSubIndex` must be set to `std::nullopt`.
+    //
+    // * If a symbol is defined in a `variadic` DAG, `dag` specifies the DAG
+    //   of the parent operation, `operandIndexOrNumValues` specifies the
+    //   declared operand index of the variadic operand in the parent
+    //   operation.
+    //
+    //   - If the symbol is defined as a result of `variadic` DAG, the
+    //     `variadicSubIndex` must be set to `std::nullopt`, which means that
+    //     the symbol binds to the full operand range.
+    //
+    //   - If the symbol is defined as a operand, the `variadicSubIndex` must
+    //     be set to the index within the variadic sub-operand list.
+    //
+    // * If a symbol is defined in a `either` DAG, `dag` specifies the DAG
+    //   of the parent operation, `operandIndexOrNumValues` specifies the
+    //   operand index in the parent operation (not necessary the index in the
+    //   DAG).
+    //
+    // * If a symbol is defined as a result, specifies the number of returning
+    //   value.
+    //
+    // Example 1:
+    //
+    //   def : Pat<(OpA $input0, $input1), ...>;
+    //
+    //   $input0: (OpA, 0, nullopt)
+    //   $input1: (OpA, 1, nullopt)
+    //
+    // Example 2:
+    //
+    //   def : Pat<(OpB (variadic:$input0 $input0a, $input0b),
+    //                  (variadic:$input1 $input1a, $input1b, $input1c)),
+    //             ...>;
+    //
+    //   $input0:  (OpB, 0, nullopt)
+    //   $input0a: (OpB, 0, 0)
+    //   $input0b: (OpB, 0, 1)
+    //   $input1:  (OpB, 1, nullopt)
+    //   $input1a: (OpB, 1, 0)
+    //   $input1b: (OpB, 1, 1)
+    //   $input1c: (OpB, 1, 2)
+    //
+    // Example 3:
+    //
+    //   def : Pat<(OpC $input0, (either $input1, $input2)), ...>;
+    //
+    //   $input0: (OpC, 0, nullopt)
+    //   $input1: (OpC, 1, nullopt)
+    //   $input2: (OpC, 2, nullopt)
+    //
+    // Example 4:
+    //
+    //   def ThreeResultOp : TEST_Op<...> {
+    //     let results = (outs
+    //       AnyType:$result1,
+    //       AnyType:$result2,
+    //       AnyType:$result3
+    //     );
+    //   }
+    //
+    //   def : Pat<...,
+    //             (ThreeResultOp:$result ...)>;
+    //
+    //   $result: (nullptr, 3, nullopt)
+    //
+    struct DagAndConstant {
+      // DagNode and DagLeaf are accessed by value which means it can't be used
+      // as identifier here. Use an opaque pointer type instead.
+      const void *dag;
+      int operandIndexOrNumValues;
+      std::optional<int> variadicSubIndex;
+
+      DagAndConstant(const void *dag, int operandIndexOrNumValues,
+                     std::optional<int> variadicSubIndex)
+          : dag(dag), operandIndexOrNumValues(operandIndexOrNumValues),
+            variadicSubIndex(variadicSubIndex) {}
+
+      bool operator==(const DagAndConstant &rhs) const {
+        return dag == rhs.dag &&
+               operandIndexOrNumValues == rhs.operandIndexOrNumValues &&
+               variadicSubIndex == rhs.variadicSubIndex;
+      }
+    };
 
     // What kind of entity this symbol represents:
     // * Attr: op attribute
     // * Operand: op operand
     // * Result: op result
     // * Value: a value not attached to an op (e.g., from NativeCodeCall)
     // * MultipleValues: a pack of values not attached to an op (e.g., from
     //   NativeCodeCall). This kind supports indexing.
     enum class Kind : uint8_t { Attr, Operand, Result, Value, MultipleValues };
 
     // Creates a SymbolInfo instance. `dagAndConstant` is only used for `Attr`
     // and `Operand` so should be std::nullopt for `Result` and `Value` kind.
     SymbolInfo(const Operator *op, Kind kind,
                std::optional<DagAndConstant> dagAndConstant);
 
     // Static methods for creating SymbolInfo.
     static SymbolInfo getAttr(const Operator *op, int index) {
-      return SymbolInfo(op, Kind::Attr, DagAndConstant(nullptr, index));
+      return SymbolInfo(op, Kind::Attr,
+                        DagAndConstant(nullptr, index, std::nullopt));
     }
     static SymbolInfo getAttr() {
       return SymbolInfo(nullptr, Kind::Attr, std::nullopt);
     }
-    static SymbolInfo getOperand(DagNode node, const Operator *op, int index) {
+    static SymbolInfo
+    getOperand(DagNode node, const Operator *op, int operandIndex,
+               std::optional<int> variadicSubIndex = std::nullopt) {
       return SymbolInfo(op, Kind::Operand,
-                        DagAndConstant(node.getAsOpaquePointer(), index));
+                        DagAndConstant(node.getAsOpaquePointer(), operandIndex,
+                                       variadicSubIndex));
     }
     static SymbolInfo getResult(const Operator *op) {
       return SymbolInfo(op, Kind::Result, std::nullopt);
     }
     static SymbolInfo getValue() {
       return SymbolInfo(nullptr, Kind::Value, std::nullopt);
     }
     static SymbolInfo getMultipleValues(int numValues) {
       return SymbolInfo(nullptr, Kind::MultipleValues,
-                        DagAndConstant(nullptr, numValues));
+                        DagAndConstant(nullptr, numValues, std::nullopt));
     }
 
     // Returns the number of static values this symbol corresponds to.
     // A static value is an operand/result declared in ODS. Normally a symbol
     // only represents one static value, but symbols bound to op results can
     // represent more than one if the op is a multi-result op.
     int getStaticValueCount() const;
 
private:
     // represents. `name` is the name of the C++ variable that this symbol
     // bounds to. `index` should only be used for indexing results. `fmt` is
     // used to format each value. `separator` is used to separate values in the
     // range.
     std::string getAllRangeUse(StringRef name, int index, const char *fmt,
                                const char *separator) const;
 
     // The argument index (for `Attr` and `Operand` only)
-    int getArgIndex() const { return (*dagAndConstant).second; }
+    int getArgIndex() const { return dagAndConstant->operandIndexOrNumValues; }
 
     // The number of values in the MultipleValue
-    int getSize() const { return (*dagAndConstant).second; }
+    int getSize() const { return dagAndConstant->operandIndexOrNumValues; }
+
+    // The variadic sub-operands index (for variadic `Operand` only)
+    std::optional<int> getVariadicSubIndex() const {
+      return dagAndConstant->variadicSubIndex;
+    }
 
     const Operator *op; // The op where the bound entity belongs
     Kind kind;          // The kind of the bound entity
 
-    // The pair of DagNode pointer and constant value (for `Attr`, `Operand` and
-    // the size of MultipleValue symbol). Note that operands may be bound to the
-    // same symbol, use the DagNode and index to distinguish them. For `Attr`
-    // and MultipleValue, the Dag part will be nullptr.
+    // The tuple of DagNode pointer and two constant values (for `Attr`,
+    // `Operand` and the size of MultipleValue symbol). Note that operands may
+    // be bound to the same symbol, use the DagNode and index to distinguish
+    // them. For `Attr` and MultipleValue, the Dag part will be nullptr.
     std::optional<DagAndConstant> dagAndConstant;
 
     // Alternative name for the symbol. It is used in case the name
     // is not unique. Applicable for `Operand` only.
     std::optional<std::string> alternativeName;
   };
 
   using BaseT = std::unordered_multimap<std::string, SymbolInfo>;
 
   // Iterators for accessing all symbols.
   using iterator = BaseT::iterator;
   iterator begin() { return symbolInfoMap.begin(); }
   iterator end() { return symbolInfoMap.end(); }
 
   // Const iterators for accessing all symbols.
   using const_iterator = BaseT::const_iterator;
   const_iterator begin() const { return symbolInfoMap.begin(); }
   const_iterator end() const { return symbolInfoMap.end(); }
 
   // Binds the given `symbol` to the `argIndex`-th argument to the given `op`.
   // Returns false if `symbol` is already bound and symbols are not operands.
   bool bindOpArgument(DagNode node, StringRef symbol, const Operator &op,
-                      int argIndex);
+                      int argIndex,
+                      std::optional<int> variadicSubIndex = std::nullopt);
 
   // Binds the given `symbol` to the results the given `op`. Returns false if
   // `symbol` is already bound.
   bool bindOpResult(StringRef symbol, const Operator &op);
 
   // A helper function for dispatching target value binding functions.
   bool bindValues(StringRef symbol, int numValues = 1);
 
public:
   bool contains(StringRef symbol) const;
 
   // Returns an iterator to the information of the given symbol named as `key`.
   const_iterator find(StringRef key) const;
 
   // Returns an iterator to the information of the given symbol named as `key`,
   // with index `argIndex` for operator `op`.
   const_iterator findBoundSymbol(StringRef key, DagNode node,
-                                 const Operator &op, int argIndex) const;
+                                 const Operator &op, int argIndex,
+                                 std::optional<int> variadicSubIndex) const;
   const_iterator findBoundSymbol(StringRef key,
                                  const SymbolInfo &symbolInfo) const;
 
   // Returns the bounds of a range that includes all the elements which
   // bind to the `key`.
   std::pair<iterator, iterator> getRangeOfEqualElements(StringRef key);
 
   // Returns number of times symbol named as `key` was used.

mlir/lib/TableGen/Pattern.cpp

 bool DagNode::isNativeCodeCall() const {
   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(node->getOperator()))
     return defInit->getDef()->isSubClassOf("NativeCodeCall");
   return false;
 }
 
 bool DagNode::isOperation() const {
   return !isNativeCodeCall() && !isReplaceWithValue() &&
-         !isLocationDirective() && !isReturnTypeDirective() && !isEither();
+         !isLocationDirective() && !isReturnTypeDirective() && !isEither() &&
+         !isVariadic();
 }
 
 llvm::StringRef DagNode::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsString("expression");
 }
bool DagNode::isReturnTypeDirective() const {
   return dagOpDef->getName() == "returnType";
 }
 
 bool DagNode::isEither() const {
   auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
   return dagOpDef->getName() == "either";
 }
 
+bool DagNode::isVariadic() const {
+  auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
+  return dagOpDef->getName() == "variadic";
+}
+
 void DagNode::print(raw_ostream &os) const {
   if (node)
     node->print(os);
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolInfoMap
 //===----------------------------------------------------------------------===//
std::string SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
   case Kind::Attr: {
     assert(index < 0);
     auto repl = formatv(fmt, name);
     LLVM_DEBUG(llvm::dbgs() << repl << " (Attr)\n");
     return std::string(repl);
   }
   case Kind::Operand: {
     assert(index < 0);
     auto *operand = op->getArg(getArgIndex()).get<NamedTypeConstraint *>();
-    // If this operand is variadic, then return a range. Otherwise, return the
-    // value itself.
-    if (operand->isVariableLength()) {
+    // If this operand is variadic and this SymbolInfo doesn't have a range
+    // index, then return the full variadic operand_range. Otherwise, return
+    // the value itself.
+    if (operand->isVariableLength() && !getVariadicSubIndex().has_value()) {
       auto repl = formatv(fmt, name);
       LLVM_DEBUG(llvm::dbgs() << repl << " (VariadicOperand)\n");
       return std::string(repl);
     }
     auto repl = formatv(fmt, formatv("(*{0}.begin())", name));
     LLVM_DEBUG(llvm::dbgs() << repl << " (SingleOperand)\n");
     return std::string(repl);
   }
case Kind::MultipleValues: {
     LLVM_DEBUG(llvm::dbgs() << repl << " (MultipleValues)\n");
     return std::string(repl);
   }
   }
   llvm_unreachable("unknown kind");
 }
 
 bool SymbolInfoMap::bindOpArgument(DagNode node, StringRef symbol,
-                                   const Operator &op, int argIndex) {
+                                   const Operator &op, int argIndex,
+                                   std::optional<int> variadicSubIndex) {
   StringRef name = getValuePackName(symbol);
   if (name != symbol) {
     auto error = formatv(
         "symbol '{0}' with trailing index cannot bind to op argument", symbol);
     PrintFatalError(loc, error);
   }
 
-  auto symInfo = op.getArg(argIndex).is<NamedAttribute *>()
-                     ? SymbolInfo::getAttr(&op, argIndex)
-                     : SymbolInfo::getOperand(node, &op, argIndex);
+  auto symInfo =
+      op.getArg(argIndex).is<NamedAttribute *>()
+          : SymbolInfo::getOperand(node, &op, argIndex, variadicSubIndex);
 
   std::string key = symbol.str();
   if (symbolInfoMap.count(key)) {
     // Only non unique name for the operand is supported.
     if (symInfo.kind != SymbolInfo::Kind::Operand) {
       return false;
     }
 

 SymbolInfoMap::const_iterator SymbolInfoMap::find(StringRef key) const {
   std::string name = getValuePackName(key).str();
 
   return symbolInfoMap.find(name);
 }
 
 SymbolInfoMap::const_iterator
 SymbolInfoMap::findBoundSymbol(StringRef key, DagNode node, const Operator &op,
-                               int argIndex) const {
-  return findBoundSymbol(key, SymbolInfo::getOperand(node, &op, argIndex));
+                               int argIndex,
+                               std::optional<int> variadicSubIndex) const {
+  return findBoundSymbol(
+      key, SymbolInfo::getOperand(node, &op, argIndex, variadicSubIndex));
 }
 
 SymbolInfoMap::const_iterator
 SymbolInfoMap::findBoundSymbol(StringRef key,
                                const SymbolInfo &symbolInfo) const {
   std::string name = getValuePackName(key).str();
   auto range = symbolInfoMap.equal_range(name);
 
auto collectSymbolInEither = [&](DagNode parent, DagNode tree,
           if (!argName.empty() && argName != "_") {
             verifyBind(infoMap.bindOpArgument(parent, argName, op, opArgIdx),
                        argName);
           }
         }
       }
     };
 
+    // The operand in `variadic` DAG should be bound to the operation in the
+    // parent DagNode. The range index must be included as well to distinguish
+    // (potentially) repeating argName within the `variadic` DAG.
+    auto collectSymbolInVariadic = [&](DagNode parent, DagNode tree,
+                                       int opArgIdx) {
+      auto treeName = tree.getSymbol();
+      if (!treeName.empty()) {
+        // If treeName is specified, bind to the full variadic operand_range.
+        verifyBind(infoMap.bindOpArgument(parent, treeName, op, opArgIdx,
+                                          std::nullopt),
+                   treeName);
+      }
+
+      for (int i = 0; i < tree.getNumArgs(); ++i) {
+        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,
+                                              /*variadicSubIndex=*/i),
+                       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 if (treeArg.isVariadic()) {
+          collectSymbolInVariadic(tree, treeArg, 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

 // Generates an op with variadic results using nested pattern.
 def : Pat<(OneI32ResultOp),
           (MixedVOperandOp3
               (MixedVResultOp3:$results__0 ConstantAttr<I32Attr, "2">),
               (replaceWithValue $results__1),
               (replaceWithValue $results__2),
               ConstantAttr<I32Attr, "2">)>;
 
+// Variadic structured matching
+def MixedVOperandOp4 : TEST_Op<"mixed_variadic_in4"> {
+  let arguments = (ins
+    Variadic<I32>:$input1,
+    I32:$input2,
+    I32Attr:$attr1
+  );
+}
+
+def MixedVOperandOp5 : TEST_Op<"mixed_variadic_in5"> {
+  let arguments = (ins
+    I32:$input1,
+    I32:$input2,
+    I32:$input3,
+    I32Attr:$attr1,
+    StrAttr:$pattern_name
+  );
+}
+
+// Helper op to test variadic recursive pattern matching
+def MixedVOperandInOutI32Op : TEST_Op<"mixed_variadic_in_out_i32"> {
+  let arguments = (ins
+    I32:$input
+  );
+  let results = (outs
+    I32:$output
+  );
+}
+
+def : Pat<
+  (MixedVOperandOp4 (variadic $input1a, $input1b), $input2,
+                    ConstantAttr<I32Attr, "0">:$attr1),
+  (MixedVOperandOp5 $input1a, $input1b, $input2, $attr1,
+                    ConstantStrAttr<StrAttr, "MatchVariadic">)>;
+
+def : Pat<
+  (MixedVOperandOp4 (variadic (MixedVOperandInOutI32Op $input1a),
+                              (MixedVOperandInOutI32Op $input1b)),
+                    $input2, ConstantAttr<I32Attr, "1">:$attr1),
+  (MixedVOperandOp5 $input1a, $input1b, $input2, $attr1,
+                    ConstantStrAttr<StrAttr, "MatchVariadicSubDag">)>;
+
+def : Pat<
+  (MixedVOperandOp4 (variadic $input1, $input1), $input2,
+                    ConstantAttr<I32Attr, "2">:$attr1),
+  (MixedVOperandOp5 $input1, $input1, $input2, $attr1,
+                    ConstantStrAttr<StrAttr, "MatchVariadicSameSymbol">)>;
+
+def MixedVOperandOp6 : TEST_Op<"mixed_variadic_in6",
+                               [SameVariadicOperandSize]> {
+  let arguments = (ins
+    Variadic<I32>:$input1,
+    Variadic<I32>:$input2,
+    I32Attr:$attr1
+  );
+}
+
+def : Pat<
+  (MixedVOperandOp6 (variadic:$input1 $input1a, $input1b),
+                    (variadic:$input2 $input2a, $input2b),
+                    ConstantAttr<I32Attr, "1">:$attr1),
+  (MixedVOperandOp6 $input2, $input1, ConstantAttr<I32Attr, "-1">)>;
+
+def : Pat<
+  (MixedVOperandOp6 (variadic $input1a, $input1b),
+                    (variadic $input2a, $input2b),
+                    ConstantAttr<I32Attr, "2">:$attr1),
+  (MixedVOperandOp5 $input2a, $input2b, $input1b, $attr1,
+                    ConstantStrAttr<StrAttr, "MatchMultiVariadicSubSymbol">)>;
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (either)
 
 def TestEitherOpA : TEST_Op<"either_op_a"> {
   let arguments = (ins AnyInteger:$arg0, AnyInteger:$arg1, AnyInteger:$arg2);
   let results = (outs I32:$output);
 }
 

mlir/test/mlir-tblgen/pattern.mlir

func.func @generateVariadicOutputOpInNestedPattern() -> (i32) {
   // CHECK: %[[cnt5_a:.*]], %[[cnt5_b:.*]]:2, %[[cnt5_c:.*]]:2 = "test.mixed_variadic_out3"()
   // CHECK: %[[res:.*]] = "test.mixed_variadic_in3"(%[[cnt5_a]], %[[cnt5_b]]#0, %[[cnt5_b]]#1, %[[cnt5_c]]#0, %[[cnt5_c]]#1)
   // CHECK: return %[[res]]
 
   %0 = "test.one_i32_out"() : () -> (i32)
   return %0 : i32
 }
 
+// CHECK-LABEL: @testMatchVariadic
+func.func @testMatchVariadic(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32) -> () {
+  // CHECK: "test.mixed_variadic_in5"(%arg0, %arg1, %arg2) <{attr1 = 0 : i32, pattern_name = "MatchVariadic"}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%arg0, %arg1, %arg2) {attr1 = 0 : i32} : (i32, i32, i32) -> ()
+
+  // Note: Not rewritten because variadic operand size mismatches.
+  // CHECK: "test.mixed_variadic_in4"(%arg0, %arg1, %arg2, %arg3) <{attr1 = 0 : i32}> : (i32, i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%arg0, %arg1, %arg2, %arg3) {attr1 = 0 : i32} : (i32, i32, i32, i32) -> ()
+
+  return
+}
+
+// CHECK-LABEL: @testMatchVariadicSubDag
+func.func @testMatchVariadicSubDag(%arg0: i32, %arg1: i32, %arg2: i32) -> () {
+  // CHECK: %[[IN0:.*]] = "test.mixed_variadic_in_out_i32"(%arg0) : (i32) -> i32
+  %0 = "test.mixed_variadic_in_out_i32"(%arg0) : (i32) -> i32
+  // CHECK: %[[IN1:.*]] = "test.mixed_variadic_in_out_i32"(%arg1) : (i32) -> i32
+  %1 = "test.mixed_variadic_in_out_i32"(%arg1) : (i32) -> i32
+
+  // CHECK: "test.mixed_variadic_in5"(%arg0, %arg1, %arg2) <{attr1 = 1 : i32, pattern_name = "MatchVariadicSubDag"}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%0, %1, %arg2) {attr1 = 1 : i32} : (i32, i32, i32) -> ()
+
+  // Note: MatchVariadicSubDag doesn't apply
+  // CHECK: "test.mixed_variadic_in4"(%arg0, %arg1, %arg2) <{attr1 = 1 : i32}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%arg0, %arg1, %arg2) {attr1 = 1 : i32} : (i32, i32, i32) -> ()
+
+  return
+}
+
+// CHECK-LABEL: @testMatchVariadicSameSymbol
+func.func @testMatchVariadicSameSymbol(%arg0: i32, %arg1: i32, %arg2: i32) -> () {
+  // CHECK: "test.mixed_variadic_in5"(%arg0, %arg0, %arg2) <{attr1 = 2 : i32, pattern_name = "MatchVariadicSameSymbol"}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%arg0, %arg0, %arg2) {attr1 = 2 : i32} : (i32, i32, i32) -> ()
+
+  // Note: MatchVariadicSameSymbol doesn't apply.
+  // CHECK: "test.mixed_variadic_in4"(%arg0, %arg1, %arg2) <{attr1 = 2 : i32}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in4"(%arg0, %arg1, %arg2) {attr1 = 2 : i32} : (i32, i32, i32) -> ()
+
+  return
+}
+
+// CHECK-LABEL: @testMatchAndRewriteVariadicFullRange
+func.func @testMatchAndRewriteVariadicFullRange(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32) -> () {
+  // CHECK: "test.mixed_variadic_in6"(%arg2, %arg3, %arg0, %arg1) <{attr1 = -1 : i32}> : (i32, i32, i32, i32) -> ()
+  "test.mixed_variadic_in6"(%arg0, %arg1, %arg2, %arg3) {attr1 = 1 : i32} : (i32, i32, i32, i32) -> ()
+
+  // Note: MatchAndRewriteVariadicFullRange doesn't apply because the length of each variadic operand is not equal to 2.
+  // CHECK: "test.mixed_variadic_in6"(%arg0, %arg1) <{attr1 = 1 : i32}> : (i32, i32) -> ()
+  "test.mixed_variadic_in6"(%arg0, %arg1) {attr1 = 1 : i32} : (i32, i32) -> ()
+
+  return
+}
+
+// CHECK-LABEL: @testMatchMultiVariadicSubSymbol
+func.func @testMatchMultiVariadicSubSymbol(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: i32) -> () {
+  // CHECK: "test.mixed_variadic_in5"(%arg2, %arg3, %arg1) <{attr1 = 2 : i32, pattern_name = "MatchMultiVariadicSubSymbol"}> : (i32, i32, i32) -> ()
+  "test.mixed_variadic_in6"(%arg0, %arg1, %arg2, %arg3) {attr1 = 2 : i32} : (i32, i32, i32, i32) -> ()
+
+  return
+}
+
 //===----------------------------------------------------------------------===//
 // Test that natives calls are only called once during rewrites.
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: redundantTest
 func.func @redundantTest(%arg0: i32) -> i32 {
   %0 = "test.op_m"(%arg0) : (i32) -> i32
   // CHECK: "test.op_m"(%arg0) <{optional_attr = 314159265 : i32}> : (i32) -> i32

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
   // 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. `operandName` and `operandMatcher` indicate the
   // bound name and the constraint of the operand respectively.
   void emitOperandMatch(DagNode tree, StringRef opName, StringRef operandName,
-                        DagLeaf operandMatcher, StringRef argName,
-                        int argIndex);
+                        int operandIndex, DagLeaf operandMatcher,
+                        StringRef argName, int argIndex,
+                        std::optional<int> variadicSubIndex);
 
   // Emits C++ statements for matching the operands which can be matched in
   // either order.
   void emitEitherOperandMatch(DagNode tree, DagNode eitherArgTree,
                               StringRef opName, int argIndex, int &operandIndex,
                               int depth);
 
+  // Emits C++ statements for matching a variadic operand.
+  void emitVariadicOperandMatch(DagNode tree, DagNode variadicArgTree,
+                                StringRef opName, int argIndex,
+                                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.
   void emitMatchCheck(StringRef opName, const FmtObjectBase &matchFmt,

 class StaticMatcherHelper {
 public:
   StaticMatcherHelper(raw_ostream &os, const RecordKeeper &recordKeeper,
                       RecordOperatorMap &mapper);
 
   // Determine if we should inline the match logic or delegate to a static
   // function.
   bool useStaticMatcher(DagNode node) {
+    // either/variadic node must be associated to the parentOp, thus we can't
+    // emit a static matcher rooted at them.
+    if (node.isEither() || node.isVariadic())
+      return false;
+
     return refStats[node] > kStaticMatcherThreshold;
   }
 
   // Get the name of the static DAG matcher function corresponding to the node.
   std::string getMatcherName(DagNode node) {
     assert(useStaticMatcher(node));
     return matcherNames[node];
   }
void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,
 
   raw_indented_ostream::DelimitedScope scope(os);
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     std::string argName = formatv("arg{0}_{1}", depth, i);
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (argTree.isEither())
         PrintFatalError(loc, "NativeCodeCall cannot have `either` operands");
+      if (argTree.isVariadic())
+        PrintFatalError(loc, "NativeCodeCall cannot have `variadic` operands");
 
       os << "::mlir::Value " << argName << ";\n";
     } else {
       auto leaf = tree.getArgAsLeaf(i);
       if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {
         os << "::mlir::Attribute " << argName << ";\n";
       } else {
         os << "::mlir::Value " << argName << ";\n";
for (int i = 0, opArgIdx = 0, e = tree.getNumArgs(), nextOperand = 0; i != e;
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (argTree.isEither()) {
         emitEitherOperandMatch(tree, argTree, castedName, opArgIdx, nextOperand,
                                depth);
         ++opArgIdx;
         continue;
       }
       if (auto *operand = llvm::dyn_cast_if_present<NamedTypeConstraint *>(opArg)) {
+        if (argTree.isVariadic()) {
+          if (!operand->isVariadic()) {
+            auto error = formatv("variadic DAG construct can't match op {0}'s "
+                                 "non-variadic operand #{1}",
+                                 op.getOperationName(), opArgIdx);
+            PrintFatalError(loc, error);
+          }
+          emitVariadicOperandMatch(tree, argTree, castedName, opArgIdx,
+                                   nextOperand, depth);
+          ++nextOperand;
+          continue;
+        }
         if (operand->isVariableLength()) {
           auto error = formatv("use nested DAG construct to match op {0}'s "
                                "variadic operand #{1} unsupported now",
                                op.getOperationName(), opArgIdx);
           PrintFatalError(loc, error);
         }
       }
 
if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       os.unindent() << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
       auto operandName =
           formatv("{0}.getODSOperands({1})", castedName, nextOperand);
-      emitOperandMatch(tree, castedName, operandName.str(),
+      emitOperandMatch(tree, castedName, operandName.str(), opArgIdx,
                        /*operandMatcher=*/tree.getArgAsLeaf(i),
-                       /*argName=*/tree.getArgName(i), opArgIdx);
+                       /*argName=*/tree.getArgName(i), opArgIdx,
+                       /*variadicSubIndex=*/std::nullopt);
       ++nextOperand;
     } else if (opArg.is<NamedAttribute *>()) {
       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');
 }
 
 void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,
-                                      StringRef operandName,
+                                      StringRef operandName, int operandIndex,
                                       DagLeaf operandMatcher, StringRef argName,
-                                      int argIndex) {
+                                      int argIndex,
+                                      std::optional<int> variadicSubIndex) {
   Operator &op = tree.getDialectOp(opMap);
-  auto *operand = op.getArg(argIndex).get<NamedTypeConstraint *>();
+  auto *operand = op.getArg(operandIndex).get<NamedTypeConstraint *>();
 
   // If a constraint is specified, we need to generate C++ statements to
   // check the constraint.
   if (!operandMatcher.isUnspecified()) {
     if (!operandMatcher.isOperandMatcher())
       PrintFatalError(
           loc, formatv("the {1}-th argument of op '{0}' should be an operand",
                        op.getOperationName(), argIndex + 1));
if (operand->constraint != constraint) {
               escapeString(constraint.getSummary()))
               .str());
     }
   }
 
   // Capture the value
   // `$_` is a special symbol to ignore op argument matching.
   if (!argName.empty() && argName != "_") {
-    auto res = symbolInfoMap.findBoundSymbol(argName, tree, op, argIndex);
+    auto res = symbolInfoMap.findBoundSymbol(argName, tree, op, operandIndex,
+                                             variadicSubIndex);
+    if (res == symbolInfoMap.end())
+      PrintFatalError(loc, formatv("symbol not found: {0}", argName));
+
     os << formatv("{0} = {1};\n", res->second.getVarName(argName), operandName);
   }
 }
 
 void PatternEmitter::emitEitherOperandMatch(DagNode tree, DagNode eitherArgTree,
                                             StringRef opName, int argIndex,
                                             int &operandIndex, int depth) {
   constexpr int numEitherArgs = 2;
if (DagNode argTree = eitherArgTree.getArgAsNestedDag(i)) {
       os.unindent() << "}\n";
 
       // `tblgen_ops` is used to collect the matched operations. In either, we
       // need to queue the operation only if the matching success. Thus we emit
       // the code at the end.
       tblgenOps << formatv("tblgen_ops.push_back({0});\n", argName);
     } else if (op.getArg(argIndex).is<NamedTypeConstraint *>()) {
       emitOperandMatch(tree, opName, /*operandName=*/formatv("v{0}", i).str(),
+                       operandIndex,
                        /*operandMatcher=*/eitherArgTree.getArgAsLeaf(i),
-                       /*argName=*/eitherArgTree.getArgName(i), argIndex);
+                       /*argName=*/eitherArgTree.getArgName(i), argIndex,
+                       /*variadicSubIndex=*/std::nullopt);
       ++operandIndex;
     } else {
       PrintFatalError(loc, "either can only be applied on operand");
     }
   }
 
   os << tblgenOps.str();
   os << "return ::mlir::success();\n";
os << formatv("if(::mlir::failed({0}(eitherOperand0, eitherOperand1)) && "
                 "::mlir::failed({0}(eitherOperand1, "
                 "eitherOperand0)))\n",
                 lambda);
   os.indent() << "return ::mlir::failure();\n";
 
   os.unindent().unindent() << "}\n";
 }
 
+void PatternEmitter::emitVariadicOperandMatch(DagNode tree,
+                                              DagNode variadicArgTree,
+                                              StringRef opName, int argIndex,
+                                              int &operandIndex, int depth) {
+  Operator &op = tree.getDialectOp(opMap);
+
+  os << "{\n";
+  os.indent();
+
+  os << formatv("auto variadic_operand_range = {0}.getODSOperands({1});\n",
+                opName, operandIndex);
+  os << formatv("if (variadic_operand_range.size() != {0}) "
+                "return ::mlir::failure();\n",
+                variadicArgTree.getNumArgs());
+
+  StringRef variadicTreeName = variadicArgTree.getSymbol();
+  if (!variadicTreeName.empty()) {
+    auto res =
+        symbolInfoMap.findBoundSymbol(variadicTreeName, tree, op, operandIndex,
+                                      /*variadicSubIndex=*/std::nullopt);
+    if (res == symbolInfoMap.end())
+      PrintFatalError(loc, formatv("symbol not found: {0}", variadicTreeName));
+
+    os << formatv("{0} = variadic_operand_range;\n",
+                  res->second.getVarName(variadicTreeName));
+  }
+
+  for (int i = 0; i < variadicArgTree.getNumArgs(); ++i) {
+    if (DagNode argTree = variadicArgTree.getArgAsNestedDag(i)) {
+      if (!argTree.isOperation())
+        PrintFatalError(loc, "variadic only accepts operation sub-dags");
+
+      os << "{\n";
+      os.indent();
+
+      std::string argName = formatv("local_op_{0}", i).str();
+      os << formatv("auto *{0} = "
+                    "variadic_operand_range[{1}].getDefiningOp();\n",
+                    argName, i);
+      emitMatchCheck(
+          opName, /*matchStr=*/argName,
+          formatv("\"There's no operation that defines variadic operand "
+                  "{0} (variadic sub-opearnd #{1}) of {2}\"",
+                  operandIndex, i, opName));
+      emitMatch(argTree, argName, depth + 1);
+      os << formatv("tblgen_ops.push_back({0});\n", argName);
+
+      os.unindent() << "}\n";
+    } else if (op.getArg(argIndex).is<NamedTypeConstraint *>()) {
+      auto operandName = formatv("variadic_operand_range.slice({0}, 1)", i);
+      emitOperandMatch(tree, opName, operandName.str(), operandIndex,
+                       /*operandMatcher=*/variadicArgTree.getArgAsLeaf(i),
+                       /*argName=*/variadicArgTree.getArgName(i), argIndex, i);
+    } else {
+      PrintFatalError(loc, "variadic can only be applied on operand");
+    }
+  }
+
+  os.unindent() << "}\n";
+}
+
 void PatternEmitter::emitAttributeMatch(DagNode tree, StringRef opName,
                                         int argIndex, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *namedAttr = op.getArg(argIndex).get<NamedAttribute *>();
   const auto &attr = namedAttr->attr;
 
   os << "{\n";
   os.indent() << formatv("auto tblgen_attr = {0}->getAttrOfType<{1}>(\"{2}\");"