Support NativeCodeCall binding in rewrite pattern.

[Chia-hungDuan]

Summary

We are able to bind the result from native function while rewriting
pattern. In matching pattern, if we want to get some values back, we can
do that by passing parameter as return value placeholder. Beside, add
the semantic of '$_self' in NativeCodeCall while matching, it'll be the
operation that defines certain operand.

[jpienaar]

Is there perhaps some sensible structure that could be verified from the emitted C++ code? Not a change test but something capturing the ordering, matching & rewrite?

mlir/include/mlir/IR/OpBase.td

2504

Just to verify: the below is now consistent with this and https://mlir.llvm.org/docs/DeclarativeRewrites/ seems to also be up to date (although the sections both are more rewrite based.

mlir/lib/TableGen/Pattern.cpp

675

What is difference between this and line 670?

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

851

OOC what happens if this was

def : Pat<(SomeBinaryOp (WriteReturnValueToArgument $ret)), $ret), ... > ? (I'm assuming the equality checking would kick in but not sure, and esp if reversed)

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

38

Would a Value& also work?

38

"write" is a bit confusing to me, how about getFirstResult?

39

Lets create a negative result too. E.g., if op has even number of results return false, so that we can check the filtering.

mlir/test/mlir-tblgen/pattern.mlir

94

Nit: capture this one too

mlir/test/mlir-tblgen/rewriter-errors.td

28

And this would be (effectively) $_self ? (well if the operation had single result)

mlir/tools/mlir-tblgen/RewriterGen.cpp

282

operand

[bondhugula]

Wouldn't this need an update to the DRR documentation examples?
https://mlir.llvm.org/docs/DeclarativeRewrites/

[Chia-hungDuan]

Applied review comments and updated the document.

[Chia-hungDuan]

mlir/include/mlir/IR/OpBase.td

2504

Yes, updated the description about $_self above to make this more clear.

mlir/lib/TableGen/Pattern.cpp

675

Line 670 is to have an Attribute variable bound with symbol $a
Line 661 and Line 675 are Value variable bound with symbol $c and $b respectively. The difference between them is $b has type constraint on it.

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

38

We can do that. the reason I use pointer type is just to align with ConstantLikeMatcher's behavior. Do you think using reference will be better?

39

Done. Added another case which will fail rewriting.

mlir/test/mlir-tblgen/rewriter-errors.td

28

Yes, the difference is if we want to use the $_self later, we should set it up in $val in this case

[Chia-hungDuan]

mlir/docs/DeclarativeRewrites.md

409–410 ↗

(On Diff #339962)

This is not supported now and I removed it. Actually we can do something like

NativeCodeCall<"cast<ArrayAttr>($0)[" # n # "]">

to get the same function. If we want to make simpler to use, I can do it in the later CL. Feel free to let me know.

[bondhugula]

A few minor comments related the documentation, etc. I'll defer to @jpienaar to make a pass on the doc addition.

mlir/docs/DeclarativeRewrites.md

395 ↗

(On Diff #339962)

Nit: in a source pattern ?

413 ↗

(On Diff #339962)

-> ... the names can be put in ...

415–416 ↗

(On Diff #339962)

These lines appear broken - missing an article or other things punctuation-wise. Can you please rephrase?

mlir/lib/TableGen/Pattern.cpp

659

then block should use braces if the else block does.

[Chia-hungDuan]

Revised the comments

[Chia-hungDuan]

@bondhugula

Thanks for reviewing the sentences and sorry for my bad writing. Can you take a look the updated one and see if it's more clear now?

Thanks

Chia-hung

[jpienaar]

mlir/docs/DeclarativeRewrites.md

395 ↗

(On Diff #339997)

Lets keep the style consistent here

$_self will be replaced by the defining operation in a source pattern.

396 ↗

(On Diff #339997)

Why the defining operation vs the Value matched? Operands are Values rather than Operations.

413–414 ↗

(On Diff #339997)

I'm not sure what this sentence means.

414 ↗

(On Diff #339997)

Nit: space before (

416 ↗

(On Diff #339997)

Is reference actually required?

419–420 ↗

(On Diff #339997)

Could you phrase this more directly? Names with attribute constraints will be captured as Attributes while everything else will be treated as Value. ?

mlir/include/mlir/IR/OpBase.td

2501

If you update the above, update this too to match

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

38

I think so. At least this isn't a hard requirement. And makes it simpler to write ODS side I feel.

[Chia-hungDuan]

mlir/docs/DeclarativeRewrites.md

396 ↗

(On Diff #339997)

This is defined by the implementation of tblgen, we generate code like,

auto *op1 = (*castedOp0.getODSOperands(0).begin()).getDefiningOp();

I have a concern about this, not all Values have a defining operation, e.g., function arguments, right? I'm not sure if we had any reason to generate this pattern. Do you think we need to change it to Value? (This may slightly impact people who use NativeCodeCall in source pattern)

413–414 ↗

(On Diff #339997)

Rephrase to "To carry some return values from helper function"

416 ↗

(On Diff #339997)

Yes, this is also defined by implementation. I'm thinking maybe passing by reference by default will be better, but leave it to the user will be easier for tblgen(not complicate to change though).

I'm thinking to keep the same logic as mush as possible so I didn't change this one and the above one.

419–420 ↗

(On Diff #339997)

Thanks, this is much better!

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

851

Yes, we will get the message like "symbol ret bound more than once" because it causes ambiguity

[Chia-hungDuan]

Addressed review comments

[Chia-hungDuan]

mlir/docs/DeclarativeRewrites.md

396 ↗

(On Diff #339997)

One thing to add, for pattern like,

def : Pat<(Op1 (Op2 $arg1), $arg2), ...>

we are generating code like
auto castedOp0 = ::llvm::dyn_cast_or_null<Op1>(op0);
auto *op1 = (*castedOp0.getODSOperands(0).begin()).getDefiningOp();
auto castedOp1 = ::llvm::dyn_cast_or_null<Op2>(op1);

So match the operand is matching on the defining op, it's not bound to the NativeCodeCall only. it applies to all the nested structure.

416 ↗

(On Diff #339997)

Please ignore the above reply, I messed different things together. It shouldn't be 'reference'. I would like to say either pass-by-reference or pass-by-pointer should be fine. Fixed.

[Chia-hungDuan]

Addressed missed/wrong replied review comments

[jpienaar]

Thanks!

mlir/docs/DeclarativeRewrites.md

414 ↗

(On Diff #342288)

s/can/must/ ? (part of this is just C++, this is only way to set variable passed in)

must be either passed by reference or a pointer to variable used as argument so that the matched value can be returned.

417 ↗

(On Diff #342288)

argument

396 ↗

(On Diff #339997)

I think it is actually always on defining op at the moment - as these correspond to nodes of the dag pattern. I think we will probably need to get back to this indeed. So for now this seems accurate.

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

851

It would be good to file a bug for this, we do allow (Foo $ret, $ret) to mean match Foo, bind $ret to first arg and verify that matches second arg

mlir/tools/mlir-tblgen/RewriterGen.cpp

282

What does "of certain operand" mean here? Do we not need $_self if no operands are referenced? (else I'd just remove the end)

[Chia-hungDuan]

mlir/docs/DeclarativeRewrites.md

396 ↗

(On Diff #339997)

Added a TODO in RewriterGen.cpp

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

851

https://bugs.llvm.org/show_bug.cgi?id=50283

mlir/tools/mlir-tblgen/RewriterGen.cpp

282

Yes, no operands we still need it, it's just like this pointer. trailing words removed.

[Chia-hungDuan]

Addressed review comments

mlir/include/mlir/IR/OpBase.td

 // ## Placeholders
 //
 // If used as a DAG leaf, i.e., `(... NativeCodeCall<"...">:$arg, ...)`,
 // the wrapped expression can take special placeholders listed below:
 //
 // * `$_builder` will be replaced by the current `mlir::PatternRewriter`.
 // * `$_self` will be replaced with the entity this transformer is attached to.
 //   E.g., with the definition `def transform : NativeCodeCall<"$_self...">`,
 //   `$_self` in `transform:$attr` will be replaced by the value for `$attr`.
 //
 // 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>`.
 
 class NativeCodeCall<string expr> {
   string expression = expr;
 }
 
-def ConstantLikeMatcher : NativeCodeCall<"success(matchPattern($0->getResult(0), m_Constant(&$1)))">;
+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
 // DAG construct. The arguments to location must be previously captured symbol.

mlir/lib/TableGen/Pattern.cpp

std::string SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
   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("::llvm::ArrayRef<::mlir::Value> {0};\n", name));
+    return std::string(formatv("::mlir::Value {0};\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");
 }

 
 void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
                                   bool isSrcPattern) {
   auto treeName = tree.getSymbol();
   auto numTreeArgs = tree.getNumArgs();
 
   if (tree.isNativeCodeCall()) {
     if (!treeName.empty()) {
-      PrintFatalError(
-          &def,
-          formatv(
-              "binding symbol '{0}' to native code call unsupported right now",
+      if (!isSrcPattern) {
+        LLVM_DEBUG(llvm::dbgs() << "found symbol bound to NativeCodeCall: "
+                                << treeName << '\n');
+        verifyBind(infoMap.bindValue(treeName), treeName);
+      } else
+        PrintFatalError(&def,
+                        formatv("binding symbol '{0}' to NativecodeCall in "
+                                "MatchPattern is not supported",
                                 treeName));
     }
 
     for (int i = 0; i != numTreeArgs; ++i) {
       if (auto treeArg = tree.getArgAsNestedDag(i)) {
         // This DAG node argument is a DAG node itself. Go inside recursively.
         collectBoundSymbols(treeArg, infoMap, isSrcPattern);
         continue;
       }
 
       if (!isSrcPattern)
         continue;
 
       // We can only bind symbols to 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 != "_") {
-        if (tree.isNestedDagArg(i)) {
-          auto err = formatv("cannot bind '{0}' for nested native call arg",
-                             treeArgName);
-          PrintFatalError(&def, err);
-        }
-
         DagLeaf leaf = tree.getArgAsLeaf(i);
+
+        // In (NativeCodeCall<"Foo($_self, $0, $1, $2)"> I8Attr:$a, I8:$b, $c),
+        if (leaf.isUnspecified())
+          // This is case of $c, a Value without any constraints.
+          verifyBind(infoMap.bindValue(treeArgName), treeArgName);
+        else {
           auto constraint = leaf.getAsConstraint();
           bool isAttr = leaf.isAttrMatcher() || leaf.isEnumAttrCase() ||
                         leaf.isConstantAttr() ||
                         constraint.getKind() == Constraint::Kind::CK_Attr;
 
           if (isAttr) {
+            // This is case of $a, a binding to a certain attribute.
             verifyBind(infoMap.bindAttr(treeArgName), treeArgName);
             continue;
           }
 
+          // This is case of $b, a binding to a certain type.
           verifyBind(infoMap.bindValue(treeArgName), treeArgName);
         }
       }
+    }
 
     return;
   }
 
   if (tree.isOperation()) {
     auto &op = getDialectOp(tree);
     auto numOpArgs = op.getNumArgs();
 

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

def OpNativeCodeCall3 : TEST_Op<"native_code_call3"> {
   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),
                (OpK)]>;
 
+def OpNativeCodeCall4 : TEST_Op<"native_code_call4"> {
+  let arguments = (ins I32:$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 WriteReturnValueToArgument : NativeCodeCall<"success(writeReturnValueToArgument($_self, &$0))">;
+def BindNativeCodeCallResult : NativeCodeCall<"bindNativeCodeCallResult($0)">;
+def : Pat<(OpNativeCodeCall4 (WriteReturnValueToArgument $ret)),
+          (OpNativeCodeCall5 (BindNativeCodeCallResult:$native $ret), $native)>;
+
 // 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);

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

 
 static void handleNoResultOp(PatternRewriter &rewriter,
                              OpSymbolBindingNoResult op) {
   // Turn the no result op to a one-result op.
   rewriter.create<OpSymbolBindingB>(op.getLoc(), op.operand().getType(),
                                     op.operand());
 }
 
+static bool writeReturnValueToArgument(Operation *op, Value *value) {
+  *value = op->getResult(0);
+  return true;
+}
+
+static Value bindNativeCodeCallResult(Value value) { return value; }
+
 // 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

 // CHECK-LABEL: verifyAuxiliaryNativeCodeCall
 func @verifyAuxiliaryNativeCodeCall(%arg0: i32) -> (i32) {
   // CHECK: test.op_i
   // CHECK: test.op_k
   %0 = "test.native_code_call3"(%arg0) : (i32) -> (i32)
   return %0 : i32
 }
 
+// CHECK-LABEL: verifyNativeCodeCallBinding
+func @verifyNativeCodeCallBinding(%arg0 : i32) -> (i32) {
+  %0 = "test.op_k"() : () -> (i32)
+  // CHECK: %[[A:.*]], %[[B:.*]] = "test.native_code_call5"(%0, %0)
+  // CHECK: return %[[A]]
+  %1, %2 = "test.native_code_call4"(%0) : (i32) -> (i32, i32)
+  return %1 : 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/test/mlir-tblgen/rewriter-errors.td

 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR1 %s 2>&1 | FileCheck --check-prefix=ERROR1 %s
 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR2 %s 2>&1 | FileCheck --check-prefix=ERROR2 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR3 %s 2>&1 | FileCheck --check-prefix=ERROR3 %s
 
 include "mlir/IR/OpBase.td"
 
 // Check using the dialect name as the namespace
 def A_Dialect : Dialect {
   let name = "a";
 }
 
 class A_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<A_Dialect, mnemonic, traits>;
 
 def OpA : A_Op<"op_a">, Arguments<(ins AnyInteger, AnyInteger)>, Results<(outs AnyInteger)>;
 def OpB : A_Op<"op_b">, Arguments<(ins AnyInteger, AnyAttr:$value)>, Results<(outs AnyInteger)>;
 
 #ifdef ERROR1
 def NativeMatcher : NativeCodeCall<"success(nativeCall($0, $1))">;
-// ERROR1: [[@LINE+1]]:1: error: binding symbol 'error' to native code call unsupported right now
-def : Pat<(OpA (NativeMatcher:$error $val), AnyI32Attr:$arg),
+// ERROR1: [[@LINE+1]]:1: error: NativeCodeCall must have $_self as argument for passing the Operation of certain oprand
+def : Pat<(OpA (NativeMatcher $val), AnyI32Attr:$arg),
           (OpB $val, $arg)>;
 #endif
 
 #ifdef ERROR2
-def NativeMatcher : NativeCodeCall<"success(nativeCall($0, $1))">;
-// ERROR2: [[@LINE+1]]:1: error: Matching nested tree in NativeCodecall not support for 
+def NativeMatcher : NativeCodeCall<"success(nativeCall($_self, $0, $1))">;
+// ERROR2: [[@LINE+1]]:1: error: binding symbol 'error' to NativecodeCall in MatchPattern is not supported
+def : Pat<(OpA (NativeMatcher:$error $val), AnyI32Attr:$arg),
+          (OpB $val, $arg)>;
+#endif
+
+#ifdef ERROR3
+def NativeMatcher : NativeCodeCall<"success(nativeCall($_self, $0, $1))">;
+// ERROR3: [[@LINE+1]]:1: error: Matching nested tree in NativeCodecall not support for
 def : Pat<(OpA (NativeMatcher (OpB $val, $unused)), AnyI32Attr:$arg),
           (OpB $val, $arg)>;
 #endif

mlir/tools/mlir-tblgen/RewriterGen.cpp

void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,
                                          int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   // TODO(suderman): iterate through arguments, determine their types, output
   // names.
   SmallVector<std::string, 8> capture;
-  capture.push_back(opName.str());
 
   raw_indented_ostream::DelimitedScope scope(os);
 
   os << "if(!" << opName << ") return ::mlir::failure();\n";
   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)) {
       os << "Value " << argName << ";\n";
     } else {
       auto leaf = tree.getArgAsLeaf(i);
       if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {
         os << "Attribute " << argName << ";\n";
-      } else if (leaf.isOperandMatcher()) {
-        os << "Operation " << argName << ";\n";
+      } else {
+        os << "Value " << argName << ";\n";
       }
     }
 
     capture.push_back(std::move(argName));
   }
 
   bool hasLocationDirective;
   std::string locToUse;
   std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
   auto fmt = tree.getNativeCodeTemplate();
-  auto nativeCodeCall =
-      std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), capture));
+  if (fmt.count("$_self") != 1) {
+    PrintFatalError(loc, "NativeCodeCall must have $_self as argument for "
+                         "passing the Operation of certain oprand");
+  }
+
+  auto nativeCodeCall = std::string(tgfmt(
+      fmt, &fmtCtx.addSubst("_loc", locToUse).withSelf(opName.str()), capture));
 
   os << "if (failed(" << nativeCodeCall << ")) return ::mlir::failure();\n";
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     auto name = tree.getArgName(i);
     if (!name.empty() && name != "_") {
-      os << formatv("{0} = {1};\n", name, capture[i + 1]);
+      os << formatv("{0} = {1};\n", name, capture[i]);
     }
   }
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
-    std::string argName = capture[i + 1];
+    std::string argName = capture[i];
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       PrintFatalError(
           loc, formatv("Matching nested tree in NativeCodecall not support for "
                        "{0} as arg {1}",
                        argName, i));
     }
 
     DagLeaf leaf = tree.getArgAsLeaf(i);
+
+    // The parameter for native function doesn't bind any constraints.
+    if (leaf.isUnspecified())
+      continue;
+
     auto constraint = leaf.getAsConstraint();
 
-    auto self = formatv("{0}", argName);
+    std::string self;
+    if (leaf.isAttrMatcher() || leaf.isConstantAttr())
+      self = argName;
+    else
+      self = formatv("{0}.getType()", argName);
     emitMatchCheck(
         opName,
         tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),
         formatv("\"operand {0} of native code call '{1}' failed to satisfy "
                 "constraint: "
                 "'{2}'\"",
                 i, tree.getNativeCodeTemplate(), constraint.getSummary()));
   }
for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
     } else {
       attrs.push_back(
           handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)));
     }
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
                             << " replacement: " << attrs[i] << "\n");
   }
 
-  return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs));
+  std::string symbol = tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs);
+  if (!tree.getSymbol().empty()) {
+    os << formatv("auto {0} = {1};\n", tree.getSymbol(), symbol);
+    symbol = tree.getSymbol().str();
+  }
+
+  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()) {