[mlir-tblgen] Support binding multi-results of NativeCodeCall

[Chia-hungDuan]

Summary

We are able to bind NativeCodeCall result as binding operation. To make
table-gen have better understanding in the form of helper function,
we need to specify the number of return values in the NativeCodeCall
template. A VoidNativeCodeCall is added for void case.

[jpienaar]

Nice, thanks, small nits.

mlir/lib/TableGen/Pattern.cpp

339

How about

std::string repl = formatv(...);

? Then you don't need the std::string cast below (same in the other site below which looks identical below - should these always be the same and if so should we use helper?)

449

I think .str() is not needed as you assign directly to a variable of type std::string

mlir/test/lib/Dialect/Test/TestPatterns.cpp

49

How about flipping the order here to ensure that this is called?

mlir/tools/mlir-tblgen/RewriterGen.cpp

952–953

Could you expand here as to what is the incorrect form and how it is captured?

959

This should not be needed as varName is of type std::string

[Chia-hungDuan]

Addressed the review comments

[jpienaar]

Didn't realize this was still pending another review.

mlir/docs/DeclarativeRewrites.md

436

Mention 0 result one here?

mlir/include/mlir/IR/OpBase.td

2580

I was wondering if void should be at the back, that way we have common prefix, but not sure if better.

mlir/include/mlir/TableGen/Pattern.h

325

This name is a bit confusing. Seems a bit more general than its use. Would argIndexOrSize work?

[Chia-hungDuan]

mlir/include/mlir/IR/OpBase.td

2580

After reading it several times, I took your suggestion and switched it to NativeCodeCallVoid.

mlir/include/mlir/TableGen/Pattern.h

325

The other patch(D105677) also changed the same field. I'll merge that first and change this name to DagOrConstant. What do you think? Or the name seems to be to long if we put all the possible meanings?

mlir/lib/TableGen/Pattern.cpp

339

Done. Replaced the auto with std::string(Just tried to keep the same style as other switch-case).

I was thinking about getAllRangeUse, we could support unpacking the return values only, but that means we can't use things like NativeCodeCall<...>:$bind__0 in rewrite pattern. OTOH, for getValueAndRangeUse without indexing, in Kind::Result, it unpacks the values and I just align its semantic. With these two considerations, they just fall into same implementation.

The reason I didn't use a helper for that is, I'm thinking they still mean the different thing. If we want to change the semantic for multiple NativeCodeCall binding, these two cases should be considered separately. I don't to let people think they need to have the same logic here.

BTW, for Kind::Result, they have pretty similar concept as well(in getValueAndRangeUse and getAllRangeUse).

What do you think?

449

The type of getValuePackName() is StringRef and it doesn't seem to have implicit conversion to std::string.

mlir/test/lib/Dialect/Test/TestPatterns.cpp

49

Good idea, done.

[Chia-hungDuan]

Addressed review comment and fix a bug that handleReplaceWithNativeCodeCall
Should be replaced with childNodeNames lookup when preparing arguments for a
call.

[Chia-hungDuan]

Rebase and update the variable name

[jpienaar]

mlir/docs/DeclarativeRewrites.md

446

s/Beside/Additionally/

448

Lets spell this out. if a NativeCodeCall that doesn't return a result, isn't labeled with 0 returns.

But independent of 0, if you report incorrect number of results, then there would probably be OOB accesses that would just fail at runtime. So pretty much, don't do that :)

mlir/include/mlir/IR/OpBase.td

2572

And this is a constraint on the function invoked by NativeCodeCall ?

mlir/include/mlir/TableGen/Pattern.h

263

Nit: add 2 spaces before to make it clearer part of list item

325

SG, if changed in follow up. I think focusing on what it provides in the name would be most important.

mlir/lib/TableGen/Pattern.cpp

339

SG

[Chia-hungDuan]

Addressed review comment

mlir/docs/DeclarativeRewrites.md

 And make sure the generated C++ code from the above pattern has access to the
 definition of the C++ helper function.
 
 In the above example, we are using a string to specialize the `NativeCodeCall`
 template. The string can be an arbitrary C++ expression that evaluates into some
 C++ object expected at the `NativeCodeCall` site (here it would be expecting an
 array attribute). Typically the string should be a function call.
 
