[mlir-tblgen] Fix failed matching when binds same operand of an op in different depth

[Chia-hungDuan]

Summary

For example, we will generate incorrect code for the pattern,

def : Pat<((FooOp (FooOp, $a, $b), $b)), (...)>;

We didn't allow $b to be bond twice with same operand of same op.

[Chia-hungDuan]

Fix comments and argument spelling

[jpienaar]

Just did a quick skim/will look more later

mlir/include/mlir/TableGen/Pattern.h

243

Why not just use a pointer to DagNode?

[Chia-hungDuan]

Fix variable typo

[Chia-hungDuan]

Failed upload the change in previous patches.

[jpienaar]

mlir/include/mlir/TableGen/Pattern.h

243

Just commenting to ensure you didn't miss this before.

304

bound

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

719

So this would have generated incorrect code before, what would the generated code have done?

[rdzhabarov]

mlir/include/mlir/TableGen/Pattern.h

188

nit: const llvm::DagInit *.

254–257

could you update comments on methods everywhere when it's applicable.

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

719

I think it would be beneficial to have a link to paste with the generated code before and after.

mlir/test/mlir-tblgen/pattern.mlir

164

comment seems to be slightly different from the actual pattern

def TestNestedSameOpAndSameArgEqualityPattern :
  Pat<(OpN (OpN $_, $x), $x), (replaceWithValue $x)>;

[mehdi_amini]

Nit: typo in the title falied

[Chia-hungDuan]

Addressed review comments

[Chia-hungDuan]

In D105677#2877917, @mehdi_amini wrote:

Nit: typo in the title falied

Done.

mlir/include/mlir/TableGen/Pattern.h

188

This is supposed to be an opaque pointer identifier, will it be better not to expose the type? Or do you mean let's cast the type on the caller side? I.e., the SymbolInfo will do the casting.

243

Sorry, I was supposed to cover this during our chat.

DagNode/DagLeaf are accessed by value (like if (DagNode argTree = tree.getArgAsNestedDag(i)) ) so we can't use DagNode* as a identifier.

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

719

Please check,
https://gist.githubusercontent.com/ChiaHungDuan/312cd6e47e0ff17183d1d0dfbf236b37/raw/c629bf7e4fe164646a8d9170fcde14d8ce92eac5/Operand%2520match%2520diff

Brief summary, there's only x being assigned to the 2nd operand, x0 has the default value(which is incorrect) in 'before'

mlir/test/mlir-tblgen/pattern.mlir

164

Sorry, changed the case and forgot to update it.

[jpienaar]

mlir/include/mlir/TableGen/Pattern.h

243

Lets add a comment to this point here. Then it is locally clear why a less specific type is used.

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

719

That last sentence is an important one to add to description as that makes it clear what was wrong and how it would have failed.

[Chia-hungDuan]

Addressed review comments

[rdzhabarov]

mlir/include/mlir/TableGen/Pattern.h

188

nvm. Yeah, i think it's better to have it here, explicit. I somehow overlooked the intent of the method.

mlir/include/mlir/TableGen/Pattern.h

public:
 
   // Returns the native code call template inside this DAG node.
   // Precondition: isNativeCodeCall()
   StringRef getNativeCodeTemplate() 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
 };
 
 // A class for maintaining information for symbols bound in patterns and
 // provides methods for resolving them according to specific use cases.
 //
 // Symbols can be bound to
 //
public:
 
     // 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 DagAndIndex = 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 };
 
-    // Creates a SymbolInfo instance. `index` is only used for `Attr` and
-    // `Operand` so should be negative for `Result` and `Value` kind.
-    SymbolInfo(const Operator *op, Kind kind, Optional<int> index);
+    // Creates a SymbolInfo instance. `dagAndIndex` 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);
 
     // Static methods for creating SymbolInfo.
     static SymbolInfo getAttr(const Operator *op, int index) {
-      return SymbolInfo(op, Kind::Attr, index);
+      return SymbolInfo(op, Kind::Attr, DagAndIndex(nullptr, index));
     }
     static SymbolInfo getAttr() {
       return SymbolInfo(nullptr, Kind::Attr, llvm::None);
     }
