[mlir][DRR] Add location directive

[jpienaar]

Summary

Add directive to indicate the location to give to op being created. This
directive is optional and if unused the location will still be the fused
location of all source operations.

Currently this directive only works with other op locations, reusing an
existing op location or a fusion of op locations. But doesn't yet support
supplying metadata for the FusedLoc.

Based off initial revision by antiagainst@ and effectively mirrors GlobalIsel
debug_locations directive.

[rriddle]

mlir/docs/DeclarativeRewrites.md

677

Why the last argument? May be weird to come from the C++ world where it is always the first argument.

684

nit: typo the matched the location

mlir/tools/mlir-tblgen/RewriterGen.cpp

696

nit: Remove trivial braces.

[jpienaar]

Fixed typo and removed trivial braces.

[antiagainst]

Cool! Thanks for pushing forward on this! :)

mlir/docs/DeclarativeRewrites.md

671

typo: tablegen

mlir/include/mlir/IR/OpBase.td

2187

symbols

mlir/docs/DeclarativeRewrites.md

 
 *   Larger matches are more beneficial than smaller ones.
 *   If a smaller one is applied first the larger one may not apply anymore.
 
 
 The fourth parameter to `Pattern` (and `Pat`) allows to manually tweak a
 pattern's benefit. Just supply `(addBenefit N)` to add `N` to the benefit value.
 
-## Special directives
+## Rewrite directives
 
-[TODO]
+### `location`
+
+By default the C++ pattern expanded from a DRR pattern uses the fused location
+of all source ops as the location for all generated ops. This is not always the
+best location mapping relationship. For such cases, DRR provides the `location`
+directive to provide finer control.
+
+`location` is of the following syntax:
+
+```tablgen
+(location $symbol0, $symbol1, ...)
+```
+
+where all `$symbol` should be bound previously in the pattern.
+
+`location` must be used as the last argument to an op creation. For example,
+
+```tablegen
+def : Pat<(LocSrc1Op:$src1 (LocSrc2Op:$src2 ...),
+          (LocDst1Op (LocDst2Op ..., (location $src2)))>;
+```
+
+In the above pattern, the generated `LocDst2Op` will use the matched location
+of `LocSrc2Op` while the root `LocDst1Op` node will still se the fused location
+of all source Ops.
+
+### `replaceWithValue`
+
+The `replaceWithValue` directive is used to eliminate a matched op by replacing
+all of it uses with a captured value. It is of the following syntax:
+
+```tablegen
+(replaceWithValue $symbol)
+```
+
+where `$symbol` should be a symbol bound previously in the pattern.
+
+For example,
+
+```tablegen
+def : Pat<(Foo $input), (replaceWithValue $input)>;
+```
+
+The above pattern removes the `Foo` and replaces all uses of `Foo` with
+`$input`.
 
 ## Debugging Tips
 
 ### Run `mlir-tblgen` to see the generated content
 
 TableGen syntax sometimes can be obscure; reading the generated content can be
 a very helpful way to understand and debug issues. To build `mlir-tblgen`, run
 `cmake --build . --target mlir-tblgen` in your build directory and find the

mlir/include/mlir/IR/OpBase.td

 // 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;
 }
 
 //===----------------------------------------------------------------------===//
-// Common directives
+// 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.
+def location;
+
 // Directive used in result pattern to indicate that no new op are generated,
 // so to replace the matched DAG with an existing SSA value.
 def replaceWithValue;
 
 #endif // OP_BASE

mlir/include/mlir/TableGen/Pattern.h

public:
 
   // Returns the specified name of the `index`-th argument.
   StringRef getArgName(unsigned index) const;
 
   // Returns true if this DAG construct means to replace with an existing SSA
   // value.
   bool isReplaceWithValue() const;
 
+  // Returns whether this DAG represents the location of an op creation.
+  bool isLocationDirective() const;
+
   // Returns true if this DAG node is wrapping native code call.
   bool isNativeCodeCall() 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()

mlir/lib/TableGen/Pattern.cpp

 
 bool tblgen::DagNode::isNativeCodeCall() const {
   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(node->getOperator()))
     return defInit->getDef()->isSubClassOf("NativeCodeCall");
   return false;
 }
 
 bool tblgen::DagNode::isOperation() const {
-  return !(isNativeCodeCall() || isReplaceWithValue());
+  return !isNativeCodeCall() && !isReplaceWithValue() && !isLocationDirective();
 }
 
 llvm::StringRef tblgen::DagNode::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsString("expression");
 }