-Note that currently `NativeCodeCall` must return no more than one value or
-attribute. This might change in the future.
-
 ##### `NativeCodeCall` placeholders
 
 In `NativeCodeCall`, we can use placeholders like `$_builder`, `$N` and `$N...`.
 The former is called _special placeholder_, while the latter is called
 _positional placeholder_ and _positional range placeholder_.
 
 `NativeCodeCall` right now only supports three special placeholders:
 `$_builder`, `$_loc`, and `$_self`:

 NativeCodeCall<"someFn($1, $2, $0)">` and use it like `(SomeCall $in0, $in1,
 $in2)`, then this will be translated into C++ call `someFn($in1, $in2, $in0)`.
 
 Positional range placeholders will be substituted by multiple `dag` object
 parameters at the `NativeCodeCall` use site. For example, if we define
 `SomeCall : NativeCodeCall<"someFn($1...)">` and use it like `(SomeCall $in0,
 $in1, $in2)`, then this will be translated into C++ call `someFn($in1, $in2)`.
 
+##### `NativeCodeCall` binding multi-results
+
+To bind multi-results and access the N-th result with `$<name>__N`, specify the
+number of return values in the template. Note that only `Value` type is
+supported for multiple results binding. For example,
+
+```tablegen
+
+def PackAttrs : NativeCodeCall<"packAttrs($0, $1)", 2>;
+def : Pattern<(TwoResultOp $attr1, $attr2),
+              [(OneResultOp (PackAttr:$res__0, $attr1, $attr2)),
+               (OneResultOp $res__1)]>;
+
+```
+
+Use `NativeCodeCallVoid` for case has no return value.
+
+The correct number of returned value specified in NativeCodeCall is important.
+It will be used to verify the consistency of the number of result values.
+Additionally, `mlir-tblgen` will try to capture the return value of
+NativeCodeCall in the generated code so that it will trigger a later compilation
+error if a NativeCodeCall that doesn't return a result isn't labeled with 0
+returns.
+
 ##### Customizing entire op building
 
 `NativeCodeCall` is not only limited to transforming arguments for building an
 op; it can be also used to specify how to build an op entirely. An example:
 
 If we have a C++ function for building an op:
 
 ```c++

mlir/include/mlir/IR/OpBase.td

 // * `$_builder` will be replaced by the current `mlir::PatternRewriter`.
 // * `$_self` will be replaced by the defining operation in a source pattern.
 //   E.g., `NativeCodeCall<"Foo($_self, &$0)> I32Attr:$attr)>`, `$_self` will be
 //   replaced with the defining operation of the first operand of OneArgOp.
 //
 // If used as a DAG node, i.e., `(NativeCodeCall<"..."> <arg0>, ..., <argN>)`,
 // then positional placeholders are also supported; placeholder `$N` in the
 // wrapped C++ expression will be replaced by `<argN>`.
+//
+// ## Bind multiple results
+//
+// To bind multi-results and access the N-th result with `$<name>__N`, specify
+// the number of return values in the template. Note that only `Value` type is
+// supported for multiple results binding.
 
-class NativeCodeCall<string expr> {
+class NativeCodeCall<string expr, int returns = 1> {
   string expression = expr;
+  int numReturns = returns;
 }
 
+class NativeCodeCallVoid<string expr> : NativeCodeCall<expr, 0>;
+
 def ConstantLikeMatcher : NativeCodeCall<"success(matchPattern($_self->getResult(0), m_Constant(&$0)))">;
 
 //===----------------------------------------------------------------------===//
 // Rewrite directives
 //===----------------------------------------------------------------------===//
 
 // Directive used in result pattern to specify the location of the generated
 // op. This directive must be used as the last argument to the op creation

mlir/include/mlir/TableGen/Pattern.h

public:
   // Returns the matching condition template inside this DAG leaf. Assumes the
   // leaf is an operand/attribute matcher and asserts otherwise.
   std::string getConditionTemplate() const;
 
   // Returns the native code call template inside this DAG leaf.
   // Precondition: isNativeCodeCall()
   StringRef getNativeCodeTemplate() const;
 
+  // Returns the number of values will be returned by the native helper
+  // function.
+  // Precondition: isNativeCodeCall()
+  int getNumReturnsOfNativeCode() const;
+
   // Returns the string associated with the leaf.
   // Precondition: isStringAttr()
   std::string getStringAttr() const;
 
   void print(raw_ostream &os) const;
 
 private:
   // Returns true if the TableGen Init `def` in this DagLeaf is a DefInit and
public:
 
   // 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
+  // function.
+  // Precondition: isNativeCodeCall()
+  int getNumReturnsOfNativeCode() const;
+
   void print(raw_ostream &os) const;
 
 private:
   friend class SymbolInfoMap;
   const void *getAsOpaquePointer() const { return node; }
 
   const llvm::DagInit *node; // nullptr means null DagNode
 };