-    static SymbolInfo getOperand(const Operator *op, int index) {
-      return SymbolInfo(op, Kind::Operand, index);
+    static SymbolInfo getOperand(DagNode node, const Operator *op, int index) {
+      return SymbolInfo(op, Kind::Operand,
+                        DagAndIndex(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);
     }
 
private:
     // 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;
 
     const Operator *op; // The op where the bound entity belongs
     Kind kind;          // The kind of the bound entity
-    // The argument index (for `Attr` and `Operand` only)
-    Optional<int> argIndex;
+    // 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;
     // 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.
   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(StringRef symbol, const Operator &op, int argIndex);
+  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.
   bool bindValue(StringRef symbol);
 
   // 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`.
   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, const Operator &op,
-                                 int argIndex) const;
+  const_iterator findBoundSymbol(StringRef key, DagNode node,
+                                 const Operator &op, int argIndex) 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.
   int count(StringRef key) const;
 

mlir/lib/TableGen/Pattern.cpp

StringRef SymbolInfoMap::getValuePackName(StringRef symbol, int *index) {
   }
   if (index) {
     *index = idx;
   }
   return name;
 }
 
 SymbolInfoMap::SymbolInfo::SymbolInfo(const Operator *op, SymbolInfo::Kind kind,
-                                      Optional<int> index)
-    : op(op), kind(kind), argIndex(index) {}
+                                      Optional<DagAndIndex> dagAndIndex)
+    : op(op), kind(kind), dagAndIndex(dagAndIndex) {}
 
 int SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
   switch (kind) {
   case Kind::Attr:
   case Kind::Operand:
   case Kind::Value:
     return 1;
   case Kind::Result:
     return op->getNumResults();
   }
   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(*argIndex).get<NamedAttribute *>()->attr.getStorageType();
+      auto type = op->getArg((*dagAndIndex).second)
+                      .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).
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(*argIndex).get<NamedTypeConstraint *>();
+    auto *operand =
+        op->getArg((*dagAndIndex).second).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));
case Kind::Value: {
     auto repl = formatv(fmt, formatv("{{{0}}", name));
     LLVM_DEBUG(llvm::dbgs() << repl << " (Value)\n");
     return std::string(repl);
   }
   }
   llvm_unreachable("unknown kind");
 }
 
-bool SymbolInfoMap::bindOpArgument(StringRef symbol, const Operator &op,
-                                   int argIndex) {
+bool SymbolInfoMap::bindOpArgument(DagNode node, StringRef symbol,
+                                   const Operator &op, int argIndex) {
   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(&op, argIndex);
+                     : SymbolInfo::getOperand(node, &op, argIndex);
 
   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, const Operator &op,
+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.op == &op && it->second.argIndex == argIndex) {
+    if (it->second.dagAndIndex == symbolInfo.dagAndIndex) {
       return it;
     }
   }
 
   return symbolInfoMap.end();
 }
 
 std::pair<SymbolInfoMap::iterator, SymbolInfoMap::iterator>
for (int i = 0; i != numTreeArgs; ++i) {
       if (isSrcPattern) {
         // We can only bind symbols to op arguments in source pattern. Those
         // symbols are referenced in result patterns.
         auto treeArgName = tree.getArgName(i);
         // `$_` is a special symbol meaning ignore the current argument.
         if (!treeArgName.empty() && treeArgName != "_") {
           LLVM_DEBUG(llvm::dbgs() << "found symbol bound to op argument: "
                                   << treeArgName << '\n');
-          verifyBind(infoMap.bindOpArgument(treeArgName, op, i), treeArgName);
+          verifyBind(infoMap.bindOpArgument(tree, treeArgName, op, i),
+                     treeArgName);
         }
       }
     }
     return;
   }
 
   if (!treeName.empty()) {
     PrintFatalError(
         &def, formatv("binding symbol '{0}' to non-operation/native code call "
                       "unsupported right now",
                       treeName));
   }
   return;
 }

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

 
 // Test same operand name enforces equality condition check.
 def TestEqualArgsPattern : Pat<(OpN $a, $a), (OpO $a)>;
 
 // Test when equality is enforced at different depth.
 def TestNestedOpEqualArgsPattern :
   Pat<(OpN $b, (OpP $a, $b, $c, $d, $e, $f)), (replaceWithValue $b)>;
 