StringRef tblgen::DagNode::getArgName(unsigned index) const {
   return node->getArgNameStr(index);
 }
 
 bool tblgen::DagNode::isReplaceWithValue() const {
   auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
   return dagOpDef->getName() == "replaceWithValue";
 }
 
+bool tblgen::DagNode::isLocationDirective() const {
+  auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
+  return dagOpDef->getName() == "location";
+}
+
 void tblgen::DagNode::print(raw_ostream &os) const {
   if (node)
     node->print(os);
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolInfoMap
 //===----------------------------------------------------------------------===//
if (!tree.isOperation()) {
     }
     return;
   }
 
   auto &op = getDialectOp(tree);
   auto numOpArgs = op.getNumArgs();
   auto numTreeArgs = tree.getNumArgs();
 
-  if (numOpArgs != numTreeArgs) {
+  // The pattern might have the last argument specifying the location.
+  bool hasLocDirective = false;
+  if (numTreeArgs != 0) {
+    if (auto lastArg = tree.getArgAsNestedDag(numTreeArgs - 1))
+      hasLocDirective = lastArg.isLocationDirective();
+  }
+
+  if (numOpArgs != numTreeArgs - hasLocDirective) {
     auto err = formatv("op '{0}' argument number mismatch: "
                        "{1} in pattern vs. {2} in definition",
                        op.getOperationName(), numTreeArgs, numOpArgs);
     PrintFatalError(def.getLoc(), err);
   }
 
   // The name attached to the DAG node's operator is for representing the
   // results generated from this op. It should be remembered as bound results.

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

: TEST_Op<"string_attr_pretty_name",
   let arguments = (ins StrArrayAttr:$names);
   let results = (outs Variadic<I32>:$r);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 //===----------------------------------------------------------------------===//
+// Test Locations
+//===----------------------------------------------------------------------===//
+
+def TestLocationSrcOp : TEST_Op<"loc_src"> {
+  let arguments = (ins I32:$input);
+  let results = (outs I32:$output);
+}
+
+def TestLocationDstOp : TEST_Op<"loc_dst", [SameOperandsAndResultType]> {
+  let arguments = (ins I32:$input);
+  let results = (outs I32:$output);
+}
+
+//===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//
 
 def OpA : TEST_Op<"op_a"> {
   let arguments = (ins I32, I32Attr:$attr);
   let results = (outs I32);
 }
 

 def : Pat<(OneI32ResultOp),
           (MixedVOperandOp3
               (MixedVResultOp3:$results__0 ConstantAttr<I32Attr, "2">),
               (replaceWithValue $results__1),
               (replaceWithValue $results__2),
               ConstantAttr<I32Attr, "2">)>;
 
 //===----------------------------------------------------------------------===//
+// Test Patterns (Location)
+
+// Test that we can specify locations for generated ops.
+def : Pat<(TestLocationSrcOp:$res1
+           (TestLocationSrcOp:$res2
+            (TestLocationSrcOp:$res3 $input))),
+          (TestLocationDstOp
+            (TestLocationDstOp
+              (TestLocationDstOp $input, (location $res1))),
+      (location $res2, $res3))>;
+
+//===----------------------------------------------------------------------===//
 // Test Legalization
 //===----------------------------------------------------------------------===//
 
 def Test_LegalizerEnum_Success : StrEnumAttrCase<"Success">;
 def Test_LegalizerEnum_Failure : StrEnumAttrCase<"Failure">;
 
 def Test_LegalizerEnum : StrEnumAttr<"Success", "Failure",
   [Test_LegalizerEnum_Success, Test_LegalizerEnum_Failure]>;

mlir/test/mlir-tblgen/pattern.mlir

-// RUN: mlir-opt -test-patterns -mlir-print-debuginfo %s | FileCheck %s
+// RUN: mlir-opt -test-patterns -mlir-print-debuginfo %s | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: verifyFusedLocs
 func @verifyFusedLocs(%arg0 : i32) -> i32 {
   %0 = "test.op_a"(%arg0) {attr = 10 : i32} : (i32) -> i32 loc("a")
   %result = "test.op_a"(%0) {attr = 20 : i32} : (i32) -> i32 loc("b")
 
   // CHECK: "test.op_b"(%arg0) {attr = 10 : i32} : (i32) -> i32 loc("a")
   // CHECK: "test.op_b"(%arg0) {attr = 20 : i32} : (i32) -> i32 loc(fused["b", "a"])
   return %result : i32
 }
 
+// CHECK-LABEL: verifyDesignatedLoc
+func @verifyDesignatedLoc(%arg0 : i32) -> i32 {
+  %0 = "test.loc_src"(%arg0) : (i32) -> i32 loc("loc3")
+  %1 = "test.loc_src"(%0) : (i32) -> i32 loc("loc2")
+  %2 = "test.loc_src"(%1) : (i32) -> i32 loc("loc1")
+
+  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc("loc1")
+  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused[
+  // CHECK-SAME: "loc1"
+  // CHECK-SAME: "loc3"
+  // CHECK-SAME: "loc2"
+  // CHECK: "test.loc_dst"({{.*}}) : (i32) -> i32 loc(fused["loc2", "loc3"])
+  return %1 : i32
+}
+
 // CHECK-LABEL: verifyZeroResult
 func @verifyZeroResult(%arg0 : i32) {
   // CHECK: "test.op_i"(%arg0) : (i32) -> ()
   "test.op_h"(%arg0) : (i32) -> ()
   return
 }
 
 // CHECK-LABEL verifyZeroArg

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
   // pattern.
   std::string handleResultPattern(DagNode resultTree, int resultIndex,
                                   int depth);
 
   // 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 C++ expression referencing the old value serving as the
-  // replacement.
-  std::string handleReplaceWithValue(DagNode tree);
+  // Returns the symbol of the old value serving as the replacement.
+  StringRef handleReplaceWithValue(DagNode tree);
+
+  // Returns the symbol of the value whose location to use.
+  std::string handleUseLocationOf(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);
 
for (int i = numResultPatterns - 1; i >= 0; --i) {
     }
   }
 
   if (offsets.front() > 0) {
     const char error[] = "no enough values generated to replace the matched op";
     PrintFatalError(loc, error);
   }
 
-  os.indent(4) << "auto loc = rewriter.getFusedLoc({";
+  os.indent(4) << "auto odsLoc = rewriter.getFusedLoc({";
   for (int i = 0, e = pattern.getSourcePattern().getNumOps(); i != e; ++i) {
     os << (i ? ", " : "") << "tblgen_ops[" << i << "]->getLoc()";
   }
-  os << "}); (void)loc;\n";
+  os << "}); (void)odsLoc;\n";
 
   // 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

 }
 
 std::string PatternEmitter::handleResultPattern(DagNode resultTree,
                                                 int resultIndex, int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle result pattern: ");
   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);
     symbolInfoMap.bindValue(symbol);
     return symbol;
   }
 
