[drr] Add $_loc special directive for NativeCodeCall

[jpienaar]

Summary

Allows propagating the location to ops created via NativeCodeCall.

[jurahul]

2 documentation typos, otherwise looks good to me.

mlir/docs/DeclarativeRewrites.md

388

$_loc

391

typo: $_loc?

[jpienaar]

Fixing typos

mlir/docs/DeclarativeRewrites.md

 attribute. This might change in the future.
 
 ##### `NativeCodeCall` placeholders
 
 In `NativeCodeCall`, we can use placeholders like `$_builder`, `$N`. The former
 is called _special placeholder_, while the latter is called _positional
 placeholder_.
 
-`NativeCodeCall` right now only supports two special placeholders: `$_builder`
-and `$_self`:
+`NativeCodeCall` right now only supports three special placeholders:
+`$_builder`, `$_loc`, and `$_self`:
 
 *   `$_builder` will be replaced by the current `mlir::PatternRewriter`.
+*   `$_loc` will be replaced by the fused location or custom location (as
+    determined by location directive).
 *   `$_self` will be replaced with the entity `NativeCodeCall` is attached to.
 
 We have seen how `$_builder` can be used in the above; it allows us to pass a
 `mlir::Builder` (`mlir::PatternRewriter` is a subclass of `mlir::OpBuilder`,
 which is a subclass of `mlir::Builder`) to the C++ helper function to use the
 handy methods on `mlir::Builder`.
 
 `$_self` is useful when we want to write something in the form of

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, $0)"> $input), (OpK)]>;
+              [(NativeCodeCall<"createOpI($_builder, $_loc, $0)"> $input),
+               (OpK)]>;
 
 // 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);

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

 
 using namespace mlir;
 
 // Native function for testing NativeCodeCall
 static Value chooseOperand(Value input1, Value input2, BoolAttr choice) {
   return choice.getValue() ? input1 : input2;
 }
 
-static void createOpI(PatternRewriter &rewriter, Value input) {
-  rewriter.create<OpI>(rewriter.getUnknownLoc(), input);
+static void createOpI(PatternRewriter &rewriter, Location loc, Value input) {
+  rewriter.create<OpI>(loc, input);
 }
 
 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());
 }

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
   // Emits the C++ statement to replace the matched DAG with a value built via
   // calling native C++ code.
   std::string handleReplaceWithNativeCodeCall(DagNode resultTree);
 
   // Returns the symbol of the old value serving as the replacement.
   StringRef handleReplaceWithValue(DagNode tree);
 
   // Returns the location value to use.
+  std::pair<bool, std::string> getLocation(DagNode tree);
+
+  // Returns the location value to use.
   std::string handleLocationDirective(DagNode tree);
 
   // Emits the C++ statement to build a new op out of the given DAG `tree` and
   // returns the variable name that this op is assigned to. If the root op in
   // DAG `tree` has a specified name, the created op will be assigned to a
   // variable of the given name. Otherwise, a unique name will be used as the
   // result value name.
   std::string handleOpCreation(DagNode tree, int resultIndex, int depth);
std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree) {
 
   auto fmt = tree.getNativeCodeTemplate();
   // TODO: replace formatv arguments with the exact specified args.
   SmallVector<std::string, 8> attrs(8);
   if (tree.getNumArgs() > 8) {
     PrintFatalError(loc, "unsupported NativeCodeCall argument numbers: " +
                              Twine(tree.getNumArgs()));
   }
-  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+  bool hasLocationDirective;
+  std::string locToUse;
+  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
+
+  for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
     attrs[i] = handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i));
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
                             << " replacement: " << attrs[i] << "\n");
   }
-  return std::string(tgfmt(fmt, &fmtCtx, attrs[0], attrs[1], attrs[2], attrs[3],
-                           attrs[4], attrs[5], attrs[6], attrs[7]));
+  return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs[0],
+                           attrs[1], attrs[2], attrs[3], attrs[4], attrs[5],
+                           attrs[6], attrs[7]));
 }
 
 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.
   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))
+      if (lastArg.isLocationDirective()) {
+        return std::make_pair(true, handleLocationDirective(lastArg));
+      }
+  }
+
+  // If no explicit location is given, use the default, all fused, location.
+  return std::make_pair(false, "odsLoc");
+}
+
 std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
                                              int depth) {
   LLVM_DEBUG(llvm::dbgs() << "create op for pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   Operator &resultOp = tree.getDialectOp(opMap);
   auto numOpArgs = resultOp.getNumArgs();
   auto numPatArgs = tree.getNumArgs();
 
-  // Get the location for this operation if explicitly provided.
+  bool hasLocationDirective;
   std::string locToUse;
-  if (numPatArgs != 0) {
-    if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
-      if (lastArg.isLocationDirective())
-        locToUse = handleLocationDirective(lastArg);
-  }
+  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
-  auto inPattern = numPatArgs - !locToUse.empty();
+  auto inPattern = numPatArgs - hasLocationDirective;
   if (numOpArgs != inPattern) {
     PrintFatalError(loc,
                     formatv("resultant op '{0}' argument number mismatch: "
                             "{1} in pattern vs. {2} in definition",
                             resultOp.getOperationName(), inPattern, numOpArgs));
   }
 
-  // If no explicit location is given, use the default, all fused, location.
-  if (locToUse.empty())
-    locToUse = "odsLoc";
-
   // A map to collect all nested DAG child nodes' names, with operand index as
   // the key. This includes both bound and unbound child nodes.
   ChildNodeIndexNameMap childNodeNames;
 
   // First go through all the child nodes who are nested DAG constructs to
   // create ops for them and remember the symbol names for them, so that we can
   // use the results in the current node. This happens in a recursive manner.
   for (int i = 0, e = resultOp.getNumOperands(); i != e; ++i) {