public:
     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 DagAndIndex = std::pair<const void *, int>;
+    using DagAndConstant = std::pair<const void *, int>;
 
     // 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)
-    enum class Kind : uint8_t { Attr, Operand, Result, Value };
+    // * 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. `dagAndIndex` is only used for `Attr` and
-    // `Operand` so should be llvm::None for `Result` and `Value` kind.
+    // Creates a SymbolInfo instance. `dagAndConstant` is only used for `Attr`
+    // and `Operand` so should be llvm::None for `Result` and `Value` kind.
     SymbolInfo(const Operator *op, Kind kind,
-               Optional<DagAndIndex> dagAndIndex);
+               Optional<DagAndConstant> dagAndConstant);
 
     // Static methods for creating SymbolInfo.
     static SymbolInfo getAttr(const Operator *op, int index) {
-      return SymbolInfo(op, Kind::Attr, DagAndIndex(nullptr, index));
+      return SymbolInfo(op, Kind::Attr, DagAndConstant(nullptr, index));
     }
     static SymbolInfo getAttr() {
       return SymbolInfo(nullptr, Kind::Attr, llvm::None);
     }
     static SymbolInfo getOperand(DagNode node, const Operator *op, int index) {
       return SymbolInfo(op, Kind::Operand,
-                        DagAndIndex(node.getAsOpaquePointer(), index));
+                        DagAndConstant(node.getAsOpaquePointer(), index));
     }
     static SymbolInfo getResult(const Operator *op) {
       return SymbolInfo(op, Kind::Result, llvm::None);
     }
     static SymbolInfo getValue() {
       return SymbolInfo(nullptr, Kind::Value, llvm::None);
     }
+    static SymbolInfo getMultipleValues(int numValues) {
+      return SymbolInfo(nullptr, Kind::MultipleValues,
+                        DagAndConstant(nullptr, numValues));
+    }
 
     // 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;
 
     // Returns a string containing the C++ expression for referencing this
private:
     // symbol as a value range regardless of how many static values this symbol
     // 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; }
+
+    // The number of values in the MultipleValue
+    int getSize() const { return (*dagAndConstant).second; }
+
     const Operator *op; // The op where the bound entity belongs
     Kind kind;          // The kind of the bound entity
-    // The pair of DagNode pointer and argument index (for `Attr` and `Operand`
-    // only). Note that operands may be bound to the same symbol, use the
-    // DagNode and index to distinguish them. For `Attr`, the Dag part will be
-    // nullptr.
-    Optional<DagAndIndex> dagAndIndex;
+
+    // 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.
+    Optional<DagAndConstant> dagAndConstant;
+
     // Alternative name for the symbol. It is used in case the name
     // is not unique. Applicable for `Operand` only.
     Optional<std::string> alternativeName;
   };
 
   using BaseT = std::unordered_multimap<std::string, SymbolInfo>;
 
   // Iterators for accessing all symbols.
public:
   // Returns false if `symbol` is already bound and symbols are not operands.
   bool bindOpArgument(DagNode node, StringRef symbol, const Operator &op,
                       int argIndex);
 
   // Binds the given `symbol` to the results the given `op`. Returns false if
   // `symbol` is already bound.
   bool bindOpResult(StringRef symbol, const Operator &op);
 
-  // Registers the given `symbol` as bound to a value. Returns false if `symbol`
-  // is already bound.
+  // A helper function for dispatching target value binding functions.
+  bool bindValues(StringRef symbol, int numValues = 1);
+
+  // Registers the given `symbol` as bound to the Value(s). Returns false if
+  // `symbol` is already bound.
   bool bindValue(StringRef symbol);
 
+  // Registers the given `symbol` as bound to a MultipleValue. Return false if
+  // `symbol` is already bound.
+  bool bindMultipleValues(StringRef symbol, int numValues);
+
   // Registers the given `symbol` as bound to an attr. Returns false if `symbol`
   // is already bound.
   bool bindAttr(StringRef symbol);
 
   // Returns true if the given `symbol` is bound.
   bool contains(StringRef symbol) const;
 
   // Returns an iterator to the information of the given symbol named as `key`.

mlir/lib/TableGen/Pattern.cpp

std::string DagLeaf::getConditionTemplate() const {
   return getAsConstraint().getConditionTemplate();
 }
 
 llvm::StringRef DagLeaf::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(def)->getDef()->getValueAsString("expression");
 }
 