-  if (resultTree.isReplaceWithValue()) {
-    return handleReplaceWithValue(resultTree);
-  }
+  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.
     // Register the auto-generated symbol for it.
     symbolInfoMap.bindOpResult(symbol, pattern.getDialectOp(resultTree));
   }
   return symbol;
 }
 
-std::string PatternEmitter::handleReplaceWithValue(DagNode tree) {
+StringRef PatternEmitter::handleReplaceWithValue(DagNode tree) {
   assert(tree.isReplaceWithValue());
 
   if (tree.getNumArgs() != 1) {
     PrintFatalError(
         loc, "replaceWithValue directive must take exactly one argument");
   }
 
   if (!tree.getSymbol().empty()) {
     PrintFatalError(loc, "cannot bind symbol to replaceWithValue");
   }
 
-  return std::string(tree.getArgName(0));
+  return tree.getArgName(0);
+}
+
+std::string PatternEmitter::handleUseLocationOf(DagNode tree) {
+  assert(tree.isLocationDirective());
+  auto lookUpArgLoc = [this, &tree](int idx) {
+    const auto *const lookupFmt = "(*{0}.begin()).getLoc()";
+    return symbolInfoMap.getAllRangeUse(tree.getArgName(idx), lookupFmt);
+  };
+
+  if (tree.getNumArgs() != 1) {
+    std::string ret;
+    llvm::raw_string_ostream os(ret);
+    os << "rewriter.getFusedLoc({";
+    for (int i = 0, e = tree.getNumArgs(); i != e; ++i)
+      os << (i ? ", " : "") << lookUpArgLoc(i);
+    os << "})";
+    return os.str();
+  }
+
+  if (!tree.getSymbol().empty())
+    PrintFatalError(loc, "cannot bind symbol to location");
+
+  return lookUpArgLoc(0);
 }
 
 std::string PatternEmitter::handleOpArgument(DagLeaf leaf,
                                              StringRef patArgName) {
   if (leaf.isConstantAttr()) {
     auto constAttr = leaf.getAsConstantAttr();
     return handleConstantAttr(constAttr.getAttribute(),
                               constAttr.getConstantValue());

 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();
 
-  if (numOpArgs != tree.getNumArgs()) {
-    PrintFatalError(loc, formatv("resultant op '{0}' argument number mismatch: "
-                                 "{1} in pattern vs. {2} in definition",
-                                 resultOp.getOperationName(), tree.getNumArgs(),
-                                 numOpArgs));
+  // Get the location for this operation if explicitly provided.
+  std::string locToUse;
+  if (numPatArgs != 0) {
+    if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
+      if (lastArg.isLocationDirective())
+        locToUse = handleUseLocationOf(lastArg);
+  }
+
+  auto inPattern = numPatArgs - !locToUse.empty();
+  if (numOpArgs != inPattern) {
+    PrintFatalError(loc,
+                    formatv("resultant op '{0}' argument number mismatch: "
+                            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) {
-    if (auto child = tree.getArgAsNestedDag(i)) {
+    if (auto child = tree.getArgAsNestedDag(i))
       childNodeNames[i] = handleResultPattern(child, i, depth + 1);
-    }
-  }
 
   // The name of the local variable holding this op.
   std::string valuePackName;
   // The symbol for holding the result of this pattern. Note that the result of
   // this pattern is not necessarily the same as the variable created by this
   // pattern because we can use `__N` suffix to refer only a specific result if
   // the generated op is a multi-result op.
   std::string resultValue;
std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
 
   if (isSameOperandsAndResultType || useFirstAttr) {
     // We know how to deduce the result type for ops with these traits and we've
     // generated builders taking aggregate parameters. Use those builders to
     // create the ops.
 
     // First prepare local variables for op arguments used in builder call.
     createAggregateLocalVarsForOpArgs(tree, childNodeNames);
+
     // Then create the op.
     os.indent(6) << formatv(
-        "{0} = rewriter.create<{1}>(loc, tblgen_values, tblgen_attrs);\n",
-        valuePackName, resultOp.getQualCppClassName());
+        "{0} = rewriter.create<{1}>({2}, tblgen_values, tblgen_attrs);\n",
+        valuePackName, resultOp.getQualCppClassName(), locToUse);
     os.indent(4) << "}\n";
     return resultValue;
   }
 
   bool usePartialResults = valuePackName != resultValue;
 
   if (usePartialResults || depth > 0 || resultIndex < 0) {
     // For these cases (broadcastable ops, op results used both as auxiliary
     // values and replacement values, ops in nested patterns, auxiliary ops), we
     // still need to supply the result types when building the op. But because
     // we don't generate a builder automatically with ODS for them, it's the
     // developer's responsibility to make sure such a builder (with result type
     // deduction ability) exists. We go through the separate-parameter builder
     // here given that it's easier for developers to write compared to
     // aggregate-parameter builders.
     createSeparateLocalVarsForOpArgs(tree, childNodeNames);
-    os.indent(6) << formatv("{0} = rewriter.create<{1}>(loc", valuePackName,
-                            resultOp.getQualCppClassName());
+
+    os.indent(6) << formatv("{0} = rewriter.create<{1}>({2}", valuePackName,
+                            resultOp.getQualCppClassName(), locToUse);
     supplyValuesForOpArgs(tree, childNodeNames);
     os << "\n      );\n";
     os.indent(4) << "}\n";
     return resultValue;
   }
 
   // If depth == 0 and resultIndex >= 0, it means we are replacing the values
   // generated from the source pattern root op. Then we can use the source
os.indent(6) << formatv(
       "SmallVector<Type, 4> tblgen_types; (void)tblgen_types;\n");
   int numResults = resultOp.getNumResults();
   if (numResults != 0) {
     for (int i = 0; i < numResults; ++i)
       os.indent(6) << formatv("for (auto v : castedOp0.getODSResults({0})) {{"
                               "tblgen_types.push_back(v.getType()); }\n",
                               resultIndex + i);
   }
-  os.indent(6) << formatv("{0} = rewriter.create<{1}>(loc, tblgen_types, "
+  os.indent(6) << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, "
                           "tblgen_values, tblgen_attrs);\n",
-                          valuePackName, resultOp.getQualCppClassName());
+                          valuePackName, resultOp.getQualCppClassName(),
+                          locToUse);
   os.indent(4) << "}\n";
   return resultValue;
 }
 
 void PatternEmitter::createSeparateLocalVarsForOpArgs(
     DagNode node, ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);