+// Test when equality is enforced on same op and same operand but at different
+// depth. We only bound one of the $x to the second operand of outer OpN and
+// left another be the default value (which is the value of first operand of
+// outer OpN). As a result, it ended up comparing wrong values in some cases.
+def TestNestedSameOpAndSameArgEqualityPattern :
+  Pat<(OpN (OpN $_, $x), $x), (replaceWithValue $x)>;
+
 // Test multiple equal arguments check enforced.
 def TestMultipleEqualArgsPattern :
   Pat<(OpP $a, $b, $a, $a, $b, $c), (OpN $c, $b)>;
 
 // Test for memrefs normalization of an op with normalizable memrefs.
 def OpNorm : TEST_Op<"op_norm", [MemRefsNormalizable]> {
   let arguments = (ins AnyMemRef:$X, AnyMemRef:$Y);
 }

mlir/test/mlir-tblgen/pattern.mlir

func @verifyNestedOpEqualArgs(
   // CHECK: test.op_n
   %2 = "test.op_p"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5)
     : (i32, i32, i32, i32, i32, i32) -> (i32)
   %3 = "test.op_n"(%arg0, %2) : (i32, i32) -> (i32)
 
   return
 }
 
+// CHECK-LABEL: verifyNestedSameOpAndSameArgEquality
+func @verifyNestedSameOpAndSameArgEquality(%arg0: i32, %arg1: i32) -> i32 {
+  // def TestNestedSameOpAndSameArgEqualityPattern:
+  //   Pat<(OpN (OpN $_, $x), $x), (replaceWithValue $x)>;
+
+  %0 = "test.op_n"(%arg1, %arg0) : (i32, i32) -> (i32)
+  %1 = "test.op_n"(%0, %arg0) : (i32, i32) -> (i32)
+  // CHECK: return %arg0 : i32
+  return %1 : i32
+}
+
 // CHECK-LABEL: verifyMultipleEqualArgs
 func @verifyMultipleEqualArgs(
   %arg0: i32, %arg1 : i32, %arg2 : i32, %arg3 : i32, %arg4 : i32) {
   // def TestMultipleEqualArgsPattern :
   //   Pat<(OpP $a, $b, $a, $a, $b, $c), (OpN $c, $b)>;
 
   // CHECK: "test.op_n"(%arg2, %arg1) : (i32, i32) -> i32
   "test.op_p"(%arg0, %arg1, %arg0, %arg0, %arg1, %arg2) :

mlir/tools/mlir-tblgen/RewriterGen.cpp

void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     // We need to subtract the number of attributes before this operand to get
     // the index in the operand list.
     auto numPrevAttrs = std::count_if(
         op.arg_begin(), op.arg_begin() + argIndex,
         [](const Argument &arg) { return arg.is<NamedAttribute *>(); });
 
-    auto res = symbolInfoMap.findBoundSymbol(name, op, argIndex);
+    auto res = symbolInfoMap.findBoundSymbol(name, tree, op, argIndex);
     os << formatv("{0} = {1}.getODSOperands({2});\n",
                   res->second.getVarName(name), opName,
                   argIndex - numPrevAttrs);
   }
 }
 
 void PatternEmitter::emitAttributeMatch(DagNode tree, StringRef opName,
                                         int argIndex, int depth) {