+int DagLeaf::getNumReturnsOfNativeCode() const {
+  assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
+  return cast<llvm::DefInit>(def)->getDef()->getValueAsInt("numReturns");
+}
+
 std::string DagLeaf::getStringAttr() const {
   assert(isStringAttr() && "the DAG leaf must be string attribute");
   return def->getAsUnquotedString();
 }
 bool DagLeaf::isSubClassOf(StringRef superclass) const {
   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(def))
     return defInit->getDef()->isSubClassOf(superclass);
   return false;

 
 llvm::StringRef DagNode::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsString("expression");
 }
 
+int DagNode::getNumReturnsOfNativeCode() const {
+  assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
+  return cast<llvm::DefInit>(node->getOperator())
+      ->getDef()
+      ->getValueAsInt("numReturns");
+}
+
 llvm::StringRef DagNode::getSymbol() const { return node->getNameStr(); }
 
 Operator &DagNode::getDialectOp(RecordOperatorMap *mapper) const {
   llvm::Record *opDef = cast<llvm::DefInit>(node->getOperator())->getDef();
   auto it = mapper->find(opDef);
   if (it != mapper->end())
     return *it->second;
   return *mapper->try_emplace(opDef, std::make_unique<Operator>(opDef))
StringRef SymbolInfoMap::getValuePackName(StringRef symbol, int *index) {
   }
   if (index) {
     *index = idx;
   }
   return name;
 }
 
 SymbolInfoMap::SymbolInfo::SymbolInfo(const Operator *op, SymbolInfo::Kind kind,
-                                      Optional<DagAndIndex> dagAndIndex)
-    : op(op), kind(kind), dagAndIndex(dagAndIndex) {}
+                                      Optional<DagAndConstant> dagAndConstant)
+    : op(op), kind(kind), dagAndConstant(dagAndConstant) {}
 
 int SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
   switch (kind) {
   case Kind::Attr:
   case Kind::Operand:
   case Kind::Value:
     return 1;
   case Kind::Result:
     return op->getNumResults();
+  case Kind::MultipleValues:
+    return getSize();
   }
   llvm_unreachable("unknown kind");
 }
 
 std::string SymbolInfoMap::SymbolInfo::getVarName(StringRef name) const {
   return alternativeName.hasValue() ? alternativeName.getValue() : name.str();
 }
 
 std::string SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
   LLVM_DEBUG(llvm::dbgs() << "getVarDecl for '" << name << "': ");
   switch (kind) {
   case Kind::Attr: {
     if (op) {
-      auto type = op->getArg((*dagAndIndex).second)
+      auto type = op->getArg(getArgIndex())
                       .get<NamedAttribute *>()
                       ->attr.getStorageType();
       return std::string(formatv("{0} {1};\n", type, name));
     }
     // TODO(suderman): Use a more exact type when available.
     return std::string(formatv("Attribute {0};\n", name));
   }
   case Kind::Operand: {
     // Use operand range for captured operands (to support potential variadic
     // operands).
     return std::string(
         formatv("::mlir::Operation::operand_range {0}(op0->getOperands());\n",
                 getVarName(name)));
   }
   case Kind::Value: {
     return std::string(formatv("::mlir::Value {0};\n", name));
   }
+  case Kind::MultipleValues: {
+    // This is for the variable used in the source pattern. Each named value in
+    // source pattern will only be bound to a Value. The others in the result
+    // pattern may be associated with multiple Values as we will use `auto` to
+    // do the type inference.
+    return std::string(formatv(
+        "::mlir::Value {0}_raw; ::mlir::ValueRange {0}({0}_raw);\n", name));
+  }
   case Kind::Result: {
     // Use the op itself for captured results.
     return std::string(formatv("{0} {1};\n", op->getQualCppClassName(), name));
   }
   }
   llvm_unreachable("unknown kind");
 }
 
 std::string SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
     StringRef name, int index, const char *fmt, const char *separator) const {
   LLVM_DEBUG(llvm::dbgs() << "getValueAndRangeUse for '" << name << "': ");
   switch (kind) {
   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((*dagAndIndex).second).get<NamedTypeConstraint *>();
+    auto *operand = op->getArg(getArgIndex()).get<NamedTypeConstraint *>();
     // If this operand is variadic, then return a range. Otherwise, return the
     // value itself.
     if (operand->isVariableLength()) {
       auto repl = formatv(fmt, name);
       LLVM_DEBUG(llvm::dbgs() << repl << " (VariadicOperand)\n");
       return std::string(repl);
     }
     auto repl = formatv(fmt, formatv("(*{0}.begin())", name));
std::string SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
   }
   case Kind::Value: {
     assert(index < 0);
     assert(op == nullptr);
     auto repl = formatv(fmt, name);
     LLVM_DEBUG(llvm::dbgs() << repl << " (Value)\n");
     return std::string(repl);
   }
+  case Kind::MultipleValues: {
+    assert(op == nullptr);
+    assert(index < getSize());
+    if (index >= 0) {
+      std::string repl =
+          formatv(fmt, std::string(formatv("{0}[{1}]", name, index)));
+      LLVM_DEBUG(llvm::dbgs() << repl << " (MultipleValues)\n");
+      return repl;
+    }
+    // If it doesn't specify certain element, unpack them all.
+    auto repl =
+        formatv(fmt, std::string(formatv("{0}.begin(), {0}.end()", name)));
+    LLVM_DEBUG(llvm::dbgs() << repl << " (MultipleValues)\n");
+    return std::string(repl);
+  }
   }
   llvm_unreachable("unknown kind");
 }
 
 std::string SymbolInfoMap::SymbolInfo::getAllRangeUse(
     StringRef name, int index, const char *fmt, const char *separator) const {
   LLVM_DEBUG(llvm::dbgs() << "getAllRangeUse for '" << name << "': ");
   switch (kind) {
std::string SymbolInfoMap::SymbolInfo::getAllRangeUse(
   }
   case Kind::Value: {
     assert(index < 0 && "only allowed for symbol bound to result");
     assert(op == nullptr);
     auto repl = formatv(fmt, formatv("{{{0}}", name));
     LLVM_DEBUG(llvm::dbgs() << repl << " (Value)\n");
     return std::string(repl);
   }
+  case Kind::MultipleValues: {
+    assert(op == nullptr);
+    assert(index < getSize());
+    if (index >= 0) {
+      std::string repl =
+          formatv(fmt, std::string(formatv("{0}[{1}]", name, index)));
+      LLVM_DEBUG(llvm::dbgs() << repl << " (MultipleValues)\n");
+      return repl;
+    }
+    auto repl =
+        formatv(fmt, std::string(formatv("{0}.begin(), {0}.end()", name)));
+    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) {
   StringRef name = getValuePackName(symbol);
   if (name != symbol) {

 
 bool SymbolInfoMap::bindOpResult(StringRef symbol, const Operator &op) {
   std::string name = getValuePackName(symbol).str();
   auto inserted = symbolInfoMap.emplace(name, SymbolInfo::getResult(&op));
 
   return symbolInfoMap.count(inserted->first) == 1;
 }
 
+bool SymbolInfoMap::bindValues(StringRef symbol, int numValues) {
+  std::string name = getValuePackName(symbol).str();
+  if (numValues > 1)
+    return bindMultipleValues(name, numValues);
+  return bindValue(name);
+}
+
 bool SymbolInfoMap::bindValue(StringRef symbol) {
   auto inserted = symbolInfoMap.emplace(symbol.str(), SymbolInfo::getValue());
   return symbolInfoMap.count(inserted->first) == 1;
 }
 
+bool SymbolInfoMap::bindMultipleValues(StringRef symbol, int numValues) {
+  std::string name = getValuePackName(symbol).str();
+  auto inserted =
+      symbolInfoMap.emplace(name, SymbolInfo::getMultipleValues(numValues));
+  return symbolInfoMap.count(inserted->first) == 1;
+}
+
 bool SymbolInfoMap::bindAttr(StringRef symbol) {
   auto inserted = symbolInfoMap.emplace(symbol.str(), SymbolInfo::getAttr());
   return symbolInfoMap.count(inserted->first) == 1;
 }
 
 bool SymbolInfoMap::contains(StringRef symbol) const {
   return find(symbol) != symbolInfoMap.end();
 }
 
 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 {
   std::string name = getValuePackName(key).str();
   auto range = symbolInfoMap.equal_range(name);
 
   const auto symbolInfo = SymbolInfo::getOperand(node, &op, argIndex);
 
-  for (auto it = range.first; it != range.second; ++it) {
-    if (it->second.dagAndIndex == symbolInfo.dagAndIndex) {
+  for (auto it = range.first; it != range.second; ++it)
+    if (it->second.dagAndConstant == symbolInfo.dagAndConstant)
       return it;
-    }
-  }
 
   return symbolInfoMap.end();
 }
 
 std::pair<SymbolInfoMap::iterator, SymbolInfoMap::iterator>
 SymbolInfoMap::getRangeOfEqualElements(StringRef key) {
   std::string name = getValuePackName(key).str();
 
void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
   auto treeName = tree.getSymbol();
   auto numTreeArgs = tree.getNumArgs();
 
   if (tree.isNativeCodeCall()) {
     if (!treeName.empty()) {
       if (!isSrcPattern) {
         LLVM_DEBUG(llvm::dbgs() << "found symbol bound to NativeCodeCall: "
                                 << treeName << '\n');
-        verifyBind(infoMap.bindValue(treeName), treeName);
+        verifyBind(
+            infoMap.bindValues(treeName, tree.getNumReturnsOfNativeCode()),
+            treeName);
       } else {
         PrintFatalError(&def,
                         formatv("binding symbol '{0}' to NativecodeCall in "
                                 "MatchPattern is not supported",
                                 treeName));
       }
     }
 

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

 
 def OpNativeCodeCall3 : TEST_Op<"native_code_call3"> {
   let arguments = (ins I32:$input);
   let results = (outs I32);
 }
 // Test that NativeCodeCall is not ignored if it is not used to directly
 // replace the matched root op.
 def : Pattern<(OpNativeCodeCall3 $input),
-              [(NativeCodeCall<"createOpI($_builder, $_loc, $0)"> $input),
+              [(NativeCodeCallVoid<"createOpI($_builder, $_loc, $0)"> $input),
                (OpK)]>;
 
 def OpNativeCodeCall4 : TEST_Op<"native_code_call4"> {
   let arguments = (ins AnyType:$input1);
   let results = (outs I32:$output1, I32:$output2);
 }
 def OpNativeCodeCall5 : TEST_Op<"native_code_call5"> {
   let arguments = (ins I32:$input1, I32:$input2);
   let results = (outs I32:$output1, I32:$output2);
 }
 
 def GetFirstI32Result : NativeCodeCall<"success(getFirstI32Result($_self, $0))">;
 def BindNativeCodeCallResult : NativeCodeCall<"bindNativeCodeCallResult($0)">;
 def : Pat<(OpNativeCodeCall4 (GetFirstI32Result $ret)),
           (OpNativeCodeCall5 (BindNativeCodeCallResult:$native $ret), $native)>;
 
+def OpNativeCodeCall6 : TEST_Op<"native_code_call6"> {
+  let arguments = (ins I32:$input1, I32:$input2);
+  let results = (outs I32:$output1, I32:$output2);
+}
+def OpNativeCodeCall7 : TEST_Op<"native_code_call7"> {
+  let arguments = (ins I32:$input);
+  let results = (outs I32);
+}
+def BindMultipleNativeCodeCallResult : NativeCodeCall<"bindMultipleNativeCodeCallResult($0, $1)", 2>;
+def : Pattern<(OpNativeCodeCall6 $arg1, $arg2),
+              [(OpNativeCodeCall7 (BindMultipleNativeCodeCallResult:$native__0 $arg1, $arg2)),
+               (OpNativeCodeCall7 $native__1)]>;
+
 // Test AllAttrConstraintsOf.
 def OpAllAttrConstraint1 : TEST_Op<"all_attr_constraint_of1"> {
   let arguments = (ins I64ArrayAttr:$attr);
   let results = (outs I32);
 }
 def OpAllAttrConstraint2 : TEST_Op<"all_attr_constraint_of2"> {
   let arguments = (ins I64ArrayAttr:$attr);
   let results = (outs I32);
def : Pattern<
     [(HasOneUse $res_a)]>;
 
 def OpSymbolBindingNoResult : TEST_Op<"symbol_binding_no_result", []> {
   let arguments = (ins I32:$operand);
 }
 
 // Test that we can bind to an op without results and reference it later.
 def : Pat<(OpSymbolBindingNoResult:$op $operand),
-          (NativeCodeCall<"handleNoResultOp($_builder, $0)"> $op)>;
+          (NativeCodeCallVoid<"handleNoResultOp($_builder, $0)"> $op)>;
 
 //===----------------------------------------------------------------------===//
 // Test Patterns (Attributes)
 
 // Test matching against op attributes.
 def OpAttrMatch1 : TEST_Op<"match_op_attribute1"> {
   let arguments = (ins
     I32Attr:$required_attr,

mlir/test/lib/Dialect/Test/TestPatterns.cpp

static bool getFirstI32Result(Operation *op, Value &value) {
   if (!Type(op->getResult(0).getType()).isSignlessInteger(32))
     return false;
   value = op->getResult(0);
   return true;
 }
 
 static Value bindNativeCodeCallResult(Value value) { return value; }
 
+static SmallVector<Value, 2> bindMultipleNativeCodeCallResult(Value input1,
+                                                              Value input2) {
+  return SmallVector<Value, 2>({input2, input1});
+}
+
 // Test that natives calls are only called once during rewrites.
 // OpM_Test will return Pi, increased by 1 for each subsequent calls.
 // This let us check the number of times OpM_Test was called by inspecting
 // the returned value in the MLIR output.
 static int64_t opMIncreasingValue = 314159265;
 static Attribute OpMTest(PatternRewriter &rewriter, Value val) {
   int64_t i = opMIncreasingValue++;
   return rewriter.getIntegerAttr(rewriter.getIntegerType(32), i);

mlir/test/mlir-tblgen/pattern.mlir

func @verifyNativeCodeCallBinding(%arg0 : i32) -> (i32) {
   // %3 is i8 so it'll fail at GetFirstI32Result match. The operation should
   // keep the same form.
   // CHECK: %{{.*}}, %{{.*}} = "test.native_code_call4"({{%.*}}) : (i8) -> (i32, i32)
   %4, %5 = "test.native_code_call4"(%3) : (i8) -> (i32, i32)
   // CHECK: return %[[A]]
   return %1 : i32
 }
 
+// CHECK-LABEL: verifyMultipleNativeCodeCallBinding
+func@verifyMultipleNativeCodeCallBinding(%arg0 : i32) -> (i32) {
+  %0 = "test.op_k"() : () -> (i32)
+  %1 = "test.op_k"() : () -> (i32)
+  // CHECK: %[[A:.*]] = "test.native_code_call7"(%1) : (i32) -> i32
+  // CHECK: %[[A:.*]] = "test.native_code_call7"(%0) : (i32) -> i32
+  %2, %3 = "test.native_code_call6"(%0, %1) : (i32, i32) -> (i32, i32)
+  return %2 : i32
+}
+
 // CHECK-LABEL: verifyAllAttrConstraintOf
 func @verifyAllAttrConstraintOf() -> (i32, i32, i32) {
   // CHECK: "test.all_attr_constraint_of2"
   %0 = "test.all_attr_constraint_of1"() {attr = [0, 1]} : () -> (i32)
   // CHECK: "test.all_attr_constraint_of1"
   %1 = "test.all_attr_constraint_of1"() {attr = [0, 2]} : () -> (i32)
   // CHECK: "test.all_attr_constraint_of1"
   %2 = "test.all_attr_constraint_of1"() {attr = [-1, 1]} : () -> (i32)

mlir/tools/mlir-tblgen/RewriterGen.cpp

void PatternEmitter::emitRewriteLogic() {
   // Process auxiliary result patterns.
   for (int i = 0; i < replStartIndex; ++i) {
     DagNode resultTree = pattern.getResultPattern(i);
     auto val = handleResultPattern(resultTree, offsets[i], 0);
     // Normal op creation will be streamed to `os` by the above call; but
     // NativeCodeCall will only be materialized to `os` if it is used. Here
     // we are handling auxiliary patterns so we want the side effect even if
     // NativeCodeCall is not replacing matched root op's results.
-    if (resultTree.isNativeCodeCall())
+    if (resultTree.isNativeCodeCall() &&
+        resultTree.getNumReturnsOfNativeCode() == 0)
       os << val << ";\n";
   }
 
   if (numExpectedResults == 0) {
     assert(replStartIndex >= numResultPatterns &&
            "invalid auxiliary vs. replacement pattern division!");
     // No result to replace. Just erase the op.
     os << "rewriter.eraseOp(op0);\n";
std::string PatternEmitter::handleResultPattern(DagNode resultTree,
   LLVM_DEBUG(resultTree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   if (resultTree.isLocationDirective()) {
     PrintFatalError(loc,
                     "location directive can only be used with op creation");
   }
 
-  if (resultTree.isNativeCodeCall()) {
-    auto symbol = handleReplaceWithNativeCodeCall(resultTree, depth);
-    symbolInfoMap.bindValue(symbol);
-    return symbol;
-  }
+  if (resultTree.isNativeCodeCall())
+    return handleReplaceWithNativeCodeCall(resultTree, depth);
 
   if (resultTree.isReplaceWithValue())
     return handleReplaceWithValue(resultTree).str();
 
   // Normal op creation.
   auto symbol = handleOpCreation(resultTree, resultIndex, depth);
   if (resultTree.getSymbol().empty()) {
     // This is an op not explicitly bound to a symbol in the rewrite rule.
if (tree.isNestedDagArg(i)) {
       attrs.push_back(
           handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)));
     }
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
                             << " replacement: " << attrs[i] << "\n");
   }
 
   std::string symbol = tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs);
-  if (!tree.getSymbol().empty()) {
-    os << formatv("auto {0} = {1};\n", tree.getSymbol(), symbol);
+
+  // In general, NativeCodeCall without naming binding don't need this. To
+  // ensure void helper function has been correctly labeled, i.e., use
+  // NativeCodeCallVoid, we cache the result to a local variable so that we will
+  // get a compilation error in the auto-generated file.
+  // Example.
+  //   // In the td file
+  //   Pat<(...), (NativeCodeCall<Foo> ...)>
+  //
+  //   ---
+  //
+  //   // In the auto-generated .cpp
+  //   ...
+  //   // Causes compilation error if Foo() returns void.
+  //   auto nativeVar = Foo();
+  //   ...
+  if (tree.getNumReturnsOfNativeCode() != 0) {
+    // Determine the local variable name for return value.
+    std::string varName =
+        SymbolInfoMap::getValuePackName(tree.getSymbol()).str();
+    if (varName.empty()) {
+      varName = formatv("nativeVar_{0}", nextValueId++);
+      // Register the local variable for later uses.
+      symbolInfoMap.bindValues(varName, tree.getNumReturnsOfNativeCode());
+    }
+
+    // Catch the return value of helper function.
+    os << formatv("auto {0} = {1}; (void){0};\n", varName, symbol);
+
+    if (!tree.getSymbol().empty())
       symbol = tree.getSymbol().str();
+    else
+      symbol = varName;
   }
 
   return symbol;
 }
 
 int PatternEmitter::getNodeValueCount(DagNode node) {
   if (node.isOperation()) {
     // If the op is bound to a symbol in the rewrite rule, query its result
     // count from the symbol info map.
     auto symbol = node.getSymbol();
     if (!symbol.empty()) {
       return symbolInfoMap.getStaticValueCount(symbol);
     }
     // Otherwise this is an unbound op; we will use all its results.
     return pattern.getDialectOp(node).getNumResults();
   }
-  // TODO: This considers all NativeCodeCall as returning one
-  // value. Enhance if multi-value ones are needed.
+
+  if (node.isNativeCodeCall())
+    return node.getNumReturnsOfNativeCode();
+
   return 1;
 }
 
 std::pair<bool, std::string> PatternEmitter::getLocation(DagNode tree) {
   auto numPatArgs = tree.getNumArgs();
 
   if (numPatArgs != 0) {
     if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();
     }
 
     // The argument in the op definition.
     auto opArgName = resultOp.getArgName(argIndex);
     if (auto subTree = node.getArgAsNestedDag(argIndex)) {
       if (!subTree.isNativeCodeCall())
         PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                              "for creating attribute");
-      os << formatv("/*{0}=*/{1}", opArgName,
-                    handleReplaceWithNativeCodeCall(subTree, depth));
+      os << formatv("/*{0}=*/{1}", opArgName, childNodeNames.lookup(argIndex));
     } else {
       auto leaf = node.getArgAsLeaf(argIndex);
       // The argument in the result DAG pattern.
       auto patArgName = node.getArgName(argIndex);
       if (leaf.isConstantAttr() || leaf.isEnumAttrCase()) {
         // TODO: Refactor out into map to avoid recomputing these.
         if (!opArg.is<NamedAttribute *>())
           PrintFatalError(loc, Twine("expected attribute ") + Twine(argIndex));
void PatternEmitter::createAggregateLocalVarsForOpArgs(
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
       // The argument in the op definition.
       auto opArgName = resultOp.getArgName(argIndex);
       if (auto subTree = node.getArgAsNestedDag(argIndex)) {
         if (!subTree.isNativeCodeCall())
           PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                                "for creating attribute");
-        os << formatv(addAttrCmd, opArgName,
-                      handleReplaceWithNativeCodeCall(subTree, depth + 1));
+        os << formatv(addAttrCmd, opArgName, childNodeNames.lookup(argIndex));
       } else {
         auto leaf = node.getArgAsLeaf(argIndex);
         // The argument in the result DAG pattern.
         auto patArgName = node.getArgName(argIndex);
         os << formatv(addAttrCmd, opArgName,
                       handleOpArgument(leaf, patArgName));
       }
       continue;